Blame for tests/test_simulator.py

"""Known-answer tests for the beam simulator."""

from dataclasses import replace

import numpy as np
import pytest

from beampilot.simulator import (
    Aperture,
    BeamParams,
    CameraConfig,
    Ghost,
    HotPixels,
    beam_diameter,
    dark_frame,
    rayleigh_length,
    render_frame,
)
from beampilot.widths import Roi, measure_widths, second_moments

# Astigmatic test beam: different waist sizes, positions and M² per axis.
BEAM = BeamParams(
    wavelength_m=1064e-9,
    d0x_m=0.5e-3,
    d0y_m=0.4e-3,
    z0x_m=0.10,
    z0y_m=0.12,
    m2x=1.2,
    m2y=1.5,
)
CAMERA = CameraConfig(
    pixel_size_m=5.5e-6,
    nx=512,
    ny=512,
    bit_depth=12,
    offset_counts=100.0,
    read_noise_counts=5.0,
    signal_counts_per_ms=1e6,
)
NOISELESS = replace(CAMERA, offset_counts=0.0, read_noise_counts=0.0, shot_noise=False)


def _centroid(frame: np.ndarray) -> tuple[float, float]:
    """Return the centroid (x_c, y_c) in pixels of an offset-free frame."""
    image = frame.astype(float)
    y, x = np.indices(image.shape)
    total = image.sum()
    return (image * x).sum() / total, (image * y).sum() / total


def _second_moment_diameters(

    frame: np.ndarray, pixel_size_m: float

) -> tuple[float, float]:
    """Return 4-sigma diameters (d_x, d_y) in m of an offset-free frame."""
    image = frame.astype(float)
    y, x = np.indices(image.shape)
    total = image.sum()
    x_c = (image * x).sum() / total
    y_c = (image * y).sum() / total
    var_x = (image * (x - x_c) ** 2).sum() / total
    var_y = (image * (y - y_c) ** 2).sum() / total
    return 4.0 * np.sqrt(var_x) * pixel_size_m, 4.0 * np.sqrt(var_y) * pixel_size_m


def test_diameter_at_waist_equals_d0() -> None:
    d = beam_diameter(BEAM.z0x_m, BEAM.d0x_m, BEAM.z0x_m, BEAM.m2x, BEAM.wavelength_m)
    assert d == pytest.approx(BEAM.d0x_m, rel=1e-12)


def test_diameter_one_rayleigh_length_from_waist() -> None:
    z_r = rayleigh_length(BEAM.d0x_m, BEAM.m2x, BEAM.wavelength_m)
    d = beam_diameter(
        BEAM.z0x_m + z_r, BEAM.d0x_m, BEAM.z0x_m, BEAM.m2x, BEAM.wavelength_m
    )
    assert d == pytest.approx(BEAM.d0x_m * np.sqrt(2.0), rel=1e-12)


def test_noiseless_frame_matches_diameters() -> None:
    z_m = 0.25  # away from both waists, so d_x and d_y differ clearly
    frame = render_frame(BEAM, NOISELESS, z_m, 20.0, np.random.default_rng(0))
    assert frame.max() < NOISELESS.full_scale_counts  # unsaturated

    d_x, d_y = _second_moment_diameters(frame, NOISELESS.pixel_size_m)
    expected_x = beam_diameter(z_m, BEAM.d0x_m, BEAM.z0x_m, BEAM.m2x, BEAM.wavelength_m)
    expected_y = beam_diameter(z_m, BEAM.d0y_m, BEAM.z0y_m, BEAM.m2y, BEAM.wavelength_m)
    assert d_x == pytest.approx(expected_x, rel=0.002)
    assert d_y == pytest.approx(expected_y, rel=0.002)


def test_frame_sum_is_proportional_to_exposure() -> None:
    z_m = 0.25  # beam edge is > 7 sigma from the sensor edge: nothing truncated
    single = render_frame(BEAM, NOISELESS, z_m, 5.0, np.random.default_rng(0))
    double = render_frame(BEAM, NOISELESS, z_m, 10.0, np.random.default_rng(0))
    assert double.max() < NOISELESS.full_scale_counts  # both unsaturated

    expected = NOISELESS.signal_counts_per_ms * 5.0
    assert single.sum() == pytest.approx(expected, rel=0.005)
    assert double.sum() == pytest.approx(2.0 * single.sum(), rel=0.005)


