Add Gaussian beam example

docs/examples/wavefront/Gaussian_beams.ipynb

@@ -0,0 +1,196 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "69b8ec98-574e-49d2-b9b7-81e81c051bd3",
+   "metadata": {},
+   "source": [
+    "# Gaussian beams"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8f647cea-891e-41b6-9e2e-57fd6f5c07e7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2de2c5c9-8fa2-4c7c-bd71-c66560c0f357",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pmd_beamphysics.wavefront.wavefront import Wavefront\n",
+    "\n",
+    "from pmd_beamphysics.wavefront.gaussian import add_gaussian\n",
+    "\n",
+    "from pmd_beamphysics.wavefront.propagators import drift_wavefront\n",
+    "\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "25ad7285-980d-4526-a20a-04aad44beeb0",
+   "metadata": {},
+   "source": [
+    "## Beam at a waist (z=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f95ed503-75c1-4890-bfb3-23e8091171f0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W0 = Wavefront(\n",
+    "    Ex=np.zeros((101, 101, 51)),\n",
+    "    dx=10e-6,\n",
+    "    dy=10e-6,\n",
+    "    dz=10e-6,\n",
+    "    wavelength=1e-9,\n",
+    ")\n",
+    "\n",
+    "W1 = W0.copy()\n",
+    "add_gaussian(W1, z=0, w0=100e-6, energy=1.2345, sigma_z=50e-6)\n",
+    "\n",
+    "W1.plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "929bbfef-7083-45dc-9955-024f77724031",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W1k = W1.to_kspace()\n",
+    "W1k.plot()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "82e435ce-ef26-474e-8b7d-e5c29b85bc65",
+   "metadata": {},
+   "source": [
+    "## Alternatively defining zR"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0b62c1e5-aff5-42a6-8191-9fc17249a780",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zR = 10\n",
+    "\n",
+    "W2 = W0.copy()\n",
+    "add_gaussian(W2, z=0, zR=zR, energy=1.2345)\n",
+    "\n",
+    "# Check size\n",
+    "W2.sigma_x"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "15e2efec-2a81-4ee7-b449-5385d9ba5372",
+   "metadata": {},
+   "source": [
+    "## Beam at z=zR"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "487468e1-4b9a-451f-b419-1f8688f3145b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W3 = W0.copy()\n",
+    "zR = 10\n",
+    "add_gaussian(W3, z=zR, zR=zR, energy=1.2345)\n",
+    "\n",
+    "# Check size\n",
+    "W3.sigma_x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6064795c-1bf2-4375-b83e-16ceef7c3926",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Check  w(zR) = sqrt(2) * w0\n",
+    "assert np.isclose(W3.sigma_x, np.sqrt(2) * W2.sigma_x)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "099f0c7c-c689-44d6-86db-3d4c67477bb4",
+   "metadata": {},
+   "source": [
+    "## Drift "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8960bb27-31d3-42c3-93b5-4adc0a492ac2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W4 = drift_wavefront(W2, zR)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b38c0b53-282c-4689-9b16-290f92812b4a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert np.isclose(W4.sigma_x, W3.sigma_x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c7e04f33-6aba-4308-8600-aac29b9cc12d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W4.plot()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}