Add golden-output regression tests for mPIE and qNewton

tests/regression/conftest.py

@@ -0,0 +1,58 @@
+"""Fixtures for the numerical regression suite.
+
+The datasets here are synthesized in-test with a fixed seed rather than loaded
+from ``example_data/``. That is deliberate: locally ``example_data/simu.hdf5``
+is a real 100x128x128 measurement, while CI synthesizes a 20x64x64 stand-in, so
+a golden recorded against "simu.hdf5" would not be reproducible across the two.
+Everything the goldens depend on is generated here.
+"""
+
+import h5py
+import numpy as np
+import pytest
+
+# Small enough that a handful of iterations of every engine runs in ~seconds,
+# large enough that the object is meaningfully bigger than the probe (so the
+# patch-extraction and scatter-back indexing is actually exercised).
+ND = 32
+GRID = 4 # GRID x GRID raster scan -> 16 frames
+SEED = 20240607
+
+
+def _simulate_dataset(path, nd=ND, grid=GRID, seed=SEED):
+ """Write a deterministic CPM dataset to ``path``."""
+ rng = np.random.default_rng(seed)
+ n_frames = grid * grid
+
+ # A raster scan, in metres, with a slight jitter so positions are not
+ # perfectly degenerate.
+ step = 3e-6
+ coords = (np.arange(grid) - (grid - 1) / 2) * step
+ yy, xx = np.meshgrid(coords, coords, indexing="ij")
+ encoder = np.stack([yy.ravel(), xx.ravel()], axis=1)
+ encoder = encoder + rng.normal(0, step / 50, encoder.shape)
+
+ # Smooth-ish speckle: random field low-pass filtered in Fourier space, so
+ # the diffraction patterns look like data rather than white noise.
+ field = rng.random((n_frames, nd, nd))
+ spectrum = np.fft.fftshift(np.fft.fft2(field), axes=(-2, -1))
+ ky = np.fft.fftshift(np.fft.fftfreq(nd))
+ mask = (ky[:, None] ** 2 + ky[None, :] ** 2) < 0.25**2
+ ptychogram = np.abs(np.fft.ifft2(np.fft.ifftshift(spectrum * mask, axes=(-2, -1)))) ** 2
+ ptychogram = (ptychogram / ptychogram.max()).astype(np.float32)
+
+ with h5py.File(path, "w") as hf:
+ hf.create_dataset("ptychogram", data=ptychogram, dtype="f")
+ hf.create_dataset("encoder", data=encoder, dtype="f")
+ hf.create_dataset("dxd", data=np.array(75e-6))
+ hf.create_dataset("zo", data=np.array(0.05))
+ hf.create_dataset("wavelength", data=np.array(632.8e-9))
+ hf.create_dataset("entrancePupilDiameter", data=np.array(400e-6))
+ return path
+
+
+@pytest.fixture(scope="session")
+def regression_dataset(tmp_path_factory):
+ """Path to a deterministic synthetic CPM dataset."""
+ path = tmp_path_factory.mktemp("regression_data") / "regression_cpm.hdf5"
+ return _simulate_dataset(path)

tests/regression/test_engine_regression.py