def test_beam_position_moves_centroid() -> None:
    x0_m, y0_m = 200e-6, -100e-6
    rng = np.random.default_rng(0)
    centred = render_frame(BEAM, NOISELESS, 0.25, 20.0, rng)
    shifted = render_frame(BEAM, NOISELESS, 0.25, 20.0, rng, x0_m=x0_m, y0_m=y0_m)

    (xc0, yc0), (xc1, yc1) = _centroid(centred), _centroid(shifted)
    assert xc0 == pytest.approx((NOISELESS.nx - 1) / 2.0, abs=0.01)
    assert yc0 == pytest.approx((NOISELESS.ny - 1) / 2.0, abs=0.01)
    assert xc1 - xc0 == pytest.approx(x0_m / NOISELESS.pixel_size_m, abs=0.01)
    assert yc1 - yc0 == pytest.approx(y0_m / NOISELESS.pixel_size_m, abs=0.01)


def test_saturation_depends_on_exposure_and_position() -> None:
    full_scale = CAMERA.full_scale_counts
    near_waist = BEAM.z0x_m
    far = 0.60

    high = render_frame(BEAM, CAMERA, near_waist, 100.0, np.random.default_rng(1))
    low = render_frame(BEAM, CAMERA, near_waist, 1.0, np.random.default_rng(1))
    high_far = render_frame(BEAM, CAMERA, far, 100.0, np.random.default_rng(1))

    assert (high == full_scale).any()
    assert not (low == full_scale).any()
    assert not (high_far == full_scale).any()


def test_shot_noise_variance_matches_signal() -> None:
    # A 50 mm beam is flat to < 0.01 % over the central 64 x 64 pixels.
    flat_beam = replace(
        BEAM, d0x_m=50e-3, d0y_m=50e-3, z0x_m=0.0, z0y_m=0.0, m2x=1.0, m2y=1.0
    )
    bright = replace(CAMERA, signal_counts_per_ms=6.5e10)  # ~2000 counts per pixel
    patch = np.s_[224:288, 224:288]
    rng = np.random.default_rng(5)

    with_shot = render_frame(flat_beam, bright, 0.0, 1.0, rng)[patch]
    no_shot_camera = replace(bright, shot_noise=False)
    without_shot = render_frame(flat_beam, no_shot_camera, 0.0, 1.0, rng)[patch]

    signal = with_shot.mean() - bright.offset_counts
    read_variance = bright.read_noise_counts**2
    assert with_shot.var() == pytest.approx(signal + read_variance, rel=0.1)
    assert without_shot.var() == pytest.approx(read_variance, rel=0.1)


def test_same_seed_gives_identical_frame() -> None:
    a = render_frame(BEAM, CAMERA, 0.2, 5.0, np.random.default_rng(42))
    b = render_frame(BEAM, CAMERA, 0.2, 5.0, np.random.default_rng(42))
    c = render_frame(BEAM, CAMERA, 0.2, 5.0, np.random.default_rng(43))
    np.testing.assert_array_equal(a, b)
    assert not np.array_equal(a, c)


@pytest.mark.parametrize(

    ("field", "value"),

    [

        ("wavelength_m", 0.0),

        ("d0x_m", 0.0),

        ("d0y_m", -1e-3),

        ("m2x", 0.99),

        ("m2y", float("nan")),

    ],

)
def test_invalid_beam_params_raise(field: str, value: float) -> None:
    with pytest.raises(ValueError, match=field):
        replace(BEAM, **{field: value})


def test_m2_of_one_is_valid() -> None:
    beam = replace(BEAM, m2x=1.0, m2y=1.0)
    assert beam.m2x == beam.m2y == 1.0


@pytest.mark.parametrize(

    ("field", "value"),

    [("pixel_size_m", 0.0), ("pixel_size_m", -5e-6), ("bit_depth", 0)],

)
def test_invalid_camera_config_raises(field: str, value: float) -> None:
    with pytest.raises(ValueError, match=field):
        replace(CAMERA, **{field: value})


@pytest.mark.parametrize("exposure_ms", [0.0, -1.0])
def test_non_positive_exposure_raises(exposure_ms: float) -> None:
    with pytest.raises(ValueError, match="exposure_ms"):
        render_frame(BEAM, CAMERA, 0.2, exposure_ms, np.random.default_rng(0))


