From fabb3d4efcff5b89b894c62491ac51e98b193400 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 11:30:27 +0200 Subject: [PATCH 01/12] Add implementation plan for camera geometry & measurement uncertainty redesign Covers CameraModel/CameraModelTolerance, tolerance-unified refinement of camera pose/intrinsics and per-plane z, 2D beam pointing, and updates to every affected module, tests, docs, and the example script. Co-Authored-By: Claude Sonnet 5 --- ...026-07-03-camera-geometry-redesign-plan.md | 3361 +++++++++++++++++ 1 file changed, 3361 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-03-camera-geometry-redesign-plan.md 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. From 4f65c2ce4f735eec0e73156bfcb7551b99148179 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 11:32:31 +0200 Subject: [PATCH 02/12] 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 --- he11lib/__init__.py | 4 +- he11lib/geometry.py | 262 ++++++++++++++++++++++++++++++++++------- tests/test_geometry.py | 182 +++++++++++++++++++--------- 3 files changed, 352 insertions(+), 96 deletions(-) diff --git a/he11lib/__init__.py b/he11lib/__init__.py index 790d428..b5ec294 100644 --- a/he11lib/__init__.py +++ b/he11lib/__init__.py @@ -6,7 +6,7 @@ See docs/ for the full API and design documentation. from .data import MeasurementPlane, ReconstructionResult from .deconvolution import DiffusionDeconvolver from .fitting import ModalFitter, generate_mode_shells -from .geometry import GeometryCalibration +from .geometry import CameraModel, CameraModelTolerance, GeometryCalibration from .modes import LGBasis from .noise import NoiseEstimator from .phase_retrieval import PhaseRetrievalResult, PhaseRetriever, propagate_angular_spectrum @@ -21,6 +21,8 @@ __all__ = [ "DiffusionDeconvolver", "ModalFitter", "generate_mode_shells", + "CameraModel", + "CameraModelTolerance", "GeometryCalibration", "LGBasis", "NoiseEstimator", diff --git a/he11lib/geometry.py b/he11lib/geometry.py index 0b423e5..edd25b4 100644 --- a/he11lib/geometry.py +++ b/he11lib/geometry.py @@ -1,64 +1,244 @@ -"""Camera geometry correction: pixel-to-physical scale and viewing angle. +"""Camera geometry: a shared pinhole camera model and pixel<->physical mapping. -Converts a MeasurementPlane's pixel grid into physical (x, y) coordinates in -the beam's transverse plane, compensating for an oblique camera viewing -angle (which compresses the image along the tilt axis by cos(angle)). -Known calibration values (on the MeasurementPlane) are used directly; when a -value is unknown, an override must be supplied (e.g. by ModalFitter while -exploring it as a free parameter). +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 -from .data import MeasurementPlane +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 pixel scale / viewing angle and builds a physical coordinate grid.""" + """Resolves the pixel<->physical mapping for a shared pinhole `CameraModel`.""" - def __init__(self, plane: MeasurementPlane): - self.plane = plane + def __init__(self, camera: CameraModel): + self.camera = camera + self._rotation = _rotation_matrix(*camera.orientation_deg) - @property - def pixel_scale_known(self) -> bool: - return self.plane.pixel_scale is not None + 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 - @property - def viewing_angle_known(self) -> bool: - return self.plane.viewing_angle_deg is not None + 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, - pixel_scale: float | None = None, - viewing_angle_deg: float | None = None, + self, image_shape: tuple[int, int], z: float ) -> tuple[np.ndarray, np.ndarray]: - """Physical (x, y) grid matching the plane's flux array shape. + """Inverse pinhole projection: pixel grid at depth z -> physical (x, y). - Known values on the MeasurementPlane take precedence; overrides are - only used to fill in values that are not known/calibrated. + 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. """ - scale = self.plane.pixel_scale if self.pixel_scale_known else pixel_scale - angle_deg = ( - self.plane.viewing_angle_deg if self.viewing_angle_known else viewing_angle_deg - ) - - if scale is None: - raise ValueError( - "pixel_scale is not known for this MeasurementPlane and no override was given" - ) - if angle_deg is None: - raise ValueError( - "viewing_angle_deg is not known for this MeasurementPlane and no override was given" - ) - - rows, cols = self.plane.flux.shape + 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) - cos_angle = np.cos(np.deg2rad(angle_deg)) - x = col_grid * scale / cos_angle - y = row_grid * scale + 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) diff --git a/tests/test_geometry.py b/tests/test_geometry.py index 6fdd399..6179841 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -1,77 +1,151 @@ import numpy as np import pytest -from he11lib.data import MeasurementPlane -from he11lib.geometry import GeometryCalibration +from he11lib.geometry import CameraModel, CameraModelTolerance, GeometryCalibration -def test_pixel_scale_known_reflects_plane(): - plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, pixel_scale=1e-4) - plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) - - assert GeometryCalibration(plane_known).pixel_scale_known is True - assert GeometryCalibration(plane_unknown).pixel_scale_known is False +def test_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_viewing_angle_known_reflects_plane(): - plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, viewing_angle_deg=10.0) - plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) - - assert GeometryCalibration(plane_known).viewing_angle_known is True - assert GeometryCalibration(plane_unknown).viewing_angle_known is False +def test_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_physical_coordinates_uses_known_calibration(): - plane = MeasurementPlane( - flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=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), ) - calib = GeometryCalibration(plane) - x, y = calib.physical_coordinates() - - row_idx = np.arange(5) - 2 - col_idx = np.arange(5) - 2 - expected_x = col_idx * 2e-4 - expected_y = row_idx * 2e-4 - np.testing.assert_allclose(x[2, :], expected_x) - np.testing.assert_allclose(y[:, 2], expected_y) -def test_physical_coordinates_compresses_x_for_viewing_angle(): - plane = MeasurementPlane( - flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=60.0 +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), ) - calib = GeometryCalibration(plane) - x, y = calib.physical_coordinates() - - col_idx = np.arange(5) - 2 - expected_x = col_idx * 2e-4 / np.cos(np.deg2rad(60.0)) - np.testing.assert_allclose(x[2, :], expected_x) -def test_physical_coordinates_raises_without_calibration_or_override(): - plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) - calib = GeometryCalibration(plane) +@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) - with pytest.raises(ValueError, match="pixel_scale"): - calib.physical_coordinates() + 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_physical_coordinates_accepts_override_for_unknown_values(): - plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) - calib = GeometryCalibration(plane) +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 - x, y = calib.physical_coordinates(pixel_scale=1e-4, viewing_angle_deg=0.0) - col_idx = np.arange(5) - 2 - np.testing.assert_allclose(x[2, :], col_idx * 1e-4) + 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_known_calibration_takes_precedence_over_override(): - plane = MeasurementPlane( - flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=0.0 +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(plane) + calib = GeometryCalibration(camera) - # override should be ignored since plane already specifies calibration - x, _ = calib.physical_coordinates(pixel_scale=999.0, viewing_angle_deg=45.0) - col_idx = np.arange(5) - 2 - np.testing.assert_allclose(x[2, :], col_idx * 2e-4) + 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) From dffca62f817c3b77235fc1d11b10973d50a12d1c Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 11:39:16 +0200 Subject: [PATCH 03/12] 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 --- he11lib/data.py | 29 +++++++++++++++++------------ he11lib/plotting.py | 6 +++++- tests/test_data.py | 22 +++++++++++++--------- tests/test_plotting.py | 3 ++- 4 files changed, 37 insertions(+), 23 deletions(-) diff --git a/he11lib/data.py b/he11lib/data.py index b88b889..c0cfbc4 100644 --- a/he11lib/data.py +++ b/he11lib/data.py @@ -9,24 +9,22 @@ import numpy as np @dataclass class MeasurementPlane: - """A single thermal (flux) image at a known distance from the output window. + """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 : distance from the output window, in meters. Must be positive. - pixel_scale : known mm/pixel scale, if calibrated. If None, treated as an - unknown to be estimated jointly during reconstruction. - viewing_angle_deg : known camera viewing angle relative to the beam axis, - in degrees, if calibrated. If None, treated as an unknown. + 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 - pixel_scale: float | None = None - viewing_angle_deg: float | None = None + z_tolerance: float = 0.0 label: str | None = None def __post_init__(self) -> None: @@ -36,6 +34,10 @@ class MeasurementPlane: ) 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: @@ -68,9 +70,11 @@ class ReconstructionResult: 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_deg : fitted shared beam pointing angle (tilt), in degrees. - geometry : geometry parameters used or fitted (e.g. pixel_scale, - viewing_angle_deg), keyed by name. + 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. @@ -81,7 +85,8 @@ class ReconstructionResult: purity: dict[tuple[int, int], tuple[float, float]] reconstructed_field: np.ndarray centers: list[tuple[float, float]] - pointing_angle_deg: 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) diff --git a/he11lib/plotting.py b/he11lib/plotting.py index ccb3138..0422b85 100644 --- a/he11lib/plotting.py +++ b/he11lib/plotting.py @@ -41,7 +41,11 @@ def plot_center_trace( ax_y.plot(z_values, y_values, marker="o") ax_y.set_xlabel("z (m)") ax_y.set_ylabel("center y (m)") - fig.suptitle(f"Beam center (pointing angle {result.pointing_angle_deg:.3g} deg)") + fig.suptitle( + "Beam center (pointing angle " + f"h={result.pointing_angle_horizontal_deg:.3g} deg, " + f"v={result.pointing_angle_vertical_deg:.3g} deg)" + ) return fig diff --git a/tests/test_data.py b/tests/test_data.py index b4d2a52..d96d3f1 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -10,19 +10,15 @@ def test_measurement_plane_stores_fields(): assert plane.z == 0.3 assert np.array_equal(plane.flux, flux) - assert plane.pixel_scale is None - assert plane.viewing_angle_deg is None + assert plane.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, pixel_scale=0.1, viewing_angle_deg=5.0, label="plane_40cm" - ) + plane = MeasurementPlane(flux=flux, z=0.4, z_tolerance=0.01, label="plane_40cm") - assert plane.pixel_scale == 0.1 - assert plane.viewing_angle_deg == 5.0 + assert plane.z_tolerance == 0.01 assert plane.label == "plane_40cm" @@ -39,6 +35,11 @@ def test_measurement_plane_rejects_non_positive_z(): 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), @@ -82,12 +83,15 @@ def test_reconstruction_result_stores_fields(): purity={(0, 0): (1.0, 0.0)}, reconstructed_field=np.ones((4, 4), dtype=complex), centers=[(0.0, 0.0)], - pointing_angle_deg=0.0, - geometry={"pixel_scale": 0.1, "viewing_angle_deg": 2.0}, + 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 diff --git a/tests/test_plotting.py b/tests/test_plotting.py index a376bb9..04265d7 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -11,7 +11,8 @@ def make_result(**overrides): purity={(0, 0): (0.9, 0.1), (1, 0): (0.1, -0.2)}, reconstructed_field=np.zeros((5, 5), dtype=complex), centers=[(0.0, 0.0), (1e-4, -1e-4), (2e-4, -2e-4)], - pointing_angle_deg=0.5, + 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) From c2ae95e01f4192cff4028d896e139444b37157ba Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 11:43:20 +0200 Subject: [PATCH 04/12] 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 --- he11lib/synthetic.py | 64 ++++++++++++------------- tests/test_synthetic.py | 103 ++++++++++++++++++++++++---------------- 2 files changed, 93 insertions(+), 74 deletions(-) diff --git a/he11lib/synthetic.py b/he11lib/synthetic.py index 215b34d..b9f2baf 100644 --- a/he11lib/synthetic.py +++ b/he11lib/synthetic.py @@ -10,6 +10,7 @@ from __future__ import annotations import numpy as np from .data import MeasurementPlane +from .geometry import CameraModel, GeometryCalibration from .modes import LGBasis @@ -19,65 +20,60 @@ class SyntheticBeamGenerator: Parameters ---------- basis : LGBasis defining the reference w0, z0, wavelength. - image_shape : (rows, cols) pixel shape of generated images. - pixel_scale : physical size of one pixel, in meters, along the - non-tilted (y) axis. The tilt/projection axis is assumed to be x. + camera : ground-truth CameraModel (position/orientation/intrinsics) used + to render each plane via true perspective projection. """ - def __init__(self, basis: LGBasis, image_shape: tuple[int, int], pixel_scale: float): + def __init__(self, basis: LGBasis, camera: CameraModel): self.basis = basis - self.image_shape = image_shape - self.pixel_scale = pixel_scale - - def _pixel_grid(self, center: tuple[float, float], viewing_angle_deg: float): - rows, cols = self.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) - - cos_angle = np.cos(np.deg2rad(viewing_angle_deg)) - x = col_grid * self.pixel_scale / cos_angle - center[0] - y = row_grid * self.pixel_scale - center[1] - return x, y + 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_deg: float = 0.0, - viewing_angle_deg: float = 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 z distance. + """Generate one MeasurementPlane per requested (true) z distance. - The beam transverse center drifts linearly with z according to - pointing_angle_deg (tilt of the beam axis along x), starting from - `center` at the basis's reference z0. + 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_rad = np.deg2rad(pointing_angle_deg) + 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_rad) - plane_center = (center[0] + drift_x, center[1]) + 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._pixel_grid(plane_center, viewing_angle_deg) - field = self.basis.field_superposition(x, y, z, coefficients) + 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=z, - pixel_scale=self.pixel_scale, - viewing_angle_deg=viewing_angle_deg, - ) + MeasurementPlane(flux=flux, z=nominal_z, z_tolerance=z_tolerance) ) return planes diff --git a/tests/test_synthetic.py b/tests/test_synthetic.py index 90e0e9d..00d69d2 100644 --- a/tests/test_synthetic.py +++ b/tests/test_synthetic.py @@ -1,6 +1,7 @@ import numpy as np import pytest +from he11lib.geometry import CameraModel from he11lib.modes import LGBasis from he11lib.synthetic import SyntheticBeamGenerator @@ -8,21 +9,29 @@ from he11lib.synthetic import SyntheticBeamGenerator W0 = 5e-3 Z0 = 0.5 WAVELENGTH = 1.76e-3 -PIXEL_SCALE = 2e-4 # 0.2 mm/px +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, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE - ) + 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) + 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) @@ -30,7 +39,9 @@ def test_generate_returns_planes_with_requested_z(): def test_generate_pure_mode_peak_at_image_center_when_centered(): gen = make_generator() - planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(0.0, 0.0)) + 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) @@ -40,9 +51,9 @@ def test_generate_pure_mode_peak_at_image_center_when_centered(): def test_generate_applies_center_offset(): gen = make_generator() - offset_m = 20 * PIXEL_SCALE # 20 pixels + offset_m = 20 * PIXEL_SCALE # ~20 pixels at z=Z0 planes = gen.generate( - coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(offset_m, 0.0) + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, center=(offset_m, 0.0) ) flux = planes[0].flux @@ -53,35 +64,41 @@ def test_generate_applies_center_offset(): assert peak_idx[1] == pytest.approx(center_col + 20, abs=1) -def test_generate_applies_pointing_angle_as_linear_drift(): +def test_generate_applies_pointing_angles_as_2d_linear_drift(): gen = make_generator() - pointing_angle_deg = 1.0 # small tilt z_list = [Z0, Z0 + 0.2] planes = gen.generate( coefficients={(0, 0): 1 + 0j}, z_list=z_list, + image_shape=IMAGE_SHAPE, center=(0.0, 0.0), - pointing_angle_deg=pointing_angle_deg, + pointing_angle_horizontal_deg=1.0, + pointing_angle_vertical_deg=0.5, ) - peaks_col = [] + peaks = [] for plane in planes: peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape) - peaks_col.append(peak_idx[1]) + peaks.append(peak_idx) - expected_shift_m = 0.2 * np.tan(np.deg2rad(pointing_angle_deg)) - expected_shift_px = expected_shift_m / PIXEL_SCALE - actual_shift_px = peaks_col[1] - peaks_col[0] - assert actual_shift_px == pytest.approx(expected_shift_px, abs=1) + 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], noise_std=0.01, seed=42 + 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], noise_std=0.01, seed=42 + 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) @@ -90,34 +107,40 @@ def test_generate_noise_std_matches_requested_level(): gen = make_generator() noise_std = 0.02 planes_noisy = gen.generate( - coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=noise_std, seed=1 + 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 ) - planes_clean = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.0) diff = planes_noisy[0].flux - planes_clean[0].flux assert np.std(diff) == pytest.approx(noise_std, rel=0.15) -def test_generate_viewing_angle_compresses_tilt_axis(): +def test_generate_applies_z_tolerance_to_every_plane(): gen = make_generator() - planes_straight = gen.generate( - coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=0.0 + planes = gen.generate( + coefficients={(0, 0): 1 + 0j}, + z_list=[0.3, 0.4, 0.5], + image_shape=IMAGE_SHAPE, + z_tolerance=0.02, ) - planes_tilted = gen.generate( - coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=60.0 + 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, ) - def width_along_axis(flux, axis): - profile = flux[flux.shape[0] // 2, :] if axis == 1 else flux[:, flux.shape[1] // 2] - half_max = profile.max() / 2 - above = np.where(profile >= half_max)[0] - return above[-1] - above[0] - - width_straight_x = width_along_axis(planes_straight[0].flux, axis=1) - width_tilted_x = width_along_axis(planes_tilted[0].flux, axis=1) - width_straight_y = width_along_axis(planes_straight[0].flux, axis=0) - width_tilted_y = width_along_axis(planes_tilted[0].flux, axis=0) - - # tilt compresses the viewed beam along the tilt (x) axis, y unaffected - assert width_tilted_x < width_straight_x - assert width_tilted_y == pytest.approx(width_straight_y, abs=1) + 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. From ef57ec81e45ff7475104abdb5fef7a518ecf86e6 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 12:04:47 +0200 Subject: [PATCH 05/12] 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. Also seeds the first mode's coefficient with a small imaginary offset (1.0 + 0.05j instead of 1.0 + 0j) rather than exactly on the real axis: for a single mode, intensity depends only on |c| (phase is unobservable), so Im=0 sits exactly on that flat/degenerate valley, aligned with a coordinate axis, giving a zero-gradient Jacobian column that destabilized trf's trust-region step and prevented pointing-angle recovery from a default (0, 0) initial guess. Co-Authored-By: Claude Sonnet 5 --- he11lib/fitting.py | 161 ++++++++++++++++++++++------------- tests/test_fitting.py | 189 +++++++++++++++++++++++++++++++++++------- 2 files changed, 265 insertions(+), 85 deletions(-) diff --git a/he11lib/fitting.py b/he11lib/fitting.py index 4c3184f..37bb6f7 100644 --- a/he11lib/fitting.py +++ b/he11lib/fitting.py @@ -8,7 +8,14 @@ import numpy as np from scipy.optimize import least_squares from .data import MeasurementPlane, ReconstructionResult, validate_planes -from .geometry import GeometryCalibration +from .geometry import ( + CameraModel, + CameraModelTolerance, + GeometryCalibration, + camera_from_values, + camera_to_values, + tolerance_to_values, +) from .modes import LGBasis from .noise import NoiseEstimator @@ -35,19 +42,31 @@ class ModalFitter: 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_tilt_deg: tuple[float, float] = (0.0, 0.0), - initial_pixel_scale: float | None = None, - initial_viewing_angle_deg: float = 0.0, + initial_pointing_deg: tuple[float, float] = (0.0, 0.0), ) -> ReconstructionResult: - """Jointly fit complex coefficients for `modes` plus center/tilt/geometry.""" - validate_planes(planes) + """Jointly fit complex coefficients for `modes` plus center/pointing/geometry. - unknown_scale_idx = [i for i, p in enumerate(planes) if p.pixel_scale is None] - unknown_angle_idx = [i for i, p in enumerate(planes) if p.viewing_angle_deg is None] + 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): @@ -56,98 +75,126 @@ class ModalFitter: # 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 + # A purely-real seed (Im=0) sits exactly on the flat/ + # degenerate valley of a single mode's phase (for one + # mode alone, |c|^2 -- not arg(c) -- is all that's + # observable in intensity), giving a zero-gradient + # column that destabilizes trf's trust-region step; + # a small imaginary offset avoids landing on that axis. + c = 1.0 + 0.05j if i == 0 else 0.1 + 0.05j x += [c.real, c.imag] - x += [initial_center[0], initial_center[1], initial_tilt_deg[0], initial_tilt_deg[1]] - for _ in unknown_scale_idx: - x.append(initial_pixel_scale if initial_pixel_scale is not None else 1e-4) - for _ in unknown_angle_idx: - x.append(initial_viewing_angle_deg) + 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) - n_modes = len(modes) + 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_x_deg, tilt_y_deg = x[offset : offset + 4] + x0, y0, tilt_h_deg, tilt_v_deg = x[offset : offset + 4] offset += 4 - scales = {} - for idx in unknown_scale_idx: - scales[idx] = x[offset] - offset += 1 - angles = {} - for idx in unknown_angle_idx: - angles[idx] = x[offset] - offset += 1 - return coeffs, (x0, y0), (tilt_x_deg, tilt_y_deg), scales, angles - def plane_center(x0: float, y0: float, tilt_deg: tuple[float, float], z: float): - drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[0])) - drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[1])) + 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(i: int, plane: MeasurementPlane, coeffs, center0, tilt_deg, scales, angles): - scale = plane.pixel_scale if plane.pixel_scale is not None else scales[i] - angle = plane.viewing_angle_deg if plane.viewing_angle_deg is not None else angles[i] - calib = GeometryCalibration(plane) - x_grid, y_grid = calib.physical_coordinates(pixel_scale=scale, viewing_angle_deg=angle) - cx, cy = plane_center(center0[0], center0[1], tilt_deg, plane.z) - field = self.basis.field_superposition(x_grid - cx, y_grid - cy, plane.z, coeffs) + 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, tilt_deg, scales, angles = unpack(x) + coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(x) parts = [] for i, plane in enumerate(planes): - model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles) + 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() - # 'trf' + x_scale='jac' handles the very different natural magnitudes - # of these parameters (coefficients ~O(1), pixel_scale ~O(1e-3), - # angles ~O(1-90)); plain 'lm' can terminate prematurely on 'xtol' - # because its unscaled step-size test is dominated by the largest - # parameters. + 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", max_nfev=5000 + residuals, x0_vec, method="trf", x_scale="jac", bounds=(lower, upper), max_nfev=5000 ) - coeffs, center0, tilt_deg, scales, angles = unpack(opt_result.x) + 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], tilt_deg, p.z) for p in planes] - pointing_angle_deg = float(np.hypot(tilt_deg[0], tilt_deg[1])) + centers = [ + plane_center(center0[0], center0[1], pointing_deg, z_values[i]) + for i in range(len(planes)) + ] - geometry: dict[str, float] = {} + 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"pixel_scale_{i}"] = ( - planes[i].pixel_scale if planes[i].pixel_scale is not None else scales[i] - ) - geometry[f"viewing_angle_deg_{i}"] = ( - planes[i].viewing_angle_deg if planes[i].viewing_angle_deg is not None else angles[i] - ) + geometry[f"z_{i}"] = z_values[i] residual_maps = [] for i, plane in enumerate(planes): - model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles) + 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_z = min(planes, key=lambda p: abs(p.z - self.basis.z0)).z - field_at_reference = self._field_on_default_grid(coeffs, reference_z) + 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_deg=pointing_angle_deg, + pointing_angle_horizontal_deg=pointing_deg[0], + pointing_angle_vertical_deg=pointing_deg[1], geometry=geometry, residuals=residual_maps, coefficient_uncertainty=coefficient_uncertainty, diff --git a/tests/test_fitting.py b/tests/test_fitting.py index ff0c170..cd58cdb 100644 --- a/tests/test_fitting.py +++ b/tests/test_fitting.py @@ -3,6 +3,7 @@ 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 @@ -10,6 +11,7 @@ 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] @@ -18,8 +20,21 @@ def make_basis(): return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) -def make_generator(basis): - return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) +def 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(): @@ -31,13 +46,14 @@ def test_generate_mode_shells_orders_by_2p_plus_abs_l(): def test_fit_recovers_pure_fundamental_mode(): basis = make_basis() - gen = make_generator(basis) + camera = make_camera() + gen = make_generator(basis, camera) planes = gen.generate( - coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0 + 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)]) + 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) @@ -48,12 +64,17 @@ def test_fit_recovers_pure_fundamental_mode(): def test_fit_recovers_two_mode_purity_ratio(): basis = make_basis() - gen = make_generator(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, noise_std=1e-4, seed=1) + 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())) + 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(): @@ -64,49 +85,161 @@ def test_fit_recovers_two_mode_purity_ratio(): def test_fit_recovers_center_offset(): basis = make_basis() - gen = make_generator(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)], initial_center=true_center) + 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_unknown_pixel_scale(): - # Use a coarser pixel scale so the (much wider, far-field) beam at the - # outer z distances still fits within the frame -- otherwise pixel scale - # becomes unobservable from clipped images. +def test_fit_recovers_pointing_angles_independently(): basis = make_basis() - local_pixel_scale = 1.5e-3 - gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=local_pixel_scale) - planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=3) + 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, + ) - # hide the known calibration to force the fitter to solve for it - for plane in planes: - plane.pixel_scale = None - plane.viewing_angle_deg = None + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], 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)], - initial_pixel_scale=local_pixel_scale * 1.1, - initial_viewing_angle_deg=0.0, + planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=zero_tolerance() ) - fitted_scales = [result.geometry[f"pixel_scale_{i}"] for i in range(len(planes))] - for scale in fitted_scales: - assert scale == pytest.approx(local_pixel_scale, rel=0.05) + 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) def test_fit_auto_does_not_add_modes_for_pure_fundamental(): From b81964fbada7f25e821d3a01cf2e5658ff0f19a2 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 12:14:07 +0200 Subject: [PATCH 06/12] 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 --- he11lib/fitting.py | 56 +++++++++++++++++++++++++++++++++++----- tests/test_fitting.py | 59 ++++++++++++++++++++++++++++++++++++++----- 2 files changed, 102 insertions(+), 13 deletions(-) diff --git a/he11lib/fitting.py b/he11lib/fitting.py index 37bb6f7..e43c3fd 100644 --- a/he11lib/fitting.py +++ b/he11lib/fitting.py @@ -204,23 +204,28 @@ class ModalFitter: 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) - best_bic = self._bic(planes, best_result, current_modes) + 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, initial_coefficients=warm_start) - trial_bic = self._bic(planes, trial_result, trial_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 @@ -250,10 +255,47 @@ class ModalFitter: coeffs[mode] = amplitude * np.exp(1j * phase) return coeffs - def _bic(self, planes: list[MeasurementPlane], result: ReconstructionResult, modes: list[tuple[int, int]]) -> float: - chi2 = sum(np.sum((r * np.sqrt(self.noise_estimator.weights(p.flux))) ** 2) for r, p in zip(result.residuals, planes)) + 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, + ) + + 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) - n_params = 2 * len(modes) + 4 + 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)) def _estimate_uncertainty(self, opt_result, modes, coeffs, total_power): diff --git a/tests/test_fitting.py b/tests/test_fitting.py index cd58cdb..096f32a 100644 --- a/tests/test_fitting.py +++ b/tests/test_fitting.py @@ -1,3 +1,5 @@ +import warnings + import numpy as np import pytest @@ -244,24 +246,69 @@ def test_fit_recovers_offset_z_within_tolerance(): def test_fit_auto_does_not_add_modes_for_pure_fundamental(): basis = make_basis() - gen = make_generator(basis) + camera = make_camera() + gen = make_generator(basis, camera) planes = gen.generate( - coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=4 + 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, max_order=2) + 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() - gen = make_generator(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, noise_std=1e-4, seed=5) + 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, max_order=2) + 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) From 1d399ccc9b9544055607d0af449bfdf6f77d72e5 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 12:19:34 +0200 Subject: [PATCH 07/12] 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 --- he11lib/phase_retrieval.py | 17 ++++++++--------- tests/test_phase_retrieval.py | 36 +++++++++++++++++++++++++++-------- 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/he11lib/phase_retrieval.py b/he11lib/phase_retrieval.py index daa9c87..7fd4e16 100644 --- a/he11lib/phase_retrieval.py +++ b/he11lib/phase_retrieval.py @@ -15,7 +15,7 @@ from dataclasses import dataclass import numpy as np from .data import MeasurementPlane, validate_planes -from .geometry import GeometryCalibration +from .geometry import CameraModel, GeometryCalibration def propagate_angular_spectrum( @@ -79,22 +79,21 @@ class PhaseRetriever: def retrieve( self, planes: list[MeasurementPlane], - pixel_scale: float | None = None, - viewing_angle_deg: float | None = None, + camera: CameraModel, 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. + 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(ordered[0]).physical_coordinates( - pixel_scale=pixel_scale, viewing_angle_deg=viewing_angle_deg - ) + x, y = GeometryCalibration(camera).physical_coordinates(ordered[0].flux.shape, ordered[0].z) dx = float(x[0, 1] - x[0, 0]) amplitudes = [np.sqrt(np.clip(p.flux, 0, None)) for p in ordered] diff --git a/tests/test_phase_retrieval.py b/tests/test_phase_retrieval.py index ad66a5d..5808be7 100644 --- a/tests/test_phase_retrieval.py +++ b/tests/test_phase_retrieval.py @@ -1,6 +1,7 @@ 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 @@ -9,6 +10,7 @@ W0 = 5e-3 Z0 = 0.5 WAVELENGTH = 1.76e-3 PIXEL_SCALE = 3e-4 +CAMERA_DISTANCE = 5.0 IMAGE_SHAPE = (121, 121) @@ -16,6 +18,15 @@ 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) @@ -42,7 +53,6 @@ def test_propagate_matches_lgbasis_analytic_evolution(): propagated = propagate_angular_spectrum(field_at_waist, PIXEL_SCALE, dz=dz, wavelength=WAVELENGTH) analytic = basis.field(x, y, Z0 + dz, p=0, l=0) - # compare intensity profiles (phase reference/global constant may differ) np.testing.assert_allclose( np.abs(propagated) ** 2, np.abs(analytic) ** 2, atol=1e-2 * np.max(np.abs(analytic) ** 2) ) @@ -53,15 +63,19 @@ def test_retrieve_recovers_pure_mode_purity(): # within the frame -- otherwise FFT wraparound/clipping at the edges # degrades angular-spectrum propagation accuracy. basis = make_basis() - gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) + 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, noise_std=1e-5, seed=0) + 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, viewing_angle_deg=0.0, max_iterations=100) + 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, PIXEL_SCALE, result.z, modes=[(0, 0), (1, 0), (0, 1)] + 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 @@ -70,15 +84,21 @@ def test_retrieve_recovers_pure_mode_purity(): def test_retrieve_estimates_beam_center(): basis = make_basis() - gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) + 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, center=true_center, noise_std=1e-5, seed=1 + 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, viewing_angle_deg=0.0, max_iterations=100) + 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) From 57dc7d743ebc506361f34be3c83ed2898f5bc5d2 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 12:26:18 +0200 Subject: [PATCH 08/12] 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 --- he11lib/reconstruct.py | 29 +++++--- tests/test_reconstruct.py | 144 +++++++++++++++++++++++++++++++++----- 2 files changed, 145 insertions(+), 28 deletions(-) diff --git a/he11lib/reconstruct.py b/he11lib/reconstruct.py index 0d22b4e..be4ab9a 100644 --- a/he11lib/reconstruct.py +++ b/he11lib/reconstruct.py @@ -9,6 +9,7 @@ import numpy as np 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 @@ -25,12 +26,15 @@ class BeamReconstructor: 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 (its - `pixel_scale` must be known) before fitting. + 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` @@ -43,6 +47,8 @@ class BeamReconstructor: w0: float, z0: float, wavelength: float, + camera: CameraModel, + camera_tolerance: CameraModelTolerance, max_order: int = 4, noise_estimator: NoiseEstimator | None = None, deconvolver: DiffusionDeconvolver | None = None, @@ -51,6 +57,8 @@ class BeamReconstructor: ): 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 @@ -63,7 +71,9 @@ class BeamReconstructor: planes = self._deconvolve(planes) fitter = ModalFitter(self.basis, self.noise_estimator) - result = fitter.fit_auto(planes, max_order=self.max_order) + 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) @@ -73,13 +83,11 @@ class BeamReconstructor: def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]: if self.deconvolver is None: return planes + calib = GeometryCalibration(self.camera) 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) + 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 @@ -101,7 +109,7 @@ class BeamReconstructor: self, planes: list[MeasurementPlane] ) -> ReconstructionResult: retriever = PhaseRetriever(self.wavelength) - pr_result = retriever.retrieve(planes) + 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]) @@ -116,7 +124,8 @@ class BeamReconstructor: purity=purity, reconstructed_field=pr_result.field, centers=[pr_result.center for _ in planes], - pointing_angle_deg=float("nan"), + pointing_angle_horizontal_deg=float("nan"), + pointing_angle_vertical_deg=float("nan"), geometry={}, residuals=[], coefficient_uncertainty={mode: float("nan") for mode in modes}, diff --git a/tests/test_reconstruct.py b/tests/test_reconstruct.py index 2c20c91..86ea393 100644 --- a/tests/test_reconstruct.py +++ b/tests/test_reconstruct.py @@ -4,6 +4,7 @@ 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 @@ -12,6 +13,7 @@ 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] @@ -20,16 +22,36 @@ def make_basis(): return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) -def make_generator(basis): - return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) +def 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() - gen = make_generator(basis) - planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0) + 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, max_order=2) + 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)] @@ -39,13 +61,21 @@ def test_reconstruct_recovers_pure_mode_purity(): def test_reconstruct_recovers_center_offset(): basis = make_basis() - gen = make_generator(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, center=true_center, noise_std=1e-4, seed=1 + 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, max_order=2) + 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: @@ -55,21 +85,34 @@ def test_reconstruct_recovers_center_offset(): def test_reconstruct_with_deconvolution_corrects_blur(): basis = make_basis() - gen = make_generator(basis) - planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=2) + 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, p.pixel_scale)) for p in 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)]) + 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, max_order=2, deconvolver=deconvolver + 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)] @@ -80,12 +123,21 @@ def test_reconstruct_with_deconvolution_corrects_blur(): def test_reconstruct_forces_phase_retrieval_fallback(): basis = make_basis() - gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4) + 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, noise_std=1e-5, seed=3) + 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, max_order=2, force_phase_retrieval=True + w0=W0, + z0=Z0, + wavelength=WAVELENGTH, + camera=camera, + camera_tolerance=zero_tolerance(), + max_order=2, + force_phase_retrieval=True, ) result = reconstructor.reconstruct(planes) @@ -96,17 +148,73 @@ def test_reconstruct_forces_phase_retrieval_fallback(): def test_reconstruct_falls_back_automatically_on_high_residual(): basis = make_basis() - gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4) + 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, noise_std=1e-5, seed=4) + 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) From c6b824660d1531572ba6b49585927c666349bbfa Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 12:33:00 +0200 Subject: [PATCH 09/12] 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 --- docs/api.md | 180 ++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 132 insertions(+), 48 deletions(-) diff --git a/docs/api.md b/docs/api.md index 8aea1f3..3f54769 100644 --- a/docs/api.md +++ b/docs/api.md @@ -15,11 +15,33 @@ Every class/function below is exported from the top-level `he11lib` package ## Quick start ```python -from he11lib import BeamReconstructor, MeasurementPlane +from he11lib import ( + BeamReconstructor, + CameraModel, + CameraModelTolerance, + 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) + +# 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(): @@ -28,19 +50,22 @@ for mode, (power_fraction, phase_rad) in result.purity.items(): ## `data` — `MeasurementPlane`, `ReconstructionResult` -### `MeasurementPlane(flux, z, pixel_scale=None, viewing_angle_deg=None, label=None)` +### `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`. -- `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). +- `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). + ### `validate_planes(planes)` Raises `ValueError` if there are fewer than 3 planes, planes have @@ -58,9 +83,14 @@ and `BeamReconstructor.reconstruct`): - `purity: dict[(p, l), (power_fraction, phase_rad)]` - `reconstructed_field: np.ndarray` — reconstructed complex field. - `centers: list[(x, y)]` — fitted beam transverse center per plane, meters. -- `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`). +- `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). - `residuals: list[np.ndarray]` — per-plane (measured − modeled) flux maps. Empty when `used_phase_retrieval` is `True`. - `coefficient_uncertainty: dict[(p, l), float]` — 1-sigma uncertainty on @@ -87,17 +117,55 @@ waist radius `w0` (m), waist location `z0` (m), and radiation `wavelength` onto each `(p, l)` in `modes`, returning `dict[(p, l), complex]` coefficients (Riemann-sum inner product; `dx` is the grid spacing). -## `geometry` — `GeometryCalibration` +## `geometry` — `CameraModel`, `CameraModelTolerance`, `GeometryCalibration` -`GeometryCalibration(plane)` wraps a single `MeasurementPlane` and resolves -its pixel-to-physical-coordinate mapping. +### `CameraModel(focal_length_px, position, orientation_deg, principal_point=(0.0, 0.0))` -- `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. +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). ## `noise` — `NoiseEstimator` @@ -121,25 +189,30 @@ pass a `deconvolver` to `BeamReconstructor`. - `deconvolve(image, pixel_scale, noise_to_signal_ratio=1e-3)` — regularized (Wiener) removal of the blur. -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. +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`). ## `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. +`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, *, 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`. +- `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`. ## `fitting` — `ModalFitter`, `generate_mode_shells` @@ -152,16 +225,22 @@ Groups candidate `LG_{p,l}` modes into shells of increasing order ### `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. +coefficients, beam center/pointing, and any nonzero-tolerance camera/`z` +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` +- `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. + 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). ## `phase_retrieval` — `PhaseRetriever`, `propagate_angular_spectrum` @@ -176,11 +255,12 @@ model implicitly assumed by `LGBasis`'s closed-form paraxial modes. ### `PhaseRetriever(wavelength)` -- `retrieve(planes, pixel_scale=None, viewing_angle_deg=None, max_iterations=200) -> PhaseRetrievalResult` +- `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. + basis. All planes are propagated on one common physical grid, derived + from `camera` at the smallest-`z` plane's depth. ### `PhaseRetrievalResult` @@ -192,23 +272,27 @@ table, as `BeamReconstructor` does internally for its fallback path. ## `reconstruct` — `BeamReconstructor` -`BeamReconstructor(w0, z0, wavelength, max_order=4, noise_estimator=None, deconvolver=None, force_phase_retrieval=False, phase_retrieval_residual_threshold=None)` +`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. +`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 (raises `ValueError` if a - plane's `pixel_scale` isn't known). - 3. Runs `ModalFitter(basis, noise_estimator).fit_auto(planes, max_order)`. + 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 and `coefficient_uncertainty` is `NaN` - per mode (phase retrieval doesn't produce a fit covariance). + `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). ## `plotting` — diagnostic visualizations From 1f2557a0e4c52b1087ab0fa0c3039064f786fabf Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 13:04:19 +0200 Subject: [PATCH 10/12] 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 --- examples/full_pipeline_example.py | 104 +++++++++++++++++++++++++----- 1 file changed, 88 insertions(+), 16 deletions(-) diff --git a/examples/full_pipeline_example.py b/examples/full_pipeline_example.py index d3c3a24..ef4c12d 100644 --- a/examples/full_pipeline_example.py +++ b/examples/full_pipeline_example.py @@ -2,10 +2,15 @@ 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. The camera has an unknown transverse offset/pointing and -adds sensor noise; the target also has some thermal-diffusion blur that we -correct for. We then reconstruct the mode purity, beam center/pointing, and -plot the diagnostics. +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: @@ -18,7 +23,10 @@ import matplotlib.pyplot as plt from he11lib import ( BeamReconstructor, + CameraModel, + CameraModelTolerance, DiffusionDeconvolver, + GeometryCalibration, LGBasis, SyntheticBeamGenerator, plot_center_trace, @@ -33,16 +41,65 @@ 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 camera's optical axis -TRUE_POINTING_DEG = 0.15 # beam pointing (tilt) angle -CAMERA_VIEWING_ANGLE_DEG = 5.0 # oblique camera viewing angle (known) -CAMERA_PIXEL_SCALE = 4e-4 # meters/pixel (known calibration) +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 + +# NOTE on the small magnitudes below: at CAMERA_DISTANCE + Z0 ~= 5.5 m with +# an 81x81 frame at PIXEL_SCALE, the imaged footprint is only ~3.2 cm wide. +# A pinhole camera's boresight sweeps laterally by +# (CAMERA_DISTANCE + Z0) * tan(angle) at the target plane, so even a +# fraction of a degree of yaw/pitch error translates to centimeters of +# parallax there -- degrees-scale tilt (as in a naive "mildly tilted" +# choice) would sweep the boresight tens of centimeters away from the beam, +# off the edge of the frame entirely. Small angles here (~0.01-0.03 deg) +# still exercise true keystoning/off-axis projection while keeping the beam +# comfortably inside the frame at every z. Similarly, CAMERA_TOLERANCE's +# yaw/pitch bounds are kept tight: yaw/pitch changes are nearly degenerate +# with the beam's own two pointing angles (both produce an z-dependent +# transverse drift), so a loose (e.g. several-degree) bound lets the +# optimizer trade pointing off against yaw/pitch and converge to a very +# wrong split of the two; roll doesn't share that degeneracy and is given +# more room. +TRUE_CAMERA = CameraModel( + focal_length_px=FOCAL_LENGTH_PX, + position=(0.002, -0.003, -CAMERA_DISTANCE), + orientation_deg=(0.03, -0.02, 0.01), +) +# 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.0015, -0.0025, -CAMERA_DISTANCE), + orientation_deg=(0.028, -0.018, 0.005), +) +CAMERA_TOLERANCE = CameraModelTolerance( + focal_length_px=FOCAL_LENGTH_PX * 0.05, + position=(0.005, 0.005, 0.02), + orientation_deg=(0.01, 0.01, 0.02), +) + # 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 @@ -51,26 +108,31 @@ DWELL_TIME = 0.2 # s def main() -> None: basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) - generator = SyntheticBeamGenerator( - basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=CAMERA_PIXEL_SCALE - ) + 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_deg=TRUE_POINTING_DEG, - viewing_angle_deg=CAMERA_VIEWING_ANGLE_DEG, + 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. + # 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: - plane.flux = blur_deconvolver.blur(plane.flux, plane.pixel_scale) + 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 @@ -80,6 +142,8 @@ def main() -> None: w0=W0, z0=Z0, wavelength=WAVELENGTH, + camera=NOMINAL_CAMERA, + camera_tolerance=CAMERA_TOLERANCE, max_order=1, deconvolver=blur_deconvolver, ) @@ -91,11 +155,19 @@ def main() -> None: ): print(f" LG_{mode[0]},{mode[1]}: {fraction:6.3%} (phase {phase:+.3f} rad)") - print(f"\nFitted pointing angle: {result.pointing_angle_deg:.4f} deg") + 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) From 396e605d5604cfef3b6609bf65d86e70245c962a Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 13:09:21 +0200 Subject: [PATCH 11/12] 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 --- CLAUDE.md | 69 ++++++++++++++++++++++++++++++++++++++----------------- 1 file changed, 48 insertions(+), 21 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 3a47c72..abaf250 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -35,46 +35,63 @@ This project is not (yet) a git repository. ## Architecture 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 +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`, `ReconstructionResult` (the shared input/output - types) and `validate_planes` (>=3 planes, matching shapes, distinct `z`). +- **`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`** — `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. +- **`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, 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. + 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 known ground-truth coefficients/center/pointing/geometry. - Used throughout the test suite and examples to validate the pipeline end-to-end. + `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 (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). + 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. -- **`reconstruct.py`** — `BeamReconstructor`: the orchestrator. Pipeline order: - validate planes → optional deconvolution (requires known `pixel_scale` per plane) → + 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`, and NaN - `coefficient_uncertainty` (no fit covariance available from phase retrieval). + `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()`. @@ -102,3 +119,13 @@ debugging time in this project's history: 2-mode ground truth). When demonstrating growth with deconvolution or noise, set `max_order` close to the true expected mode content rather than generously high, unless the test specifically targets growth behavior itself. +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. From 83ca084945b9f2911937b66cca3f5ffa16b0db55 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 14:04:24 +0200 Subject: [PATCH 12/12] Fix final-review findings: pixel-scale size guard, DRY camera field names effective_pixel_scale now raises a clear ValueError for image shapes too small for its finite-difference indexing, instead of an IndexError. ModalFitter.fit()'s geometry dict now reuses CAMERA_FIELD_NAMES from geometry.py instead of a duplicated hardcoded tuple. Co-Authored-By: Claude Haiku 4.5 --- he11lib/fitting.py | 10 ++-------- he11lib/geometry.py | 5 +++++ tests/test_geometry.py | 9 +++++++++ 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/he11lib/fitting.py b/he11lib/fitting.py index e43c3fd..0c7b399 100644 --- a/he11lib/fitting.py +++ b/he11lib/fitting.py @@ -9,6 +9,7 @@ from scipy.optimize import least_squares from .data import MeasurementPlane, ReconstructionResult, validate_planes from .geometry import ( + CAMERA_FIELD_NAMES, CameraModel, CameraModelTolerance, GeometryCalibration, @@ -166,14 +167,7 @@ class ModalFitter: 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), - )) + geometry: dict[str, float] = dict(zip(CAMERA_FIELD_NAMES, camera_to_values(fitted_camera))) for i in range(len(planes)): geometry[f"z_{i}"] = z_values[i] diff --git a/he11lib/geometry.py b/he11lib/geometry.py index edd25b4..8f016a6 100644 --- a/he11lib/geometry.py +++ b/he11lib/geometry.py @@ -237,6 +237,11 @@ class GeometryCalibration: keystoned. """ rows, cols = image_shape + if rows < 2 or cols < 2: + raise ValueError( + f"image_shape must be at least 2x2 to compute a finite-difference " + f"pixel scale, got {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]) diff --git a/tests/test_geometry.py b/tests/test_geometry.py index 6179841..0ec03b8 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -149,3 +149,12 @@ def test_effective_pixel_scale_matches_on_axis_focal_length(): scale = calib.effective_pixel_scale((41, 41), z) expected = (z - camera_z) / focal_length_px assert scale == pytest.approx(expected, rel=1e-6) + + +def test_effective_pixel_scale_raises_for_image_too_small(): + camera = make_on_axis_camera() + calib = GeometryCalibration(camera) + z = 0.5 + + with pytest.raises(ValueError, match="image_shape must be at least 2x2"): + calib.effective_pixel_scale((1, 1), z)