From 2187b25d25e46b77bd8666b023759e5728f6228a Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:36:47 +0000
Subject: [PATCH 01/34] Check the fitted detector rotation in the overfocus
 optimizer test

The test checked scan_rotation, which the fit doesn't change.

- Fix docstring of make_overfocus_loss_function, which named scan_rotation
- Raise ValueError for a starting model in focus instead of dividing by zero

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/util/optimize.py | 17 ++++++++++-------
 tests/test_optimize.py                      | 17 +++++++++++++++--
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/microscope_calibration/util/optimize.py b/src/microscope_calibration/util/optimize.py
index ebec9f0..fd41bed 100644
--- a/src/microscope_calibration/util/optimize.py
+++ b/src/microscope_calibration/util/optimize.py
@@ -72,10 +72,11 @@ def make_overfocus_loss_function(
     """
     Build a parameter mapping and loss function to optimize
 
-    This maps the :code:`scan_rotation` and :code:`overfocus` parameters
+    This maps the :code:`detector_rotation` and :code:`overfocus` parameters
     for :class:`OverfocusUDF` in such a way that the starting value is at
-    0 and a sensible interval for optimization is :code`[-10, 10]` based
-    on the starting parameters passed as an argument.
+    0 and a sensible interval for optimization is :code:`[-10, 10]` based
+    on the starting parameters passed as an argument: The detector rotation
+    changes by +-10 deg, and the overfocus by +-25 % of its starting value.
 
     It also constructs a loss function that updates the parameters of
     the :class:`OverfocusUDF` passed as parameter :code:`overfocus_udf`,
@@ -85,7 +86,7 @@ def make_overfocus_loss_function(
     Parameters
     ----------
     model
-        Starting value for optimization
+        Starting value for optimization. Its overfocus must not be zero.
     ctx
         Run UDF using this Context
     dataset
@@ -117,8 +118,10 @@ def make_overfocus_loss_function(
         Function that can be called with :func:`scipy.optimize.minimize`. Starting value is
         [0, 0], sensible range is ([-10, 10], [-10, 10]).
     """
+    if model.overfocus == 0:
+        raise ValueError("The overfocus of the starting model must not be zero.")
     # Rotate and scale the angle so that the optimizer works between +-10,
-    # corresponding to +- 5 deg
+    # corresponding to +- 10 deg
     rotation_diff = model.detector_rotation * 180 / np.pi
     rotation_scale = 1
     # Values to shift and scale the overfocus so that the optimizer works between +-10
@@ -136,7 +139,7 @@ def make_new_model(args) -> Model4DSTEM:
         ----------
 
         args
-            (scan_rotation, overfocused) mapped to +- 10
+            (detector_rotation, overfocus) mapped to +-10
         """
         transformed_rotation, transformed_overfocus = args
         rotation = transformed_rotation / rotation_scale + rotation_diff
@@ -158,7 +161,7 @@ def loss(args) -> float:
         ----------
 
         args
-            (scan_rotation, overfocused) mapped to +-10 range
+            (detector_rotation, overfocus) mapped to +-10
         """
         model = make_new_model(args)
         # Hack to make parameter update work
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index 7e3cb5a..1214b3c 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -88,9 +88,11 @@ def test_optimize():
         dataset=ds,
         overfocus_udf=udf,
     )
+    # The fit starts at the true values, i.e. this tests the plumbing, not the
+    # convergence: The blur of a back-projected point object is flat.
     res = optimize(loss=loss)
     res_model = make_new_model(res.x)
-    assert_allclose(res_model.scan_rotation, model.scan_rotation, atol=0.1)
+    assert_allclose(res_model.detector_rotation, model.detector_rotation, atol=0.1)
     assert_allclose(res_model.overfocus, model.overfocus, rtol=0.1)
 
     valdict = {"val": False}
@@ -120,10 +122,21 @@ def callback(args, new_model, udf_results, current_loss):
         bounds=[(-10, 10), (-10, 10)],
     )
     res_model = make_new_model(res.x)
-    assert_allclose(res_model.scan_rotation, model.scan_rotation, atol=0.1)
+    assert_allclose(res_model.detector_rotation, model.detector_rotation, atol=0.1)
     assert_allclose(res_model.overfocus, model.overfocus, rtol=0.1)
 
 
+def test_overfocus_loss_in_focus():
+    model = Model4DSTEM.default(dataset_shape=(8, 8, 8, 8))
+    with pytest.raises(ValueError, match="overfocus"):
+        make_overfocus_loss_function(
+            model=model,
+            ctx=None,
+            dataset=None,
+            overfocus_udf=OverfocusUDF(overfocus_model={"model": model}),
+        )
+
+
 def test_descan_error():
     scan_pixel_pitch = 0.1
     detector_pixel_pitch = 0.2

From 3f10a32f3369da76c8937c2587dfe01ef91ae065 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:37:24 +0000
Subject: [PATCH 02/34] Compare inverted rays with the inverted model's center
 in test_invert_focus

The center ray of the inverted model was traced with the original model.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/test_model.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/tests/test_model.py b/tests/test_model.py
index 9425db9..7341efc 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1082,10 +1082,6 @@ def test_adjust_camera_length(random_model):
 
 
 def test_invert_focus(random_model):
-    # model = Model4DSTEM.default().derive(
-    #     overfocus=0.1,
-    #     descan_error=DescanError().derive(syo_pxi=0.1),
-    # )
     model = random_model
     trace = lambdify_trace_for(np)
     inverted_model = model.invert_focus().normalize_types()
@@ -1093,7 +1089,6 @@ def test_invert_focus(random_model):
     assert_allclose(
         np.abs(model.detector_rotation - inverted_model.detector_rotation), np.pi, atol=1e-12
     )
-    print(model)
 
     assert_no_descan_deviation_cl(model, inverted_model)
 
@@ -1118,7 +1113,7 @@ def slope(res, center, dim1, dim2):
                     model=model, scan_pos=PixelYX(*scan_pos), source_dy=0.0, source_dx=0.0
                 )
                 inverted_center = trace(
-                    model=model, scan_pos=PixelYX(*scan_pos), source_dy=0.0, source_dx=0.0
+                    model=inverted_model, scan_pos=PixelYX(*scan_pos), source_dy=0.0, source_dx=0.0
                 )
                 res = trace(model=model, scan_pos=PixelYX(*scan_pos), source_dy=dy, source_dx=dx)
                 inv_res = trace(

From f22dd3a0ad87245a6d44a8c0af8b4a42d37b6259 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:47:03 +0000
Subject: [PATCH 03/34] Dispatch cos/sin on the argument type instead of always
 using SymPy

Float models got SymPy numbers from adjust_*() and invert_focus(), and
jax.grad() through adjust_*() failed.

- Add util/backend.py, used by rotate(), generated code and runtime derivations
- Regenerate _generated.py
- Add docstring to invert_focus()

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CLAUDE.md                                     |  2 +
 .../common/_generated.py                      |  2 +-
 src/microscope_calibration/common/model.py    | 18 +++++--
 src/microscope_calibration/util/backend.py    | 48 +++++++++++++++++++
 src/microscope_calibration/util/codegen.py    | 20 ++++----
 tests/test_codegen.py                         | 16 +++++--
 tests/test_model.py                           | 37 ++++++++++++++
 7 files changed, 127 insertions(+), 16 deletions(-)
 create mode 100644 src/microscope_calibration/util/backend.py

diff --git a/CLAUDE.md b/CLAUDE.md
index 5ac22d1..61a877a 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -80,6 +80,8 @@ Key layers, bottom-up:
    `wrapt`) that recurses into nested dataclasses/pytrees (`PixelYX`, `Model4DSTEM`,
    `DescanError`) so a symbolic trace can be turned into a numeric function. Touch with care;
    the three-backend contract depends on it.
+   `util/backend.py` provides math functions such as `cos`/`sin` that dispatch on the argument
+   type; use them instead of `sympy`/`numpy`/`jax.numpy` functions in model code.
 
 3. **`common/stem_overfocus.py`** — pixel-to-pixel transformation matrices (backward
    projection, detector correction) calculated from the generated `ray_transfer_matrix()`;
diff --git a/src/microscope_calibration/common/_generated.py b/src/microscope_calibration/common/_generated.py
index 2bf8a7e..26c6397 100644
--- a/src/microscope_calibration/common/_generated.py
+++ b/src/microscope_calibration/common/_generated.py
@@ -4,7 +4,7 @@
 # microscope_calibration.common.generate.GENERATED_FUNCTIONS.
 # tests/test_codegen.py verifies that this file is up to date.
 
-from sympy import cos, sin
+from microscope_calibration.util.backend import cos, sin
 from temgym_core.components import DescanError
 
 
diff --git a/src/microscope_calibration/common/model.py b/src/microscope_calibration/common/model.py
index d3ba27c..106f62d 100644
--- a/src/microscope_calibration/common/model.py
+++ b/src/microscope_calibration/common/model.py
@@ -10,6 +10,7 @@
 jax.config.update("jax_enable_x64", True)  # noqa: E702
 # ruff: disable[E402]
 import jax_dataclasses as jdc
+import numpy as np
 import sympy as sym
 from sympy.solvers.solveset import NonlinearError
 from jax.errors import TracerBoolConversionError
@@ -20,6 +21,7 @@
 from temgym_core.run import run_iter
 from temgym_core.source import PointSource, Source
 
+from microscope_calibration.util import backend
 from microscope_calibration.util.sympy import lambdify, lambdify_tree, symbol_maker
 # ruff: enable[E402]
 
@@ -40,8 +42,8 @@ def scale(c: CoordXY | PixelYX, factor: float) -> CoordXY:
 
 
 def rotate(c: CoordXY | PixelYX, radians: float) -> CoordXY:
