Adding Gaussian Beam summation

src/temgym_core/__init__.py

@@ -76,3 +76,11 @@ class PixelsYX(NamedTuple):
     """
     y: NDArray[np.integer]
     x: NDArray[np.integer]
+
+
+# Convenience re-exports
+try:
+    from .plotting import plot_model, PlotParams  # noqa: F401
+except Exception:
+    # Plotting has optional dependencies (matplotlib); ignore import errors at package import time
+    pass

src/temgym_core/aberrations.py

@@ -0,0 +1,108 @@
+import jax.numpy as jnp
+from dataclasses import dataclass
+
+
+@dataclass
+class KrivanekCoeffs:
+    C10: float = 0.0
+    C12: float = 0.0
+    phi12: float = 0.0
+    C21: float = 0.0
+    phi21: float = 0.0
+    C23: float = 0.0
+    phi23: float = 0.0
+    C30: float = 0.0
+    C32: float = 0.0
+    phi32: float = 0.0
+    C34: float = 0.0
+    phi34: float = 0.0
+    C41: float = 0.0
+    phi41: float = 0.0
+    C43: float = 0.0
+    phi43: float = 0.0
+    C45: float = 0.0
+    phi45: float = 0.0
+    C50: float = 0.0
+    C52: float = 0.0
+    phi52: float = 0.0
+    C54: float = 0.0
+    phi54: float = 0.0
+    C56: float = 0.0
+    phi56: float = 0.0
+
+
+def _cos_kriv(m, ph, ph0):
+    return jnp.cos(m * (ph - ph0))
+
+
+def _sin_kriv(m, ph, ph0):
+    return jnp.sin(m * (ph - ph0))
+
+
+def krivanek_coeff_brackets(phi, p: KrivanekCoeffs):
+    B2 = p.C10 + p.C12 * _cos_kriv(2, phi, p.phi12)
+    B3 = p.C21 * _cos_kriv(1, phi, p.phi21) + p.C23 * _cos_kriv(3, phi, p.phi23)
+    B4 = p.C30 + p.C32 * _cos_kriv(2, phi, p.phi32) + p.C34 * _cos_kriv(4, phi, p.phi34)
+    B5 = p.C41 * _cos_kriv(1, phi, p.phi41) + p.C43 * _cos_kriv(3, phi, p.phi43) + p.C45 * _cos_kriv(5, phi, p.phi45)  # noqa: E501
+    B6 = p.C50 + p.C52 * _cos_kriv(2, phi, p.phi52) + p.C54 * _cos_kriv(4, phi, p.phi54) + p.C56 * _cos_kriv(6, phi, p.phi56)  # noqa: E501
+    return B2, B3, B4, B5, B6
+
+
+def W_krivanek(alpha, phi, p):
+
+    B2, B3, B4, B5, B6 = krivanek_coeff_brackets(phi, p)
+    a = alpha
+    a2 = a * a
+    a3 = a2 * a
+    a4 = a2 * a2
+    a6 = a3 * a3
+
+    return 0.5 * a2 * B2 + (a3 / 3.0) * B3 + 0.25 * a4 * B4 + 0.2 * a4 * a * B5 + (a6 / 6.0) * B6
+
+
+def grad_W_krivanek(alpha_x, alpha_y, p):
+    ax, ay = alpha_x, alpha_y
+    alpha = jnp.hypot(ax, ay)
+    phi = jnp.arctan2(ay, ax)
+
+    B2, B3, B4, B5, B6 = krivanek_coeff_brackets(phi, p)
+
+    a = alpha
+    a2 = a * a
+    a3 = a2 * a
+    a4 = a2 * a2
+    a5 = a4 * a
+    a6 = a3 * a3
+
+    dW_dalpha = a * B2 + a2 * B3 + a3 * B4 + a4 * B5 + a5 * B6
+
+    dW_dphi = (0.5 * a2) * (
+        -2.0 * p.C12 * _sin_kriv(2, phi, p.phi12)
+    )
+    dW_dphi += (a3 / 3.0) * (
+        -1.0 * p.C21 * _sin_kriv(1, phi, p.phi21)
+        - 3.0 * p.C23 * _sin_kriv(3, phi, p.phi23)
+    )
+    dW_dphi += (0.25 * a4) * (
+        -2.0 * p.C32 * _sin_kriv(2, phi, p.phi32)
+        - 4.0 * p.C34 * _sin_kriv(4, phi, p.phi34)
+    )
+    dW_dphi += (0.2 * a4 * a) * (
+        -1.0 * p.C41 * _sin_kriv(1, phi, p.phi41)
+        - 3.0 * p.C43 * _sin_kriv(3, phi, p.phi43)
+        - 5.0 * p.C45 * _sin_kriv(5, phi, p.phi45)
+    )
+    dW_dphi += (a6 / 6.0) * (
+        -2.0 * p.C52 * _sin_kriv(2, phi, p.phi52)
+        - 4.0 * p.C54 * _sin_kriv(4, phi, p.phi54)
+        - 6.0 * p.C56 * _sin_kriv(6, phi, p.phi56)
+    )
+
+    eps = 1e-30
+    a_safe = jnp.where(a == 0, eps, a)
+    inv_a = 1.0 / a_safe
+    inv_a2 = inv_a * inv_a
+
+    dWx = dW_dalpha * (ax * inv_a) + dW_dphi * (-ay * inv_a2)
+    dWy = dW_dalpha * (ay * inv_a) + dW_dphi * (ax * inv_a2)
+    return dWx, dWy

src/temgym_core/components.py

@@ -1,12 +1,12 @@
-from typing import NamedTuple
+from typing import NamedTuple, Dict
 import jax_dataclasses as jdc
 import jax.numpy as jnp
 
 from .ray import Ray
 from .grid import Grid
 from . import Degrees, CoordsXY, ScaleYX, ShapeYX
 from .tree_utils import HasParamsMixin
-
+from .aberrations import grad_W_krivanek, W_krivanek
 
 class Component(HasParamsMixin):
     """Base component that transforms a ray without side effects.
