Initial commit: he11lib mode-purity reconstruction library

Full implementation of Laguerre-Gauss modal reconstruction for gyrotron
beam diagnostics, per the approved design spec, plus tests, docs, and
a runnable end-to-end example.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tests/conftest.py

@@ -0,0 +1,3 @@
+import matplotlib
+
+matplotlib.use("Agg")

tests/test_data.py

@@ -0,0 +1,93 @@
+import numpy as np
+import pytest
+
+from he11lib.data import MeasurementPlane, ReconstructionResult, validate_planes
+
+
+def test_measurement_plane_stores_fields():
+    flux = np.ones((4, 4))
+    plane = MeasurementPlane(flux=flux, z=0.3)
+
+    assert plane.z == 0.3
+    assert np.array_equal(plane.flux, flux)
+    assert plane.pixel_scale is None
+    assert plane.viewing_angle_deg is None
+    assert plane.label is None
+
+
+def test_measurement_plane_stores_optional_fields():
+    flux = np.ones((4, 4))
+    plane = MeasurementPlane(
+        flux=flux, z=0.4, pixel_scale=0.1, viewing_angle_deg=5.0, label="plane_40cm"
+    )
+
+    assert plane.pixel_scale == 0.1
+    assert plane.viewing_angle_deg == 5.0
+    assert plane.label == "plane_40cm"
+
+
+def test_measurement_plane_rejects_non_2d_flux():
+    with pytest.raises(ValueError, match="2D"):
+        MeasurementPlane(flux=np.ones((4, 4, 3)), z=0.3)
+
+
+def test_measurement_plane_rejects_non_positive_z():
+    with pytest.raises(ValueError, match="positive"):
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.0)
+
+    with pytest.raises(ValueError, match="positive"):
+        MeasurementPlane(flux=np.ones((4, 4)), z=-0.1)
+
+
+def test_validate_planes_rejects_fewer_than_three():
+    planes = [
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.4),
+    ]
+    with pytest.raises(ValueError, match="[Aa]t least 3"):
+        validate_planes(planes)
+
+
+def test_validate_planes_rejects_mismatched_shapes():
+    planes = [
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
+        MeasurementPlane(flux=np.ones((5, 5)), z=0.4),
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
+    ]
+    with pytest.raises(ValueError, match="same shape"):
+        validate_planes(planes)
+
+
+def test_validate_planes_rejects_duplicate_z():
+    planes = [
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
+    ]
+    with pytest.raises(ValueError, match="distinct"):
+        validate_planes(planes)
+
+
+def test_validate_planes_accepts_valid_list():
+    planes = [
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.4),
+        MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
+    ]
+    validate_planes(planes)  # should not raise
+
+
+def test_reconstruction_result_stores_fields():
+    result = ReconstructionResult(
+        purity={(0, 0): (1.0, 0.0)},
+        reconstructed_field=np.ones((4, 4), dtype=complex),
+        centers=[(0.0, 0.0)],
+        pointing_angle_deg=0.0,
+        geometry={"pixel_scale": 0.1, "viewing_angle_deg": 2.0},
+        residuals=[np.zeros((4, 4))],
+        coefficient_uncertainty={(0, 0): 0.01},
+        used_phase_retrieval=False,
+    )
+
+    assert result.purity[(0, 0)] == (1.0, 0.0)
+    assert result.used_phase_retrieval is False

tests/test_deconvolution.py

