Arbitrary aberration support for initial probe

phaser/execute.py

@@ -473,7 +473,7 @@ def _as_dict(val) -> t.Optional[dict]:
             return right
         d = merge(left.props or {}, right.props or {})
         d['type'] = right.type if right.type is not None else right.ref
-        return pane.from_data(d, right.__class__)
+        return pane.convert(d, right.__class__)
 
     if (left_d := _as_dict(left)) is not None and (right_d := _as_dict(right)) is not None:
         keys = set(left_d.keys()) | set(right_d.keys())

phaser/hooks/__init__.py

@@ -8,6 +8,7 @@
 
 from ..types import Dataclass, Slices
 from .hook import Hook
+from ..utils.optics import Aberration
 
 if t.TYPE_CHECKING:
     from phaser.utils.num import Sampling
@@ -84,8 +85,6 @@ class RawDataHook(Hook[None, RawData]):
     }
 
 
-
-
 class ProbeHookArgs(t.TypedDict):
     sampling: 'Sampling'
     wavelength: float
@@ -97,6 +96,7 @@ class ProbeHookArgs(t.TypedDict):
 class FocusedProbeProps(Dataclass):
     defocus: t.Optional[float] = None  # defocus, + is overfocus [A]
     conv_angle: t.Optional[float] = None  # semiconvergence angle [mrad]
+    aberrations: t.Sequence[Aberration] = ()
 
 
 class ProbeHook(Hook[ProbeHookArgs, 'ProbeState']):

phaser/hooks/probe.py

@@ -14,11 +14,14 @@ def focused_probe(args: ProbeHookArgs, props: FocusedProbeProps) -> ProbeState:
         raise ValueError("Probe 'defocus' must be specified by metadata or manually")
 
     logger.info(f"Making probe, conv_angle {props.conv_angle} mrad, defocus {props.defocus} A")
+    if len(props.aberrations):
+        s = '\n'.join(f"  {ab!r}" for ab in props.aberrations)
+        logger.info(f"Aberrations:\n{s}")
 
     sampling = args['sampling']
     ky, kx = sampling.recip_grid(dtype=args['dtype'], xp=args['xp'])
     probe = make_focused_probe(
         ky, kx, args['wavelength'],
-        props.conv_angle, defocus=props.defocus
+        props.conv_angle, defocus=props.defocus, aberrations=props.aberrations
     )
     return ProbeState(sampling, probe)

phaser/utils/optics.py

@@ -7,28 +7,141 @@
 
 import numpy
 from numpy.typing import NDArray, ArrayLike
+import pane
+from pane.annotations import Condition
+from pane.util import pluralize
 
 from .num import get_array_module, ifft2, abs2, NumT, ufunc_outer, is_jax, cast_array_module
 from .num import Float, Sampling, to_complex_dtype, to_real_dtype, split_array, to_numpy
 
 
