From 2e00c45b62591501fdff215c89bcd11cf4b9664d Mon Sep 17 00:00:00 2001 From: Colin Gilgenbach Date: Sat, 20 Dec 2025 16:50:17 -0500 Subject: [PATCH 1/5] Beginning arbitrary aberration support --- phaser/utils/optics.py | 92 ++++++++++++++++++++++++++++++++++++++++++ tests/test_optics.py | 38 ++++++++++++++++- 2 files changed, 129 insertions(+), 1 deletion(-) diff --git a/phaser/utils/optics.py b/phaser/utils/optics.py index cdb406c..cb1a1c7 100644 --- a/phaser/utils/optics.py +++ b/phaser/utils/optics.py @@ -7,11 +7,103 @@ 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, 4.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)) + + @t.overload def make_focused_probe(ky: NDArray[numpy.float64], kx: NDArray[numpy.float64], wavelength: Float, aperture: Float, *, defocus: Float = 0.) -> NDArray[numpy.complex128]: diff --git a/tests/test_optics.py b/tests/test_optics.py index cc85354..b0fc573 100644 --- a/tests/test_optics.py +++ b/tests/test_optics.py @@ -1,10 +1,15 @@ import numpy +import pytest from .utils import with_backends, check_array_equals_file from phaser.utils.num import get_backend_module, BackendName, Sampling, to_numpy, fft2, ifft2 -from phaser.utils.optics import make_focused_probe, fresnel_propagator +from phaser.utils.optics import ( + make_focused_probe, fresnel_propagator, + Aberration, AberrationList, _normalize_aberrations, + Krivanek, Cartesian, Polar, KrivanekComplex, KrivanekCartesian, KrivanekPolar, +) @with_backends('numpy', 'jax', 'cupy', 'torch') @@ -56,3 +61,34 @@ def test_propagator_sign(backend: BackendName) -> numpy.ndarray: probe = ifft2(fft2(probe) * prop) return to_numpy(xp.abs(probe)) + + +def test_parse_aberrations(): + import pane + result = pane.convert([ + {'c3': 5.0}, # haider complex + {'b2': {'a': 5.0, 'b': -2.0}}, # haider cartesian + {'a1': {'mag': 5.0, 'angle': 90.0}}, # haider polar + {'n': 4, 'm': 1, 'val': 1+1.j}, # krivanek complex + {'n': 1, 'm': 0, 'a': 5.0}, # krivanek cartesian + {'n': 5, 'm': 0, 'mag': 5.0}, # krivanek polar + ], AberrationList) + + assert result == [ + {'c3': complex(5.0)}, + {'b2': Cartesian(a=5.0, b=-2.0)}, + {'a1': Polar(mag=5.0, angle=90.0)}, + KrivanekComplex(4, 1, val=1+1.j), + KrivanekCartesian(1, 0, a=5.0, b=0.0), + KrivanekPolar(5, 0, mag=5.0, angle=0.0), + ] + + assert list(_normalize_aberrations(result)) == [ + KrivanekComplex.make_unchecked(3, 0, val=complex(5.0)), + KrivanekComplex.make_unchecked(2, 1, val=15.0-6.0j), + KrivanekComplex.make_unchecked(1, 2, val=pytest.approx(5.0j)), + KrivanekComplex.make_unchecked(4, 1, val=1+1.j), + KrivanekComplex.make_unchecked(1, 0, val=complex(5.0)), + KrivanekComplex.make_unchecked(5, 0, val=complex(5.0)), + ] + # TODO test failures \ No newline at end of file From 6b96810fe17fb1a2de42e3af6616cc2043c538e7 Mon Sep 17 00:00:00 2001 From: Colin Gilgenbach Date: Sat, 20 Dec 2025 17:26:47 -0500 Subject: [PATCH 2/5] Aberration suport in make_focused_probe --- phaser/utils/optics.py | 37 ++++++++++++++++--- .../expected/probe_25mrad_aberrated_mag.tiff | 3 ++ .../probe_25mrad_aberrated_phase.tiff | 3 ++ tests/test_optics.py | 22 ++++++++++- 4 files changed, 58 insertions(+), 7 deletions(-) create mode 100644 tests/expected/probe_25mrad_aberrated_mag.tiff create mode 100644 tests/expected/probe_25mrad_aberrated_phase.tiff diff --git a/phaser/utils/optics.py b/phaser/utils/optics.py index cb1a1c7..04eb3aa 100644 --- a/phaser/utils/optics.py +++ b/phaser/utils/optics.py @@ -104,23 +104,44 @@ def _normalize_aberrations(aberrations: t.Iterable[Aberration]) -> t.Iterator[Kr 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). @@ -131,8 +152,12 @@ 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 diff --git a/tests/expected/probe_25mrad_aberrated_mag.tiff b/tests/expected/probe_25mrad_aberrated_mag.tiff new file mode 100644 index 0000000..ca9dedc --- /dev/null +++ b/tests/expected/probe_25mrad_aberrated_mag.tiff @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fb529979078b577fde574b06577a489c8bf94022f7aafd970ee49bb955ae80fb +size 8388880 diff --git a/tests/expected/probe_25mrad_aberrated_phase.tiff b/tests/expected/probe_25mrad_aberrated_phase.tiff new file mode 100644 index 0000000..3d0cae2 --- /dev/null +++ b/tests/expected/probe_25mrad_aberrated_phase.tiff @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:853533607f978cc777454c8fc751f7668f105c15666d1510ee0d0aff794095f3 +size 8388880 diff --git a/tests/test_optics.py b/tests/test_optics.py index b0fc573..44e1dcf 100644 --- a/tests/test_optics.py +++ b/tests/test_optics.py @@ -2,9 +2,11 @@ import numpy import pytest +from phaser.utils.physics import Electron + from .utils import with_backends, check_array_equals_file -from phaser.utils.num import get_backend_module, BackendName, Sampling, to_numpy, fft2, ifft2 +from phaser.utils.num import abs2, get_backend_module, BackendName, Sampling, to_numpy, fft2, ifft2 from phaser.utils.optics import ( make_focused_probe, fresnel_propagator, Aberration, AberrationList, _normalize_aberrations, @@ -34,6 +36,24 @@ def test_defocused_probe(backend: BackendName) -> numpy.ndarray: return to_numpy(probe) +@with_backends('numpy', 'jax', 'cupy', 'torch') +@check_array_equals_file('probe_25mrad_aberrated.tiff', decimal=5) +def test_aberrated_probe(backend: BackendName) -> numpy.ndarray: + xp = get_backend_module(backend) + sampling = Sampling((1024, 1024), extent=(39.05, 39.05)) + wavelength = Electron(200e3).wavelength + + probe = make_focused_probe( + *sampling.recip_grid(dtype=numpy.float32, xp=xp), wavelength, aperture=25., defocus=10., + aberrations=[ + {'a1': -12.0+5.0j}, + KrivanekCartesian(2, 3, a=300., b=-400.), + KrivanekComplex(3, 0, val=-500_000.0), + ] + ) + return to_numpy(probe) + + @with_backends('numpy', 'jax', 'cupy', 'torch') @check_array_equals_file('fresnel_200kV_1nm_phase.tiff', decimal=5) def test_fresnel_propagator(backend: BackendName) -> numpy.ndarray: From e7a11c571907a2ceda1957f6651cf149d8f1c5d1 Mon Sep 17 00:00:00 2001 From: Colin Gilgenbach Date: Sat, 20 Dec 2025 17:36:57 -0500 Subject: [PATCH 3/5] Add aberration support to init probe hook --- phaser/execute.py | 2 +- phaser/hooks/__init__.py | 4 ++-- phaser/hooks/probe.py | 5 ++++- tests/test_initialization.py | 1 + 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/phaser/execute.py b/phaser/execute.py index dc7c8a2..593a71c 100644 --- a/phaser/execute.py +++ b/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()) diff --git a/phaser/hooks/__init__.py b/phaser/hooks/__init__.py index 6371748..82841f8 100644 --- a/phaser/hooks/__init__.py +++ b/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']): diff --git a/phaser/hooks/probe.py b/phaser/hooks/probe.py index c875c47..d39e513 100644 --- a/phaser/hooks/probe.py +++ b/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) \ No newline at end of file diff --git a/tests/test_initialization.py b/tests/test_initialization.py index a422656..b8494b4 100644 --- a/tests/test_initialization.py +++ b/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 From 85e3b0d9deaef42e0ca1d91096d203e16228596d Mon Sep 17 00:00:00 2001 From: Colin Gilgenbach Date: Mon, 22 Dec 2025 14:18:40 -0600 Subject: [PATCH 4/5] Aberration fix (still needs more validation) --- phaser/utils/optics.py | 2 +- tests/expected/probe_25mrad_aberrated_mag.tiff | 2 +- tests/expected/probe_25mrad_aberrated_phase.tiff | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/phaser/utils/optics.py b/phaser/utils/optics.py index 04eb3aa..c01016d 100644 --- a/phaser/utils/optics.py +++ b/phaser/utils/optics.py @@ -115,7 +115,7 @@ def aberration_surface( for ab in _normalize_aberrations(aberrations): p = (ab.n + 1 + ab.m) // 2 q = ab.n + 1 - p - prod = omega**p + omega.conj()**q + prod = omega**p * omega.conj()**q chi += (prod.real * ab.val.real + prod.imag * ab.val.imag) / (ab.n+1) return chi diff --git a/tests/expected/probe_25mrad_aberrated_mag.tiff b/tests/expected/probe_25mrad_aberrated_mag.tiff index ca9dedc..34be019 100644 --- a/tests/expected/probe_25mrad_aberrated_mag.tiff +++ b/tests/expected/probe_25mrad_aberrated_mag.tiff @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:fb529979078b577fde574b06577a489c8bf94022f7aafd970ee49bb955ae80fb +oid sha256:95fc648c939eb8ffd51bdaca5af3226921656795a9e436e97082931d4811a4b9 size 8388880 diff --git a/tests/expected/probe_25mrad_aberrated_phase.tiff b/tests/expected/probe_25mrad_aberrated_phase.tiff index 3d0cae2..ac27d17 100644 --- a/tests/expected/probe_25mrad_aberrated_phase.tiff +++ b/tests/expected/probe_25mrad_aberrated_phase.tiff @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:853533607f978cc777454c8fc751f7668f105c15666d1510ee0d0aff794095f3 +oid sha256:4cc43fdb3d06d2f03e835777bd982f5238fc9cde99ae29c7b704b2ad00caedb6 size 8388880 From 28d5013f155395ad26a257f4442b4ca4fa4712c8 Mon Sep 17 00:00:00 2001 From: Colin Gilgenbach Date: Tue, 23 Dec 2025 17:03:29 -0600 Subject: [PATCH 5/5] Validate aberrations against Kirkland --- phaser/utils/optics.py | 8 +-- tests/expected/probe_15mrad_spherical_info.md | 65 ++++++++++++++++++ .../expected/probe_15mrad_spherical_mag.tiff | 3 + .../probe_15mrad_spherical_phase.tiff | 3 + .../expected/probe_25mrad_aberrated_mag.tiff | 3 - .../probe_25mrad_aberrated_phase.tiff | 3 - tests/expected/probe_30mrad_aberrated_info.md | 68 +++++++++++++++++++ .../expected/probe_30mrad_aberrated_mag.tiff | 3 + .../probe_30mrad_aberrated_phase.tiff | 3 + tests/test_optics.py | 39 ++++++++--- 10 files changed, 179 insertions(+), 19 deletions(-) create mode 100644 tests/expected/probe_15mrad_spherical_info.md create mode 100644 tests/expected/probe_15mrad_spherical_mag.tiff create mode 100644 tests/expected/probe_15mrad_spherical_phase.tiff delete mode 100644 tests/expected/probe_25mrad_aberrated_mag.tiff delete mode 100644 tests/expected/probe_25mrad_aberrated_phase.tiff create mode 100644 tests/expected/probe_30mrad_aberrated_info.md create mode 100644 tests/expected/probe_30mrad_aberrated_mag.tiff create mode 100644 tests/expected/probe_30mrad_aberrated_phase.tiff diff --git a/phaser/utils/optics.py b/phaser/utils/optics.py index c01016d..2cb8e05 100644 --- a/phaser/utils/optics.py +++ b/phaser/utils/optics.py @@ -42,7 +42,7 @@ def from_known(s: str) -> 'Krivanek': '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, 4.0), # C_32 = 3*S3 + '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 @@ -157,7 +157,7 @@ def make_focused_probe(ky: NDArray[numpy.floating], kx: NDArray[numpy.floating], if len(aberrations) > 0: chi += aberration_surface(thetay, thetax, aberrations) - probe = xp.exp(-2.j*numpy.pi/wavelength * chi) + probe = xp.exp(2.j*numpy.pi/wavelength * chi) mask = theta2 <= (aperture * 1e-3)**2 probe *= mask @@ -165,7 +165,7 @@ def make_focused_probe(ky: NDArray[numpy.floating], kx: NDArray[numpy.floating], # normalize intensity of probe probe /= xp.sqrt(xp.sum(abs2(probe))) - return ifft2(probe) + return ifft2(probe).conj() def make_hermetian_modes( @@ -342,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 diff --git a/tests/expected/probe_15mrad_spherical_info.md b/tests/expected/probe_15mrad_spherical_info.md new file mode 100644 index 0000000..7ff3eb1 --- /dev/null +++ b/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)) +``` \ No newline at end of file diff --git a/tests/expected/probe_15mrad_spherical_mag.tiff b/tests/expected/probe_15mrad_spherical_mag.tiff new file mode 100644 index 0000000..fa45742 --- /dev/null +++ b/tests/expected/probe_15mrad_spherical_mag.tiff @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d4afc40db352525f7d782b7f211b41bc093e162b97df3d3cb16911d686447920 +size 4194576 diff --git a/tests/expected/probe_15mrad_spherical_phase.tiff b/tests/expected/probe_15mrad_spherical_phase.tiff new file mode 100644 index 0000000..98e6ab6 --- /dev/null +++ b/tests/expected/probe_15mrad_spherical_phase.tiff @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b91bf5578286c8f367207ad03a04606c8055d18fe859f12e0a89a77c5d6dcd21 +size 4194576 diff --git a/tests/expected/probe_25mrad_aberrated_mag.tiff b/tests/expected/probe_25mrad_aberrated_mag.tiff deleted file mode 100644 index 34be019..0000000 --- a/tests/expected/probe_25mrad_aberrated_mag.tiff +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:95fc648c939eb8ffd51bdaca5af3226921656795a9e436e97082931d4811a4b9 -size 8388880 diff --git a/tests/expected/probe_25mrad_aberrated_phase.tiff b/tests/expected/probe_25mrad_aberrated_phase.tiff deleted file mode 100644 index ac27d17..0000000 --- a/tests/expected/probe_25mrad_aberrated_phase.tiff +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4cc43fdb3d06d2f03e835777bd982f5238fc9cde99ae29c7b704b2ad00caedb6 -size 8388880 diff --git a/tests/expected/probe_30mrad_aberrated_info.md b/tests/expected/probe_30mrad_aberrated_info.md new file mode 100644 index 0000000..5a7bd80 --- /dev/null +++ b/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)) +``` \ No newline at end of file diff --git a/tests/expected/probe_30mrad_aberrated_mag.tiff b/tests/expected/probe_30mrad_aberrated_mag.tiff new file mode 100644 index 0000000..b5c85b8 --- /dev/null +++ b/tests/expected/probe_30mrad_aberrated_mag.tiff @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6b617ecb73a4c0d2ec1bca0806beb44a6457cbffd6b3deaadb0b3c5c4e0b353d +size 4194576 diff --git a/tests/expected/probe_30mrad_aberrated_phase.tiff b/tests/expected/probe_30mrad_aberrated_phase.tiff new file mode 100644 index 0000000..77f0d2a --- /dev/null +++ b/tests/expected/probe_30mrad_aberrated_phase.tiff @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eede767e50a0bc50b89d41a09193429a61fb5a0ef91d9d52406129ecce81f19d +size 4194576 diff --git a/tests/test_optics.py b/tests/test_optics.py index 44e1dcf..8866bc3 100644 --- a/tests/test_optics.py +++ b/tests/test_optics.py @@ -6,11 +6,11 @@ from .utils import with_backends, check_array_equals_file -from phaser.utils.num import abs2, get_backend_module, BackendName, Sampling, to_numpy, fft2, ifft2 +from phaser.utils.num import get_backend_module, BackendName, Sampling, to_numpy, fft2, ifft2 from phaser.utils.optics import ( make_focused_probe, fresnel_propagator, - Aberration, AberrationList, _normalize_aberrations, - Krivanek, Cartesian, Polar, KrivanekComplex, KrivanekCartesian, KrivanekPolar, + AberrationList, _normalize_aberrations, + Cartesian, Polar, KrivanekComplex, KrivanekCartesian, KrivanekPolar, ) @@ -37,18 +37,39 @@ def test_defocused_probe(backend: BackendName) -> numpy.ndarray: @with_backends('numpy', 'jax', 'cupy', 'torch') -@check_array_equals_file('probe_25mrad_aberrated.tiff', decimal=5) +@check_array_equals_file('probe_30mrad_aberrated.tiff', decimal=5, + out_name='probe_30mrad_aberrated_{backend}.tiff') def test_aberrated_probe(backend: BackendName) -> numpy.ndarray: xp = get_backend_module(backend) - sampling = Sampling((1024, 1024), extent=(39.05, 39.05)) + sampling = Sampling((1024, 1024), extent=(50.0, 50.0)) + wavelength = Electron(300e3).wavelength + + probe = make_focused_probe( + *sampling.recip_grid(dtype=numpy.float32, xp=xp), wavelength, aperture=30., defocus=0., + aberrations=[ + {'a1': 10.0+10.0j}, + {'b2': (1e3+2e3j) / 3}, + KrivanekComplex(3, 2, val=50e3j), + ] + ) + return to_numpy(probe) + + +@with_backends('numpy', 'jax', 'cupy', 'torch') +@check_array_equals_file('probe_15mrad_spherical.tiff', decimal=5, + out_name='probe_15mrad_spherical_{backend}.tiff') +def test_spherical_probe(backend: BackendName) -> numpy.ndarray: + xp = get_backend_module(backend) + sampling = Sampling((1024, 1024), extent=(50.0, 50.0)) wavelength = Electron(200e3).wavelength probe = make_focused_probe( - *sampling.recip_grid(dtype=numpy.float32, xp=xp), wavelength, aperture=25., defocus=10., + *sampling.recip_grid(dtype=numpy.float32, xp=xp), wavelength, aperture=15., + defocus=-578.266, aberrations=[ - {'a1': -12.0+5.0j}, - KrivanekCartesian(2, 3, a=300., b=-400.), - KrivanekComplex(3, 0, val=-500_000.0), + {'c3': 1.0e+7}, + {'a1': 20.0+20.0j}, + KrivanekCartesian(3, 2, a=1.5e6, b=2.0e6), ] ) return to_numpy(probe)