@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+
+from he11lib.deconvolution import DiffusionDeconvolver
+
+
+def gaussian_bump(n, sigma_px):
+    coords = np.arange(n) - n // 2
+    xx, yy = np.meshgrid(coords, coords)
+    return np.exp(-(xx**2 + yy**2) / (2 * sigma_px**2))
+
+
+def profile_std(image):
+    n = image.shape[0]
+    coords = np.arange(n) - n // 2
+    weights = image[n // 2, :]
+    weights = np.clip(weights, 0, None)
+    mean = np.sum(coords * weights) / np.sum(weights)
+    var = np.sum(weights * (coords - mean) ** 2) / np.sum(weights)
+    return np.sqrt(var)
+
+
+def test_blur_sigma_m_from_diffusivity_and_dwell_time():
+    diffusivity = 1e-6  # m^2/s
+    dwell_time = 0.5  # s
+    deconvolver = DiffusionDeconvolver(thermal_diffusivity=diffusivity, dwell_time=dwell_time)
+
+    expected_sigma = np.sqrt(2 * diffusivity * dwell_time)
+    assert deconvolver.blur_sigma_m() == pytest.approx(expected_sigma)
+
+
+def test_blur_widens_a_sharp_peak():
+    deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.5)
+    pixel_scale = 2e-4
+    sharp = gaussian_bump(101, sigma_px=3)
+
+    blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale)
+
+    assert profile_std(blurred) > profile_std(sharp)
+
+
+def test_deconvolve_reduces_error_relative_to_blurred():
+    deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.3)
+    pixel_scale = 2e-4
+    sharp = gaussian_bump(101, sigma_px=4)
+
+    blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale)
+    deconvolved = deconvolver.deconvolve(blurred, pixel_scale=pixel_scale)
+
+    error_blurred = np.sum((blurred - sharp) ** 2)
+    error_deconvolved = np.sum((deconvolved - sharp) ** 2)
+    assert error_deconvolved < error_blurred
+
+
+def test_deconvolve_narrows_width_back_toward_original():
+    deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.3)
+    pixel_scale = 2e-4
+    sharp = gaussian_bump(101, sigma_px=4)
+
+    blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale)
+    deconvolved = deconvolver.deconvolve(blurred, pixel_scale=pixel_scale)
+
+    std_sharp = profile_std(sharp)
+    std_blurred = profile_std(blurred)
+    std_deconvolved = profile_std(deconvolved)
+
+    assert std_sharp < std_deconvolved < std_blurred

tests/test_fitting.py

@@ -0,0 +1,134 @@
+import numpy as np
+import pytest
+
+from he11lib.data import validate_planes
+from he11lib.fitting import ModalFitter, generate_mode_shells
+from he11lib.modes import LGBasis
+from he11lib.synthetic import SyntheticBeamGenerator
+
+W0 = 5e-3
+Z0 = 0.5
+WAVELENGTH = 1.76e-3
+PIXEL_SCALE = 4e-4
+IMAGE_SHAPE = (61, 61)
+Z_LIST = [0.35, 0.5, 0.65, 0.8]
+
+
+def make_basis():
+    return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+
+
+def make_generator(basis):
+    return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE)
+
+
+def test_generate_mode_shells_orders_by_2p_plus_abs_l():
+    shells = generate_mode_shells(max_order=2)
+    assert shells[0] == [(0, 0)]
+    assert set(shells[1]) == {(0, 1), (0, -1)}
+    assert set(shells[2]) == {(0, 2), (0, -2), (1, 0)}
+
+
+def test_fit_recovers_pure_fundamental_mode():
+    basis = make_basis()
+    gen = make_generator(basis)
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0
+    )
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit(planes, modes=[(0, 0)])
+
+    power_fraction, _ = result.purity[(0, 0)]
+    assert power_fraction == pytest.approx(1.0, abs=1e-6)
+    for cx, cy in result.centers:
+        assert cx == pytest.approx(0.0, abs=2 * PIXEL_SCALE)
+        assert cy == pytest.approx(0.0, abs=2 * PIXEL_SCALE)
+
+
+def test_fit_recovers_two_mode_purity_ratio():
+    basis = make_basis()
+    gen = make_generator(basis)
+    true_coeffs = {(0, 0): 0.9 + 0j, (1, 0): 0.3 + 0.1j}
+    planes = gen.generate(coefficients=true_coeffs, z_list=Z_LIST, noise_std=1e-4, seed=1)
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit(planes, modes=list(true_coeffs.keys()))
+
+    true_total = sum(abs(c) ** 2 for c in true_coeffs.values())
+    for mode, c in true_coeffs.items():
+        expected_fraction = abs(c) ** 2 / true_total
+        recovered_fraction, _ = result.purity[mode]
+        assert recovered_fraction == pytest.approx(expected_fraction, abs=0.03)
+
+
+def test_fit_recovers_center_offset():
+    basis = make_basis()
+    gen = make_generator(basis)
+    true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j},
+        z_list=Z_LIST,
+        center=true_center,
+        noise_std=1e-4,
+        seed=2,
+    )
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit(planes, modes=[(0, 0)], initial_center=true_center)
+
+    for cx, cy in result.centers:
+        assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE)
+        assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE)
+
+
+def test_fit_recovers_unknown_pixel_scale():
+    # Use a coarser pixel scale so the (much wider, far-field) beam at the
+    # outer z distances still fits within the frame -- otherwise pixel scale
+    # becomes unobservable from clipped images.
+    basis = make_basis()
+    local_pixel_scale = 1.5e-3
+    gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=local_pixel_scale)
+    planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=3)
+
+    # hide the known calibration to force the fitter to solve for it
+    for plane in planes:
+        plane.pixel_scale = None
+        plane.viewing_angle_deg = None
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit(
+        planes,
+        modes=[(0, 0)],
+        initial_pixel_scale=local_pixel_scale * 1.1,
+        initial_viewing_angle_deg=0.0,
+    )
+
+    fitted_scales = [result.geometry[f"pixel_scale_{i}"] for i in range(len(planes))]
+    for scale in fitted_scales:
+        assert scale == pytest.approx(local_pixel_scale, rel=0.05)
+
+
+def test_fit_auto_does_not_add_modes_for_pure_fundamental():
+    basis = make_basis()
+    gen = make_generator(basis)
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=4
+    )
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit_auto(planes, max_order=2)
+
+    assert set(result.purity.keys()) == {(0, 0)}
+
+
+def test_fit_auto_grows_to_include_second_mode():
+    basis = make_basis()
+    gen = make_generator(basis)
+    true_coeffs = {(0, 0): 0.9 + 0j, (0, 1): 0.4 + 0j}
+    planes = gen.generate(coefficients=true_coeffs, z_list=Z_LIST, noise_std=1e-4, seed=5)
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit_auto(planes, max_order=2)
+
+    assert (0, 1) in result.purity or (0, -1) in result.purity

