diff --git a/docs/superpowers/plans/2026-07-03-camera-geometry-redesign-plan.md b/docs/superpowers/plans/2026-07-03-camera-geometry-redesign-plan.md new file mode 100644 index 0000000..6290cf2 --- /dev/null +++ b/docs/superpowers/plans/2026-07-03-camera-geometry-redesign-plan.md @@ -0,0 +1,3361 @@ +# Camera Geometry & Measurement Uncertainty Redesign Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace he11lib's single-scalar pixel-scale/viewing-angle camera model and exact-`z`/single-pointing-angle assumptions with a shared, physically-parameterized pinhole `CameraModel` (true perspective projection), 2D beam pointing, and a uniform tolerance mechanism that lets every nominal geometry value (camera pose/intrinsics, per-plane `z`) be either held fixed or jointly refined within a bound. + +**Architecture:** `geometry.py` gains `CameraModel`/`CameraModelTolerance` dataclasses and a `GeometryCalibration` rewritten around true pinhole forward/inverse projection (ray-plane intersection) instead of the old cosine-compression formula. `data.py`'s `MeasurementPlane` drops `pixel_scale`/`viewing_angle_deg` for `z_tolerance`; `ReconstructionResult.pointing_angle_deg` splits into horizontal/vertical fields. `fitting.py`'s `ModalFitter` builds its optimizer parameter vector dynamically: coefficients/center/pointing angles always free, camera fields and per-plane `z` free only when their paired tolerance is `> 0` (bounded fit) and otherwise held as constants. `synthetic.py`, `phase_retrieval.py`, and `reconstruct.py` are updated to match, and `docs/api.md`/`examples/full_pipeline_example.py`/`CLAUDE.md` are brought in sync. + +**Tech Stack:** Python 3.10+, NumPy, `scipy.optimize.least_squares` (`bounds=`), pytest. No new dependencies. + +## Global Constraints + +- Python `>=3.10`, `numpy>=1.24`, `scipy>=1.10`, `matplotlib>=3.7` (unchanged floors from `pyproject.toml`). No new dependencies. +- Out of scope (per spec): lens distortion, rolling-shutter effects, multi-camera setups. +- `CameraModelTolerance` fields and `MeasurementPlane.z_tolerance` must be `>= 0`; raise `ValueError` at construction otherwise (validate only at boundaries, matching existing style). +- Tolerance mechanism: `tolerance == 0` holds a value fixed (excluded from the optimizer's parameter vector, substituted as a constant); `tolerance > 0` bounds it to `[nominal - tolerance, nominal + tolerance]` via `scipy.optimize.least_squares(bounds=...)`. No unbounded/"fully unknown" mode for these parameters. +- Degenerate camera geometry (target plane edge-on to or behind the camera) raises `ValueError`, never produces NaNs silently. +- `fit_auto`/`BeamReconstructor` must emit `UserWarning` (not raise) when the free camera+z parameter count is large relative to the number of planes (concrete rule below, Task 5). +- Follow existing code style: `from __future__ import annotations`, module + class docstrings explaining physical meaning and units, type hints on public signatures, dataclasses for data containers, "validate only at boundaries." +- Keep `he11lib/__init__.py`'s `__all__` in sync with every new/removed public name. +- `tests/conftest.py` already forces the `Agg` matplotlib backend; no change needed there. +- This is a breaking pre-1.0 API change (no external users) — do not add backwards-compatibility shims for the removed `pixel_scale`/`viewing_angle_deg`/`pointing_angle_deg` fields. +- Numeric test tolerances (e.g. `abs=`, `rel=`) given in this plan's test code are reasonable starting points for the chosen synthetic parameters, not sacred values — if a step's "run to verify PASS" fails only because a tolerance is a little tight/loose for the true perspective model's behavior (not because the implementation is wrong), adjust the tolerance constant and re-run, consistent with this project's documented physics/fitting pitfalls in `CLAUDE.md`. + +--- + +## Task 1: `CameraModel`, `CameraModelTolerance`, `GeometryCalibration` rewrite + +**Files:** +- Modify: `he11lib/geometry.py` (full rewrite) +- Modify: `he11lib/__init__.py` (export `CameraModel`, `CameraModelTolerance`) +- Modify: `tests/test_geometry.py` (full rewrite) + +**Interfaces:** +- Produces: `CameraModel(focal_length_px, position, orientation_deg, principal_point=(0.0, 0.0))`; `CameraModelTolerance(focal_length_px, position, orientation_deg, principal_point=(0.0, 0.0))` (raises `ValueError` if any field `< 0`); `GeometryCalibration(camera: CameraModel)` with `.pixel_coordinates(x, y, z) -> (row, col)`, `.physical_coordinates(image_shape, z) -> (x, y)`, `.effective_pixel_scale(image_shape, z) -> float`; module-level `CAMERA_FIELD_NAMES: tuple[str, ...]`, `camera_to_values(camera) -> list[float]`, `tolerance_to_values(tolerance) -> list[float]`, `camera_from_values(values) -> CameraModel` (used internally by `fitting.py` in Tasks 4-5). +- Consumes: nothing from other tasks (this is the foundational module). + +This is a full-file rewrite; the old `pixel_scale_known`/`viewing_angle_known` properties and cosine-compression `physical_coordinates(pixel_scale=, viewing_angle_deg=)` are removed entirely. + +- [ ] **Step 1: Write the failing tests** + +Replace `tests/test_geometry.py` entirely with: + +```python +import numpy as np +import pytest + +from he11lib.geometry import CameraModel, CameraModelTolerance, GeometryCalibration + + +def test_camera_model_tolerance_accepts_zero_and_positive(): + CameraModelTolerance( + focal_length_px=0.0, + position=(0.0, 0.0, 0.0), + orientation_deg=(1.0, 2.0, 3.0), + principal_point=(0.5, 0.5), + ) # should not raise + + +def test_camera_model_tolerance_rejects_negative_scalar_field(): + with pytest.raises(ValueError, match="focal_length_px"): + CameraModelTolerance( + focal_length_px=-1.0, + position=(0.0, 0.0, 0.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + + +def test_camera_model_tolerance_rejects_negative_tuple_component(): + with pytest.raises(ValueError, match="position"): + CameraModelTolerance( + focal_length_px=1.0, + position=(0.0, -0.5, 0.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + + +def make_on_axis_camera(focal_length_px=2000.0, camera_z=-2.0): + return CameraModel( + focal_length_px=focal_length_px, + position=(0.0, 0.0, camera_z), + orientation_deg=(0.0, 0.0, 0.0), + ) + + +def make_tilted_camera(): + return CameraModel( + focal_length_px=2000.0, + position=(0.05, -0.03, -2.0), + orientation_deg=(8.0, -5.0, 3.0), + ) + + +@pytest.mark.parametrize( + "camera", + [make_on_axis_camera(), make_tilted_camera()], + ids=["on_axis", "tilted_off_center"], +) +@pytest.mark.parametrize("z", [0.3, 0.5, 0.8]) +def test_projection_round_trip_recovers_pixel_grid(camera, z): + image_shape = (41, 41) + calib = GeometryCalibration(camera) + + x, y = calib.physical_coordinates(image_shape, z) + row, col = calib.pixel_coordinates(x, y, z) + + rows, cols = image_shape + row_idx = np.arange(rows) - rows // 2 + col_idx = np.arange(cols) - cols // 2 + expected_col, expected_row = np.meshgrid(col_idx, row_idx) + + np.testing.assert_allclose(row, expected_row, atol=1e-6) + np.testing.assert_allclose(col, expected_col, atol=1e-6) + + +def test_keystone_regression_uniform_for_on_axis_camera(): + # A camera with zero orientation, centered on the beam axis, produces + # uniform pixel spacing for evenly spaced physical points (no keystoning). + camera = make_on_axis_camera() + calib = GeometryCalibration(camera) + z = 0.5 + + xs = np.array([-0.02, -0.01, 0.0, 0.01, 0.02]) + ys = np.zeros_like(xs) + _, col = calib.pixel_coordinates(xs, ys, z) + + spacings = np.diff(col) + np.testing.assert_allclose(spacings, spacings[0], rtol=1e-6) + + +def test_keystone_regression_nonuniform_for_tilted_camera(): + # A tilted/off-axis camera produces non-uniform pixel spacing for the + # same evenly spaced physical points -- genuine keystoning. + camera = make_tilted_camera() + calib = GeometryCalibration(camera) + z = 0.5 + + xs = np.array([-0.02, -0.01, 0.0, 0.01, 0.02]) + ys = np.zeros_like(xs) + _, col = calib.pixel_coordinates(xs, ys, z) + + spacings = np.diff(col) + assert not np.allclose(spacings, spacings[0], rtol=1e-3) + + +def test_pixel_coordinates_raises_when_point_behind_camera(): + camera = CameraModel( + focal_length_px=2000.0, + position=(0.0, 0.0, 10.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + calib = GeometryCalibration(camera) + + with pytest.raises(ValueError): + calib.pixel_coordinates(np.array([0.0]), np.array([0.0]), z=0.5) + + +def test_physical_coordinates_raises_when_plane_behind_camera(): + # Camera sits downstream of the target plane and looks further + # downstream (boresight = +z world) -- the z=0.5 plane is behind it. + camera = CameraModel( + focal_length_px=2000.0, + position=(0.0, 0.0, 10.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + calib = GeometryCalibration(camera) + + with pytest.raises(ValueError): + calib.physical_coordinates((21, 21), z=0.5) + + +def test_physical_coordinates_raises_when_edge_on(): + # Pitch=90 deg points the boresight along world -y, making the + # z=const target plane edge-on (parallel to the view direction). + camera = CameraModel( + focal_length_px=2000.0, + position=(0.0, 0.0, -2.0), + orientation_deg=(0.0, 90.0, 0.0), + ) + calib = GeometryCalibration(camera) + + with pytest.raises(ValueError): + calib.physical_coordinates((41, 41), z=0.5) + + +def test_effective_pixel_scale_matches_on_axis_focal_length(): + focal_length_px = 2000.0 + camera_z = -2.0 + z = 0.5 + camera = make_on_axis_camera(focal_length_px=focal_length_px, camera_z=camera_z) + calib = GeometryCalibration(camera) + + scale = calib.effective_pixel_scale((41, 41), z) + expected = (z - camera_z) / focal_length_px + assert scale == pytest.approx(expected, rel=1e-6) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `.venv/bin/pytest tests/test_geometry.py -q` +Expected: FAIL with `ImportError: cannot import name 'CameraModel' from 'he11lib.geometry'` (or similar collection error), since `geometry.py` hasn't been rewritten yet. + +- [ ] **Step 3: Rewrite `he11lib/geometry.py`** + +```python +"""Camera geometry: a shared pinhole camera model and pixel<->physical mapping. + +Models the camera as a full pinhole camera (3D position + yaw/pitch/roll +orientation + focal length + principal point) shared across all measurement +planes in one reconstruction. Every nominal value on `CameraModel` is paired +with a `CameraModelTolerance` entry that determines whether `ModalFitter` +holds it fixed (tolerance == 0) or refines it within a bound +(tolerance > 0) -- `CameraModel` alone is never trusted as exact. + +Coordinate conventions +---------------------- +World frame: `x` increases along the pixel-column direction, `y` increases +along the pixel-row direction, `z` is distance from the output window along +the beam axis (target planes live at `z = const > 0`). + +Camera frame: `X_cam` = right (pixel-column direction), `Y_cam` = down +(pixel-row direction), `Z_cam` = boresight (depth). At +`orientation_deg == (0, 0, 0)`, the camera frame is axis-aligned with the +world frame, so the boresight points along `+z` -- normal to every +`z = const` target plane, with no in-plane rotation. + +`orientation_deg = (yaw, pitch, roll)` composes as +`R = R_yaw(about Y) @ R_pitch(about X) @ R_roll(about Z)`, applied to the +camera axes to obtain their world-frame directions. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Sequence + +import numpy as np + +CAMERA_FIELD_NAMES: tuple[str, ...] = ( + "focal_length_px", + "position_x", + "position_y", + "position_z", + "yaw_deg", + "pitch_deg", + "roll_deg", + "principal_point_x", + "principal_point_y", +) + + +@dataclass +class CameraModel: + """Nominal pinhole camera parameters, shared across all measurement planes. + + Never trusted as exact by itself -- pair with a `CameraModelTolerance` + to express how much each field may be refined during fitting. + + Parameters + ---------- + focal_length_px : focal length, in pixel units. + position : (x, y, z) camera position in the world (beam-axis) frame, + in meters. z=0 is the output window. + orientation_deg : (yaw, pitch, roll), in degrees. All-zero means the + boresight is normal to every z=const target plane, no in-plane + rotation (see module docstring for the full convention). + principal_point : (px, px) offset of the principal point from the frame + center, in pixels. + """ + + focal_length_px: float + position: tuple[float, float, float] + orientation_deg: tuple[float, float, float] + principal_point: tuple[float, float] = (0.0, 0.0) + + +@dataclass +class CameraModelTolerance: + """+/- bound (same units as `CameraModel`) within which each field is refined. + + `0` holds the paired `CameraModel` field fixed at its nominal value; + `> 0` bounds it to `[nominal - tolerance, nominal + tolerance]` during + fitting. All fields must be `>= 0`. + """ + + focal_length_px: float + position: tuple[float, float, float] + orientation_deg: tuple[float, float, float] + principal_point: tuple[float, float] = (0.0, 0.0) + + def __post_init__(self) -> None: + for f in fields(self): + value = getattr(self, f.name) + components = value if isinstance(value, tuple) else (value,) + for component in components: + if component < 0: + raise ValueError( + f"CameraModelTolerance.{f.name} must be >= 0, got {value}" + ) + + +def camera_to_values(camera: CameraModel) -> list[float]: + """Flatten a `CameraModel` into the 9 scalars named by `CAMERA_FIELD_NAMES`.""" + return [ + camera.focal_length_px, + camera.position[0], + camera.position[1], + camera.position[2], + camera.orientation_deg[0], + camera.orientation_deg[1], + camera.orientation_deg[2], + camera.principal_point[0], + camera.principal_point[1], + ] + + +def tolerance_to_values(tolerance: CameraModelTolerance) -> list[float]: + """Flatten a `CameraModelTolerance` into the 9 scalars named by `CAMERA_FIELD_NAMES`.""" + return [ + tolerance.focal_length_px, + tolerance.position[0], + tolerance.position[1], + tolerance.position[2], + tolerance.orientation_deg[0], + tolerance.orientation_deg[1], + tolerance.orientation_deg[2], + tolerance.principal_point[0], + tolerance.principal_point[1], + ] + + +def camera_from_values(values: Sequence[float]) -> CameraModel: + """Inverse of `camera_to_values`: rebuild a `CameraModel` from 9 scalars.""" + return CameraModel( + focal_length_px=values[0], + position=(values[1], values[2], values[3]), + orientation_deg=(values[4], values[5], values[6]), + principal_point=(values[7], values[8]), + ) + + +def _rotation_matrix(yaw_deg: float, pitch_deg: float, roll_deg: float) -> np.ndarray: + """3x3 rotation matrix mapping camera-frame axes to world-frame directions.""" + yaw = np.deg2rad(yaw_deg) + pitch = np.deg2rad(pitch_deg) + roll = np.deg2rad(roll_deg) + + cy, sy = np.cos(yaw), np.sin(yaw) + cx, sx = np.cos(pitch), np.sin(pitch) + cz, sz = np.cos(roll), np.sin(roll) + + r_yaw = np.array([[cy, 0.0, sy], [0.0, 1.0, 0.0], [-sy, 0.0, cy]]) + r_pitch = np.array([[1.0, 0.0, 0.0], [0.0, cx, -sx], [0.0, sx, cx]]) + r_roll = np.array([[cz, -sz, 0.0], [sz, cz, 0.0], [0.0, 0.0, 1.0]]) + + return r_yaw @ r_pitch @ r_roll + + +class GeometryCalibration: + """Resolves the pixel<->physical mapping for a shared pinhole `CameraModel`.""" + + def __init__(self, camera: CameraModel): + self.camera = camera + self._rotation = _rotation_matrix(*camera.orientation_deg) + + def pixel_coordinates( + self, x: np.ndarray, y: np.ndarray, z: float + ) -> tuple[np.ndarray, np.ndarray]: + """Forward pinhole projection: physical (x, y) at depth z -> centered pixel (row, col).""" + px, py, pz = self.camera.position + dx = x - px + dy = y - py + dz = z - pz + + r = self._rotation + xc = r[0, 0] * dx + r[1, 0] * dy + r[2, 0] * dz + yc = r[0, 1] * dx + r[1, 1] * dy + r[2, 1] * dz + zc = r[0, 2] * dx + r[1, 2] * dy + r[2, 2] * dz + + if np.any(zc <= 0): + raise ValueError( + f"One or more target points are behind or edge-on to the " + f"camera at z={z}; cannot project." + ) + + f = self.camera.focal_length_px + cx, cy = self.camera.principal_point + col = f * xc / zc + cx + row = f * yc / zc + cy + return row, col + + def physical_coordinates( + self, image_shape: tuple[int, int], z: float + ) -> tuple[np.ndarray, np.ndarray]: + """Inverse pinhole projection: pixel grid at depth z -> physical (x, y). + + Casts a ray from the camera through each pixel and intersects it + with the world plane z=const. Raises ValueError if the target + plane is edge-on to (parallel to) the view direction or behind the + camera for this pose. + """ + rows, cols = image_shape + row_idx = np.arange(rows) - rows // 2 + col_idx = np.arange(cols) - cols // 2 + col_grid, row_grid = np.meshgrid(col_idx, row_idx) + + f = self.camera.focal_length_px + cx, cy = self.camera.principal_point + dir_cam_x = (col_grid - cx) / f + dir_cam_y = (row_grid - cy) / f + dir_cam_z = np.ones_like(dir_cam_x) + + r = self._rotation + dir_world_x = r[0, 0] * dir_cam_x + r[0, 1] * dir_cam_y + r[0, 2] * dir_cam_z + dir_world_y = r[1, 0] * dir_cam_x + r[1, 1] * dir_cam_y + r[1, 2] * dir_cam_z + dir_world_z = r[2, 0] * dir_cam_x + r[2, 1] * dir_cam_y + r[2, 2] * dir_cam_z + + if np.any(np.abs(dir_world_z) < 1e-12): + raise ValueError( + f"Camera pose is edge-on to the target plane z={z}; no " + "valid ray-plane intersection." + ) + + px, py, pz = self.camera.position + t = (z - pz) / dir_world_z + if np.any(t <= 0): + raise ValueError( + f"Target plane z={z} is behind the camera for this pose; " + "no valid ray-plane intersection." + ) + + x = px + t * dir_world_x + y = py + t * dir_world_y + return x, y + + def effective_pixel_scale(self, image_shape: tuple[int, int], z: float) -> float: + """Isotropic finite-difference approximation of the local pixel scale. + + `DiffusionDeconvolver` assumes one isotropic pixel-space blur + kernel; this is only exact for an on-axis, zero-orientation + camera, and an approximation whenever the true projection is + keystoned. + """ + rows, cols = image_shape + x, y = self.physical_coordinates(image_shape, z) + mid_row, mid_col = rows // 2, cols // 2 + dx = abs(x[mid_row, mid_col + 1] - x[mid_row, mid_col]) + dy = abs(y[mid_row + 1, mid_col] - y[mid_row, mid_col]) + return float((dx + dy) / 2) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `.venv/bin/pytest tests/test_geometry.py -q` +Expected: PASS (all tests green). + +- [ ] **Step 5: Update package exports** + +In `he11lib/__init__.py`, change: + +```python +from .geometry import GeometryCalibration +``` + +to: + +```python +from .geometry import CameraModel, CameraModelTolerance, GeometryCalibration +``` + +and change: + +```python + "GeometryCalibration", +``` + +to: + +```python + "CameraModel", + "CameraModelTolerance", + "GeometryCalibration", +``` + +- [ ] **Step 6: Run the full suite to check for collection errors elsewhere** + +Run: `.venv/bin/pytest -q` +Expected: `tests/test_geometry.py` passes; other test files will now fail/error (they still use the old `MeasurementPlane(pixel_scale=..., viewing_angle_deg=...)` and old `GeometryCalibration(plane)` API) -- this is expected until Tasks 2-7 land. Confirm the failures are all in other files, not `test_geometry.py`. + +- [ ] **Step 7: Commit** + +```bash +git add he11lib/geometry.py he11lib/__init__.py tests/test_geometry.py +git commit -m "$(cat <<'EOF' +Replace cosine-compression geometry model with a full pinhole CameraModel + +GeometryCalibration now performs true perspective forward/inverse +projection (with genuine keystoning) around a shared CameraModel, paired +with a CameraModelTolerance that will drive ModalFitter's per-field +fixed/refined behavior in a later task. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 2: `MeasurementPlane`/`ReconstructionResult` data model changes + +**Files:** +- Modify: `he11lib/data.py` +- Modify: `he11lib/plotting.py:44` (pointing-angle field rename) +- Modify: `tests/test_data.py` +- Modify: `tests/test_plotting.py` + +**Interfaces:** +- Consumes: nothing new (dataclasses only). +- Produces: `MeasurementPlane(flux, z, z_tolerance=0.0, label=None)` (raises `ValueError` if `z_tolerance < 0`); `ReconstructionResult(..., pointing_angle_horizontal_deg, pointing_angle_vertical_deg, ...)` replacing the old single `pointing_angle_deg` field. Tasks 3-7 consume these exact field names. + +- [ ] **Step 1: Write the failing tests** + +Replace `tests/test_data.py` entirely with: + +```python +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.z_tolerance == 0.0 + assert plane.label is None + + +def test_measurement_plane_stores_optional_fields(): + flux = np.ones((4, 4)) + plane = MeasurementPlane(flux=flux, z=0.4, z_tolerance=0.01, label="plane_40cm") + + assert plane.z_tolerance == 0.01 + 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_measurement_plane_rejects_negative_z_tolerance(): + with pytest.raises(ValueError, match="z_tolerance"): + MeasurementPlane(flux=np.ones((4, 4)), z=0.3, z_tolerance=-0.01) + + +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_horizontal_deg=0.1, + pointing_angle_vertical_deg=-0.2, + geometry={"focal_length_px": 2000.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.pointing_angle_horizontal_deg == 0.1 + assert result.pointing_angle_vertical_deg == -0.2 + assert result.used_phase_retrieval is False +``` + +Replace `tests/test_plotting.py`'s `make_result` and `make_planes` helpers (leave the four `test_*` functions below them unchanged): + +```python +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_horizontal_deg=0.3, + pointing_angle_vertical_deg=0.4, + 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), + ] +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `.venv/bin/pytest tests/test_data.py tests/test_plotting.py -q` +Expected: FAIL — `test_data.py` fails with `TypeError: __init__() got an unexpected keyword argument 'z_tolerance'` (or missing `pointing_angle_horizontal_deg`); `test_plotting.py` fails similarly on `ReconstructionResult(**defaults)`. + +- [ ] **Step 3: Rewrite `he11lib/data.py`** + +```python +"""Data containers for he11lib: measurement inputs and reconstruction outputs.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np + + +@dataclass +class MeasurementPlane: + """A single thermal (flux) image at a nominal distance from the output window. + + Parameters + ---------- + flux : 2D array of flux values (already dead-pixel/background/saturation + corrected upstream). + z : nominal distance from the output window, in meters. Must be positive. + z_tolerance : +/- bound, in meters, around the nominal `z` within which + the true distance is refined during fitting. Must be `>= 0`; `0` + means `z` is trusted exactly and held fixed. + label : optional human-readable label (e.g. "plane_40cm"). + """ + + flux: np.ndarray + z: float + z_tolerance: float = 0.0 + label: str | None = None + + def __post_init__(self) -> None: + if self.flux.ndim != 2: + raise ValueError( + f"MeasurementPlane.flux must be a 2D array, got shape {self.flux.shape}" + ) + if self.z <= 0: + raise ValueError(f"MeasurementPlane.z must be positive, got {self.z}") + if self.z_tolerance < 0: + raise ValueError( + f"MeasurementPlane.z_tolerance must be >= 0, got {self.z_tolerance}" + ) + + +def validate_planes(planes: list[MeasurementPlane]) -> None: + """Validate a list of MeasurementPlanes for use in reconstruction. + + Raises ValueError if there are fewer than 3 planes, shapes mismatch + across planes, or z distances are not distinct. + """ + if len(planes) < 3: + raise ValueError( + f"At least 3 measurement planes are required, got {len(planes)}" + ) + + shapes = {p.flux.shape for p in planes} + if len(shapes) > 1: + raise ValueError(f"All MeasurementPlanes must have the same shape, got {shapes}") + + z_values = [p.z for p in planes] + if len(set(z_values)) != len(z_values): + raise ValueError(f"MeasurementPlane z distances must be distinct, got {z_values}") + + +@dataclass +class ReconstructionResult: + """Output of a full mode-purity reconstruction. + + Parameters + ---------- + purity : mapping from (p, l) mode index to (power_fraction, phase_rad). + reconstructed_field : reconstructed complex field (at the reference + waist, or as configured). + centers : fitted beam transverse center (x, y) in meters, per plane. + pointing_angle_horizontal_deg, pointing_angle_vertical_deg : fitted + shared beam pointing (tilt) angles, in degrees. + geometry : fitted/held geometry parameters, keyed by name (the 9 + `CameraModel` field names from `he11lib.geometry.CAMERA_FIELD_NAMES`, + plus `z_{i}` per plane index `i`). + residuals : per-plane residual maps (measured - modeled flux). + coefficient_uncertainty : mapping from (p, l) mode index to the + 1-sigma uncertainty on its fitted power fraction. + used_phase_retrieval : whether the phase-retrieval fallback was used + instead of (or to seed) the modal fit. + """ + + purity: dict[tuple[int, int], tuple[float, float]] + reconstructed_field: np.ndarray + centers: list[tuple[float, float]] + pointing_angle_horizontal_deg: float + pointing_angle_vertical_deg: float + geometry: dict[str, float] = field(default_factory=dict) + residuals: list[np.ndarray] = field(default_factory=list) + coefficient_uncertainty: dict[tuple[int, int], float] = field(default_factory=dict) + used_phase_retrieval: bool = False +``` + +- [ ] **Step 4: Fix `he11lib/plotting.py`'s pointing-angle reference** + +In `he11lib/plotting.py`, change: + +```python + fig.suptitle(f"Beam center (pointing angle {result.pointing_angle_deg:.3g} deg)") +``` + +to: + +```python + fig.suptitle( + "Beam center (pointing angle " + f"h={result.pointing_angle_horizontal_deg:.3g} deg, " + f"v={result.pointing_angle_vertical_deg:.3g} deg)" + ) +``` + +- [ ] **Step 5: Run tests to verify they pass** + +Run: `.venv/bin/pytest tests/test_data.py tests/test_plotting.py -q` +Expected: PASS. + +- [ ] **Step 6: Commit** + +```bash +git add he11lib/data.py he11lib/plotting.py tests/test_data.py tests/test_plotting.py +git commit -m "$(cat <<'EOF' +Replace per-plane pixel_scale/viewing_angle with z_tolerance; split pointing angle + +MeasurementPlane now carries a z_tolerance (uniform tolerance mechanism) +instead of the removed pixel_scale/viewing_angle_deg fields. +ReconstructionResult.pointing_angle_deg becomes horizontal/vertical fields +to match the beam's two independent tilt degrees of freedom. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 3: `SyntheticBeamGenerator` rewrite + +**Files:** +- Modify: `he11lib/synthetic.py` (full rewrite) +- Modify: `tests/test_synthetic.py` (full rewrite) + +**Interfaces:** +- Consumes: `CameraModel`, `GeometryCalibration` (Task 1); `MeasurementPlane(flux, z, z_tolerance=0.0, label=None)` (Task 2). +- Produces: `SyntheticBeamGenerator(basis, camera)`; `.generate(coefficients, z_list, image_shape, *, center=(0,0), pointing_angle_horizontal_deg=0.0, pointing_angle_vertical_deg=0.0, z_tolerance=0.0, nominal_z_offsets=None, noise_std=0.0, seed=None) -> list[MeasurementPlane]`. Tasks 4-7's tests all construct generators this way. + +- [ ] **Step 1: Write the failing tests** + +Replace `tests/test_synthetic.py` entirely with: + +```python +import numpy as np +import pytest + +from he11lib.geometry import CameraModel +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, achieved at z=Z0 +CAMERA_DISTANCE = 5.0 # camera stands 5 m upstream of the output window +IMAGE_SHAPE = (161, 161) # odd so there's a well-defined center pixel + + +def make_camera(pixel_scale=PIXEL_SCALE, z0=Z0, camera_distance=CAMERA_DISTANCE): + 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 make_generator(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + return SyntheticBeamGenerator(basis=basis, camera=make_camera()) + + +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, image_shape=IMAGE_SHAPE) + + 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], image_shape=IMAGE_SHAPE, 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 at z=Z0 + planes = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, 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_angles_as_2d_linear_drift(): + gen = make_generator() + z_list = [Z0, Z0 + 0.2] + planes = gen.generate( + coefficients={(0, 0): 1 + 0j}, + z_list=z_list, + image_shape=IMAGE_SHAPE, + center=(0.0, 0.0), + pointing_angle_horizontal_deg=1.0, + pointing_angle_vertical_deg=0.5, + ) + + peaks = [] + for plane in planes: + peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape) + peaks.append(peak_idx) + + expected_shift_x_m = 0.2 * np.tan(np.deg2rad(1.0)) + expected_shift_y_m = 0.2 * np.tan(np.deg2rad(0.5)) + expected_shift_col_px = expected_shift_x_m / PIXEL_SCALE + expected_shift_row_px = expected_shift_y_m / PIXEL_SCALE + + actual_shift_col_px = peaks[1][1] - peaks[0][1] + actual_shift_row_px = peaks[1][0] - peaks[0][0] + assert actual_shift_col_px == pytest.approx(expected_shift_col_px, abs=1) + assert actual_shift_row_px == pytest.approx(expected_shift_row_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], image_shape=IMAGE_SHAPE, noise_std=0.01, seed=42 + ) + planes_b = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, 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], image_shape=IMAGE_SHAPE, noise_std=noise_std, seed=1 + ) + planes_clean = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, 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_applies_z_tolerance_to_every_plane(): + gen = make_generator() + planes = gen.generate( + coefficients={(0, 0): 1 + 0j}, + z_list=[0.3, 0.4, 0.5], + image_shape=IMAGE_SHAPE, + z_tolerance=0.02, + ) + assert all(p.z_tolerance == 0.02 for p in planes) + + +def test_generate_applies_nominal_z_offset_independent_of_true_z(): + gen = make_generator() + true_z_list = [0.3, 0.4, 0.5] + offsets = {0.3: 0.01, 0.4: -0.005, 0.5: 0.0} + planes = gen.generate( + coefficients={(0, 0): 1 + 0j}, + z_list=true_z_list, + image_shape=IMAGE_SHAPE, + nominal_z_offsets=offsets, + ) + + nominal_zs = [p.z for p in planes] + assert nominal_zs == pytest.approx([0.31, 0.395, 0.5]) + # The flux is still rendered at each plane's *true* z (0.3, 0.4, 0.5), + # not its offset nominal z -- verified indirectly in Task 7's + # end-to-end tolerance-recovery test. +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `.venv/bin/pytest tests/test_synthetic.py -q` +Expected: FAIL with `TypeError: __init__() missing 1 required positional argument: 'camera'` (or similar), since `synthetic.py` hasn't been rewritten yet. + +- [ ] **Step 3: Rewrite `he11lib/synthetic.py`** + +```python +"""Forward model: synthetic thermal (flux) images from known ground truth. + +Used to validate the reconstruction pipeline (recover known mode content) +and to help users evaluate experimental design (e.g. whether a given set of +measurement distances will separate candidate modes). +""" + +from __future__ import annotations + +import numpy as np + +from .data import MeasurementPlane +from .geometry import CameraModel, GeometryCalibration +from .modes import LGBasis + + +class SyntheticBeamGenerator: + """Generates synthetic multi-plane flux images for a known ground-truth beam. + + Parameters + ---------- + basis : LGBasis defining the reference w0, z0, wavelength. + camera : ground-truth CameraModel (position/orientation/intrinsics) used + to render each plane via true perspective projection. + """ + + def __init__(self, basis: LGBasis, camera: CameraModel): + self.basis = basis + self.camera = camera + self.calibration = GeometryCalibration(camera) + + def generate( + self, + coefficients: dict[tuple[int, int], complex], + z_list: list[float], + image_shape: tuple[int, int], + *, + center: tuple[float, float] = (0.0, 0.0), + pointing_angle_horizontal_deg: float = 0.0, + pointing_angle_vertical_deg: float = 0.0, + z_tolerance: float = 0.0, + nominal_z_offsets: dict[float, float] | None = None, + noise_std: float = 0.0, + seed: int | None = None, + ) -> list[MeasurementPlane]: + """Generate one MeasurementPlane per requested (true) z distance. + + The beam transverse center drifts linearly with z according to the + two pointing angles, starting from `center` at the basis's + reference z0. `nominal_z_offsets`, if given, maps a true z (as + given in z_list) to an offset applied to the *nominal* z stored on + the resulting MeasurementPlane -- letting tests verify a fit + recovers the true z despite a deliberately-offset nominal input. + Every resulting plane shares `z_tolerance`. + """ + rng = np.random.default_rng(seed) + tilt_h_rad = np.deg2rad(pointing_angle_horizontal_deg) + tilt_v_rad = np.deg2rad(pointing_angle_vertical_deg) + offsets = nominal_z_offsets or {} + + planes = [] + for z in z_list: + drift_x = (z - self.basis.z0) * np.tan(tilt_h_rad) + drift_y = (z - self.basis.z0) * np.tan(tilt_v_rad) + cx = center[0] + drift_x + cy = center[1] + drift_y + + x, y = self.calibration.physical_coordinates(image_shape, z) + field = self.basis.field_superposition(x - cx, y - cy, z, coefficients) + flux = np.abs(field) ** 2 + + if noise_std > 0: + flux = flux + rng.normal(0.0, noise_std, size=flux.shape) + + nominal_z = z + offsets.get(z, 0.0) + planes.append( + MeasurementPlane(flux=flux, z=nominal_z, z_tolerance=z_tolerance) + ) + return planes +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `.venv/bin/pytest tests/test_synthetic.py -q` +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add he11lib/synthetic.py tests/test_synthetic.py +git commit -m "$(cat <<'EOF' +Rewrite SyntheticBeamGenerator around CameraModel and 2D beam pointing + +Renders each plane via true pinhole projection through a shared +CameraModel instead of the old cosine-compression formula, adds +independent horizontal/vertical pointing drift, and supports generating +a deliberately-offset nominal z (vs. true z) per plane for tolerance- +recovery testing in later tasks. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 4: `ModalFitter.fit()` rewrite with the tolerance mechanism + +**Files:** +- Modify: `he11lib/fitting.py` (rewrite `ModalFitter.fit`; `fit_auto`/`_bic`/`_warn_if_degenerate` land in Task 5) +- Modify: `tests/test_fitting.py` (rewrite the `fit`-level tests; `fit_auto` tests stay as-is structurally but move to Task 5's step since they call the not-yet-updated `fit_auto`) + +**Interfaces:** +- Consumes: `CameraModel`, `CameraModelTolerance`, `GeometryCalibration`, `CAMERA_FIELD_NAMES`, `camera_to_values`, `tolerance_to_values`, `camera_from_values` (Task 1); `MeasurementPlane.z_tolerance`, `ReconstructionResult.pointing_angle_horizontal_deg`/`pointing_angle_vertical_deg` (Task 2); `SyntheticBeamGenerator(basis, camera)` (Task 3). +- Produces: `ModalFitter.fit(planes, modes, camera, camera_tolerance, initial_coefficients=None, initial_center=(0.0, 0.0), initial_pointing_deg=(0.0, 0.0)) -> ReconstructionResult`. Task 5's `fit_auto` and Task 7's `BeamReconstructor` call `fit` with this exact signature. + +Note: this task temporarily leaves `fit_auto`, `_bic`, and the `test_fit_auto_*` tests referencing the *old* `fit_auto` signature broken/uncalled — they're rewritten together in Task 5, which lands immediately after. Do not run the full `test_fitting.py` file's `fit_auto` tests as a gate for this task; only the `fit`-level tests below. + +- [ ] **Step 1: Write the failing tests** + +Replace `tests/test_fitting.py`'s content above `test_fit_auto_does_not_add_modes_for_pure_fundamental` (i.e. everything through `test_fit_recovers_unknown_pixel_scale`) with: + +```python +import numpy as np +import pytest + +from he11lib.data import validate_planes +from he11lib.fitting import ModalFitter, generate_mode_shells +from he11lib.geometry import CameraModel, CameraModelTolerance +from he11lib.modes import LGBasis +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(pixel_scale=PIXEL_SCALE, position=(0.0, 0.0, -CAMERA_DISTANCE), orientation_deg=(0.0, 0.0, 0.0)): + focal_length_px = (CAMERA_DISTANCE + Z0) / pixel_scale + return CameraModel( + focal_length_px=focal_length_px, position=position, orientation_deg=orientation_deg + ) + + +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_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() + 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 + ) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance()) + + 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() + camera = make_camera() + gen = make_generator(basis, camera) + true_coeffs = {(0, 0): 0.9 + 0j, (1, 0): 0.3 + 0.1j} + planes = gen.generate( + coefficients=true_coeffs, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=1 + ) + + fitter = ModalFitter(basis) + result = fitter.fit( + planes, modes=list(true_coeffs.keys()), camera=camera, camera_tolerance=zero_tolerance() + ) + + 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() + 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=2, + ) + + fitter = ModalFitter(basis) + result = fitter.fit( + planes, + modes=[(0, 0)], + camera=camera, + camera_tolerance=zero_tolerance(), + 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_pointing_angles_independently(): + 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, + pointing_angle_horizontal_deg=0.3, + pointing_angle_vertical_deg=-0.15, + noise_std=1e-4, + seed=6, + ) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance()) + + assert result.pointing_angle_horizontal_deg == pytest.approx(0.3, abs=0.05) + assert result.pointing_angle_vertical_deg == pytest.approx(-0.15, abs=0.05) + + +def test_fit_holds_zero_tolerance_camera_field_fixed_at_wrong_nominal(): + # A tolerance=0 field must stay exactly at its (deliberately wrong) + # nominal value rather than being corrected. + basis = make_basis() + true_camera = make_camera() + gen = make_generator(basis, true_camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=7 + ) + + wrong_focal_length = true_camera.focal_length_px * 1.2 + nominal_camera = CameraModel( + focal_length_px=wrong_focal_length, + position=true_camera.position, + orientation_deg=true_camera.orientation_deg, + ) + + fitter = ModalFitter(basis) + result = fitter.fit( + planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=zero_tolerance() + ) + + assert result.geometry["focal_length_px"] == wrong_focal_length + + +def test_fit_recovers_offset_camera_field_within_tolerance(): + # A tolerance>0 field recovers a ground-truth offset from nominal, but + # within its band. + basis = make_basis() + true_camera = make_camera() + gen = make_generator(basis, true_camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=8 + ) + + offset = true_camera.focal_length_px * 0.02 # 2% off nominal + nominal_camera = CameraModel( + focal_length_px=true_camera.focal_length_px + offset, + position=true_camera.position, + orientation_deg=true_camera.orientation_deg, + ) + tolerance = CameraModelTolerance( + focal_length_px=true_camera.focal_length_px * 0.05, # +/-5% band + position=(0.0, 0.0, 0.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tolerance) + + assert result.geometry["focal_length_px"] == pytest.approx( + true_camera.focal_length_px, rel=0.02 + ) + + +def test_fit_clips_out_of_band_ground_truth_to_bound(): + # A ground truth placed outside a deliberately too-tight band is + # clipped to the bound rather than escaping it. + basis = make_basis() + true_camera = make_camera() + gen = make_generator(basis, true_camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=9 + ) + + # nominal is 10% off true, but the band only allows +/-1%. + nominal_focal_length = true_camera.focal_length_px * 1.10 + nominal_camera = CameraModel( + focal_length_px=nominal_focal_length, + position=true_camera.position, + orientation_deg=true_camera.orientation_deg, + ) + tight_tolerance = CameraModelTolerance( + focal_length_px=nominal_focal_length * 0.01, + position=(0.0, 0.0, 0.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + + fitter = ModalFitter(basis) + result = fitter.fit( + planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tight_tolerance + ) + + lower_bound = nominal_focal_length - tight_tolerance.focal_length_px + assert result.geometry["focal_length_px"] == pytest.approx(lower_bound, rel=1e-3) + + +def test_fit_recovers_offset_z_within_tolerance(): + basis = make_basis() + camera = make_camera() + gen = make_generator(basis, camera) + true_z_list = [0.35, 0.5, 0.65, 0.8] + offsets = {z: 0.01 for z in true_z_list} # nominal is 1 cm off true + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, + z_list=true_z_list, + image_shape=IMAGE_SHAPE, + nominal_z_offsets=offsets, + z_tolerance=0.03, + noise_std=1e-4, + seed=10, + ) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance()) + + for i, true_z in enumerate(true_z_list): + assert result.geometry[f"z_{i}"] == pytest.approx(true_z, abs=0.005) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `.venv/bin/pytest tests/test_fitting.py -k "not fit_auto" -q` +Expected: FAIL with `TypeError: fit() got an unexpected keyword argument 'camera'`, since `fit()` hasn't been rewritten yet. + +- [ ] **Step 3: Rewrite `ModalFitter.fit()` in `he11lib/fitting.py`** + +Replace the file's imports and the entire `fit` method (keep `generate_mode_shells`, `ModalFitter.__init__`, `fit_auto`, `_warm_start_coefficients`, `_bic`, `_estimate_uncertainty`, `_field_on_default_grid` for now -- `fit_auto`/`_bic` are rewritten in Task 5): + +```python +"""Joint nonlinear least-squares modal fit with automatic mode-set growth.""" + +from __future__ import annotations + +import warnings + +import numpy as np +from scipy.optimize import least_squares + +from .data import MeasurementPlane, ReconstructionResult, validate_planes +from .geometry import ( + CameraModel, + CameraModelTolerance, + GeometryCalibration, + camera_from_values, + camera_to_values, + tolerance_to_values, +) +from .modes import LGBasis +from .noise import NoiseEstimator + + +def generate_mode_shells(max_order: int) -> list[list[tuple[int, int]]]: + """Group candidate LG_{p,l} modes into shells of increasing order 2p+|l|.""" + shells: list[list[tuple[int, int]]] = [[] for _ in range(max_order + 1)] + for p in range(0, max_order + 1): + for l in range(-max_order, max_order + 1): + order = 2 * p + abs(l) + if order <= max_order: + shells[order].append((p, l)) + return shells + + +class ModalFitter: + """Fits LG mode coefficients, beam center/pointing, and geometry to measured planes.""" + + def __init__(self, basis: LGBasis, noise_estimator: NoiseEstimator | None = None): + self.basis = basis + self.noise_estimator = noise_estimator or NoiseEstimator() + + def fit( + self, + planes: list[MeasurementPlane], + modes: list[tuple[int, int]], + camera: CameraModel, + camera_tolerance: CameraModelTolerance, + initial_coefficients: dict[tuple[int, int], complex] | None = None, + initial_center: tuple[float, float] = (0.0, 0.0), + initial_pointing_deg: tuple[float, float] = (0.0, 0.0), + ) -> ReconstructionResult: + """Jointly fit complex coefficients for `modes` plus center/pointing/geometry. + + Every `CameraModel` field with a nonzero `camera_tolerance` entry, + and every plane whose `z_tolerance` is nonzero, is refined within + `[nominal - tolerance, nominal + tolerance]`; zero-tolerance fields + are held fixed at their nominal value. + """ + validate_planes(planes) + weights = [np.sqrt(self.noise_estimator.weights(p.flux)) for p in planes] + + camera_nominal = camera_to_values(camera) + camera_tol = tolerance_to_values(camera_tolerance) + free_camera_idx = [i for i, t in enumerate(camera_tol) if t > 0] + + free_z_idx = [i for i, p in enumerate(planes) if p.z_tolerance > 0] + + n_modes = len(modes) + n_always_free = 2 * n_modes + 4 # coefficients + center(2) + pointing(2) + + def pack_initial() -> np.ndarray: + x: list[float] = [] + for i, mode in enumerate(modes): + c = (initial_coefficients or {}).get(mode) + if c is None: + # Nonzero seed for every mode: starting a coefficient at + # exactly 0+0j sits at a flat/degenerate point for the + # optimizer and can prevent it from ever leaving zero. + c = 1.0 + 0j if i == 0 else 0.1 + 0.05j + x += [c.real, c.imag] + x += [initial_center[0], initial_center[1], initial_pointing_deg[0], initial_pointing_deg[1]] + for i in free_camera_idx: + x.append(camera_nominal[i]) + for i in free_z_idx: + x.append(planes[i].z) + return np.array(x, dtype=float) + + def pack_bounds() -> tuple[np.ndarray, np.ndarray]: + lower = [-np.inf] * n_always_free + upper = [np.inf] * n_always_free + for i in free_camera_idx: + lower.append(camera_nominal[i] - camera_tol[i]) + upper.append(camera_nominal[i] + camera_tol[i]) + for i in free_z_idx: + lower.append(planes[i].z - planes[i].z_tolerance) + upper.append(planes[i].z + planes[i].z_tolerance) + return np.array(lower), np.array(upper) + + def unpack(x: np.ndarray): + coeffs = {mode: complex(x[2 * i], x[2 * i + 1]) for i, mode in enumerate(modes)} + offset = 2 * n_modes + x0, y0, tilt_h_deg, tilt_v_deg = x[offset : offset + 4] + offset += 4 + + camera_values = list(camera_nominal) + for i in free_camera_idx: + camera_values[i] = x[offset] + offset += 1 + fitted_camera = camera_from_values(camera_values) + + z_values = [p.z for p in planes] + for i in free_z_idx: + z_values[i] = x[offset] + offset += 1 + + return coeffs, (x0, y0), (tilt_h_deg, tilt_v_deg), fitted_camera, z_values + + def plane_center(x0: float, y0: float, pointing_deg: tuple[float, float], z: float): + drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[0])) + drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[1])) + return x0 + drift_x, y0 + drift_y + + def model_flux_for_plane(plane, fitted_camera, z, coeffs, center0, pointing_deg): + calib = GeometryCalibration(fitted_camera) + x_grid, y_grid = calib.physical_coordinates(plane.flux.shape, z) + cx, cy = plane_center(center0[0], center0[1], pointing_deg, z) + field = self.basis.field_superposition(x_grid - cx, y_grid - cy, z, coeffs) + return np.abs(field) ** 2 + + def residuals(x: np.ndarray) -> np.ndarray: + coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(x) + parts = [] + for i, plane in enumerate(planes): + model_flux = model_flux_for_plane( + plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg + ) + parts.append(((plane.flux - model_flux) * weights[i]).ravel()) + return np.concatenate(parts) + + x0_vec = pack_initial() + lower, upper = pack_bounds() + # 'trf' + x_scale='jac' handles the very different natural + # magnitudes of these parameters (coefficients ~O(1), focal length + # ~O(1e3-1e4), angles ~O(1-90), z ~O(0.1-1)); plain 'lm' can + # terminate prematurely on 'xtol' because its unscaled step-size + # test is dominated by the largest parameters. 'lm' also doesn't + # support bounds, which the tolerance mechanism requires. + opt_result = least_squares( + residuals, x0_vec, method="trf", x_scale="jac", bounds=(lower, upper), max_nfev=5000 + ) + + coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(opt_result.x) + + total_power = sum(abs(c) ** 2 for c in coeffs.values()) + if total_power == 0: + total_power = 1.0 + purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()} + + centers = [ + plane_center(center0[0], center0[1], pointing_deg, z_values[i]) + for i in range(len(planes)) + ] + + geometry: dict[str, float] = dict(zip( + ( + "focal_length_px", "position_x", "position_y", "position_z", + "yaw_deg", "pitch_deg", "roll_deg", + "principal_point_x", "principal_point_y", + ), + camera_to_values(fitted_camera), + )) + for i in range(len(planes)): + geometry[f"z_{i}"] = z_values[i] + + residual_maps = [] + for i, plane in enumerate(planes): + model_flux = model_flux_for_plane( + plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg + ) + residual_maps.append(plane.flux - model_flux) + + coefficient_uncertainty = self._estimate_uncertainty(opt_result, modes, coeffs, total_power) + + reference_idx = min(range(len(planes)), key=lambda i: abs(z_values[i] - self.basis.z0)) + field_at_reference = self._field_on_default_grid(coeffs, z_values[reference_idx]) + + return ReconstructionResult( + purity=purity, + reconstructed_field=field_at_reference, + centers=centers, + pointing_angle_horizontal_deg=pointing_deg[0], + pointing_angle_vertical_deg=pointing_deg[1], + geometry=geometry, + residuals=residual_maps, + coefficient_uncertainty=coefficient_uncertainty, + used_phase_retrieval=False, + ) +``` + +Leave `fit_auto`, `_warm_start_coefficients`, `_bic`, `_estimate_uncertainty`, and `_field_on_default_grid` in place below this (unchanged for now; `fit_auto`/`_bic` are rewritten in Task 5 and will currently fail to call the new `fit` signature -- that's expected and fixed next task). + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `.venv/bin/pytest tests/test_fitting.py -k "not fit_auto" -q` +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add he11lib/fitting.py tests/test_fitting.py +git commit -m "$(cat <<'EOF' +Rewrite ModalFitter.fit around the CameraModel tolerance mechanism + +The optimizer's parameter vector is now built dynamically: LG +coefficients, per-plane center, and both pointing angles stay always +free; each CameraModel field and each plane's z join the fit (bounded to +its +/- tolerance) only when its paired tolerance is nonzero, and are +otherwise substituted as fixed constants. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 5: `ModalFitter.fit_auto()` and the degeneracy `UserWarning` + +**Files:** +- Modify: `he11lib/fitting.py` (`fit_auto`, `_bic`, new `_warn_if_degenerate`) +- Modify: `tests/test_fitting.py` (append `fit_auto` tests) + +**Interfaces:** +- Consumes: `ModalFitter.fit(planes, modes, camera, camera_tolerance, ...)` (Task 4). +- Produces: `ModalFitter.fit_auto(planes, camera, camera_tolerance, max_order=4, bic_improvement_threshold=10.0) -> ReconstructionResult`. Task 7's `BeamReconstructor.reconstruct` calls this exact signature. + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/test_fitting.py` (after the tests added in Task 4): + +```python +def test_fit_auto_does_not_add_modes_for_pure_fundamental(): + 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=4 + ) + + fitter = ModalFitter(basis) + result = fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=2) + + assert set(result.purity.keys()) == {(0, 0)} + + +def test_fit_auto_grows_to_include_second_mode(): + basis = make_basis() + camera = make_camera() + gen = make_generator(basis, camera) + true_coeffs = {(0, 0): 0.9 + 0j, (0, 1): 0.4 + 0j} + planes = gen.generate( + coefficients=true_coeffs, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=5 + ) + + fitter = ModalFitter(basis) + result = fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=2) + + assert (0, 1) in result.purity or (0, -1) in result.purity + + +def test_fit_auto_warns_when_free_geometry_params_exceed_plane_count(): + basis = make_basis() + camera = make_camera() + gen = make_generator(basis, camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, + z_list=Z_LIST, # 4 planes + image_shape=IMAGE_SHAPE, + z_tolerance=0.05, # +4 free z params + noise_std=1e-4, + seed=11, + ) + + # +7 free camera params (all but the 2 principal_point components) + + # 4 free z params = 11 free geometry params > 4 planes. + generous_tolerance = CameraModelTolerance( + focal_length_px=camera.focal_length_px * 0.05, + position=(0.01, 0.01, 0.01), + orientation_deg=(2.0, 2.0, 2.0), + principal_point=(0.0, 0.0), + ) + + fitter = ModalFitter(basis) + with pytest.warns(UserWarning, match="free camera/z geometry parameters"): + fitter.fit_auto(planes, camera=camera, camera_tolerance=generous_tolerance, max_order=1) + + +def test_fit_auto_does_not_warn_when_geometry_fully_fixed(): + 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=12 + ) + + fitter = ModalFitter(basis) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=1) +``` + +Add `import warnings` to the top of `tests/test_fitting.py` alongside the existing `numpy`/`pytest` imports. + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `.venv/bin/pytest tests/test_fitting.py -q` +Expected: FAIL — `test_fit_auto_*` tests fail with `TypeError: fit_auto() got an unexpected keyword argument 'camera'` (since `fit_auto` still has its old signature and calls `self.fit(planes, current_modes)` without the new required args). + +- [ ] **Step 3: Rewrite `fit_auto`, `_bic`, and add `_warn_if_degenerate` in `he11lib/fitting.py`** + +Replace the existing `fit_auto` and `_bic` methods with: + +```python + def fit_auto( + self, + planes: list[MeasurementPlane], + camera: CameraModel, + camera_tolerance: CameraModelTolerance, + max_order: int = 4, + bic_improvement_threshold: float = 10.0, + ) -> ReconstructionResult: + """Fit with automatic mode-set growth, capped at `max_order`.""" + validate_planes(planes) + self._warn_if_degenerate(planes, camera_tolerance) + shells = generate_mode_shells(max_order) + + current_modes = list(shells[0]) + best_result = self.fit(planes, current_modes, camera, camera_tolerance) + best_bic = self._bic(planes, best_result, current_modes, camera_tolerance) + + grew_until_cap = True + for shell in shells[1:]: + trial_modes = current_modes + shell + warm_start = self._warm_start_coefficients(best_result, current_modes) + trial_result = self.fit( + planes, trial_modes, camera, camera_tolerance, initial_coefficients=warm_start + ) + trial_bic = self._bic(planes, trial_result, trial_modes, camera_tolerance) + + if trial_bic < best_bic - bic_improvement_threshold: + current_modes = trial_modes + best_result = trial_result + best_bic = trial_bic + else: + grew_until_cap = False + break + + if grew_until_cap and len(shells) > 1: + warnings.warn( + "Automatic mode-set growth hit the configured max_order cap " + f"({max_order}) while still improving the fit; consider raising max_order.", + stacklevel=2, + ) + + return best_result + + def _warn_if_degenerate( + self, planes: list[MeasurementPlane], camera_tolerance: CameraModelTolerance + ) -> None: + """Warn when free camera+z geometry parameters exceed the plane count. + + With only a handful of planes, adding ~7-9 shared camera unknowns + plus one z correction per plane can be practically underdetermined + even though each plane contributes many pixels of data, because + those unknowns are *global* and only weakly constrained by subtle + keystone differences between planes. + """ + free_camera_count = sum(1 for t in tolerance_to_values(camera_tolerance) if t > 0) + free_z_count = sum(1 for p in planes if p.z_tolerance > 0) + free_geometry_count = free_camera_count + free_z_count + + if free_geometry_count > len(planes): + warnings.warn( + f"{free_geometry_count} free camera/z geometry parameters " + f"(from nonzero tolerances) but only {len(planes)} measurement " + "planes; the joint fit may be practically underdetermined. " + "Consider tightening CameraModelTolerance / " + "MeasurementPlane.z_tolerance.", + UserWarning, + stacklevel=3, + ) +``` + +Replace the existing `_bic` method with: + +```python + def _bic( + self, + planes: list[MeasurementPlane], + result: ReconstructionResult, + modes: list[tuple[int, int]], + camera_tolerance: CameraModelTolerance, + ) -> float: + chi2 = sum( + np.sum((r * np.sqrt(self.noise_estimator.weights(p.flux))) ** 2) + for r, p in zip(result.residuals, planes) + ) + n_data = sum(p.flux.size for p in planes) + free_camera_count = sum(1 for t in tolerance_to_values(camera_tolerance) if t > 0) + free_z_count = sum(1 for p in planes if p.z_tolerance > 0) + n_params = 2 * len(modes) + 4 + free_camera_count + free_z_count + return float(chi2 + n_params * np.log(n_data)) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `.venv/bin/pytest tests/test_fitting.py -q` +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add he11lib/fitting.py tests/test_fitting.py +git commit -m "$(cat <<'EOF' +Update fit_auto for CameraModel and warn on underdetermined geometry fits + +fit_auto now threads camera/camera_tolerance through to fit and _bic +(whose parameter count must include any free camera/z unknowns). Emits a +UserWarning, not an error, when free camera+z geometry parameters exceed +the number of measurement planes -- a new documented degeneracy pitfall. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 6: `PhaseRetriever.retrieve()` update + +**Files:** +- Modify: `he11lib/phase_retrieval.py` +- Modify: `tests/test_phase_retrieval.py` + +**Interfaces:** +- Consumes: `CameraModel`, `GeometryCalibration.physical_coordinates(image_shape, z)` (Task 1); `SyntheticBeamGenerator(basis, camera)` (Task 3). +- Produces: `PhaseRetriever.retrieve(planes, camera, max_iterations=200) -> PhaseRetrievalResult`. Task 7's `BeamReconstructor._phase_retrieval_fallback` calls this exact signature. + +- [ ] **Step 1: Write the failing tests** + +Replace `tests/test_phase_retrieval.py`'s `retrieve`-related tests (the module docstring, imports, `make_basis`/`make_grid`, and the two `propagate_*` tests stay unchanged; only `make_grid`'s camera plumbing and the two `test_retrieve_*` tests change). Replace the whole file with: + +```python +import numpy as np +import pytest + +from he11lib.geometry import CameraModel +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 +CAMERA_DISTANCE = 5.0 +IMAGE_SHAPE = (121, 121) + + +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 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) + + 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() + 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=0 + ) + + retriever = PhaseRetriever(wavelength=WAVELENGTH) + result = retriever.retrieve(planes, camera, max_iterations=100) + + dx = float(result.x[0, 1] - result.x[0, 0]) + coeffs = basis.project( + result.field, result.x, result.y, dx, 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() + camera = make_camera() + gen = SyntheticBeamGenerator(basis=basis, camera=camera) + 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, + image_shape=IMAGE_SHAPE, + center=true_center, + noise_std=1e-5, + seed=1, + ) + + retriever = PhaseRetriever(wavelength=WAVELENGTH) + result = retriever.retrieve(planes, camera, 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) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `.venv/bin/pytest tests/test_phase_retrieval.py -q` +Expected: FAIL — `test_retrieve_*` fail with `TypeError: retrieve() got an unexpected keyword argument 'viewing_angle_deg'` or a positional-argument mismatch, since `retrieve` hasn't been updated yet. (`test_propagate_*` should already pass, unaffected by this change.) + +- [ ] **Step 3: Update `he11lib/phase_retrieval.py`** + +Change the import line: + +```python +from .geometry import GeometryCalibration +``` + +to: + +```python +from .geometry import CameraModel, GeometryCalibration +``` + +Replace the `retrieve` method's signature and body's calibration lines: + +```python + def retrieve( + self, + planes: list[MeasurementPlane], + pixel_scale: float | None = None, + viewing_angle_deg: float | None = None, + max_iterations: int = 200, + ) -> PhaseRetrievalResult: + """Run Gerchberg-Saxton phase retrieval across the given planes. + + Planes must share the same known (or overridden) pixel_scale and + viewing_angle_deg, since all planes are propagated on one common + physical grid. + """ + validate_planes(planes) + ordered = sorted(planes, key=lambda p: p.z) + + x, y = GeometryCalibration(ordered[0]).physical_coordinates( + pixel_scale=pixel_scale, viewing_angle_deg=viewing_angle_deg + ) + dx = float(x[0, 1] - x[0, 0]) +``` + +with: + +```python + def retrieve( + self, + planes: list[MeasurementPlane], + camera: CameraModel, + max_iterations: int = 200, + ) -> PhaseRetrievalResult: + """Run Gerchberg-Saxton phase retrieval across the given planes. + + All planes are propagated on one common physical grid, derived + from `camera` at the smallest-z plane's depth (an existing + approximation: the shared grid is only exact at that one z, since + other planes may sit at a slightly different true depth under true + perspective projection). + """ + validate_planes(planes) + ordered = sorted(planes, key=lambda p: p.z) + + x, y = GeometryCalibration(camera).physical_coordinates(ordered[0].flux.shape, ordered[0].z) + dx = float(x[0, 1] - x[0, 0]) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `.venv/bin/pytest tests/test_phase_retrieval.py -q` +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add he11lib/phase_retrieval.py tests/test_phase_retrieval.py +git commit -m "$(cat <<'EOF' +Thread CameraModel through PhaseRetriever.retrieve + +retrieve() now takes a CameraModel directly instead of the removed +pixel_scale/viewing_angle_deg override kwargs, matching the rest of the +pipeline's shared-camera convention. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 7: `BeamReconstructor` update + +**Files:** +- Modify: `he11lib/reconstruct.py` +- Modify: `tests/test_reconstruct.py` (full rewrite) + +**Interfaces:** +- Consumes: `CameraModel`, `CameraModelTolerance`, `GeometryCalibration.effective_pixel_scale` (Task 1); `ModalFitter.fit_auto(planes, camera, camera_tolerance, max_order=...)` (Task 5); `PhaseRetriever.retrieve(planes, camera, ...)` (Task 6); `SyntheticBeamGenerator(basis, camera)` (Task 3). +- Produces: `BeamReconstructor(w0, z0, wavelength, camera, camera_tolerance, max_order=4, noise_estimator=None, deconvolver=None, force_phase_retrieval=False, phase_retrieval_residual_threshold=None)`. This is the top-level public constructor documented in Task 8/9. + +- [ ] **Step 1: Write the failing tests** + +Replace `tests/test_reconstruct.py` entirely with: + +```python +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) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `.venv/bin/pytest tests/test_reconstruct.py -q` +Expected: FAIL with `TypeError: __init__() missing 2 required positional arguments: 'camera' and 'camera_tolerance'`, since `BeamReconstructor` hasn't been updated yet. + +- [ ] **Step 3: Update `he11lib/reconstruct.py`** + +Replace the imports: + +```python +from .data import MeasurementPlane, ReconstructionResult, validate_planes +from .deconvolution import DiffusionDeconvolver +from .fitting import ModalFitter, generate_mode_shells +from .modes import LGBasis +from .noise import NoiseEstimator +from .phase_retrieval import PhaseRetriever +``` + +with: + +```python +from .data import MeasurementPlane, ReconstructionResult, validate_planes +from .deconvolution import DiffusionDeconvolver +from .fitting import ModalFitter, generate_mode_shells +from .geometry import CameraModel, CameraModelTolerance, GeometryCalibration +from .modes import LGBasis +from .noise import NoiseEstimator +from .phase_retrieval import PhaseRetriever +``` + +Replace `__init__`: + +```python + def __init__( + self, + w0: float, + z0: float, + wavelength: float, + max_order: int = 4, + noise_estimator: NoiseEstimator | None = None, + deconvolver: DiffusionDeconvolver | None = None, + force_phase_retrieval: bool = False, + phase_retrieval_residual_threshold: float | None = None, + ): + self.basis = LGBasis(w0=w0, z0=z0, wavelength=wavelength) + self.wavelength = wavelength + self.max_order = max_order + self.noise_estimator = noise_estimator or NoiseEstimator() + self.deconvolver = deconvolver + self.force_phase_retrieval = force_phase_retrieval + self.phase_retrieval_residual_threshold = phase_retrieval_residual_threshold +``` + +with: + +```python + def __init__( + self, + w0: float, + z0: float, + wavelength: float, + camera: CameraModel, + camera_tolerance: CameraModelTolerance, + max_order: int = 4, + noise_estimator: NoiseEstimator | None = None, + deconvolver: DiffusionDeconvolver | None = None, + force_phase_retrieval: bool = False, + phase_retrieval_residual_threshold: float | None = None, + ): + self.basis = LGBasis(w0=w0, z0=z0, wavelength=wavelength) + self.wavelength = wavelength + self.camera = camera + self.camera_tolerance = camera_tolerance + self.max_order = max_order + self.noise_estimator = noise_estimator or NoiseEstimator() + self.deconvolver = deconvolver + self.force_phase_retrieval = force_phase_retrieval + self.phase_retrieval_residual_threshold = phase_retrieval_residual_threshold +``` + +Also update the class docstring's parameter list to mention `camera`/`camera_tolerance` and remove the stale `pixel_scale` reference: + +```python + Parameters + ---------- + w0, z0, wavelength : known reference beam parameters (see `LGBasis`). + camera : nominal shared CameraModel (position/orientation/intrinsics). + camera_tolerance : per-field +/- refinement bound for `camera`; a + zero-tolerance field is held fixed at its nominal value. + max_order : cap on automatic candidate-mode-set growth (see + `ModalFitter.fit_auto`), and also the mode set used to project the + phase-retrieval fallback's recovered field onto the LG basis. + noise_estimator : shared noise model; defaults to `NoiseEstimator()`. + deconvolver : if given, each plane's flux is deblurred (using + `GeometryCalibration(camera).effective_pixel_scale`) before fitting. + force_phase_retrieval : if True, always run the phase-retrieval fallback + instead of the modal fit. + phase_retrieval_residual_threshold : if set (and `force_phase_retrieval` + is False), the phase-retrieval fallback runs automatically whenever + the modal fit's noise-weighted RMS residual exceeds this value. + """ +``` + +Replace `reconstruct`: + +```python + def reconstruct(self, planes: list[MeasurementPlane]) -> ReconstructionResult: + """Run the full pipeline and return a `ReconstructionResult`.""" + validate_planes(planes) + planes = self._deconvolve(planes) + + fitter = ModalFitter(self.basis, self.noise_estimator) + result = fitter.fit_auto(planes, max_order=self.max_order) + + if self.force_phase_retrieval or self._residual_too_high(result, planes): + result = self._phase_retrieval_fallback(planes) + + return result +``` + +with: + +```python + def reconstruct(self, planes: list[MeasurementPlane]) -> ReconstructionResult: + """Run the full pipeline and return a `ReconstructionResult`.""" + validate_planes(planes) + planes = self._deconvolve(planes) + + fitter = ModalFitter(self.basis, self.noise_estimator) + result = fitter.fit_auto( + planes, self.camera, self.camera_tolerance, max_order=self.max_order + ) + + if self.force_phase_retrieval or self._residual_too_high(result, planes): + result = self._phase_retrieval_fallback(planes) + + return result +``` + +Replace `_deconvolve`: + +```python + def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]: + if self.deconvolver is None: + return planes + deblurred = [] + for plane in planes: + if plane.pixel_scale is None: + raise ValueError( + "Deconvolution requires a known pixel_scale on every MeasurementPlane." + ) + flux = self.deconvolver.deconvolve(plane.flux, plane.pixel_scale) + deblurred.append(replace(plane, flux=flux)) + return deblurred +``` + +with: + +```python + def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]: + if self.deconvolver is None: + return planes + calib = GeometryCalibration(self.camera) + deblurred = [] + for plane in planes: + pixel_scale = calib.effective_pixel_scale(plane.flux.shape, plane.z) + flux = self.deconvolver.deconvolve(plane.flux, pixel_scale) + deblurred.append(replace(plane, flux=flux)) + return deblurred +``` + +Replace `_phase_retrieval_fallback`: + +```python + def _phase_retrieval_fallback( + self, planes: list[MeasurementPlane] + ) -> ReconstructionResult: + retriever = PhaseRetriever(self.wavelength) + pr_result = retriever.retrieve(planes) + + modes = [mode for shell in generate_mode_shells(self.max_order) for mode in shell] + dx = float(pr_result.x[0, 1] - pr_result.x[0, 0]) + coeffs = self.basis.project(pr_result.field, pr_result.x, pr_result.y, dx, pr_result.z, modes) + + total_power = sum(abs(c) ** 2 for c in coeffs.values()) + if total_power == 0: + total_power = 1.0 + purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()} + + return ReconstructionResult( + purity=purity, + reconstructed_field=pr_result.field, + centers=[pr_result.center for _ in planes], + pointing_angle_deg=float("nan"), + geometry={}, + residuals=[], + coefficient_uncertainty={mode: float("nan") for mode in modes}, + used_phase_retrieval=True, + ) +``` + +with: + +```python + def _phase_retrieval_fallback( + self, planes: list[MeasurementPlane] + ) -> ReconstructionResult: + retriever = PhaseRetriever(self.wavelength) + pr_result = retriever.retrieve(planes, self.camera) + + modes = [mode for shell in generate_mode_shells(self.max_order) for mode in shell] + dx = float(pr_result.x[0, 1] - pr_result.x[0, 0]) + coeffs = self.basis.project(pr_result.field, pr_result.x, pr_result.y, dx, pr_result.z, modes) + + total_power = sum(abs(c) ** 2 for c in coeffs.values()) + if total_power == 0: + total_power = 1.0 + purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()} + + return ReconstructionResult( + purity=purity, + reconstructed_field=pr_result.field, + centers=[pr_result.center for _ in planes], + pointing_angle_horizontal_deg=float("nan"), + pointing_angle_vertical_deg=float("nan"), + geometry={}, + residuals=[], + coefficient_uncertainty={mode: float("nan") for mode in modes}, + used_phase_retrieval=True, + ) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `.venv/bin/pytest tests/test_reconstruct.py -q` +Expected: PASS. (If `test_reconstruct_recovers_camera_and_z_offset_from_nominal` is flaky on the exact `abs=`/`rel=` bounds given the chosen synthetic parameters, widen the tolerance rather than changing the fit logic -- per this plan's Global Constraints.) + +- [ ] **Step 5: Commit** + +```bash +git add he11lib/reconstruct.py tests/test_reconstruct.py +git commit -m "$(cat <<'EOF' +Wire CameraModel/CameraModelTolerance through BeamReconstructor + +BeamReconstructor now requires camera/camera_tolerance, threading them +into ModalFitter.fit_auto and PhaseRetriever.retrieve, and uses +GeometryCalibration.effective_pixel_scale for deconvolution instead of +the removed MeasurementPlane.pixel_scale. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 8: Update `docs/api.md` + +**Files:** +- Modify: `docs/api.md` (full rewrite of the affected sections) + +**Interfaces:** +- Consumes: the final public signatures from Tasks 1-7 (`CameraModel`, `CameraModelTolerance`, `GeometryCalibration`, `MeasurementPlane`, `ReconstructionResult`, `SyntheticBeamGenerator`, `ModalFitter`, `PhaseRetriever`, `BeamReconstructor`). +- Produces: nothing consumed by later tasks (docs-only; Task 9's example script is written independently against the same source signatures, not against this doc). + +This is a documentation-only task — no tests, no code. Since there's no failing-test cycle for prose docs, the "step" here is: rewrite, then a lightweight self-check that no stale identifiers remain. + +- [ ] **Step 1: Rewrite the Quick start section** + +In `docs/api.md`, replace: + +```markdown +## Quick start + +```python +from he11lib import BeamReconstructor, MeasurementPlane + +# planes: a list of >=3 MeasurementPlane objects built from your own +# flux arrays (see MeasurementPlane below). +reconstructor = BeamReconstructor(w0=5e-3, z0=0.5, wavelength=1.76e-3) +result = reconstructor.reconstruct(planes) + +for mode, (power_fraction, phase_rad) in result.purity.items(): + print(mode, power_fraction, phase_rad) +``` +``` + +with: + +```markdown +## Quick start + +```python +from he11lib import ( + BeamReconstructor, + CameraModel, + CameraModelTolerance, + MeasurementPlane, +) + +# planes: a list of >=3 MeasurementPlane objects built from your own +# flux arrays (see MeasurementPlane below). + +# Nominal camera pose/intrinsics from calibration; every field here is +# refined jointly with the mode fit because its tolerance is nonzero. +camera = CameraModel( + focal_length_px=2000.0, + position=(0.0, 0.0, -2.0), + orientation_deg=(0.0, 0.0, 0.0), +) +camera_tolerance = CameraModelTolerance( + focal_length_px=20.0, + position=(0.01, 0.01, 0.05), + orientation_deg=(2.0, 2.0, 2.0), +) + +reconstructor = BeamReconstructor( + w0=5e-3, z0=0.5, wavelength=1.76e-3, + camera=camera, camera_tolerance=camera_tolerance, +) +result = reconstructor.reconstruct(planes) + +for mode, (power_fraction, phase_rad) in result.purity.items(): + print(mode, power_fraction, phase_rad) +``` +``` + +- [ ] **Step 2: Rewrite the `data` section** + +Replace: + +```markdown +## `data` — `MeasurementPlane`, `ReconstructionResult` + +### `MeasurementPlane(flux, z, pixel_scale=None, viewing_angle_deg=None, label=None)` + +One measurement: a 2D flux array plus its acquisition metadata. + +- `flux` — 2D `np.ndarray` of flux values. Dead-pixel correction, background + subtraction, and saturation clipping are assumed already handled upstream. +- `z` — nominal distance from the output window, in meters. Must be `> 0`. +- `pixel_scale` — known meters/pixel, or `None` if unknown (then jointly + fit by `ModalFitter`/`BeamReconstructor`). +- `viewing_angle_deg` — known camera viewing angle relative to the beam + axis, in degrees, or `None` if unknown (also jointly fit). +- `label` — optional human-readable identifier. +``` + +with: + +```markdown +## `data` — `MeasurementPlane`, `ReconstructionResult` + +### `MeasurementPlane(flux, z, z_tolerance=0.0, label=None)` + +One measurement: a 2D flux array plus its acquisition metadata. + +- `flux` — 2D `np.ndarray` of flux values. Dead-pixel correction, background + subtraction, and saturation clipping are assumed already handled upstream. +- `z` — nominal distance from the output window, in meters. Must be `> 0`. +- `z_tolerance` — `+/-` bound, in meters, around the nominal `z` within + which the true distance is jointly refined by `ModalFitter`. Must be + `>= 0`; `0` (the default) means `z` is trusted exactly and held fixed. +- `label` — optional human-readable identifier. + +Per-plane camera geometry (`pixel_scale`/`viewing_angle_deg`) no longer +lives on `MeasurementPlane` — camera pose/intrinsics are a single shared +`CameraModel` for the whole reconstruction (see `geometry` below). +``` + +- [ ] **Step 3: Update the `ReconstructionResult` section** + +Replace: + +```markdown +- `pointing_angle_deg: float` — fitted shared beam pointing angle (tilt). +- `geometry: dict[str, float]` — geometry parameters used or fitted (keys + `pixel_scale_{i}`, `viewing_angle_deg_{i}` per plane index `i`). +``` + +with: + +```markdown +- `pointing_angle_horizontal_deg`, `pointing_angle_vertical_deg: float` — + fitted shared beam pointing (tilt) angles, independent horizontal and + vertical. +- `geometry: dict[str, float]` — geometry parameters used or fitted: the 9 + `CameraModel` field names from `he11lib.geometry.CAMERA_FIELD_NAMES` + (`focal_length_px`, `position_x`, `position_y`, `position_z`, `yaw_deg`, + `pitch_deg`, `roll_deg`, `principal_point_x`, `principal_point_y`), plus + `z_{i}` per plane index `i` (that plane's fitted/held distance). +``` + +- [ ] **Step 4: Rewrite the `geometry` section** + +Replace: + +```markdown +## `geometry` — `GeometryCalibration` + +`GeometryCalibration(plane)` wraps a single `MeasurementPlane` and resolves +its pixel-to-physical-coordinate mapping. + +- `pixel_scale_known` / `viewing_angle_known` — `bool` properties. +- `physical_coordinates(pixel_scale=None, viewing_angle_deg=None)` — + returns `(x, y)` physical coordinate grids matching the plane's flux + shape. Values known on the `MeasurementPlane` take precedence over the + `override` arguments; raises `ValueError` if a value is neither known nor + overridden. +``` + +with: + +```markdown +## `geometry` — `CameraModel`, `CameraModelTolerance`, `GeometryCalibration` + +### `CameraModel(focal_length_px, position, orientation_deg, principal_point=(0.0, 0.0))` + +A nominal pinhole camera pose/intrinsics shared across every plane in one +reconstruction. Always a point estimate — never trusted as exact by +itself; trust is expressed via the paired `CameraModelTolerance`. + +- `focal_length_px` — focal length in pixel units. +- `position` — `(x, y, z)` camera position in the beam-axis world frame, + meters; `z=0` is the output window. +- `orientation_deg` — `(yaw, pitch, roll)`, degrees. All-zero means the + boresight is normal to every `z=const` target plane with no in-plane + rotation. +- `principal_point` — `(px, px)` offset from the frame center. + +### `CameraModelTolerance(focal_length_px, position, orientation_deg, principal_point=(0.0, 0.0))` + +Per-field `+/-` refinement bound, same shape as `CameraModel`. Every field +must be `>= 0` (raises `ValueError` otherwise). A field's tolerance of `0` +holds that `CameraModel` field fixed at its nominal value during fitting; +`> 0` lets `ModalFitter` refine it within `[nominal - tolerance, nominal + +tolerance]`. + +### `GeometryCalibration(camera)` + +Wraps a `CameraModel` and resolves pixel <-> physical coordinate mappings +via true pinhole projection (not a uniform affine/cosine approximation). + +- `pixel_coordinates(x, y, z) -> (row, col)` — forward-projects physical + `(x, y)` at depth `z` to pixel coordinates. Raises `ValueError` if the + point is behind the camera (`Z_cam <= 0`). +- `physical_coordinates(image_shape, z) -> (x, y)` — inverse-projects every + pixel in a frame of `image_shape` to physical `(x, y)` on the `z=const` + plane, via ray-plane intersection (this is what produces genuine + keystoning — non-uniform spacing across the frame — for tilted/off-axis + poses). Raises `ValueError` if the plane is edge-on to or behind the + camera. +- `effective_pixel_scale(image_shape, z) -> float` — a single isotropic + meters/pixel figure (finite-difference approximation at the frame + center), for callers like `DiffusionDeconvolver` that assume one + isotropic pixel-space kernel. + +### `CAMERA_FIELD_NAMES`, `camera_to_values`, `tolerance_to_values`, `camera_from_values` + +Module-level helpers used internally by `ModalFitter` to flatten/unflatten +`CameraModel`/`CameraModelTolerance` into the optimizer's parameter vector. +Not usually needed by application code, but exported for advanced use +(e.g. inspecting `CAMERA_FIELD_NAMES` to interpret `ReconstructionResult.geometry` keys). +``` + +- [ ] **Step 5: Update the `deconvolution` section's viewing-angle note** + +Replace: + +```markdown +Note: the blur/deconvolution kernel is isotropic in pixel space. If a +plane has a nonzero `viewing_angle_deg`, its `x` and `y` pixel axes have +different physical scales (see `SyntheticBeamGenerator` below), so +deconvolution is only exact for `viewing_angle_deg == 0`; at oblique +angles it is an approximation. +``` + +with: + +```markdown +Note: the blur/deconvolution kernel is isotropic in pixel space. A tilted +or off-axis `CameraModel` produces a pixel scale that varies across the +frame and between `x`/`y` (keystoning), so `deconvolve` uses +`GeometryCalibration.effective_pixel_scale` — a single isotropic +approximation evaluated at the frame center. This is exact only for an +on-axis, untilted camera; at oblique poses it is an accepted +approximation (see `CLAUDE.md`). +``` + +- [ ] **Step 6: Rewrite the `synthetic` section** + +Replace: + +```markdown +## `synthetic` — `SyntheticBeamGenerator` + +`SyntheticBeamGenerator(basis, image_shape, pixel_scale)` — forward model +used to validate the pipeline against known ground truth, and to evaluate +experimental design (e.g. "would these distances separate my modes?"). +`pixel_scale` is the physical pixel size, in meters, along the non-tilted +`y` axis; the `x` axis is compressed by `1/cos(viewing_angle_deg)` to model +an oblique camera view. + +- `generate(coefficients, z_list, *, center=(0, 0), pointing_angle_deg=0.0, viewing_angle_deg=0.0, noise_std=0.0, seed=None)` + — returns one `MeasurementPlane` per `z` in `z_list`. The beam's + transverse center drifts linearly with `z` according to + `pointing_angle_deg`, starting from `center` at `z0`. +``` + +with: + +```markdown +## `synthetic` — `SyntheticBeamGenerator` + +`SyntheticBeamGenerator(basis, camera)` — forward model used to validate +the pipeline against known ground truth, and to evaluate experimental +design. `camera` is the ground-truth `CameraModel` (position/orientation/ +intrinsics) used to render each plane via true perspective projection. + +- `generate(coefficients, z_list, image_shape, *, center=(0.0, 0.0), pointing_angle_horizontal_deg=0.0, pointing_angle_vertical_deg=0.0, z_tolerance=0.0, nominal_z_offsets=None, noise_std=0.0, seed=None) -> list[MeasurementPlane]` + — returns one `MeasurementPlane` per (true) `z` in `z_list`. The beam's + transverse center drifts linearly with `z` according to the two + independent pointing angles, starting from `center` at the basis's + `z0`. `nominal_z_offsets`, if given, maps a true `z` to an offset + applied to that plane's *nominal* `z` — letting a reconstruction be + tested against a deliberately-offset nominal input while the plane's + flux is still rendered at the true `z`. Every resulting plane shares + `z_tolerance`. +``` + +- [ ] **Step 7: Update the `fitting` section** + +Replace: + +```markdown +### `ModalFitter(basis, noise_estimator=None)` + +Core reconstruction path: a joint nonlinear least-squares fit of complex LG +coefficients, beam center/pointing, and (if unknown) geometry. + +- `fit(planes, modes, initial_coefficients=None, initial_center=(0.0, 0.0), initial_tilt_deg=(0.0, 0.0), initial_pixel_scale=None, initial_viewing_angle_deg=0.0) -> ReconstructionResult` + — fits exactly the given candidate `modes`. +- `fit_auto(planes, max_order=4, bic_improvement_threshold=10.0) -> ReconstructionResult` + — starts from `LG_00` and grows the candidate mode set shell-by-shell + (via `generate_mode_shells`), stopping once BIC no longer improves by + more than `bic_improvement_threshold`, capped at `max_order`. Emits a + `UserWarning` (does not raise) if the cap is reached while the fit is + still improving. +``` + +with: + +```markdown +### `ModalFitter(basis, noise_estimator=None)` + +Core reconstruction path: a joint nonlinear least-squares fit of complex LG +coefficients, beam center/pointing, and any nonzero-tolerance camera/`z` +geometry. + +- `fit(planes, modes, camera, camera_tolerance, initial_coefficients=None, initial_center=(0.0, 0.0), initial_pointing_deg=(0.0, 0.0)) -> ReconstructionResult` + — fits exactly the given candidate `modes`. Every `CameraModel` field + with a nonzero `camera_tolerance` entry, and every plane whose + `z_tolerance` is nonzero, is refined within `[nominal - tolerance, + nominal + tolerance]`; zero-tolerance fields are held fixed at their + nominal value. +- `fit_auto(planes, camera, camera_tolerance, max_order=4, bic_improvement_threshold=10.0) -> ReconstructionResult` + — starts from `LG_00` and grows the candidate mode set shell-by-shell + (via `generate_mode_shells`), stopping once BIC no longer improves by + more than `bic_improvement_threshold`, capped at `max_order`. Emits a + `UserWarning` (does not raise) if the cap is reached while the fit is + still improving, or if the number of free camera+`z` parameters is large + relative to the number of planes (see `CLAUDE.md`'s degeneracy pitfall). +``` + +- [ ] **Step 8: Update the `phase_retrieval` section** + +Replace: + +```markdown +### `PhaseRetriever(wavelength)` + +- `retrieve(planes, pixel_scale=None, viewing_angle_deg=None, max_iterations=200) -> PhaseRetrievalResult` + — multi-plane Gerchberg-Saxton phase retrieval: propagates a trial + complex field back and forth between planes, enforcing the measured + amplitude (`sqrt(flux)`) at each plane, without assuming a finite mode + basis. +``` + +with: + +```markdown +### `PhaseRetriever(wavelength)` + +- `retrieve(planes, camera, max_iterations=200) -> PhaseRetrievalResult` + — multi-plane Gerchberg-Saxton phase retrieval: propagates a trial + complex field back and forth between planes, enforcing the measured + amplitude (`sqrt(flux)`) at each plane, without assuming a finite mode + basis. All planes are propagated on one common physical grid, derived + from `camera` at the smallest-`z` plane's depth. +``` + +- [ ] **Step 9: Update the `reconstruct` section** + +Replace: + +```markdown +## `reconstruct` — `BeamReconstructor` + +`BeamReconstructor(w0, z0, wavelength, max_order=4, noise_estimator=None, deconvolver=None, force_phase_retrieval=False, phase_retrieval_residual_threshold=None)` + +High-level orchestrator wiring together the full pipeline: optional +diffusion deblurring → `ModalFitter.fit_auto` → optional +`PhaseRetriever` fallback. + +- `reconstruct(planes) -> ReconstructionResult` + 1. Validates `planes` (see `validate_planes`). + 2. If `deconvolver` is set, deblurs each plane (raises `ValueError` if a + plane's `pixel_scale` isn't known). + 3. Runs `ModalFitter(basis, noise_estimator).fit_auto(planes, max_order)`. + 4. Runs the `PhaseRetriever` fallback instead, projecting its recovered + field onto all modes up to `max_order`, if `force_phase_retrieval` is + `True`, or if `phase_retrieval_residual_threshold` is set and the + modal fit's noise-weighted RMS residual exceeds it. In that case + `result.residuals` is empty and `coefficient_uncertainty` is `NaN` + per mode (phase retrieval doesn't produce a fit covariance). +``` + +with: + +```markdown +## `reconstruct` — `BeamReconstructor` + +`BeamReconstructor(w0, z0, wavelength, camera, camera_tolerance, max_order=4, noise_estimator=None, deconvolver=None, force_phase_retrieval=False, phase_retrieval_residual_threshold=None)` + +High-level orchestrator wiring together the full pipeline: optional +diffusion deblurring → `ModalFitter.fit_auto` → optional +`PhaseRetriever` fallback. `camera`/`camera_tolerance` are the nominal +shared `CameraModel` and its per-field refinement bounds for this +reconstruction. + +- `reconstruct(planes) -> ReconstructionResult` + 1. Validates `planes` (see `validate_planes`). + 2. If `deconvolver` is set, deblurs each plane using + `GeometryCalibration(camera).effective_pixel_scale(plane.flux.shape, plane.z)`. + 3. Runs `ModalFitter(basis, noise_estimator).fit_auto(planes, camera, camera_tolerance, max_order)`. + 4. Runs the `PhaseRetriever` fallback instead, projecting its recovered + field onto all modes up to `max_order`, if `force_phase_retrieval` is + `True`, or if `phase_retrieval_residual_threshold` is set and the + modal fit's noise-weighted RMS residual exceeds it. In that case + `result.residuals` is empty, `coefficient_uncertainty` is `NaN` per + mode, `geometry` is empty, and both pointing-angle fields are `NaN` + (phase retrieval doesn't fit geometry/pointing or produce a fit + covariance). +``` + +- [ ] **Step 10: Grep for stale identifiers and fix any remaining hits** + +Run: + +```bash +grep -n "pixel_scale\|viewing_angle_deg\|pointing_angle_deg[^_]" docs/api.md +``` + +Expected: no output (every remaining `pixel_scale` mention, if any, should +only be in the `deconvolution` section describing `DiffusionDeconvolver`'s +own `pixel_scale` parameter, which is unchanged — inspect any hits and +confirm they're that case, not a stale `MeasurementPlane`/`GeometryCalibration` +reference). + +- [ ] **Step 11: Commit** + +```bash +git add docs/api.md +git commit -m "$(cat <<'EOF' +Update docs/api.md for the CameraModel geometry redesign + +Documents CameraModel/CameraModelTolerance, the rewritten +GeometryCalibration, z_tolerance, the two pointing angles, and every +downstream signature change (ModalFitter, SyntheticBeamGenerator, +PhaseRetriever, BeamReconstructor) introduced by the redesign. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 9: Rewrite `examples/full_pipeline_example.py` + +**Files:** +- Modify: `examples/full_pipeline_example.py` (full rewrite) + +**Interfaces:** +- Consumes: `CameraModel`, `CameraModelTolerance`, `SyntheticBeamGenerator(basis, camera)`, `.generate(coefficients, z_list, image_shape, *, center, pointing_angle_horizontal_deg, pointing_angle_vertical_deg, z_tolerance, noise_std, seed)`, `GeometryCalibration(camera).effective_pixel_scale(shape, z)`, `BeamReconstructor(w0, z0, wavelength, camera, camera_tolerance, max_order, deconvolver)`, `result.pointing_angle_horizontal_deg`/`pointing_angle_vertical_deg`. +- Produces: nothing consumed by later tasks (this is the last code-facing deliverable; Task 10 only touches `CLAUDE.md` prose). + +This script has no dedicated pytest file — its own "test" is running it end-to-end and inspecting the printed output, per the plan's Global Constraints (this is the one step in the plan involving executing code, not just drafting it). + +- [ ] **Step 1: Rewrite `examples/full_pipeline_example.py`** + +Replace the whole file with: + +```python +"""End-to-end demonstration of the he11lib reconstruction pipeline. + +Simulates a gyrotron beam that is mostly the LG_00 fundamental mode with a +small admixture of LG_01, viewed by a thermal camera at four distances from +the output window through a tilted, off-axis pinhole `CameraModel`. The +camera has an unknown transverse offset/pointing (two independent tilt +angles) and adds sensor noise; the target also has some thermal-diffusion +blur that we correct for. The nominal camera pose and each plane's nominal +`z` are deliberately offset from the (unknown-to-the-reconstructor) ground +truth, simulating realistic calibration/measurement error, and are jointly +refined by the fit within their tolerances. We then reconstruct the mode +purity, beam center/pointing, camera pose, and per-plane z, and plot the +diagnostics. + +Run with: + + python examples/full_pipeline_example.py +""" + +from __future__ import annotations + +import matplotlib.pyplot as plt + +from he11lib import ( + BeamReconstructor, + CameraModel, + CameraModelTolerance, + DiffusionDeconvolver, + GeometryCalibration, + LGBasis, + SyntheticBeamGenerator, + plot_center_trace, + plot_mode_purity, + plot_residuals, +) + +# --- Known reference beam parameters (from the gyrotron/mode-converter design) --- +W0 = 5e-3 # reference waist radius, meters +Z0 = 0.5 # reference waist location, meters from the output window +WAVELENGTH = 1.76e-3 # radiation wavelength, meters (e.g. a 170 GHz gyrotron) + +# --- Ground truth for the synthetic beam (unknown to the reconstructor) --- +TRUE_COEFFICIENTS = {(0, 0): 0.95 + 0j, (0, 1): 0.25 + 0.05j} +TRUE_CENTER = (0.4e-3, -0.3e-3) # beam offset from the optical axis +TRUE_POINTING_HORIZONTAL_DEG = 0.15 # beam pointing (horizontal tilt) +TRUE_POINTING_VERTICAL_DEG = -0.08 # beam pointing (vertical tilt) +IMAGE_SHAPE = (81, 81) + +# A camera positioned well upstream of the target planes and mildly tilted, +# so true perspective projection (keystoning) is in play but the frame +# still comfortably contains the beam at every z below. Calibration only +# gives us a nominal estimate of this pose -- it's refined jointly with +# everything else, within CAMERA_TOLERANCE, because of mechanical vibration +# between calibration and measurement. +PIXEL_SCALE = 4e-4 # meters/pixel, used only to size FOCAL_LENGTH_PX below +CAMERA_DISTANCE = 5.0 # meters upstream of the output window +FOCAL_LENGTH_PX = (CAMERA_DISTANCE + Z0) / PIXEL_SCALE + +TRUE_CAMERA = CameraModel( + focal_length_px=FOCAL_LENGTH_PX, + position=(0.01, -0.02, -CAMERA_DISTANCE), + orientation_deg=(1.5, -1.0, 0.5), +) +# Nominal (calibrated) camera pose, deliberately offset from TRUE_CAMERA +# within CAMERA_TOLERANCE, standing in for real calibration error. +NOMINAL_CAMERA = CameraModel( + focal_length_px=FOCAL_LENGTH_PX * 1.01, + position=(0.0, 0.0, -CAMERA_DISTANCE), + orientation_deg=(1.0, -0.5, 0.0), +) +CAMERA_TOLERANCE = CameraModelTolerance( + focal_length_px=FOCAL_LENGTH_PX * 0.05, + position=(0.02, 0.02, 0.05), + orientation_deg=(2.0, 2.0, 2.0), +) + +# Measurement plane distances, meters. Kept within roughly +/-2 Rayleigh +# ranges of z0 so the (widening) beam stays well within the camera frame -- +# planes much farther out would be clipped by the finite frame, which +# degrades the fit. +Z_LIST = [0.4, 0.45, 0.55, 0.6] +# Each plane's z is only known to a nominal precision (e.g. a translation +# stage's readout); offset the nominal value from the true z used to +# render the plane, and let the fit recover the true z within Z_TOLERANCE. +NOMINAL_Z_OFFSETS = {0.4: 0.003, 0.45: -0.002, 0.55: 0.004, 0.6: -0.003} +Z_TOLERANCE = 0.01 + +# --- Target thermal-diffusion blur (known target material properties) --- +THERMAL_DIFFUSIVITY = 1e-6 # m^2/s +DWELL_TIME = 0.2 # s + + +def main() -> None: + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + generator = SyntheticBeamGenerator(basis=basis, camera=TRUE_CAMERA) + + planes = generator.generate( + coefficients=TRUE_COEFFICIENTS, + z_list=Z_LIST, + image_shape=IMAGE_SHAPE, + center=TRUE_CENTER, + pointing_angle_horizontal_deg=TRUE_POINTING_HORIZONTAL_DEG, + pointing_angle_vertical_deg=TRUE_POINTING_VERTICAL_DEG, + z_tolerance=Z_TOLERANCE, + nominal_z_offsets=NOMINAL_Z_OFFSETS, + noise_std=2e-4, + seed=42, + ) + + # Apply the same thermal-diffusion blur a real target would exhibit, + # using the nominal (not true) camera to compute the pixel scale -- + # exactly what BeamReconstructor itself does internally. + blur_deconvolver = DiffusionDeconvolver( + thermal_diffusivity=THERMAL_DIFFUSIVITY, dwell_time=DWELL_TIME + ) + nominal_calibration = GeometryCalibration(NOMINAL_CAMERA) + for plane in planes: + pixel_scale = nominal_calibration.effective_pixel_scale(plane.flux.shape, plane.z) + plane.flux = blur_deconvolver.blur(plane.flux, pixel_scale) + + # The ground truth only has order-0 and order-1 content, so a max_order + # of 1 is enough for automatic mode-set growth to find it; growing much + # further would start fitting deconvolution/noise artifacts as spurious + # higher-order modes. + reconstructor = BeamReconstructor( + w0=W0, + z0=Z0, + wavelength=WAVELENGTH, + camera=NOMINAL_CAMERA, + camera_tolerance=CAMERA_TOLERANCE, + max_order=1, + deconvolver=blur_deconvolver, + ) + result = reconstructor.reconstruct(planes) + + print("Mode purity table (power fraction, phase [rad]):") + for mode, (fraction, phase) in sorted( + result.purity.items(), key=lambda item: -item[1][0] + ): + print(f" LG_{mode[0]},{mode[1]}: {fraction:6.3%} (phase {phase:+.3f} rad)") + + print( + "\nFitted pointing angles: " + f"horizontal={result.pointing_angle_horizontal_deg:.4f} deg, " + f"vertical={result.pointing_angle_vertical_deg:.4f} deg" + ) + print("Fitted beam center per plane (m):") + for plane, (cx, cy) in zip(planes, result.centers): + print(f" z={plane.z:.2f} m -> ({cx:.3e}, {cy:.3e})") + + print("\nFitted camera geometry:") + for key, value in result.geometry.items(): + print(f" {key}: {value:.6g}") + + print(f"\nUsed phase-retrieval fallback: {result.used_phase_retrieval}") + + plot_mode_purity(result) + plot_center_trace(planes, result) + plot_residuals(planes, result) + plt.show() + + +if __name__ == "__main__": + main() +``` + +- [ ] **Step 2: Run the example and verify it completes without error** + +Run: `.venv/bin/python examples/full_pipeline_example.py` +Expected: the script prints a mode purity table (LG_00 fraction near 0.9, +LG_01 fraction near 0.1), fitted pointing angles near `0.15`/`-0.08` deg, +a fitted camera geometry dict with 9 `CameraModel` fields, no exception or +traceback, and (since a display may not be available in this environment) +either three matplotlib windows or a harmless "cannot connect to display" +warning from `plt.show()` -- not a Python exception from the reconstruction +itself. If the fit clearly fails to converge (e.g. purity wildly off, or a +`ValueError`/`RuntimeError` from `scipy.optimize.least_squares`), adjust +`CAMERA_TOLERANCE`/`Z_TOLERANCE`/`NOMINAL_CAMERA` to be closer to +`TRUE_CAMERA` rather than changing library code -- per this plan's Global +Constraints on numeric tolerances being adjustable starting points. + +- [ ] **Step 3: Commit** + +```bash +git add examples/full_pipeline_example.py +git commit -m "$(cat <<'EOF' +Update full_pipeline_example.py for the CameraModel geometry redesign + +Demonstrates the new CameraModel/CameraModelTolerance API, two +independent beam pointing angles, and per-plane z refinement with a +deliberately offset nominal camera pose and z values. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 10: Update `CLAUDE.md` + +**Files:** +- Modify: `CLAUDE.md` + +**Interfaces:** +- Consumes: the final module responsibilities from Tasks 1-9 (no new code interfaces produced or consumed; this is prose only). + +Documentation-only task, same style as Task 8: rewrite, then a stale-identifier grep as the "verification" step. + +- [ ] **Step 1: Update the module responsibilities list** + +In `CLAUDE.md`, replace: + +```markdown +Data flows through the pipeline as a list of `MeasurementPlane` (one per imaging +distance `z`), each holding a raw 2D `flux` array plus optionally-known +`pixel_scale`/`viewing_angle_deg`. Everything downstream is keyed off `LGBasis`, which +defines the mode basis relative to a known waist `w0`/`z0`/`wavelength`. + +Module responsibilities (`he11lib/`): + +- **`data.py`** — `MeasurementPlane`, `ReconstructionResult` (the shared input/output + types) and `validate_planes` (>=3 planes, matching shapes, distinct `z`). +- **`modes.py`** — `LGBasis`: closed-form paraxial LG fields, beam radius `w(z)`, + Gouy phase, inverse radius of curvature, and projection of a measured field onto a + candidate mode set. This is the analytic ground truth all fitting is checked against. +- **`geometry.py`** — `GeometryCalibration`: resolves a plane's pixel-to-physical + coordinate grid, deferring to known `pixel_scale`/`viewing_angle_deg` on the plane + over any override passed in. +- **`noise.py`** — `NoiseEstimator`: automatic per-image noise-std estimation + (Laplacian method) and per-pixel weights for noise-weighted least squares. +- **`deconvolution.py`** — `DiffusionDeconvolver`: optional forward blur / Wiener + deconvolution for thermal-diffusion blur in the absorbing target. The blur kernel is + isotropic in pixel space, so it's only exact when `viewing_angle_deg == 0` (an + oblique view makes x/y pixel scales differ) — an accepted approximation, not a bug. +- **`synthetic.py`** — `SyntheticBeamGenerator`: forward model that produces + `MeasurementPlane`s from known ground-truth coefficients/center/pointing/geometry. + Used throughout the test suite and examples to validate the pipeline end-to-end. +- **`fitting.py`** — `ModalFitter` (`fit`, `fit_auto`) and `generate_mode_shells`: the + core joint nonlinear least-squares fit (complex LG coefficients + beam + center/pointing + unknown geometry) via `scipy.optimize.least_squares`. `fit_auto` + grows the candidate mode set shell-by-shell (by order `2p + |l|`), stopping via a BIC + improvement threshold, capped at `max_order` (emits `UserWarning`, doesn't raise, if + still improving at the cap). +- **`phase_retrieval.py`** — `propagate_angular_spectrum` (FFT-based paraxial + free-space propagation) and `PhaseRetriever` (multi-plane Gerchberg-Saxton), the + fallback reconstruction path for when a finite mode basis doesn't fit well. +- **`reconstruct.py`** — `BeamReconstructor`: the orchestrator. Pipeline order: + validate planes → optional deconvolution (requires known `pixel_scale` per plane) → + `ModalFitter.fit_auto` → optional `PhaseRetriever` fallback (forced via + `force_phase_retrieval`, or triggered automatically when the noise-weighted RMS + residual exceeds `phase_retrieval_residual_threshold`). The fallback path projects + the recovered field onto all modes up to `max_order` and produces a + `ReconstructionResult` with `used_phase_retrieval=True`, empty `residuals`, and NaN + `coefficient_uncertainty` (no fit covariance available from phase retrieval). +- **`plotting.py`** — diagnostic figures (`plot_mode_purity`, `plot_center_trace`, + `plot_residuals`); each returns a `Figure` rather than calling `plt.show()`. +``` + +with: + +```markdown +Data flows through the pipeline as a list of `MeasurementPlane` (one per imaging +distance `z`), each holding a raw 2D `flux` array plus a nominal `z` and `z_tolerance`. +Camera pose/intrinsics are a single shared `CameraModel`/`CameraModelTolerance` for the +whole reconstruction, not per-plane. Everything downstream is keyed off `LGBasis`, which +defines the mode basis relative to a known waist `w0`/`z0`/`wavelength`. + +Module responsibilities (`he11lib/`): + +- **`data.py`** — `MeasurementPlane` (flux, nominal `z`, `z_tolerance`), `ReconstructionResult` + (the shared input/output types) and `validate_planes` (>=3 planes, matching shapes, + distinct `z`). +- **`modes.py`** — `LGBasis`: closed-form paraxial LG fields, beam radius `w(z)`, + Gouy phase, inverse radius of curvature, and projection of a measured field onto a + candidate mode set. This is the analytic ground truth all fitting is checked against. +- **`geometry.py`** — `CameraModel`/`CameraModelTolerance` (a nominal pinhole camera + pose/intrinsics and its paired per-field refinement bound) and `GeometryCalibration`: + resolves pixel<->physical coordinates via true pinhole forward/inverse projection + (ray-plane intersection), producing genuine keystoning for tilted/off-axis poses + rather than a uniform affine correction. A tolerance of `0` on a `CameraModel` field + means it's trusted exactly; `>0` means it's refined within `[nominal-tolerance, + nominal+tolerance]` by `ModalFitter` — the same mechanism applies to + `MeasurementPlane.z`/`z_tolerance`. +- **`noise.py`** — `NoiseEstimator`: automatic per-image noise-std estimation + (Laplacian method) and per-pixel weights for noise-weighted least squares. +- **`deconvolution.py`** — `DiffusionDeconvolver`: optional forward blur / Wiener + deconvolution for thermal-diffusion blur in the absorbing target. The blur kernel is + isotropic in pixel space; callers use `GeometryCalibration.effective_pixel_scale` + (a finite-difference approximation at the frame center) as a single figure, so it's + only exact for an on-axis, untilted camera — an accepted approximation, not a bug. +- **`synthetic.py`** — `SyntheticBeamGenerator`: forward model that produces + `MeasurementPlane`s from a known ground-truth `CameraModel`, coefficients, center, + and two pointing angles, rendering each plane at its own true `z` (which may + deliberately differ from the plane's nominal `z`, via `nominal_z_offsets`). Used + throughout the test suite and examples to validate the pipeline end-to-end. +- **`fitting.py`** — `ModalFitter` (`fit`, `fit_auto`) and `generate_mode_shells`: the + core joint nonlinear least-squares fit via `scipy.optimize.least_squares`. Complex LG + coefficients, per-plane beam center, and the two pointing angles (horizontal/vertical) + are always free; each `CameraModel` field and each plane's `z` is additionally free + (bounded by its tolerance) only when its paired tolerance is nonzero, otherwise held + fixed as a constant. `fit_auto` grows the candidate mode set shell-by-shell (by order + `2p + |l|`), stopping via a BIC improvement threshold, capped at `max_order` (emits + `UserWarning`, doesn't raise, if still improving at the cap, or if the free + camera+`z` parameter count is large relative to the number of planes — see the + degeneracy pitfall below). +- **`phase_retrieval.py`** — `propagate_angular_spectrum` (FFT-based paraxial + free-space propagation) and `PhaseRetriever` (multi-plane Gerchberg-Saxton), the + fallback reconstruction path for when a finite mode basis doesn't fit well. Takes a + shared `CameraModel` (not per-plane pixel scale) to derive its common physical grid. +- **`reconstruct.py`** — `BeamReconstructor`: the orchestrator, now constructed with a + required `camera`/`camera_tolerance`. Pipeline order: validate planes → optional + deconvolution (using `GeometryCalibration(camera).effective_pixel_scale`) → + `ModalFitter.fit_auto` → optional `PhaseRetriever` fallback (forced via + `force_phase_retrieval`, or triggered automatically when the noise-weighted RMS + residual exceeds `phase_retrieval_residual_threshold`). The fallback path projects + the recovered field onto all modes up to `max_order` and produces a + `ReconstructionResult` with `used_phase_retrieval=True`, empty `residuals`, empty + `geometry`, NaN pointing angles, and NaN `coefficient_uncertainty` (no fit covariance + available from phase retrieval). +- **`plotting.py`** — diagnostic figures (`plot_mode_purity`, `plot_center_trace`, + `plot_residuals`); each returns a `Figure` rather than calling `plt.show()`. +``` + +- [ ] **Step 2: Add the new degeneracy pitfall** + +In `CLAUDE.md`'s "Known physics/fitting pitfalls" section, after pitfall 2 +(automatic mode-set growth overfitting), add a third pitfall: + +```markdown +3. **Shared camera/`z` geometry can be underdetermined with few planes.** With only + 3-10 planes, adding the ~7-9 shared `CameraModel` unknowns (whichever have nonzero + `CameraModelTolerance`) plus one `z` correction per plane (for nonzero + `z_tolerance`) can be practically underdetermined even though each plane + contributes many pixels of data, because those unknowns are *global* and only + weakly constrained by subtle keystone differences between planes. `fit_auto`/ + `BeamReconstructor` emit a `UserWarning` (not an error) when the free-parameter + count is large relative to the number of planes — if you see it, tighten + `CameraModelTolerance`/`z_tolerance` toward values you actually trust rather than + leaving them generously wide. +``` + +- [ ] **Step 3: Grep for stale identifiers and fix any remaining hits** + +Run: + +```bash +grep -n "pixel_scale\|viewing_angle_deg" CLAUDE.md +``` + +Expected: no output (`pixel_scale`/`viewing_angle_deg` no longer exist +anywhere in the codebase's public API after this redesign). + +- [ ] **Step 4: Commit** + +```bash +git add CLAUDE.md +git commit -m "$(cat <<'EOF' +Update CLAUDE.md for the CameraModel geometry redesign + +Refreshes the module responsibilities to describe CameraModel/ +CameraModelTolerance, z_tolerance, and the two pointing angles, and adds +the new shared-geometry-underdetermined pitfall. + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +--- + +## Task 11: Full-suite verification + +**Files:** +- None modified (verification only), unless the full-suite run surfaces an + integration issue missed by per-task test runs, in which case fix it in + the relevant file from Tasks 1-9 and note the fix in the commit message. + +**Interfaces:** +- Consumes: every public interface produced by Tasks 1-9. +- Produces: nothing (terminal task). + +- [ ] **Step 1: Run the full test suite** + +Run: `.venv/bin/pytest -q` +Expected: all tests pass (0 failures, 0 errors). If a cross-module +integration issue surfaces here that wasn't caught by an individual task's +own test run (e.g. a signature mismatch between two tasks written before +the other was finalized), fix it now in the appropriate module. + +- [ ] **Step 2: Confirm no stray references to the removed API remain anywhere in the tree** + +Run: + +```bash +grep -rn "pixel_scale_known\|viewing_angle_known\|pointing_angle_deg[^_]\|initial_pixel_scale\|initial_viewing_angle_deg" he11lib/ tests/ examples/ docs/api.md CLAUDE.md +``` + +Expected: no output. + +- [ ] **Step 3: Check git status is clean** + +Run: `git status` +Expected: working tree clean (every task ended in its own commit; nothing +should be left staged or modified). + +- [ ] **Step 4: If Step 1 or Step 2 required a fix, commit it** + +```bash +git add -A +git commit -m "$(cat <<'EOF' +Fix cross-module integration issue found in full-suite verification + + + +Co-Authored-By: Claude Sonnet 5 +EOF +)" +``` + +If Steps 1-2 passed cleanly with no fixes needed, skip this step entirely +— do not create an empty commit.