| """Known-answer tests for the caustic fit, ISO coverage and astigmatism check."""
|
|
|
| import numpy as np
|
| import pytest
|
|
|
| from beampilot.caustic import coverage, fit_caustic, general_astigmatism
|
| from beampilot.simulator import (
|
| BeamParams,
|
| CameraConfig,
|
| beam_diameter,
|
| dark_frame,
|
| rayleigh_length,
|
| render_frame,
|
| )
|
| from beampilot.widths import (
|
| mean_dark,
|
| measure_widths,
|
| subtract_background,
|
| width_statistics,
|
| )
|
|
|
| WAVELENGTH_M = 1064e-9
|
| D0_M = 0.5e-3
|
| Z0_M = 0.10
|
|
|
|
|
| ISO_OFFSETS = np.array([-3.0, -2.2, -0.8, -0.4, 0.0, 0.4, 0.8, 2.2, 3.0, 3.8])
|
|
|
|
|
| NEAR_ONLY_OFFSETS = np.linspace(-0.5, 0.5, 10)
|
|
|
|
|
| def _diameters(offsets: np.ndarray, m2: float) -> tuple[np.ndarray, np.ndarray]:
|
| """Stage positions and true diameters in m for offsets in units of zR."""
|
| z_r = rayleigh_length(D0_M, m2, WAVELENGTH_M)
|
| z = Z0_M + offsets * z_r
|
| return z, beam_diameter(z, D0_M, Z0_M, m2, WAVELENGTH_M)
|
|
|
|
|
| @pytest.mark.parametrize("m2", [1.0, 1.5])
|
| def test_noiseless_caustic_recovers_beam(m2: float) -> None:
|
| z, d = _diameters(ISO_OFFSETS, m2)
|
| fit = fit_caustic(z, d, wavelength_m=WAVELENGTH_M)
|
|
|
| assert fit.success
|
| assert fit.n_points == 10
|
| assert fit.d0_m == pytest.approx(D0_M, rel=1e-6)
|
| assert fit.z0_m == pytest.approx(Z0_M, rel=1e-6)
|
| assert fit.m2 == pytest.approx(m2, rel=1e-6)
|
| z_r = rayleigh_length(D0_M, m2, WAVELENGTH_M)
|
| assert fit.z_r_m == pytest.approx(z_r, rel=1e-6)
|
|
|
| theta = 4 * m2 * WAVELENGTH_M / (np.pi * D0_M)
|
| assert fit.theta_rad == pytest.approx(theta, rel=1e-6)
|
| assert fit.residual_rms < 1e-9
|
|
|
|
|
| def test_noiseless_caustic_with_repeat_std() -> None:
|
| m2 = 1.5
|
| z, d = _diameters(ISO_OFFSETS, m2)
|
| fit = fit_caustic(z, d, 0.01 * d, wavelength_m=WAVELENGTH_M)
|
| assert fit.success
|
| assert fit.m2 == pytest.approx(m2, rel=1e-6)
|
|
|
|
|
| @pytest.mark.parametrize("m2", [1.0, 1.5])
|
| def test_two_percent_noise_mean_m2_unbiased(m2: float) -> None:
|
| z, d_true = _diameters(ISO_OFFSETS, m2)
|
| m2_fits, m2_stds = [], []
|
| for seed in range(100):
|
| rng = np.random.default_rng(seed)
|
| d = d_true * (1.0 + 0.02 * rng.standard_normal(d_true.size))
|
| fit = fit_caustic(z, d, wavelength_m=WAVELENGTH_M)
|
| assert fit.success, f"seed {seed}"
|
| m2_fits.append(fit.m2)
|
| m2_stds.append(fit.m2_std)
|
| m2_fits = np.array(m2_fits)
|
| print(
|
| f"\nM² = {m2}, 2 % noise, 100 seeds: mean {m2_fits.mean():.4f} "
|
| f"(error {m2_fits.mean() / m2 - 1:+.2%}), "
|
| f"scatter std {m2_fits.std(ddof=1):.4f}, "
|
| f"mean fit uncertainty {np.mean(m2_stds):.4f}"
|
| )
|
| assert m2_fits.mean() == pytest.approx(m2, rel=0.01)
|
|
|
| assert np.mean(m2_stds) == pytest.approx(m2_fits.std(ddof=1), rel=0.3)
|
|
|
|
|
| def test_near_waist_only_sampling_inflates_m2_uncertainty() -> None:
|
| """Without far-field points the divergence, and so M², is badly constrained."""
|
| m2 = 1.2
|
| z_iso, d_iso = _diameters(ISO_OFFSETS, m2)
|
| z_near, d_near = _diameters(NEAR_ONLY_OFFSETS, m2)
|
| iso_stds, near_stds = [], []
|
| for seed in range(20):
|
| rng = np.random.default_rng(seed)
|
| noise = 1.0 + 0.02 * rng.standard_normal(10)
|
| iso = fit_caustic(z_iso, d_iso * noise, wavelength_m=WAVELENGTH_M)
|
| near = fit_caustic(z_near, d_near * noise, wavelength_m=WAVELENGTH_M)
|
| iso_stds.append(iso.m2_std)
|
| near_stds.append(near.m2_std)
|
| ratio = np.median(near_stds) / np.median(iso_stds)
|
| print(f"\nnear-waist / ISO median M² uncertainty: {ratio:.1f}x")
|
| assert ratio >= 3.0
|
|
|
|
|
| def test_fit_needs_four_points() -> None:
|
| z, d = _diameters(ISO_OFFSETS[:3], 1.0)
|
| with pytest.raises(ValueError, match="at least 4"):
|
| fit_caustic(z, d, wavelength_m=WAVELENGTH_M)
|
|
|
|
|
| def test_fit_rejects_bad_input() -> None:
|
| z, d = _diameters(ISO_OFFSETS, 1.0)
|
| with pytest.raises(ValueError, match="wavelength"):
|
| fit_caustic(z, d, wavelength_m=0.0)
|
| with pytest.raises(ValueError, match="> 0"):
|
| fit_caustic(z, -d, wavelength_m=WAVELENGTH_M)
|
| with pytest.raises(ValueError, match="d_std_m"):
|
| fit_caustic(z, d, d[:5], wavelength_m=WAVELENGTH_M)
|
| with pytest.raises(ValueError, match="distinct"):
|
| fit_caustic([0.1, 0.1, 0.2, 0.2], d[:4], wavelength_m=WAVELENGTH_M)
|
|
|
|
|
| def test_tiny_repeat_std_is_floored() -> None:
|
| """A zero std must not make one point dominate (floor at 1e-3 * d)."""
|
| m2 = 1.0
|
| z, d_true = _diameters(ISO_OFFSETS, m2)
|
| d = d_true * (1.0 + 0.02 * np.random.default_rng(3).standard_normal(10))
|
| d_std = 0.02 * d
|
| d_std[4] = 0.0
|
| fit = fit_caustic(z, d, d_std, wavelength_m=WAVELENGTH_M)
|
| assert fit.success
|
| assert np.isfinite(fit.m2_std)
|
|
|
|
|
| def test_coverage_counts_distinct_positions() -> None:
|
|
|
| z0, z_r = 0.5, 0.125
|
| z = [
|
| 0.5,
|
| 0.5,
|
| 0.375,
|
| 0.625,
|
| 0.4375,
|
| 0.6875,
|
| 0.25,
|
| 0.75,
|
| 0.0,
|
| 1.0,
|
| 1.0,
|
| ]
|
| cov = coverage(z, z0, z_r)
|
| assert (cov.n_near, cov.n_far, cov.n_between) == (4, 4, 1)
|
| assert not cov.compliant
|
| assert coverage(z, z0, z_r, min_near=4, min_far=4).compliant
|
| assert coverage([*z, 0.5625, 1.125], z0, z_r).compliant
|
|
|
|
|
| def test_coverage_rejects_non_positive_rayleigh_length() -> None:
|
| with pytest.raises(ValueError, match="z_r_m"):
|
| coverage([0.0, 0.1], 0.05, 0.0)
|
|
|
|
|
| @pytest.mark.parametrize(
|
| ("azimuths_deg", "ellipticities", "expected"),
|
| [
|
| ([5.0, 20.0], [0.5, 0.5], True),
|
| ([5.0, 12.0], [0.5, 0.5], False),
|
| ([0.0, 45.0], [0.5, 0.95], False),
|
| ([0.0, 45.0, 70.0], [0.95, 0.9, 0.5], False),
|
| ([89.0, -89.0], [0.5, 0.5], False),
|
| ([80.0, -80.0], [0.5, 0.5], True),
|
| ([0.0, 90.0], [0.5, 0.5], False),
|
| ([0.0, 45.0], [0.5, 0.5], True),
|
| ([0.0, 30.0], [0.5, np.nan], False),
|
| ([], [], False),
|
| ],
|
| )
|
| def test_general_astigmatism(
|
| azimuths_deg: list[float], ellipticities: list[float], expected: bool
|
| ) -> None:
|
| assert general_astigmatism(np.radians(azimuths_deg), ellipticities) is expected
|
|
|
|
|
| def test_simple_astigmatic_beam_is_not_general_astigmatism() -> None:
|
| """x and y waists 2 zR apart: the major axis swaps from y to x along z."""
|
| z_r = rayleigh_length(0.2e-3, 1.0, WAVELENGTH_M)
|
| beam = BeamParams(
|
| wavelength_m=WAVELENGTH_M,
|
| d0x_m=0.2e-3,
|
| d0y_m=0.2e-3,
|
| z0x_m=0.0,
|
| z0y_m=2 * z_r,
|
| m2x=1.0,
|
| m2y=1.0,
|
| )
|
| camera = CameraConfig(
|
| pixel_size_m=5.5e-6,
|
| nx=512,
|
| ny=512,
|
| bit_depth=12,
|
| offset_counts=0.0,
|
| read_noise_counts=0.0,
|
| signal_counts_per_ms=1e6,
|
| shot_noise=False,
|
| )
|
| rng = np.random.default_rng(0)
|
| results = [
|
| measure_widths(render_frame(beam, camera, z, 1.0, rng), camera.pixel_size_m)
|
| for z in np.array([-1.0, -0.5, 0.0, 1.0, 2.0, 2.5, 3.0]) * z_r
|
| ]
|
| azimuths = np.array([r.azimuth_rad for r in results])
|
| ellipticities = np.array([r.ellipticity for r in results])
|
| non_circular = ellipticities <= 0.87
|
|
|
| assert np.any(np.abs(azimuths[non_circular]) < 0.01)
|
| assert np.any(np.abs(azimuths[non_circular]) > np.radians(89.0))
|
| assert not general_astigmatism(azimuths, ellipticities)
|
|
|
|
|
| def test_general_astigmatism_shape_mismatch_raises() -> None:
|
| with pytest.raises(ValueError):
|
| general_astigmatism([0.0, 0.1], [0.5])
|
|
|
|
|
| def test_end_to_end_simulated_frames_give_m2_per_axis() -> None:
|
| """Simulator -> dark subtraction -> widths -> repeat stats -> fit, per axis.
|
|
|
| The waists are sized so that both axes have the same zR (~24.6 mm), so one
|
| set of 10 positions satisfies the ISO point rule for x and y, and the
|
| 3 x d rectangle still fits on the 512 px sensor at the farthest position.
|
| """
|
| m2x, m2y = 1.2, 1.5
|
| d0x = 0.2e-3
|
| d0y = d0x * np.sqrt(m2y / m2x)
|
| beam = BeamParams(
|
| wavelength_m=WAVELENGTH_M,
|
| d0x_m=d0x,
|
| d0y_m=d0y,
|
| z0x_m=0.100,
|
| z0y_m=0.105,
|
| m2x=m2x,
|
| m2y=m2y,
|
| )
|
| 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,
|
| )
|
| z_r = rayleigh_length(d0x, m2x, WAVELENGTH_M)
|
| z_mid = 0.5 * (beam.z0x_m + beam.z0y_m)
|
| z_positions = (
|
| z_mid + np.array([-3.0, -2.3, -0.8, -0.4, 0, 0.4, 0.8, 2.3, 3.0, 3.6]) * z_r
|
| )
|
|
|
| rng = np.random.default_rng(2026)
|
| dark = mean_dark([dark_frame(camera, rng) for _ in range(10)])
|
| d_x, d_x_std, d_y, d_y_std = [], [], [], []
|
| for z in z_positions:
|
|
|
| dx = beam_diameter(z, d0x, beam.z0x_m, m2x, WAVELENGTH_M)
|
| dy = beam_diameter(z, d0y, beam.z0y_m, m2y, WAVELENGTH_M)
|
| peak_per_ms = (
|
| camera.signal_counts_per_ms
|
| * camera.pixel_size_m**2
|
| / (2 * np.pi * (dx / 4) * (dy / 4))
|
| )
|
| exposure_ms = 0.75 * camera.full_scale_counts / peak_per_ms
|
| results = []
|
| for _ in range(5):
|
| frame = render_frame(beam, camera, z, exposure_ms, rng)
|
| image, _ = subtract_background(frame, dark)
|
| result = measure_widths(image, camera.pixel_size_m)
|
| assert result.converged and not result.roi_clipped
|
| results.append(result)
|
| stats = width_statistics(results)
|
| d_x.append(stats.d_x_mean_m)
|
| d_x_std.append(stats.d_x_std_m)
|
| d_y.append(stats.d_y_mean_m)
|
| d_y_std.append(stats.d_y_std_m)
|
|
|
| fit_x = fit_caustic(z_positions, d_x, d_x_std, wavelength_m=WAVELENGTH_M)
|
| fit_y = fit_caustic(z_positions, d_y, d_y_std, wavelength_m=WAVELENGTH_M)
|
| print(
|
| f"\nend-to-end M²x {fit_x.m2:.4f} +- {fit_x.m2_std:.4f} "
|
| f"(error {fit_x.m2 / m2x - 1:+.2%}), "
|
| f"M²y {fit_y.m2:.4f} +- {fit_y.m2_std:.4f} (error {fit_y.m2 / m2y - 1:+.2%})"
|
| )
|
| print(
|
| f"d0x error {fit_x.d0_m / d0x - 1:+.2%}, "
|
| f"d0y error {fit_y.d0_m / d0y - 1:+.2%}, "
|
| f"z0x error {(fit_x.z0_m - beam.z0x_m) * 1e3:+.3f} mm, "
|
| f"z0y error {(fit_y.z0_m - beam.z0y_m) * 1e3:+.3f} mm"
|
| )
|
| assert fit_x.success and fit_y.success
|
| assert fit_x.m2 == pytest.approx(m2x, rel=0.03)
|
| assert fit_y.m2 == pytest.approx(m2y, rel=0.03)
|
|
|
| assert coverage(z_positions, fit_x.z0_m, fit_x.z_r_m).compliant
|
| assert coverage(z_positions, fit_y.z0_m, fit_y.z_r_m).compliant
|
|
|