tests/test_geometry.py

@@ -0,0 +1,77 @@
+import numpy as np
+import pytest
+
+from he11lib.data import MeasurementPlane
+from he11lib.geometry import GeometryCalibration
+
+
+def test_pixel_scale_known_reflects_plane():
+    plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, pixel_scale=1e-4)
+    plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)
+
+    assert GeometryCalibration(plane_known).pixel_scale_known is True
+    assert GeometryCalibration(plane_unknown).pixel_scale_known is False
+
+
+def test_viewing_angle_known_reflects_plane():
+    plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, viewing_angle_deg=10.0)
+    plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)
+
+    assert GeometryCalibration(plane_known).viewing_angle_known is True
+    assert GeometryCalibration(plane_unknown).viewing_angle_known is False
+
+
+def test_physical_coordinates_uses_known_calibration():
+    plane = MeasurementPlane(
+        flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=0.0
+    )
+    calib = GeometryCalibration(plane)
+    x, y = calib.physical_coordinates()
+
+    row_idx = np.arange(5) - 2
+    col_idx = np.arange(5) - 2
+    expected_x = col_idx * 2e-4
+    expected_y = row_idx * 2e-4
+    np.testing.assert_allclose(x[2, :], expected_x)
+    np.testing.assert_allclose(y[:, 2], expected_y)
+
+
+def test_physical_coordinates_compresses_x_for_viewing_angle():
+    plane = MeasurementPlane(
+        flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=60.0
+    )
+    calib = GeometryCalibration(plane)
+    x, y = calib.physical_coordinates()
+
+    col_idx = np.arange(5) - 2
+    expected_x = col_idx * 2e-4 / np.cos(np.deg2rad(60.0))
+    np.testing.assert_allclose(x[2, :], expected_x)
+
+
+def test_physical_coordinates_raises_without_calibration_or_override():
+    plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)
+    calib = GeometryCalibration(plane)
+
+    with pytest.raises(ValueError, match="pixel_scale"):
+        calib.physical_coordinates()
+
+
+def test_physical_coordinates_accepts_override_for_unknown_values():
+    plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)
+    calib = GeometryCalibration(plane)
+
+    x, y = calib.physical_coordinates(pixel_scale=1e-4, viewing_angle_deg=0.0)
+    col_idx = np.arange(5) - 2
+    np.testing.assert_allclose(x[2, :], col_idx * 1e-4)
+
+
+def test_known_calibration_takes_precedence_over_override():
+    plane = MeasurementPlane(
+        flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=0.0
+    )
+    calib = GeometryCalibration(plane)
+
+    # override should be ignored since plane already specifies calibration
+    x, _ = calib.physical_coordinates(pixel_scale=999.0, viewing_angle_deg=45.0)
+    col_idx = np.arange(5) - 2
+    np.testing.assert_allclose(x[2, :], col_idx * 2e-4)

