feat(phase_thick_3d_tilt): Re-factor of ops_process PR#99 into OPS agnostic waveorder optimizations.

tests/models/test_isotropic_thin_3d.py

@@ -41,3 +41,155 @@ def test_reconstruct():
     assert phase.shape == yx_shape
     assert np.all(np.isfinite(absorption.numpy()))
     assert np.all(np.isfinite(phase.numpy()))
+
+
+_WRAP_KWARGS = dict(
+    yx_shape=(64, 64),
+    yx_pixel_size=6.5 / 40,
+    z_position_list=[-1.0, 0.0, 1.0],
+    wavelength_illumination=0.532,
+    index_of_refraction_media=1.33,
+    numerical_aperture_illumination=0.4,
+    numerical_aperture_detection=0.55,
+    invert_phase_contrast=False,
+    tilt_angle_zenith=0.1,
+    tilt_angle_azimuth=0.2,
+    pupil_steepness=1e4,
+)
+
+
+def test_thin_3d_angle_z_split_composes_to_wrap_unsafe():
+    """The thin-3d angle/z optics split composes back to bit-identical legacy output."""
+    legacy_Hu, legacy_Hp = isotropic_thin_3d._calculate_wrap_unsafe_transfer_function(**_WRAP_KWARGS)
+
+    angle_optics = isotropic_thin_3d._compute_angle_optics(
+        _WRAP_KWARGS["yx_shape"],
+        _WRAP_KWARGS["yx_pixel_size"],
+        _WRAP_KWARGS["wavelength_illumination"],
+        _WRAP_KWARGS["index_of_refraction_media"],
+        _WRAP_KWARGS["numerical_aperture_illumination"],
+        _WRAP_KWARGS["numerical_aperture_detection"],
+        tilt_angle_zenith=_WRAP_KWARGS["tilt_angle_zenith"],
+        tilt_angle_azimuth=_WRAP_KWARGS["tilt_angle_azimuth"],
+        pupil_steepness=_WRAP_KWARGS["pupil_steepness"],
+    )
+    det_prop = isotropic_thin_3d._compute_z_propagation(
+        angle_optics,
+        _WRAP_KWARGS["z_position_list"],
+        invert_phase_contrast=_WRAP_KWARGS["invert_phase_contrast"],
+    )
+    Hu, Hp = isotropic_thin_3d._wotf_from_split_optics(angle_optics, det_prop)
+    assert torch.equal(legacy_Hu, Hu)
+    assert torch.equal(legacy_Hp, Hp)
+
+
+def test_thin_3d_angle_optics_cached_across_z_changes():
+    """Cached angle optics give bit-identical WOTFs when only z changes.
+
+    This is the actual FREEZE_ANGLES workflow: build angle optics ONCE,
+    re-call _compute_z_propagation per optimizer iter with new
+    z_position_list. Compare against the legacy single-call path
+    invoked fresh for each z.
+    """
+    z_lists = [[-1.0, 0.0, 1.0], [-0.8, 0.0, 0.8], [-1.5, 0.0, 1.5]]
+    base = dict(_WRAP_KWARGS)
+
+    angle_optics = isotropic_thin_3d._compute_angle_optics(
+        base["yx_shape"],
+        base["yx_pixel_size"],
+        base["wavelength_illumination"],
+        base["index_of_refraction_media"],
+        base["numerical_aperture_illumination"],
+        base["numerical_aperture_detection"],
+        tilt_angle_zenith=base["tilt_angle_zenith"],
+        tilt_angle_azimuth=base["tilt_angle_azimuth"],
+        pupil_steepness=base["pupil_steepness"],
+    )
+    for z_list in z_lists:
+        kw = {**base, "z_position_list": z_list}
+        legacy_Hu, legacy_Hp = isotropic_thin_3d._calculate_wrap_unsafe_transfer_function(**kw)
+        det_prop = isotropic_thin_3d._compute_z_propagation(
+            angle_optics, z_list, invert_phase_contrast=base["invert_phase_contrast"]
+        )
+        Hu, Hp = isotropic_thin_3d._wotf_from_split_optics(angle_optics, det_prop)
+        assert torch.equal(legacy_Hu, Hu), f"abs TF mismatch at z={z_list}"
+        assert torch.equal(legacy_Hp, Hp), f"phase TF mismatch at z={z_list}"
+
+
+def test_thin_3d_angle_optics_batched_tilt():
+    """Batched (B,) tilt angles produce the same split output as legacy."""
+    kw = dict(_WRAP_KWARGS)
+    kw["tilt_angle_zenith"] = torch.tensor([0.0, 0.1, 0.2])
+    kw["tilt_angle_azimuth"] = torch.tensor([0.0, 0.5, 1.0])
+
+    legacy_Hu, legacy_Hp = isotropic_thin_3d._calculate_wrap_unsafe_transfer_function(**kw)
+    assert legacy_Hu.shape[0] == 3
+
+    angle_optics = isotropic_thin_3d._compute_angle_optics(
+        kw["yx_shape"],
+        kw["yx_pixel_size"],
+        kw["wavelength_illumination"],
+        kw["index_of_refraction_media"],
+        kw["numerical_aperture_illumination"],
+        kw["numerical_aperture_detection"],
+        tilt_angle_zenith=kw["tilt_angle_zenith"],
+        tilt_angle_azimuth=kw["tilt_angle_azimuth"],
+        pupil_steepness=kw["pupil_steepness"],
+    )
+    assert angle_optics["batched"]
+    det_prop = isotropic_thin_3d._compute_z_propagation(angle_optics, kw["z_position_list"])
+    Hu, Hp = isotropic_thin_3d._wotf_from_split_optics(angle_optics, det_prop)
+    assert torch.equal(legacy_Hu, Hu)
+    assert torch.equal(legacy_Hp, Hp)
+
+
+def test_cached_tilt_optics_matches_legacy():
+    """CachedTiltOptics produces bit-identical TFs to the legacy fresh build.
+
+    The FREEZE_ANGLES workflow: build cache once, call transfer_functions()
+    each iter. The output must match what a fresh single-shot
+    `_calculate_wrap_unsafe_transfer_function` would produce for the same
+    inputs.
+    """
+    cache = isotropic_thin_3d.CachedTiltOptics(
+        yx_shape=_WRAP_KWARGS["yx_shape"],
+        yx_pixel_size=_WRAP_KWARGS["yx_pixel_size"],
+        wavelength_illumination=_WRAP_KWARGS["wavelength_illumination"],
+        index_of_refraction_media=_WRAP_KWARGS["index_of_refraction_media"],
+        numerical_aperture_illumination=_WRAP_KWARGS["numerical_aperture_illumination"],
+        numerical_aperture_detection=_WRAP_KWARGS["numerical_aperture_detection"],
+        tilt_angle_zenith=_WRAP_KWARGS["tilt_angle_zenith"],
+        tilt_angle_azimuth=_WRAP_KWARGS["tilt_angle_azimuth"],
+        pupil_steepness=_WRAP_KWARGS["pupil_steepness"],
+    )
+    legacy_Hu, legacy_Hp = isotropic_thin_3d._calculate_wrap_unsafe_transfer_function(
+        **_WRAP_KWARGS
+    )
+    Hu, Hp = cache.transfer_functions(
+        _WRAP_KWARGS["z_position_list"],
+        invert_phase_contrast=_WRAP_KWARGS["invert_phase_contrast"],
+    )
+    assert torch.equal(legacy_Hu, Hu)
+    assert torch.equal(legacy_Hp, Hp)
+
+
+def test_cached_tilt_optics_reusable_across_z_iterations():
+    """Calling transfer_functions() repeatedly with different z lists works
+    and produces the same outputs as legacy fresh builds each time."""
+    cache = isotropic_thin_3d.CachedTiltOptics(
+        yx_shape=_WRAP_KWARGS["yx_shape"],
+        yx_pixel_size=_WRAP_KWARGS["yx_pixel_size"],
+        wavelength_illumination=_WRAP_KWARGS["wavelength_illumination"],
+        index_of_refraction_media=_WRAP_KWARGS["index_of_refraction_media"],
+        numerical_aperture_illumination=_WRAP_KWARGS["numerical_aperture_illumination"],
+        numerical_aperture_detection=_WRAP_KWARGS["numerical_aperture_detection"],
+        tilt_angle_zenith=_WRAP_KWARGS["tilt_angle_zenith"],
+        tilt_angle_azimuth=_WRAP_KWARGS["tilt_angle_azimuth"],
+        pupil_steepness=_WRAP_KWARGS["pupil_steepness"],
+    )
+    for z_list in ([-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-0.5, 0.0, 0.5]):
+        kw = {**_WRAP_KWARGS, "z_position_list": z_list}
+        legacy_Hu, legacy_Hp = isotropic_thin_3d._calculate_wrap_unsafe_transfer_function(**kw)
+        Hu, Hp = cache.transfer_functions(z_list)
+        assert torch.equal(legacy_Hu, Hu), f"abs TF mismatch at z={z_list}"
+        assert torch.equal(legacy_Hp, Hp), f"phase TF mismatch at z={z_list}"