def test_dark_frame_has_offset_and_read_noise() -> None:
    frame = dark_frame(CAMERA, np.random.default_rng(3))
    assert frame.shape == (CAMERA.ny, CAMERA.nx)
    assert frame.mean() == pytest.approx(CAMERA.offset_counts, abs=0.1)
    # Rounding adds a variance of 1/12 count², negligible against 5 counts.
    assert frame.std() == pytest.approx(CAMERA.read_noise_counts, rel=0.02)


# Ghost reflection


GHOST_Z_M = 0.25  # d_x ~ 698 um, d_y ~ 772 um


def _frame(ghost=None, z_m=GHOST_Z_M) -> np.ndarray:
    return render_frame(
        BEAM, NOISELESS, z_m, 10.0, np.random.default_rng(0), ghost=ghost
    ).astype(np.float64)


def test_ghost_second_moments_match_the_two_spot_mixture() -> None:
    """Two equal-width spots, power ratio r, offset dx: the mixture has

    sigma_x² = sigma² + w1 w2 dx² (w1 = 1/(1+r), w2 = r/(1+r)), its centroid

    moves by w2 dx, and the y moment is unchanged."""
    r, dx = 0.1, 0.3e-3
    full = Roi(0, NOISELESS.nx, 0, NOISELESS.ny)
    clean = second_moments(_frame(), NOISELESS.pixel_size_m, full)
    ghosted = second_moments(
        _frame(Ghost(relative_power=r, offset_x_m=dx, offset_y_m=0.0)),
        NOISELESS.pixel_size_m,
        full,
    )
    w1, w2 = 1 / (1 + r), r / (1 + r)
    # Integer rounding of the faint tails costs ~0.1 % of the shift.
    assert ghosted.x_c_m - clean.x_c_m == pytest.approx(w2 * dx, rel=2e-3)
    assert ghosted.var_x_m2 == pytest.approx(clean.var_x_m2 + w1 * w2 * dx**2, rel=1e-3)
    assert ghosted.var_y_m2 == pytest.approx(clean.var_y_m2, rel=1e-3)
    assert ghosted.total_counts == pytest.approx(clean.total_counts * (1 + r), rel=1e-3)


@pytest.mark.parametrize(

    ("offset_x_m", "offset_y_m"), [(0.3e-3, 0.0), (0.0, -0.5e-3), (0.4e-3, 0.4e-3)]

)
def test_ghost_makes_iso_widths_larger(offset_x_m: float, offset_y_m: float) -> None:
    pixel = NOISELESS.pixel_size_m
    clean = measure_widths(_frame(), pixel)
    ghosted = measure_widths(
        _frame(Ghost(relative_power=0.1, offset_x_m=offset_x_m, offset_y_m=offset_y_m)),
        pixel,
    )
    if offset_x_m:
        assert ghosted.d_x_m > 1.01 * clean.d_x_m
    if offset_y_m:
        assert ghosted.d_y_m > 1.01 * clean.d_y_m
    if not offset_x_m:
        assert ghosted.d_x_m == pytest.approx(clean.d_x_m, rel=2e-3)


def test_ghost_only_inside_its_z_window_and_offset_follows_z() -> None:
    ghost = Ghost(
        relative_power=0.2,
        offset_x_m=0.1e-3,
        offset_y_m=0.0,
        offset_x_per_m=2e-3,  # +2 um per mm of z
        z_min_m=0.20,
        z_max_m=0.30,
    )
    assert np.array_equal(_frame(ghost, z_m=0.15), _frame(None, z_m=0.15))
    assert not np.array_equal(_frame(ghost, z_m=0.25), _frame(None, z_m=0.25))
    assert ghost.offset_at(0.25) == pytest.approx((0.1e-3 + 0.5e-3, 0.0))


@pytest.mark.parametrize(

    "kwargs",

    [

        {"relative_power": 0.0},

        {"relative_power": 0.1, "size_factor": 0.0},

        {"relative_power": 0.1, "z_min_m": 0.2, "z_max_m": 0.1},

        {"relative_power": 0.1, "offset_x_m": float("nan")},

        {"relative_power": 0.1, "offset_y_per_m": float("inf")},

    ],

)
def test_ghost_rejects_bad_parameters(kwargs) -> None:
    kwargs = {"offset_x_m": 0.0, "offset_y_m": 0.0, **kwargs}
    with pytest.raises(ValueError):
        Ghost(**kwargs)