tests/test_modes.py

@@ -0,0 +1,103 @@
+import numpy as np
+import pytest
+
+from he11lib.modes import LGBasis
+
+
+W0 = 5e-3  # 5 mm waist
+Z0 = 0.5  # waist at 0.5 m
+WAVELENGTH = 1.76e-3  # ~170 GHz gyrotron
+
+
+def make_grid(w, half_widths_in_w=6.0, n=300):
+    extent = half_widths_in_w * w
+    coords = np.linspace(-extent, extent, n)
+    dx = coords[1] - coords[0]
+    x, y = np.meshgrid(coords, coords)
+    return x, y, dx
+
+
+def test_beam_radius_at_waist_equals_w0():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    assert basis.beam_radius(Z0) == pytest.approx(W0)
+
+
+def test_beam_radius_at_one_rayleigh_range():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    zR = np.pi * W0**2 / WAVELENGTH
+    assert basis.beam_radius(Z0 + zR) == pytest.approx(W0 * np.sqrt(2))
+
+
+def test_gouy_phase_zero_at_waist():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    assert basis.gouy_phase(Z0, p=0, l=0) == pytest.approx(0.0)
+    assert basis.gouy_phase(Z0, p=1, l=2) == pytest.approx(0.0)
+
+
+def test_gouy_phase_at_one_rayleigh_range_for_fundamental():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    zR = np.pi * W0**2 / WAVELENGTH
+    # order (2p+|l|+1) = 1 for p=0,l=0; atan(1) = pi/4
+    assert basis.gouy_phase(Z0 + zR, p=0, l=0) == pytest.approx(np.pi / 4)
+
+
+def test_fundamental_mode_matches_analytic_gaussian_at_waist():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    x, y, _ = make_grid(W0, n=50)
+    r2 = x**2 + y**2
+
+    field = basis.field(x, y, Z0, p=0, l=0)
+    expected_intensity_shape = np.exp(-2 * r2 / W0**2)
+    intensity = np.abs(field) ** 2
+
+    # shapes should match up to a constant normalization factor
+    ratio = intensity / expected_intensity_shape
+    ratio_center = ratio[len(ratio) // 2, len(ratio) // 2]
+    np.testing.assert_allclose(ratio / ratio_center, 1.0, atol=1e-6)
+
+
+def test_mode_is_normalized_at_waist():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    x, y, dx = make_grid(W0, n=400)
+
+    field = basis.field(x, y, Z0, p=0, l=0)
+    total_power = np.sum(np.abs(field) ** 2) * dx**2
+    assert total_power == pytest.approx(1.0, rel=2e-3)
+
+
+def test_mode_is_normalized_away_from_waist():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    z = Z0 + 0.2
+    w_z = basis.beam_radius(z)
+    x, y, dx = make_grid(w_z, n=400)
+
+    field = basis.field(x, y, z, p=1, l=2)
+    total_power = np.sum(np.abs(field) ** 2) * dx**2
+    assert total_power == pytest.approx(1.0, rel=2e-3)
+
+
+def test_modes_are_orthogonal():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    z = Z0 + 0.1
+    w_z = basis.beam_radius(z)
+    x, y, dx = make_grid(w_z, n=400)
+
+    field_a = basis.field(x, y, z, p=0, l=0)
+    field_b = basis.field(x, y, z, p=1, l=0)
+    inner_product = np.sum(field_a * np.conj(field_b)) * dx**2
+    assert abs(inner_product) < 1e-3
+
+
+def test_project_recovers_known_coefficients():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    z = Z0 + 0.15
+    w_z = basis.beam_radius(z)
+    x, y, dx = make_grid(w_z, n=400)
+
+    true_coeffs = {(0, 0): 0.8 + 0.1j, (1, 0): 0.2 - 0.3j}
+    field = basis.field_superposition(x, y, z, true_coeffs)
+
+    recovered = basis.project(field, x, y, dx, z, modes=list(true_coeffs.keys()))
+
+    for mode, coeff in true_coeffs.items():
+        assert recovered[mode] == pytest.approx(coeff, abs=5e-3)

tests/test_noise.py

@@ -0,0 +1,46 @@
+import numpy as np
+import pytest
+
+from he11lib.noise import NoiseEstimator
+
+
+def test_estimate_std_recovers_known_noise_on_flat_image():
+    rng = np.random.default_rng(0)
+    true_std = 0.05
+    image = np.ones((200, 200)) * 10.0 + rng.normal(0, true_std, size=(200, 200))
+
+    estimated = NoiseEstimator().estimate_std(image)
+    assert estimated == pytest.approx(true_std, rel=0.15)
+
+
+def test_estimate_std_recovers_known_noise_on_smooth_bump():
+    rng = np.random.default_rng(1)
+    x = np.linspace(-3, 3, 200)
+    xx, yy = np.meshgrid(x, x)
+    smooth = np.exp(-(xx**2 + yy**2))
+    true_std = 0.01
+    image = smooth + rng.normal(0, true_std, size=smooth.shape)
+
+    estimated = NoiseEstimator().estimate_std(image)
+    assert estimated == pytest.approx(true_std, rel=0.35)
+
+
+def test_estimate_std_near_zero_for_noise_free_image():
+    x = np.linspace(-3, 3, 100)
+    xx, yy = np.meshgrid(x, x)
+    smooth = np.exp(-(xx**2 + yy**2))
+
+    estimated = NoiseEstimator().estimate_std(smooth)
+    assert estimated < 1e-6
+
+
+def test_weights_are_uniform_and_match_shape():
+    rng = np.random.default_rng(2)
+    image = np.ones((50, 50)) * 5.0 + rng.normal(0, 0.1, size=(50, 50))
+
+    weights = NoiseEstimator().weights(image)
+    assert weights.shape == image.shape
+    assert np.allclose(weights, weights.flat[0])
+
+    expected_std = NoiseEstimator().estimate_std(image)
+    assert weights.flat[0] == pytest.approx(1.0 / expected_std**2, rel=1e-6)

tests/test_phase_retrieval.py

@@ -0,0 +1,84 @@
+import numpy as np
+import pytest
+
+from he11lib.modes import LGBasis
+from he11lib.phase_retrieval import PhaseRetriever, propagate_angular_spectrum
+from he11lib.synthetic import SyntheticBeamGenerator
+
+W0 = 5e-3
+Z0 = 0.5
+WAVELENGTH = 1.76e-3
+PIXEL_SCALE = 3e-4
+IMAGE_SHAPE = (121, 121)
+
+
+def make_basis():
+    return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+
+
+def make_grid():
+    coords = (np.arange(IMAGE_SHAPE[0]) - IMAGE_SHAPE[0] // 2) * PIXEL_SCALE
+    x, y = np.meshgrid(coords, coords)
+    return x, y
+
+
+def test_propagate_round_trip_recovers_original_field():
+    basis = make_basis()
+    x, y = make_grid()
+    field = basis.field(x, y, Z0, p=0, l=0)
+
+    forward = propagate_angular_spectrum(field, PIXEL_SCALE, dz=0.05, wavelength=WAVELENGTH)
+    back = propagate_angular_spectrum(forward, PIXEL_SCALE, dz=-0.05, wavelength=WAVELENGTH)
+
+    np.testing.assert_allclose(back, field, atol=1e-3 * np.max(np.abs(field)))
+
+
+def test_propagate_matches_lgbasis_analytic_evolution():
+    basis = make_basis()
+    x, y = make_grid()
+    field_at_waist = basis.field(x, y, Z0, p=0, l=0)
+
+    dz = 0.05
+    propagated = propagate_angular_spectrum(field_at_waist, PIXEL_SCALE, dz=dz, wavelength=WAVELENGTH)
+    analytic = basis.field(x, y, Z0 + dz, p=0, l=0)
+
+    # compare intensity profiles (phase reference/global constant may differ)
+    np.testing.assert_allclose(
+        np.abs(propagated) ** 2, np.abs(analytic) ** 2, atol=1e-2 * np.max(np.abs(analytic) ** 2)
+    )
+
+
+def test_retrieve_recovers_pure_mode_purity():
+    # Keep z distances close to the waist so the (widening) beam stays well
+    # within the frame -- otherwise FFT wraparound/clipping at the edges
+    # degrades angular-spectrum propagation accuracy.
+    basis = make_basis()
+    gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE)
+    z_list = [0.47, 0.5, 0.53]
+    planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=0)
+
+    retriever = PhaseRetriever(wavelength=WAVELENGTH)
+    result = retriever.retrieve(planes, viewing_angle_deg=0.0, max_iterations=100)
+
+    coeffs = basis.project(
+        result.field, result.x, result.y, PIXEL_SCALE, result.z, modes=[(0, 0), (1, 0), (0, 1)]
+    )
+    total_power = sum(abs(c) ** 2 for c in coeffs.values())
+    purity_00 = abs(coeffs[(0, 0)]) ** 2 / total_power
+    assert purity_00 > 0.9
+
+
+def test_retrieve_estimates_beam_center():
+    basis = make_basis()
+    gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE)
+    z_list = [0.47, 0.5, 0.53]
+    true_center = (15 * PIXEL_SCALE, -8 * PIXEL_SCALE)
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, center=true_center, noise_std=1e-5, seed=1
+    )
+
+    retriever = PhaseRetriever(wavelength=WAVELENGTH)
+    result = retriever.retrieve(planes, viewing_angle_deg=0.0, max_iterations=100)
+
+    assert result.center[0] == pytest.approx(true_center[0], abs=3 * PIXEL_SCALE)
+    assert result.center[1] == pytest.approx(true_center[1], abs=3 * PIXEL_SCALE)

