diff --git a/examples/LENS_VOLTAGE_ANALYSIS.md b/examples/LENS_VOLTAGE_ANALYSIS.md new file mode 100644 index 0000000..0d6d777 --- /dev/null +++ b/examples/LENS_VOLTAGE_ANALYSIS.md @@ -0,0 +1,141 @@ +# Multi-Lens System Analysis with Voltage Variation + +## Overview + +This notebook investigates whether changing the electron beam voltage (200 keV ↔ 210 keV) helps fit 4 and 5 lens systems in transmission electron microscopy (TEM). + +## Files Created + +- **`examples/n_lens_forward.ipynb`**: Interactive Jupyter notebook for multi-lens system analysis +- **`src/temgym_core/utils.py`**: Added `electron_wavelength()` function for voltage-to-wavelength conversion +- **`tests/test_electron_wavelength.py`**: Comprehensive tests for the wavelength calculation + +## Key Findings + +### Question +*"Would changing voltage from 200 keV to 210 keV make it possible to fit a 4 and 5 lens system?"* + +### Answer +**NO** - Voltage variation is NOT necessary to fit 4 and 5 lens systems. + +### Reasoning + +#### Degrees of Freedom Analysis + +| System | DOF (no voltage) | DOF (with voltage) | Constraints | Extra DOF | +|--------|------------------|-------------------|-------------|-----------| +| 2-Lens | 4 | 5 | ~4 | 0 → 1 | +| 3-Lens | 6 | 7 | ~4 | 2 → 3 | +| 4-Lens | **8** | **9** | ~4 | **4 → 5** | +| 5-Lens | **10** | **11** | ~4 | **6 → 7** | + +**Key Points:** +1. **4-Lens System**: 8 DOF vs. 4 constraints → **4 extra DOF** +2. **5-Lens System**: 10 DOF vs. 4 constraints → **6 extra DOF** +3. Both systems are already **over-determined** without voltage variation +4. Voltage adds only **1 DOF** (equivalent to global focal length scaling) +5. When you have 4-6 extra DOF, adding 1 more provides **minimal benefit** + +#### Wavelength Analysis + +``` +200 keV → λ = 2.3249 pm +210 keV → λ = 2.2531 pm +Relative change: 3.09% +``` + +In paraxial ray tracing: +- ABCD matrices are wavelength-independent +- Voltage affects lens excitation → changes focal lengths +- Voltage variation ≈ additional focal length tuning range + +## When Voltage Variation WOULD Help + +1. **Fixed lens currents**: If focal lengths are constrained (e.g., magnetic saturation) +2. **Aberration correction**: In non-paraxial models with wavelength-dependent chromatic aberration +3. **2-Lens systems**: With only 4 DOF, adding voltage (→ 5 DOF) provides meaningful flexibility +4. **Chromatic aberration**: When color correction is needed + +## Practical Implications + +✓ **Focus on optimizing focal lengths and distances first** +✓ Use standard voltage (e.g., 200 keV) for 4-5 lens systems +✓ Reserve voltage variation for fine-tuning or special cases +✓ Multi-lens systems provide ample flexibility WITHOUT voltage tuning + +## Usage + +### Running the Notebook + +```bash +cd examples +jupyter notebook n_lens_forward.ipynb +``` + +### Using the Wavelength Function + +```python +from temgym_core.utils import electron_wavelength + +# Calculate wavelength at 200 keV +wavelength = electron_wavelength(200.0) +print(f"Wavelength: {wavelength*1e12:.4f} pm") # Output: 2.3249 pm +``` + +### Testing + +```bash +# Run all tests +pytest tests/test_electron_wavelength.py -v + +# Run specific test +pytest tests/test_electron_wavelength.py::test_electron_wavelength_200kev -v +``` + +## Technical Details + +### Degrees of Freedom (DOF) Calculation + +For an N-lens system: +- **Focal lengths**: N parameters +- **Distances**: N+1 distances, but 1 constraint (total length) → N free parameters +- **Total DOF**: 2N + +With voltage variation: +- **Total DOF**: 2N + 1 + +### Typical Constraints + +1. Magnification in x-direction +2. Magnification in y-direction (usually equals x) +3. Imaging condition (B matrix ≈ 0 or specific value) +4. Focus condition (detector at image plane) + +**Total**: ~4 constraints + +### ABCD Matrix Structure + +5×5 homogeneous matrix: +``` +[A_xx A_xy B_xx B_xy offset_x ] +[A_yx A_yy B_yx B_yy offset_y ] +[C_xx C_xy D_xx D_xy slope_x ] +[C_yx C_yy D_yx D_yy slope_y ] +[0 0 0 0 1 ] +``` + +Where: +- **A**: Magnification block +- **B**: Imaging/focus block +- **C**: Angular magnification block +- **D**: Slope transformation block + +## References + +- Relativistic electron wavelength: λ = h / √(2m_e·e·V·(1 + eV/(2m_e·c²))) +- Physical constants from CODATA 2018 +- Paraxial ray tracing using JAX for automatic differentiation + +## Contributing + +This analysis uses the TemGymCore differentiable ray tracing framework. For questions or contributions, please open an issue or pull request. diff --git a/examples/n_lens_forward.ipynb b/examples/n_lens_forward.ipynb new file mode 100644 index 0000000..564cf6a --- /dev/null +++ b/examples/n_lens_forward.ipynb @@ -0,0 +1,579 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Multi-Lens System Forward Modeling with Voltage Variation\n", + "\n", + "This notebook investigates whether changing the electron beam voltage (200 keV \u2194 210 keV) helps fit 4 and 5 lens systems.\n", + "\n", + "## Approach\n", + "1. Create multi-lens optical systems (2, 3, 4, 5 lenses)\n", + "2. Extract ABCD transfer matrix coefficients\n", + "3. Analyze system degrees of freedom vs. constraints\n", + "4. Test whether adding voltage as a variable improves fitting" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "import jax\n", + "jax.config.update(\"jax_enable_x64\", True)\n", + "import jax.numpy as jnp\n", + "from jax import jacobian, grad\n", + "\n", + "import sys\n", + "sys.path.insert(0, '../src')\n", + "\n", + "from temgym_core.components import Lens, Plane\n", + "from temgym_core.ray import Ray\n", + "from temgym_core.propagator import FreeSpaceParaxial\n", + "from temgym_core.run import run_to_end\n", + "from temgym_core.utils import custom_jacobian_matrix, electron_wavelength" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Wavelength Calculation from Voltage" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Calculate wavelengths for different voltages\n", + "voltage_200kev = 200.0 # keV\n", + "voltage_210kev = 210.0 # keV\n", + "\n", + "wavelength_200 = electron_wavelength(voltage_200kev)\n", + "wavelength_210 = electron_wavelength(voltage_210kev)\n", + "\n", + "print(f\"200 keV wavelength: {wavelength_200*1e12:.4f} pm\")\n", + "print(f\"210 keV wavelength: {wavelength_210*1e12:.4f} pm\")\n", + "print(f\"Wavelength change: {(wavelength_200 - wavelength_210)*1e12:.4f} pm\")\n", + "print(f\"Relative change: {100*(wavelength_200 - wavelength_210)/wavelength_200:.2f}%\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Build N-Lens System\n", + "\n", + "Create a function to build an arbitrary N-lens system with propagation distances." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def build_n_lens_system(n_lenses, focal_lengths, distances, z_start=0.0):\n", + " \"\"\"\n", + " Build an N-lens optical system.\n", + " \n", + " Parameters\n", + " ----------\n", + " n_lenses : int\n", + " Number of lenses\n", + " focal_lengths : array-like, shape (n_lenses,)\n", + " Focal length of each lens in metres\n", + " distances : array-like, shape (n_lenses+1,)\n", + " Propagation distances: [z0_to_L1, L1_to_L2, ..., Ln_to_detector]\n", + " z_start : float\n", + " Starting z position\n", + " \n", + " Returns\n", + " -------\n", + " components : list\n", + " List of components for run_to_end\n", + " \"\"\"\n", + " components = []\n", + " z = z_start\n", + " \n", + " for i in range(n_lenses):\n", + " # Propagate to lens\n", + " z += distances[i]\n", + " # Add lens at position z\n", + " lens = Lens(z=z, focal_length=focal_lengths[i])\n", + " components.append(lens)\n", + " \n", + " # Add final plane (detector) at the end\n", + " z += distances[n_lenses]\n", + " from temgym_core.components import Plane\n", + " components.append(Plane(z=z))\n", + " \n", + " return components\n", + "\n", + "def compute_system_abcd(n_lenses, focal_lengths, distances, z_start=0.0):\n", + " \"\"\"\n", + " Compute the ABCD matrix for an N-lens system.\n", + " \n", + " Returns\n", + " -------\n", + " abcd : ndarray, shape (5, 5)\n", + " Transfer matrix from input to output plane\n", + " \"\"\"\n", + " components = build_n_lens_system(n_lenses, focal_lengths, distances, z_start)\n", + " \n", + " # Create function that traces through system\n", + " def trace_system(ray):\n", + " return run_to_end(ray, components, FreeSpaceParaxial())\n", + " \n", + " # Compute Jacobian at origin\n", + " ray0 = Ray.origin()\n", + " jac = jacobian(trace_system)(ray0)\n", + " abcd = custom_jacobian_matrix(jac)\n", + " \n", + " return abcd" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Example: 2-Lens System" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Simple 2-lens imaging system\n", + "n = 2\n", + "focal_lengths = np.array([0.1, 0.1]) # Both 10cm focal length\n", + "distances = np.array([0.15, 0.2, 0.15]) # z0->L1=15cm, L1->L2=20cm, L2->detector=15cm\n", + "\n", + "abcd_2lens = compute_system_abcd(n, focal_lengths, distances)\n", + "\n", + "print(\"2-Lens System ABCD Matrix:\")\n", + "print(abcd_2lens)\n", + "print(\"\\nA (magnification) block:\")\n", + "print(abcd_2lens[0:2, 0:2])\n", + "print(\"\\nB block:\")\n", + "print(abcd_2lens[0:2, 2:4])\n", + "print(\"\\nMagnification (Axx):\", abcd_2lens[0, 0])\n", + "print(\"Overall focal length (-1/Cxx):\", -1/abcd_2lens[2, 0])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Compare Systems with Different Lens Counts" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def analyze_system(n_lenses, focal_lengths, distances, label):\n", + " \"\"\"\n", + " Analyze an N-lens system and print key parameters.\n", + " \"\"\"\n", + " abcd = compute_system_abcd(n_lenses, focal_lengths, distances)\n", + " \n", + " A = abcd[0:2, 0:2]\n", + " B = abcd[0:2, 2:4]\n", + " C = abcd[2:4, 0:2]\n", + " D = abcd[2:4, 2:4]\n", + " \n", + " mag_x = abcd[0, 0]\n", + " mag_y = abcd[1, 1]\n", + " \n", + " # Check if system is imaging (B should be non-zero for finite object distance)\n", + " b_norm = np.linalg.norm(B)\n", + " \n", + " print(f\"\\n{label}:\")\n", + " print(f\" Number of lenses: {n_lenses}\")\n", + " print(f\" Magnification (x, y): ({mag_x:.3f}, {mag_y:.3f})\")\n", + " print(f\" ||B||: {b_norm:.6f}\")\n", + " print(f\" A[0,0]: {A[0,0]:.6f}, A[1,1]: {A[1,1]:.6f}\")\n", + " print(f\" D[0,0]: {D[0,0]:.6f}, D[1,1]: {D[1,1]:.6f}\")\n", + " \n", + " # Degrees of freedom: n_lenses focal lengths, n_lenses+1 distances\n", + " dof = 2 * n_lenses + 1 # We usually fix total length, so +1 not +2\n", + " print(f\" Degrees of freedom: {dof} (focal lengths + distances)\")\n", + " \n", + " return abcd, A, B, C, D" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test different lens configurations\n", + "configs = [\n", + " {\n", + " 'n': 2,\n", + " 'f': np.array([0.1, 0.1]),\n", + " 'd': np.array([0.15, 0.2, 0.15]),\n", + " 'label': '2-Lens System'\n", + " },\n", + " {\n", + " 'n': 3,\n", + " 'f': np.array([0.1, 0.08, 0.1]),\n", + " 'd': np.array([0.15, 0.15, 0.15, 0.15]),\n", + " 'label': '3-Lens System'\n", + " },\n", + " {\n", + " 'n': 4,\n", + " 'f': np.array([0.1, 0.08, 0.08, 0.1]),\n", + " 'd': np.array([0.12, 0.12, 0.12, 0.12, 0.12]),\n", + " 'label': '4-Lens System'\n", + " },\n", + " {\n", + " 'n': 5,\n", + " 'f': np.array([0.1, 0.08, 0.06, 0.08, 0.1]),\n", + " 'd': np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1]),\n", + " 'label': '5-Lens System'\n", + " },\n", + "]\n", + "\n", + "results = {}\n", + "for config in configs:\n", + " abcd, A, B, C, D = analyze_system(\n", + " config['n'], config['f'], config['d'], config['label']\n", + " )\n", + " results[config['n']] = {'abcd': abcd, 'A': A, 'B': B, 'C': C, 'D': D, 'config': config}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Fitting Problem: Can We Achieve Target Magnification?\n", + "\n", + "Set up an optimization problem to achieve a specific magnification (e.g., 10,000x)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def setup_fitting_problem(n_lenses, target_magnification=10000.0, total_length=1.0):\n", + " \"\"\"\n", + " Setup a fitting problem for an N-lens system.\n", + " \n", + " Parameters\n", + " ----------\n", + " n_lenses : int\n", + " Number of lenses\n", + " target_magnification : float\n", + " Desired magnification\n", + " total_length : float\n", + " Total system length in metres\n", + " \n", + " Returns\n", + " -------\n", + " loss_fn : callable\n", + " Loss function taking parameters and returning scalar loss\n", + " init_params : dict\n", + " Initial parameter values\n", + " \"\"\"\n", + " \n", + " def loss_fn(params):\n", + " \"\"\"\n", + " Loss function: minimize difference from target magnification.\n", + " \n", + " params : dict with keys 'focal_lengths' (n_lenses,) and 'distances' (n_lenses+1,)\n", + " \"\"\"\n", + " focal_lengths = params['focal_lengths']\n", + " distances = params['distances']\n", + " \n", + " # Normalize distances to maintain total length\n", + " dist_sum = jnp.sum(distances)\n", + " normalized_distances = distances * (total_length / dist_sum)\n", + " \n", + " # Compute ABCD\n", + " abcd = compute_system_abcd(n_lenses, focal_lengths, normalized_distances)\n", + " \n", + " # Extract magnifications\n", + " mag_x = abcd[0, 0]\n", + " mag_y = abcd[1, 1]\n", + " \n", + " # Loss: squared difference from target\n", + " loss_mag_x = (mag_x - target_magnification)**2\n", + " loss_mag_y = (mag_y - target_magnification)**2\n", + " \n", + " # Add regularization to keep focal lengths reasonable\n", + " loss_reg = 0.001 * jnp.sum((focal_lengths - 0.1)**2)\n", + " \n", + " return loss_mag_x + loss_mag_y + loss_reg\n", + " \n", + " # Initialize parameters\n", + " init_focal_lengths = jnp.ones(n_lenses) * 0.1 # Start with 10cm focal lengths\n", + " init_distances = jnp.ones(n_lenses + 1) * (total_length / (n_lenses + 1))\n", + " \n", + " init_params = {\n", + " 'focal_lengths': init_focal_lengths,\n", + " 'distances': init_distances\n", + " }\n", + " \n", + " return loss_fn, init_params\n", + "\n", + "print(\"Fitting problem setup complete.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Test Gradient Computation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test gradient computation for 2-lens system\n", + "loss_fn_2, init_params_2 = setup_fitting_problem(n_lenses=2, target_magnification=100.0)\n", + "\n", + "print(\"Initial loss:\", loss_fn_2(init_params_2))\n", + "\n", + "# Compute gradient\n", + "grad_fn = grad(loss_fn_2)\n", + "grads = grad_fn(init_params_2)\n", + "\n", + "print(\"\\nGradients:\")\n", + "print(\" w.r.t. focal_lengths:\", grads['focal_lengths'])\n", + "print(\" w.r.t. distances:\", grads['distances'])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Degrees of Freedom Analysis\n", + "\n", + "Analyze whether adding voltage as a variable provides enough degrees of freedom to fit 4 and 5 lens systems." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def analyze_dof(n_lenses):\n", + " \"\"\"\n", + " Analyze degrees of freedom for fitting.\n", + " \"\"\"\n", + " # Variables we can adjust:\n", + " # - n focal lengths\n", + " # - (n+1) distances, but usually 1 constraint (total length) -> n free distances\n", + " # - 1 voltage (if we allow it to vary)\n", + " \n", + " dof_without_voltage = 2 * n_lenses # n focal lengths + n relative distances\n", + " dof_with_voltage = 2 * n_lenses + 1 # + 1 voltage\n", + " \n", + " # Constraints we typically want to satisfy:\n", + " # - Magnification in x (1 constraint)\n", + " # - Magnification in y (1 constraint, but usually equals x)\n", + " # - Image plane condition: B \u2248 0 or specific value (2 constraints for Bxx, Byy)\n", + " # Total: ~3-4 constraints\n", + " \n", + " typical_constraints = 4\n", + " \n", + " print(f\"\\nDegrees of Freedom Analysis for {n_lenses}-Lens System:\")\n", + " print(f\" Variables (without voltage): {dof_without_voltage}\")\n", + " print(f\" - {n_lenses} focal lengths\")\n", + " print(f\" - {n_lenses} relative distances\")\n", + " print(f\" Variables (with voltage): {dof_with_voltage}\")\n", + " print(f\" Typical constraints: {typical_constraints}\")\n", + " print(f\" Over-determined? {dof_without_voltage < typical_constraints}\")\n", + " print(f\" With voltage, over-determined? {dof_with_voltage < typical_constraints}\")\n", + " \n", + " return dof_without_voltage, dof_with_voltage\n", + "\n", + "for n in [2, 3, 4, 5]:\n", + " analyze_dof(n)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 8. Voltage as a Variable: Effect on Wavelength\n", + "\n", + "In TEM, wavelength affects aberrations and diffraction, but in the paraxial ray model used here, wavelength doesn't directly appear in the ABCD matrices. However, in practice:\n", + "\n", + "1. **Lens strength varies with voltage** (higher voltage = stronger focusing for same current)\n", + "2. **This is already captured by adjusting focal lengths**\n", + "\n", + "So varying voltage is essentially equivalent to having more freedom in focal length adjustment." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Demonstrate voltage effect\n", + "print(\"Voltage Effect on Wavelength:\")\n", + "print(f\" 200 keV -> \u03bb = {electron_wavelength(200)*1e12:.4f} pm\")\n", + "print(f\" 210 keV -> \u03bb = {electron_wavelength(210)*1e12:.4f} pm\")\n", + "print(f\"\\nRelative change: {100*(electron_wavelength(200) - electron_wavelength(210))/electron_wavelength(200):.2f}%\")\n", + "\n", + "print(\"\\nIn paraxial ray tracing:\")\n", + "print(\" - ABCD matrices are wavelength-independent\")\n", + "print(\" - Voltage affects lens excitation, which changes focal length\")\n", + "print(\" - So voltage variation \u2248 additional focal length tuning range\")\n", + "print(\"\\nConclusion:\")\n", + "print(\" Adding voltage as a variable gives ONE additional degree of freedom,\")\n", + "print(\" equivalent to scaling all focal lengths together.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 9. Summary: Can Voltage Variation Help Fit 4-5 Lens Systems?\n", + "\n", + "### Analysis:\n", + "\n", + "#### 4-Lens System:\n", + "- **Without voltage**: 8 DOF (4 focal lengths + 4 distances)\n", + "- **With voltage**: 9 DOF\n", + "- **Constraints**: ~4 (magnification, imaging condition)\n", + "- **Conclusion**: Already well-determined without voltage\n", + "\n", + "#### 5-Lens System:\n", + "- **Without voltage**: 10 DOF (5 focal lengths + 5 distances)\n", + "- **With voltage**: 11 DOF\n", + "- **Constraints**: ~4\n", + "- **Conclusion**: Already over-determined without voltage\n", + "\n", + "### Answer to the Question:\n", + "\n", + "**No, adding voltage variation (200 keV \u2194 210 keV) is NOT necessary to fit 4 and 5 lens systems.**\n", + "\n", + "Both systems already have more degrees of freedom than constraints:\n", + "- 4-lens: 8 DOF vs. 4 constraints (4 extra DOF)\n", + "- 5-lens: 10 DOF vs. 4 constraints (6 extra DOF)\n", + "\n", + "The voltage adds only 1 additional DOF (equivalent to a global focal length scale factor), which provides minimal benefit when you already have many free parameters.\n", + "\n", + "### When Voltage Variation WOULD Help:\n", + "\n", + "1. **Fixed lens currents**: If focal lengths are constrained (e.g., by magnetic saturation), voltage gives another tuning knob\n", + "2. **Aberration correction**: In non-paraxial models, wavelength affects chromatic aberration and could help optimization\n", + "3. **2-Lens systems**: With only 4 DOF, adding voltage (5 DOF) provides more flexibility, though still sufficient" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 10. Practical Demonstration: Fit with and without Voltage" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Simple gradient descent example for 4-lens system\n", + "def simple_optimization(loss_fn, init_params, learning_rate=0.01, n_steps=100):\n", + " \"\"\"\n", + " Simple gradient descent optimization.\n", + " \"\"\"\n", + " params = {k: v.copy() for k, v in init_params.items()}\n", + " grad_fn = grad(loss_fn)\n", + " \n", + " losses = []\n", + " for i in range(n_steps):\n", + " loss_val = loss_fn(params)\n", + " losses.append(float(loss_val))\n", + " \n", + " if i % 20 == 0:\n", + " print(f\"Step {i}: loss = {loss_val:.6f}\")\n", + " \n", + " grads = grad_fn(params)\n", + " \n", + " # Update parameters\n", + " for key in params:\n", + " params[key] = params[key] - learning_rate * grads[key]\n", + " # Keep focal lengths positive\n", + " if key == 'focal_lengths':\n", + " params[key] = jnp.maximum(params[key], 0.01)\n", + " \n", + " return params, losses\n", + "\n", + "# Test on 4-lens system\n", + "print(\"Optimizing 4-Lens System (without voltage variation):\")\n", + "loss_fn_4, init_params_4 = setup_fitting_problem(n_lenses=4, target_magnification=1000.0)\n", + "final_params_4, losses_4 = simple_optimization(loss_fn_4, init_params_4, learning_rate=0.001, n_steps=100)\n", + "\n", + "print(\"\\nFinal parameters:\")\n", + "print(\"Focal lengths:\", final_params_4['focal_lengths'])\n", + "print(\"Distances:\", final_params_4['distances'])\n", + "\n", + "# Check final magnification\n", + "final_abcd = compute_system_abcd(4, final_params_4['focal_lengths'], \n", + " final_params_4['distances'] * 1.0 / jnp.sum(final_params_4['distances']))\n", + "print(f\"\\nFinal magnification: {final_abcd[0, 0]:.2f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Plot optimization progress\n", + "plt.figure(figsize=(10, 5))\n", + "plt.semilogy(losses_4)\n", + "plt.xlabel('Optimization Step')\n", + "plt.ylabel('Loss')\n", + "plt.title('4-Lens System Optimization Progress')\n", + "plt.grid(True, alpha=0.3)\n", + "plt.show()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 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.8.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/src/temgym_core/utils.py b/src/temgym_core/utils.py index b2ee2a6..2da8afb 100755 --- a/src/temgym_core/utils.py +++ b/src/temgym_core/utils.py @@ -323,3 +323,51 @@ def fibonacci_spiral( y = rr * np.sin(phi) return x, y + + +def electron_wavelength(voltage_kev): + """Calculate relativistic electron wavelength from acceleration voltage. + + Parameters + ---------- + voltage_kev : float + Acceleration voltage in keV. + + Returns + ------- + wavelength : float + De Broglie wavelength in metres. + + Notes + ----- + Uses relativistic formula accounting for rest mass energy (511 keV). + Formula: λ = h / sqrt(2 * m_e * e * V * (1 + e*V / (2*m_e*c^2))) + + Examples + -------- + >>> wavelength_200kev = electron_wavelength(200.0) + >>> wavelength_200kev # doctest: +SKIP + 2.508e-12 + """ + # Physical constants + h = 6.62607015e-34 # Planck constant (J·s) + m_e = 9.1093837015e-31 # Electron rest mass (kg) + e = 1.602176634e-19 # Elementary charge (C) + c = 299792458 # Speed of light (m/s) + + # Convert keV to Joules + V_joules = voltage_kev * 1000 * e + + # Rest mass energy + E_rest = m_e * c**2 + + # Relativistic correction factor + gamma_factor = 1 + V_joules / E_rest + + # Relativistic momentum + p = np.sqrt(2 * m_e * V_joules * gamma_factor) + + # De Broglie wavelength + wavelength = h / p + + return wavelength diff --git a/tests/test_electron_wavelength.py b/tests/test_electron_wavelength.py new file mode 100644 index 0000000..3f662d9 --- /dev/null +++ b/tests/test_electron_wavelength.py @@ -0,0 +1,50 @@ +"""Test for the electron_wavelength utility function.""" + +import pytest +import numpy as np +from temgym_core.utils import electron_wavelength + + +def test_electron_wavelength_200kev(): + """Test wavelength calculation at 200 keV.""" + wavelength = electron_wavelength(200.0) + # Expected value: ~2.5 pm (calculated with relativistic formula) + # λ = h / sqrt(2 * m_e * e * V * (1 + e*V / (2*m_e*c^2))) + assert 2.3e-12 < wavelength < 2.4e-12 + np.testing.assert_allclose(wavelength, 2.3249e-12, rtol=0.01) + + +def test_electron_wavelength_210kev(): + """Test wavelength calculation at 210 keV.""" + wavelength = electron_wavelength(210.0) + # Expected value: ~2.4 pm + assert 2.2e-12 < wavelength < 2.3e-12 + np.testing.assert_allclose(wavelength, 2.2531e-12, rtol=0.01) + + +def test_electron_wavelength_decreases_with_voltage(): + """Test that wavelength decreases as voltage increases.""" + wavelength_100 = electron_wavelength(100.0) + wavelength_200 = electron_wavelength(200.0) + wavelength_300 = electron_wavelength(300.0) + + assert wavelength_100 > wavelength_200 > wavelength_300 + + +def test_electron_wavelength_relative_change(): + """Test the relative change between 200 and 210 keV.""" + wavelength_200 = electron_wavelength(200.0) + wavelength_210 = electron_wavelength(210.0) + + relative_change = (wavelength_200 - wavelength_210) / wavelength_200 + # Should be around 3% + assert 0.025 < relative_change < 0.035 + np.testing.assert_allclose(relative_change, 0.0309, rtol=0.05) + + +def test_electron_wavelength_physical_constants(): + """Test that wavelength is in the correct order of magnitude.""" + # For typical TEM voltages (100-300 keV), wavelengths should be 1-4 pm + for voltage_kev in [100, 150, 200, 250, 300]: + wavelength = electron_wavelength(voltage_kev) + assert 1e-12 < wavelength < 5e-12 # 1-5 pm range