# Aperture (iris) and hot pixels


def test_aperture_moments_match_a_truncated_gaussian() -> None:
    """Circular Gaussian I ~ exp(-2 r²/w²) cut at radius a, with A = 2a²/w²:

    sigma_x² = (w²/4) (1 - A e^-A / (1 - e^-A)), and the power drops to

    1 - e^-A of the full beam."""
    w = 0.4e-3  # 1/e² radius at this z (d = 2w), 3.5 w to the sensor edge
    round_beam = BeamParams(1064e-9, 2 * w, 2 * w, 0.1, 0.1, 1.0, 1.0)
    a = 0.4e-3
    big_a = 2 * a**2 / w**2
    full = Roi(0, NOISELESS.nx, 0, NOISELESS.ny)

    def moments(aperture):
        frame = render_frame(
            round_beam,
            NOISELESS,
            0.1,
            10.0,
            np.random.default_rng(0),
            aperture=aperture,
        ).astype(np.float64)
        return second_moments(frame, NOISELESS.pixel_size_m, full)

    clean = moments(None)
    cut = moments(Aperture(radius_m=a, z_min_m=0.0, z_max_m=0.2))
    expected = (w**2 / 4) * (1 - big_a * np.exp(-big_a) / (1 - np.exp(-big_a)))
    # Integer rounding of the faint tails lowers the full-beam moment ~0.2 %.
    assert clean.var_x_m2 == pytest.approx(w**2 / 4, rel=5e-3)
    assert cut.var_x_m2 == pytest.approx(expected, rel=5e-3)
    assert cut.var_y_m2 == pytest.approx(expected, rel=5e-3)
    assert cut.total_counts == pytest.approx(
        clean.total_counts * (1 - np.exp(-big_a)), rel=2e-3
    )


def test_aperture_makes_iso_widths_smaller_only_inside_its_range() -> None:
    pixel = NOISELESS.pixel_size_m
    iris = Aperture(radius_m=0.4e-3, z_min_m=0.20, z_max_m=0.30)
    clean = measure_widths(_frame(), pixel)
    clipped = measure_widths(
        render_frame(
            BEAM, NOISELESS, GHOST_Z_M, 10.0, np.random.default_rng(0), aperture=iris
        ).astype(np.float64),
        pixel,
    )
    assert clipped.d_x_m < 0.9 * clean.d_x_m
    assert clipped.d_y_m < 0.9 * clean.d_y_m
    outside = render_frame(
        BEAM, NOISELESS, 0.15, 10.0, np.random.default_rng(0), aperture=iris
    )
    assert np.array_equal(outside, _frame(None, z_m=0.15).astype(np.int32))


def test_hot_pixels_add_counts_to_light_frames_only() -> None:
    hot = HotPixels(pixels=((40, 470), (41, 471)), extra_counts=2500.0)
    frame = render_frame(
        BEAM, NOISELESS, GHOST_Z_M, 10.0, np.random.default_rng(0), hot_pixels=hot
    ).astype(np.float64)
    difference = frame - _frame()
    assert difference[40, 470] == pytest.approx(2500.0, abs=1)
    assert difference[41, 471] == pytest.approx(2500.0, abs=1)
    difference[40, 470] = difference[41, 471] = 0.0
    assert not difference.any()


@pytest.mark.parametrize(

    ("factory", "kwargs"),

    [

        (Aperture, {"radius_m": 0.0, "z_min_m": 0.0, "z_max_m": 1.0}),

        (Aperture, {"radius_m": 1e-3, "z_min_m": 0.2, "z_max_m": 0.1}),

        (

            Aperture,

            {"radius_m": 1e-3, "z_min_m": 0, "z_max_m": 1, "center_x_m": np.nan},

        ),

        (Aperture, {"radius_m": 1e-3, "z_min_m": 0.0, "z_max_m": np.inf}),

        (HotPixels, {"pixels": (), "extra_counts": 10.0}),

        (HotPixels, {"pixels": ((1, 1),), "extra_counts": 0.0}),

        (HotPixels, {"pixels": ((-1, 3),), "extra_counts": 10.0}),

    ],

)
def test_aperture_and_hot_pixels_reject_bad_parameters(factory, kwargs) -> None:
    with pytest.raises(ValueError):
        factory(**kwargs)