tests/test_plotting.py

@@ -0,0 +1,64 @@
+import matplotlib.figure
+import numpy as np
+import pytest
+
+from he11lib.data import MeasurementPlane, ReconstructionResult
+from he11lib.plotting import plot_center_trace, plot_mode_purity, plot_residuals
+
+
+def make_result(**overrides):
+    defaults = dict(
+        purity={(0, 0): (0.9, 0.1), (1, 0): (0.1, -0.2)},
+        reconstructed_field=np.zeros((5, 5), dtype=complex),
+        centers=[(0.0, 0.0), (1e-4, -1e-4), (2e-4, -2e-4)],
+        pointing_angle_deg=0.5,
+        residuals=[np.ones((5, 5)), np.ones((5, 5)) * 2, np.ones((5, 5)) * 3],
+    )
+    defaults.update(overrides)
+    return ReconstructionResult(**defaults)
+
+
+def make_planes():
+    return [
+        MeasurementPlane(flux=np.zeros((5, 5)), z=0.3),
+        MeasurementPlane(flux=np.zeros((5, 5)), z=0.5),
+        MeasurementPlane(flux=np.zeros((5, 5)), z=0.7),
+    ]
+
+
+def test_plot_mode_purity_draws_one_bar_per_mode():
+    result = make_result()
+    fig = plot_mode_purity(result)
+
+    assert isinstance(fig, matplotlib.figure.Figure)
+    ax = fig.axes[0]
+    assert len(ax.patches) == len(result.purity)
+
+
+def test_plot_center_trace_plots_one_point_per_plane():
+    planes = make_planes()
+    result = make_result()
+    fig = plot_center_trace(planes, result)
+
+    assert isinstance(fig, matplotlib.figure.Figure)
+    ax = fig.axes[0]
+    line = ax.lines[0]
+    assert len(line.get_xdata()) == len(planes)
+
+
+def test_plot_residuals_draws_one_axes_per_plane():
+    planes = make_planes()
+    result = make_result()
+    fig = plot_residuals(planes, result)
+
+    assert isinstance(fig, matplotlib.figure.Figure)
+    image_axes = [ax for ax in fig.axes if ax.images]
+    assert len(image_axes) == len(planes)
+
+
+def test_plot_residuals_raises_when_phase_retrieval_used_without_residuals():
+    planes = make_planes()
+    result = make_result(residuals=[], used_phase_retrieval=True)
+
+    with pytest.raises(ValueError, match="residuals"):
+        plot_residuals(planes, result)

