fabb3d4efc
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 <noreply@anthropic.com>
3362 lines
127 KiB
Markdown
3362 lines
127 KiB
Markdown
# 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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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 <noreply@anthropic.com>
|
|
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
|
|
|
|
<describe the specific issue found and fixed here>
|
|
|
|
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
|
EOF
|
|
)"
|
|
```
|
|
|
|
If Steps 1-2 passed cleanly with no fixes needed, skip this step entirely
|
|
— do not create an empty commit.
|