TemGym · Copilot · Sep 18, 2025 · Sep 19, 2025 · Sep 20, 2025 · Sep 20, 2025
diff --git a/archived/aberrated_probe.ipynb b/archived/aberrated_probe.ipynb
diff --git a/archived/aberrated_probe_further_testing.ipynb b/archived/aberrated_probe_further_testing.ipynb
diff --git a/archived/aberrations_visualisation.ipynb b/archived/aberrations_visualisation.ipynb
diff --git a/archived/aperture_image.ipynb b/archived/aperture_image.ipynb
diff --git a/archived/approximating_functions_with_gaussians_example.ipynb b/archived/approximating_functions_with_gaussians_example.ipynb
diff --git a/archived/fit_three_gaussians.ipynb b/archived/fit_three_gaussians.ipynb
diff --git a/archived/local_gaussians_hessian.ipynb b/archived/local_gaussians_hessian.ipynb
diff --git a/archived/local_gaussians_parabola.ipynb b/archived/local_gaussians_parabola.ipynb
diff --git a/examples/aberrated_probe.ipynb b/examples/aberrated_probe.ipynb
diff --git a/examples/aperture_diffraction.ipynb b/examples/aperture_diffraction.ipynb
diff --git a/examples/aperture_image.ipynb b/examples/aperture_image.ipynb
diff --git a/examples/atoms/abtem_HRTEM_example.ipynb b/examples/atoms/abtem_HRTEM_example.ipynb
diff --git a/examples/atoms/glauber_ray_example.ipynb b/examples/atoms/glauber_ray_example.ipynb
diff --git a/examples/atoms/multiple_rays_single_atom.ipynb b/examples/atoms/multiple_rays_single_atom.ipynb
diff --git a/examples/atoms/single_ray_single_atom_baseline.ipynb b/examples/atoms/single_ray_single_atom_baseline.ipynb
diff --git a/examples/biprism.ipynb b/examples/biprism.ipynb
diff --git a/examples/decompose_image_into_gaussians.ipynb b/examples/decompose_image_into_gaussians.ipynb
diff --git a/examples/fit_three_gaussians.ipynb b/examples/fit_three_gaussians.ipynb
diff --git a/examples/gaussians/aberrated_aperture_probe.ipynb b/examples/gaussians/aberrated_aperture_probe.ipynb
diff --git a/examples/gaussians/aberrated_krivanek_image.ipynb b/examples/gaussians/aberrated_krivanek_image.ipynb
diff --git a/examples/gaussians/aberrated_krivanek_lens.ipynb b/examples/gaussians/aberrated_krivanek_lens.ipynb
diff --git a/examples/gaussians/aperture.ipynb b/examples/gaussians/aperture.ipynb
diff --git a/examples/gaussians/aperture_probe.ipynb b/examples/gaussians/aperture_probe.ipynb
diff --git a/examples/gaussians/biprism.ipynb b/examples/gaussians/biprism.ipynb
diff --git a/examples/gaussians/discontinuous_phase_error_demo.ipynb b/examples/gaussians/discontinuous_phase_error_demo.ipynb
diff --git a/examples/gaussians/distorted_projector_lens.ipynb b/examples/gaussians/distorted_projector_lens.ipynb
diff --git a/examples/gaussians/lens_biprism.ipynb b/examples/gaussians/lens_biprism.ipynb
diff --git a/examples/gaussians/magnetic_phase_sample.ipynb b/examples/gaussians/magnetic_phase_sample.ipynb
diff --git a/examples/gaussians/sigmoid_aperture.ipynb b/examples/gaussians/sigmoid_aperture.ipynb
diff --git a/examples/gaussians/single_gaussian.ipynb b/examples/gaussians/single_gaussian.ipynb
@@ -0,0 +1,167 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d5e74cc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ['JAX_ENABLE_X64'] = '1'\n",
+    "os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.1' \n",
+    "\n",
+    "%matplotlib widget\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "\n",
+    "import numpy as np\n",
+    "from temgym_core.components import Detector\n",
+    "from temgym_core.gaussian import (\n",
+    "    Lens,\n",
+    "    run_to_end,\n",
+    "    make_gaussian,\n",
+    ")\n",
+    "from temgym_core.constants import energy2wavelength\n",
+    "from temgym_core.evaluate import evaluate_gaussians_gpu_kernel_wrapper, evaluate_gaussians_for, evaluate_gaussians_jax_scan\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "38bc5771",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "jax.config.update(\"jax_enable_x64\", True)\n",
+    "W = 100e-9\n",
+    "\n",
+    "Nx = Ny = 512\n",
+    "dx = W/Nx\n",
+    "dy = W/Ny\n",
+    "input_grid = Detector(z=0.0, pixel_size=(dx, dy), shape=(Nx, Ny))\n",
+    "input_extent = input_grid.extent\n",
+    "coords = input_grid.coords\n",
+    "X, Y = coords[:, 0].reshape(input_grid.shape), coords[:, 1].reshape(input_grid.shape)\n",
+    "\n",
+    "distance = 1e-3\n",
+    "output_grid = Detector(z=distance, pixel_size=(dx * 4, dy * 4), shape=(Nx, Ny))\n",
+    "output_extent = output_grid.extent\n",
+    "\n",
+    "voltage = 100e3\n",
+    "wavelength = energy2wavelength(voltage)  # 200 keV\n",
+    "k0 = 2 * jnp.pi / wavelength\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "43c6bbf2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "w0 = 1e-8  # beam waist\n",
+    "\n",
+    "rays_in = make_gaussian(waist_x=w0, waist_y=w0, voltage=voltage)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "73963965",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lens = Lens(focal_length=1e-5, z=0.0)\n",
+    "rays_out = run_to_end(rays_in, [lens])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "773df429",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "TypeError",
+     "evalue": "evaluate_gaussians_reliable() missing 5 required positional arguments: 'amplitude', 'pathlength', 'Q_inv', 'k', and 'r2'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+      "\u001b[31mTypeError\u001b[39m                                 Traceback (most recent call last)",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[5]\u001b[39m\u001b[32m, line 2\u001b[39m\n\u001b[32m      1\u001b[39m E_in = evaluate_gaussians_for(rays_out, input_grid)\n\u001b[32m----> \u001b[39m\u001b[32m2\u001b[39m E_gpu = \u001b[43mevaluate_gaussians_reliable\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrays_out\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minput_grid\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m      3\u001b[39m E_scan = evaluate_gaussians_jax_scan(rays_out, input_grid)\n\u001b[32m      5\u001b[39m \u001b[38;5;66;03m# assert jnp.allclose(E_in, E_gpu), \"E_in and E_gpu are not the same\"\u001b[39;00m\n",
+      "\u001b[31mTypeError\u001b[39m: evaluate_gaussians_reliable() missing 5 required positional arguments: 'amplitude', 'pathlength', 'Q_inv', 'k', and 'r2'"
+     ]
+    }
+   ],
+   "source": [
+    "E_in = evaluate_gaussians_for(rays_out, input_grid)\n",
+    "E_gpu = evaluate_gaussians_gpu_kernel_wrapper(rays_out, input_grid)\n",
+    "E_scan = evaluate_gaussians_jax_scan(rays_out, input_grid)\n",
+    "\n",
+    "# assert jnp.allclose(E_in, E_gpu), \"E_in and E_gpu are not the same\"\n",
+    "assert jnp.allclose(E_in, E_scan), \"E_in and E_scan are not the same\"\n",
+    "# assert jnp.allclose(E_gpu, E_scan), \"E_gpu and E_scan are not the same\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "02ad1737",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axs = plt.subplots(1, 2, figsize=(12, 5))\n",
+    "\n",
+    "# Phase of gaussian beam\n",
+    "axs[0].imshow(np.angle(E_in), extent=input_extent, cmap='twilight')\n",
+    "axs[0].set_title('Phase')\n",
+    "axs[0].set_xlabel('x (m)')\n",
+    "axs[0].set_ylabel('y (m)')\n",
+    "axs[0].axis('equal')\n",
+    "\n",
+    "quadratic_phase = jnp.exp(-1j * k0 / (2 * lens.focal_length) * (X**2 + Y**2))\n",
+    "axs[1].imshow(np.angle(quadratic_phase), extent=input_extent, cmap='twilight')\n",
+    "axs[1].set_title('Phase')\n",
+    "axs[1].set_xlabel('x (m)')\n",
+    "axs[1].set_ylabel('y (m)')\n",
+    "axs[1].axis('equal')\n",
+    "\n",
+    "# Plot cross section\n",
+    "plt.figure()\n",
+    "plt.plot(np.angle(E_in[Nx//2,:]), label='Gaussian beam')\n",
+    "plt.plot(np.angle(quadratic_phase[Nx//2,:]), label='Quadratic phase')\n",
+    "plt.xlabel('x (nm)')\n",
+    "plt.ylabel('Phase (radians)')\n",
+    "plt.legend()\n",
+    "plt.title('Phase cross-section at y=0')\n",
+    "\n",
+    "plt.figure()\n",
+    "plt.imshow(np.abs(E_in), extent=input_extent, cmap='inferno')\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "temgym_core (3.12.3)",
+   "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/gaussians/two_gaussians_interfering.ipynb b/examples/gaussians/two_gaussians_interfering.ipynb
diff --git a/examples/holography_distortion/3D_sample_phase_shift.ipynb b/examples/holography_distortion/3D_sample_phase_shift.ipynb
diff --git a/examples/holography_distortion/amorphous_carbon.ipynb b/examples/holography_distortion/amorphous_carbon.ipynb