-    cosr = sym.cos(radians)
-    sinr = sym.sin(radians)
+    cosr = backend.cos(radians)
+    sinr = backend.sin(radians)
     return CoordXY(
         x=c.x * cosr - c.y * sinr,
         y=c.y * cosr + c.x * sinr,
@@ -600,9 +602,17 @@ def adjust_camera_length(self, camera_length: float) -> Model4DSTEM:
         return self._adjust_compensating_descan_error("camera_length", camera_length)
 
     def invert_focus(self) -> Self:
+        """
+        Model with inverted sign of the overfocus that projects the specimen
+        to the same detector pixels
+
+        Over- and underfocus can't be distinguished from the shadow image
+        alone. This allows switching between the two interpretations.
+        """
+        pi = sym.pi if isinstance(self.detector_rotation, sym.Basic) else np.pi
         return (
             self.derive(overfocus=-self.overfocus)
-            .adjust_detector_rotation(detector_rotation=self.detector_rotation + sym.pi)
+            .adjust_detector_rotation(detector_rotation=self.detector_rotation + pi)
             .adjust_camera_length(camera_length=self.camera_length + 2 * self.overfocus)
         )
 
@@ -652,7 +662,7 @@ def _runtime_descan_compensation(model_type: type[Model4DSTEM], param: str):
     warn_runtime_derivation(f"descan_compensation_{param}", model_type)
     args, output = _derive_descan_compensation(param, model_type=model_type)
     check_pytree(args["model"])
-    return lambdify_tree(args, output, modules=sym)
+    return lambdify_tree(args, output, modules=[backend.NAMESPACE, "sympy"])
 
 
 def _derive_descan_compensation(param: str, model_type: type[Model4DSTEM] = Model4DSTEM):
diff --git a/src/microscope_calibration/util/backend.py b/src/microscope_calibration/util/backend.py
new file mode 100644
index 0000000..67ac842
--- /dev/null
+++ b/src/microscope_calibration/util/backend.py
@@ -0,0 +1,48 @@
+"""
+Math functions that work with SymPy, JAX and plain numbers
+
+The functions dispatch on the type of their argument so that code using them,
+e.g. :meth:`~microscope_calibration.common.model.Model4DSTEM.trace` and the
+generated code in :mod:`microscope_calibration.common._generated`, keeps the
+type of its inputs: SymPy expressions stay symbolic, JAX arrays and tracers
+stay differentiable and traceable, and other numbers are calculated with NumPy.
+"""
+
+from collections.abc import Callable
+
+import jax
+
+jax.config.update("jax_enable_x64", True)  # noqa: E702
+# ruff: disable[E402]
+import jax.numpy as jnp
+import numpy as np
+import sympy as sym
+# ruff: enable[E402]
+
+__all__ = ["NAMESPACE", "cos", "sin"]
+
+
+def _dispatch(name: str) -> Callable:
+    sym_f = getattr(sym, name)
+    jnp_f = getattr(jnp, name)
+    np_f = getattr(np, name)
+
+    def f(x):
+        if isinstance(x, sym.Basic):
+            return sym_f(x)
+        # Includes tracers
+        elif isinstance(x, jax.Array):
+            return jnp_f(x)
+        else:
+            return np_f(x)
+
+    f.__name__ = f.__qualname__ = name
+    f.__doc__ = f"{name}() for SymPy expressions, JAX arrays and other numbers"
+    return f
+
+
+cos = _dispatch("cos")
+sin = _dispatch("sin")
+
+# Namespace for :func:`sympy.lambdify`, e.g. :code:`modules=[NAMESPACE, "sympy"]`
+NAMESPACE = {name: globals()[name] for name in ("cos", "sin")}
diff --git a/src/microscope_calibration/util/codegen.py b/src/microscope_calibration/util/codegen.py
index 1b2ea0d..ed63187 100644
--- a/src/microscope_calibration/util/codegen.py
+++ b/src/microscope_calibration/util/codegen.py
@@ -15,6 +15,8 @@
 import sympy as sym
 from sympy.printing.pycode import SymPyPrinter
 
+from microscope_calibration.util import backend
+
 
 def symbol_paths(tree: Any, path: tuple[str, ...]) -> Iterator[tuple[tuple[str, ...], sym.Symbol]]:
     """
@@ -183,8 +185,8 @@ def render_module(functions: dict[str, tuple[dict[str, Any], Any]], header: str)
     Render a module with the given functions
 
     Names that the printed SymPy expressions use, e.g. :code:`cos`, are imported
-    from :code:`sympy`, so that the result works with SymPy and plain Python
-    number types, like :func:`sympy.lambdify` with :code:`modules=sympy`.
+    from :mod:`microscope_calibration.util.backend`, so that the result works
+    with SymPy expressions, JAX arrays and tracers, and plain numbers.
 
     Parameters
     ----------
@@ -205,12 +207,14 @@ def render_module(functions: dict[str, tuple[dict[str, Any], Any]], header: str)
     for t in types:
         imports.setdefault(t.__module__, set()).add(t.__name__)
     type_names = {t.__name__ for t in types}
-    sympy_names = _free_names(body) - type_names
-    for n in sympy_names:
-        if not hasattr(sym, n):
-            raise ValueError(f"Generated code uses name {n} that is not available in sympy.")
-    if sympy_names:
-        imports.setdefault("sympy", set()).update(sympy_names)
+    math_names = _free_names(body) - type_names
+    for n in math_names:
+        if n not in backend.NAMESPACE:
+            raise ValueError(
+                f"Generated code uses name {n} that is not available in {backend.__name__}."
+            )
+    if math_names:
+        imports.setdefault(backend.__name__, set()).update(math_names)
 
     import_lines = "".join(
         f"from {module} import {', '.join(sorted(names))}\n"
diff --git a/tests/test_codegen.py b/tests/test_codegen.py
index 25d2265..048ec64 100644
--- a/tests/test_codegen.py
+++ b/tests/test_codegen.py
@@ -1,5 +1,6 @@
 import inspect
 
+import jax
 import numpy as np
 import pytest
 import sympy as sym
@@ -53,12 +54,15 @@ def test_render_module():
 
     res = f(p=PixelYX(y=1.0, x=2.0), angle=0.3)
     assert isinstance(res, PixelYX)
-    # Functions like cos() come from SymPy and return SymPy numbers,
-    # like lambdify(modules=sympy)
-    assert_allclose([float(v) for v in res], (np.cos(0.3) + 2.0 * np.sin(0.3) + np.pi, 1.0))
+    assert_allclose(res, (np.cos(0.3) + 2.0 * np.sin(0.3) + np.pi, 1.0))
+    # Functions like cos() keep the backend of their argument
+    assert not any(isinstance(v, sym.Basic) for v in res)
     # Plain arithmetic keeps plain types
     assert type(res.x) is float
 
+    grad = jax.grad(lambda angle: f(p=PixelYX(y=1.0, x=2.0), angle=angle).y)(0.3)
+    assert_allclose(grad, -np.sin(0.3) + 2.0 * np.cos(0.3))
+
     res = f(p=p, angle=angle)
     for r, o in zip(res, output, strict=True):
         # Constants like pi are printed as floats, like lambdify(modules=sympy)
@@ -69,3 +73,9 @@ def test_render_module_unknown_symbol():
     a, b = sym.symbols("a b")
     with pytest.raises(ValueError, match="not in args"):
         render_module({"f": ({"a": a}, a * b)}, header="")
+
+
+def test_render_module_unknown_function():
+    a = sym.Symbol("a")
+    with pytest.raises(ValueError, match="not available"):
+        render_module({"f": ({"a": a}, sym.tan(a))}, header="")
diff --git a/tests/test_model.py b/tests/test_model.py
index 7341efc..bb29cc2 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1081,6 +1081,43 @@ def test_adjust_camera_length(random_model):
     assert modified.camera_length == camera_length
 
 
+ADJUST_VALUES = {
+    "scan_rotation": 0.3,
+    "scan_pixel_pitch": 0.7,
+    "scan_center": PixelYX(y=1.0, x=-2.0),
+    "detector_rotation": -0.4,
+    "flip_factor": -1.0,
+    "detector_center": PixelYX(y=-3.0, x=4.0),
+    "detector_pixel_pitch": 0.6,
+    "camera_length": 1.3,
+}
+
+
+@pytest.mark.parametrize("param", sorted(ADJUST_VALUES))
+def test_adjust_keeps_float_types(random_model, param):
+    adjusted = getattr(random_model, f"adjust_{param}")(ADJUST_VALUES[param])
+    for leaf in jax.tree.leaves(adjusted):
+        assert not isinstance(leaf, sym.Basic), leaf
+
+
+def test_invert_focus_keeps_float_types(random_model):
+    for leaf in jax.tree.leaves(random_model.invert_focus()):
+        assert not isinstance(leaf, sym.Basic), leaf
+
+
+@pytest.mark.parametrize(
+    "param", ["scan_rotation", "scan_pixel_pitch", "detector_rotation", "camera_length"]
+)
+def test_adjust_grad(random_model, param):
+    def f(value):
+        return sum(getattr(random_model, f"adjust_{param}")(value).descan_error)
+
+    value = ADJUST_VALUES[param]
+    eps = 1e-6
+    expected = (f(value + eps) - f(value - eps)) / (2 * eps)
+    assert_allclose(jax.grad(f)(value), expected, rtol=1e-5, atol=1e-8)
+
+
 def test_invert_focus(random_model):
     model = random_model
     trace = lambdify_trace_for(np)

From 3342ff2c489c4ddb65b2828d2b8eda8010268616 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:48:19 +0000
Subject: [PATCH 04/34] Remove value check from adjust_flip_factor() so that it
 works with jax.jit()

The derived compensation is the identity for an unchanged flip factor anyway.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/common/model.py | 6 +-----
 tests/test_model.py                        | 8 ++++++++
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/microscope_calibration/common/model.py b/src/microscope_calibration/common/model.py
index 106f62d..c3f1908 100644
--- a/src/microscope_calibration/common/model.py
+++ b/src/microscope_calibration/common/model.py
@@ -564,12 +564,8 @@ def adjust_flip_factor(self, flip_factor: float) -> Model4DSTEM:
         compensation the same.
 
         This allows first compensating descan error and then adjusting other parameters.
-        Returns this model unchanged if the flip factor is the same.
         """
-        if flip_factor != self.flip_factor:
-            return self._adjust_compensating_descan_error("flip_factor", flip_factor)
-        else:
-            return self
+        return self._adjust_compensating_descan_error("flip_factor", flip_factor)
 
     def adjust_detector_center(self, detector_center: PixelYX) -> Model4DSTEM:
         """
diff --git a/tests/test_model.py b/tests/test_model.py
index bb29cc2..f40d0f5 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1118,6 +1118,14 @@ def f(value):
     assert_allclose(jax.grad(f)(value), expected, rtol=1e-5, atol=1e-8)
 
 
+@pytest.mark.parametrize("flip_factor", [-1.0, 1.0])
+def test_adjust_flip_factor_jit(random_model, flip_factor):
+    def f(flip_factor):
+        return random_model.adjust_flip_factor(flip_factor).descan_error
+
+    assert_allclose(jax.jit(f)(flip_factor), f(flip_factor), rtol=1e-12, atol=1e-12)
+
+
 def test_invert_focus(random_model):
     model = random_model
     trace = lambdify_trace_for(np)

From dd69cc8d1f6ace2435ada0471b8155c6f3f5e3d8 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:48:43 +0000
Subject: [PATCH 05/34] Fix lambdify(args=...), which always raised a TypeError

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/util/sympy.py |  2 +-
 tests/test_sympy.py                      | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_sympy.py

diff --git a/src/microscope_calibration/util/sympy.py b/src/microscope_calibration/util/sympy.py
index 98f1d09..0ca4bb4 100644
--- a/src/microscope_calibration/util/sympy.py
+++ b/src/microscope_calibration/util/sympy.py
@@ -236,7 +236,7 @@ def lambdified(wrapped, instance, args, kwargs):
         f = _outer_lambdify(
             wrapped=wrapped,
             modules=modules,
-            sample_args=frozendict(sample_args),
+            sample_args=tuple(sample_args),
             sample_kwargs=frozendict(sample_kwargs),
             sym_kwargs=frozendict(sym_kwargs),
             recurse_for=recurse_for,
diff --git a/tests/test_sympy.py b/tests/test_sympy.py
new file mode 100644
index 0000000..ace2f6c
--- /dev/null
+++ b/tests/test_sympy.py
@@ -0,0 +1,20 @@
+import numpy as np
+import pytest
+import sympy as sym
+
+from microscope_calibration.util.sympy import lambdify
+
+
+@pytest.mark.parametrize("sample", ["args", "kwargs"])
+def test_lambdify_sample_input(sample):
+    a = sym.Symbol("a")
+    sample_input = {"args": (a, 2.0)} if sample == "args" else {"kwargs": {"a": a, "b": 2.0}}
+
+    @lambdify(modules=np, **sample_input)
+    def f(a: float, b: float):
+        return a * b
+
+    assert f(3.0, 2.0) == 6.0
+    # The sample value for b is a constant
+    with pytest.raises(ValueError, match="Constant value"):
+        f(3.0, 4.0)

From b08c68dcd9932904ac54f177ce62660d470aaa65 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:51:21 +0000
Subject: [PATCH 06/34] Implement Model4DSTEM.derive() with
 dataclasses.replace()

normalize_types() maps float() over the pytree.

- Fields of subclasses are kept instead of being reset to their defaults
- derive() only accepts keyword arguments
- Specifying both flip_y and flip_factor raises ValueError instead of AssertionError

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/common/model.py | 92 +++++++---------------
 tests/test_model.py                        | 17 ++++
 tests/test_subclass.py                     | 13 +++
 3 files changed, 57 insertions(+), 65 deletions(-)

diff --git a/src/microscope_calibration/common/model.py b/src/microscope_calibration/common/model.py
index c3f1908..1a0352c 100644
--- a/src/microscope_calibration/common/model.py
+++ b/src/microscope_calibration/common/model.py
@@ -136,9 +136,8 @@ class Model4DSTEM:
     Subclasses have to be decorated with
     :func:`jax_dataclasses.pytree_dataclass` like :class:`Model4DSTEM` so that
     JAX and :func:`~microscope_calibration.util.sympy.lambdify` see their
-    parameters. They can't add fields: the symbolic derivation only creates
-    symbols for the fields of :class:`Model4DSTEM`, and :meth:`derive` only
-    passes them on.
+    parameters. They can't add fields that affect :meth:`trace`: the symbolic
+    derivation only creates symbols for the fields of :class:`Model4DSTEM`.
     """
 
     overfocus: float  # m
@@ -205,6 +204,7 @@ def default(cls, dataset_shape: tuple[int, int, int, int] = (0, 0, 0, 0)):
 
     def derive(
         self,
+        *,
         overfocus: float | None = None,  # m
         scan_pixel_pitch: float | None = None,  # m
         scan_center: PixelYX | None = None,
@@ -217,70 +217,32 @@ def derive(
         flip_y: bool | None = None,
         flip_factor: float | None = None,
         descan_error: DescanError | None = None,
-    ) -> Model4DSTEM:
-        if flip_factor is not None:
-            assert flip_y is None
-        if flip_y is not None:
-            flip_factor = -1.0 if flip_y else 1.0
+    ) -> Self:
+        """
+        Copy of this model with the given parameters changed
 
-        # Keep the class of subclasses
-        return type(self)(
-            overfocus=overfocus if overfocus is not None else self.overfocus,
-            scan_pixel_pitch=(
-                scan_pixel_pitch if scan_pixel_pitch is not None else self.scan_pixel_pitch
-            ),
-            scan_center=scan_center if scan_center is not None else self.scan_center,
-            scan_rotation=scan_rotation if scan_rotation is not None else self.scan_rotation,
-            camera_length=camera_length if camera_length is not None else self.camera_length,
-            detector_pixel_pitch=(
-                detector_pixel_pitch
-                if detector_pixel_pitch is not None
-                else self.detector_pixel_pitch
-            ),
-            detector_center=(
-                detector_center if detector_center is not None else self.detector_center
-            ),
-            detector_rotation=(
-                detector_rotation if detector_rotation is not None else self.detector_rotation
-            ),
-            semiconv=semiconv if semiconv is not None else self.semiconv,
-            flip_factor=flip_factor if flip_factor is not None else self.flip_factor,
-            descan_error=descan_error if descan_error is not None else self.descan_error,
-        )
+        Unlike the :code:`adjust_*()` methods, this doesn't compensate the
+        effect of the change with the descan error. :code:`flip_y` is an
+        alternative to :code:`flip_factor`.
+        """
+        changes = {
+            name: value
+            for name, value in locals().items()
+            if name not in ("self", "flip_y") and value is not None
+        }
+        if flip_y is not None:
+            if flip_factor is not None:
+                raise ValueError("Specify either flip_y or flip_factor, not both.")
+            changes["flip_factor"] = -1.0 if flip_y else 1.0
+        # Keeps the class and any other fields of subclasses
+        return dataclasses.replace(self, **changes)
 
-    def normalize_types(self) -> Model4DSTEM:
-        return self.derive(
-            overfocus=float(self.overfocus),
-            scan_pixel_pitch=float(self.scan_pixel_pitch),
-            scan_center=PixelYX(
-                y=float(self.scan_center.y),
-                x=float(self.scan_center.x),
-            ),
-            scan_rotation=float(self.scan_rotation),
-            camera_length=float(self.camera_length),
-            detector_pixel_pitch=float(self.detector_pixel_pitch),
-            detector_center=PixelYX(
-                y=float(self.detector_center.y),
-                x=float(self.detector_center.x),
-            ),
-            detector_rotation=float(self.detector_rotation),
-            semiconv=float(self.semiconv),
-            flip_factor=float(self.flip_factor),
-            descan_error=DescanError(
-                pxo_pyi=float(self.descan_error.pxo_pyi),
-                pyo_pyi=float(self.descan_error.pyo_pyi),
-                pxo_pxi=float(self.descan_error.pxo_pxi),
-                pyo_pxi=float(self.descan_error.pyo_pxi),
-                sxo_pyi=float(self.descan_error.sxo_pyi),
-                syo_pyi=float(self.descan_error.syo_pyi),
-                sxo_pxi=float(self.descan_error.sxo_pxi),
-                syo_pxi=float(self.descan_error.syo_pxi),
-                offpxi=float(self.descan_error.offpxi),
-                offpyi=float(self.descan_error.offpyi),
-                offsxi=float(self.descan_error.offsxi),
-                offsyi=float(self.descan_error.offsyi),
-            ),
-        )
+    def normalize_types(self) -> Self:
+        """
+        Copy of this model with all parameters converted to :code:`float`,
+        e.g. from SymPy numbers or JAX and NumPy scalars
+        """
+        return jax.tree.map(float, self)
 
     def scan_to_real(self, pixels: PixelYX, _one: float = 1.0) -> CoordXY:
         return rotate(
diff --git a/tests/test_model.py b/tests/test_model.py
index f40d0f5..ae9b5bb 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1126,6 +1126,23 @@ def f(flip_factor):
     assert_allclose(jax.jit(f)(flip_factor), f(flip_factor), rtol=1e-12, atol=1e-12)
 
 
+def test_derive_flip_y(random_model):
+    assert random_model.derive(flip_y=True).flip_factor == -1.0
+    assert random_model.derive(flip_y=False).flip_factor == 1.0
+    with pytest.raises(ValueError, match="flip_y or flip_factor"):
+        random_model.derive(flip_y=True, flip_factor=-1.0)
+
+
+def test_normalize_types(random_model):
+    model = random_model.derive(
+        overfocus=jnp.float64(0.5),
+        scan_center=PixelYX(y=np.float32(1.0), x=sym.Float(2.0)),
+    ).normalize_types()
+    leaves = jax.tree.leaves(model)
+    assert len(leaves) == 24
+    assert all(type(leaf) is float for leaf in leaves)
+
+
 def test_invert_focus(random_model):
     model = random_model
     trace = lambdify_trace_for(np)
diff --git a/tests/test_subclass.py b/tests/test_subclass.py
index d33c85e..c39293c 100644
--- a/tests/test_subclass.py
+++ b/tests/test_subclass.py
@@ -53,6 +53,12 @@ class UndecoratedModel(Model4DSTEM):
     pass
 
 
+@jdc.pytree_dataclass
+class LabeledModel(CompatibleModel):
+    # Static field that doesn't affect trace()
+    label: jdc.Static[str] = "unlabeled"
+
+
 def make_model(model_type=Model4DSTEM) -> Model4DSTEM:
     return model_type(
         overfocus=0.7,
@@ -115,6 +121,13 @@ def test_derive_keeps_class():
     assert type(model.normalize_types()) is ShiftedDetectorModel
 
 
+def test_derive_keeps_extra_fields():
+    model = LabeledModel(**vars(make_model()), label="sample A")
+    assert model.derive(overfocus=0.1).label == "sample A"
+    assert model.normalize_types().label == "sample A"
+    assert model.adjust_scan_rotation(0.1).label == "sample A"
+
+
 def test_ray_transfer_runtime_derivation():
     model = make_model(ShiftedDetectorModel)
     with pytest.warns(RuntimeDerivationWarning, match="ray_transfer_matrix"):

From 7f3ce2f5d8af76a94a9672f994bc97bed32c33f0 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:51:57 +0000
Subject: [PATCH 07/34] Pass _one from trace() to _components() for consistency

No change in results since the components scale offsets with Ray._one.
The new test checks that trace() is homogeneous in _one.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/common/model.py |  1 +
 tests/test_model.py                        | 12 ++++++++++++
 2 files changed, 13 insertions(+)

diff --git a/src/microscope_calibration/common/model.py b/src/microscope_calibration/common/model.py
index 1a0352c..c06510e 100644
--- a/src/microscope_calibration/common/model.py
+++ b/src/microscope_calibration/common/model.py
@@ -354,6 +354,7 @@ def trace(
         components = self._components(
             scan_pos=scan_pos,
             specimen=specimen,
+            _one=_one,
         )
         source = components["source"]
         ray = Ray(
diff --git a/tests/test_model.py b/tests/test_model.py
index ae9b5bb..0141444 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1143,6 +1143,18 @@ def test_normalize_types(random_model):
     assert all(type(leaf) is float for leaf in leaves)
 
 
+def test_trace_homogeneous(random_model):
+    # With _one as homogeneous coordinate, all constant offsets scale with it
+    scan_y, scan_x, dy, dx, one = sym.symbols("scan_y scan_x dy dx one")
+    res = random_model.trace(
+        scan_pos=PixelYX(y=scan_y, x=scan_x), source_dy=dy, source_dx=dx, _one=one
+    )
+    zero = {scan_y: 0, scan_x: 0, dy: 0, dx: 0, one: 0}
+    for px in (res["specimen"].sampling["scan_px"], res["detector"].sampling["detector_px"]):
+        for coord in px:
+            assert sym.expand(coord).subs(zero) == 0
+
+
 def test_invert_focus(random_model):
     model = random_model
     trace = lambdify_trace_for(np)

From bf4a7dfbd33145c5ad4da21bfe788b2370833144 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:53:16 +0000
Subject: [PATCH 08/34] Pass fix points to solvers as float64 instead of
 float32

The UI downcast the point tables. The solvers now convert their input points to float64.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/ui.py            | 4 ++--
 src/microscope_calibration/util/optimize.py | 2 ++
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/microscope_calibration/ui.py b/src/microscope_calibration/ui.py
index 5545d1a..736455f 100644
--- a/src/microscope_calibration/ui.py
+++ b/src/microscope_calibration/ui.py
@@ -964,7 +964,7 @@ def detector_center_update(detector_center: PixelYX):
         }
 
     def perform_descan_update(self):
-        points = self.descan_fixpoint_table.value.reset_index().to_numpy().astype(np.float32)
+        points = self.descan_fixpoint_table.value.reset_index().to_numpy(dtype=np.float64)
         if len(points):
             new_model, residual = solve_tilt_descan_error_points(
                 ref_model=self.model, points=jnp.array(points)
@@ -972,7 +972,7 @@ def perform_descan_update(self):
             self.update_with_force(self.model_model, self.model_update(new_model))
 
     def perform_coord_update(self):
-        points = self.coord_fixpoint_table.value.reset_index().to_numpy().astype(np.float32)
+        points = self.coord_fixpoint_table.value.reset_index().to_numpy(dtype=np.float64)
         if len(points):
             new_model, residual = solve_coords_points(
                 ref_model=self.model, points=jnp.array(points)
diff --git a/src/microscope_calibration/util/optimize.py b/src/microscope_calibration/util/optimize.py
index fd41bed..7d1708a 100644
--- a/src/microscope_calibration/util/optimize.py
+++ b/src/microscope_calibration/util/optimize.py
@@ -609,6 +609,7 @@ def _de_tilt_point_loss(y, args: _DETiltPointArgs):
 
 def solve_tilt_descan_error_points(ref_model: Model4DSTEM, points: jnp.ndarray):
     check_float64()
+    points = jnp.asarray(points, dtype=jnp.float64)
     args = _DETiltPointArgs(
         model=ref_model,
         points=points,
@@ -699,6 +700,7 @@ def _coords_point_loss(y, args: _CoordPointArgs):
 
 def solve_coords_points(ref_model: Model4DSTEM, points: jnp.ndarray):
     check_float64()
+    points = jnp.asarray(points, dtype=jnp.float64)
     args = _CoordPointArgs(
         model=ref_model,
         points=points,

From 7be6c929005c21ab64bfba33d36bf2eb9a38b975 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:55:33 +0000
Subject: [PATCH 09/34] Fit both flip factors in solve_coords_points() instead
 of a continuous one

The flip factor came back as e.g. 0.99, and a start with the wrong sign
could converge to the wrong orientation.

- Return the residual without the flip penalty term
- Raise RuntimeError if the better fit didn't converge
- Add docstring and a round-trip test

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/util/optimize.py | 71 +++++++++++++--------
 tests/test_optimize.py                      | 30 +++++++++
 2 files changed, 76 insertions(+), 25 deletions(-)

diff --git a/src/microscope_calibration/util/optimize.py b/src/microscope_calibration/util/optimize.py
index 7d1708a..a155c83 100644
--- a/src/microscope_calibration/util/optimize.py
+++ b/src/microscope_calibration/util/optimize.py
@@ -638,7 +638,8 @@ def solve_tilt_descan_error_points(ref_model: Model4DSTEM, points: jnp.ndarray):
 class _CoordPointArgs(NamedTuple):
     model: Model4DSTEM
     # [(scan_y, scan_x, specimen_y (px), specimen_x (px), detector_y, detector_x) * n]
-    points: jnp.array
+    points: jax.Array
+    flip_factor: float
 
 
 @jax.jit
@@ -661,10 +662,10 @@ def _coords_point_modelupdate(
 
 @jax.jit
 def _coords_point_loss(y, args: _CoordPointArgs):
-    detector_rotation, overfocus, flip_factor = y[:3]
+    detector_rotation, overfocus = y[:2]
     # We need to co-optimize the source tilts that hit the detector pixel for
     # each point set because we can't easily calculate that directly
-    tilts = y[3:]
+    tilts = y[2:]
 
     assert len(tilts) == 2 * len(args.points)
 
@@ -672,7 +673,7 @@ def _coords_point_loss(y, args: _CoordPointArgs):
         model=args.model,
         overfocus=overfocus,
         detector_rotation=detector_rotation,
-        flip_factor=flip_factor,
+        flip_factor=args.flip_factor,
     )
 
     detector_distances = []
@@ -693,33 +694,53 @@ def _coords_point_loss(y, args: _CoordPointArgs):
             )
         )
 
-    # Strongly encourage flip_factor in (-1., 1.)
-    aspect = jnp.array((jnp.abs(jnp.abs(flip_factor) - 1) * 100,))
-    return jnp.concatenate((jnp.array(detector_distances), jnp.array(specimen_distances), aspect))
+    return jnp.concatenate((jnp.array(detector_distances), jnp.array(specimen_distances)))
 
 
 def solve_coords_points(ref_model: Model4DSTEM, points: jnp.ndarray):
+    """
+    Fit overfocus, detector rotation and flip factor to matched points
+
+    The flip factor is either 1 or -1. Both are fitted, and the better fit is returned.
+
+    Parameters
+    ----------
+    ref_model
+        Starting values. The descan error is adjusted to compensate the
+        changes, see :meth:`Model4DSTEM.adjust_detector_rotation`.
+    points
+        Array with rows :code:`(scan_y, scan_x, specimen_y, specimen_x,
+        detector_y, detector_x)` in pixels: A ray through scan position
+        :code:`(scan_y, scan_x)` that passes the specimen at the given scan
+        pixel coordinates hits the detector at the given detector pixel.
+
+    Returns
+    -------
+    model, residual
+        Fitted model and the residual of the fitted points in pixels
+    """
     check_float64()
     points = jnp.asarray(points, dtype=jnp.float64)
-    args = _CoordPointArgs(
-        model=ref_model,
-        points=points,
-    )
-
-    start = jnp.array(
-        [ref_model.detector_rotation, ref_model.overfocus, ref_model.flip_factor]
-        + [0.0, 0.0] * len(points)
-    )
-    opt_res = optimistix.least_squares(
-        fn=_coords_point_loss,
-        args=args,
-        solver=optimistix.BFGS(atol=1e-12, rtol=1e-12),
-        y0=start,
-        max_steps=10000,
-    )
-    residual = _coords_point_loss(opt_res.value, args)
+    start = jnp.array([ref_model.detector_rotation, ref_model.overfocus] + [0.0, 0.0] * len(points))
+    results = []
+    for flip_factor in (1.0, -1.0):
+        args = _CoordPointArgs(model=ref_model, points=points, flip_factor=flip_factor)
+        opt_res = optimistix.least_squares(
+            fn=_coords_point_loss,
+            args=args,
+            solver=optimistix.BFGS(atol=1e-12, rtol=1e-12),
+            y0=start,
+            max_steps=10000,
+            # Can fail to converge for the wrong flip factor
+            throw=False,
+        )
+        residual = _coords_point_loss(opt_res.value, args)
+        results.append((float(jnp.sum(residual**2)), flip_factor, opt_res, residual))
+    _, flip_factor, opt_res, residual = min(results, key=lambda r: r[0])
+    if not bool(opt_res.result == optimistix.RESULTS.successful):
+        raise RuntimeError(f"Fit did not converge: {optimistix.RESULTS[opt_res.result]}")
     # Write the new values into the previous parameters
-    detector_rotation, overfocus, flip_factor = opt_res.value[:3]
+    detector_rotation, overfocus = opt_res.value[:2]
     res_model = _coords_point_modelupdate(
         model=ref_model,
         overfocus=overfocus,
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index 1214b3c..6c4ca56 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -15,6 +15,7 @@
     normalize_descan_error,
     optimize,
     solve_camera_length,
+    solve_coords_points,
     solve_full_descan_error,
     solve_scan_pixel_pitch,
     solve_tilt_descan_error,
@@ -700,3 +701,32 @@ def test_solver_checks_float64(jax_float32):
     model = Model4DSTEM.default(dataset_shape=(32, 32, 32, 32)).derive(overfocus=0.1)
     with pytest.raises(RuntimeError, match="No float64 support"):
         solve_camera_length(ref_model=model, diffraction_angle=0.01, radius_px=10.0)
+
+
+@pytest.mark.parametrize("flip_factor", [1.0, -1.0])
+@pytest.mark.parametrize("ref_flip_factor", [1.0, -1.0])
+def test_coords_points(flip_factor, ref_flip_factor):
+    ref_model = Model4DSTEM.default(dataset_shape=(32, 32, 64, 64)).derive(
+        overfocus=0.01,
+        scan_rotation=0.4,
+        flip_factor=ref_flip_factor,
+        descan_error=DescanError(pxo_pxi=0.1, syo_pyi=-0.2, offpxi=1e-5),
+    )
+    true_model = (
+        ref_model.derive(overfocus=0.012)
+        .adjust_detector_rotation(0.3)
+        .adjust_flip_factor(flip_factor)
+    )
+    points = []
+    for scan_y, scan_x, dy, dx in ((3, 4, 1e-3, -2e-3), (20, 7, -1e-3, 0), (12, 28, 2e-3, 1e-3)):
+        res = trace(true_model, scan_pos=PixelYX(y=scan_y, x=scan_x), source_dy=dy, source_dx=dx)
+        spec = res["specimen"].sampling["scan_px"]
+        det = res["detector"].sampling["detector_px"]
+        points.append((scan_y, scan_x, spec.y, spec.x, det.y, det.x))
+
+    res_model, residual = solve_coords_points(ref_model=ref_model, points=np.array(points))
+    assert res_model.flip_factor == flip_factor
+    assert_allclose(residual, 0.0, atol=1e-6)
+    assert_allclose(res_model.overfocus, true_model.overfocus, rtol=1e-6)
+    assert_allclose(res_model.detector_rotation, true_model.detector_rotation, atol=1e-6)
+    assert_allclose(res_model.descan_error, true_model.descan_error, atol=1e-9)

From 269ad1c0731efe138cbed32f3c7d3debdec4b069 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 14:58:25 +0000
Subject: [PATCH 10/34] Fit camera length with the model's diffraction geometry
 in solve_camera_length()

It traced tilted rays with zero overfocus, so the fitted camera length
included the overfocus.

Behavior changes:
- The parameter diffraction_angle is renamed to twotheta
- The result is the specimen-to-detector distance for any overfocus

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/util/optimize.py | 56 +++++++++++----------
 tests/test_optimize.py                      | 29 +++++------
 2 files changed, 42 insertions(+), 43 deletions(-)

diff --git a/src/microscope_calibration/util/optimize.py b/src/microscope_calibration/util/optimize.py
index a155c83..0058ca4 100644
--- a/src/microscope_calibration/util/optimize.py
+++ b/src/microscope_calibration/util/optimize.py
@@ -19,6 +19,7 @@
     PixelYX,
     lambdify_trace_for,
 )
+from microscope_calibration.common.stem_overfocus import get_diffraction_pixel_radius
 from microscope_calibration.util.sympy import lambdify
 # ruff: enable[E402]
 
@@ -195,47 +196,50 @@ def optimize(loss, bounds=None, minimizer_kwargs=None, **kwargs):
 
 class _CLArgs(NamedTuple):
     ref_model: Model4DSTEM
-    test_dx: float
+    twotheta: float
     radius_px: float
 
 
 @jax.jit
 def _cl_loss(y, args: _CLArgs):
-    opt_model = args.ref_model.derive(camera_length=y[0], overfocus=0.0)
-    opt_res_1 = trace(
-        opt_model,
-        scan_pos=PixelYX(y=0.0, x=0.0),
-        source_dx=args.test_dx,
-        source_dy=0.0,
-    )
-    opt_res_2 = trace(
-        opt_model,
-        scan_pos=PixelYX(y=0.0, x=0.0),
-        source_dx=-args.test_dx,
-        source_dy=0.0,
-    )
-    px_1 = opt_res_1["detector"].sampling["detector_px"]
-    px_2 = opt_res_2["detector"].sampling["detector_px"]
-    distance = jnp.linalg.norm(jnp.array(px_2) - jnp.array(px_1))
-    return distance - 2 * args.radius_px
+    opt_model = args.ref_model.derive(camera_length=y[0])
+    radius = get_diffraction_pixel_radius(model=opt_model, twotheta=args.twotheta)
+    return jnp.array((radius - args.radius_px,))
 
 
-# FIXME include wavelength calculation etc for more practical
-# input parameters
-def solve_camera_length(ref_model: Model4DSTEM, diffraction_angle, radius_px):
+def solve_camera_length(ref_model: Model4DSTEM, twotheta: float, radius_px: float):
+    """
+    Fit the camera length to the radius of a diffraction ring
+
+    Parameters
+    ----------
+    ref_model
+        Starting value. The other parameters are kept.
+    twotheta
+        Diffraction angle 2θ in radians, e.g. from
+        :func:`microscope_calibration.util.diffraction.get_twothetas`
+    radius_px
+        Observed distance between primary beam and diffracted beam on the
+        detector in pixels, see
+        :func:`~microscope_calibration.common.stem_overfocus.get_diffraction_pixel_radius`
+
+    Returns
+    -------
+    model, residual
+        Fitted model and the residual of the radius in pixels
+    """
     check_float64()
-    args = _CLArgs(radius_px=radius_px, test_dx=jnp.tan(diffraction_angle), ref_model=ref_model)
+    args = _CLArgs(ref_model=ref_model, twotheta=twotheta, radius_px=radius_px)
     start = jnp.array((ref_model.camera_length,))
     opt_res = optimistix.least_squares(
         fn=_cl_loss, args=args, solver=optimistix.BFGS(atol=1e-12, rtol=1e-12), y0=start
     )
-    residual = _cl_loss(opt_res.value, args)
     # The loss function has minima at camera_length and -camera_length.
     # we take the positive side since a negative camera length doesn't make sense
     # for a classical TEM, only for reflection.
-    return ref_model.derive(
-        camera_length=jnp.abs(opt_res.value[0]),
-    ).normalize_types(), residual
+    camera_length = jnp.abs(opt_res.value[0])
+    residual = _cl_loss(jnp.array((camera_length,)), args)
+    return ref_model.derive(camera_length=camera_length).normalize_types(), residual
 
 
 class _SPPArgs(NamedTuple):
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index 6c4ca56..c671441 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -242,39 +242,34 @@ def optstep(optargs, opt_state):
     assert_allclose(optargs, correct)
 
 
-def test_camera_length():
+@pytest.mark.parametrize("overfocus", [0.0, 0.01, -0.02])
+def test_camera_length(overfocus):
     # Determine camera length from a known diffraction angle in radians,
     # corresponding detector pixel offset, and detector pixel pitch
-    scan_pixel_pitch = 0.1
     detector_pixel_pitch = 0.2
-    overfocus = 0.01
     camera_length = 1.234
-    propagation_distance = overfocus + camera_length
-    obj_half_size = 16
     # This is known, e.g. from crystal structure, diffraction order and
     # wavelength
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.00314157, propagation_distance)
-    model = Model4DSTEM(
+    twotheta = 0.05
+    model = Model4DSTEM.default(dataset_shape=(32, 32, 64, 64)).derive(
         overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
-        scan_rotation=0.0,
-        flip_factor=1.0,
-        detector_center=PixelYX(x=2 * obj_half_size, y=2 * obj_half_size),
+        scan_rotation=0.3,
+        detector_rotation=-0.2,
+        descan_error=DescanError(pxo_pxi=0.1, sxo_pyi=0.02, offsyi=0.01),
     )
     # This is observed on the detector
-    px_radius = jnp.tan(angle) * propagation_distance / detector_pixel_pitch
+    px_radius = np.tan(twotheta) * camera_length / detector_pixel_pitch
 
     res, residual = solve_camera_length(
         # Start with a negative value on purpose
         ref_model=model.derive(camera_length=-2 * camera_length),
-        diffraction_angle=angle,
+        twotheta=twotheta,
         radius_px=px_radius,
     )
-    assert_allclose(res.camera_length, propagation_distance)
+    assert_allclose(res.camera_length, camera_length)
+    assert res.overfocus == overfocus
     assert_allclose(residual, 0.0, atol=1e-12)
 
 
@@ -700,7 +695,7 @@ def test_check_float64_float32(jax_float32):
 def test_solver_checks_float64(jax_float32):
     model = Model4DSTEM.default(dataset_shape=(32, 32, 32, 32)).derive(overfocus=0.1)
     with pytest.raises(RuntimeError, match="No float64 support"):
-        solve_camera_length(ref_model=model, diffraction_angle=0.01, radius_px=10.0)
+        solve_camera_length(ref_model=model, twotheta=0.01, radius_px=10.0)
 
 
 @pytest.mark.parametrize("flip_factor", [1.0, -1.0])

From 05e4121841cd08ec1857d1b42519f0e363650b9c Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 15:32:29 +0000
Subject: [PATCH 11/34] Reject strings in Model4DSTEM.normalize_types() and
 name the field in errors

float() accepted strings like "1.5", and its errors didn't say which parameter
couldn't be converted. NaN and infinity are kept.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/common/model.py | 38 +++++++++++++++++++++-
 tests/test_model.py                        | 27 +++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/microscope_calibration/common/model.py b/src/microscope_calibration/common/model.py
index c06510e..b345576 100644
--- a/src/microscope_calibration/common/model.py
+++ b/src/microscope_calibration/common/model.py
@@ -33,6 +33,27 @@ def equals(ray1: Ray, ray2: Ray) -> bool:
     return True
 
 
+def _to_float(name: str, value) -> float:
+    if isinstance(value, (str, bytes)):
+        raise TypeError(f"{name} is a string: {value!r}")
+    try:
+        return float(value)
+    except TypeError as e:
+        raise TypeError(f"Can't convert {name} = {value!r} to float: {e}") from e
+
+
+def _to_float_tree(name: str, value):
+    """
+    Convert :code:`value` or the leaves of a NamedTuple like :class:`PixelYX`
+    to :code:`float`, see :meth:`Model4DSTEM.normalize_types`
+    """
+    if isinstance(value, tuple) and hasattr(value, "_fields"):
+        return type(value)(
+            **{key: _to_float_tree(f"{name}.{key}", getattr(value, key)) for key in value._fields}
+        )
+    return _to_float(name, value)
+
+
 # Simple transformation function similar to propagation by components
 def scale(c: CoordXY | PixelYX, factor: float) -> CoordXY:
     return CoordXY(
@@ -241,8 +262,23 @@ def normalize_types(self) -> Self:
         """
         Copy of this model with all parameters converted to :code:`float`,
         e.g. from SymPy numbers or JAX and NumPy scalars
+
+        Fields that subclasses add are kept as they are.
+
+        Raises
+        ------
+        TypeError
+            If a parameter can't be converted to :code:`float`, e.g. a SymPy
+            expression with free symbols, a complex number, an array or a
+            string. NaN and infinity are kept.
         """
-        return jax.tree.map(float, self)
+        return dataclasses.replace(
+            self,
+            **{
+                field.name: _to_float_tree(field.name, getattr(self, field.name))
+                for field in dataclasses.fields(Model4DSTEM)
+            },
+        )
 
     def scan_to_real(self, pixels: PixelYX, _one: float = 1.0) -> CoordXY:
         return rotate(
diff --git a/tests/test_model.py b/tests/test_model.py
index 0141444..b232ac6 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1155,6 +1155,33 @@ def test_trace_homogeneous(random_model):
             assert sym.expand(coord).subs(zero) == 0
 
 
+@pytest.mark.parametrize(
+    "changes, exception, match",
+    [
+        ({"overfocus": sym.Symbol("x")}, TypeError, "overfocus = x"),
+        ({"scan_center": PixelYX(y=0.0, x=1 + sym.I)}, TypeError, "scan_center.x"),
+        ({"semiconv": np.array([1.0, 2.0])}, TypeError, "semiconv"),
+        ({"camera_length": "1.5"}, TypeError, "camera_length is a string"),
+        ({"descan_error": DescanError(offpxi=sym.zoo)}, TypeError, "descan_error.offpxi"),
+    ],
+)
+def test_normalize_types_invalid(random_model, changes, exception, match):
+    with pytest.raises(exception, match=match):
+        random_model.derive(**changes).normalize_types()
+
+
+def test_normalize_types_non_finite(random_model):
+    model = random_model.derive(
+        detector_rotation=np.nan,
+        scan_rotation=sym.oo,
+        descan_error=DescanError(sxo_pyi=-np.inf),
+    ).normalize_types()
+    assert np.isnan(model.detector_rotation)
+    assert model.scan_rotation == np.inf
+    assert model.descan_error.sxo_pyi == -np.inf
+    assert all(type(leaf) is float for leaf in jax.tree.leaves(model))
+
+
 def test_invert_focus(random_model):
     model = random_model
     trace = lambdify_trace_for(np)

From 1659c817bfac92d68c262ea9fbc6fbadce54ea94 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 15:46:15 +0000
Subject: [PATCH 12/34] Seed random test inputs per test and relax descan
 deviation tolerance

Failures with random models couldn't be reproduced. The seed depends on the
test ID. The pixel tolerance was at the level of float64 rounding.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/conftest.py   | 8 ++++++++
 tests/test_model.py | 8 ++++----
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index b0c8117..d91b2ad 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,9 +1,17 @@
+import zlib
+
 import numpy as np
 import pytest
 
 from microscope_calibration.common.model import DescanError, Model4DSTEM, PixelYX
 
 
+@pytest.fixture(autouse=True)
+def _seed_random(request):
+    # Reproducible random test inputs with a different seed for each test
+    np.random.seed(zlib.crc32(request.node.nodeid.encode()))
+
+
 @pytest.fixture
 def random_model() -> Model4DSTEM:
     return Model4DSTEM(
diff --git a/tests/test_model.py b/tests/test_model.py
index b232ac6..ee47df4 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -966,8 +966,8 @@ def distance(scan_y, scan_x, cl):
     for scan_y in (0, 1):
         for scan_x in (0, 1):
             for cl in (0, 1):
-                print(scan_y, scan_x, cl)
-                assert_allclose(0, distance(scan_y=scan_y, scan_x=scan_x, cl=cl), atol=1e-12)
+                # Pixel coordinates can be large, so allow for rounding errors
+                assert_allclose(0, distance(scan_y=scan_y, scan_x=scan_x, cl=cl), atol=1e-10)
 
 
 def assert_no_descan_deviation_cl(model, target_model):
@@ -992,8 +992,8 @@ def distance(scan_y, scan_x, cl):
     for scan_y in (0, 1):
         for scan_x in (0, 1):
             for cl in (0, 1):
-                print(scan_y, scan_x, cl)
-                assert_allclose(0, distance(scan_y=scan_y, scan_x=scan_x, cl=cl), atol=1e-12)
+                # Pixel coordinates can be large, so allow for rounding errors
+                assert_allclose(0, distance(scan_y=scan_y, scan_x=scan_x, cl=cl), atol=1e-10)
 
 
 def test_adjust_scan_rotation(random_model: Model4DSTEM):

From f0646b4d1fc22a3b2ae2f674fb5cf4003ef1e2b9 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 15:50:30 +0000
Subject: [PATCH 13/34] Share the least-squares call of the solvers in one
 helper

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/util/optimize.py | 85 ++++++++-------------
 1 file changed, 30 insertions(+), 55 deletions(-)

diff --git a/src/microscope_calibration/util/optimize.py b/src/microscope_calibration/util/optimize.py
index 0058ca4..1a3a272 100644
--- a/src/microscope_calibration/util/optimize.py
+++ b/src/microscope_calibration/util/optimize.py
@@ -60,6 +60,26 @@ def check_float64():
         )
 
 
+def _least_squares(fn, args, y0, tol: float = 1e-12, **kwargs):
+    """
+    Minimize the sum of squares of :code:`fn(y, args)` starting from :code:`y0`
+
+    Raises :class:`RuntimeError` if JAX doesn't use float64, see
+    :func:`check_float64`. Additional :code:`kwargs` are passed to
+    :func:`optimistix.least_squares`.
+
+    Returns
+    -------
+    solution, residual
+        Optimization result and :code:`fn(y, args)` for it
+    """
+    check_float64()
+    opt_res = optimistix.least_squares(
+        fn=fn, args=args, solver=optimistix.BFGS(atol=tol, rtol=tol), y0=y0, **kwargs
+    )
+    return opt_res, fn(opt_res.value, args)
+
+
 def make_overfocus_loss_function(
     model: Model4DSTEM,
     ctx: Context,
@@ -228,12 +248,9 @@ def solve_camera_length(ref_model: Model4DSTEM, twotheta: float, radius_px: floa
     model, residual
         Fitted model and the residual of the radius in pixels
     """
-    check_float64()
     args = _CLArgs(ref_model=ref_model, twotheta=twotheta, radius_px=radius_px)
     start = jnp.array((ref_model.camera_length,))
-    opt_res = optimistix.least_squares(
-        fn=_cl_loss, args=args, solver=optimistix.BFGS(atol=1e-12, rtol=1e-12), y0=start
-    )
+    opt_res, _ = _least_squares(fn=_cl_loss, args=args, y0=start)
     # The loss function has minima at camera_length and -camera_length.
     # we take the positive side since a negative camera length doesn't make sense
     # for a classical TEM, only for reflection.
@@ -266,7 +283,6 @@ def solve_scan_pixel_pitch(
     point_2: PixelYX,
     physical_distance: float,
 ):
-    check_float64()
     args = _SPPArgs(
         ref_model=ref_model,
         point_1=point_1,
@@ -274,13 +290,7 @@ def solve_scan_pixel_pitch(
         physical_distance=physical_distance,
     )
     start = jnp.array((ref_model.scan_pixel_pitch,))
-    opt_res = optimistix.least_squares(
-        fn=_spp_loss,
-        args=args,
-        solver=optimistix.BFGS(atol=1e-12, rtol=1e-12),
-        y0=start,
-    )
-    residual = _spp_loss(opt_res.value, args)
+    opt_res, residual = _least_squares(fn=_spp_loss, args=args, y0=start)
     # The loss function has minima at scan_pixel_pitch and -scan_pixel_pitch. we
     # take the positive side since the inversion can be better expressed with a
     # scan rotation.
@@ -344,7 +354,6 @@ def _de_full_loss(y, args: _DEFullArgs):
 
 
 def solve_full_descan_error(ref_model: Model4DSTEM, regressions: CoMRegressions):
-    check_float64()
     # Caveat: scan and detector center of ref_model and of regressions should
     # match.
 
@@ -363,13 +372,7 @@ def solve_full_descan_error(ref_model: Model4DSTEM, regressions: CoMRegressions)
     # Start with a small epsilon to prevent NaN results of yet unknown origin
     # for some parameter combinations
     start = jnp.full(shape=(len(DescanError()),), fill_value=1e-6)
-    opt_res = optimistix.least_squares(
-        fn=_de_full_loss,
-        args=args,
-        solver=optimistix.BFGS(atol=1e-12, rtol=1e-12),
-        y0=start,
-    )
-    residual = _de_full_loss(opt_res.value, args)
+    opt_res, residual = _least_squares(fn=_de_full_loss, args=args, y0=start)
 
     # Bring descan error back to original coordinate system
     res_model = (
@@ -448,7 +451,6 @@ def _norm_loss(y, args: _NormArgs):
 
 
 def normalize_descan_error(ref_model: Model4DSTEM):
-    check_float64()
     args = _NormArgs(
         ref_model=ref_model,
     )
@@ -460,13 +462,7 @@ def normalize_descan_error(ref_model: Model4DSTEM):
             ref_model.detector_center.x,
         )
     )
-    opt_res = optimistix.least_squares(
-        fn=_norm_loss,
-        args=args,
-        solver=optimistix.BFGS(atol=1e-12, rtol=1e-12),
-        y0=start,
-    )
-    residual = _norm_loss(opt_res.value, args)
+    opt_res, residual = _least_squares(fn=_norm_loss, args=args, y0=start)
     scy, scx, dcy, dcx = opt_res.value
     res_model = ref_model.derive(
         scan_center=PixelYX(y=scy, x=scx),
@@ -534,7 +530,6 @@ def _de_tilt_loss(y, args: _DETiltArgs):
 
 
 def solve_tilt_descan_error(ref_model: Model4DSTEM, regression: CoMRegression):
-    check_float64()
     # Caveat: scan and detector center of ref_model and of regressions should
     # match.
 
@@ -563,14 +558,7 @@ def solve_tilt_descan_error(ref_model: Model4DSTEM, regression: CoMRegression):
     args = _DETiltArgs(
         aligned_model=args.aligned_model.normalize_types(), regression=args.regression
     )
-    opt_res = optimistix.least_squares(
-        fn=_de_tilt_loss,
-        args=args,
-        solver=optimistix.BFGS(atol=1e-12, rtol=1e-12),
-        y0=start,
-        max_steps=10000,
-    )
-    residual = _de_tilt_loss(opt_res.value, args)
+    opt_res, residual = _least_squares(fn=_de_tilt_loss, args=args, y0=start, max_steps=10000)
 
     # Bring descan error back to original coordinate system
     res_model = (
@@ -612,7 +600,6 @@ def _de_tilt_point_loss(y, args: _DETiltPointArgs):
 
 
 def solve_tilt_descan_error_points(ref_model: Model4DSTEM, points: jnp.ndarray):
-    check_float64()
     points = jnp.asarray(points, dtype=jnp.float64)
     args = _DETiltPointArgs(
         model=ref_model,
@@ -622,16 +609,15 @@ def solve_tilt_descan_error_points(ref_model: Model4DSTEM, points: jnp.ndarray):
     # Start with a small epsilon to prevent NaN results of yet unknown origin
     # for some parameter combinations
     start = jnp.full(shape=(6,), fill_value=1e-6)
-    opt_res = optimistix.least_squares(
+    opt_res, residual = _least_squares(
         fn=_de_tilt_point_loss,
         args=args,
-        # FIXME Doesn't reach 1e-12 like the others, for unknown reasons
-        solver=optimistix.BFGS(atol=1e-11, rtol=1e-11),
         y0=start,
+        # FIXME Doesn't reach 1e-12 like the others, for unknown reasons
+        tol=1e-11,
         # FIXME needs more steps than others, for unknown reasons
         max_steps=10000,
     )
-    residual = _de_tilt_point_loss(opt_res.value, args)
     # Write the new tilts and offsets into the previous descan error
     new_descan = _tilt_descan(ref_model.descan_error, opt_res.value)
     res_model = ref_model.derive(descan_error=new_descan).normalize_types()
@@ -723,22 +709,19 @@ def solve_coords_points(ref_model: Model4DSTEM, points: jnp.ndarray):
     model, residual
         Fitted model and the residual of the fitted points in pixels
     """
-    check_float64()
     points = jnp.asarray(points, dtype=jnp.float64)
     start = jnp.array([ref_model.detector_rotation, ref_model.overfocus] + [0.0, 0.0] * len(points))
     results = []
     for flip_factor in (1.0, -1.0):
         args = _CoordPointArgs(model=ref_model, points=points, flip_factor=flip_factor)
-        opt_res = optimistix.least_squares(
+        opt_res, residual = _least_squares(
             fn=_coords_point_loss,
             args=args,
-            solver=optimistix.BFGS(atol=1e-12, rtol=1e-12),
             y0=start,
             max_steps=10000,
             # Can fail to converge for the wrong flip factor
             throw=False,
         )
-        residual = _coords_point_loss(opt_res.value, args)
         results.append((float(jnp.sum(residual**2)), flip_factor, opt_res, residual))
     _, flip_factor, opt_res, residual = min(results, key=lambda r: r[0])
     if not bool(opt_res.result == optimistix.RESULTS.successful):
@@ -777,7 +760,6 @@ class SlopeYX(NamedTuple):
 
 
 def solve_hit_specimen(model: Model4DSTEM, scan_pos: PixelYX, specimen_px: PixelYX):
-    check_float64()
     args = _HitSpecimenArgs(
         model=model,
         scan_pos=scan_pos,
@@ -785,14 +767,7 @@ def solve_hit_specimen(model: Model4DSTEM, scan_pos: PixelYX, specimen_px: Pixel
     )
 
     start = jnp.array([0.0, 0.0])
-    opt_res = optimistix.least_squares(
-        fn=_hit_specimen_loss,
-        args=args,
-        solver=optimistix.BFGS(atol=1e-12, rtol=1e-12),
-        y0=start,
-        max_steps=10000,
-    )
-    residual = _hit_specimen_loss(opt_res.value, args)
+    opt_res, residual = _least_squares(fn=_hit_specimen_loss, args=args, y0=start, max_steps=10000)
     # Write the new values into the previous parameters
     dy, dx = opt_res.value
 

From e391d0cb09d78d97259a131108a3425278d06aeb Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 15:52:06 +0000
Subject: [PATCH 14/34] Remove duplicated regression residual and descan error
 code from the solvers

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/util/optimize.py | 125 +++++++-------------
 1 file changed, 43 insertions(+), 82 deletions(-)

diff --git a/src/microscope_calibration/util/optimize.py b/src/microscope_calibration/util/optimize.py
index 1a3a272..259c3a6 100644
--- a/src/microscope_calibration/util/optimize.py
+++ b/src/microscope_calibration/util/optimize.py
@@ -312,9 +312,34 @@ def solve_scan_pixel_pitch(
 CoMRegressions = dict[float, CoMRegression]
 
 
+def _regression_residuals(model: Model4DSTEM, reg: CoMRegression) -> list:
+    """
+    Deviation of the central ray on the detector from a CoM regression, in
+    pixels, at the corners of a scan pixel
+    """
+    (dy, dx), (dydy, dxdy), (dydx, dxdx) = reg
+    distances = []
+    for scan_y in (0.0, 1.0):
+        for scan_x in (0.0, 1.0):
+            det_y = model.detector_center.y + (dy + dydy * scan_y + dydx * scan_x)
+            det_x = model.detector_center.x + (dx + dxdy * scan_y + dxdx * scan_x)
+            res = trace(
+                model,
+                scan_pos=PixelYX(y=scan_y, x=scan_x),
+                source_dx=0.0,
+                source_dy=0.0,
+            )
+            distances.extend(
+                (
+                    det_y - res["detector"].sampling["detector_px"].y,
+                    det_x - res["detector"].sampling["detector_px"].x,
+                )
+            )
+    return distances
+
+
 class _DEFullArgs(NamedTuple):
-    # Aligned with the CoM regression coordinate system.
-    # Currently only tested for no scan rotation and no flip_y
+    # Aligned with the CoM regression coordinate system
     aligned_model: Model4DSTEM
     regressions: CoMRegressions
 
@@ -328,28 +353,7 @@ def _de_full_loss(y, args: _DEFullArgs):
             camera_length=cl,
             descan_error=de,
         )
-        for scan_y in (0.0, 1.0):
-            for scan_x in (0.0, 1.0):
-                dy = reg[0, 0]
-                dx = reg[0, 1]
-                dydy = reg[1, 0]
-                dxdy = reg[1, 1]
-                dydx = reg[2, 0]
-                dxdx = reg[2, 1]
-                det_y = opt_model.detector_center.y + (dy + dydy * scan_y + dydx * scan_x)
-                det_x = opt_model.detector_center.x + (dx + dxdy * scan_y + dxdx * scan_x)
-                res = trace(
-                    opt_model,
-                    scan_pos=PixelYX(y=scan_y, x=scan_x),
-                    source_dx=0.0,
-                    source_dy=0.0,
-                )
-                distances.extend(
-                    (
-                        det_y - res["detector"].sampling["detector_px"].y,
-                        det_x - res["detector"].sampling["detector_px"].x,
-                    )
-                )
+        distances.extend(_regression_residuals(opt_model, reg))
     return jnp.array(distances)
 
 
@@ -391,20 +395,7 @@ class _NormArgs(NamedTuple):
 
 
 def _zero_const(de: DescanError) -> DescanError:
-    return DescanError(
-        pxo_pxi=de.pxo_pxi,
-        pxo_pyi=de.pxo_pyi,
-        pyo_pxi=de.pyo_pxi,
-        pyo_pyi=de.pyo_pyi,
-        sxo_pxi=de.sxo_pxi,
-        sxo_pyi=de.sxo_pyi,
-        syo_pxi=de.syo_pxi,
-        syo_pyi=de.syo_pyi,
-        offpxi=0.0,
-        offpyi=0.0,
-        offsxi=0.0,
-        offsyi=0.0,
-    )
+    return de._replace(offpxi=0.0, offpyi=0.0, offsxi=0.0, offsyi=0.0)
 
 
 @jax.jit
@@ -473,26 +464,23 @@ def normalize_descan_error(ref_model: Model4DSTEM):
 
 
 class _DETiltArgs(NamedTuple):
-    # Aligned with the CoM regression coordinate system.
-    # Currently only tested for no scan rotation and no flip_y
+    # Aligned with the CoM regression coordinate system
     aligned_model: Model4DSTEM
     regression: CoMRegression
 
 
 def _tilt_descan(de: DescanError, y) -> DescanError:
-    return DescanError(
-        pxo_pxi=de.pxo_pxi,
-        pxo_pyi=de.pxo_pyi,
-        pyo_pxi=de.pyo_pxi,
-        pyo_pyi=de.pyo_pyi,
-        sxo_pxi=y[0],
-        sxo_pyi=y[1],
-        syo_pxi=y[2],
-        syo_pyi=y[3],
-        offpxi=de.offpxi,
-        offpyi=de.offpyi,
-        offsxi=y[4],
-        offsyi=y[5],
+    """
+    Replace the tilt components of :code:`de` with :code:`y`
+    """
+    sxo_pxi, sxo_pyi, syo_pxi, syo_pyi, offsxi, offsyi = y
+    return de._replace(
+        sxo_pxi=sxo_pxi,
+        sxo_pyi=sxo_pyi,
+        syo_pxi=syo_pxi,
+        syo_pyi=syo_pyi,
+        offsxi=offsxi,
+        offsyi=offsyi,
     )
 
 
@@ -502,31 +490,7 @@ def _de_tilt_loss(y, args: _DETiltArgs):
         descan_error=_tilt_descan(de=args.aligned_model.descan_error, y=y)
     )
 
-    distances = []
-    reg = args.regression
-    for scan_y in (0.0, 1.0):
-        for scan_x in (0.0, 1.0):
-            dy = reg[0, 0]
-            dx = reg[0, 1]
-            dydy = reg[1, 0]
-            dxdy = reg[1, 1]
-            dydx = reg[2, 0]
-            dxdx = reg[2, 1]
-            det_y = opt_model.detector_center.y + (dy + dydy * scan_y + dydx * scan_x)
-            det_x = opt_model.detector_center.x + (dx + dxdy * scan_y + dxdx * scan_x)
-            res = trace(
-                opt_model,
-                scan_pos=PixelYX(y=scan_y, x=scan_x),
-                source_dx=0.0,
-                source_dy=0.0,
-            )
-            distances.extend(
-                (
-                    det_y - res["detector"].sampling["detector_px"].y,
-                    det_x - res["detector"].sampling["detector_px"].x,
-                )
-            )
-    return jnp.array(distances)
+    return jnp.array(_regression_residuals(opt_model, args.regression))
 
 
 def solve_tilt_descan_error(ref_model: Model4DSTEM, regression: CoMRegression):
@@ -548,16 +512,13 @@ def solve_tilt_descan_error(ref_model: Model4DSTEM, regression: CoMRegression):
         )
     )
     args = _DETiltArgs(
-        aligned_model=aligned_model,
+        aligned_model=aligned_model.normalize_types(),
         regression=regression,
     )
 
     # Start with a small epsilon to prevent NaN results of yet unknown origin
     # for some parameter combinations
     start = jnp.full(shape=(6,), fill_value=1e-6)
-    args = _DETiltArgs(
-        aligned_model=args.aligned_model.normalize_types(), regression=args.regression
-    )
     opt_res, residual = _least_squares(fn=_de_tilt_loss, args=args, y0=start, max_steps=10000)
 
     # Bring descan error back to original coordinate system

From 1216c7abc32bf2d9d13766183f6ac625e7759fc1 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 15:55:10 +0000
Subject: [PATCH 15/34] Reject too few points in the point-based solvers
 instead of returning arbitrary values

Behavior changes:
- solve_tilt_descan_error_points() raises ValueError for fewer than 3 points
- solve_coords_points() raises ValueError without points

The UI test now adds three distinct descan points; it fitted from two before.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/util/optimize.py |  9 ++++--
 tests/test_optimize.py                      | 33 +++++++++++----------
 tests/test_ui.py                            |  4 +++
 3 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/src/microscope_calibration/util/optimize.py b/src/microscope_calibration/util/optimize.py
index 259c3a6..7fc8e5a 100644
--- a/src/microscope_calibration/util/optimize.py
+++ b/src/microscope_calibration/util/optimize.py
@@ -561,7 +561,10 @@ def _de_tilt_point_loss(y, args: _DETiltPointArgs):
 
 
 def solve_tilt_descan_error_points(ref_model: Model4DSTEM, points: jnp.ndarray):
-    points = jnp.asarray(points, dtype=jnp.float64)
+    points = jnp.asarray(points, dtype=jnp.float64).reshape((-1, 4))
+    # Six unknowns, two equations per point
+    if len(points) < 3:
+        raise ValueError(f"Need at least 3 points to fit the descan error, got {len(points)}.")
     args = _DETiltPointArgs(
         model=ref_model,
         points=points,
@@ -670,7 +673,9 @@ def solve_coords_points(ref_model: Model4DSTEM, points: jnp.ndarray):
     model, residual
         Fitted model and the residual of the fitted points in pixels
     """
-    points = jnp.asarray(points, dtype=jnp.float64)
+    points = jnp.asarray(points, dtype=jnp.float64).reshape((-1, 6))
+    if len(points) < 1:
+        raise ValueError("Need at least one point.")
     start = jnp.array([ref_model.detector_rotation, ref_model.overfocus] + [0.0, 0.0] * len(points))
     results = []
     for flip_factor in (1.0, -1.0):
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index c671441..d77301b 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -647,23 +647,18 @@ def test_tilt_descan_error_points(
         )
         detector_center = res["detector"].sampling["detector_px"]
         points.append((scan_y, scan_x, detector_center.y, detector_center.x))
-    opt_res, residual = solve_tilt_descan_error_points(
-        ref_model=model.derive(
-            # Blank out the tilt parts of the descan error
-            descan_error=_tilt_descan(de=model.descan_error, y=np.zeros(6)),
-        ),
-        points=points,
+    ref_model = model.derive(
+        # Blank out the tilt parts of the descan error
+        descan_error=_tilt_descan(de=model.descan_error, y=np.zeros(6)),
     )
-    default_attrs = ("pxo_pxi", "pxo_pyi", "pyo_pxi", "pyo_pyi", "offpxi", "offpyi")
-    opt_attrs = ("sxo_pxi", "syo_pyi", "syo_pxi", "sxo_pyi")
-    if works:
-        attrs = default_attrs + opt_attrs
-        # For some reason less accurate than in other tests
-        assert_allclose(residual, 0.0, atol=1e-8)
-    else:
-        attrs = default_attrs
-    for key in attrs:
-        print(key)
+    if not works:
+        with pytest.raises(ValueError, match="at least 3 points"):
+            solve_tilt_descan_error_points(ref_model=ref_model, points=points)
+        return
+    opt_res, residual = solve_tilt_descan_error_points(ref_model=ref_model, points=points)
+    # For some reason less accurate than in other tests
+    assert_allclose(residual, 0.0, atol=1e-8)
+    for key in model.descan_error._fields:
         assert_allclose(
             getattr(model.descan_error, key),
             getattr(opt_res.descan_error, key),
@@ -674,6 +669,12 @@ def test_tilt_descan_error_points(
         assert isinstance(getattr(opt_res.descan_error, key), float)
 
 
+def test_coords_points_no_points():
+    model = Model4DSTEM.default(dataset_shape=(8, 8, 8, 8)).derive(overfocus=0.01)
+    with pytest.raises(ValueError, match="at least one point"):
+        solve_coords_points(ref_model=model, points=[])
+
+
 def test_check_float64():
     check_float64()
 
diff --git a/tests/test_ui.py b/tests/test_ui.py
index 49b0daa..9a6a3d8 100644
--- a/tests/test_ui.py
+++ b/tests/test_ui.py
@@ -64,10 +64,14 @@ def test_events(kwargs):
     layout.move_feature_select_to(MockEvent(y=4, x=5))
 
     layout.descan_add_row()
+    # Replaces the previous row at the same scan position
     layout.descan_add_row()
     layout.move_scan_to(MockEvent(y=0, x=1))
     layout.move_rings_to(MockEvent(y=3, x=4))
     layout.descan_add_row()
+    layout.move_scan_to(MockEvent(y=3, x=0))
+    layout.descan_add_row()
+    assert len(layout.descan_fixpoint_table.value) == 3
     layout.descan_move_to_row(0)
     layout.perform_descan_update()
     layout.center_correlation_regression()

From 38f5d4649b4fed8d8bc82dabe067b7b4732af40c Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 15:59:00 +0000
Subject: [PATCH 16/34] Document the solvers and annotate their return types

- solve_hit_specimen() returns the residual in YX order and the slope as floats
- Add a round-trip test for solve_hit_specimen()

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/util/optimize.py | 136 ++++++++++++++++++--
 tests/test_optimize.py                      |  13 ++
 2 files changed, 136 insertions(+), 13 deletions(-)

diff --git a/src/microscope_calibration/util/optimize.py b/src/microscope_calibration/util/optimize.py
index 7fc8e5a..d1730b3 100644
--- a/src/microscope_calibration/util/optimize.py
+++ b/src/microscope_calibration/util/optimize.py
@@ -227,7 +227,9 @@ def _cl_loss(y, args: _CLArgs):
     return jnp.array((radius - args.radius_px,))
 
 
-def solve_camera_length(ref_model: Model4DSTEM, twotheta: float, radius_px: float):
+def solve_camera_length(
+    ref_model: Model4DSTEM, twotheta: float, radius_px: float
+) -> tuple[Model4DSTEM, jax.Array]:
     """
     Fit the camera length to the radius of a diffraction ring
 
@@ -282,7 +284,24 @@ def solve_scan_pixel_pitch(
     point_1: PixelYX,
     point_2: PixelYX,
     physical_distance: float,
-):
+) -> tuple[Model4DSTEM, jax.Array]:
+    """
+    Fit the scan pixel pitch to the known distance between two scan positions
+
+    Parameters
+    ----------
+    ref_model
+        Starting value. The other parameters are kept.
+    point_1, point_2
+        Scan positions in pixels
+    physical_distance
+        Distance between the two scan positions on the specimen in m
+
+    Returns
+    -------
+    model, residual
+        Fitted model and the residuals of the fit in m
+    """
     args = _SPPArgs(
         ref_model=ref_model,
         point_1=point_1,
@@ -357,7 +376,26 @@ def _de_full_loss(y, args: _DEFullArgs):
     return jnp.array(distances)
 
 
-def solve_full_descan_error(ref_model: Model4DSTEM, regressions: CoMRegressions):
+def solve_full_descan_error(
+    ref_model: Model4DSTEM, regressions: CoMRegressions
+) -> tuple[Model4DSTEM, jax.Array]:
+    """
+    Fit all components of the descan error to CoM regressions
+
+    Parameters
+    ----------
+    ref_model
+        Model with the same scan and detector center as the regressions. The
+        other parameters are kept.
+    regressions
+        CoM regressions of the central beam in detector pixels for at least
+        two different camera lengths, see :data:`CoMRegressions`
+
+    Returns
+    -------
+    model, residual
+        Fitted model and the residuals of the fit in detector pixels
+    """
     # Caveat: scan and detector center of ref_model and of regressions should
     # match.
 
@@ -441,7 +479,19 @@ def _norm_loss(y, args: _NormArgs):
     return jnp.array(distances)
 
 
-def normalize_descan_error(ref_model: Model4DSTEM):
+def normalize_descan_error(ref_model: Model4DSTEM) -> tuple[Model4DSTEM, jax.Array]:
+    """
+    Express the constant offsets of the descan error as scan and detector center
+
+    The descan error is overparameterized: Its constant offsets have the same
+    effect on the central beam as shifts of the scan and detector center.
+    This returns an equivalent model with the constant offsets set to zero.
+
+    Returns
+    -------
+    model, residual
+        Fitted model and the residuals of the fit in detector pixels
+    """
     args = _NormArgs(
         ref_model=ref_model,
     )
@@ -493,7 +543,29 @@ def _de_tilt_loss(y, args: _DETiltArgs):
     return jnp.array(_regression_residuals(opt_model, args.regression))
 
 
-def solve_tilt_descan_error(ref_model: Model4DSTEM, regression: CoMRegression):
+def solve_tilt_descan_error(
+    ref_model: Model4DSTEM, regression: CoMRegression
+) -> tuple[Model4DSTEM, jax.Array]:
+    """
+    Fit the tilt components of the descan error to a CoM regression
+
+    This allows a preliminary calibration of the descan error for a single
+    camera length. The position components of the descan error are kept.
+
+    Parameters
+    ----------
+    ref_model
+        Model with the same scan and detector center as the regression. The
+        other parameters are kept.
+    regression
+        CoM regression of the central beam in detector pixels, see
+        :data:`CoMRegression`
+
+    Returns
+    -------
+    model, residual
+        Fitted model and the residuals of the fit in detector pixels
+    """
     # Caveat: scan and detector center of ref_model and of regressions should
     # match.
 
@@ -536,7 +608,7 @@ def solve_tilt_descan_error(ref_model: Model4DSTEM, regression: CoMRegression):
 class _DETiltPointArgs(NamedTuple):
     model: Model4DSTEM
     # [(scan_y, scan_x, detector_cy, detector_cx) * n]
-    points: jnp.ndarray
+    points: jax.Array
 
 
 @jax.jit
@@ -560,7 +632,28 @@ def _de_tilt_point_loss(y, args: _DETiltPointArgs):
     return jnp.array(distances)
 
 
-def solve_tilt_descan_error_points(ref_model: Model4DSTEM, points: jnp.ndarray):
+def solve_tilt_descan_error_points(
+    ref_model: Model4DSTEM, points: jax.Array
+) -> tuple[Model4DSTEM, jax.Array]:
+    """
+    Fit the tilt components of the descan error to central beam positions
+
+    Like :func:`solve_tilt_descan_error`, but with the position of the
+    central beam at individual scan positions instead of a regression.
+
+    Parameters
+    ----------
+    ref_model
+        Starting value. The other parameters are kept.
+    points
+        Array with at least three rows :code:`(scan_y, scan_x, detector_y,
+        detector_x)` in pixels
+
+    Returns
+    -------
+    model, residual
+        Fitted model and the residuals of the fit in detector pixels
+    """
     points = jnp.asarray(points, dtype=jnp.float64).reshape((-1, 4))
     # Six unknowns, two equations per point
     if len(points) < 3:
@@ -651,7 +744,7 @@ def _coords_point_loss(y, args: _CoordPointArgs):
     return jnp.concatenate((jnp.array(detector_distances), jnp.array(specimen_distances)))
 
 
-def solve_coords_points(ref_model: Model4DSTEM, points: jnp.ndarray):
+def solve_coords_points(ref_model: Model4DSTEM, points: jax.Array) -> tuple[Model4DSTEM, jax.Array]:
     """
     Fit overfocus, detector rotation and flip factor to matched points
 
@@ -717,7 +810,7 @@ def _hit_specimen_loss(y, args: _HitSpecimenArgs):
 
     res = trace(opt_model, scan_pos=args.scan_pos, source_dx=dx, source_dy=dy)
     specimen_px = res["specimen"].sampling["scan_px"]
-    return jnp.array((specimen_px.x - args.specimen_px.x, specimen_px.y - args.specimen_px.y))
+    return jnp.array((specimen_px.y - args.specimen_px.y, specimen_px.x - args.specimen_px.x))
 
 
 class SlopeYX(NamedTuple):
@@ -725,7 +818,26 @@ class SlopeYX(NamedTuple):
     dx: float
 
 
-def solve_hit_specimen(model: Model4DSTEM, scan_pos: PixelYX, specimen_px: PixelYX):
+def solve_hit_specimen(
+    model: Model4DSTEM, scan_pos: PixelYX, specimen_px: PixelYX
+) -> tuple[SlopeYX, jax.Array]:
+    """
+    Find the source tilt of the ray that passes the specimen at a given point
+
+    Parameters
+    ----------
+    model
+        Model to trace
+    scan_pos
+        Scan position in pixels
+    specimen_px
+        Point on the specimen in scan pixel coordinates
+
+    Returns
+    -------
+    slope, residual
+        Source tilt and the residual on the specimen in scan pixels
+    """
     args = _HitSpecimenArgs(
         model=model,
         scan_pos=scan_pos,
@@ -734,7 +846,5 @@ def solve_hit_specimen(model: Model4DSTEM, scan_pos: PixelYX, specimen_px: Pixel
 
     start = jnp.array([0.0, 0.0])
     opt_res, residual = _least_squares(fn=_hit_specimen_loss, args=args, y0=start, max_steps=10000)
-    # Write the new values into the previous parameters
     dy, dx = opt_res.value
-
-    return SlopeYX(dy=dy, dx=dx), residual
+    return SlopeYX(dy=float(dy), dx=float(dx)), residual
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index d77301b..16ba99a 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -17,6 +17,7 @@
     solve_camera_length,
     solve_coords_points,
     solve_full_descan_error,
+    solve_hit_specimen,
     solve_scan_pixel_pitch,
     solve_tilt_descan_error,
     solve_tilt_descan_error_points,
@@ -726,3 +727,15 @@ def test_coords_points(flip_factor, ref_flip_factor):
     assert_allclose(res_model.overfocus, true_model.overfocus, rtol=1e-6)
     assert_allclose(res_model.detector_rotation, true_model.detector_rotation, atol=1e-6)
     assert_allclose(res_model.descan_error, true_model.descan_error, atol=1e-9)
+
+
+def test_hit_specimen(random_model):
+    scan_pos = PixelYX(y=3.0, x=-2.0)
+    dy, dx = 1e-3, -2e-3
+    res = trace(random_model, scan_pos=scan_pos, source_dy=dy, source_dx=dx)
+    specimen_px = res["specimen"].sampling["scan_px"]
+    slope, residual = solve_hit_specimen(
+        model=random_model, scan_pos=scan_pos, specimen_px=specimen_px
+    )
+    assert_allclose(residual, 0.0, atol=1e-9)
+    assert_allclose((slope.dy, slope.dx), (dy, dx), rtol=1e-6)

From 86c6fa39a304f1863162a41b019dc7f266415ec1 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 16:02:23 +0000
Subject: [PATCH 17/34] Calculate the UDF transformation matrices per task
 instead of in __init__

The matrices were recalculated in __init__ when the model changed, which
relied on how LiberTEM re-creates UDFs.

Behavior changes:
- The UDFs no longer accept back_mat, corr_mat and ref_model
- Models with JAX or NumPy scalars work instead of raising TypeError
- A singular matrix, e.g. the back-projection in focus, logs a warning
- Cached matrices are read-only

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CLAUDE.md                                     |   4 +-
 .../udf/stem_overfocus.py                     | 119 +++++++++---------
 tests/test_overfocus_udf.py                   |  47 ++++++-
 tests/test_transformation_matrices.py         |   6 +-
 4 files changed, 107 insertions(+), 69 deletions(-)

diff --git a/CLAUDE.md b/CLAUDE.md
index 61a877a..81bc896 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -90,8 +90,8 @@ Key layers, bottom-up:
 4. **`udf/stem_overfocus.py`** — LiberTEM UDFs that apply (3) over a real dataset:
    `OverfocusUDF` (reconstruct an overfocus/real-space image) and `CorrectedPickUDF`. Models
    are passed wrapped in a `dict` so they can be mutated in place (LiberTEM issue #1780); the
-   transformation matrices are cached via `lru_cache` and only recomputed when the model
-   changes.
+   transformation matrices are calculated in `get_task_data()` and cached per model via
+   `lru_cache`.
 
 5. **`util/stem_overfocus_sim.py`** — forward simulator (e.g. `smiley()` test object,
    `project()`) that generates synthetic 4D-STEM data from a model, used heavily in tests and
diff --git a/src/microscope_calibration/udf/stem_overfocus.py b/src/microscope_calibration/udf/stem_overfocus.py
index 4d1448b..115274a 100644
--- a/src/microscope_calibration/udf/stem_overfocus.py
+++ b/src/microscope_calibration/udf/stem_overfocus.py
@@ -19,48 +19,55 @@
 log = logging.getLogger(__name__)
 
 
+@lru_cache
+def _backward_matrix(model: Model4DSTEM) -> np.ndarray | None:
+    try:
+        return _read_only(get_backward_transformation_matrix(rec_model=model))
+    except FittingError as e:
+        log.warning("Skipping back-projection: %s", e)
+        return None
+
+
+@lru_cache
+def _correction_matrix(model: Model4DSTEM) -> np.ndarray | None:
+    try:
+        return _read_only(get_detector_correction_matrix(rec_model=model))
+    except FittingError as e:
+        log.warning("Skipping detector correction: %s", e)
+        return None
+
+
+def _read_only(mat: np.ndarray) -> np.ndarray:
+    # Cached and shared between calls
+    mat.setflags(write=False)
+    return mat
+
+
 class BaseCorrectionUDF(UDF):
-    def __init__(
-        self,
-        overfocus_model: dict,
-        *args,
-        back_mat=None,
-        corr_mat=None,
-        ref_model=None,
-        **kwargs,
-    ):
-        """
-        model are wrapped in a dict since that allows in-place updates, see
-        https://github.com/LiberTEM/LiberTEM/issues/1780
-        """
-        # Detect changes so that the mapping matrices are recalculated
-        if overfocus_model["model"] != ref_model:
-            back_mat = self._back_mat(rec_model=overfocus_model["model"])
-            corr_mat = self._corr_mat(rec_model=overfocus_model["model"])
-        super().__init__(
-            *args,
-            overfocus_model=overfocus_model,
-            back_mat=back_mat,
-            corr_mat=corr_mat,
-            ref_model=overfocus_model["model"],
-            **kwargs,
-        )
-
-    @staticmethod
-    @lru_cache
-    def _back_mat(*args, **kwargs):
-        try:
-            return get_backward_transformation_matrix(*args, **kwargs)
-        except FittingError:
-            return None
-
-    @staticmethod
-    @lru_cache
-    def _corr_mat(*args, **kwargs):
-        try:
-            return get_detector_correction_matrix(*args, **kwargs)
-        except FittingError:
-            return None
+    """
+    Base class for UDFs that use the transformation matrices of a model
+
+    Parameters
+    ----------
+    overfocus_model
+        :code:`{"model": model}` with a :class:`Model4DSTEM`. The model is
+        wrapped in a dict so that it can be updated in place, see
+        https://github.com/LiberTEM/LiberTEM/issues/1780. The matrices are
+        calculated from the current model for each task.
+
+    If a matrix can't be calculated, e.g. the back-projection in focus, the
+    UDF logs a warning and skips the corresponding results.
+    """
+
+    def __init__(self, overfocus_model: dict, **kwargs):
+        super().__init__(overfocus_model=overfocus_model, **kwargs)
+
+    def get_task_data(self):
+        model = self.params.overfocus_model["model"].normalize_types()
+        return {
+            "back_mat": _backward_matrix(model),
+            "corr_mat": _correction_matrix(model),
+        }
 
     def _get_fov(self):
         fov_size_y = int(self.meta.dataset_shape.nav[0])
@@ -69,11 +76,11 @@ def _get_fov(self):
 
     @property
     def has_correction(self):
-        return self.params.corr_mat is not None
+        return self.task_data.corr_mat is not None
 
     @property
     def has_backprojection(self):
-        return self.params.back_mat is not None
+        return self.task_data.back_mat is not None
 
 
 class OverfocusUDF(BaseCorrectionUDF):
@@ -97,7 +104,7 @@ def process_frame(self, frame):
                 source_semiconv=overfocus_model.semiconv,
                 scan_y=scan_y,
                 scan_x=scan_x,
-                mat=self.params.back_mat,
+                mat=self.task_data.back_mat,
                 image_out=self.results.backprojected_sum,
             )
         center = overfocus_model.detector_center
@@ -107,14 +114,14 @@ def process_frame(self, frame):
                 det_corr_x=center.x,
                 scan_y=scan_y,
                 scan_x=scan_x,
-                mat=self.params.corr_mat,
+                mat=self.task_data.corr_mat,
             )
             det_x = corrected_det_x(
                 det_corr_y=center.y,
                 det_corr_x=center.x,
                 scan_y=scan_y,
                 scan_x=scan_x,
-                mat=self.params.corr_mat,
+                mat=self.task_data.corr_mat,
             )
             det_y = int(np.round(det_y))
             det_x = int(np.round(det_x))
@@ -125,7 +132,7 @@ def process_frame(self, frame):
                 frame=frame,
                 scan_y=scan_y,
                 scan_x=scan_x,
-                mat=self.params.corr_mat,
+                mat=self.task_data.corr_mat,
                 detector_out=self.results.corrected_sum,
             )
 
@@ -139,20 +146,6 @@ def merge(self, dest, src):
 # to allow re-using the correction basics and
 # not run into inheritance complications
 class CorrectedPickUDF(BaseCorrectionUDF):
-    def __init__(
-        self,
-        overfocus_model,
-        back_mat=None,
-        corr_mat=None,
-        ref_model=None,
-    ):
-        super().__init__(
-            overfocus_model,
-            back_mat=back_mat,
-            corr_mat=corr_mat,
-            ref_model=ref_model,
-        )
-
     def get_preferred_input_dtype(self):
         ""
         return self.USE_NATIVE_DTYPE
@@ -196,7 +189,7 @@ def process_frame(self, frame):
                 frame=frame,
                 scan_y=scan_y,
                 scan_x=scan_x,
-                mat=self.params.corr_mat,
+                mat=self.task_data.corr_mat,
                 detector_out=self.results.corrected[sl][0],
             )
         if self.has_backprojection:
@@ -205,7 +198,7 @@ def process_frame(self, frame):
                 source_semiconv=overfocus_model.semiconv,
                 scan_y=scan_y,
                 scan_x=scan_x,
-                mat=self.params.back_mat,
+                mat=self.task_data.back_mat,
                 image_out=self.results.backprojected[sl][0],
             )
 
diff --git a/tests/test_overfocus_udf.py b/tests/test_overfocus_udf.py
index aaf18e2..47d7f41 100644
--- a/tests/test_overfocus_udf.py
+++ b/tests/test_overfocus_udf.py
@@ -1,5 +1,6 @@
 import jax
 import numpy as np
+import pytest
 from libertem.api import Context
 from numpy.testing import assert_allclose
 
@@ -15,7 +16,12 @@
     get_detector_correction_matrix,
     project_frame_backwards,
 )
-from microscope_calibration.udf.stem_overfocus import OverfocusUDF, CorrectedPickUDF
+from microscope_calibration.udf.stem_overfocus import (
+    CorrectedPickUDF,
+    OverfocusUDF,
+    _backward_matrix,
+    _correction_matrix,
+)
 
 trace = lambdify_trace_for(jax.numpy)
 
@@ -154,3 +160,42 @@ def test_corrected_pick_delayed():
     for key, value in res.items():
         print(key)
         assert_allclose(value.delayed_data.compute(), ref[key].data)
+
+
+def test_udf_model_update():
+    ctx = Context.make_with("inline")
+    data = np.random.random((4, 5, 6, 7))
+    ds = ctx.load("memory", data=data)
+    model = Model4DSTEM.default(dataset_shape=ds.shape).derive(overfocus=0.1)
+    overfocus_model = {"model": model}
+    udf = OverfocusUDF(overfocus_model=overfocus_model)
+    res_1 = ctx.run_udf(dataset=ds, udf=udf)
+    # Update in place, with JAX scalars that aren't hashable
+    overfocus_model["model"] = model.derive(scan_rotation=jax.numpy.float64(0.5))
+    res_2 = ctx.run_udf(dataset=ds, udf=udf)
+    ref = ctx.run_udf(
+        dataset=ds, udf=OverfocusUDF(overfocus_model={"model": model.derive(scan_rotation=0.5)})
+    )
+    assert not np.allclose(res_1["backprojected_sum"].data, res_2["backprojected_sum"].data)
+    assert_allclose(res_2["backprojected_sum"].data, ref["backprojected_sum"].data)
+
+
+def test_udf_in_focus(caplog):
+    _backward_matrix.cache_clear()
+    ctx = Context.make_with("inline")
+    data = np.random.random((4, 5, 6, 7))
+    ds = ctx.load("memory", data=data)
+    model = Model4DSTEM.default(dataset_shape=ds.shape)
+    assert model.overfocus == 0
+    res = ctx.run_udf(dataset=ds, udf=OverfocusUDF(overfocus_model={"model": model}))
+    assert "Skipping back-projection" in caplog.text
+    assert np.all(res["backprojected_sum"].data == 0)
+    # The detector correction works in focus
+    assert np.any(res["corrected_sum"].data != 0)
+
+
+def test_cached_matrices_read_only():
+    model = Model4DSTEM.default(dataset_shape=(4, 5, 6, 7)).derive(overfocus=0.1)
+    for mat in (_backward_matrix(model), _correction_matrix(model)):
+        with pytest.raises(ValueError, match="read-only"):
+            mat[0, 0] = 1.0
diff --git a/tests/test_transformation_matrices.py b/tests/test_transformation_matrices.py
index 49e7132..ca1f17b 100644
--- a/tests/test_transformation_matrices.py
+++ b/tests/test_transformation_matrices.py
@@ -14,7 +14,7 @@
     get_detector_correction_matrix,
     ray_transfer_matrix,
 )
-from microscope_calibration.udf.stem_overfocus import BaseCorrectionUDF
+from microscope_calibration.udf.stem_overfocus import _backward_matrix
 from microscope_calibration.util.stem_overfocus_sim import get_forward_transformation_matrix
 from microscope_calibration.util.sympy import lambdify
 
@@ -272,8 +272,8 @@ def test_udf_no_backprojection_in_focus(overfocus):
     # The UDF continues without backprojection for singular transformations.
     # Before deriving the matrix symbolically, near-singular cases raised a
     # RuntimeError instead.
-    BaseCorrectionUDF._back_mat.cache_clear()
-    assert BaseCorrectionUDF._back_mat(rec_model=default_model(overfocus=overfocus)) is None
+    _backward_matrix.cache_clear()
+    assert _backward_matrix(default_model(overfocus=overfocus)) is None
 
 
 @pytest.mark.parametrize("delta", (0.0, 1e-12))

From f1828a8755ecd3031320c9f21892345c207a10c4 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Mon, 12 Oct 2026 16:04:01 +0000
Subject: [PATCH 18/34] Clean up the overfocus UDFs

- Raise ValueError for scans that aren't 2D
- Sum the results of CorrectedPickUDF without stacking all chunks
- Remove empty docstrings and the ineffective where="device"

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../udf/stem_overfocus.py                     | 31 ++++++++-----------
 tests/test_overfocus_udf.py                   |  9 ++++++
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/src/microscope_calibration/udf/stem_overfocus.py b/src/microscope_calibration/udf/stem_overfocus.py
index 115274a..5b4cc50 100644
--- a/src/microscope_calibration/udf/stem_overfocus.py
+++ b/src/microscope_calibration/udf/stem_overfocus.py
@@ -69,10 +69,11 @@ def get_task_data(self):
             "corr_mat": _correction_matrix(model),
         }
 
-    def _get_fov(self):
-        fov_size_y = int(self.meta.dataset_shape.nav[0])
-        fov_size_x = int(self.meta.dataset_shape.nav[1])
-        return fov_size_y, fov_size_x
+    def _get_fov(self) -> tuple[int, int]:
+        nav_shape = tuple(self.meta.dataset_shape.nav)
+        if len(nav_shape) != 2:
+            raise ValueError(f"Only 2D scans are supported, got navigation shape {nav_shape}.")
+        return nav_shape
 
     @property
     def has_correction(self):
@@ -88,11 +89,9 @@ def get_result_buffers(self):
         fov = self._get_fov()
         dtype = np.result_type(self.meta.input_dtype, np.float32)
         return {
-            "backprojected_sum": self.buffer(
-                kind="single", dtype=dtype, extra_shape=fov, where="device"
-            ),
-            "corrected_point": self.buffer(kind="nav", dtype=dtype, where="device"),
-            "corrected_sum": self.buffer(kind="sig", dtype=dtype, where="device"),
+            "backprojected_sum": self.buffer(kind="single", dtype=dtype, extra_shape=fov),
+            "corrected_point": self.buffer(kind="nav", dtype=dtype),
+            "corrected_sum": self.buffer(kind="sig", dtype=dtype),
         }
 
     def process_frame(self, frame):
@@ -147,11 +146,9 @@ def merge(self, dest, src):
 # not run into inheritance complications
 class CorrectedPickUDF(BaseCorrectionUDF):
     def get_preferred_input_dtype(self):
-        ""
         return self.USE_NATIVE_DTYPE
 
     def get_result_buffers(self):
-        ""
         dtype = self.meta.input_dtype
         sigshape = tuple(self.meta.dataset_shape.sig)
         backprojection_fov = self._get_fov()
@@ -165,11 +162,13 @@ def get_result_buffers(self):
             log.warning(
                 "CorrectedPickUDF is loading %s bytes per buffer, exceeding warning limit %s. "
                 "Consider using or implementing an UDF to process data on the worker "
-                "nodes instead." % (loaded_size, warn_limit)
+                "nodes instead.",
+                loaded_size,
+                warn_limit,
             )
         # We are using a "single" buffer since we mostly load single frames. A
         # "sig" buffer would work as well, but would require a transpose to
-        # accomodate multiple frames in the last and not first dimension.
+        # accommodate multiple frames in the last and not first dimension.
         # A "nav" buffer would allocate a NaN-filled buffer for the whole dataset.
         return {
             "corrected": self.buffer(kind="single", extra_shape=(navsize,) + sigshape, dtype=dtype),
@@ -179,7 +178,6 @@ def get_result_buffers(self):
         }
 
     def process_frame(self, frame):
-        ""
         scan_y, scan_x = self.meta.coordinates[0]
         overfocus_model: Model4DSTEM = self.params.overfocus_model["model"]
         # We work in flattened nav space with ROI applied
@@ -203,7 +201,6 @@ def process_frame(self, frame):
             )
 
     def merge(self, dest, src):
-        ""
         # We receive full-size buffers from each node that
         # contributes at least one frame and rely on the rest being filled
         # with zeros correctly.
@@ -211,13 +208,11 @@ def merge(self, dest, src):
         dest.backprojected[:] += src.backprojected
 
     def merge_all(self, ordered_results):
-        ""
         res = {}
         for attr in ("corrected", "backprojected"):
             chunks = [getattr(b, attr) for b in ordered_results.values()]
             # We receive full-size buffers from each node that
             # contributes at least one frame and rely on the rest being filled
             # with zeros correctly.
-            ssum = np.stack(chunks, axis=0).sum(axis=0)
-            res[attr] = ssum
+            res[attr] = sum(chunks[1:], start=chunks[0].copy())
         return res
diff --git a/tests/test_overfocus_udf.py b/tests/test_overfocus_udf.py
index 47d7f41..9a7f634 100644
--- a/tests/test_overfocus_udf.py
+++ b/tests/test_overfocus_udf.py
@@ -199,3 +199,12 @@ def test_cached_matrices_read_only():
     for mat in (_backward_matrix(model), _correction_matrix(model)):
         with pytest.raises(ValueError, match="read-only"):
             mat[0, 0] = 1.0
+
+
+@pytest.mark.parametrize("udf_class", [OverfocusUDF, CorrectedPickUDF])
+def test_udf_requires_2d_scan(udf_class):
+    ctx = Context.make_with("inline")
+    ds = ctx.load("memory", data=np.zeros((2, 3, 4, 5, 6)), sig_dims=2)
+    model = Model4DSTEM.default(dataset_shape=(3, 4, 5, 6)).derive(overfocus=0.1)
+    with pytest.raises(ValueError, match="Only 2D scans"):
+        ctx.run_udf(dataset=ds, udf=udf_class(overfocus_model={"model": model}))

From c2e741cee792304d0b549d4d24d906b6ca3395b2 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 07:54:52 +0000
Subject: [PATCH 19/34] Replace the backend dispatch of cos/sin with lambdify()
 at the call sites

Model functions stay SymPy-compatible. Numeric NumPy or JAX versions are
created with lambdify() once where they are needed, see CLAUDE.md.

- Remove util/backend.py and regenerate _generated.py
- Test adjust_*() and invert_focus() through lambdified NumPy and JAX functions

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CLAUDE.md                                     |  7 ++-
 .../common/_generated.py                      |  2 +-
 src/microscope_calibration/common/model.py    | 11 ++---
 src/microscope_calibration/util/backend.py    | 48 -------------------
 src/microscope_calibration/util/codegen.py    | 20 ++++----
 tests/test_codegen.py                         | 16 ++-----
 tests/test_model.py                           | 48 +++++++++++++++----
 tests/test_optimize.py                        |  1 +
 8 files changed, 61 insertions(+), 92 deletions(-)
 delete mode 100644 src/microscope_calibration/util/backend.py

diff --git a/CLAUDE.md b/CLAUDE.md
index 81bc896..456fe3d 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -80,8 +80,6 @@ Key layers, bottom-up:
    `wrapt`) that recurses into nested dataclasses/pytrees (`PixelYX`, `Model4DSTEM`,
    `DescanError`) so a symbolic trace can be turned into a numeric function. Touch with care;
    the three-backend contract depends on it.
-   `util/backend.py` provides math functions such as `cos`/`sin` that dispatch on the argument
-   type; use them instead of `sympy`/`numpy`/`jax.numpy` functions in model code.
 
 3. **`common/stem_overfocus.py`** — pixel-to-pixel transformation matrices (backward
    projection, detector correction) calculated from the generated `ray_transfer_matrix()`;
@@ -117,6 +115,11 @@ Key layers, bottom-up:
   you see `try/except TracerBoolConversionError` and `_one`/`flip_factor` style tricks
   (constants threaded through so operations stay differentiable instead of using Python
   conditionals). Preserve that pattern.
+- Model functions are always SymPy-compatible and use SymPy math (`sym.cos`, `sym.pi`, ...);
+  called directly on floats they may return SymPy numbers. For NumPy or JAX, wrap the call
+  site with `@lambdify(modules=np|jnp, recurse_for=...)` (or use `lambdify_trace_for`). Create
+  the lambdified function once, e.g. at module level: each new function object triggers a new
+  symbolic derivation, so don't decorate inside loops or frequently called functions.
 - `trace()` is the single source of truth. Derive relations from it symbolically instead of
   hard-coding them: register the derivation in a `GENERATED_FUNCTIONS` registry (collected in
   `common/generate.py`) and run the generator. Never edit `common/_generated.py` by hand.
diff --git a/src/microscope_calibration/common/_generated.py b/src/microscope_calibration/common/_generated.py
index 26c6397..2bf8a7e 100644
--- a/src/microscope_calibration/common/_generated.py
+++ b/src/microscope_calibration/common/_generated.py
@@ -4,7 +4,7 @@
 # microscope_calibration.common.generate.GENERATED_FUNCTIONS.
 # tests/test_codegen.py verifies that this file is up to date.
 
-from microscope_calibration.util.backend import cos, sin
+from sympy import cos, sin
 from temgym_core.components import DescanError
 
 
diff --git a/src/microscope_calibration/common/model.py b/src/microscope_calibration/common/model.py
index b345576..365bba3 100644
--- a/src/microscope_calibration/common/model.py
+++ b/src/microscope_calibration/common/model.py
@@ -10,7 +10,6 @@
 jax.config.update("jax_enable_x64", True)  # noqa: E702
 # ruff: disable[E402]
 import jax_dataclasses as jdc
-import numpy as np
 import sympy as sym
 from sympy.solvers.solveset import NonlinearError
 from jax.errors import TracerBoolConversionError
@@ -21,7 +20,6 @@
 from temgym_core.run import run_iter
 from temgym_core.source import PointSource, Source
 
-from microscope_calibration.util import backend
 from microscope_calibration.util.sympy import lambdify, lambdify_tree, symbol_maker
 # ruff: enable[E402]
 
@@ -63,8 +61,8 @@ def scale(c: CoordXY | PixelYX, factor: float) -> CoordXY:
 
 
 def rotate(c: CoordXY | PixelYX, radians: float) -> CoordXY:
-    cosr = backend.cos(radians)
-    sinr = backend.sin(radians)
+    cosr = sym.cos(radians)
+    sinr = sym.sin(radians)
     return CoordXY(
         x=c.x * cosr - c.y * sinr,
         y=c.y * cosr + c.x * sinr,
@@ -604,10 +602,9 @@ def invert_focus(self) -> Self:
         Over- and underfocus can't be distinguished from the shadow image
         alone. This allows switching between the two interpretations.
         """
-        pi = sym.pi if isinstance(self.detector_rotation, sym.Basic) else np.pi
         return (
             self.derive(overfocus=-self.overfocus)
-            .adjust_detector_rotation(detector_rotation=self.detector_rotation + pi)
+            .adjust_detector_rotation(detector_rotation=self.detector_rotation + sym.pi)
             .adjust_camera_length(camera_length=self.camera_length + 2 * self.overfocus)
         )
 
@@ -657,7 +654,7 @@ def _runtime_descan_compensation(model_type: type[Model4DSTEM], param: str):
     warn_runtime_derivation(f"descan_compensation_{param}", model_type)
     args, output = _derive_descan_compensation(param, model_type=model_type)
     check_pytree(args["model"])
-    return lambdify_tree(args, output, modules=[backend.NAMESPACE, "sympy"])
+    return lambdify_tree(args, output, modules=sym)
 
 
 def _derive_descan_compensation(param: str, model_type: type[Model4DSTEM] = Model4DSTEM):
diff --git a/src/microscope_calibration/util/backend.py b/src/microscope_calibration/util/backend.py
deleted file mode 100644
index 67ac842..0000000
--- a/src/microscope_calibration/util/backend.py
+++ /dev/null
@@ -1,48 +0,0 @@
-"""
-Math functions that work with SymPy, JAX and plain numbers
-
-The functions dispatch on the type of their argument so that code using them,
-e.g. :meth:`~microscope_calibration.common.model.Model4DSTEM.trace` and the
-generated code in :mod:`microscope_calibration.common._generated`, keeps the
-type of its inputs: SymPy expressions stay symbolic, JAX arrays and tracers
-stay differentiable and traceable, and other numbers are calculated with NumPy.
-"""
-
-from collections.abc import Callable
-
-import jax
-
-jax.config.update("jax_enable_x64", True)  # noqa: E702
-# ruff: disable[E402]
-import jax.numpy as jnp
-import numpy as np
-import sympy as sym
-# ruff: enable[E402]
-
-__all__ = ["NAMESPACE", "cos", "sin"]
-
-
-def _dispatch(name: str) -> Callable:
-    sym_f = getattr(sym, name)
-    jnp_f = getattr(jnp, name)
-    np_f = getattr(np, name)
-
-    def f(x):
-        if isinstance(x, sym.Basic):
-            return sym_f(x)
-        # Includes tracers
-        elif isinstance(x, jax.Array):
-            return jnp_f(x)
-        else:
-            return np_f(x)
-
-    f.__name__ = f.__qualname__ = name
-    f.__doc__ = f"{name}() for SymPy expressions, JAX arrays and other numbers"
-    return f
-
-
-cos = _dispatch("cos")
-sin = _dispatch("sin")
-
-# Namespace for :func:`sympy.lambdify`, e.g. :code:`modules=[NAMESPACE, "sympy"]`
-NAMESPACE = {name: globals()[name] for name in ("cos", "sin")}
diff --git a/src/microscope_calibration/util/codegen.py b/src/microscope_calibration/util/codegen.py
index ed63187..1b2ea0d 100644
--- a/src/microscope_calibration/util/codegen.py
+++ b/src/microscope_calibration/util/codegen.py
@@ -15,8 +15,6 @@
 import sympy as sym
 from sympy.printing.pycode import SymPyPrinter
 
-from microscope_calibration.util import backend
-
 
 def symbol_paths(tree: Any, path: tuple[str, ...]) -> Iterator[tuple[tuple[str, ...], sym.Symbol]]:
     """
@@ -185,8 +183,8 @@ def render_module(functions: dict[str, tuple[dict[str, Any], Any]], header: str)
     Render a module with the given functions
 
     Names that the printed SymPy expressions use, e.g. :code:`cos`, are imported
-    from :mod:`microscope_calibration.util.backend`, so that the result works
-    with SymPy expressions, JAX arrays and tracers, and plain numbers.
+    from :code:`sympy`, so that the result works with SymPy and plain Python
+    number types, like :func:`sympy.lambdify` with :code:`modules=sympy`.
 
     Parameters
     ----------
@@ -207,14 +205,12 @@ def render_module(functions: dict[str, tuple[dict[str, Any], Any]], header: str)
     for t in types:
         imports.setdefault(t.__module__, set()).add(t.__name__)
     type_names = {t.__name__ for t in types}
-    math_names = _free_names(body) - type_names
-    for n in math_names:
-        if n not in backend.NAMESPACE:
-            raise ValueError(
-                f"Generated code uses name {n} that is not available in {backend.__name__}."
-            )
-    if math_names:
-        imports.setdefault(backend.__name__, set()).update(math_names)
+    sympy_names = _free_names(body) - type_names
+    for n in sympy_names:
+        if not hasattr(sym, n):
+            raise ValueError(f"Generated code uses name {n} that is not available in sympy.")
+    if sympy_names:
+        imports.setdefault("sympy", set()).update(sympy_names)
 
     import_lines = "".join(
         f"from {module} import {', '.join(sorted(names))}\n"
diff --git a/tests/test_codegen.py b/tests/test_codegen.py
index 048ec64..25d2265 100644
--- a/tests/test_codegen.py
+++ b/tests/test_codegen.py
@@ -1,6 +1,5 @@
 import inspect
 
-import jax
 import numpy as np
 import pytest
 import sympy as sym
@@ -54,15 +53,12 @@ def test_render_module():
 
     res = f(p=PixelYX(y=1.0, x=2.0), angle=0.3)
     assert isinstance(res, PixelYX)
-    assert_allclose(res, (np.cos(0.3) + 2.0 * np.sin(0.3) + np.pi, 1.0))
-    # Functions like cos() keep the backend of their argument
-    assert not any(isinstance(v, sym.Basic) for v in res)
+    # Functions like cos() come from SymPy and return SymPy numbers,
+    # like lambdify(modules=sympy)
+    assert_allclose([float(v) for v in res], (np.cos(0.3) + 2.0 * np.sin(0.3) + np.pi, 1.0))
     # Plain arithmetic keeps plain types
     assert type(res.x) is float
 
-    grad = jax.grad(lambda angle: f(p=PixelYX(y=1.0, x=2.0), angle=angle).y)(0.3)
-    assert_allclose(grad, -np.sin(0.3) + 2.0 * np.cos(0.3))
-
     res = f(p=p, angle=angle)
     for r, o in zip(res, output, strict=True):
         # Constants like pi are printed as floats, like lambdify(modules=sympy)
@@ -73,9 +69,3 @@ def test_render_module_unknown_symbol():
     a, b = sym.symbols("a b")
     with pytest.raises(ValueError, match="not in args"):
         render_module({"f": ({"a": a}, a * b)}, header="")
-
-
-def test_render_module_unknown_function():
-    a = sym.Symbol("a")
-    with pytest.raises(ValueError, match="not available"):
-        render_module({"f": ({"a": a}, sym.tan(a))}, header="")
diff --git a/tests/test_model.py b/tests/test_model.py
index ee47df4..cd5f64c 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1093,16 +1093,47 @@ def test_adjust_camera_length(random_model):
 }
 
 
+def _lambdify_adjust(param: str, modules):
+    def adjust(model: Model4DSTEM, value):
+        return getattr(model, f"adjust_{param}")(value)
+
+    # lambdify() generates symbols for the arguments according to their annotations
+    adjust.__annotations__["value"] = Model4DSTEM.__annotations__[param]
+    return lambdify(recurse_for=(PixelYX, Model4DSTEM, DescanError), modules=modules)(adjust)
+
+
+# The model functions are SymPy-compatible. Numeric versions for a backend are
+# created with lambdify() once, not for each call.
+ADJUST_NUMPY = {param: _lambdify_adjust(param, np) for param in ADJUST_VALUES}
+ADJUST_JAX = {param: _lambdify_adjust(param, jnp) for param in ADJUST_VALUES}
+
+
+@lambdify(recurse_for=(PixelYX, Model4DSTEM, DescanError), modules=np)
+def invert_focus_numpy(model: Model4DSTEM):
+    return model.invert_focus()
+
+
 @pytest.mark.parametrize("param", sorted(ADJUST_VALUES))
-def test_adjust_keeps_float_types(random_model, param):
-    adjusted = getattr(random_model, f"adjust_{param}")(ADJUST_VALUES[param])
+def test_adjust_numpy(random_model, param):
+    adjusted = ADJUST_NUMPY[param](random_model, ADJUST_VALUES[param])
+    direct = getattr(random_model, f"adjust_{param}")(ADJUST_VALUES[param])
+    assert type(adjusted) is Model4DSTEM
     for leaf in jax.tree.leaves(adjusted):
         assert not isinstance(leaf, sym.Basic), leaf
+    assert_allclose(
+        jax.tree.leaves(adjusted), jax.tree.leaves(direct.normalize_types()), rtol=1e-12
+    )
 
 
-def test_invert_focus_keeps_float_types(random_model):
-    for leaf in jax.tree.leaves(random_model.invert_focus()):
+def test_invert_focus_numpy(random_model):
+    inverted = invert_focus_numpy(random_model)
+    for leaf in jax.tree.leaves(inverted):
         assert not isinstance(leaf, sym.Basic), leaf
+    assert_allclose(
+        jax.tree.leaves(inverted),
+        jax.tree.leaves(random_model.invert_focus().normalize_types()),
+        rtol=1e-12,
+    )
 
 
 @pytest.mark.parametrize(
@@ -1110,7 +1141,7 @@ def test_invert_focus_keeps_float_types(random_model):
 )
 def test_adjust_grad(random_model, param):
     def f(value):
-        return sum(getattr(random_model, f"adjust_{param}")(value).descan_error)
+        return sum(ADJUST_JAX[param](random_model, value).descan_error)
 
     value = ADJUST_VALUES[param]
     eps = 1e-6
@@ -1120,10 +1151,9 @@ def f(value):
 
 @pytest.mark.parametrize("flip_factor", [-1.0, 1.0])
 def test_adjust_flip_factor_jit(random_model, flip_factor):
-    def f(flip_factor):
-        return random_model.adjust_flip_factor(flip_factor).descan_error
-
-    assert_allclose(jax.jit(f)(flip_factor), f(flip_factor), rtol=1e-12, atol=1e-12)
+    res = jax.jit(ADJUST_JAX["flip_factor"])(random_model, flip_factor)
+    ref = ADJUST_NUMPY["flip_factor"](random_model, flip_factor)
+    assert_allclose(res.descan_error, ref.descan_error, rtol=1e-12, atol=1e-12)
 
 
 def test_derive_flip_y(random_model):
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index 16ba99a..802e161 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -713,6 +713,7 @@ def test_coords_points(flip_factor, ref_flip_factor):
         ref_model.derive(overfocus=0.012)
         .adjust_detector_rotation(0.3)
         .adjust_flip_factor(flip_factor)
+        .normalize_types()
     )
     points = []
     for scan_y, scan_x, dy, dx in ((3, 4, 1e-3, -2e-3), (20, 7, -1e-3, 0), (12, 28, 2e-3, 1e-3)):

From 3089b0c5eebd21fa330c0020cdd40bc81b10b16b Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 14:48:07 +0000
Subject: [PATCH 20/34] Update the Apptainer image and its pinned requirements
 to Python 3.14

The pins were for Python 3.11 and lacked the extras of the package, e.g. JAX
and optimistix.

- Re-enable the CI job that tests the pinned requirements

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .github/workflows/ci.yml         |   5 +-
 packaging/overfocus.def          |   2 +-
 packaging/requirements.txt       | 643 ++++++++++++++++++++++---------
 packaging/update-requirements.sh |   5 +-
 4 files changed, 474 insertions(+), 181 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 99b9b9b..202f90f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -37,14 +37,13 @@ jobs:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
 
   pinned-requirements:
-    if: false  # skip job
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v7
-      - name: Choose Python version ${{ matrix.python_version }}
+      - name: 'Choose Python version 3.14'
         uses: actions/setup-python@v7
         with:
-          python-version: '3.11'
+          python-version: '3.14'
       - name: install the package
         run: python3 -m pip install uv && uv pip install --system .[test,common,diffraction,gui] -r packaging/requirements.txt
       - name: run pytest
diff --git a/packaging/overfocus.def b/packaging/overfocus.def
index 67e0f07..9aff109 100644
--- a/packaging/overfocus.def
+++ b/packaging/overfocus.def
@@ -1,5 +1,5 @@
 bootstrap: docker
-from: python:3.11-slim
+from: python:3.14-slim
 
 %files
         ./Microscope-Calibration        /Microscope-Calibration
diff --git a/packaging/requirements.txt b/packaging/requirements.txt
index 1d71105..c721eac 100644
--- a/packaging/requirements.txt
+++ b/packaging/requirements.txt
@@ -1,124 +1,176 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile --no-cache --python-version=3.11 requirements.in ../pyproject.toml
-anyio==4.8.0
+#    uv pip compile --no-cache --python-version=3.14 --extra common --extra diffraction --extra gui requirements.in ../pyproject.toml
+absl-py==2.5.1
+    # via
+    #   optax
+    #   orbax-checkpoint
+aiofiles==25.1.0
+    # via orbax-checkpoint
+anyio==4.15.1
     # via
     #   httpx
     #   jupyter-server
-argon2-cffi==23.1.0
+argon2-cffi==25.1.0
     # via jupyter-server
-argon2-cffi-bindings==21.2.0
+argon2-cffi-bindings==26.1.0
     # via argon2-cffi
-arrow==1.3.0
+arrow==1.4.0
     # via isoduration
-asttokens==3.0.0
+asttokens==3.0.2
     # via stack-data
-async-lru==2.0.4
+async-lru==2.4.0
     # via jupyterlab
-attrs==24.3.0
+attrs==26.1.0
     # via
     #   jsonschema
     #   referencing
-autopep8==2.3.1
+autopep8==2.3.2
     # via libertem
-babel==2.16.0
+babel==2.18.0
     # via jupyterlab-server
-beautifulsoup4==4.12.3
+beautifulsoup4==4.15.0
     # via nbconvert
-bleach==6.2.0
+bidict==0.24.1
+    # via libertem-ui
+bleach==6.4.0
     # via nbconvert
-bqplot==0.12.44
+bokeh==3.9.2
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   libertem-ui
+    #   panel
+    #   panel-material-ui
+bqplot==0.13.1
     # via
     #   bqplot-image-gl
     #   libertem
-bqplot-image-gl==1.6.0
+bqplot-image-gl==1.10.2
     # via libertem
-certifi==2024.12.14
+bqscales==0.3.7
+    # via bqplot
+certifi==2026.7.22
     # via
     #   httpcore
     #   httpx
     #   requests
-cffi==1.17.1
+cffi==2.1.1
     # via argon2-cffi-bindings
-charset-normalizer==3.4.1
+charset-normalizer==3.5.2
     # via requests
-click==8.1.8
+click==8.5.0
     # via
     #   dask
     #   distributed
     #   libertem
-cloudpickle==3.1.0
+cloudpickle==3.1.2
     # via
     #   dask
     #   distributed
+    #   joblib
     #   libertem
-comm==0.2.2
+colorama==0.4.6
+    # via pyqtgraph
+colorcet==3.2.1
+    # via
+    #   datashader
+    #   libertem
+    #   libertem-ui
+comm==0.2.3
     # via
     #   ipykernel
     #   ipywidgets
-contourpy==1.3.1
-    # via matplotlib
+contourpy==1.4.0
+    # via
+    #   bokeh
+    #   matplotlib
 cycler==0.12.1
     # via matplotlib
-dask==2024.12.1
+dask==2026.8.0
     # via
     #   distributed
     #   libertem
-debugpy==1.8.11
+    #   orix
+datashader==0.19.1
+    # via libertem-ui
+debugpy==1.8.22
     # via ipykernel
-decorator==5.1.1
-    # via ipython
 defusedxml==0.7.1
     # via
     #   libertem
     #   nbconvert
-deprecated==1.2.15
-    # via opentelemetry-api
-distributed==2024.12.1
-    # via libertem
-empyre==0.3.3
+diffpy-structure==3.5.0
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   diffsims
+    #   orix
+diffpy-utils==3.7.2
+    # via diffpy-structure
+diffsims @ git+https://github.com/pyxem/diffsims@2fc7216d2a9871360bdba042e2db15c530b5f8d6
+    # via microscope-calibration (../pyproject.toml)
+distributed==2026.8.0
     # via libertem
-executing==2.1.0
+equinox==0.13.8
+    # via
+    #   lineax
+    #   optimistix
+etils==1.14.0
+    # via orbax-checkpoint
+executing==2.3.0
     # via stack-data
-fastjsonschema==2.21.1
+fastjsonschema==2.22.2
     # via nbformat
-fonttools==4.55.3
+flax==0.12.10
+    # via microscope-calibration (../pyproject.toml)
+flexcache==0.3
+    # via pint
+flexparser==0.4
+    # via pint
+fonttools==4.66.1
     # via matplotlib
-fqdn==1.5.1
+fqdn==1.6.0
     # via jsonschema
-fsspec==2024.12.0
-    # via dask
-h11==0.14.0
+frozendict==2.4.7
+    # via microscope-calibration (../pyproject.toml)
+fsspec==2026.9.0
+    # via
+    #   dask
+    #   etils
+h11==0.16.0
     # via httpcore
-h5py==3.12.1
+h5py==3.16.0
     # via
+    #   hdf5plugin
     #   libertem
     #   ncempy
-httpcore==1.0.7
+    #   orix
+hdf5plugin==7.1.0
+    # via ncempy
+httpcore==1.0.9
     # via httpx
 httpx==0.28.1
     # via jupyterlab
-idna==3.10
+humanize==4.16.0
+    # via
+    #   libertem-ui
+    #   orbax-checkpoint
+idna==3.20
     # via
     #   anyio
     #   httpx
     #   jsonschema
     #   requests
-imageio==2.36.1
+imageio==2.38.1
     # via scikit-image
-importlib-metadata==8.5.0
-    # via
-    #   dask
-    #   opentelemetry-api
-ipykernel==6.29.5
+ipykernel==7.4.0
     # via
     #   jupyter
     #   jupyter-console
     #   jupyterlab
-ipympl==0.9.6
+ipympl==0.10.0
     # via
     #   -r requirements.in
     #   libertem
-ipython==8.31.0
+ipython==9.17.1
     # via
     #   ipykernel
     #   ipympl
@@ -126,41 +178,75 @@ ipython==8.31.0
     #   jupyter-console
     #   jupyter-ui-poll
     #   libertem
-ipywidgets==8.1.5
+ipython-pygments-lexers==1.1.1
+    # via ipython
+ipywidgets==8.1.9
     # via
     #   -r requirements.in
+    #   microscope-calibration (../pyproject.toml)
     #   bqplot
     #   bqplot-image-gl
+    #   bqscales
     #   ipympl
     #   jupyter
 isoduration==20.11.0
     # via jsonschema
-jedi==0.19.2
+jax==0.11.2
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   equinox
+    #   flax
+    #   jax-dataclasses
+    #   lineax
+    #   optax
+    #   optimistix
+    #   orbax-checkpoint
+    #   temgym-core
+jax-dataclasses==1.6.3
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   temgym-core
+jaxlib==0.11.2
+    # via
+    #   jax
+    #   jax-dataclasses
+    #   optax
+jaxtyping==0.3.11
+    # via
+    #   equinox
+    #   lineax
+    #   optimistix
+jedi==0.20.1
     # via ipython
-jinja2==3.1.5
+jinja2==3.1.6
     # via
+    #   bokeh
     #   distributed
     #   jupyter-server
     #   jupyterlab
     #   jupyterlab-server
     #   nbconvert
-joblib==1.4.2
+joblib==1.6.0
     # via scikit-learn
-json5==0.10.0
+json5==0.17.3
     # via jupyterlab-server
-jsonpointer==3.0.0
+jsonpointer==3.2.1
     # via jsonschema
-jsonschema==4.23.0
+jsonschema==4.26.0
     # via
     #   jupyter-events
     #   jupyterlab-server
     #   libertem
     #   nbformat
-jsonschema-specifications==2024.10.1
+jsonschema-specifications==2025.9.1
     # via jsonschema
 jupyter==1.1.1
     # via -r requirements.in
-jupyter-client==8.6.3
+jupyter-builder==1.2.3
+    # via
+    #   jupyterlab
+    #   notebook
+jupyter-client==8.10.0
     # via
     #   ipykernel
     #   jupyter-console
@@ -168,9 +254,10 @@ jupyter-client==8.6.3
     #   nbclient
 jupyter-console==6.6.3
     # via jupyter
-jupyter-core==5.7.2
+jupyter-core==5.9.1
     # via
     #   ipykernel
+    #   jupyter-builder
     #   jupyter-client
     #   jupyter-console
     #   jupyter-server
@@ -178,88 +265,133 @@ jupyter-core==5.7.2
     #   nbclient
     #   nbconvert
     #   nbformat
-jupyter-events==0.11.0
+jupyter-events==0.12.1
     # via jupyter-server
-jupyter-lsp==2.2.5
+jupyter-lsp==2.3.1
     # via jupyterlab
-jupyter-server==2.15.0
+jupyter-server==2.21.1
     # via
     #   jupyter-lsp
     #   jupyterlab
     #   jupyterlab-server
     #   notebook
     #   notebook-shim
-jupyter-server-terminals==0.5.3
+jupyter-server-terminals==0.5.4
     # via jupyter-server
-jupyter-ui-poll==1.0.0
+jupyter-ui-poll==1.1.0
     # via
     #   -r requirements.in
     #   libertem
-jupyterlab==4.3.4
+jupyterlab==4.6.4
     # via
     #   jupyter
     #   notebook
 jupyterlab-pygments==0.3.0
     # via nbconvert
-jupyterlab-server==2.27.3
+jupyterlab-server==2.28.1
     # via
     #   jupyterlab
     #   notebook
-jupyterlab-widgets==3.0.13
+jupyterlab-widgets==3.0.17
     # via ipywidgets
-kiwisolver==1.4.8
+kiwisolver==1.5.1
     # via matplotlib
-lazy-loader==0.4
-    # via scikit-image
-libertem==0.14.1
+lark==1.3.1
+    # via rfc3987-syntax
+lazy-loader==0.6
+    # via
+    #   orix
+    #   scikit-image
+libertem==0.16.0
     # via
     #   -r requirements.in
     #   microscope-calibration (../pyproject.toml)
-llvmlite==0.43.0
+libertem-blobfinder==0.6.1
+    # via microscope-calibration (../pyproject.toml)
+libertem-ui @ git+https://github.com/LiberTEM/LiberTEM-panel-ui@c46591b90f08d3a7018a06f6c5cd8c57ccaaeb59
+    # via microscope-calibration (../pyproject.toml)
+lineax==0.1.1
+    # via optimistix
+linkify-it-py==2.2.0
+    # via panel
+llvmlite==0.50.0
     # via numba
 locket==1.0.0
     # via
     #   distributed
     #   partd
-markupsafe==3.0.2
+markdown==3.11
+    # via panel
+markdown-it-py==4.2.0
+    # via
+    #   mdit-py-plugins
+    #   panel
+    #   rich
+markupsafe==3.0.4
     # via
     #   jinja2
     #   nbconvert
-matplotlib==3.10.0
+matplotlib==3.11.2
     # via
     #   -r requirements.in
-    #   empyre
+    #   diffsims
     #   ipympl
     #   libertem
     #   ncempy
-matplotlib-inline==0.1.7
+    #   orix
+matplotlib-inline==0.2.2
     # via
     #   ipykernel
     #   ipython
-mistune==3.1.0
+mdit-py-plugins==0.6.1
+    # via panel
+mdurl==0.1.2
+    # via markdown-it-py
+mistune==3.3.4
     # via nbconvert
-msgpack==1.1.0
-    # via distributed
-nbclient==0.10.2
+ml-dtypes==0.6.0
+    # via
+    #   jax
+    #   jaxlib
+    #   tensorstore
+mpmath==1.3.0
+    # via sympy
+msgpack==1.2.3
+    # via
+    #   distributed
+    #   flax
+    #   orbax-checkpoint
+multipledispatch==1.0.0
+    # via datashader
+narwhals==2.27.1
+    # via
+    #   bokeh
+    #   panel
+    #   scikit-learn
+nbclient==0.11.0
     # via nbconvert
-nbconvert==7.16.5
+nbconvert==7.17.2
     # via
     #   jupyter
     #   jupyter-server
     #   libertem
-nbformat==5.10.4
+nbformat==5.11.1
     # via
     #   jupyter-server
     #   libertem
     #   nbclient
     #   nbconvert
-ncempy==1.12
-    # via libertem
-nest-asyncio==1.6.0
+ncempy==1.16
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   libertem
+nest-asyncio2==1.7.4
     # via ipykernel
-networkx==3.4.2
+networkx==3.7
     # via scikit-image
-notebook==7.3.2
+nh3==0.3.7
+    # via panel
+notebook==7.6.3
     # via
     #   -r requirements.in
     #   jupyter
@@ -267,106 +399,197 @@ notebook-shim==0.2.4
     # via
     #   jupyterlab
     #   notebook
-numba==0.60.0
+numba==0.68.0
     # via
     #   microscope-calibration (../pyproject.toml)
+    #   datashader
+    #   diffsims
     #   libertem
+    #   libertem-blobfinder
+    #   libertem-ui
+    #   orix
     #   sparse
-    #   temgymbasic
-numexpr==2.10.2
+    #   temgym-core
+numexpr==2.14.2
     # via libertem
-numpy==2.0.2
+numpy==2.5.4
     # via
     #   -r requirements.in
     #   microscope-calibration (../pyproject.toml)
+    #   bokeh
     #   bqplot
+    #   bqscales
     #   contourpy
-    #   empyre
+    #   dask
+    #   datashader
+    #   diffpy-structure
+    #   diffpy-utils
+    #   diffsims
+    #   flax
     #   h5py
     #   imageio
     #   ipympl
+    #   jax
+    #   jaxlib
     #   libertem
+    #   libertem-blobfinder
+    #   libertem-ui
     #   matplotlib
+    #   ml-dtypes
     #   ncempy
     #   numba
     #   numexpr
+    #   optax
+    #   orbax-checkpoint
+    #   orix
     #   pandas
+    #   pycifrw
     #   pyqtgraph
     #   scikit-image
     #   scikit-learn
     #   scipy
     #   sparse
     #   sparseconverter
-    #   temgymbasic
+    #   temgym-core
+    #   tensorstore
     #   tifffile
-opentelemetry-api==1.29.0
+    #   transforms3d
+    #   treescope
+    #   xarray
+    #   xraydb
+opentelemetry-api==1.45.1
     # via libertem
-overrides==7.7.0
-    # via jupyter-server
-packaging==24.2
+opt-einsum==3.4.0
+    # via jax
+optax==0.2.8
     # via
+    #   microscope-calibration (../pyproject.toml)
+    #   flax
+optimistix==0.1.0
+    # via microscope-calibration (../pyproject.toml)
+orbax-checkpoint==0.12.7
+    # via flax
+orix==0.15.0
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   diffsims
+packaging==26.3
+    # via
+    #   bokeh
     #   dask
+    #   datashader
     #   distributed
     #   ipykernel
+    #   jupyter-events
     #   jupyter-server
     #   jupyterlab
     #   jupyterlab-server
     #   lazy-loader
     #   matplotlib
     #   nbconvert
+    #   panel
+    #   panel-material-ui
+    #   pooch
     #   scikit-image
-pandas==2.2.3
-    # via bqplot
+    #   xarray
+pandas==3.0.6
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   bqplot
+    #   datashader
+    #   libertem-ui
+    #   panel
+    #   xarray
 pandocfilters==1.5.1
     # via nbconvert
-parso==0.8.4
+panel==1.9.4
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   libertem-ui
+    #   panel-material-ui
+panel-material-ui==0.16.2
+    # via panel
+param==2.4.2
+    # via
+    #   datashader
+    #   panel
+    #   pyct
+    #   pyviz-comms
+parso==0.8.7
     # via jedi
 partd==1.4.2
     # via dask
 pexpect==4.9.0
     # via ipython
-pillow==11.1.0
+pillow==12.3.0
     # via
+    #   bokeh
     #   bqplot-image-gl
-    #   empyre
     #   imageio
     #   ipympl
     #   libertem
     #   matplotlib
     #   scikit-image
-platformdirs==4.3.6
-    # via jupyter-core
-prometheus-client==0.21.1
-    # via jupyter-server
-prompt-toolkit==3.0.48
+pint==0.26.1
+    # via microscope-calibration (../pyproject.toml)
+platformdirs==4.13.0
+    # via
+    #   jupyter-core
+    #   pint
+    #   pooch
+    #   xraydb
+ply==3.11
+    # via pycifrw
+pooch==1.9.0
+    # via orix
+prettytable==3.18.0
+    # via pycifrw
+prometheus-client==0.26.0
+    # via
+    #   jupyter-server
+    #   orbax-checkpoint
+prompt-toolkit==3.0.53
     # via
     #   ipython
     #   jupyter-console
-psutil==6.1.1
+protobuf==7.36.2
+    # via orbax-checkpoint
+psutil==7.2.2
     # via
+    #   diffsims
     #   distributed
-    #   ipykernel
+    #   ipython
     #   libertem
+    #   orbax-checkpoint
 ptyprocess==0.7.0
     # via
     #   pexpect
     #   terminado
-pure-eval==0.2.3
+pure-eval==0.2.4
     # via stack-data
-pycodestyle==2.12.1
+pycifrw==5.0.1
+    # via
+    #   microscope-calibration (../pyproject.toml)
+    #   diffpy-structure
+    #   orix
+pycodestyle==2.15.0
     # via autopep8
-pycparser==2.22
+pycparser==3.11
     # via cffi
-pygments==2.19.0
+pyct==0.6.0
+    # via datashader
+pygments==2.21.0
     # via
     #   ipython
+    #   ipython-pygments-lexers
     #   jupyter-console
     #   nbconvert
-pyopengl==3.1.7
+    #   rich
+pyopengl==3.1.10
     # via -r requirements.in
-pyparsing==3.2.1
+pyparsing==3.3.3
     # via matplotlib
-pyqtgraph==0.13.7
+pyqtgraph==0.14.0
     # via -r requirements.in
 python-dateutil==2.9.0.post0
     # via
@@ -374,28 +597,35 @@ python-dateutil==2.9.0.post0
     #   jupyter-client
     #   matplotlib
     #   pandas
-python-json-logger==3.2.1
+python-json-logger==4.2.0
     # via jupyter-events
-pytz==2024.2
-    # via pandas
-pyyaml==6.0.2
+pyviz-comms==3.0.6
+    # via panel
+pyyaml==6.0.3
     # via
+    #   bokeh
     #   dask
     #   distributed
+    #   flax
     #   jupyter-events
-pyzmq==26.2.0
+    #   orbax-checkpoint
+pyzmq==27.2.0
     # via
     #   ipykernel
     #   jupyter-client
     #   jupyter-console
     #   jupyter-server
-referencing==0.35.1
+referencing==0.37.0
     # via
     #   jsonschema
     #   jsonschema-specifications
     #   jupyter-events
-requests==2.32.3
-    # via jupyterlab-server
+requests==2.34.2
+    # via
+    #   datashader
+    #   jupyterlab-server
+    #   panel
+    #   pooch
 rfc3339-validator==0.1.4
     # via
     #   jsonschema
@@ -404,77 +634,103 @@ rfc3986-validator==0.1.1
     # via
     #   jsonschema
     #   jupyter-events
-rpds-py==0.22.3
+rfc3987-syntax==1.1.0
+    # via jsonschema
+rich==15.0.0
+    # via flax
+rpds-py==2026.9.1
     # via
     #   jsonschema
     #   referencing
-scikit-image==0.25.0
+scikit-image==0.26.0
     # via
     #   -r requirements.in
     #   microscope-calibration (../pyproject.toml)
-    #   empyre
     #   libertem
-scikit-learn==1.6.0
+    #   libertem-blobfinder
+    #   libertem-ui
+scikit-learn==1.9.1
     # via libertem
-scipy==1.15.0
+scipy==1.18.1
     # via
     #   microscope-calibration (../pyproject.toml)
-    #   empyre
+    #   datashader
+    #   diffpy-utils
+    #   diffsims
+    #   jax
+    #   jaxlib
     #   libertem
+    #   libertem-ui
     #   ncempy
+    #   orix
     #   scikit-image
     #   scikit-learn
-    #   sparse
     #   sparseconverter
-send2trash==1.8.3
+    #   xraydb
+send2trash==2.1.0
     # via jupyter-server
-setuptools==75.7.0
-    # via jupyterlab
+simplejson==4.2.0
+    # via orbax-checkpoint
 six==1.17.0
     # via
     #   python-dateutil
     #   rfc3339-validator
-sniffio==1.3.1
-    # via anyio
 sortedcontainers==2.4.0
     # via distributed
-soupsieve==2.6
+soupsieve==3.0.3
     # via beautifulsoup4
-sparse==0.15.4
+sparse==0.19.2
     # via
+    #   microscope-calibration (../pyproject.toml)
     #   libertem
     #   sparseconverter
-sparseconverter==0.4.0
-    # via libertem
+sparseconverter==0.6.0
+    # via
+    #   libertem
+    #   libertem-blobfinder
+sqlalchemy==2.1.4
+    # via xraydb
 stack-data==0.6.3
     # via ipython
-tblib==3.0.0
+strenum==0.4.15
+    # via libertem-ui
+sympy==1.14.0
+    # via microscope-calibration (../pyproject.toml)
+tblib==3.2.2
     # via
     #   distributed
     #   libertem
-temgymbasic @ git+https://github.com/TemGym/TemGym@c7a851ce2d6719acc47d36d31d91df27b1e9590b
+temgym-core==0.1.1
     # via microscope-calibration (../pyproject.toml)
+tensorstore==0.1.85
+    # via
+    #   flax
+    #   orbax-checkpoint
 terminado==0.18.1
     # via
     #   jupyter-server
     #   jupyter-server-terminals
-threadpoolctl==3.5.0
+threadpoolctl==3.7.0
     # via
     #   libertem
     #   scikit-learn
-tifffile==2024.12.12
-    # via scikit-image
-tinycss2==1.4.0
+tifffile==2026.9.20
+    # via
+    #   libertem-ui
+    #   scikit-image
+tinycss2==1.5.1
     # via bleach
-tomli==2.2.1
+tomli==2.5.0
     # via libertem
-toolz==1.0.0
+toolz==1.2.0
     # via
     #   dask
+    #   datashader
     #   distributed
     #   partd
-tornado==6.4.2
+tornado==6.5.10
     # via
+    #   bokeh
     #   distributed
     #   ipykernel
     #   jupyter-client
@@ -483,18 +739,21 @@ tornado==6.4.2
     #   libertem
     #   notebook
     #   terminado
-tqdm==4.67.1
+tqdm==4.70.1
     # via
-    #   empyre
+    #   diffsims
     #   libertem
-traitlets==5.14.3
+    #   orix
+    #   panel
+traitlets==5.16.1
     # via
     #   bqplot
-    #   comm
+    #   bqscales
     #   ipykernel
     #   ipympl
     #   ipython
     #   ipywidgets
+    #   jupyter-builder
     #   jupyter-client
     #   jupyter-console
     #   jupyter-core
@@ -506,41 +765,73 @@ traitlets==5.14.3
     #   nbconvert
     #   nbformat
     #   traittypes
-traittypes==0.2.1
-    # via bqplot
-types-python-dateutil==2.9.0.20241206
-    # via arrow
-typing-extensions==4.12.2
+traittypes==0.2.3
+    # via
+    #   bqplot
+    #   bqscales
+transforms3d==0.4.2
+    # via diffsims
+treescope==0.1.10
+    # via flax
+typing-extensions==4.16.0
     # via
-    #   microscope-calibration (../pyproject.toml)
     #   anyio
-    #   ipython
+    #   beautifulsoup4
+    #   equinox
+    #   etils
+    #   flax
+    #   flexcache
+    #   flexparser
+    #   jax-dataclasses
+    #   jupyter-client
     #   libertem
+    #   libertem-ui
+    #   lineax
+    #   opentelemetry-api
+    #   optimistix
+    #   orbax-checkpoint
+    #   panel
+    #   pint
     #   sparseconverter
-    #   temgymbasic
-tzdata==2024.2
-    # via pandas
+    #   sqlalchemy
+    #   temgym-core
+tzdata==2026.5
+    # via arrow
 uri-template==1.3.0
     # via jsonschema
-urllib3==2.3.0
+urllib3==2.8.0
+    # via requests
+uvloop==0.23.0
+    # via orbax-checkpoint
+wadler-lindig==0.1.7
     # via
-    #   distributed
-    #   requests
-wcwidth==0.2.13
-    # via prompt-toolkit
-webcolors==24.11.1
+    #   equinox
+    #   jaxtyping
+wcwidth==0.9.2
+    # via
+    #   prettytable
+    #   prompt-toolkit
+webcolors==25.10.0
     # via jsonschema
-webencodings==0.5.1
+webencodings==0.6.1
     # via
     #   bleach
     #   tinycss2
-websocket-client==1.8.0
+websocket-client==1.9.2
     # via jupyter-server
-widgetsnbextension==4.0.13
+widgetsnbextension==4.0.16
     # via ipywidgets
-wrapt==1.17.0
-    # via deprecated
+wrapt==2.5.0
+    # via microscope-calibration (../pyproject.toml)
+xarray==2026.9.0
+    # via
+    #   datashader
+    #   libertem-ui
+xraydb==4.5.8
+    # via diffpy-utils
+xyzservices==2026.9.1
+    # via bokeh
 zict==3.0.0
     # via distributed
-zipp==3.21.0
-    # via importlib-metadata
+zipp==4.1.1
+    # via etils
diff --git a/packaging/update-requirements.sh b/packaging/update-requirements.sh
index 1f15b1f..498e4f9 100755
--- a/packaging/update-requirements.sh
+++ b/packaging/update-requirements.sh
@@ -1,2 +1,5 @@
 #!/bin/sh
-uv pip compile --no-cache --python-version=3.11 requirements.in ../pyproject.toml > requirements.txt
+# Pin all dependencies of the image, including the extras of the package
+uv pip compile --no-cache --python-version=3.14 \
+    --extra common --extra diffraction --extra gui \
+    requirements.in ../pyproject.toml > requirements.txt

From 2abde36cd9c441db792d62dc0ef9a4c58688696d Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 14:48:42 +0000
Subject: [PATCH 21/34] Move the coverage configuration to pyproject.toml and
 remove stale files

- Remove setup.cfg with the unused flake8 configuration
- Remove test_requirements.txt, which duplicated part of the test extra
- Coverage no longer warns that include is ignored

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyproject.toml        | 21 ++++++++++++++++++++-
 setup.cfg             | 33 ---------------------------------
 test_requirements.txt |  1 -
 3 files changed, 20 insertions(+), 35 deletions(-)
 delete mode 100644 setup.cfg
 delete mode 100644 test_requirements.txt

diff --git a/pyproject.toml b/pyproject.toml
index 4284474..f425d64 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,6 @@ include = [
   "tests",
   "CITATION.cff",
   "pytest.ini",
-  "setup.cfg",
 ]
 
 # Version is derived from git tags like v0.1.0
@@ -107,3 +106,23 @@ force-exclude = true
 [tool.ruff.lint]
 select = ["E", "F", "B"]
 ignore = ["E24", "E226", "E241", "E242", "B007"]
+
+[tool.coverage.run]
+branch = true
+
+[tool.coverage.report]
+exclude_also = [
+    # Don't complain about missing debug-only code:
+    "def __repr__",
+    "if self\\.debug",
+    # Don't complain about typing branches:
+    "if TYPE_CHECKING",
+    "if typing.TYPE_CHECKING",
+    # Don't complain if tests don't hit defensive assertion code:
+    "raise AssertionError",
+    "raise NotImplementedError",
+    # Don't complain if non-runnable code isn't run:
+    "if 0:",
+    "if False:",
+    "if __name__ == .__main__.:",
+]
diff --git a/setup.cfg b/setup.cfg
deleted file mode 100644
index 5a44d38..0000000
--- a/setup.cfg
+++ /dev/null
@@ -1,33 +0,0 @@
-[flake8]
-max-line-length = 100
-ignore = E24,E121,E123,E126,E128,E133,E226,E241,E242,E704,W503
-exclude = .git,__pycache__,.tox,build,dist,node_modules,TOXENV
-
-[coverage:run]
-branch = True
-include = src/
-
-[coverage:report]
-# Regexes for lines to exclude from consideration
-exclude_lines =
-    # Have to re-enable the standard pragma
-    pragma: no cover
-
-    # Don't complain about missing debug-only code:
-    def __repr__
-    if self\.debug
-
-    # Don't complain about typing branches:
-    if TYPE_CHECKING
-    if typing.TYPE_CHECKING
-
-    # Don't complain if tests don't hit defensive assertion code:
-    raise AssertionError
-    raise NotImplementedError
-
-    # Don't complain if non-runnable code isn't run:
-    if 0:
-    if False:
-    if __name__ == .__main__.:
-
-# note: pytest config moved to pytest.ini beause of https://github.com/pypa/pip/issues/5182 and https://github.com/pytest-dev/pytest/issues/3062
diff --git a/test_requirements.txt b/test_requirements.txt
deleted file mode 100644
index 462c26e..0000000
--- a/test_requirements.txt
+++ /dev/null
@@ -1 +0,0 @@
-optax

From bf582e6d0cfa732fa6cbc3079418e949513fbe4f Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 14:49:06 +0000
Subject: [PATCH 22/34] Move optax to the test extra and drop Numba constraints
 for old Python versions

optax is only used in a test. The project requires Python 3.14.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 packaging/requirements.txt |  4 +---
 pyproject.toml             | 14 ++++----------
 2 files changed, 5 insertions(+), 13 deletions(-)

diff --git a/packaging/requirements.txt b/packaging/requirements.txt
index c721eac..6bbe47e 100644
--- a/packaging/requirements.txt
+++ b/packaging/requirements.txt
@@ -462,9 +462,7 @@ opentelemetry-api==1.45.1
 opt-einsum==3.4.0
     # via jax
 optax==0.2.8
-    # via
-    #   microscope-calibration (../pyproject.toml)
-    #   flax
+    # via flax
 optimistix==0.1.0
     # via microscope-calibration (../pyproject.toml)
 orbax-checkpoint==0.12.7
diff --git a/pyproject.toml b/pyproject.toml
index f425d64..612fd6c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -62,21 +62,15 @@ authors = [
 repository = "https://github.com/LiberTEM/Microscope-Calibration"
 
 [project.optional-dependencies]
-test = ["pytest", "pytest-cov"]
+test = ["pytest", "pytest-cov", "optax"]
 common = [
     "libertem",
-    # Minimum constraints of numba for all Python versions we support
+    # Constrain numba to a recent version, otherwise version resolution might
+    # try to install an ancient version that isn't constrained properly.
     # See https://numba.readthedocs.io/en/stable/release-notes-overview.html
-    "numba>=0.55;python_version < '3.11'",
-    "numba>=0.57;python_version < '3.12'",
-    "numba>=0.59;python_version < '3.13'",
-    # for any future Python release, constrain numba to a recent version,
-    # otherwise, version resolution might try to install an ancient version
-    # that isn't constrained properly:
-    "numba>=0.61;python_version >= '3.13'",
+    "numba>=0.61",
     "scipy",
     "scikit-image",
-    "optax",
     "optimistix",
     "numpy",
 ]

From ea483295ffc671b839808f5da8f65706e2fb7489 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 15:03:25 +0000
Subject: [PATCH 23/34] Read DM4 tags once per file in jeol_dm4 and test it
 with fake tags

Files were opened repeatedly and never closed.

Behavior changes:
- The model functions take the tags of a DM4 file instead of an open file, see read_tags()
- CalibratedJEOLDM4.dm4 is removed, use CalibratedJEOLDM4.tags
- microscope_info() only includes tags of the 4D STEM data, not of the thumbnail

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 examples/stem_overfocus.ipynb               |   2 +-
 src/microscope_calibration/util/jeol_dm4.py | 220 ++++++++++++--------
 tests/test_jeol_dm4.py                      | 165 +++++++++++++++
 3 files changed, 296 insertions(+), 91 deletions(-)
 create mode 100644 tests/test_jeol_dm4.py

diff --git a/examples/stem_overfocus.ipynb b/examples/stem_overfocus.ipynb
index 860ae99..8884285 100644
--- a/examples/stem_overfocus.ipynb
+++ b/examples/stem_overfocus.ipynb
@@ -1997,7 +1997,7 @@
    "outputs": [],
    "source": [
     "# We assume it is constant over the series\n",
-    "acceleration_voltage = jeol_dm4.acceleration_from_dm4(descan.dm4)"
+    "acceleration_voltage = jeol_dm4.acceleration_from_dm4(descan.tags)"
    ]
   },
   {
diff --git a/src/microscope_calibration/util/jeol_dm4.py b/src/microscope_calibration/util/jeol_dm4.py
index 44a1fef..34a3d6a 100644
--- a/src/microscope_calibration/util/jeol_dm4.py
+++ b/src/microscope_calibration/util/jeol_dm4.py
@@ -5,8 +5,10 @@
 To be tested and adapted to other microscopes!
 """
 
+import dataclasses
 import json
-from typing import Self
+from collections.abc import Mapping
+from typing import Any, Self
 
 import flax
 import numpy as np
@@ -17,36 +19,43 @@
 from microscope_calibration.common.model import DescanError, Model4DSTEM
 from microscope_calibration.ui import CalibratedDataset
 
+# Tags of the 4D STEM data. ImageList.1 is the thumbnail.
+IMAGE = ".ImageList.2"
+
+DM4Tags = Mapping[str, Any]
+
+
+def read_tags(path: str) -> DM4Tags:
+    """
+    All tags of a DM4 file, see :attr:`ncempy.io.dm.fileDM.allTags`
+    """
+    with fileDM(path) as handle:
+        return handle.allTags
+
 
 class CalibratedJEOLDM4(CalibratedDataset):
     def __init__(self, path: str, model: Model4DSTEM, ctx: Context | None = None):
-        handle = fileDM(path)
         self.path = path
-        self._tags = handle.allTags
+        self._tags = read_tags(path)
         if ctx is None:
             ctx = Context.make_with("inline")
         super().__init__(dataset=ctx.load("auto", path), model=model)
 
     @classmethod
-    def new(cls, path: str, model: Model4DSTEM = None, ctx: Context | None = None) -> Self:
-        handle = fileDM(path)
-        model = derive_model_from_dm4(handle, model)
+    def new(cls, path: str, model: Model4DSTEM | None = None, ctx: Context | None = None) -> Self:
+        model = derive_model_from_dm4(read_tags(path), model)
         return cls(path=path, model=model, ctx=ctx)
 
     def derive_relative_for(self, path: str, ctx: Context | None = None) -> Self:
-        old_handle = fileDM(self.path)
-        new_handle = fileDM(path)
-        model = derive_model_relative_dm4(old_dm4=old_handle, new_dm4=new_handle, model=self.model)
+        model = derive_model_relative_dm4(
+            old_tags=self.tags, new_tags=read_tags(path), model=self.model
+        )
         return self.__class__(
             path=path,
             model=model,
             ctx=ctx,
         )
 
-    @property
-    def dm4(self) -> fileDM:
-        return fileDM(self.path)
-
     def save(self, path: str):
         statedict = flax.serialization.to_state_dict((self.path, self.model.normalize_types()))
         with open(path, mode="w") as f:
@@ -62,104 +71,130 @@ def load(cls, path: str, ctx: Context | None = None) -> Self:
         return cls(path=dm4path, model=model, ctx=ctx)
 
     @property
-    def tags(self):
+    def tags(self) -> DM4Tags:
         return self._tags
 
     def microscope_info(self) -> dict:
-        result = {}
-        search = "ImageTags.Microscope Info."
-        for tag, value in self.tags.items():
-            if search in tag:
-                subtag = tag.rsplit(search, 1)[-1]
-                if isinstance(value, np.number):
-                    value = value.item()
-                elif isinstance(value, np.ndarray):
-                    if len(value.shape) != 1:
-                        raise RuntimeError("Only works for 1D shapes for now")
-                    value = tuple(it.item() for it in value)
-                result[subtag] = value
-        return result
+        return microscope_info(self.tags)
 
     def compare(self, other: CalibratedJEOLDM4) -> dict:
-        old_info = self.microscope_info()
-        new_info = other.microscope_info()
-        tmp_for_keys = old_info.copy()
-        tmp_for_keys.update(new_info)
-
-        dm4 = {}
-        for key in tmp_for_keys:
-            old_val = old_info.get(key)
-            new_val = new_info.get(key)
-            if old_val != new_val:
-                dm4[key] = (old_val, new_val)
-
-        model_attrs = list(Model4DSTEM.__dataclass_fields__.keys())
-        model_attrs.remove("descan_error")
-
-        old_model = self.model
-        new_model = other.model
-
-        model = {}
-
-        for attr in model_attrs:
-            old_val = getattr(old_model, attr)
-            new_val = getattr(new_model, attr)
-            if not np.allclose(old_val, new_val, rtol=1e-5, atol=1e-15):
-                model[attr] = (old_val, new_val)
-
-        descan_error_attrs = list(DescanError.__annotations__.keys())
-        descan_error = {}
-        for attr in descan_error_attrs:
-            old_val = getattr(old_model.descan_error, attr)
-            new_val = getattr(new_model.descan_error, attr)
-            if not np.allclose(old_val, new_val, rtol=1e-5, atol=1e-15):
-                descan_error[attr] = (old_val, new_val)
-
-        return {
-            "dm4": dm4,
-            "model": model,
-            "descan_error": descan_error,
-        }
+        return compare(
+            old_tags=self.tags, old_model=self.model, new_tags=other.tags, new_model=other.model
+        )
 
     def calibrated(self, model: Model4DSTEM, ctx: Context | None = None) -> CalibratedJEOLDM4:
         return self.__class__(path=self.path, model=model, ctx=ctx)
 
 
-def acceleration_from_dm4(dm4file: fileDM) -> pint.Quantity:
-    tags = dm4file.allTags
-
-    acceleration_voltage_V = tags[".ImageList.2.ImageTags.Microscope Info.Voltage"]
+def microscope_info(tags: DM4Tags) -> dict:
+    """
+    Microscope Info tags of the 4D STEM data, with keys relative to it
+    """
+    result = {}
+    search = f"{IMAGE}.ImageTags.Microscope Info."
+    for tag, value in tags.items():
+        if tag.startswith(search):
+            subtag = tag.removeprefix(search)
+            if isinstance(value, np.number):
+                value = value.item()
+            elif isinstance(value, np.ndarray):
+                if len(value.shape) != 1:
+                    raise RuntimeError("Only works for 1D shapes for now")
+                value = tuple(it.item() for it in value)
+            result[subtag] = value
+    return result
+
+
+def compare(
+    old_tags: DM4Tags, old_model: Model4DSTEM, new_tags: DM4Tags, new_model: Model4DSTEM
+) -> dict:
+    """
+    Differences of microscope info, model parameters and descan error
+
+    Returns
+    -------
+    differences
+        :code:`{"dm4": ..., "model": ..., "descan_error": ...}`, each a dict
+        of :code:`(old, new)` values for the keys that differ
+    """
+    old_info = microscope_info(old_tags)
+    new_info = microscope_info(new_tags)
+    dm4 = {}
+    for key in old_info | new_info:
+        old_val = old_info.get(key)
+        new_val = new_info.get(key)
+        if old_val != new_val:
+            dm4[key] = (old_val, new_val)
+
+    model = {}
+    for field in dataclasses.fields(Model4DSTEM):
+        if field.name == "descan_error":
+            continue
+        old_val = getattr(old_model, field.name)
+        new_val = getattr(new_model, field.name)
+        if not np.allclose(old_val, new_val, rtol=1e-5, atol=1e-15):
+            model[field.name] = (old_val, new_val)
+
+    descan_error = {}
+    for attr in DescanError._fields:
+        old_val = getattr(old_model.descan_error, attr)
+        new_val = getattr(new_model.descan_error, attr)
+        if not np.allclose(old_val, new_val, rtol=1e-5, atol=1e-15):
+            descan_error[attr] = (old_val, new_val)
+
+    return {
+        "dm4": dm4,
+        "model": model,
+        "descan_error": descan_error,
+    }
+
+
+def acceleration_from_dm4(tags: DM4Tags) -> pint.Quantity:
+    """
+    Acceleration voltage, see :func:`read_tags`
+    """
+    acceleration_voltage_V = tags[f"{IMAGE}.ImageTags.Microscope Info.Voltage"]
     return pint.Quantity(acceleration_voltage_V, "V")
 
 
-def derive_model_from_dm4(dm4file: fileDM, model: Model4DSTEM | None = None) -> Model4DSTEM:
-    tags = dm4file.allTags
+def derive_model_from_dm4(tags: DM4Tags, model: Model4DSTEM | None = None) -> Model4DSTEM:
+    """
+    Model with scan and detector calibration from DM4 tags, see :func:`read_tags`
+
+    Parameters
+    ----------
+    tags
+        Tags of a DM4 file
+    model
+        Model to start from. By default, :meth:`Model4DSTEM.default` for the
+        dataset shape.
+    """
     # 4D STEM DM4 was traditionally transposed so order of sig and nav is swapped
     shape_keys = [
         # This is nav!
-        ".ImageList.2.ImageData.Dimensions.3",
-        ".ImageList.2.ImageData.Dimensions.4",
+        f"{IMAGE}.ImageData.Dimensions.3",
+        f"{IMAGE}.ImageData.Dimensions.4",
         # this is sig!
-        ".ImageList.2.ImageData.Dimensions.1",
-        ".ImageList.2.ImageData.Dimensions.2",
+        f"{IMAGE}.ImageData.Dimensions.1",
+        f"{IMAGE}.ImageData.Dimensions.2",
     ]
     shape = tuple(int(tags[key]) for key in shape_keys)
 
     if model is None:
         model = Model4DSTEM.default(dataset_shape=shape)
 
-    scan_step_y_number = tags[".ImageList.2.ImageData.Calibrations.Dimension.3.Scale"]
-    scan_step_y_unit = tags[".ImageList.2.ImageData.Calibrations.Dimension.3.Units"]
+    scan_step_y_number = tags[f"{IMAGE}.ImageData.Calibrations.Dimension.3.Scale"]
+    scan_step_y_unit = tags[f"{IMAGE}.ImageData.Calibrations.Dimension.3.Units"]
     scan_step_y = pint.Quantity(scan_step_y_number, scan_step_y_unit)
 
-    scan_step_x_number = tags[".ImageList.2.ImageData.Calibrations.Dimension.4.Scale"]
-    scan_step_x_unit = tags[".ImageList.2.ImageData.Calibrations.Dimension.4.Units"]
+    scan_step_x_number = tags[f"{IMAGE}.ImageData.Calibrations.Dimension.4.Scale"]
+    scan_step_x_unit = tags[f"{IMAGE}.ImageData.Calibrations.Dimension.4.Units"]
     scan_step_x = pint.Quantity(scan_step_x_number, scan_step_x_unit)
     if scan_step_x != scan_step_y:
         raise ValueError("Requires uniform scan step in Y and X for the time being")
 
     # Note camera raw pixel pitch vs effective one (binning!)
-    cam_pixel_pitches = tags[".ImageList.2.ImageTags.Acquisition.Frame.CCD.Pixel Size (um)"]
+    cam_pixel_pitches = tags[f"{IMAGE}.ImageTags.Acquisition.Frame.CCD.Pixel Size (um)"]
     if len(cam_pixel_pitches) > 2:
         raise ValueError("Camera pixel pitch should have two dimensions")
     if len(cam_pixel_pitches) == 2 and cam_pixel_pitches[0] != cam_pixel_pitches[1]:
@@ -167,13 +202,13 @@ def derive_model_from_dm4(dm4file: fileDM, model: Model4DSTEM | None = None) ->
     cam_pixel_pitch = pint.Quantity(cam_pixel_pitches[0], "um")
 
     camera_length = pint.Quantity(
-        tags[".ImageList.2.ImageTags.Microscope Info.STEM Camera Length"], "mm"
+        tags[f"{IMAGE}.ImageTags.Microscope Info.STEM Camera Length"], "mm"
     )
     # Seems to be opposite of what Model4DSTEM works with
-    scan_rotation = pint.Quantity(-tags[".ImageList.2.ImageTags.DigiScan.Rotation"], "degree")
+    scan_rotation = pint.Quantity(-tags[f"{IMAGE}.ImageTags.DigiScan.Rotation"], "degree")
     # Seems to be opposite of what Model4DSTEM works with
     detector_rotation = pint.Quantity(
-        -tags[".ImageList.2.ImageTags.DigiScan.Rotation Offset"], "degree"
+        -tags[f"{IMAGE}.ImageTags.DigiScan.Rotation Offset"], "degree"
     )
     return model.derive(
         scan_pixel_pitch=scan_step_x.to("m").magnitude,
@@ -184,16 +219,21 @@ def derive_model_from_dm4(dm4file: fileDM, model: Model4DSTEM | None = None) ->
     )
 
 
-def derive_model_relative_dm4(old_dm4: fileDM, new_dm4: fileDM, model: Model4DSTEM) -> Model4DSTEM:
-    old_valmodel = derive_model_from_dm4(old_dm4)
-    new_valmodel = derive_model_from_dm4(new_dm4)
+def derive_model_relative_dm4(
+    old_tags: DM4Tags, new_tags: DM4Tags, model: Model4DSTEM
+) -> Model4DSTEM:
+    """
+    Apply the changes between two DM4 files to a model calibrated for the old one
 
-    old_tags = old_dm4.allTags
-    new_tags = new_dm4.allTags
+    Scan and detector pixel pitch and camera length are scaled, scan rotation
+    and overfocus (from the stage Z position) are shifted by the differences.
+    """
+    old_valmodel = derive_model_from_dm4(old_tags)
+    new_valmodel = derive_model_from_dm4(new_tags)
 
     # assume overfocus is adjusted with stage
     # TODO also calibrate focus length scale?
-    stage_z_key = ".ImageList.2.ImageTags.Microscope Info.Stage Position.Stage Z"
+    stage_z_key = f"{IMAGE}.ImageTags.Microscope Info.Stage Position.Stage Z"
     old_z = pint.Quantity(old_tags[stage_z_key], "um")
     new_z = pint.Quantity(new_tags[stage_z_key], "um")
 
diff --git a/tests/test_jeol_dm4.py b/tests/test_jeol_dm4.py
new file mode 100644
index 0000000..6211682
--- /dev/null
+++ b/tests/test_jeol_dm4.py
@@ -0,0 +1,165 @@
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+
+from microscope_calibration.common.model import DescanError, Model4DSTEM
+from microscope_calibration.util import jeol_dm4
+from microscope_calibration.util.jeol_dm4 import (
+    IMAGE,
+    CalibratedJEOLDM4,
+    acceleration_from_dm4,
+    compare,
+    derive_model_from_dm4,
+    derive_model_relative_dm4,
+    microscope_info,
+)
+
+
+def make_tags(**changes):
+    # Subset of the tags of a 4D STEM DM4 file, keys relative to IMAGE
+    tags = {
+        # sig
+        "ImageData.Dimensions.1": 64,
+        "ImageData.Dimensions.2": 48,
+        # nav
+        "ImageData.Dimensions.3": 32,
+        "ImageData.Dimensions.4": 16,
+        "ImageData.Calibrations.Dimension.3.Scale": 2.0,
+        "ImageData.Calibrations.Dimension.3.Units": "nm",
+        "ImageData.Calibrations.Dimension.4.Scale": 2.0,
+        "ImageData.Calibrations.Dimension.4.Units": "nm",
+        "ImageTags.Acquisition.Frame.CCD.Pixel Size (um)": np.array([55.0, 55.0]),
+        "ImageTags.Microscope Info.STEM Camera Length": np.float32(80.0),
+        "ImageTags.Microscope Info.Voltage": 200000.0,
+        "ImageTags.Microscope Info.Stage Position.Stage Z": 30.0,
+        "ImageTags.DigiScan.Rotation": 10.0,
+        "ImageTags.DigiScan.Rotation Offset": 3.0,
+    }
+    tags.update(changes)
+    res = {f"{IMAGE}.{key}": value for key, value in tags.items()}
+    # Thumbnail
+    res[".ImageList.1.ImageTags.Microscope Info.Voltage"] = 1.0
+    return res
+
+
+def test_derive_model():
+    model = derive_model_from_dm4(make_tags())
+    assert model.scan_center == (16.0, 8.0)
+    assert model.detector_center == (32.0, 24.0)
+    assert_allclose(model.scan_pixel_pitch, 2e-9)
+    assert_allclose(model.detector_pixel_pitch, 55e-6)
+    assert_allclose(model.camera_length, 0.08)
+    # Opposite sign of the DM4 tags
+    assert_allclose(model.scan_rotation, np.deg2rad(-10.0))
+    assert_allclose(model.detector_rotation, np.deg2rad(-3.0))
+
+
+def test_derive_model_start():
+    start = Model4DSTEM.default(dataset_shape=(1, 1, 1, 1)).derive(overfocus=0.1)
+    model = derive_model_from_dm4(make_tags(), model=start)
+    assert model.overfocus == 0.1
+    assert model.scan_center == start.scan_center
+
+
+@pytest.mark.parametrize(
+    "changes, match",
+    [
+        ({"ImageData.Calibrations.Dimension.4.Scale": 3.0}, "uniform scan step"),
+        ({"ImageTags.Acquisition.Frame.CCD.Pixel Size (um)": np.array([55.0, 50.0])}, "uniform"),
+        ({"ImageTags.Acquisition.Frame.CCD.Pixel Size (um)": np.ones(3)}, "two dimensions"),
+    ],
+)
+def test_derive_model_invalid(changes, match):
+    with pytest.raises(ValueError, match=match):
+        derive_model_from_dm4(make_tags(**changes))
+
+
+def test_acceleration():
+    assert acceleration_from_dm4(make_tags()).to("kV").magnitude == 200
+
+
+def test_derive_model_relative():
+    model = Model4DSTEM.default(dataset_shape=(32, 16, 64, 48)).derive(
+        overfocus=1e-3,
+        scan_pixel_pitch=3e-9,
+        detector_pixel_pitch=50e-6,
+        camera_length=0.1,
+        scan_rotation=0.2,
+    )
+    new_tags = make_tags(
+        **{
+            "ImageData.Calibrations.Dimension.3.Scale": 4.0,
+            "ImageData.Calibrations.Dimension.4.Scale": 4.0,
+            "ImageTags.Microscope Info.STEM Camera Length": 120.0,
+            "ImageTags.DigiScan.Rotation": 15.0,
+            "ImageTags.Microscope Info.Stage Position.Stage Z": 50.0,
+        }
+    )
+    res = derive_model_relative_dm4(old_tags=make_tags(), new_tags=new_tags, model=model)
+    assert_allclose(res.scan_pixel_pitch, 6e-9)
+    assert_allclose(res.detector_pixel_pitch, 50e-6)
+    assert_allclose(res.camera_length, 0.15)
+    assert_allclose(res.scan_rotation, 0.2 + np.deg2rad(-5.0))
+    # Stage Z up by 20 um reduces the overfocus
+    assert_allclose(res.overfocus, 1e-3 - 20e-6)
+
+
+def test_microscope_info():
+    info = microscope_info(make_tags())
+    assert info["Voltage"] == 200000.0
+    assert info["STEM Camera Length"] == 80.0
+    assert type(info["STEM Camera Length"]) is float
+    # Only the 4D STEM data, not the thumbnail
+    assert len(info) == 3
+
+
+def test_compare():
+    old_model = derive_model_from_dm4(make_tags())
+    new_tags = make_tags(**{"ImageTags.Microscope Info.Voltage": 300000.0})
+    new_model = old_model.derive(
+        overfocus=old_model.overfocus + 1.0, descan_error=DescanError(offpxi=1.0)
+    )
+    res = compare(old_tags=make_tags(), old_model=old_model, new_tags=new_tags, new_model=new_model)
+    assert res == {
+        "dm4": {"Voltage": (200000.0, 300000.0)},
+        "model": {"overfocus": (old_model.overfocus, new_model.overfocus)},
+        "descan_error": {"offpxi": (0.0, 1.0)},
+    }
+
+
+def test_read_tags_closes_file(monkeypatch):
+    events = []
+
+    class FakeFileDM:
+        def __init__(self, path):
+            self.allTags = {"path": path}
+
+        def __enter__(self):
+            events.append("enter")
+            return self
+
+        def __exit__(self, *args):
+            events.append("exit")
+
+    monkeypatch.setattr(jeol_dm4, "fileDM", FakeFileDM)
+    assert jeol_dm4.read_tags("test.dm4") == {"path": "test.dm4"}
+    assert events == ["enter", "exit"]
+
+
+class FakeContext:
+    def load(self, filetype, path):
+        return ("dataset", path)
+
+
+def test_save_load(monkeypatch, tmp_path):
+    monkeypatch.setattr(jeol_dm4, "read_tags", lambda path: make_tags())
+    ctx = FakeContext()
+    calibrated = CalibratedJEOLDM4.new("data.dm4", ctx=ctx)
+    assert calibrated.dataset == ("dataset", "data.dm4")
+    model = calibrated.model.derive(overfocus=0.3, descan_error=DescanError(sxo_pyi=0.5))
+    calibrated = calibrated.calibrated(model, ctx=ctx)
+    calibrated.save(tmp_path / "calibration.json")
+    loaded = CalibratedJEOLDM4.load(tmp_path / "calibration.json", ctx=ctx)
+    assert loaded.path == "data.dm4"
+    assert loaded.model == model.normalize_types()
+    assert loaded.compare(calibrated) == {"dm4": {}, "model": {}, "descan_error": {}}

From d2ea398f5b280d093fefda0926c9087071b4a116 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 15:05:54 +0000
Subject: [PATCH 24/34] Calculate the affine maps of the Numba kernels with one
 shared helper

The same dot product was written out seven times.

- Remove the project_tilt_*() and project_det_*() helpers, use affine_column()
- Run test_project_frame_forward() without JIT for coverage, too

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../common/stem_overfocus.py                  | 93 ++++---------------
 .../util/stem_overfocus_sim.py                | 42 +--------
 tests/test_sim.py                             |  1 +
 3 files changed, 25 insertions(+), 111 deletions(-)

diff --git a/src/microscope_calibration/common/stem_overfocus.py b/src/microscope_calibration/common/stem_overfocus.py
index 1c7d9dc..0eed665 100644
--- a/src/microscope_calibration/common/stem_overfocus.py
+++ b/src/microscope_calibration/common/stem_overfocus.py
@@ -367,49 +367,20 @@ def get_backward_transformation_matrix(
     )
 
 
-# Separate functions spun out to facilitate re-use of coordinate calculations
-# for other purposes
 @numba.njit(inline="always", cache=True)
-def project_tilt_y(image_y, image_x, scan_y, scan_x, mat):
-    return (
-        scan_y * mat[0, 2]
-        + scan_x * mat[1, 2]
-        + image_y * mat[2, 2]
-        + image_x * mat[3, 2]
-        + mat[4, 2]
-    )
-
-
-@numba.njit(inline="always", cache=True)
-def project_tilt_x(image_y, image_x, scan_y, scan_x, mat):
-    return (
-        scan_y * mat[0, 3]
-        + scan_x * mat[1, 3]
-        + image_y * mat[2, 3]
-        + image_x * mat[3, 3]
-        + mat[4, 3]
-    )
-
-
-@numba.njit(inline="always", cache=True)
-def project_det_y(image_y, image_x, scan_y, scan_x, mat):
-    return (
-        scan_y * mat[0, 0]
-        + scan_x * mat[1, 0]
-        + image_y * mat[2, 0]
-        + image_x * mat[3, 0]
-        + mat[4, 0]
-    )
-
+def affine_column(scan_y, scan_x, in_y, in_x, mat, column):
+    """
+    Column :code:`column` of :code:`(scan_y, scan_x, in_y, in_x, 1) @ mat`
 
-@numba.njit(inline="always", cache=True)
-def project_det_x(image_y, image_x, scan_y, scan_x, mat):
+    Calculating single columns of the matrix-vector product allows skipping
+    the rest early and facilitates auto-vectorization of loops.
+    """
     return (
-        scan_y * mat[0, 1]
-        + scan_x * mat[1, 1]
-        + image_y * mat[2, 1]
-        + image_x * mat[3, 1]
-        + mat[4, 1]
+        scan_y * mat[0, column]
+        + scan_x * mat[1, column]
+        + in_y * mat[2, column]
+        + in_x * mat[3, column]
+        + mat[4, column]
     )
 
 
@@ -418,22 +389,11 @@ def project_frame_backwards(frame, source_semiconv, mat, scan_y, scan_x, image_o
     limit = np.abs(np.tan(source_semiconv)) ** 2
     for image_y in range(image_out.shape[0]):
         for image_x in range(image_out.shape[1]):
-            # Manually unrolled matrix-vector product to allow skipping before
-            # calculating all values and facilitate auto-vectorization of the
-            # loop
-
-            # _one = (
-            #       scan_y * mat[0, 4] + scan_x * mat[1, 4]
-            #       + det_y * mat[2, 4] + det_x * mat[3, 4] + mat[4, 4]
-            # )
-            # assert np.allclose(_one, 1)
-            tilt_y = project_tilt_y(image_y, image_x, scan_y, scan_x, mat)
-            tilt_x = project_tilt_x(image_y, image_x, scan_y, scan_x, mat)
+            tilt_y = affine_column(scan_y, scan_x, image_y, image_x, mat, 2)
+            tilt_x = affine_column(scan_y, scan_x, image_y, image_x, mat, 3)
             if np.abs(tilt_y) ** 2 + np.abs(tilt_x) ** 2 < limit:
-                det_y = project_det_y(image_y, image_x, scan_y, scan_x, mat)
-                det_x = project_det_x(image_y, image_x, scan_y, scan_x, mat)
-                det_y = int(np.round(det_y))
-                det_x = int(np.round(det_x))
+                det_y = int(np.round(affine_column(scan_y, scan_x, image_y, image_x, mat, 0)))
+                det_x = int(np.round(affine_column(scan_y, scan_x, image_y, image_x, mat, 1)))
                 if det_y >= 0 and det_y < frame.shape[0] and det_x >= 0 and det_x < frame.shape[1]:
                     image_out[image_y, image_x] += frame[det_y, det_x]
 
@@ -498,37 +458,22 @@ def get_detector_correction_matrix(
     )
 
 
-# Separate functions spun out to facilitate re-use of coordinate calculations
-# for other purposes, such as corrected virtual detectors
+# Separate functions to re-use coordinate calculations for other purposes,
+# such as corrected virtual detectors
 @numba.njit(inline="always", cache=True)
 def corrected_det_y(det_corr_y, det_corr_x, scan_y, scan_x, mat):
-    return (
-        scan_y * mat[0, 0]
-        + scan_x * mat[1, 0]
-        + det_corr_y * mat[2, 0]
-        + det_corr_x * mat[3, 0]
-        + mat[4, 0]
-    )
+    return affine_column(scan_y, scan_x, det_corr_y, det_corr_x, mat, 0)
 
 
 @numba.njit(inline="always", cache=True)
 def corrected_det_x(det_corr_y, det_corr_x, scan_y, scan_x, mat):
-    return (
-        scan_y * mat[0, 1]
-        + scan_x * mat[1, 1]
-        + det_corr_y * mat[2, 1]
-        + det_corr_x * mat[3, 1]
-        + mat[4, 1]
-    )
+    return affine_column(scan_y, scan_x, det_corr_y, det_corr_x, mat, 1)
 
 
 @numba.njit(cache=True)
 def correct_frame(frame, mat, scan_y, scan_x, detector_out):
     for det_corr_y in range(detector_out.shape[0]):
         for det_corr_x in range(detector_out.shape[1]):
-            # Manually unrolled matrix-vector product to allow skipping before
-            # calculating all values and facilitate auto-vectorization of the
-            # loop
             det_y = corrected_det_y(det_corr_y, det_corr_x, scan_y, scan_x, mat)
             det_x = corrected_det_x(det_corr_y, det_corr_x, scan_y, scan_x, mat)
             det_y = int(np.round(det_y))
diff --git a/src/microscope_calibration/util/stem_overfocus_sim.py b/src/microscope_calibration/util/stem_overfocus_sim.py
index 6dbc0fd..71fbec8 100644
--- a/src/microscope_calibration/util/stem_overfocus_sim.py
+++ b/src/microscope_calibration/util/stem_overfocus_sim.py
@@ -10,6 +10,7 @@
 from microscope_calibration.common.stem_overfocus import (
     CoordMappingT,
     _eliminate_tilt,
+    affine_column,
     ray_transfer_matrix,
 )
 # ruff: enable[E402]
@@ -118,44 +119,11 @@ def project_frame_forward(obj, source_semiconv, mat, scan_y, scan_x, out):
     limit = np.abs(np.tan(source_semiconv)) ** 2
     for det_y in range(out.shape[0]):
         for det_x in range(out.shape[1]):
-            # Manually unrolled matrix-vector product to allow skipping before
-            # calculating all values and facilitate auto-vectorization of the
-            # loop
-
-            # _one = (
-            #       scan_y * mat[0, 4] + scan_x * mat[1, 4]
-            #       + det_y * mat[2, 4] + det_x * mat[3, 4] + mat[4, 4]
-            # )
-            # assert np.allclose(_one, 1)
-            tilt_y = (
-                scan_y * mat[0, 2]
-                + scan_x * mat[1, 2]
-                + det_y * mat[2, 2]
-                + det_x * mat[3, 2]
-                + mat[4, 2]
-            )
-            tilt_x = (
-                scan_y * mat[0, 3]
-                + scan_x * mat[1, 3]
-                + det_y * mat[2, 3]
-                + det_x * mat[3, 3]
-                + mat[4, 3]
-            )
+            tilt_y = affine_column(scan_y, scan_x, det_y, det_x, mat, 2)
+            tilt_x = affine_column(scan_y, scan_x, det_y, det_x, mat, 3)
             if np.abs(tilt_y) ** 2 + np.abs(tilt_x) ** 2 < limit:
-                spec_y = (
-                    scan_y * mat[0, 0]
-                    + scan_x * mat[1, 0]
-                    + det_y * mat[2, 0]
-                    + det_x * mat[3, 0]
-                    + mat[4, 0]
-                )
-                spec_x = (
-                    scan_y * mat[0, 1]
-                    + scan_x * mat[1, 1]
-                    + det_y * mat[2, 1]
-                    + det_x * mat[3, 1]
-                    + mat[4, 1]
-                )
+                spec_y = affine_column(scan_y, scan_x, det_y, det_x, mat, 0)
+                spec_x = affine_column(scan_y, scan_x, det_y, det_x, mat, 1)
                 spec_y = int(np.round(spec_y))
                 spec_x = int(np.round(spec_x))
                 if spec_y >= 0 and spec_y < obj.shape[0] and spec_x >= 0 and spec_x < obj.shape[1]:
diff --git a/tests/test_sim.py b/tests/test_sim.py
index 8d8dae0..4aee26a 100644
--- a/tests/test_sim.py
+++ b/tests/test_sim.py
@@ -24,6 +24,7 @@
 trace = lambdify_trace_for(np)
 
 
+@pytest.mark.with_numba
 def test_project_frame_forward():
     for repeat in range(10):
         scan_y = np.random.random()

From 55c11c77f607a46278029007a5117e612fbf8678 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 15:06:43 +0000
Subject: [PATCH 25/34] Check the number of reflections in test_twothetas
 instead of truncating

The 133 reflection of the reference is outside the default reciprocal radius.

- Test that the angles scale with the electron wavelength at 300 kV

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/test_diffraction.py | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/tests/test_diffraction.py b/tests/test_diffraction.py
index 9a12588..c265ab4 100644
--- a/tests/test_diffraction.py
+++ b/tests/test_diffraction.py
@@ -22,7 +22,8 @@ def test_twothetas(request):
         20.45,  # 113, 0.1227 nm
         21.36,  # 222, 0.1175 nm
         24.67,  # 400, 0.1018 nm
-        26.88,  # 133, 0.0934 nm
+        # Outside of the default reciprocal_radius of get_twothetas():
+        # 26.88,  # 133, 0.0934 nm
     ]
 
     cif_path = os.path.join(os.path.dirname(request.path), "AuEntryWithCollCode163723.cif")
@@ -56,18 +57,25 @@ def test_twothetas(request):
 
     twothetas = get_twothetas(cif_path, 200000)
 
+    assert twothetas[0] == 0
     without_zero = twothetas[1:]
-
-    print(without_zero, expected_111_twotheta, reference_au, corrected_au)
+    assert len(without_zero) == len(corrected_au)
 
     # 111 is the first reflection with intensity
     # get_twothetas rounds to 5 decimals
     assert_almost_equal(without_zero[0], np.round(expected_111_twotheta, decimals=5), decimal=4)
 
-    if len(without_zero) > len(corrected_au):
-        without_zero = without_zero[: len(corrected_au)]
-    elif len(without_zero) < len(corrected_au):
-        corrected_au = corrected_au[: len(without_zero)]
-
     # convert reference to radian
     assert_allclose(corrected_au / 1000, without_zero, rtol=5e-3)
+
+
+def test_twothetas_voltage(request):
+    cif_path = os.path.join(os.path.dirname(request.path), "AuEntryWithCollCode163723.cif")
+    twothetas_200 = get_twothetas(cif_path, 200000)
+    twothetas_300 = get_twothetas(cif_path, 300000)
+    # Relativistic electron wavelength in pm,
+    # https://www.jeol.com/words/emterms/20121023.071258.php#gsc.tab=0
+    wavelength_200 = 2.5079
+    wavelength_300 = 1.9687
+    # Small angles scale with the wavelength, get_twothetas() rounds to 1e-5
+    assert_allclose(twothetas_300, twothetas_200 * wavelength_300 / wavelength_200, atol=4e-5)

From 0b1a538a5fbdbbd6ec614e5a6261f6ea32acea9a Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 15:12:13 +0000
Subject: [PATCH 26/34] Include underfocus in the random_model test fixture

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/conftest.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index d91b2ad..c1155f3 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -14,15 +14,21 @@ def _seed_random(request):
 
 @pytest.fixture
 def random_model() -> Model4DSTEM:
+    camera_length = np.random.uniform(0.1, 2)
+    if np.random.random() < 0.5:
+        overfocus = np.random.uniform(0.1, 2)
+    else:
+        # Underfocus, not close to a crossover at the detector
+        overfocus = -np.random.uniform(0.05, 0.9) * camera_length
     return Model4DSTEM(
-        overfocus=np.random.uniform(0.1, 2),
+        overfocus=overfocus,
         scan_pixel_pitch=np.random.uniform(0.01, 2),
         scan_center=PixelYX(
             y=np.random.uniform(-10, 10),
             x=np.random.uniform(-10, 10),
         ),
         scan_rotation=np.random.uniform(-np.pi, np.pi),
-        camera_length=np.random.uniform(0.1, 2),
+        camera_length=camera_length,
         detector_pixel_pitch=np.random.uniform(0.01, 2),
         detector_center=PixelYX(
             y=np.random.uniform(-10, 10),

From 54804e5d9f43c3917984fea5b5d47e1785af0da7 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Tue, 13 Oct 2026 15:13:38 +0000
Subject: [PATCH 27/34] Combine the test_correct variants with shared helpers
 and non-square shapes

Square shapes and centers couldn't detect mix-ups of y and x.

- test_correct_smallnumba is a parameter set of test_correct_fixed_manualref

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/test_overfocus.py | 525 +++++++++++-----------------------------
 1 file changed, 142 insertions(+), 383 deletions(-)

diff --git a/tests/test_overfocus.py b/tests/test_overfocus.py
index 81ca085..2b5d137 100644
--- a/tests/test_overfocus.py
+++ b/tests/test_overfocus.py
@@ -367,89 +367,72 @@ def test_backproject_scale_scanref():
     assert_allclose(projected[:, :, 32, 32], res)
 
 
-@pytest.mark.parametrize("scan_rotation", (0.0, np.pi / 2))
-@pytest.mark.parametrize("detector_rotation", (0.0, np.pi / 2))
-@pytest.mark.parametrize("flip_factor", (1.0, -1.0))
-@pytest.mark.parametrize("manual_reference", (False, True))
-def test_correct(scan_rotation, detector_rotation, flip_factor, manual_reference):
-    scan_pixel_pitch = 0.1
-    detector_pixel_pitch = 0.2
-    overfocus = 1.0
-    camera_length = 1.0
-    propagation_distance = overfocus + camera_length
-    obj_half_size = 16
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.00314157, propagation_distance)
+SCAN_PIXEL_PITCH = 0.1
+DETECTOR_PIXEL_PITCH = 0.2
+CAMERA_LENGTH = 1.0
 
-    model = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
-        scan_rotation=scan_rotation,
-        flip_factor=flip_factor,
-        # Simulate detector larger than object to avoid clipping at the borders
-        detector_center=PixelYX(x=obj_half_size * 2, y=obj_half_size * 2),
-        detector_rotation=detector_rotation,
-        # Descan error designed to give whole pixel shifts
-        descan_error=DescanError(
-            offpxi=detector_pixel_pitch,
-            offpyi=detector_pixel_pitch * 2,
-            offsxi=-1 * detector_pixel_pitch / camera_length,
-            offsyi=-2 * detector_pixel_pitch / camera_length,
-            pxo_pxi=2 * detector_pixel_pitch / scan_pixel_pitch,
-            pyo_pyi=3 * detector_pixel_pitch / scan_pixel_pitch,
-            sxo_pxi=-3 * detector_pixel_pitch / scan_pixel_pitch / camera_length,
-            syo_pyi=-4 * detector_pixel_pitch / scan_pixel_pitch / camera_length,
-        ),
-    )
-    # Manual reference parametes to check that code path
-    # Should be identical to the default calculated by get_detector_correction_matrix()
-    # Note that this rotates the detector to follow the scan in order to cancel out
-    # the scan rotation.
-    model_ref_manual = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
-        scan_rotation=0.0,
-        flip_factor=1.0,
-        detector_center=PixelYX(x=obj_half_size * 2, y=obj_half_size * 2),
-        detector_rotation=scan_rotation,
-        descan_error=DescanError(),
+
+def correction_shapes(obj_half_size: int) -> tuple[tuple[int, int], tuple[int, int]]:
+    # Not square to detect mix-ups of y and x
+    scan_shape = (obj_half_size * 2, obj_half_size * 2 + 2)
+    # Detector larger than object to avoid clipping at the borders
+    detector_shape = (obj_half_size * 4, obj_half_size * 4 + 6)
+    return scan_shape, detector_shape
+
+
+def correction_model(obj_half_size: int, **changes) -> Model4DSTEM:
+    """
+    Model without aberrations for the shapes of :func:`correction_shapes`
+    """
+    overfocus = 1.0
+    scan_shape, detector_shape = correction_shapes(obj_half_size)
+    angle = np.arctan2(
+        obj_half_size * DETECTOR_PIXEL_PITCH / 2 + 0.00314157, overfocus + CAMERA_LENGTH
     )
-    # Parameters for simulated result without aberrations
-    model_ref_sim = Model4DSTEM(
+    return Model4DSTEM(
         overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch,
+        scan_pixel_pitch=SCAN_PIXEL_PITCH,
+        camera_length=CAMERA_LENGTH,
+        detector_pixel_pitch=DETECTOR_PIXEL_PITCH,
         semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
+        scan_center=PixelYX(y=scan_shape[0] / 2, x=scan_shape[1] / 2),
         scan_rotation=0.0,
         flip_factor=1.0,
-        detector_center=PixelYX(x=obj_half_size * 2, y=obj_half_size * 2),
+        detector_center=PixelYX(y=detector_shape[0] / 2, x=detector_shape[1] / 2),
         detector_rotation=0.0,
         descan_error=DescanError(),
-    )
-    # Obtain correction matrix for 4D STEM dataset, i.e. transform the data as
-    # if the rotations were 0, no flip, and the descan error was 0.
-    mat = get_detector_correction_matrix(
-        rec_model=model,
-        ref_model=model_ref_manual if manual_reference else None,
-    )
-    obj = np.random.random((obj_half_size * 2, obj_half_size * 2))
+    ).derive(**changes)
+
+
+# Descan error designed to give whole pixel shifts
+WHOLE_PIXEL_DESCAN_ERROR = DescanError(
+    offpxi=DETECTOR_PIXEL_PITCH,
+    offpyi=DETECTOR_PIXEL_PITCH * 2,
+    offsxi=-1 * DETECTOR_PIXEL_PITCH / CAMERA_LENGTH,
+    offsyi=-2 * DETECTOR_PIXEL_PITCH / CAMERA_LENGTH,
+    pxo_pxi=2 * DETECTOR_PIXEL_PITCH / SCAN_PIXEL_PITCH,
+    pyo_pyi=3 * DETECTOR_PIXEL_PITCH / SCAN_PIXEL_PITCH,
+    sxo_pxi=-3 * DETECTOR_PIXEL_PITCH / SCAN_PIXEL_PITCH / CAMERA_LENGTH,
+    syo_pyi=-4 * DETECTOR_PIXEL_PITCH / SCAN_PIXEL_PITCH / CAMERA_LENGTH,
+)
+
+
+def check_correction(obj_half_size, model, ref_model, sim_ref_model, specimen_to_image=None):
+    """
+    Correct a dataset simulated with :code:`model` according to
+    :code:`ref_model` and compare with a dataset simulated with
+    :code:`sim_ref_model`
+    """
+    scan_shape, detector_shape = correction_shapes(obj_half_size)
+    mat = get_detector_correction_matrix(rec_model=model, ref_model=ref_model)
+    obj = np.random.random(scan_shape)
     projected = project(
         image=obj,
-        detector_shape=(obj_half_size * 4, obj_half_size * 4),
-        scan_shape=(obj_half_size * 2, obj_half_size * 2),
+        detector_shape=detector_shape,
+        scan_shape=scan_shape,
         sim_model=model,
+        specimen_to_image=specimen_to_image,
     )
-    # Calculate corrected 4D STEM dataset, i.e. as if the rotations were 0,
-    # no flip, and the descan error was 0.
     out = np.zeros_like(projected)
     for scan_y in range(out.shape[0]):
         for scan_x in range(out.shape[1]):
@@ -462,350 +445,126 @@ def test_correct(scan_rotation, detector_rotation, flip_factor, manual_reference
             )
     projected_ref = project(
         image=obj,
-        detector_shape=(obj_half_size * 4, obj_half_size * 4),
-        scan_shape=(obj_half_size * 2, obj_half_size * 2),
-        sim_model=model_ref_sim,
+        detector_shape=detector_shape,
+        scan_shape=scan_shape,
+        sim_model=sim_ref_model,
+        # Detector image correction doesn't interfere with
+        # how scan positions are mapped, so we have to use the same mapping here
+        specimen_to_image=specimen_to_image,
     )
-    # 100 % match between corrected frames and simulated reference without aberrations
+    # 100 % match between corrected frames and simulated reference
     assert_allclose(projected_ref, out)
-    # no descan error left: Trace of central pixel is the object
-    assert_allclose(obj, out[:, :, obj_half_size * 2, obj_half_size * 2])
-    # Counter-test: Trace of central pixel of simulate ddataset with descan
-    # error doesn't match the object
-    assert not np.allclose(obj, projected[:, :, obj_half_size * 2, obj_half_size * 2])
+    return obj, projected, out
 
 
 @pytest.mark.parametrize("scan_rotation", (0.0, np.pi / 2))
 @pytest.mark.parametrize("detector_rotation", (0.0, np.pi / 2))
-def test_correct_flip(scan_rotation, detector_rotation):
-    scan_pixel_pitch = 0.1
-    detector_pixel_pitch = 0.2
-    overfocus = 1.0
-    camera_length = 1.0
-    propagation_distance = overfocus + camera_length
+@pytest.mark.parametrize("flip_factor", (1.0, -1.0))
+@pytest.mark.parametrize("manual_reference", (False, True))
+def test_correct(scan_rotation, detector_rotation, flip_factor, manual_reference):
     obj_half_size = 16
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.00314157, propagation_distance)
-
-    model = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
+    model = correction_model(
+        obj_half_size,
         scan_rotation=scan_rotation,
-        flip_factor=1.0,
-        # Simulate detector larger than object to avoid clipping at the borders
-        detector_center=PixelYX(x=obj_half_size * 2, y=obj_half_size * 2),
+        flip_factor=flip_factor,
         detector_rotation=detector_rotation,
-        # Descan error designed to give whole pixel shifts
-        descan_error=DescanError(
-            offpxi=detector_pixel_pitch,
-            offpyi=detector_pixel_pitch * 2,
-            offsxi=-1 * detector_pixel_pitch / camera_length,
-            offsyi=-2 * detector_pixel_pitch / camera_length,
-            pxo_pxi=2 * detector_pixel_pitch / scan_pixel_pitch,
-            pyo_pyi=3 * detector_pixel_pitch / scan_pixel_pitch,
-            sxo_pxi=-3 * detector_pixel_pitch / scan_pixel_pitch / camera_length,
-            syo_pyi=-4 * detector_pixel_pitch / scan_pixel_pitch / camera_length,
-        ),
-    )
-    # Manual reference parametes that introduce flip_y
-    # and compensate the rotations
-    model_ref_manual = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
-        scan_rotation=0.0,
-        flip_factor=-1.0,
-        detector_center=PixelYX(x=obj_half_size * 2, y=obj_half_size * 2),
-        detector_rotation=scan_rotation,
-        descan_error=DescanError(),
-    )
-    # Parameters for simulated result with flip_y
-    model_ref_sim = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
-        scan_rotation=0.0,
-        flip_factor=-1.0,
-        detector_center=PixelYX(x=obj_half_size * 2, y=obj_half_size * 2),
-        detector_rotation=0.0,
-        descan_error=DescanError(),
-    )
-    # Obtain correction matrix for 4D STEM dataset that transforms the data as
-    # if the rotations were 0, flip_y, and the descan error was 0.
-    mat = get_detector_correction_matrix(
-        rec_model=model,
-        ref_model=model_ref_manual,
-    )
-    obj = np.random.random((obj_half_size * 2, obj_half_size * 2))
-    projected = project(
-        image=obj,
-        detector_shape=(obj_half_size * 4, obj_half_size * 4),
-        scan_shape=(obj_half_size * 2, obj_half_size * 2),
-        sim_model=model,
-    )
-    # Calculate corrected 4D STEM dataset with the rotations were 0,
-    # flip, and no descan error.
-    out = np.zeros_like(projected)
-    for scan_y in range(out.shape[0]):
-        for scan_x in range(out.shape[1]):
-            correct_frame(
-                frame=projected[scan_y, scan_x],
-                mat=mat,
-                scan_y=scan_y,
-                scan_x=scan_x,
-                detector_out=out[scan_y, scan_x],
-            )
-    projected_ref = project(
-        image=obj,
-        detector_shape=(obj_half_size * 4, obj_half_size * 4),
-        scan_shape=(obj_half_size * 2, obj_half_size * 2),
-        sim_model=model_ref_sim,
-    )
-    # 100 % match between corrected frames and simulated reference without aberrations
-    assert_allclose(projected_ref, out)
+        descan_error=WHOLE_PIXEL_DESCAN_ERROR,
+    )
+    # Manual reference parameters to check that code path. Should be identical
+    # to the default calculated by get_detector_correction_matrix(). Note that
+    # this rotates the detector to follow the scan in order to cancel out the
+    # scan rotation.
+    ref_model = correction_model(obj_half_size, detector_rotation=scan_rotation)
+    # Transform the data as if the rotations were 0, no flip, and the descan error was 0.
+    obj, projected, out = check_correction(
+        obj_half_size,
+        model=model,
+        ref_model=ref_model if manual_reference else None,
+        sim_ref_model=correction_model(obj_half_size),
+    )
+    center = model.detector_center
+    center = (int(center.y), int(center.x))
+    # no descan error left: Trace of central pixel is the object
+    assert_allclose(obj, out[:, :, *center])
+    # Counter-test: Trace of central pixel of simulated dataset with descan
+    # error doesn't match the object
+    assert not np.allclose(obj, projected[:, :, *center])
 
 
 @pytest.mark.parametrize("scan_rotation", (0.0, np.pi / 2))
 @pytest.mark.parametrize("detector_rotation", (0.0, np.pi / 2))
-def test_correct_fixed_manualref(scan_rotation, detector_rotation):
-    scan_pixel_pitch = 0.1
-    detector_pixel_pitch = 0.2
-    overfocus = 1.0
-    camera_length = 1.0
-    propagation_distance = overfocus + camera_length
+def test_correct_flip(scan_rotation, detector_rotation):
     obj_half_size = 16
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.00314157, propagation_distance)
-
-    # Fixed mapping from physical to image for forward simulations
-    def map_coord(inp):
-        cy = obj.shape[0] / 2
-        cx = obj.shape[1] / 2
-        inp_vec = np.array((inp.y, inp.x))
-        y, x = rotate(-np.pi / 2) @ scale(1 / scan_pixel_pitch) @ inp_vec
-        return PixelYX(y=y + cy + 2, x=x + cx - 3)
-
-    model = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
+    model = correction_model(
+        obj_half_size,
         scan_rotation=scan_rotation,
-        flip_factor=1.0,
-        # Simulate detector larger than object to avoid clipping at the borders
-        detector_center=PixelYX(x=obj_half_size * 2, y=obj_half_size * 2),
         detector_rotation=detector_rotation,
-        # Descan error designed to give whole pixel shifts
-        descan_error=DescanError(
-            offpxi=detector_pixel_pitch,
-            offpyi=detector_pixel_pitch * 2,
-            offsxi=-1 * detector_pixel_pitch / camera_length,
-            offsyi=-2 * detector_pixel_pitch / camera_length,
-            pxo_pxi=2 * detector_pixel_pitch / scan_pixel_pitch,
-            pyo_pyi=3 * detector_pixel_pitch / scan_pixel_pitch,
-            sxo_pxi=-3 * detector_pixel_pitch / scan_pixel_pitch / camera_length,
-            syo_pyi=-4 * detector_pixel_pitch / scan_pixel_pitch / camera_length,
-        ),
-    )
-    # Manual reference parametes that introduce flip_y
-    # and compensate the rotations
-    model_ref_manual = Model4DSTEM(
-        overfocus=overfocus,
-        # No impact on correction
-        scan_pixel_pitch=scan_pixel_pitch * 42,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch * 2,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
-        # Has no impact since we don't remap the scan dimension,
-        # only the projection after the specimen
-        scan_rotation=np.pi / 23,
-        flip_factor=-1.0,
-        detector_center=PixelYX(x=obj_half_size * 2 - 1, y=obj_half_size * 2 + 2),
-        detector_rotation=0.0,
-        descan_error=DescanError(),
-    )
-    # Parameters for simulated result with flip_y
-    model_ref_sim = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch * 2,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
-        # Has to match the input scan rotation since we don't
-        # remap the scan dimension
-        scan_rotation=scan_rotation,
-        flip_factor=-1.0,
-        detector_center=PixelYX(x=obj_half_size * 2 - 1, y=obj_half_size * 2 + 2),
-        detector_rotation=0.0,
-        descan_error=DescanError(),
-    )
-    # Obtain correction matrix for 4D STEM dataset that transforms the data as
-    # if the rotations were 0, flip_y, and the descan error was 0.
-    mat = get_detector_correction_matrix(
-        rec_model=model,
-        ref_model=model_ref_manual,
-    )
-    obj = np.random.random((obj_half_size * 2, obj_half_size * 2))
-    projected = project(
-        image=obj,
-        detector_shape=(obj_half_size * 4, obj_half_size * 4),
-        scan_shape=(obj_half_size * 2, obj_half_size * 2),
-        sim_model=model,
-        # Detector image correction doesn't interfere with
-        # how scan positions are mapped
-        specimen_to_image=map_coord,
-    )
-    # Calculate corrected 4D STEM dataset with the rotations were 0,
-    # flip, and no descan error.
-    out = np.zeros_like(projected)
-    for scan_y in range(out.shape[0]):
-        for scan_x in range(out.shape[1]):
-            correct_frame(
-                frame=projected[scan_y, scan_x],
-                mat=mat,
-                scan_y=scan_y,
-                scan_x=scan_x,
-                detector_out=out[scan_y, scan_x],
-            )
-    projected_ref = project(
-        image=obj,
-        detector_shape=(obj_half_size * 4, obj_half_size * 4),
-        scan_shape=(obj_half_size * 2, obj_half_size * 2),
-        sim_model=model_ref_sim,
-        # Detector image correction doesn't interfere with
-        # how scan positions are mapped, so we have to use the same mapping here
-        specimen_to_image=map_coord,
-    )
-    # 100 % match between corrected frames and simulated reference without aberrations
-    assert_allclose(projected_ref, out)
-
-
-@pytest.mark.with_numba
-def test_correct_smallnumba():
-    scan_rotation = 0.0
-    detector_rotation = 0.0
-    scan_pixel_pitch = 0.1
-    detector_pixel_pitch = 0.2
-    overfocus = 1.0
-    camera_length = 1.0
-    propagation_distance = overfocus + camera_length
-    obj_half_size = 1
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.00314157, propagation_distance)
+        descan_error=WHOLE_PIXEL_DESCAN_ERROR,
+    )
+    # Manual reference parameters that introduce flip_y and compensate the rotations
+    ref_model = correction_model(obj_half_size, flip_factor=-1.0, detector_rotation=scan_rotation)
+    # Transform the data as if the rotations were 0, flip_y, and the descan error was 0.
+    check_correction(
+        obj_half_size,
+        model=model,
+        ref_model=ref_model,
+        sim_ref_model=correction_model(obj_half_size, flip_factor=-1.0),
+    )
+
+
+@pytest.mark.parametrize(
+    "scan_rotation, detector_rotation, obj_half_size",
+    [
+        (0.0, 0.0, 16),
+        (0.0, np.pi / 2, 16),
+        (np.pi / 2, 0.0, 16),
+        (np.pi / 2, np.pi / 2, 16),
+        pytest.param(0.0, 0.0, 1, marks=pytest.mark.with_numba),
+    ],
+)
+def test_correct_fixed_manualref(scan_rotation, detector_rotation, obj_half_size):
+    scan_shape, _ = correction_shapes(obj_half_size)
 
     # Fixed mapping from physical to image for forward simulations
     def map_coord(inp):
-        cy = obj.shape[0] / 2
-        cx = obj.shape[1] / 2
+        cy = scan_shape[0] / 2
+        cx = scan_shape[1] / 2
         inp_vec = np.array((inp.y, inp.x))
-        y, x = rotate(-np.pi / 2) @ scale(1 / scan_pixel_pitch) @ inp_vec
+        y, x = rotate(-np.pi / 2) @ scale(1 / SCAN_PIXEL_PITCH) @ inp_vec
         return PixelYX(y=y + cy + 2, x=x + cx - 3)
 
-    model = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
+    model = correction_model(
+        obj_half_size,
         scan_rotation=scan_rotation,
-        flip_factor=1.0,
-        # Simulate detector larger than object to avoid clipping at the borders
-        detector_center=PixelYX(x=obj_half_size * 2, y=obj_half_size * 2),
         detector_rotation=detector_rotation,
-        # Descan error designed to give whole pixel shifts
-        descan_error=DescanError(
-            offpxi=detector_pixel_pitch,
-            offpyi=detector_pixel_pitch * 2,
-            offsxi=-1 * detector_pixel_pitch / camera_length,
-            offsyi=-2 * detector_pixel_pitch / camera_length,
-            pxo_pxi=2 * detector_pixel_pitch / scan_pixel_pitch,
-            pyo_pyi=3 * detector_pixel_pitch / scan_pixel_pitch,
-            sxo_pxi=-3 * detector_pixel_pitch / scan_pixel_pitch / camera_length,
-            syo_pyi=-4 * detector_pixel_pitch / scan_pixel_pitch / camera_length,
-        ),
+        descan_error=WHOLE_PIXEL_DESCAN_ERROR,
     )
-    # Manual reference parametes that introduce flip_y
-    # and compensate the rotations
-    model_ref_manual = Model4DSTEM(
-        overfocus=overfocus,
+    base = correction_model(obj_half_size)
+    detector_center = PixelYX(y=base.detector_center.y + 2, x=base.detector_center.x - 1)
+    # Manual reference parameters that introduce flip_y and compensate the rotations
+    ref_model = base.derive(
         # No impact on correction
-        scan_pixel_pitch=scan_pixel_pitch * 42,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch * 2,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
+        scan_pixel_pitch=SCAN_PIXEL_PITCH * 42,
+        detector_pixel_pitch=DETECTOR_PIXEL_PITCH * 2,
         # Has no impact since we don't remap the scan dimension,
         # only the projection after the specimen
         scan_rotation=np.pi / 23,
         flip_factor=-1.0,
-        detector_center=PixelYX(x=obj_half_size * 2 - 1, y=obj_half_size * 2 + 2),
-        detector_rotation=0.0,
-        descan_error=DescanError(),
+        detector_center=detector_center,
     )
-    # Parameters for simulated result with flip_y
-    model_ref_sim = Model4DSTEM(
-        overfocus=overfocus,
-        scan_pixel_pitch=scan_pixel_pitch,
-        camera_length=camera_length,
-        detector_pixel_pitch=detector_pixel_pitch * 2,
-        semiconv=angle,
-        scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
-        # Has to match the input scan rotation since we don't
-        # remap the scan dimension
+    # Simulated result with flip_y
+    sim_ref_model = base.derive(
+        detector_pixel_pitch=DETECTOR_PIXEL_PITCH * 2,
+        # Has to match the input scan rotation since we don't remap the scan dimension
         scan_rotation=scan_rotation,
         flip_factor=-1.0,
-        detector_center=PixelYX(x=obj_half_size * 2 - 1, y=obj_half_size * 2 + 2),
-        detector_rotation=0.0,
-        descan_error=DescanError(),
-    )
-    # Obtain correction matrix for 4D STEM dataset that transforms the data as
-    # if the rotations were 0, flip_y, and the descan error was 0.
-    mat = get_detector_correction_matrix(
-        rec_model=model,
-        ref_model=model_ref_manual,
-    )
-    obj = np.random.random((obj_half_size * 2, obj_half_size * 2))
-    projected = project(
-        image=obj,
-        detector_shape=(obj_half_size * 4, obj_half_size * 4),
-        scan_shape=(obj_half_size * 2, obj_half_size * 2),
-        sim_model=model,
-        # Detector image correction doesn't interfere with
-        # how scan positions are mapped
+        detector_center=detector_center,
+    )
+    # Detector image correction doesn't interfere with how scan positions are mapped
+    check_correction(
+        obj_half_size,
+        model=model,
+        ref_model=ref_model,
+        sim_ref_model=sim_ref_model,
         specimen_to_image=map_coord,
     )
-    # Calculate corrected 4D STEM dataset with the rotations were 0,
-    # flip, and no descan error.
-    out = np.zeros_like(projected)
-    for scan_y in range(out.shape[0]):
-        for scan_x in range(out.shape[1]):
-            correct_frame(
-                frame=projected[scan_y, scan_x],
-                mat=mat,
-                scan_y=scan_y,
-                scan_x=scan_x,
-                detector_out=out[scan_y, scan_x],
-            )
-    projected_ref = project(
-        image=obj,
-        detector_shape=(obj_half_size * 4, obj_half_size * 4),
-        scan_shape=(obj_half_size * 2, obj_half_size * 2),
-        sim_model=model_ref_sim,
-        # Detector image correction doesn't interfere with
-        # how scan positions are mapped, so we have to use the same mapping here
-        specimen_to_image=map_coord,
-    )
-    # 100 % match between corrected frames and simulated reference without aberrations
-    assert_allclose(projected_ref, out)

From 96c9aaee370e07c80d67ec9d297485952497ed1d Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Wed, 14 Oct 2026 07:54:57 +0000
Subject: [PATCH 28/34] Treat semiconv consistently as the slope at the source

The kernels used it as an angle, the transformation matrices and the beam
radius as a slope. See the Model4DSTEM docstring.

Behavior changes:
- Back-projection and simulation limit the source slope to semiconv instead of tan(semiconv)
- The UI still shows the convergence semi-angle and converts it

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CLAUDE.md                                     |  2 +-
 src/microscope_calibration/common/model.py    | 14 ++++++--
 .../common/stem_overfocus.py                  |  9 ++---
 src/microscope_calibration/ui.py              |  5 +--
 .../util/stem_overfocus_sim.py                |  5 +--
 tests/test_optimize.py                        | 24 ++++++-------
 tests/test_overfocus.py                       | 26 +++++++-------
 tests/test_sim.py                             | 36 +++++++++----------
 tests/test_ui.py                              |  3 +-
 9 files changed, 67 insertions(+), 57 deletions(-)

diff --git a/CLAUDE.md b/CLAUDE.md
index 456fe3d..a965c57 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -60,7 +60,7 @@ third-party dependency you introduce.
 The whole library is organized around one central object, **`Model4DSTEM`**
 (`common/model.py`), a `jax_dataclasses.pytree_dataclass` holding the full geometric
 calibration of a 4D-STEM experiment (overfocus, scan pixel pitch, scan center/rotation,
-camera length, detector pixel pitch/center/rotation, semi-convergence angle, flip factor, and
+camera length, detector pixel pitch/center/rotation, semi-convergence slope, flip factor, and
 a 12-parameter `DescanError`). Because it is a JAX pytree, it is both differentiable and
 traceable.
 
diff --git a/src/microscope_calibration/common/model.py b/src/microscope_calibration/common/model.py
index 365bba3..8226441 100644
--- a/src/microscope_calibration/common/model.py
+++ b/src/microscope_calibration/common/model.py
@@ -157,6 +157,13 @@ class Model4DSTEM:
     JAX and :func:`~microscope_calibration.util.sympy.lambdify` see their
     parameters. They can't add fields that affect :meth:`trace`: the symbolic
     derivation only creates symbols for the fields of :class:`Model4DSTEM`.
+
+    Units
+    -----
+    Lengths are in m, rotations in rad. :code:`semiconv` is the maximum
+    slope of rays at the source relative to the optical axis, i.e.
+    :code:`tan()` of the convergence semi-angle, since :meth:`trace` works with
+    slopes.
     """
 
     overfocus: float  # m
@@ -166,7 +173,8 @@ class Model4DSTEM:
     camera_length: float  # m
     detector_pixel_pitch: float  # m
     detector_center: PixelYX
-    semiconv: float  # rad
+    # Slope at the source, i.e. tan() of the convergence semi-angle
+    semiconv: float
     flip_factor: float  # 1.: no flip; -1.: flip
     descan_error: DescanError = DescanError()
     detector_rotation: float = 0.0  # rad
@@ -217,7 +225,7 @@ def default(cls, dataset_shape: tuple[int, int, int, int] = (0, 0, 0, 0)):
             camera_length=1.0,
             detector_pixel_pitch=50e-6,
             detector_center=detector_center,
-            semiconv=1e-3,  # radian
+            semiconv=1e-3,  # slope
             flip_factor=1.0,
         )
 
@@ -232,7 +240,7 @@ def derive(
         detector_pixel_pitch: float | None = None,  # m
         detector_center: PixelYX | None = None,
         detector_rotation: float | None = None,  # rad
-        semiconv: float | None = None,  # rad
+        semiconv: float | None = None,  # slope
         flip_y: bool | None = None,
         flip_factor: float | None = None,
         descan_error: DescanError | None = None,
diff --git a/src/microscope_calibration/common/stem_overfocus.py b/src/microscope_calibration/common/stem_overfocus.py
index 0eed665..c651086 100644
--- a/src/microscope_calibration/common/stem_overfocus.py
+++ b/src/microscope_calibration/common/stem_overfocus.py
@@ -287,7 +287,7 @@ def _eliminate_tilt(
     ------
     FittingError
         The rows of :code:`plane` for tilt are the change of the position in
-        the plane with source tilt. Scaled by the semi-convergence angle, the
+        the plane with source tilt. Scaled by the semi-convergence slope, the
         smallest singular value of that is the smallest semi-axis of the
         beam's footprint in the plane in pixels. If the footprint is smaller
         than :code:`min_footprint_px`, the transformation is considered (close
@@ -349,7 +349,7 @@ def get_backward_transformation_matrix(
         coefficients. By default, scan pixel coordinates of :code:`rec_model`.
     min_footprint_px
         Minimum smallest semi-axis of the beam's footprint in the specimen
-        plane in specimen pixels, given the semi-convergence angle. Below
+        plane in specimen pixels, given the semi-convergence slope. Below
         that, the transformation is considered (close to) singular, e.g. in
         focus.
 
@@ -386,7 +386,8 @@ def affine_column(scan_y, scan_x, in_y, in_x, mat, column):
 
 @numba.njit(cache=True)
 def project_frame_backwards(frame, source_semiconv, mat, scan_y, scan_x, image_out):
-    limit = np.abs(np.tan(source_semiconv)) ** 2
+    # source_semiconv is a slope, see Model4DSTEM
+    limit = np.abs(source_semiconv) ** 2
     for image_y in range(image_out.shape[0]):
         for image_x in range(image_out.shape[1]):
             tilt_y = affine_column(scan_y, scan_x, image_y, image_x, mat, 2)
@@ -434,7 +435,7 @@ def get_detector_correction_matrix(
         Model of the reference system
     min_footprint_px
         Minimum smallest semi-axis of the beam's footprint in the detector
-        plane of the reference model in pixels, given the semi-convergence
+        plane of the reference model in pixels, given the semi-convergence slope
         angle of :code:`rec_model`. Below that, the transformation is
         considered (close to) singular, e.g. with the detector in a crossover.
 
diff --git a/src/microscope_calibration/ui.py b/src/microscope_calibration/ui.py
index 736455f..e7fb183 100644
--- a/src/microscope_calibration/ui.py
+++ b/src/microscope_calibration/ui.py
@@ -374,7 +374,8 @@ def __init__(
         self.semiconv_input = pn.widgets.FloatInput(
             label="Convergence semi-angle / mrad",
             step=0.01,
-            value=self.start_model.semiconv * 1000,
+            # The model uses the slope, see Model4DSTEM
+            value=float(np.arctan(self.start_model.semiconv)) * 1000,
         )
         self.scalebar_input = pn.widgets.FloatInput(
             label="Scale bar / nm",
@@ -642,7 +643,7 @@ def update_cl(self, event):
         self.update_with_force(self.model_model, self.model_update(model))
 
     def update_semiconv(self, event):
-        model = self.model.derive(semiconv=event.new / 1000)
+        model = self.model.derive(semiconv=float(np.tan(event.new / 1000)))
         self.update_with_force(self.model_model, self.model_update(model))
 
     def update_scale(self, event):
diff --git a/src/microscope_calibration/util/stem_overfocus_sim.py b/src/microscope_calibration/util/stem_overfocus_sim.py
index 71fbec8..1f1a629 100644
--- a/src/microscope_calibration/util/stem_overfocus_sim.py
+++ b/src/microscope_calibration/util/stem_overfocus_sim.py
@@ -96,7 +96,7 @@ def get_forward_transformation_matrix(
         coefficients. By default, scan pixel coordinates of :code:`sim_model`.
     min_footprint_px
         Minimum smallest semi-axis of the beam's footprint in the detector
-        plane in pixels, given the semi-convergence angle. Below that, the
+        plane in pixels, given the semi-convergence slope. Below that, the
         transformation is considered (close to) singular, e.g. with the
         detector in a crossover.
 
@@ -116,7 +116,8 @@ def get_forward_transformation_matrix(
 
 @numba.njit(cache=True)
 def project_frame_forward(obj, source_semiconv, mat, scan_y, scan_x, out):
-    limit = np.abs(np.tan(source_semiconv)) ** 2
+    # source_semiconv is a slope, see Model4DSTEM
+    limit = np.abs(source_semiconv) ** 2
     for det_y in range(out.shape[0]):
         for det_x in range(out.shape[1]):
             tilt_y = affine_column(scan_y, scan_x, det_y, det_x, mat, 2)
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index 802e161..3a2dd7e 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -47,14 +47,14 @@ def test_optimize():
     camera_length = 1.0
     propagation_distance = overfocus + camera_length
     obj_half_size = 16
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.00314157, propagation_distance)
+    semiconv = (obj_half_size * detector_pixel_pitch / 2 + 0.00314157) / propagation_distance
 
     model = Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
         scan_rotation=scan_rotation,
         flip_factor=flip_factor,
@@ -146,14 +146,14 @@ def test_descan_error():
     camera_length = 1.0
     propagation_distance = overfocus + camera_length
     obj_half_size = 16
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.00314157, propagation_distance)
+    semiconv = (obj_half_size * detector_pixel_pitch / 2 + 0.00314157) / propagation_distance
 
     model = Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(x=0.0, y=0.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -281,13 +281,13 @@ def test_scan_pixel_pitch():
     camera_length = 1.234
     propagation_distance = overfocus + camera_length
     obj_half_size = 16
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.00314157, propagation_distance)
+    semiconv = (obj_half_size * detector_pixel_pitch / 2 + 0.00314157) / propagation_distance
     model = Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -331,14 +331,14 @@ def test_full_descan_error(scan_rotation, flip_factor, detector_rotation, descan
     propagation_distance = overfocus + camera_length
     obj_half_size = 8
     # Small epsilon to combat aliasing
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 * 2 + 0.001, propagation_distance)
+    semiconv = (obj_half_size * detector_pixel_pitch / 2 * 2 + 0.001) / propagation_distance
 
     model = Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
         scan_rotation=scan_rotation,
         flip_factor=flip_factor,
@@ -483,14 +483,14 @@ def test_tilt_descan_error(scan_rotation, flip_factor, detector_rotation, descan
     propagation_distance = overfocus + camera_length
     obj_half_size = 8
     # Small epsilon to combat aliasing
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 * 2 + 0.001, propagation_distance)
+    semiconv = (obj_half_size * detector_pixel_pitch / 2 * 2 + 0.001) / propagation_distance
 
     model = Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
         scan_rotation=scan_rotation,
         flip_factor=flip_factor,
@@ -609,14 +609,14 @@ def test_tilt_descan_error_points(
     propagation_distance = overfocus + camera_length
     obj_half_size = 8
     # Small epsilon to combat aliasing
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 * 2 + 0.001, propagation_distance)
+    semiconv = (obj_half_size * detector_pixel_pitch / 2 * 2 + 0.001) / propagation_distance
 
     model = Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
         scan_rotation=scan_rotation,
         flip_factor=flip_factor,
diff --git a/tests/test_overfocus.py b/tests/test_overfocus.py
index 2b5d137..bcb800f 100644
--- a/tests/test_overfocus.py
+++ b/tests/test_overfocus.py
@@ -151,7 +151,7 @@ def test_backproject_identity():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=7.1, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -163,7 +163,7 @@ def test_backproject_identity():
     mat = get_backward_transformation_matrix(rec_model=model)
     project_frame_backwards(
         frame=obj,
-        source_semiconv=np.pi / 2,
+        source_semiconv=np.inf,
         mat=mat,
         scan_y=16,
         scan_x=7,
@@ -181,7 +181,7 @@ def test_backproject_counterrotate():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=np.pi / 2,
         flip_factor=1.0,
@@ -194,7 +194,7 @@ def test_backproject_counterrotate():
     mat = get_backward_transformation_matrix(rec_model=model)
     project_frame_backwards(
         frame=obj,
-        source_semiconv=np.pi / 2,
+        source_semiconv=np.inf,
         mat=mat,
         scan_y=16,
         scan_x=16,
@@ -228,7 +228,7 @@ def test_backproject_rot90_flip(rotate_scan, rotate_detector, fixed_reference, f
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=scan_rotation,
         flip_factor=flip_factor,
@@ -269,7 +269,7 @@ def map_coord(inp):
             res = np.zeros_like(obj)
             project_frame_backwards(
                 frame=projected[pick_y, pick_x],
-                source_semiconv=np.pi / 2,
+                source_semiconv=np.inf,
                 mat=mat,
                 scan_y=pick_y,
                 scan_x=pick_x,
@@ -287,7 +287,7 @@ def test_backproject_scale_fixed():
         scan_pixel_pitch=2,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16.0, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -316,7 +316,7 @@ def map_coord(inp):
     mat = get_backward_transformation_matrix(rec_model=model, specimen_to_image=map_coord)
     project_frame_backwards(
         frame=projected[16, 16],
-        source_semiconv=np.pi / 2,
+        source_semiconv=np.inf,
         mat=mat,
         scan_y=16,
         scan_x=16,
@@ -334,7 +334,7 @@ def test_backproject_scale_scanref():
         scan_pixel_pitch=2,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16.0, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -356,7 +356,7 @@ def test_backproject_scale_scanref():
     mat = get_backward_transformation_matrix(rec_model=model, specimen_to_image=None)
     project_frame_backwards(
         frame=projected[16, 16],
-        source_semiconv=np.pi / 2,
+        source_semiconv=np.inf,
         mat=mat,
         scan_y=16,
         scan_x=16,
@@ -386,15 +386,13 @@ def correction_model(obj_half_size: int, **changes) -> Model4DSTEM:
     """
     overfocus = 1.0
     scan_shape, detector_shape = correction_shapes(obj_half_size)
-    angle = np.arctan2(
-        obj_half_size * DETECTOR_PIXEL_PITCH / 2 + 0.00314157, overfocus + CAMERA_LENGTH
-    )
+    semiconv = (obj_half_size * DETECTOR_PIXEL_PITCH / 2 + 0.00314157) / (overfocus + CAMERA_LENGTH)
     return Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=SCAN_PIXEL_PITCH,
         camera_length=CAMERA_LENGTH,
         detector_pixel_pitch=DETECTOR_PIXEL_PITCH,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(y=scan_shape[0] / 2, x=scan_shape[1] / 2),
         scan_rotation=0.0,
         flip_factor=1.0,
diff --git a/tests/test_sim.py b/tests/test_sim.py
index 4aee26a..c6e7b89 100644
--- a/tests/test_sim.py
+++ b/tests/test_sim.py
@@ -36,7 +36,7 @@ def ref_project(obj, source_semiconv, mat, scan_y, scan_x, out):
                 for det_x in range(out.shape[1]):
                     inp = np.array((scan_y, scan_x, det_y, det_x, 1.0))
                     spec_y, spec_x, tilt_y, tilt_x, _one = inp @ mat
-                    if np.linalg.norm((tilt_y, tilt_x)) < np.tan(source_semiconv):
+                    if np.linalg.norm((tilt_y, tilt_x)) < source_semiconv:
                         spec_y = int(np.round(spec_y))
                         spec_x = int(np.round(spec_x))
                         if (
@@ -201,7 +201,7 @@ def test_project_identity():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=6.9, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -226,7 +226,7 @@ def test_project_scale():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=1,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -250,7 +250,7 @@ def test_project_shift():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=6.9, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -274,7 +274,7 @@ def test_project_rotate():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=np.pi / 2,
         flip_factor=1.0,
@@ -298,7 +298,7 @@ def test_project_flip():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=0.0,
         flip_factor=-1.0,
@@ -322,7 +322,7 @@ def test_project_detector_rotate():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -347,7 +347,7 @@ def test_project_map_identity():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -381,7 +381,7 @@ def test_project_map_scale():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -416,7 +416,7 @@ def test_project_map_rotate():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -450,7 +450,7 @@ def test_project_map_flip():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -483,7 +483,7 @@ def test_project_fixref_scanscale():
         scan_pixel_pitch=2,  # <--
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -518,7 +518,7 @@ def test_project_fixref_scanshift():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=17, y=15.0),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -551,7 +551,7 @@ def test_project_fixref_scanrotate():
         scan_pixel_pitch=1,
         camera_length=1,
         detector_pixel_pitch=2,
-        semiconv=np.pi / 2,
+        semiconv=np.inf,
         scan_center=PixelYX(x=16, y=16.0),
         scan_rotation=np.pi / 2,
         flip_factor=1.0,
@@ -586,14 +586,14 @@ def test_project_aperture():
     propagation_distance = overfocus + camera_length
     obj_half_size = 16
     # Small epsilon to avoid hitting numerical errors at exactly the pixel boundary
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.001, propagation_distance)
+    semiconv = (obj_half_size * detector_pixel_pitch / 2 + 0.001) / propagation_distance
 
     model = Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
         scan_rotation=0.0,
         flip_factor=1.0,
@@ -627,14 +627,14 @@ def test_project_descan():
     propagation_distance = overfocus + camera_length
     obj_half_size = 16
     # Small epsilon to avoid hitting numerical errors at exactly the pixel boundary
-    angle = np.arctan2(obj_half_size * detector_pixel_pitch / 2 + 0.001, propagation_distance)
+    semiconv = (obj_half_size * detector_pixel_pitch / 2 + 0.001) / propagation_distance
 
     model = Model4DSTEM(
         overfocus=overfocus,
         scan_pixel_pitch=scan_pixel_pitch,
         camera_length=camera_length,
         detector_pixel_pitch=detector_pixel_pitch,
-        semiconv=angle,
+        semiconv=semiconv,
         scan_center=PixelYX(x=obj_half_size, y=obj_half_size),
         scan_rotation=0.0,
         flip_factor=1.0,
diff --git a/tests/test_ui.py b/tests/test_ui.py
index 9a6a3d8..49b0cad 100644
--- a/tests/test_ui.py
+++ b/tests/test_ui.py
@@ -54,7 +54,8 @@ def test_events(kwargs):
     assert calibrated.dataset == ds
     m = calibrated.model
     assert_allclose(m.camera_length, 1.2)
-    assert_allclose(m.semiconv, 2 / 1000)
+    # The UI shows the angle, the model holds the slope
+    assert_allclose(m.semiconv, np.tan(2 / 1000))
     assert_allclose(m.detector_pixel_pitch, 60e-6)
     assert_allclose(m.scan_rotation, np.deg2rad(17))
 

From 7b29c0a612fe0a01168363cb96406b34cbe05882 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Wed, 14 Oct 2026 08:05:02 +0000
Subject: [PATCH 29/34] Split CoordinateCorrectionLayout.__init__() into state,
 figures, widgets and events

No change in behavior.

- Add helpers for figures, cursors, diffraction rings and fix point tables
- Remove redundant column checks in the table click handlers

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/ui.py | 277 ++++++++++++++-----------------
 1 file changed, 122 insertions(+), 155 deletions(-)

diff --git a/src/microscope_calibration/ui.py b/src/microscope_calibration/ui.py
index e7fb183..0a7515d 100644
--- a/src/microscope_calibration/ui.py
+++ b/src/microscope_calibration/ui.py
@@ -129,7 +129,6 @@ def __init__(
         sig_mode: SigModeT = "lin",
         twothetas: np.ndarray | None = None,
     ):
-        # # Stuff
         self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
         self.calibrated_dataset = calibrated_dataset
         self.ctx = ctx
@@ -143,7 +142,15 @@ def __init__(
 
         self.coords_adjusted = False
 
-        # # GUI state
+        self._init_state(twothetas)
+        self._init_figures()
+        self._init_widgets()
+        self._connect_events()
+
+    def _init_state(self, twothetas: np.ndarray):
+        """
+        ColumnDataSources and tables that hold the GUI state
+        """
         self.model_model = ColumnDataSource(
             data=(
                 self.model_update(self.start_model)
@@ -195,74 +202,38 @@ def __init__(
             self.coord_index_cols
         )
 
-        # # GUI elements
+    def _init_figures(self):
+        """
+        Figures with cursors and other glyphs, in the order they are added
+        """
         preview = self.get_preview()
-        self.nav_fig = ApertureFigure.new(
+        self.nav_fig = self._new_figure(
             preview["nav"], title=f"Uncorrected scan map (method: {self.nav_mode})", maxdim=350
         )
-        self.adjust_layout(self.nav_fig, shape=preview["nav"].shape)
-
-        self.cursor = (
-            Cursor(cds=self.model_model, x="scan_x", y="scan_y")
-            .on(self.nav_fig.fig)
-            # this adds the necessary components to make the cursor draggable
-            .editable(selected=True)
-        )
-        # The cursor Glyph is accessed via the .cursor property, until the API is unified
-        self.cursor.cursor.line_color = "red"
+        self.cursor = self._add_cursor(self.nav_fig, self.model_model, "scan_x", "scan_y", "red")
 
         frames = self.get_frames(pos=self.scan_pos)
-        self.pick_fig = ApertureFigure.new(
+        self.pick_fig = self._new_figure(
             frames["raw"],
             title=f"Currently picked uncorrected detector frame ({self.sig_mode})",
             maxdim=350,
         )
-        self.adjust_layout(self.pick_fig, shape=frames["raw"].shape)
-
-        self.beam_centre = (
-            Cursor(cds=self.ring_model, x="detector_cx", y="detector_cy")
-            .on(self.pick_fig.fig)
-            .editable(selected=True)
-        )
-        self.beam_centre.glyph.line_color = "red"
-        self.diffraction_ringsets = []
-        for i in range(self.ring_model.data["num_rings"][0]):
-            rikey = f"ri_{i}"
-            rokey = f"ro_{i}"
-            self.diffraction_ringsets.append(
-                RingSet(
-                    cds=self.ring_model,
-                    x="detector_cx",
-                    y="detector_cy",
-                    inner_radius=rikey,
-                    outer_radius=rokey,
-                ).on(self.pick_fig.fig)
-            )
+        self.beam_centre = self._add_cursor(
+            self.pick_fig, self.ring_model, "detector_cx", "detector_cy", "red"
+        )
+        self.diffraction_ringsets = self._add_rings(self.pick_fig, self.ring_model)
 
-        self.corr_point_fig = ApertureFigure.new(
+        self.corr_point_fig = self._new_figure(
             preview["corrected_point"],
             title="Point analysis corrected for de-scan error",
             maxdim=250,
         )
-        self.adjust_layout(self.corr_point_fig, shape=preview["corrected_point"].shape)
-        self.cursor_2 = (
-            Cursor(cds=self.model_model, x="scan_x", y="scan_y")
-            .on(self.corr_point_fig.fig)
-            # this adds the necessary components to make the cursor draggable
-            .editable(selected=True)
+        self.cursor_2 = self._add_cursor(
+            self.corr_point_fig, self.model_model, "scan_x", "scan_y", "red"
         )
-        # The cursor Glyph is accessed via the .cursor property, until the API is unified
-        self.cursor_2.cursor.line_color = "red"
-
-        self.nav_feature_cursor = (
-            Cursor(cds=self.feature_model, x="scan_x", y="scan_y")
-            .on(self.corr_point_fig.fig)
-            # this adds the necessary components to make the cursor draggable
-            .editable(selected=True)
+        self.nav_feature_cursor = self._add_cursor(
+            self.corr_point_fig, self.feature_model, "scan_x", "scan_y", "lightgreen"
         )
-        # The cursor Glyph is accessed via the .cursor property, until the API is unified
-        self.nav_feature_cursor.cursor.line_color = "lightgreen"
-
         self.scalebar_handles = (
             PointSet(cds=self.scalebar_model, x="scalebar_x", y="scalebar_y")
             .on(self.corr_point_fig.fig)
@@ -275,99 +246,111 @@ def __init__(
         self.scalebar_handles.glyph.fill_color = None
         self.scalebar_line.glyph.line_color = "yellow"
 
-        self.corr_pick_fig = ApertureFigure.new(
+        self.corr_pick_fig = self._new_figure(
             frames["corrected"],
             title=f"Frame corrected for de-scan error  ({self.sig_mode})",
             maxdim=250,
         )
-        self.adjust_layout(self.corr_pick_fig, shape=frames["corrected"].shape)
-
-        self.corr_beam_centre = Cursor(
-            cds=self.centered_ring_model, x="detector_cx", y="detector_cy"
-        ).on(self.corr_pick_fig.fig)
-        self.corr_beam_centre.glyph.line_color = "red"
-        self.corr_diffraction_ringsets = []
-        for i in range(self.centered_ring_model.data["num_rings"][0]):
-            rikey = f"ri_{i}"
-            rokey = f"ro_{i}"
-            self.corr_diffraction_ringsets.append(
-                RingSet(
-                    cds=self.centered_ring_model,
-                    x="detector_cx",
-                    y="detector_cy",
-                    inner_radius=rikey,
-                    outer_radius=rokey,
-                ).on(self.corr_pick_fig.fig)
-            )
+        self.corr_beam_centre = self._add_cursor(
+            self.corr_pick_fig,
+            self.centered_ring_model,
+            "detector_cx",
+            "detector_cy",
+            "red",
+            editable=False,
+        )
+        self.corr_diffraction_ringsets = self._add_rings(
+            self.corr_pick_fig, self.centered_ring_model
+        )
 
-        self.corr_sum_fig = ApertureFigure.new(
+        self.corr_sum_fig = self._new_figure(
             preview["corrected_sum"],
             title=f"Sum of frames corrected for de-scan error ({self.sig_mode})",
             maxdim=250,
         )
-        self.adjust_layout(self.corr_sum_fig, shape=preview["corrected_sum"].shape)
-        self.corr_beam_centre_2 = Cursor(
-            cds=self.centered_ring_model, x="detector_cx", y="detector_cy"
-        ).on(self.corr_sum_fig.fig)
-        self.corr_beam_centre_2.glyph.line_color = "red"
-        self.corr_diffraction_ringsets_2 = []
-        for i in range(self.centered_ring_model.data["num_rings"][0]):
-            rikey = f"ri_{i}"
-            rokey = f"ro_{i}"
-            self.corr_diffraction_ringsets_2.append(
-                RingSet(
-                    cds=self.centered_ring_model,
-                    x="detector_cx",
-                    y="detector_cy",
-                    inner_radius=rikey,
-                    outer_radius=rokey,
-                ).on(self.corr_sum_fig.fig)
-            )
+        self.corr_beam_centre_2 = self._add_cursor(
+            self.corr_sum_fig,
+            self.centered_ring_model,
+            "detector_cx",
+            "detector_cy",
+            "red",
+            editable=False,
+        )
+        self.corr_diffraction_ringsets_2 = self._add_rings(
+            self.corr_sum_fig, self.centered_ring_model
+        )
 
-        self.nav_fig_2 = ApertureFigure.new(
+        self.nav_fig_2 = self._new_figure(
             preview["nav"], title=f"Uncorrected scan map (method: {self.nav_mode})"
         )
-        self.adjust_layout(self.nav_fig_2, shape=preview["nav"].shape)
-
-        self.cursor_3 = (
-            Cursor(cds=self.model_model, x="scan_x", y="scan_y")
-            .on(self.nav_fig_2.fig)
-            # this adds the necessary components to make the cursor draggable
-            .editable(selected=True)
+        self.cursor_3 = self._add_cursor(
+            self.nav_fig_2, self.model_model, "scan_x", "scan_y", "red"
         )
-        # The cursor Glyph is accessed via the .cursor property, until the API is unified
-        self.cursor_3.cursor.line_color = "red"
 
-        self.pick_fig_2 = ApertureFigure.new(
+        self.pick_fig_2 = self._new_figure(
             frames["raw"],
             title=f"Currently picked uncorrected detector frame ({self.sig_mode})",
             maxdim=350,
         )
-        self.adjust_layout(self.pick_fig_2, shape=frames["raw"].shape)
-
-        self.nav_feature_select_cursor = (
-            Cursor(cds=self.feature_select_model, x="detector_x", y="detector_y")
-            .on(self.pick_fig_2.fig)
-            # this adds the necessary components to make the cursor draggable
-            .editable(selected=True)
+        self.nav_feature_select_cursor = self._add_cursor(
+            self.pick_fig_2, self.feature_select_model, "detector_x", "detector_y", "lightgreen"
         )
-        # The cursor Glyph is accessed via the .cursor property, until the API is unified
-        self.nav_feature_select_cursor.cursor.line_color = "lightgreen"
 
-        self.back_sum_fig = ApertureFigure.new(
+        self.back_sum_fig = self._new_figure(
             preview["backprojected_sum"],
             title="Frames superimposed in scan coordinate system",
             maxdim=350,
         )
-        self.adjust_layout(self.back_sum_fig, shape=preview["backprojected_sum"].shape)
-
-        self.back_pick_fig = ApertureFigure.new(
+        self.back_pick_fig = self._new_figure(
             frames["backprojected"],
             title="Currently selected frame back-projected to scan coordinate system",
             maxdim=350,
         )
-        self.adjust_layout(self.back_pick_fig, shape=frames["backprojected"].shape)
 
+    @classmethod
+    def _new_figure(cls, data: np.ndarray, title: str, maxdim: int | None = None) -> ApertureFigure:
+        kwargs = {} if maxdim is None else {"maxdim": maxdim}
+        fig = ApertureFigure.new(data, title=title, **kwargs)
+        cls.adjust_layout(fig, shape=data.shape)
+        return fig
+
+    @staticmethod
+    def _add_cursor(
+        fig: ApertureFigure,
+        cds: ColumnDataSource,
+        x: str,
+        y: str,
+        color: str,
+        editable: bool = True,
+    ) -> Cursor:
+        cursor = Cursor(cds=cds, x=x, y=y).on(fig.fig)
+        if editable:
+            # this adds the necessary components to make the cursor draggable
+            cursor = cursor.editable(selected=True)
+        cursor.glyph.line_color = color
+        return cursor
+
+    @staticmethod
+    def _add_rings(fig: ApertureFigure, cds: ColumnDataSource) -> list[RingSet]:
+        """
+        Diffraction rings with the radii :code:`ri_{i}` and :code:`ro_{i}` of
+        :meth:`ring_update`
+        """
+        return [
+            RingSet(
+                cds=cds,
+                x="detector_cx",
+                y="detector_cy",
+                inner_radius=f"ri_{i}",
+                outer_radius=f"ro_{i}",
+            ).on(fig.fig)
+            for i in range(cds.data["num_rings"][0])
+        ]
+
+    def _init_widgets(self):
+        """
+        Inputs, buttons and tables
+        """
         self.cl_input = pn.widgets.FloatInput(
             label="Camera length / m", step=0.01, value=self.start_model.camera_length
         )
@@ -402,58 +385,50 @@ def __init__(
             value=self.start_model.detector_pixel_pitch * 1e6,
         )
 
-        self.descan_fixpoint_table = pn.widgets.Tabulator(
-            self.descan_fixpoints,
-            buttons={
-                "select": '<i class="fa fa-bullseye"></i>',
-                "delete": '<i class="fa fa-trash"></i>',
-            },
-            # See https://github.com/holoviz/panel/pull/8256
-            selectable=False,
-        )
+        self.descan_fixpoint_table = self._new_fixpoint_table(self.descan_fixpoints)
 
         self.record_button = pn.widgets.Button(label="Record")
         self.apply_button = pn.widgets.Button(label="Apply correction from table")
         self.clear_button = pn.widgets.Button(label="Clear")
         self.correlate_button = pn.widgets.Button(label="Refine correction with cross-correlation")
 
-        self.coord_fixpoint_table = pn.widgets.Tabulator(
-            self.coord_fixpoints,
-            buttons={
-                "select": '<i class="fa fa-bullseye"></i>',
-                "delete": '<i class="fa fa-trash"></i>',
-            },
-            # See https://github.com/holoviz/panel/pull/8256
-            selectable=False,
-        )
+        self.coord_fixpoint_table = self._new_fixpoint_table(self.coord_fixpoints)
 
         self.coord_record_button = pn.widgets.Button(label="Record")
         self.coord_apply_button = pn.widgets.Button(label="Derive coordinate system from table")
         self.coord_clear_button = pn.widgets.Button(label="Clear")
         self.optimize_button = pn.widgets.Button(label="Optimize sharpness of back-projection")
-        model_dataframe = pd.DataFrame(columns=self.model_columns)
-        descan_error_dataframe = pd.DataFrame(columns=self.descan_error_columns)
 
         self.invert_focus_button = pn.widgets.Button(label="Invert sign of overfocus")
 
         self.model_table = pn.widgets.Tabulator(
-            model_dataframe,
+            pd.DataFrame(columns=self.model_columns),
             # See https://github.com/holoviz/panel/pull/8256
             selectable=False,
         )
         self.descan_error_table = pn.widgets.Tabulator(
-            descan_error_dataframe,
+            pd.DataFrame(columns=self.descan_error_columns),
             # See https://github.com/holoviz/panel/pull/8256
             selectable=False,
         )
         self.update_model_tables(self.start_model)
 
-        # # Event handler setup
+    @staticmethod
+    def _new_fixpoint_table(df: pd.DataFrame) -> pn.widgets.Tabulator:
+        return pn.widgets.Tabulator(
+            df,
+            buttons={
+                "select": '<i class="fa fa-bullseye"></i>',
+                "delete": '<i class="fa fa-trash"></i>',
+            },
+            # See https://github.com/holoviz/panel/pull/8256
+            selectable=False,
+        )
 
+    def _connect_events(self):
         # ## state
         # trigger whenever the "data" attribute of the cursor ColumnDataSource changes
         self.model_model.on_change("data", self.on_model_model_change)
-        # self.model_model.on_change("data", lambda attr, old, new: print(attr, old, new))
         self.scalebar_model.on_change("data", self.on_scalebar_model_change)
         self.feature_model.on_change("data", self.on_feature_model_change)
 
@@ -480,12 +455,10 @@ def __init__(
 
         # ## descan fixpoints table
         self.descan_fixpoint_table.on_click(
-            lambda e: self.descan_delete_row(e.row) if e.column == "delete" else None,
-            column="delete",
+            lambda e: self.descan_delete_row(e.row), column="delete"
         )
         self.descan_fixpoint_table.on_click(
-            lambda e: self.descan_move_to_row(e.row) if e.column == "select" else None,
-            column="select",
+            lambda e: self.descan_move_to_row(e.row), column="select"
         )
 
         # ## descan correction buttons
@@ -495,16 +468,10 @@ def __init__(
         self.correlate_button.on_click(lambda e: self.center_correlation_regression())
 
         # ## coord fixpoints table
-        self.coord_fixpoint_table.on_click(
-            lambda e: self.coord_delete_row(e.row) if e.column == "delete" else None,
-            column="delete",
-        )
-        self.coord_fixpoint_table.on_click(
-            lambda e: self.coord_move_to_row(e.row) if e.column == "select" else None,
-            column="select",
-        )
+        self.coord_fixpoint_table.on_click(lambda e: self.coord_delete_row(e.row), column="delete")
+        self.coord_fixpoint_table.on_click(lambda e: self.coord_move_to_row(e.row), column="select")
 
-        # ## descan correction buttons
+        # ## coordinate system buttons
         self.coord_record_button.on_click(lambda e: self.coord_add_row())
         self.coord_apply_button.on_click(lambda e: self.perform_coord_update())
         self.coord_clear_button.on_click(lambda e: self.coord_drop())

From 1b31bd12e94596aef161969b5d1236e3a22643a8 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Wed, 14 Oct 2026 08:06:02 +0000
Subject: [PATCH 30/34] Share the row handling of the two fix point tables in
 the calibration UI

No change in behavior.

- Replace commented-out code with a note on why it doesn't work

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/ui.py | 108 ++++++++++++++-----------------
 1 file changed, 49 insertions(+), 59 deletions(-)

diff --git a/src/microscope_calibration/ui.py b/src/microscope_calibration/ui.py
index 0a7515d..c9f8c4d 100644
--- a/src/microscope_calibration/ui.py
+++ b/src/microscope_calibration/ui.py
@@ -523,28 +523,10 @@ def on_model_model_change(self, attr, old, new):
         # new will be the CDS dict after the trigger
         assert attr == "data"
         new_model = self.get_model(new)
-
-        # Doesn't work, likely because not all changes trigger this callback,
-        # see https://github.com/holoviz/panel/issues/8275
-        # If the coordinate system calibration is current
-        # if self.coords_adjusted:
-        #     # Check if parameters have changed that invalidate the
-        #     # calibration, and re-calibrate in that case
-        #     old_model = self.get_model(old)
-        #     new_model = self.get_model(new)
-        #     # These should be parameters that are NOT changed by
-        #     # self.perform_coord_update()
-        #     relevant = ('scan_pixel_pitch', 'camera_length', 'scan_rotation')
-        #     if any(getattr(old_model, r) != getattr(new_model, r) for r in relevant):
-        #         print("changed")
-        #         self.perform_coord_update()
-        #     else:
-        #         print("unchanged")
-
+        # Re-calibrating the coordinate system or updating only what changed
+        # doesn't work, likely because not all changes trigger this callback,
+        # see https://github.com/holoviz/panel/issues/8275. Update everything.
         new_pos = self.get_scan_pos(new)
-        # Somehow doesn't work with all the event handling confusion
-        # Just update always and be done with it *sigh*
-        # if old_pos != new_pos:
         frames = self.get_frames(pos=new_pos)
         self.pick_fig.update(frames["raw"])
         self.pick_fig_2.update(frames["raw"])
@@ -662,24 +644,41 @@ def move_rings_to(self, event):
             ),
         )
 
+    @staticmethod
+    def _upsert_row(table: pn.widgets.Tabulator, columns: list, index_cols: list, values: tuple):
+        """
+        Add a row to a fix point table, or update it if the index already exists
+        """
+        idx = values[: len(index_cols)]
+        new_df = pd.DataFrame([values], columns=columns).set_index(index_cols)
+        if idx in table.value.index:
+            table.value.update(new_df)
+            # Make sure the change is registered
+            table.value = table.value
+        else:
+            table.value = pd.concat((table.value, new_df))
+
+    @staticmethod
+    def _delete_row(table: pn.widgets.Tabulator, row: int):
+        df = table.value
+        table.value = df.drop(index=df.index[row])
+
+    @staticmethod
+    def _drop_rows(table: pn.widgets.Tabulator):
+        table.value = table.value.drop(table.value.index)
+
     def descan_add_row(self):
         scan_pos = self.scan_pos
         center_pos = self.get_detector_center(self.ring_model.data)
-        idx = (int(np.round(scan_pos.y)), int(np.round(scan_pos.x)))
-        new_df = pd.DataFrame([idx + (center_pos.y, center_pos.x)], columns=self.descan_columns)
-        new_df = new_df.set_index(self.descan_index_cols)
-        t = self.descan_fixpoint_table
-        if idx in t.value.index:
-            t.value.update(new_df)
-            # Make sure the change is registered
-            t.value = t.value
-        else:
-            t.value = pd.concat((t.value, new_df))
+        self._upsert_row(
+            self.descan_fixpoint_table,
+            self.descan_columns,
+            self.descan_index_cols,
+            (int(np.round(scan_pos.y)), int(np.round(scan_pos.x)), center_pos.y, center_pos.x),
+        )
 
     def descan_delete_row(self, row):
-        df = self.descan_fixpoint_table.value
-        key = df.index[row]
-        self.descan_fixpoint_table.value = df.drop(index=key)
+        self._delete_row(self.descan_fixpoint_table, row)
 
     def descan_move_to_row(self, row):
         df = self.descan_fixpoint_table.value
@@ -689,34 +688,29 @@ def descan_move_to_row(self, row):
         self.update_with_force(self.model_model, self.scan_pos_update(scan_pos))
 
     def descan_drop(self):
-        t = self.descan_fixpoint_table
-        t.value = t.value.drop(t.value.index)
+        self._drop_rows(self.descan_fixpoint_table)
 
     def coord_add_row(self):
         scan_pos = self.scan_pos
         feature_nav_pos = self.get_scan_pos(self.feature_model.data)
         feature_pos = self.get_feature(self.feature_select_model.data)
-        idx = (
-            int(np.round(scan_pos.y)),
-            int(np.round(scan_pos.x)),
-            float(feature_nav_pos.y),
-            float(feature_nav_pos.x),
-        )
-        new_df = pd.DataFrame([idx + (feature_pos.y, feature_pos.x)], columns=self.coord_columns)
-        new_df = new_df.set_index(self.coord_index_cols)
-        t = self.coord_fixpoint_table
-        if idx in t.value.index:
-            t.value.update(new_df)
-            # Make sure the change is registered
-            t.value = t.value
-        else:
-            t.value = pd.concat((t.value, new_df))
+        self._upsert_row(
+            self.coord_fixpoint_table,
+            self.coord_columns,
+            self.coord_index_cols,
+            (
+                int(np.round(scan_pos.y)),
+                int(np.round(scan_pos.x)),
+                float(feature_nav_pos.y),
+                float(feature_nav_pos.x),
+                feature_pos.y,
+                feature_pos.x,
+            ),
+        )
         self.coords_adjusted = False
 
     def coord_delete_row(self, row):
-        df = self.coord_fixpoint_table.value
-        key = df.index[row]
-        self.coord_fixpoint_table.value = df.drop(index=key)
+        self._delete_row(self.coord_fixpoint_table, row)
         self.coords_adjusted = False
 
     def coord_move_to_row(self, row):
@@ -732,8 +726,7 @@ def coord_move_to_row(self, row):
         self.update_with_force(self.feature_select_model, self.feature_update(detector_pos))
 
     def coord_drop(self):
-        t = self.coord_fixpoint_table
-        t.value = t.value.drop(t.value.index)
+        self._drop_rows(self.coord_fixpoint_table)
         self.coords_adjusted = False
 
     def get_model(self, model_data) -> Model4DSTEM:
@@ -745,7 +738,6 @@ def get_scan_pos(model_data):
 
     @staticmethod
     def get_model_twothetas(model_data):
-        # return np.array((0.01, 0.02, 0.03.))
         return model_data["twothetas"][0]
 
     @staticmethod
@@ -833,7 +825,6 @@ def get_mask():
             result["nav"] = res[0]["intensity"].data[..., 0]
             result["corrected_point"] = res[1]["corrected_point"].data
         result["backprojected_sum"] = res[1]["backprojected_sum"].data
-        # result["corrected_point"] = res[1]["corrected_point"].data
         if self.sig_mode == "lin":
             result["corrected_sum"] = res[1]["corrected_sum"].data
         elif self.sig_mode == "log":
@@ -981,7 +972,6 @@ def update_with_force(self, cds: ColumnDataSource, update):
             cds.document.callbacks.hold(holding)
 
     def center_correlation_regression(self):
-        # Delicious! 🍝🍝🍝
         model = self.model
         ri, ro = ring_radii(model, self.get_model_twothetas(self.model_model.data))
         radius_outer = 1.2 * ro[0]

From 99efb9aff2a034673852943219fab18223cd7b84 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Wed, 14 Oct 2026 08:07:59 +0000
Subject: [PATCH 31/34] Move the calculations of the calibration UI into tested
 functions

No change in behavior.

- Correlation pattern, expected beam shifts and regression of the descan refinement
- Formatting of the model and descan error tables

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/ui.py | 166 ++++++++++++++++++++-----------
 tests/test_ui_functions.py       |  70 +++++++++++++
 2 files changed, 177 insertions(+), 59 deletions(-)
 create mode 100644 tests/test_ui_functions.py

diff --git a/src/microscope_calibration/ui.py b/src/microscope_calibration/ui.py
index c9f8c4d..622db73 100644
--- a/src/microscope_calibration/ui.py
+++ b/src/microscope_calibration/ui.py
@@ -102,6 +102,105 @@ def tags(self):
         return {}
 
 
+def format_model(model: Model4DSTEM, columns: list[str]) -> dict[str, str]:
+    """
+    Model parameters as text for the model table, rotations in degrees
+    """
+    res = {}
+    for key in columns:
+        value = getattr(model, key)
+        if isinstance(value, float):
+            if "rotation" in key:
+                value = np.rad2deg(value)
+            value = f"{value:.5}"
+            if "rotation" in key:
+                value = value + "°"
+        else:
+            value = str(value)
+        res[key] = value
+    return res
+
+
+def format_descan_error(descan_error: DescanError) -> dict[str, str]:
+    """
+    Descan error components as text for the descan error table
+    """
+    return {
+        key: f"{value:.5}" if isinstance(value, float) else str(value)
+        for key, value in descan_error._asdict().items()
+    }
+
+
+def correlation_pattern(ri: np.ndarray, ro: np.ndarray) -> BackgroundSubtraction:
+    """
+    Pattern to find the primary beam with the ring radii of :func:`ring_radii`
+    """
+    radius_outer = 1.2 * ro[0]
+    if len(ri) >= 2:
+        radius_outer = max(radius_outer, ri[1] / 2)
+    # Make sure we stay away from other peaks
+    return BackgroundSubtraction(radius=ro[0], radius_outer=radius_outer)
+
+
+def expected_center_shifts(model: Model4DSTEM, nav_shape: tuple[int, int]) -> np.ndarray:
+    """
+    Shift of the primary beam from the detector center at each scan position
+    according to the descan error of :code:`model`, in detector pixels (y, x)
+    """
+    # Correction to same parameters, except descan error of 0
+    ref_model = model.derive(descan_error=DescanError())
+    mat = get_detector_correction_matrix(rec_model=model, ref_model=ref_model)
+    # Scan positions in the dataset
+    y, x = np.mgrid[: nav_shape[0], : nav_shape[1]]
+    # Calculate where the nominal center actually is on the detector
+    expected_x = corrected_det_x(
+        det_corr_x=model.detector_center.x,
+        det_corr_y=model.detector_center.y,
+        scan_x=x,
+        scan_y=y,
+        mat=mat,
+    )
+    expected_y = corrected_det_y(
+        det_corr_x=model.detector_center.x,
+        det_corr_y=model.detector_center.y,
+        scan_x=x,
+        scan_y=y,
+        mat=mat,
+    )
+    # Subtract the expected position so that only the deviation remains
+    dc = model.detector_center
+    return np.stack((expected_y - dc.y, expected_x - dc.x), axis=-1)
+
+
+def center_regression(field: np.ndarray, center: PixelYX) -> np.ndarray:
+    """
+    Linear regression of beam positions, following the CoMUDF regression
+
+    Parameters
+    ----------
+    field
+        Beam position (y, x) on the detector for each scan position (y, x)
+    center
+        Nominal detector center
+
+    Returns
+    -------
+    regression
+        Rows offset, d/dy, d/dx of the deviation from :code:`center`, with
+        columns (y, x), see :data:`~microscope_calibration.util.optimize.CoMRegression`
+    """
+    # Only keep deviation from expected value in regression
+    field = field.copy()
+    field[..., 0] -= center.y
+    field[..., 1] -= center.x
+    inp = np.ones(field.shape[:-1] + (3,))
+    y, x = np.ogrid[: field.shape[0], : field.shape[1]]
+    inp[..., 1] = y
+    inp[..., 2] = x
+    reg_res = np.linalg.lstsq(inp.reshape((-1, 3)), field.reshape((-1, 2)), rcond=None)
+    return reg_res[0]
+
+
 class CoordinateCorrectionLayout:
     descan_columns = ["scan_y", "scan_x", "detector_cy", "detector_cx"]
     descan_index_cols = descan_columns[:2]
@@ -492,26 +591,12 @@ def current_calibrated(self) -> AbstractCalibratedDataset:
 
     def update_model_tables(self, model):
         model_df = self.model_table.value
-        for key in self.model_columns:
-            value = getattr(model, key)
-            if isinstance(value, float):
-                if "rotation" in key:
-                    value = np.rad2deg(value)
-                value = f"{value:.5}"
-                if "rotation" in key:
-                    value = value + "°"
-            else:
-                value = str(value)
+        for key, value in format_model(model, self.model_columns).items():
             model_df.loc[0, key] = value
         self.model_table.value = model_df
 
         descan_df = self.descan_error_table.value
-        for key in self.descan_error_columns:
-            value = getattr(model.descan_error, key)
-            if isinstance(value, float):
-                value = f"{value:.5}"
-            else:
-                value = str(value)
+        for key, value in format_descan_error(model.descan_error).items():
             descan_df.loc[0, key] = value
         self.descan_error_table.value = descan_df
 
@@ -974,35 +1059,7 @@ def update_with_force(self, cds: ColumnDataSource, update):
     def center_correlation_regression(self):
         model = self.model
         ri, ro = ring_radii(model, self.get_model_twothetas(self.model_model.data))
-        radius_outer = 1.2 * ro[0]
-        if len(ri) >= 2:
-            radius_outer = max(radius_outer, ri[1] / 2)
-        # Make sure we stay away from other peaks
-        pattern = BackgroundSubtraction(radius=ro[0], radius_outer=radius_outer)
-        # Correction to same parameters, except descan error of 0
-        ref_model = model.derive(descan_error=DescanError())
-        mat = get_detector_correction_matrix(rec_model=model, ref_model=ref_model)
-        # Scan positions in the dataset
-        nav_shape = self.dataset.shape.nav
-        y, x = np.mgrid[: nav_shape[0], : nav_shape[1]]
-        # Calculate where the nominal center actually is on the detector
-        expected_x = corrected_det_x(
-            det_corr_x=model.detector_center.x,
-            det_corr_y=model.detector_center.y,
-            scan_x=x,
-            scan_y=y,
-            mat=mat,
-        )
-        expected_y = corrected_det_y(
-            det_corr_x=model.detector_center.x,
-            det_corr_y=model.detector_center.y,
-            scan_x=x,
-            scan_y=y,
-            mat=mat,
-        )
-        # Subtract the expected position so that only the deviation remains
-        dc = model.detector_center
-        shifts = np.stack((expected_y - dc.y, expected_x - dc.x), axis=-1)
+        shifts = expected_center_shifts(model, nav_shape=tuple(self.dataset.shape.nav))
         aux_shifts = FastCorrelationUDF.aux_data(
             data=shifts,
             kind="nav",
@@ -1013,23 +1070,14 @@ def center_correlation_regression(self):
             ctx=self.ctx,
             dataset=self.dataset,
             peaks=np.array([model.detector_center]),
-            match_pattern=pattern,
+            match_pattern=correlation_pattern(ri, ro),
             zero_shift=aux_shifts,
             upsample=True,
         )
-
-        # Following CoMUDF regression
-        field = res["refineds"].data[:, :, 0, :]
-        # Only keep deviation from expected value in regression
-        field[..., 0] -= model.detector_center.y
-        field[..., 1] -= model.detector_center.x
-
-        inp = np.ones(field.shape[:-1] + (3,))
-        y, x = np.ogrid[: field.shape[0], : field.shape[1]]
-        inp[..., 1] = y
-        inp[..., 2] = x
-        reg_res = np.linalg.lstsq(inp.reshape((-1, 3)), field.reshape((-1, 2)), rcond=None)
-        new_model, residual = solve_tilt_descan_error(ref_model=model, regression=reg_res[0])
+        regression = center_regression(
+            res["refineds"].data[:, :, 0, :], center=model.detector_center
+        )
+        new_model, residual = solve_tilt_descan_error(ref_model=model, regression=regression)
         self.model_model.data.update(**self.model_update(new_model))
         self.push()
 
diff --git a/tests/test_ui_functions.py b/tests/test_ui_functions.py
new file mode 100644
index 0000000..1db7d64
--- /dev/null
+++ b/tests/test_ui_functions.py
@@ -0,0 +1,70 @@
+import numpy as np
+from numpy.testing import assert_allclose
+
+from microscope_calibration.common.model import DescanError, Model4DSTEM, PixelYX
+from microscope_calibration.common.stem_overfocus import get_center
+from microscope_calibration.ui import (
+    center_regression,
+    correlation_pattern,
+    expected_center_shifts,
+    format_descan_error,
+    format_model,
+)
+
+
+def test_center_regression():
+    center = PixelYX(y=10.0, x=20.0)
+    y, x = np.mgrid[:5, :7]
+    field = np.stack(
+        (center.y + 1.0 + 0.2 * y - 0.3 * x, center.x - 2.0 + 0.4 * y + 0.5 * x), axis=-1
+    ).astype(np.float32)
+    regression = center_regression(field, center)
+    assert_allclose(regression, [[1.0, -2.0], [0.2, 0.4], [-0.3, 0.5]], atol=1e-5)
+    # The input is not modified
+    assert field[0, 0, 0] == np.float32(center.y + 1.0)
+
+
+def test_expected_center_shifts():
+    model = Model4DSTEM.default(dataset_shape=(4, 5, 32, 32)).derive(
+        overfocus=1e-3,
+        descan_error=DescanError(offpxi=1e-5, pxo_pxi=0.3, pyo_pyi=-0.2),
+    )
+    shifts = expected_center_shifts(model, nav_shape=(4, 5))
+    assert shifts.shape == (4, 5, 2)
+    for scan_y, scan_x in ((0, 0), (3, 1), (2, 4)):
+        center = get_center(model, PixelYX(y=float(scan_y), x=float(scan_x)))
+        assert_allclose(
+            shifts[scan_y, scan_x],
+            (center.y - model.detector_center.y, center.x - model.detector_center.x),
+            atol=1e-9,
+        )
+    no_descan = model.derive(descan_error=DescanError())
+    center = get_center(no_descan, PixelYX(y=0.0, x=0.0))
+    assert_allclose(
+        expected_center_shifts(no_descan, nav_shape=(4, 5)),
+        np.broadcast_to(
+            (center.y - model.detector_center.y, center.x - model.detector_center.x), (4, 5, 2)
+        ),
+        atol=1e-9,
+    )
+
+
+def test_correlation_pattern():
+    pattern = correlation_pattern(ri=np.array([0.0, 20.0]), ro=np.array([5.0, 30.0]))
+    assert pattern.radius == 5.0
+    assert pattern.radius_outer == 10.0
+    pattern = correlation_pattern(ri=np.array([0.0]), ro=np.array([5.0]))
+    assert pattern.radius_outer == 6.0
+
+
+def test_format_model():
+    model = Model4DSTEM.default(dataset_shape=(8, 8, 8, 8)).derive(
+        scan_rotation=np.pi / 2, descan_error=DescanError(offpxi=1.23456789)
+    )
+    res = format_model(model, ["scan_rotation", "camera_length", "scan_center"])
+    assert res == {
+        "scan_rotation": "90.0°",
+        "camera_length": "1.0",
+        "scan_center": str(model.scan_center),
+    }
+    assert format_descan_error(model.descan_error)["offpxi"] == "1.2346"

From daaec93e1db60fe0557da10e20855e1c89582b37 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Wed, 14 Oct 2026 08:13:26 +0000
Subject: [PATCH 32/34] Move the instruction texts of the calibration UI layout
 to module constants

No change in the displayed text, since Panel dedents Markdown.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/ui.py | 238 ++++++++++++++++---------------
 1 file changed, 122 insertions(+), 116 deletions(-)

diff --git a/src/microscope_calibration/ui.py b/src/microscope_calibration/ui.py
index 622db73..f115bf5 100644
--- a/src/microscope_calibration/ui.py
+++ b/src/microscope_calibration/ui.py
@@ -102,6 +102,122 @@ def tags(self):
         return {}
 
 
+# Instructions in the layout, see CoordinateCorrectionLayout.layout
+SECTION_1_TEXT = """
+# 4D STEM coordinate system calibration
+
+**Dataset**: {dataset}
+
+{tags}
+
+This tool helps to check and adjust the geometry of a 4D STEM
+experiment.
+
+## De-scan error, convergence semi-angle and effective camera length
+
+<details>
+<summary>Instructions</summary>
+
+Fisrt, enter the nominal values for the scan rotation as well as the
+detector pixel pitch in the corresponding input fields.
+
+Second, confirm or calibrate the shift of the beam center as a
+function of scan position (de-scan error).
+
+1. Change the selected scan position on the left
+1. Observe if the cursor on the right stays at the beam center.
+1. Observe if the plot "Corrected point analysis" shows a bright
+   field image of the specimen.
+1. Observe if the plots "Frame corrected for descan error" and "Sum
+   of frames corrected for de-scan error" show the primary beam
+   exactly at the cursor position.
+1. If the positions don't match or the plot is distorted, move the
+   cursor to the beam center and record the position. Do this for at
+   least three different scan positions.
+1. Press the button "Apply correction from table" to re-calibrate
+   the de-scan error.
+1. Observe if the de-scan error is calibrated correctly now by
+   repeating the first three steps.
+
+In case the frames show diffraction rings or spots and a list of
+"twotheta" angles was provided, now calibrate the camera length so
+that the ring glyphs match the peaks or rings in the detector
+frames.
+
+Finally, adjust the convergence semi-angle so that the innermost
+glyph matches the size of the diffraction disk.
+
+</details>
+"""
+
+SCAN_STEP_TEXT = """
+## Scan step calibration
+
+<details>
+
+<summary>Instructions</summary>
+
+Move the yellow line handles in the plot "Corrected point analysis"
+so that the line spans a feature of known physical size. Check the
+size in the "Scale bar" input field and adjust if necessary. Note
+that for defocused data the scale is only correct if the de-scan
+error was compensated correctly.
+
+The scale bar angle value is for information only.
+
+</details>
+"""
+
+SECTION_2_TEXT = """
+## Detector rotation, overfocus and handedness
+
+<details>
+
+<summary>Instructions</summary>
+
+These parameters can only be calibrated in a dataset that was
+recorded with a strong defocus so that the detector shows a shadow
+image projection of the specimen.
+
+1. Move the green cursor in the plot "Corrected point analysis" to a
+   prominent feature.
+1. Select a scan position where this feature is also visible on a
+   detector frame. Selecting the position of the green cursor will
+   show the feature at the beam center.
+1. Observe if the green cursor in the second plot "Picked detector
+   frame" points to the same feature.
+1. Observe if the plot "Frames back-projected to scan coordinate
+   system" shows a sharp image of the specimen that matches the plot
+   "Corrected point analysis".
+1. If the positions don't match or the plot is distorted, move the
+   green cursor in the second plot "Picked detector frame" to the
+   feature and record the position. Do this for at least three
+   different features or scan positions.
+1. Press the button "Derive coordinate system from table" to
+   re-calibrate the coordinate system.
+1. Observe if the coordinate system is calibrated correctly now by
+   repeating the first four steps.
+
+</details>
+"""
+
+
+def tags_markdown(tags: Mapping[str, Any]) -> str:
+    """
+    Collapsible HTML table of the metadata tags of a dataset
+    """
+    if not tags:
+        return "Embedded metadata tags: None"
+    rows = "".join(
+        f"<tr><td>{html.escape(str(key))}</td><td>{html.escape(str(value))}</td></tr>"
+        for key, value in tags.items()
+    )
+    return (
+        "\n<details>\n<summary>Embedded metadata tags</summary>\n"
+        f"<table>{rows}</table>\n</details>\n"
+    )
+
+
 def format_model(model: Model4DSTEM, columns: list[str]) -> dict[str, str]:
     """
     Model parameters as text for the model table, rotations in degrees
@@ -1100,69 +1216,11 @@ def callback(args, model, res, blur):
 
     @property
     def layout(self):
-
-        if self.calibrated_dataset.tags:
-            tag_html = [
-                f"<tr><td>{html.escape(str(key))}</td><td>{html.escape(str(value))}</td></tr>"
-                for key, value in self.calibrated_dataset.tags.items()
-            ]
-            tag_html = f"<table>{''.join(tag_html)}</table>"
-            tag_md = f"""
-            <details>
-            <summary>Embedded metadata tags</summary>
-            {tag_html}
-            </details>
-            """
-        else:
-            tag_md = "Embedded metadata tags: None"
-
         self.section_1_label = pn.pane.Markdown(
-            f"""
-            # 4D STEM coordinate system calibration
-
-            **Dataset**: {html.escape(str(self.calibrated_dataset.dataset))}
-
-            {tag_md}
-
-            This tool helps to check and adjust the geometry of a 4D STEM
-            experiment.
-
-            ## De-scan error, convergence semi-angle and effective camera length
-
-            <details>
-            <summary>Instructions</summary>
-
-            Fisrt, enter the nominal values for the scan rotation as well as the
-            detector pixel pitch in the corresponding input fields.
-
-            Second, confirm or calibrate the shift of the beam center as a
-            function of scan position (de-scan error).
-
-            1. Change the selected scan position on the left
-            1. Observe if the cursor on the right stays at the beam center.
-            1. Observe if the plot "Corrected point analysis" shows a bright
-               field image of the specimen.
-            1. Observe if the plots "Frame corrected for descan error" and "Sum
-               of frames corrected for de-scan error" show the primary beam
-               exactly at the cursor position.
-            1. If the positions don't match or the plot is distorted, move the
-               cursor to the beam center and record the position. Do this for at
-               least three different scan positions.
-            1. Press the button "Apply correction from table" to re-calibrate
-               the de-scan error.
-            1. Observe if the de-scan error is calibrated correctly now by
-               repeating the first three steps.
-
-            In case the frames show diffraction rings or spots and a list of
-            "twotheta" angles was provided, now calibrate the camera length so
-            that the ring glyphs match the peaks or rings in the detector
-            frames.
-
-            Finally, adjust the convergence semi-angle so that the innermost
-            glyph matches the size of the diffraction disk.
-
-            </details>
-            """,
+            SECTION_1_TEXT.format(
+                dataset=html.escape(str(self.calibrated_dataset.dataset)),
+                tags=tags_markdown(self.calibrated_dataset.tags),
+            ),
             max_width=500,
         )
         inputs_1 = pn.layout.Row(
@@ -1171,25 +1229,7 @@ def layout(self):
             self.cl_input,
             self.semiconv_input,
         )
-        self.scan_step_label = pn.pane.Markdown(
-            """
-            ## Scan step calibration
-
-            <details>
-
-            <summary>Instructions</summary>
-
-            Move the yellow line handles in the plot "Corrected point analysis"
-            so that the line spans a feature of known physical size. Check the
-            size in the "Scale bar" input field and adjust if necessary. Note
-            that for defocused data the scale is only correct if the de-scan
-            error was compensated correctly.
-
-            The scale bar angle value is for information only.
-
-            </details>
-            """
-        )
+        self.scan_step_label = pn.pane.Markdown(SCAN_STEP_TEXT)
         inputs_2 = pn.layout.Row(
             self.scalebar_input,
             self.scale_angle_display,
@@ -1201,41 +1241,7 @@ def layout(self):
             self.corr_sum_fig.layout,
         )
 
-        self.section_2_label = pn.pane.Markdown(
-            """
-            ## Detector rotation, overfocus and handedness
-
-            <details>
-
-            <summary>Instructions</summary>
-
-            These parameters can only be calibrated in a dataset that was
-            recorded with a strong defocus so that the detector shows a shadow
-            image projection of the specimen.
-
-            1. Move the green cursor in the plot "Corrected point analysis" to a
-               prominent feature.
-            1. Select a scan position where this feature is also visible on a
-               detector frame. Selecting the position of the green cursor will
-               show the feature at the beam center.
-            1. Observe if the green cursor in the second plot "Picked detector
-               frame" points to the same feature.
-            1. Observe if the plot "Frames back-projected to scan coordinate
-               system" shows a sharp image of the specimen that matches the plot
-               "Corrected point analysis".
-            1. If the positions don't match or the plot is distorted, move the
-               green cursor in the second plot "Picked detector frame" to the
-               feature and record the position. Do this for at least three
-               different features or scan positions.
-            1. Press the button "Derive coordinate system from table" to
-               re-calibrate the coordinate system.
-            1. Observe if the coordinate system is calibrated correctly now by
-               repeating the first four steps.
-
-            </details>
-            """,
-            max_width=500,
-        )
+        self.section_2_label = pn.pane.Markdown(SECTION_2_TEXT, max_width=500)
         raw_figs_2 = pn.layout.Row(self.nav_fig_2.layout, self.pick_fig_2.layout)
         back_figs = pn.layout.Row(self.back_sum_fig.layout, self.back_pick_fig.layout)
         descan_buttons = pn.layout.Row(

From aaa32f5c463042c5a322f3558944c0e900546b58 Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Wed, 14 Oct 2026 08:16:45 +0000
Subject: [PATCH 33/34] Move CalibratedDataset to its own module so jeol_dm4
 doesn't import the GUI

ui.py imports it from there, so it is still available from microscope_calibration.ui.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 CLAUDE.md                                   |  7 +--
 src/microscope_calibration/dataset.py       | 53 +++++++++++++++++++++
 src/microscope_calibration/ui.py            | 46 ++----------------
 src/microscope_calibration/util/jeol_dm4.py |  2 +-
 4 files changed, 61 insertions(+), 47 deletions(-)
 create mode 100644 src/microscope_calibration/dataset.py

diff --git a/CLAUDE.md b/CLAUDE.md
index a965c57..eb0b105 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -104,9 +104,10 @@ Key layers, bottom-up:
    structure (uses the `diffraction` extra); feeds known-geometry references into the
    optimizers.
 
-8. **`ui.py`** — interactive Panel/Bokeh calibration GUI (`gui` extra). `CalibratedDataset`
-   pairs a dataset with a model; `CoordinateCorrectionLayout` wires sliders/tables to live
-   re-projection so a user can tune calibration against a displayed dataset.
+8. **`ui.py`** — interactive Panel/Bokeh calibration GUI (`gui` extra).
+   `CoordinateCorrectionLayout` wires sliders/tables to live re-projection so a user can tune
+   calibration against a displayed dataset. `CalibratedDataset` (`dataset.py`, also importable
+   from `ui`) pairs a dataset with a model.
 
 ### Mental model for making changes
 
diff --git a/src/microscope_calibration/dataset.py b/src/microscope_calibration/dataset.py
new file mode 100644
index 0000000..4666638
--- /dev/null
+++ b/src/microscope_calibration/dataset.py
@@ -0,0 +1,53 @@
+"""
+Datasets with a calibration, see :class:`CalibratedDataset`
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from typing import Any
+
+from libertem.io.dataset.base import DataSet
+
+from microscope_calibration.common.model import Model4DSTEM
+
+
+class AbstractCalibratedDataset(ABC):
+    @property
+    @abstractmethod
+    def model(self) -> Model4DSTEM:
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def dataset(self) -> DataSet:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def calibrated(self, model: Model4DSTEM) -> AbstractCalibratedDataset:
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def tags(self) -> Mapping[str, Any]:
+        raise NotImplementedError()
+
+
+class CalibratedDataset(AbstractCalibratedDataset):
+    def __init__(self, dataset: DataSet, model: Model4DSTEM):
+        self._ds = dataset
+        self._model = model
+
+    @property
+    def dataset(self) -> DataSet:
+        return self._ds
+
+    @property
+    def model(self) -> Model4DSTEM:
+        return self._model
+
+    def calibrated(self, model) -> CalibratedDataset:
+        return CalibratedDataset(dataset=self.dataset, model=model)
+
+    @property
+    def tags(self):
+        return {}
diff --git a/src/microscope_calibration/ui.py b/src/microscope_calibration/ui.py
index f115bf5..05b9c13 100644
--- a/src/microscope_calibration/ui.py
+++ b/src/microscope_calibration/ui.py
@@ -1,5 +1,4 @@
 import concurrent.futures
-from abc import ABC, abstractmethod
 from typing import Literal, Any
 from collections.abc import Mapping
 import html
@@ -29,6 +28,9 @@
 from libertem_ui.figure import ApertureFigure
 
 from .common.model import DescanError, Model4DSTEM, PixelYX, lambdify_trace_for
+
+# Also available from this module
+from .dataset import AbstractCalibratedDataset, CalibratedDataset  # noqa: F401
 from .common.stem_overfocus import (
     corrected_det_x,
     corrected_det_y,
@@ -60,48 +62,6 @@
 trace = lambdify_trace_for(jnp)
 
 
-class AbstractCalibratedDataset(ABC):
-    @property
-    @abstractmethod
-    def model(self) -> Model4DSTEM:
-        raise NotImplementedError()
-
-    @property
-    @abstractmethod
-    def dataset(self) -> DataSet:
-        raise NotImplementedError()
-
-    @abstractmethod
-    def calibrated(self, model: Model4DSTEM) -> AbstractCalibratedDataset:
-        raise NotImplementedError()
-
-    @property
-    @abstractmethod
-    def tags(self) -> Mapping[str, Any]:
-        raise NotImplementedError()
-
-
-class CalibratedDataset(AbstractCalibratedDataset):
-    def __init__(self, dataset: DataSet, model: Model4DSTEM):
-        self._ds = dataset
-        self._model = model
-
-    @property
-    def dataset(self) -> DataSet:
-        return self._ds
-
-    @property
-    def model(self) -> Model4DSTEM:
-        return self._model
-
-    def calibrated(self, model) -> CalibratedDataset:
-        return CalibratedDataset(dataset=self.dataset, model=model)
-
-    @property
-    def tags(self):
-        return {}
-
-
 # Instructions in the layout, see CoordinateCorrectionLayout.layout
 SECTION_1_TEXT = """
 # 4D STEM coordinate system calibration
diff --git a/src/microscope_calibration/util/jeol_dm4.py b/src/microscope_calibration/util/jeol_dm4.py
index 34a3d6a..995eb7c 100644
--- a/src/microscope_calibration/util/jeol_dm4.py
+++ b/src/microscope_calibration/util/jeol_dm4.py
@@ -17,7 +17,7 @@
 from ncempy.io.dm import fileDM
 
 from microscope_calibration.common.model import DescanError, Model4DSTEM
-from microscope_calibration.ui import CalibratedDataset
+from microscope_calibration.dataset import CalibratedDataset
 
 # Tags of the 4D STEM data. ImageList.1 is the thumbnail.
 IMAGE = ".ImageList.2"

From 7daa7ae62d090e8fc43550af82044689a7a571af Mon Sep 17 00:00:00 2001
From: Dieter Weber <d.weber@fz-juelich.de>
Date: Wed, 14 Oct 2026 08:41:38 +0000
Subject: [PATCH 34/34] Fix the instructions of the calibration UI and order
 the sections accordingly

Behavior changes:
- The instructions name the plots by their titles and fix a typo
- The corrected plots follow the uncorrected ones in the first section, the scale bar
  section follows the plot it refers to
- Remove a double space from the title of the corrected frame

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/microscope_calibration/ui.py | 46 +++++++++++++++++---------------
 1 file changed, 24 insertions(+), 22 deletions(-)

diff --git a/src/microscope_calibration/ui.py b/src/microscope_calibration/ui.py
index 05b9c13..8c53caf 100644
--- a/src/microscope_calibration/ui.py
+++ b/src/microscope_calibration/ui.py
@@ -78,17 +78,18 @@
 <details>
 <summary>Instructions</summary>
 
-Fisrt, enter the nominal values for the scan rotation as well as the
+First, enter the nominal values for the scan rotation as well as the
 detector pixel pitch in the corresponding input fields.
 
 Second, confirm or calibrate the shift of the beam center as a
 function of scan position (de-scan error).
 
-1. Change the selected scan position on the left
-1. Observe if the cursor on the right stays at the beam center.
-1. Observe if the plot "Corrected point analysis" shows a bright
-   field image of the specimen.
-1. Observe if the plots "Frame corrected for descan error" and "Sum
+1. Change the selected scan position in the plot "Uncorrected scan map".
+1. Observe if the cursor in the plot "Currently picked uncorrected
+   detector frame" stays at the beam center.
+1. Observe if the plot "Point analysis corrected for de-scan error"
+   shows a bright field image of the specimen.
+1. Observe if the plots "Frame corrected for de-scan error" and "Sum
    of frames corrected for de-scan error" show the primary beam
    exactly at the cursor position.
 1. If the positions don't match or the plot is distorted, move the
@@ -117,8 +118,8 @@
 
 <summary>Instructions</summary>
 
-Move the yellow line handles in the plot "Corrected point analysis"
-so that the line spans a feature of known physical size. Check the
+Move the yellow line handles in the plot "Point analysis corrected for
+de-scan error" above so that the line spans a feature of known physical size. Check the
 size in the "Scale bar" input field and adjust if necessary. Note
 that for defocused data the scale is only correct if the de-scan
 error was compensated correctly.
@@ -139,19 +140,20 @@
 recorded with a strong defocus so that the detector shows a shadow
 image projection of the specimen.
 
-1. Move the green cursor in the plot "Corrected point analysis" to a
-   prominent feature.
-1. Select a scan position where this feature is also visible on a
-   detector frame. Selecting the position of the green cursor will
-   show the feature at the beam center.
-1. Observe if the green cursor in the second plot "Picked detector
-   frame" points to the same feature.
-1. Observe if the plot "Frames back-projected to scan coordinate
+1. Move the green cursor in the plot "Point analysis corrected for
+   de-scan error" above to a prominent feature.
+1. Select a scan position in the plot "Uncorrected scan map" below
+   where this feature is also visible on a detector frame. Selecting
+   the position of the green cursor will show the feature at the beam
+   center.
+1. Observe if the green cursor in the plot "Currently picked
+   uncorrected detector frame" below points to the same feature.
+1. Observe if the plot "Frames superimposed in scan coordinate
    system" shows a sharp image of the specimen that matches the plot
-   "Corrected point analysis".
+   "Point analysis corrected for de-scan error".
 1. If the positions don't match or the plot is distorted, move the
-   green cursor in the second plot "Picked detector frame" to the
-   feature and record the position. Do this for at least three
+   green cursor in the plot "Currently picked uncorrected detector
+   frame" below to the feature and record the position. Do this for at least three
    different features or scan positions.
 1. Press the button "Derive coordinate system from table" to
    re-calibrate the coordinate system.
@@ -423,7 +425,7 @@ def _init_figures(self):
 
         self.corr_pick_fig = self._new_figure(
             frames["corrected"],
-            title=f"Frame corrected for de-scan error  ({self.sig_mode})",
+            title=f"Frame corrected for de-scan error ({self.sig_mode})",
             maxdim=250,
         )
         self.corr_beam_centre = self._add_cursor(
@@ -1239,11 +1241,11 @@ def layout(self):
             self.section_1_label,
             inputs_1,
             raw_figs,
-            descan_section,
-            self.section_2_label,
             corrected_figs,
+            descan_section,
             self.scan_step_label,
             inputs_2,
+            self.section_2_label,
             raw_figs_2,
             coord_section,
             back_figs,