+class Krivanek(pane.PaneBase):
+    n: int
+    m: int
+    scale_factor: float = 1.0
+
+    def __post_init__(self):
+        if (
+            self.n < 0 or self.m < 0 or
+            self.m > self.n + 1 or
+            self.m % 2 + self.n % 2 != 1
+        ):
+            raise ValueError(f"Invalid Krivanek aberration n={self.n} m={self.m}")
+
+    @staticmethod
+    def from_known(s: str) -> 'Krivanek':
+        try:
+            return _KNOWN_ABERRATIONS[s.lower()]
+        except (KeyError, TypeError):
+            raise ValueError(f"Unknown aberration '{s}'") from None
+
+
+_KNOWN_ABERRATIONS: t.Dict[str, Krivanek] = {
+    'c1': Krivanek.make_unchecked(1, 0),
+    'a1': Krivanek.make_unchecked(1, 2),
+    'b2': Krivanek.make_unchecked(2, 1, 3.0),  # C_21 = 3*B2
+    'a2': Krivanek.make_unchecked(2, 3),
+    'c3': Krivanek.make_unchecked(3, 0),
+    's3': Krivanek.make_unchecked(3, 2, 3.0),  # C_32 = 3*S3
+    'a3': Krivanek.make_unchecked(3, 4),
+    'b4': Krivanek.make_unchecked(4, 1, 4.0),  # C_41 = 4*B4
+    'd4': Krivanek.make_unchecked(4, 3, 4.0),  # C_43 = 4*D4
+    'a4': Krivanek.make_unchecked(4, 5),
+    'c5': Krivanek.make_unchecked(5, 0),
+}
+
+KnownAberration: t.TypeAlias = t.Annotated[str, Condition(
+    lambda s: s.lower() in _KNOWN_ABERRATIONS,
+    'known aberration',
+    lambda exp, plural: pluralize('known aberration', plural)
+)]
+
+
+class Cartesian(pane.PaneBase, kw_only=True):
+    a: float
+    b: float = 0.0
+
+    def __complex__(self) -> complex:
+        return complex(self.a, self.b)
+
+
+class Polar(pane.PaneBase, kw_only=True):
+    mag: float
+    angle: float = 0.0  # degrees
+
+    def __complex__(self) -> complex:
+        theta = numpy.deg2rad(self.angle)
+        return self.mag * complex(numpy.cos(theta), numpy.sin(theta))
+
+
+class KrivanekComplex(Krivanek, kw_only=True):
+    val: complex
+
+    def __complex__(self) -> complex:
+        return self.val
+
+class KrivanekCartesian(Krivanek, Cartesian, kw_only=True):
+    ...
+
+class KrivanekPolar(Krivanek, Polar, kw_only=True):
+    ...
+
+
+Aberration: t.TypeAlias = t.Union[
+    t.Dict[KnownAberration, t.Union[complex, Cartesian, Polar]],
+    KrivanekComplex, KrivanekCartesian, KrivanekPolar,
+]
+AberrationList: t.TypeAlias = t.List[Aberration]
+
+
+def _normalize_aberrations(aberrations: t.Iterable[Aberration]) -> t.Iterator[KrivanekComplex]:
+    for ab in aberrations:
+        if isinstance(ab, dict):
+            for known, val in ab.items():
+                ty = Krivanek.from_known(known)
+                yield KrivanekComplex(ty.n, ty.m, val=ty.scale_factor * complex(val))
+        else:
+            yield KrivanekComplex(ab.n, ab.m, val=complex(ab))
+
+
+def aberration_surface(
+    thetay: NDArray[numpy.float64], thetax: NDArray[numpy.float64],
+    aberrations: t.Iterable[Aberration]
+) -> NDArray[numpy.floating]:
+    xp = get_array_module(thetay, thetax)
+    chi = xp.zeros_like(thetay)
+    omega = thetax + thetay*1.j
+
+    for ab in _normalize_aberrations(aberrations):
+        p = (ab.n + 1 + ab.m) // 2
+        q = ab.n + 1 - p
+        prod = omega**p * omega.conj()**q
+        chi += (prod.real * ab.val.real + prod.imag * ab.val.imag) / (ab.n+1)
+
+    return chi
+
+
 @t.overload
 def make_focused_probe(ky: NDArray[numpy.float64], kx: NDArray[numpy.float64], wavelength: Float,
-                       aperture: Float, *, defocus: Float = 0.) -> NDArray[numpy.complex128]:
+                       aperture: Float, *, defocus: Float = 0.,
+                       aberrations: t.Sequence[Aberration] = ()) -> NDArray[numpy.complex128]:
     ...
 
 @t.overload
 def make_focused_probe(ky: NDArray[numpy.float32], kx: NDArray[numpy.float32], wavelength: Float,
-                       aperture: Float, *, defocus: Float = 0.) -> NDArray[numpy.complex64]:
+                       aperture: Float, *, defocus: Float = 0.,
+                       aberrations: t.Sequence[Aberration] = ()) -> NDArray[numpy.complex64]:
     ...
 
 @t.overload
 def make_focused_probe(ky: NDArray[numpy.floating], kx: NDArray[numpy.floating], wavelength: Float,
-                       aperture: Float, *, defocus: Float = 0.) -> NDArray[numpy.complexfloating]:
+                       aperture: Float, *, defocus: Float = 0.,
+                       aberrations: t.Sequence[Aberration] = ()) -> NDArray[numpy.complexfloating]:
     ...
 
 def make_focused_probe(ky: NDArray[numpy.floating], kx: NDArray[numpy.floating], wavelength: Float,
-                       aperture: Float, *, defocus: Float = 0.) -> NDArray[numpy.complexfloating]:
+                       aperture: Float, *, defocus: Float = 0.,
+                       aberrations: t.Sequence[Aberration] = ()) -> NDArray[numpy.complexfloating]:
     """
     Create a focused probe from a circular aperture of semi-angle `aperture` (in mrad).
 
@@ -39,16 +152,20 @@ def make_focused_probe(ky: NDArray[numpy.floating], kx: NDArray[numpy.floating],
     thetay, thetax = ky * wavelength, kx * wavelength
     theta2 = thetay**2 + thetax**2
 
-    phase = (defocus/(2. * wavelength)) * theta2
-    probe = xp.exp(-2.j*numpy.pi * phase)
+    # wavefront error, length units
+    chi = defocus/2. * theta2
+    if len(aberrations) > 0:
+        chi += aberration_surface(thetay, thetax, aberrations)
+
+    probe = xp.exp(2.j*numpy.pi/wavelength * chi)
 
     mask = theta2 <= (aperture * 1e-3)**2
     probe *= mask
 
     # normalize intensity of probe
     probe /= xp.sqrt(xp.sum(abs2(probe)))
 
-    return ifft2(probe)
+    return ifft2(probe).conj()
 
 
 def make_hermetian_modes(
@@ -225,7 +342,7 @@ def estimate_probe_radius(wavelength: Float, aperture: Float, defocus: Float, *,
     aperture *= 1e-3  # mrad -> rad
 
     if threshold == 'geom':
-        return float(defocus * aperture)
+        return abs(float(defocus * aperture))
 
     rel_defocus = numpy.abs(defocus) / wavelength
 

tests/expected/probe_15mrad_spherical_info.md

@@ -0,0 +1,65 @@
+Generated from Kirkland's temsim code, with the following settings:
+
+- Voltage: 200 kV
+- 1024x1024 pixels
+- Supercell size: 50x50 angstrom
+- Convergence angle: 15 mrad
+
+Aberrations:
+```
+Haider  Krivanek value
+C3      C3,0     1 mm
+C1      C1,0     -578.266 angstrom (Scherzer underfocus)
+A1      C1,2     20.0+20.0j angstrom
+3*S3    C3,2     0.15+0.2j mm
+```
+
+This corresponds to the below input to the 'probe' program:
+```
+0
+probe_15mrad_spherical.tiff
+1024
+1024
+50.0
+50.0
+200.0
+1.0
+0.0
+578.266
+15.0
+0
+0.0
+0.0
+C12a
+0.0000020000
+C12b
+0.0000020000
+C32a
+0.1500000000
+C32b
+0.2000000000
+END
+```
+
+Then, the following post-processing can be used to output the 'mag' (really amplitude) and 'phase' images:
+
+```python
+from pathlib import Path
+import numpy
+import tifffile
+
+in_path = Path("probe_15mrad_spherical.tiff")
+# read Kirkland FloatTIFF
+img = numpy.asarray(tifffile.imread(f, series=1))
+# combine real and imaginary images
+re, im = numpy.split(img, 2, axis=-1)
+img = re + im * 1.j
+# center probe
+img = numpy.fft.fftshift(img)
+# normalize
+img /= numpy.sqrt(numpy.sum(numpy.abs(img)**2))
+
+# write output files
+tifffile.imwrite(in_path.with_stem(in_path.stem + "_mag"), numpy.abs(img))
+tifffile.imwrite(in_path.with_stem(in_path.stem + "_phase"), numpy.angle(img))
+```

tests/expected/probe_30mrad_aberrated_info.md

@@ -0,0 +1,68 @@
+Generated from Kirkland's temsim code, with the following settings:
+
+- Voltage: 300 kV
+- 1024x1024 pixels
+- Supercell size: 50x50 angstrom
+- Convergence angle: 30 mrad
+
+Aberrations:
+```
+Haider  Krivanek  value
+A1      C1,2      10+10j angstrom
+3*B2    C2,1      1e3+2e3j angstrom
+3*S3    C3,2      50e3j angstrom
+```
+
+This corresponds to the below input to the 'probe' program:
+```
+0
+probe_30mrad_aberrated.tiff
+1024
+1024
+50.0
+50.0
+300.0
+0.0
+0.0
+0.0
+30.0
+0
+0.0
+0.0
+C12a
+0.0000010000
+C12b
+0.0000010000
+C21a
+0.0001000000
+C21b
+0.0002000000
+C32a
+0.0000000000
+C32b
+0.0050000000
+END
+```
+
+Then, the following post-processing can be used to output the 'mag' (really amplitude) and 'phase' images:
+
+```python
+from pathlib import Path
+import numpy
+import tifffile
+
+in_path = Path("probe_30mrad_aberrated.tiff")
+# read Kirkland FloatTIFF
+img = numpy.asarray(tifffile.imread(f, series=1))
+# combine real and imaginary images
+re, im = numpy.split(img, 2, axis=-1)
+img = re + im * 1.j
+# center probe
+img = numpy.fft.fftshift(img)
+# normalize
+img /= numpy.sqrt(numpy.sum(numpy.abs(img)**2))
+
+# write output files
+tifffile.imwrite(in_path.with_stem(in_path.stem + "_mag"), numpy.abs(img))
+tifffile.imwrite(in_path.with_stem(in_path.stem + "_phase"), numpy.angle(img))
+```

tests/test_initialization.py

@@ -82,6 +82,7 @@ def test_load_raw_data_override():
         'type': 'focused',
         'conv_angle': 20.0,
         'defocus': 200.0,
+        'aberrations': (),
     }
 
     assert pane.into_data(raw_data['scan_hook']) == {  # type: ignore