Build propagator transfer functions on the device the data is on

PtyLab/Operators/Operators.py

@@ -268,7 +268,8 @@ def propagate_twoStepPolychrome(
         tuple(reconstruction.spectralDensity),
         reconstruction.Lp,
         reconstruction.dxp,
-        params.gpuSwitch,
+        # device placement follows the data, not the global switch
+        isGpuArray(fields),
     )
     if inverse:
         result = ifft2c(
@@ -340,7 +341,8 @@ def propagate_scaledASP(
         reconstruction.wavelength,
         reconstruction.dxo,
         reconstruction.dxd,
-        params.gpuSwitch,
+        # device placement follows the data, not the global switch
+        isGpuArray(fields),
     )
     if inverse:
         Q1, Q2 = Q1.conj(), Q2.conj()
@@ -412,7 +414,8 @@ def propagate_scaledPolychromeASP(
         tuple(reconstruction.spectralDensity),
         reconstruction.dxo,
         reconstruction.dxd,
-        params.gpuSwitch,
+        # device placement follows the data, not the global switch
+        isGpuArray(fields),
     )
     if inverse:
         Q1, Q2 = Q1.conj(), Q2.conj()
@@ -492,7 +495,8 @@ def propagate_polychromeASP(
         reconstruction.Lp,
         reconstruction.nlambda,
         tuple(reconstruction.spectralDensity),
-        params.gpuSwitch,
+        # device placement follows the data, not the global switch
+        isGpuArray(fields),
     )
 
     if inverse:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ptylab"
-version = "0.2.7"
+version = "0.2.8"
 description = "A cross-platform, open-source inverse modeling toolbox for conventional and Fourier ptychography"
 authors = [
     { name = "Lars Loetgering", email = "lars.loetgering@fulbrightmail.org" },

tests/Operators/test_operators_integration.py

@@ -44,6 +44,123 @@ def test_object2detector():
         Operators.Operators.object2detector(reconstruction.esw, params, reconstruction)
 
 
+# Unitary propagators: forward followed by inverse must return the input.
+UNITARY_PROPAGATORS = ["fraunhofer", "fresnel", "identity"]
+
+# Band-limited propagators: forward followed by inverse is *not* the identity.
+# Their transfer functions suppress out-of-band spatial frequencies, and that
+# energy is genuinely gone -- measured ||P(x)-x||/||x|| is 3.4e-01 for asp and
+# 5.4e-02 for scaledasp.
+#
+# What that leaves differs between the two, which is why neither is asserted
+# here: asp is a clean projection (||P(P(x))-P(x)||/||P(x)|| = 3.1e-07, i.e. a
+# hard 0/1 cutoff), while scaledasp only approximately so (1.6e-03) because its
+# chirp and rescaling factors are not idempotent. Linearity is the property both
+# genuinely share; the numbers themselves are pinned by
+# tests/regression/test_propagator_regression.py.
+BANDLIMITED_PROPAGATORS = ["asp", "scaledasp"]
+
+# The polychrome variants need spectralDensity with nlambda > 1; they are
+# covered end-to-end by tests/regression instead.
+ALL_ROUND_TRIP_PROPAGATORS = UNITARY_PROPAGATORS + BANDLIMITED_PROPAGATORS
+
+
+def _round_trip(field, params, reconstruction):
+    reconstruction.esw = field
+    _, forward = Operators.Operators.object2detector(field, params, reconstruction)
+    _, back = Operators.Operators.detector2object(forward, params, reconstruction)
+    return back
+
+
+@pytest.mark.parametrize("propagator", UNITARY_PROPAGATORS)
+def test_propagator_round_trip_is_identity(propagator):
+    """Forward then inverse must return the input field.
+
+    This is the property every performance change to the propagators has to
+    preserve, and it is what the previous smoke tests (which called the
+    operators but asserted nothing) failed to check.
+    """
+    experimentalData, reconstruction, params, monitor, engine = easyInitialize(
+        "example:simulation_cpm"
+    )
+    params.gpuSwitch = False
+    params.propagatorType = propagator
+    params.fftshiftSwitch = False
+    reconstruction._move_data_to_cpu()
+
+    field = reconstruction.probe.copy()
+    back = _round_trip(field, params, reconstruction)
+
+    scale = np.abs(field).max()
+    assert_allclose(
+        back, field, rtol=1e-4, atol=1e-5 * scale,
+        err_msg=f"{propagator} does not round-trip",
+    )
+
+
+@pytest.mark.parametrize("propagator", ALL_ROUND_TRIP_PROPAGATORS)
+def test_propagator_is_linear(propagator):
+    """Every propagator must be a linear operator: P(a*x) == a*P(x).
+
+    For the band-limited propagators this is the strongest property that
+    actually holds -- their transfer functions carry amplitude masks
+    (|Q| spans [0, 1]), so forward-then-inverse is neither unitary nor
+    idempotent. The exact numerical output is pinned separately by
+    tests/regression/test_propagator_regression.py.
+    """
+    experimentalData, reconstruction, params, monitor, engine = easyInitialize(
+        "example:simulation_cpm"
+    )
+    params.gpuSwitch = False
+    params.propagatorType = propagator
+    params.fftshiftSwitch = False
+    reconstruction._move_data_to_cpu()
+
+    field = reconstruction.probe.copy()
+    alpha = 3.0
+
+    reconstruction.esw = field
+    _, base = Operators.Operators.object2detector(field, params, reconstruction)
+    scaled_field = (alpha * field).astype(field.dtype)
+    reconstruction.esw = scaled_field
+    _, scaled = Operators.Operators.object2detector(
+        scaled_field, params, reconstruction
+    )
+
+    tol = 1e-5 * np.abs(alpha * base).max()
+    assert_allclose(
+        scaled, alpha * base, rtol=1e-4, atol=tol,
+        err_msg=f"{propagator} is not linear",
+    )
+
+
+@pytest.mark.parametrize("propagator", ALL_ROUND_TRIP_PROPAGATORS)
+@pytest.mark.skipif(not HAS_GPU, reason="no CUDA GPU available")
+def test_propagator_gpu_matches_cpu(propagator):
+    """The GPU propagator path must agree with the CPU one."""
+    experimentalData, reconstruction, params, monitor, engine = easyInitialize(
+        "example:simulation_cpm"
+    )
+    params.propagatorType = propagator
+    params.fftshiftSwitch = False
+
+    params.gpuSwitch = False
+    reconstruction._move_data_to_cpu()
+    field_cpu = reconstruction.probe.copy()
+    reconstruction.esw = field_cpu
+    _, out_cpu = Operators.Operators.object2detector(field_cpu, params, reconstruction)
+
+    field_gpu = cp.asarray(field_cpu)
+    reconstruction.esw = field_gpu
+    _, out_gpu = Operators.Operators.object2detector(field_gpu, params, reconstruction)
+
+    scale = np.abs(out_cpu).max()
+    assert_allclose(
+        cp.asnumpy(out_gpu), out_cpu, rtol=1e-4, atol=1e-5 * scale,
+        err_msg=f"{propagator}: GPU output disagrees with CPU",
+    )
+
+
 def test_propagate_fresnel():
     experimentalData, reconstruction, params, monitor, engine = easyInitialize(
         "example:simulation_cpm"

tests/regression/test_propagator_regression.py

@@ -0,0 +1,113 @@
+"""Golden-output regression tests for the propagators themselves.
+
+The property tests in ``tests/Operators/test_operators_integration.py`` pin the
+mathematical invariants that hold (round-trip identity for the unitary
+propagators, linearity for all of them). These pin the actual numbers, which is
+what catches a subtly wrong transfer function, a changed FFT convention, or a
+dropped ``fftshift``.
+
+Re-record with ``PTYLAB_REGEN_GOLDENS=1`` (see test_engine_regression).
+"""
+
+import numpy as np
+import pytest
+
+from PtyLab import Operators
+from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
+from PtyLab.Params.Params import Params
+from PtyLab.Reconstruction.Reconstruction import Reconstruction
+
+from test_engine_regression import GOLDEN_DIR, HAS_GPU, SEED, compare, relative_error
+
+if HAS_GPU:
+    import cupy as cp
+
+# scaledpolychromeasp/polychromeasp/twosteppolychrome need nlambda > 1; the rest
+# run single-wavelength.
+PROPAGATORS = {
+    "fraunhofer": 1,
+    "fresnel": 1,
+    "asp": 1,
+    "scaledasp": 1,
+    "identity": 1,
+    "polychromeasp": 3,
+    "scaledpolychromeasp": 3,
+    "twosteppolychrome": 3,
+}
+
+
+def _setup(dataset, propagator, nlambda):
+    data = ExperimentalData(str(dataset), operationMode="CPM")
+    params = Params()
+    params.gpuSwitch = False
+    params.propagatorType = propagator
+    params.fftshiftSwitch = False
+
+    reconstruction = Reconstruction(data, params)
+    reconstruction.nlambda = nlambda
+    if nlambda > 1:
+        base = float(np.atleast_1d(reconstruction.wavelength)[0])
+        reconstruction.spectralDensity = base * np.linspace(0.98, 1.02, nlambda)
+
+    np.random.seed(SEED)
+    reconstruction.initializeObjectProbe()
+    return data, reconstruction, params
+
+
+@pytest.mark.parametrize("propagator", list(PROPAGATORS))
+def test_propagator_output_golden(regression_dataset, propagator):
+    nlambda = PROPAGATORS[propagator]
+    _data, reconstruction, params = _setup(regression_dataset, propagator, nlambda)
+
+    field = reconstruction.probe.copy()
+    reconstruction.esw = field
+    _, forward = Operators.Operators.object2detector(field, params, reconstruction)
+    # BaseEngine.intensityProjection sets reconstruction.ESW before propagating
+    # back; propagate_twoStepPolychrome_inv reads it, so mirror that here.
+    reconstruction.ESW = forward
+    _, back = Operators.Operators.detector2object(forward, params, reconstruction)
+
+    result = {
+        "forward": np.asarray(forward),
+        "back": np.asarray(back),
+        # zero-size arrays are awkward in npz; store the norms as a cheap
+        # scalar summary that fails loudly on any global scaling change
+        "error": np.array(
+            [np.linalg.norm(forward.ravel()), np.linalg.norm(back.ravel())]
+        ),
+    }
+    compare(result, GOLDEN_DIR / f"propagator_{propagator}.npz",
+            rtol=1e-5, atol=1e-7, label=f"propagator {propagator}")
+
+
+@pytest.mark.skipif(not HAS_GPU, reason="no CUDA GPU available")
+@pytest.mark.parametrize("propagator", list(PROPAGATORS))
+def test_propagator_gpu_matches_cpu(regression_dataset, propagator):
+    """Propagating a GPU field must match the CPU result.
+
+    ``params.gpuSwitch`` is deliberately left False while GPU arrays are passed
+    in. Device placement has to follow the data, not the global switch --
+    ``BaseEngine._checkGPU`` and the engines both move arrays independently of
+    when a propagator's cached transfer function is first built, so a
+    transfer function built from ``params.gpuSwitch`` lands on the wrong device.
+
+    This covers all four call sites that read the switch, including the three
+    polychrome propagators that need nlambda > 1 and so cannot be reached from
+    the single-wavelength suite in tests/Operators.
+    """
+    nlambda = PROPAGATORS[propagator]
+    _data, reconstruction, params = _setup(regression_dataset, propagator, nlambda)
+
+    field_cpu = reconstruction.probe.copy()
+    reconstruction.esw = field_cpu
+    _, out_cpu = Operators.Operators.object2detector(field_cpu, params, reconstruction)
+
+    field_gpu = cp.asarray(field_cpu)
+    reconstruction.esw = field_gpu
+    _, out_gpu = Operators.Operators.object2detector(field_gpu, params, reconstruction)
+
+    err = relative_error(cp.asnumpy(out_gpu), np.asarray(out_cpu))
+    assert err < 1e-3, (
+        f"{propagator}: GPU output diverges from CPU by {err:.2e} "
+        f"(relative Frobenius norm, tolerance 1e-3)"
+    )