tests/test_reconstruct.py

@@ -0,0 +1,112 @@
+from dataclasses import replace
+
+import pytest
+
+from he11lib.deconvolution import DiffusionDeconvolver
+from he11lib.fitting import ModalFitter
+from he11lib.modes import LGBasis
+from he11lib.reconstruct import BeamReconstructor
+from he11lib.synthetic import SyntheticBeamGenerator
+
+W0 = 5e-3
+Z0 = 0.5
+WAVELENGTH = 1.76e-3
+PIXEL_SCALE = 4e-4
+IMAGE_SHAPE = (61, 61)
+Z_LIST = [0.35, 0.5, 0.65, 0.8]
+
+
+def make_basis():
+    return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+
+
+def make_generator(basis):
+    return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE)
+
+
+def test_reconstruct_recovers_pure_mode_purity():
+    basis = make_basis()
+    gen = make_generator(basis)
+    planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0)
+
+    reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2)
+    result = reconstructor.reconstruct(planes)
+
+    power_fraction, _ = result.purity[(0, 0)]
+    assert power_fraction == pytest.approx(1.0, abs=1e-3)
+    assert result.used_phase_retrieval is False
+
+
+def test_reconstruct_recovers_center_offset():
+    basis = make_basis()
+    gen = make_generator(basis)
+    true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, center=true_center, noise_std=1e-4, seed=1
+    )
+
+    reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2)
+    result = reconstructor.reconstruct(planes)
+
+    for cx, cy in result.centers:
+        assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE)
+        assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE)
+
+
+def test_reconstruct_with_deconvolution_corrects_blur():
+    basis = make_basis()
+    gen = make_generator(basis)
+    planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=2)
+
+    deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=30.0)
+    blurred_planes = [
+        replace(p, flux=deconvolver.blur(p.flux, p.pixel_scale)) for p in planes
+    ]
+
+    # Without deconvolution, blur should measurably hurt purity recovery.
+    fitter = ModalFitter(basis)
+    result_no_deconv = fitter.fit(blurred_planes, modes=[(0, 0), (1, 0), (0, 1)])
+    purity_no_deconv, _ = result_no_deconv.purity[(0, 0)]
+
+    reconstructor = BeamReconstructor(
+        w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, deconvolver=deconvolver
+    )
+    result = reconstructor.reconstruct(blurred_planes)
+    purity_with_deconv, _ = result.purity[(0, 0)]
+
+    assert purity_with_deconv > purity_no_deconv
+    assert purity_with_deconv > 0.9
+
+
+def test_reconstruct_forces_phase_retrieval_fallback():
+    basis = make_basis()
+    gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4)
+    z_list = [0.47, 0.5, 0.53]
+    planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=3)
+
+    reconstructor = BeamReconstructor(
+        w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, force_phase_retrieval=True
+    )
+    result = reconstructor.reconstruct(planes)
+
+    assert result.used_phase_retrieval is True
+    power_fraction, _ = result.purity[(0, 0)]
+    assert power_fraction > 0.9
+
+
+def test_reconstruct_falls_back_automatically_on_high_residual():
+    basis = make_basis()
+    gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4)
+    z_list = [0.47, 0.5, 0.53]
+    planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=4)
+
+    reconstructor = BeamReconstructor(
+        w0=W0,
+        z0=Z0,
+        wavelength=WAVELENGTH,
+        max_order=2,
+        phase_retrieval_residual_threshold=1e-8,
+    )
+    result = reconstructor.reconstruct(planes)
+
+    assert result.used_phase_retrieval is True

