diff --git a/PtyLab/Operators/Operators.py b/PtyLab/Operators/Operators.py index 7181342..419c7b8 100644 --- a/PtyLab/Operators/Operators.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index bb3e095..2210b4d 100644 --- a/pyproject.toml +++ b/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" }, diff --git a/tests/Operators/test_operators_integration.py b/tests/Operators/test_operators_integration.py index 5c63009..5d9962f 100644 --- a/tests/Operators/test_operators_integration.py +++ b/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" diff --git a/tests/regression/data/propagator_asp.npz b/tests/regression/data/propagator_asp.npz new file mode 100644 index 0000000..7429a05 Binary files /dev/null and b/tests/regression/data/propagator_asp.npz differ diff --git a/tests/regression/data/propagator_fraunhofer.npz b/tests/regression/data/propagator_fraunhofer.npz new file mode 100644 index 0000000..aa0d936 Binary files /dev/null and b/tests/regression/data/propagator_fraunhofer.npz differ diff --git a/tests/regression/data/propagator_fresnel.npz b/tests/regression/data/propagator_fresnel.npz new file mode 100644 index 0000000..e8a2ec0 Binary files /dev/null and b/tests/regression/data/propagator_fresnel.npz differ diff --git a/tests/regression/data/propagator_identity.npz b/tests/regression/data/propagator_identity.npz new file mode 100644 index 0000000..27656cd Binary files /dev/null and b/tests/regression/data/propagator_identity.npz differ diff --git a/tests/regression/data/propagator_polychromeasp.npz b/tests/regression/data/propagator_polychromeasp.npz new file mode 100644 index 0000000..c560cc1 Binary files /dev/null and b/tests/regression/data/propagator_polychromeasp.npz differ diff --git a/tests/regression/data/propagator_scaledasp.npz b/tests/regression/data/propagator_scaledasp.npz new file mode 100644 index 0000000..ae6de9a Binary files /dev/null and b/tests/regression/data/propagator_scaledasp.npz differ diff --git a/tests/regression/data/propagator_scaledpolychromeasp.npz b/tests/regression/data/propagator_scaledpolychromeasp.npz new file mode 100644 index 0000000..5161629 Binary files /dev/null and b/tests/regression/data/propagator_scaledpolychromeasp.npz differ diff --git a/tests/regression/data/propagator_twosteppolychrome.npz b/tests/regression/data/propagator_twosteppolychrome.npz new file mode 100644 index 0000000..2aed593 Binary files /dev/null and b/tests/regression/data/propagator_twosteppolychrome.npz differ diff --git a/tests/regression/test_propagator_regression.py b/tests/regression/test_propagator_regression.py new file mode 100644 index 0000000..2abe0c2 --- /dev/null +++ b/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)" + )