@@ -0,0 +1,172 @@
+"""Golden-output regression tests for the reconstruction engines.
+
+These pin the *numerical output* of each engine so that performance work
+(kernel fusion, CUDA graphs, alternative linear algebra) cannot silently change
+what the library computes.
+
+Coverage includes a mixed-state configuration, not just single-mode: that
+exercises the 6D ``(nlambda, nosm, npsm, nslice, Ny, Nx)`` broadcasting which is
+the easiest thing to get wrong when rewriting the update rules, and exactly what
+a single-mode-only suite would let through. Polychromatic and multislice
+configurations follow once the engines that need them are testable.
+
+To re-record the goldens after an *intended* numerical change::
+
+ PTYLAB_REGEN_GOLDENS=1 uv run pytest tests/regression -q
+
+Review the resulting diff in ``tests/regression/data/`` before committing it.
+"""
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
+from PtyLab.Monitor.Monitor import DummyMonitor
+from PtyLab.Params.Params import Params
+from PtyLab.Reconstruction.Reconstruction import Reconstruction
+from PtyLab.utils.gpuUtils import asNumpyArray
+
+try:
+ import cupy
+
+ HAS_GPU = cupy.cuda.is_available()
+except Exception:
+ HAS_GPU = False
+
+GOLDEN_DIR = Path(__file__).parent / "data"
+REGEN = os.environ.get("PTYLAB_REGEN_GOLDENS", "") not in ("", "0")
+SEED = 20240607
+
+# name -> engine, propagator, (nlambda, nosm, npsm, nslice), iterations
+#
+# Engines are added here as they become testable. ePIE and e3PIE are absent on
+# purpose: ePIE.reconstruct() is currently a generator that nobody iterates, so
+# it returns without doing any work, and e3PIE raises on its own betaProbe.
+# Pinning either today would record a meaningless baseline. They join this table
+# in the PRs that repair them.
+CONFIGS = {
+ "mpie_single": ("mPIE", "Fraunhofer", (1, 1, 1, 1), 3),
+ "mpie_mixed_state": ("mPIE", "Fraunhofer", (1, 2, 3, 1), 3),
+ "qnewton_single": ("qNewton", "Fraunhofer", (1, 1, 1, 1), 3),
+}
+
+
+def build(dataset, config, gpu):
+ """Construct a fully-determined reconstruction for ``config``."""
+ _engine_name, propagator, (nlambda, nosm, npsm, nslice), _iters = config
+
+ data = ExperimentalData(str(dataset), operationMode="CPM")
+ params = Params()
+ params.gpuSwitch = gpu
+ params.propagatorType = propagator
+ # 'random' shuffles positions via the global numpy RNG; pin the order so the
+ # golden does not depend on RNG call sequence elsewhere in the engine.
+ params.positionOrder = "sequential"
+
+ reconstruction = Reconstruction(data, params)
+ reconstruction.nlambda = nlambda
+ reconstruction.nosm = nosm
+ reconstruction.npsm = npsm
+ reconstruction.nslice = nslice
+
+ if nlambda > 1:
+ base = float(np.atleast_1d(reconstruction.wavelength)[0])
+ reconstruction.spectralDensity = base * np.linspace(0.98, 1.02, nlambda)
+ if nslice > 1:
+ reconstruction.dz = 1e-4
+ reconstruction.refrIndex = 1.0
+
+ # initialProbeOrObject() adds 0.001 * np.random.rand(...) noise to break mode
+ # degeneracy, so the initial guess itself needs the global seed pinned.
+ np.random.seed(SEED)
+ reconstruction.initializeObjectProbe()
+
+ return data, reconstruction, params, DummyMonitor()
+
+
+def run(dataset, config, gpu):
+ """Run ``config`` to completion and return the arrays worth pinning."""
+ from PtyLab import Engines
+
+ engine_name, _propagator, _modes, iters = config
+ data, reconstruction, params, monitor = build(dataset, config, gpu)
+
+ engine = getattr(Engines, engine_name)(reconstruction, data, params, monitor)
+ engine.numIterations = iters
+
+ # mPIE fires its momentum update on np.random.rand(1) > 0.95; seed again so
+ # the decision sequence is fixed regardless of how much RNG setup consumed.
+ np.random.seed(SEED)
+ engine.reconstruct()
+
+ return {
+ "object": asNumpyArray(reconstruction.object),
+ "probe": asNumpyArray(reconstruction.probe),
+ "error": np.asarray(asNumpyArray(reconstruction.error), dtype=np.float64),
+ }
+
+
+def compare(result, golden_path, rtol, atol, label):
+ if REGEN or not golden_path.exists():
+ golden_path.parent.mkdir(parents=True, exist_ok=True)
+ np.savez_compressed(golden_path, **result)
+ pytest.skip(f"recorded golden {golden_path.name}; re-run to verify")
+
+ golden = np.load(golden_path)
+ assert sorted(golden.files) == sorted(result), (
+ f"{label}: golden holds {sorted(golden.files)} but got {sorted(result)}; "
+ f"re-record with PTYLAB_REGEN_GOLDENS=1"
+ )
+ for key in sorted(result):
+ actual, expected = result[key], golden[key]
+ assert actual.shape == expected.shape, (
+ f"{label}: {key} shape {actual.shape} != golden {expected.shape}"
+ )
+ np.testing.assert_allclose(
+ actual, expected, rtol=rtol, atol=atol,
+ err_msg=f"{label}: {key} drifted from golden",
+ )
+
+
+def relative_error(actual, expected):
+ """Frobenius-norm relative error.
+
+ Elementwise relative error is the wrong metric here: these arrays contain
+ near-zero elements where a 1e-7 absolute wobble reads as a huge relative
+ one. The norm ratio measures what actually matters -- whether the
+ reconstruction as a whole moved.
+ """
+ denom = np.linalg.norm(np.asarray(expected).ravel())
+ return float(np.linalg.norm((np.asarray(actual) - expected).ravel()) / max(denom, 1e-30))
+
+
+@pytest.mark.parametrize("name", list(CONFIGS))
+def test_engine_cpu_golden(regression_dataset, name):
+ """CPU output must match the recorded golden."""
+ result = run(regression_dataset, CONFIGS[name], gpu=False)
+ compare(result, GOLDEN_DIR / f"{name}.npz", rtol=1e-5, atol=1e-7,
+ label=f"{name} [cpu]")
+
+
+@pytest.mark.skipif(not HAS_GPU, reason="no CUDA GPU available")
+@pytest.mark.parametrize("name", list(CONFIGS))
+def test_engine_gpu_agrees_with_cpu(regression_dataset, name):
+ """GPU and CPU backends must agree to within float32 accumulation noise.
+
+ cuFFT and CuPy reductions accumulate in a different order than NumPy, so
+ exact equality is not expected. Measured divergence across these configs is
+ 1e-8 to 2.2e-4; the 1e-3 bound leaves roughly 5x headroom while still being
+ tight enough to catch a genuinely wrong kernel.
+ """
+ cpu = run(regression_dataset, CONFIGS[name], gpu=False)
+ gpu = run(regression_dataset, CONFIGS[name], gpu=True)
+
+ for key in ("object", "probe", "error"):
+ err = relative_error(gpu[key], cpu[key])
+ assert err < 1e-3, (
+ f"{name}: GPU {key} diverges from CPU by {err:.2e} "
+ f"(relative Frobenius norm, tolerance 1e-3)"
+ )