he11lib/tests/test_reconstruct.py

from dataclasses import replace

import pytest

from he11lib.deconvolution import DiffusionDeconvolver
from he11lib.fitting import ModalFitter
from he11lib.geometry import CameraModel, CameraModelTolerance, GeometryCalibration
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
CAMERA_DISTANCE = 5.0
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_camera():
    focal_length_px = (CAMERA_DISTANCE + Z0) / PIXEL_SCALE
    return CameraModel(
        focal_length_px=focal_length_px,
        position=(0.0, 0.0, -CAMERA_DISTANCE),
        orientation_deg=(0.0, 0.0, 0.0),
    )


def zero_tolerance():
    return CameraModelTolerance(
        focal_length_px=0.0, position=(0.0, 0.0, 0.0), orientation_deg=(0.0, 0.0, 0.0)
    )


def make_generator(basis, camera):
    return SyntheticBeamGenerator(basis=basis, camera=camera)


def test_reconstruct_recovers_pure_mode_purity():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=0
    )

    reconstructor = BeamReconstructor(
        w0=W0, z0=Z0, wavelength=WAVELENGTH, camera=camera, camera_tolerance=zero_tolerance(), 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()
    camera = make_camera()
    gen = make_generator(basis, camera)
    true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j},
        z_list=Z_LIST,
        image_shape=IMAGE_SHAPE,
        center=true_center,
        noise_std=1e-4,
        seed=1,
    )

    reconstructor = BeamReconstructor(
        w0=W0, z0=Z0, wavelength=WAVELENGTH, camera=camera, camera_tolerance=zero_tolerance(), 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()
    camera = make_camera()
    gen = make_generator(basis, camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=2
    )

    deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=30.0)
    calib = GeometryCalibration(camera)
    blurred_planes = [
        replace(p, flux=deconvolver.blur(p.flux, calib.effective_pixel_scale(p.flux.shape, p.z)))
        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)], camera=camera, camera_tolerance=zero_tolerance()
    )
    purity_no_deconv, _ = result_no_deconv.purity[(0, 0)]

    reconstructor = BeamReconstructor(
        w0=W0,
        z0=Z0,
        wavelength=WAVELENGTH,
        camera=camera,
        camera_tolerance=zero_tolerance(),
        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()
    camera = make_camera()
    gen = SyntheticBeamGenerator(basis=basis, camera=camera)
    z_list = [0.47, 0.5, 0.53]
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE, noise_std=1e-5, seed=3
    )

    reconstructor = BeamReconstructor(
        w0=W0,
        z0=Z0,
        wavelength=WAVELENGTH,
        camera=camera,
        camera_tolerance=zero_tolerance(),
        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()
    camera = make_camera()
    gen = SyntheticBeamGenerator(basis=basis, camera=camera)
    z_list = [0.47, 0.5, 0.53]
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE, noise_std=1e-5, seed=4
    )

    reconstructor = BeamReconstructor(
        w0=W0,
        z0=Z0,
        wavelength=WAVELENGTH,
        camera=camera,
        camera_tolerance=zero_tolerance(),
        max_order=2,
        phase_retrieval_residual_threshold=1e-8,
    )
    result = reconstructor.reconstruct(planes)

    assert result.used_phase_retrieval is True


def test_reconstruct_recovers_camera_and_z_offset_from_nominal():
    # End-to-end: ground truth is offset from the nominal camera/z inputs
    # (within their tolerances), simulating realistic calibration error.
    basis = make_basis()
    true_camera = make_camera()
    gen = make_generator(basis, true_camera)
    true_z_list = Z_LIST
    z_offsets = {z: 0.01 for z in true_z_list}
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j},
        z_list=true_z_list,
        image_shape=IMAGE_SHAPE,
        nominal_z_offsets=z_offsets,
        z_tolerance=0.03,
        pointing_angle_horizontal_deg=0.2,
        pointing_angle_vertical_deg=-0.1,
        noise_std=1e-4,
        seed=13,
    )

    nominal_focal_offset = true_camera.focal_length_px * 0.03
    nominal_camera = CameraModel(
        focal_length_px=true_camera.focal_length_px + nominal_focal_offset,
        position=true_camera.position,
        orientation_deg=true_camera.orientation_deg,
    )
    tolerance = CameraModelTolerance(
        focal_length_px=true_camera.focal_length_px * 0.1,
        position=(0.0, 0.0, 0.0),
        orientation_deg=(0.0, 0.0, 0.0),
    )

    reconstructor = BeamReconstructor(
        w0=W0,
        z0=Z0,
        wavelength=WAVELENGTH,
        camera=nominal_camera,
        camera_tolerance=tolerance,
        max_order=1,
    )
    result = reconstructor.reconstruct(planes)

    power_fraction, _ = result.purity[(0, 0)]
    assert power_fraction > 0.95
    assert result.pointing_angle_horizontal_deg == pytest.approx(0.2, abs=0.1)
    assert result.pointing_angle_vertical_deg == pytest.approx(-0.1, abs=0.1)
    assert result.geometry["focal_length_px"] == pytest.approx(true_camera.focal_length_px, rel=0.03)
    for i, true_z in enumerate(true_z_list):
        assert result.geometry[f"z_{i}"] == pytest.approx(true_z, abs=0.005)