@@ -173,6 +173,47 @@ def __call__(self, ray: Ray):
         )
 
 
+@jdc.pytree_dataclass
+class AberratedLensKrivanek(Lens):
+    """Thin lens with Krivanek aberrations.
+
+    Parameters
+    ----------
+    z : float
+        Axial position in metres.
+    focal_length : float
+        Focal length in metres.
+    aber_coeffs : jnp.ndarray
+
+    """
+    coeffs: Dict
+
+    def __call__(self, ray: Ray):
+        f = self.focal_length
+        x, y, dx, dy = ray.x, ray.y, ray.dx, ray.dy
+        coeffs = self.coeffs
+
+        # Paraxial thin lens
+        ideal_dx = -x / f + dx
+        ideal_dy = -y / f + dy
+
+        alpha = jnp.hypot(ideal_dx, ideal_dy)    # radians
+        phi = jnp.arctan2(ideal_dy, ideal_dx)    # radians
+
+        dWx, dWy = grad_W_krivanek(ideal_dx, ideal_dy, coeffs)
+        dux, duy = -dWx / f, -dWy / f
+
+        aber_dx = ideal_dx + dux
+        aber_dy = ideal_dy + duy
+
+        pathlength = ray.pathlength - (x**2 + y**2) / (2 * f) + W_krivanek(alpha, phi, coeffs) / f
+        one = ray._one * 1.0
+
+        return Ray(
+            x=x, y=y, dx=aber_dx, dy=aber_dy, _one=one, pathlength=pathlength, z=ray.z
+        )
+
+
 @jdc.pytree_dataclass
 class ScanGrid(Component, Grid):
     """Scanning grid defining pixel-to-metre mapping at plane z.
@@ -434,8 +475,8 @@ class Deflector(Component):
     def __call__(self, ray: Ray):
         x, y, dx, dy = ray.x, ray.y, ray.dx, ray.dy
         return ray.derive(
-            dx=dx + self.def_x,
-            dy=dy + self.def_y,
+            dx=dx + self.def_x * ray._one,
+            dy=dy + self.def_y * ray._one,
             pathlength=ray.pathlength + dx * x + dy * y,
         )
 
@@ -505,53 +546,13 @@ class Biprism(Component):
     z: float
     offset: float = 0.0
     rotation: Degrees = 0.0
-    deflection: float = 0.0
-
-    def __call__(
-        self,
-        ray: Ray,
-    ) -> Ray:
-        pos_x, pos_y, dx, dy = ray.x, ray.y, ray.dx, ray.dy
-
-        deflection = self.deflection
-        offset = self.offset
-        rot = jnp.deg2rad(self.rotation)
-
-        rays_v = jnp.array([pos_x, pos_y]).T
-
-        biprism_loc_v = jnp.array([offset * jnp.cos(rot), offset * jnp.sin(rot)])
-
-        biprism_v = jnp.array([-jnp.sin(rot), jnp.cos(rot)])
-        biprism_v /= jnp.linalg.norm(biprism_v)
-
-        rays_v_centred = rays_v - biprism_loc_v
-
-        dot_product = jnp.dot(rays_v_centred, biprism_v) / jnp.dot(biprism_v, biprism_v)
-        projection = jnp.outer(dot_product, biprism_v)
-
-        rejection = rays_v_centred - projection
-        rejection = rejection / jnp.linalg.norm(rejection, axis=1, keepdims=True)
-
-        # If the ray position is located at [zero, zero], rejection_norm returns a nan,
-        # so we convert it to a zero, zero.
-        rejection = jnp.nan_to_num(rejection)
-
-        xdeflection_mag = rejection[:, 0]
-        ydeflection_mag = rejection[:, 1]
-
-        new_dx = (dx + xdeflection_mag * deflection).squeeze()
-        new_dy = (dy + ydeflection_mag * deflection).squeeze()
+    def_x: float = 0.0
+    side: int = 1
 
-        pathlength = ray.pathlength + (
-            xdeflection_mag * deflection * pos_x + ydeflection_mag * deflection * pos_y
-        )
-
-        return Ray(
-            x=pos_x.squeeze(),
-            y=pos_y.squeeze(),
-            dx=new_dx,
-            dy=new_dy,
-            _one=ray._one,
-            pathlength=pathlength,
-            z=ray.z,
+    def __call__(self, ray: Ray):
+        x, y, dx, dy = ray.x, ray.y, ray.dx, ray.dy
+        return ray.derive(
+            dx=dx + self.def_x * ray._one * jnp.sign(ray.x),
+            dy=dy,
+            pathlength=ray.pathlength + dx * x + dy * y,
         )