tests/models/test_phase_thick_3d.py

@@ -133,3 +133,167 @@ def test_reconstruct():
 
     assert result.shape == zyx_shape
     assert np.all(np.isfinite(result.numpy()))
+
+
+_SHARED_OPTICS_KWARGS = dict(
+    zyx_shape=(20, 64, 64),
+    yx_pixel_size=6.5 / 40,
+    z_pixel_size=0.25,
+    wavelength_illumination=0.532,
+    z_padding=5,
+    index_of_refraction_media=1.33,
+    numerical_aperture_detection=1.2,
+    invert_phase_contrast=False,
+    pupil_steepness=1e4,
+)
+
+
+def test_compute_shared_optics_default_is_cpu():
+    """With no device kwarg, tensors land on CPU (back-compat)."""
+    tensors = phase_thick_3d._compute_shared_optics(**_SHARED_OPTICS_KWARGS)
+    for t in tensors:
+        assert t.device.type == "cpu"
+
+
+def test_compute_shared_optics_device_str_cpu():
+    """device='cpu' string is accepted and materializes on CPU."""
+    tensors = phase_thick_3d._compute_shared_optics(device="cpu", **_SHARED_OPTICS_KWARGS)
+    for t in tensors:
+        assert t.device.type == "cpu"
+
+
+def _pearson_complex(a: torch.Tensor, b: torch.Tensor) -> float:
+    """Pearson correlation over (Re, Im) concatenated and flattened."""
+    a_flat = torch.cat([a.real.flatten(), a.imag.flatten()]).double()
+    b_flat = torch.cat([b.real.flatten(), b.imag.flatten()]).double()
+    a_c = a_flat - a_flat.mean()
+    b_c = b_flat - b_flat.mean()
+    den = torch.sqrt((a_c ** 2).sum() * (b_c ** 2).sum())
+    if den.item() == 0:
+        # Constant tensor (e.g. pure pupil support); fall back to max-abs-diff check
+        return 1.0 if torch.allclose(a_flat, b_flat) else 0.0
+    return ((a_c * b_c).sum() / den).item()
+
+
+def test_angle_z_split_composes_to_shared_optics():
+    """The angle/z optics split composes back to bit-identical _compute_shared_optics output.
+
+    Validates that callers using the split helpers
+    (:func:`_compute_angle_optics` + :func:`_compute_z_optics`) for the
+    FREEZE_ANGLES tilt-recon recipe get the same numbers as the
+    legacy single-call path.
+    """
+    legacy = phase_thick_3d._compute_shared_optics(**_SHARED_OPTICS_KWARGS)
+    legacy_fyy, legacy_fxx, legacy_det_pupil, legacy_prop, legacy_green = legacy
+
+    fyy, fxx, radial_frequencies, det_pupil = phase_thick_3d._compute_angle_optics(
+        _SHARED_OPTICS_KWARGS["zyx_shape"][1:],
+        _SHARED_OPTICS_KWARGS["yx_pixel_size"],
+        _SHARED_OPTICS_KWARGS["wavelength_illumination"],
+        _SHARED_OPTICS_KWARGS["numerical_aperture_detection"],
+        pupil_steepness=_SHARED_OPTICS_KWARGS["pupil_steepness"],
+    )
+    z_position_list = phase_thick_3d._compute_z_position_list(
+        _SHARED_OPTICS_KWARGS["zyx_shape"][0],
+        _SHARED_OPTICS_KWARGS["z_pixel_size"],
+        _SHARED_OPTICS_KWARGS["z_padding"],
+        invert_phase_contrast=_SHARED_OPTICS_KWARGS["invert_phase_contrast"],
+    )
+    prop, green = phase_thick_3d._compute_z_optics(
+        radial_frequencies,
+        det_pupil,
+        z_position_list,
+        _SHARED_OPTICS_KWARGS["wavelength_illumination"],
+        _SHARED_OPTICS_KWARGS["index_of_refraction_media"],
+    )
+    assert torch.equal(legacy_fyy, fyy)
+    assert torch.equal(legacy_fxx, fxx)
+    assert torch.equal(legacy_det_pupil, det_pupil)
+    assert torch.equal(legacy_prop, prop)
+    assert torch.equal(legacy_green, green)
+
+
+def test_angle_optics_cached_across_z_changes():
+    """Angle optics tensors don't depend on z_pixel_size or z_padding.
+
+    Concrete check: build angle optics once, then build z optics with two
+    different z configurations and confirm the angle outputs are unchanged
+    (caller can hold them as a cache).
+    """
+    angle_kwargs = dict(
+        yx_shape=(64, 64),
+        yx_pixel_size=6.5 / 40,
+        wavelength_illumination=0.532,
+        numerical_aperture_detection=1.2,
+        pupil_steepness=1e4,
+    )
+    fyy_a, fxx_a, rf_a, det_a = phase_thick_3d._compute_angle_optics(**angle_kwargs)
+    fyy_b, fxx_b, rf_b, det_b = phase_thick_3d._compute_angle_optics(**angle_kwargs)
+    assert torch.equal(fyy_a, fyy_b)
+    assert torch.equal(fxx_a, fxx_b)
+    assert torch.equal(rf_a, rf_b)
+    assert torch.equal(det_a, det_b)
+
+    z_list_1 = phase_thick_3d._compute_z_position_list(20, 0.25, 5)
+    z_list_2 = phase_thick_3d._compute_z_position_list(20, 0.30, 5)
+    prop_1, green_1 = phase_thick_3d._compute_z_optics(
+        rf_a, det_a, z_list_1,
+        wavelength_illumination=0.532,
+        index_of_refraction_media=1.33,
+    )
+    prop_2, green_2 = phase_thick_3d._compute_z_optics(
+        rf_a, det_a, z_list_2,
+        wavelength_illumination=0.532,
+        index_of_refraction_media=1.33,
+    )
+    # Different z → different propagation kernels & Green's functions
+    assert not torch.equal(prop_1, prop_2)
+    assert not torch.equal(green_1, green_2)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_compute_shared_optics_cuda_matches_cpu():
+    """Building on CUDA must yield numerically equivalent tensors to CPU.
+
+    The change should be *mechanically equivalent* to the legacy CPU-build path
+    (no math changes — same generators, same constants), but CPU and CUDA do
+    not produce bit-identical floats for transcendentals like ``torch.exp``.
+    Strand C of the OPS tilt-recon work targets Pearson ≥ 0.999999 on the
+    derived transfer functions (see ``pattern_waveorder_gpu_shared_optics.md``).
+    Max-abs-diff on float32 is gated at the ~1e-4 level which corresponds
+    to the precision of CUDA's fast transcendentals.
+    """
+    cpu_tensors = phase_thick_3d._compute_shared_optics(device="cpu", **_SHARED_OPTICS_KWARGS)
+    cuda_tensors = phase_thick_3d._compute_shared_optics(device="cuda", **_SHARED_OPTICS_KWARGS)
+    names = ["fyy", "fxx", "det_pupil", "propagation_kernel", "greens_function_z"]
+    for name, cpu_t, cuda_t in zip(names, cpu_tensors, cuda_tensors):
+        assert cuda_t.device.type == "cuda"
+        cuda_on_cpu = cuda_t.cpu()
+        max_abs = (cpu_t - cuda_on_cpu).abs().max().item()
+        p = _pearson_complex(cpu_t, cuda_on_cpu) if cpu_t.is_complex() else _pearson_complex(
+            cpu_t.to(torch.complex64), cuda_on_cpu.to(torch.complex64)
+        )
+        # FP32 transcendental drift between CPU and CUDA is bounded; numerical
+        # equivalence is gated by Pearson, not bit-identicality.
+        assert max_abs < 1e-3, f"{name} max_abs_diff {max_abs:.3e} exceeds 1e-3"
+        assert p >= 0.999999, f"{name} Pearson {p:.9f} < 0.999999"
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_calculate_transfer_function_device_threading():
+    """When tilt angles arrive on CUDA, TFs come back on CUDA without a CPU detour."""
+    cuda = torch.device("cuda")
+    H_re, H_im = phase_thick_3d.calculate_transfer_function(
+        zyx_shape=(16, 64, 64),
+        yx_pixel_size=6.5 / 40,
+        z_pixel_size=0.25,
+        z_padding=4,
+        wavelength_illumination=0.532,
+        index_of_refraction_media=1.33,
+        numerical_aperture_illumination=0.9,
+        numerical_aperture_detection=1.2,
+        tilt_angle_zenith=torch.tensor(0.1, device=cuda),
+        tilt_angle_azimuth=torch.tensor(0.2, device=cuda),
+    )
+    assert H_re.device.type == "cuda"
+    assert H_im.device.type == "cuda"