tests/test_synthetic.py

@@ -0,0 +1,123 @@
+import numpy as np
+import pytest
+
+from he11lib.modes import LGBasis
+from he11lib.synthetic import SyntheticBeamGenerator
+
+
+W0 = 5e-3
+Z0 = 0.5
+WAVELENGTH = 1.76e-3
+PIXEL_SCALE = 2e-4  # 0.2 mm/px
+IMAGE_SHAPE = (161, 161)  # odd so there's a well-defined center pixel
+
+
+def make_generator():
+    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
+    return SyntheticBeamGenerator(
+        basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE
+    )
+
+
+def test_generate_returns_planes_with_requested_z():
+    gen = make_generator()
+    z_list = [0.3, 0.4, 0.5]
+    planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=z_list)
+
+    assert [p.z for p in planes] == z_list
+    assert all(p.flux.shape == IMAGE_SHAPE for p in planes)
+
+
+def test_generate_pure_mode_peak_at_image_center_when_centered():
+    gen = make_generator()
+    planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(0.0, 0.0))
+    flux = planes[0].flux
+
+    peak_idx = np.unravel_index(np.argmax(flux), flux.shape)
+    center_idx = (IMAGE_SHAPE[0] // 2, IMAGE_SHAPE[1] // 2)
+    assert peak_idx == center_idx
+
+
+def test_generate_applies_center_offset():
+    gen = make_generator()
+    offset_m = 20 * PIXEL_SCALE  # 20 pixels
+    planes = gen.generate(
+        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(offset_m, 0.0)
+    )
+    flux = planes[0].flux
+
+    peak_idx = np.unravel_index(np.argmax(flux), flux.shape)
+    center_row = IMAGE_SHAPE[0] // 2
+    center_col = IMAGE_SHAPE[1] // 2
+    assert peak_idx[0] == center_row
+    assert peak_idx[1] == pytest.approx(center_col + 20, abs=1)
+
+
+def test_generate_applies_pointing_angle_as_linear_drift():
+    gen = make_generator()
+    pointing_angle_deg = 1.0  # small tilt
+    z_list = [Z0, Z0 + 0.2]
+    planes = gen.generate(
+        coefficients={(0, 0): 1 + 0j},
+        z_list=z_list,
+        center=(0.0, 0.0),
+        pointing_angle_deg=pointing_angle_deg,
+    )
+
+    peaks_col = []
+    for plane in planes:
+        peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape)
+        peaks_col.append(peak_idx[1])
+
+    expected_shift_m = 0.2 * np.tan(np.deg2rad(pointing_angle_deg))
+    expected_shift_px = expected_shift_m / PIXEL_SCALE
+    actual_shift_px = peaks_col[1] - peaks_col[0]
+    assert actual_shift_px == pytest.approx(expected_shift_px, abs=1)
+
+
+def test_generate_noise_is_reproducible_with_seed():
+    gen = make_generator()
+    planes_a = gen.generate(
+        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42
+    )
+    planes_b = gen.generate(
+        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42
+    )
+    np.testing.assert_array_equal(planes_a[0].flux, planes_b[0].flux)
+
+
+def test_generate_noise_std_matches_requested_level():
+    gen = make_generator()
+    noise_std = 0.02
+    planes_noisy = gen.generate(
+        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=noise_std, seed=1
+    )
+    planes_clean = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.0)
+
+    diff = planes_noisy[0].flux - planes_clean[0].flux
+    assert np.std(diff) == pytest.approx(noise_std, rel=0.15)
+
+
+def test_generate_viewing_angle_compresses_tilt_axis():
+    gen = make_generator()
+    planes_straight = gen.generate(
+        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=0.0
+    )
+    planes_tilted = gen.generate(
+        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=60.0
+    )
+
+    def width_along_axis(flux, axis):
+        profile = flux[flux.shape[0] // 2, :] if axis == 1 else flux[:, flux.shape[1] // 2]
+        half_max = profile.max() / 2
+        above = np.where(profile >= half_max)[0]
+        return above[-1] - above[0]
+
+    width_straight_x = width_along_axis(planes_straight[0].flux, axis=1)
+    width_tilted_x = width_along_axis(planes_tilted[0].flux, axis=1)
+    width_straight_y = width_along_axis(planes_straight[0].flux, axis=0)
+    width_tilted_y = width_along_axis(planes_tilted[0].flux, axis=0)
+
+    # tilt compresses the viewed beam along the tilt (x) axis, y unaffected
+    assert width_tilted_x < width_straight_x
+    assert width_tilted_y == pytest.approx(width_straight_y, abs=1)