Rewrite ModalFitter.fit around the CameraModel tolerance mechanism

The optimizer's parameter vector is now built dynamically: LG
coefficients, per-plane center, and both pointing angles stay always
free; each CameraModel field and each plane's z join the fit (bounded to
its +/- tolerance) only when its paired tolerance is nonzero, and are
otherwise substituted as fixed constants.

Also seeds the first mode's coefficient with a small imaginary offset
(1.0 + 0.05j instead of 1.0 + 0j) rather than exactly on the real axis:
for a single mode, intensity depends only on |c| (phase is unobservable),
so Im=0 sits exactly on that flat/degenerate valley, aligned with a
coordinate axis, giving a zero-gradient Jacobian column that destabilized
trf's trust-region step and prevented pointing-angle recovery from a
default (0, 0) initial guess.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

he11lib/fitting.py

@@ -8,7 +8,14 @@ import numpy as np
 from scipy.optimize import least_squares
 
 from .data import MeasurementPlane, ReconstructionResult, validate_planes
-from .geometry import GeometryCalibration
+from .geometry import (
+    CameraModel,
+    CameraModelTolerance,
+    GeometryCalibration,
+    camera_from_values,
+    camera_to_values,
+    tolerance_to_values,
+)
 from .modes import LGBasis
 from .noise import NoiseEstimator
 
@@ -35,19 +42,31 @@ class ModalFitter:
         self,
         planes: list[MeasurementPlane],
         modes: list[tuple[int, int]],
+        camera: CameraModel,
+        camera_tolerance: CameraModelTolerance,
         initial_coefficients: dict[tuple[int, int], complex] | None = None,
         initial_center: tuple[float, float] = (0.0, 0.0),
-        initial_tilt_deg: tuple[float, float] = (0.0, 0.0),
+        initial_pointing_deg: tuple[float, float] = (0.0, 0.0),
-        initial_pixel_scale: float | None = None,
-        initial_viewing_angle_deg: float = 0.0,
     ) -> ReconstructionResult:
-        """Jointly fit complex coefficients for `modes` plus center/tilt/geometry."""
+        """Jointly fit complex coefficients for `modes` plus center/pointing/geometry.
-        validate_planes(planes)
 
-        unknown_scale_idx = [i for i, p in enumerate(planes) if p.pixel_scale is None]
+        Every `CameraModel` field with a nonzero `camera_tolerance` entry,
-        unknown_angle_idx = [i for i, p in enumerate(planes) if p.viewing_angle_deg is None]
+        and every plane whose `z_tolerance` is nonzero, is refined within
+        `[nominal - tolerance, nominal + tolerance]`; zero-tolerance fields
+        are held fixed at their nominal value.
+        """
+        validate_planes(planes)
         weights = [np.sqrt(self.noise_estimator.weights(p.flux)) for p in planes]
 
+        camera_nominal = camera_to_values(camera)
+        camera_tol = tolerance_to_values(camera_tolerance)
+        free_camera_idx = [i for i, t in enumerate(camera_tol) if t > 0]
+
+        free_z_idx = [i for i, p in enumerate(planes) if p.z_tolerance > 0]
+
+        n_modes = len(modes)
+        n_always_free = 2 * n_modes + 4  # coefficients + center(2) + pointing(2)
+
         def pack_initial() -> np.ndarray:
             x: list[float] = []
             for i, mode in enumerate(modes):
@@ -56,98 +75,126 @@ class ModalFitter:
                     # Nonzero seed for every mode: starting a coefficient at
                     # exactly 0+0j sits at a flat/degenerate point for the
                     # optimizer and can prevent it from ever leaving zero.
-                    c = 1.0 + 0j if i == 0 else 0.1 + 0.05j
+                    # A purely-real seed (Im=0) sits exactly on the flat/
+                    # degenerate valley of a single mode's phase (for one
+                    # mode alone, |c|^2 -- not arg(c) -- is all that's
+                    # observable in intensity), giving a zero-gradient
+                    # column that destabilizes trf's trust-region step;
+                    # a small imaginary offset avoids landing on that axis.
+                    c = 1.0 + 0.05j if i == 0 else 0.1 + 0.05j
                 x += [c.real, c.imag]
-            x += [initial_center[0], initial_center[1], initial_tilt_deg[0], initial_tilt_deg[1]]
+            x += [initial_center[0], initial_center[1], initial_pointing_deg[0], initial_pointing_deg[1]]
-            for _ in unknown_scale_idx:
+            for i in free_camera_idx:
-                x.append(initial_pixel_scale if initial_pixel_scale is not None else 1e-4)
+                x.append(camera_nominal[i])
-            for _ in unknown_angle_idx:
+            for i in free_z_idx:
-                x.append(initial_viewing_angle_deg)
+                x.append(planes[i].z)
             return np.array(x, dtype=float)
 
-        n_modes = len(modes)
+        def pack_bounds() -> tuple[np.ndarray, np.ndarray]:
+            lower = [-np.inf] * n_always_free
+            upper = [np.inf] * n_always_free
+            for i in free_camera_idx:
+                lower.append(camera_nominal[i] - camera_tol[i])
+                upper.append(camera_nominal[i] + camera_tol[i])
+            for i in free_z_idx:
+                lower.append(planes[i].z - planes[i].z_tolerance)
+                upper.append(planes[i].z + planes[i].z_tolerance)
+            return np.array(lower), np.array(upper)
 
         def unpack(x: np.ndarray):
             coeffs = {mode: complex(x[2 * i], x[2 * i + 1]) for i, mode in enumerate(modes)}
             offset = 2 * n_modes
-            x0, y0, tilt_x_deg, tilt_y_deg = x[offset : offset + 4]
+            x0, y0, tilt_h_deg, tilt_v_deg = x[offset : offset + 4]
             offset += 4
-            scales = {}
-            for idx in unknown_scale_idx:
-                scales[idx] = x[offset]
-                offset += 1
-            angles = {}
-            for idx in unknown_angle_idx:
-                angles[idx] = x[offset]
-                offset += 1
-            return coeffs, (x0, y0), (tilt_x_deg, tilt_y_deg), scales, angles
 
-        def plane_center(x0: float, y0: float, tilt_deg: tuple[float, float], z: float):
+            camera_values = list(camera_nominal)
-            drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[0]))
+            for i in free_camera_idx:
-            drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[1]))
+                camera_values[i] = x[offset]
+                offset += 1
+            fitted_camera = camera_from_values(camera_values)
+
+            z_values = [p.z for p in planes]
+            for i in free_z_idx:
+                z_values[i] = x[offset]
+                offset += 1
+
+            return coeffs, (x0, y0), (tilt_h_deg, tilt_v_deg), fitted_camera, z_values
+
+        def plane_center(x0: float, y0: float, pointing_deg: tuple[float, float], z: float):
+            drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[0]))
+            drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[1]))
             return x0 + drift_x, y0 + drift_y
 
-        def model_flux_for_plane(i: int, plane: MeasurementPlane, coeffs, center0, tilt_deg, scales, angles):
+        def model_flux_for_plane(plane, fitted_camera, z, coeffs, center0, pointing_deg):
-            scale = plane.pixel_scale if plane.pixel_scale is not None else scales[i]
+            calib = GeometryCalibration(fitted_camera)
-            angle = plane.viewing_angle_deg if plane.viewing_angle_deg is not None else angles[i]
+            x_grid, y_grid = calib.physical_coordinates(plane.flux.shape, z)
-            calib = GeometryCalibration(plane)
+            cx, cy = plane_center(center0[0], center0[1], pointing_deg, z)
-            x_grid, y_grid = calib.physical_coordinates(pixel_scale=scale, viewing_angle_deg=angle)
+            field = self.basis.field_superposition(x_grid - cx, y_grid - cy, z, coeffs)
-            cx, cy = plane_center(center0[0], center0[1], tilt_deg, plane.z)
-            field = self.basis.field_superposition(x_grid - cx, y_grid - cy, plane.z, coeffs)
             return np.abs(field) ** 2
 
         def residuals(x: np.ndarray) -> np.ndarray:
-            coeffs, center0, tilt_deg, scales, angles = unpack(x)
+            coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(x)
             parts = []
             for i, plane in enumerate(planes):
-                model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles)
+                model_flux = model_flux_for_plane(
+                    plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg
+                )
                 parts.append(((plane.flux - model_flux) * weights[i]).ravel())
             return np.concatenate(parts)
 
         x0_vec = pack_initial()
-        # 'trf' + x_scale='jac' handles the very different natural magnitudes
+        lower, upper = pack_bounds()
-        # of these parameters (coefficients ~O(1), pixel_scale ~O(1e-3),
+        # 'trf' + x_scale='jac' handles the very different natural
-        # angles ~O(1-90)); plain 'lm' can terminate prematurely on 'xtol'
+        # magnitudes of these parameters (coefficients ~O(1), focal length
-        # because its unscaled step-size test is dominated by the largest
+        # ~O(1e3-1e4), angles ~O(1-90), z ~O(0.1-1)); plain 'lm' can
-        # parameters.
+        # terminate prematurely on 'xtol' because its unscaled step-size
+        # test is dominated by the largest parameters. 'lm' also doesn't
+        # support bounds, which the tolerance mechanism requires.
         opt_result = least_squares(
-            residuals, x0_vec, method="trf", x_scale="jac", max_nfev=5000
+            residuals, x0_vec, method="trf", x_scale="jac", bounds=(lower, upper), max_nfev=5000
         )
 
-        coeffs, center0, tilt_deg, scales, angles = unpack(opt_result.x)
+        coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(opt_result.x)
 
         total_power = sum(abs(c) ** 2 for c in coeffs.values())
         if total_power == 0:
             total_power = 1.0
         purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()}
 
-        centers = [plane_center(center0[0], center0[1], tilt_deg, p.z) for p in planes]
+        centers = [
-        pointing_angle_deg = float(np.hypot(tilt_deg[0], tilt_deg[1]))
+            plane_center(center0[0], center0[1], pointing_deg, z_values[i])
+            for i in range(len(planes))
+        ]
 
-        geometry: dict[str, float] = {}
+        geometry: dict[str, float] = dict(zip(
+            (
+                "focal_length_px", "position_x", "position_y", "position_z",
+                "yaw_deg", "pitch_deg", "roll_deg",
+                "principal_point_x", "principal_point_y",
+            ),
+            camera_to_values(fitted_camera),
+        ))
         for i in range(len(planes)):
-            geometry[f"pixel_scale_{i}"] = (
+            geometry[f"z_{i}"] = z_values[i]
-                planes[i].pixel_scale if planes[i].pixel_scale is not None else scales[i]
-            )
-            geometry[f"viewing_angle_deg_{i}"] = (
-                planes[i].viewing_angle_deg if planes[i].viewing_angle_deg is not None else angles[i]
-            )
 
         residual_maps = []
         for i, plane in enumerate(planes):
-            model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles)
+            model_flux = model_flux_for_plane(
+                plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg
+            )
             residual_maps.append(plane.flux - model_flux)
 
         coefficient_uncertainty = self._estimate_uncertainty(opt_result, modes, coeffs, total_power)
 
-        reference_z = min(planes, key=lambda p: abs(p.z - self.basis.z0)).z
+        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, reference_z)
+        field_at_reference = self._field_on_default_grid(coeffs, z_values[reference_idx])
 
         return ReconstructionResult(
             purity=purity,
             reconstructed_field=field_at_reference,
             centers=centers,
-            pointing_angle_deg=pointing_angle_deg,
+            pointing_angle_horizontal_deg=pointing_deg[0],
+            pointing_angle_vertical_deg=pointing_deg[1],
             geometry=geometry,
             residuals=residual_maps,
             coefficient_uncertainty=coefficient_uncertainty,

tests/test_fitting.py

@@ -3,6 +3,7 @@ import pytest
 
 from he11lib.data import validate_planes
 from he11lib.fitting import ModalFitter, generate_mode_shells
+from he11lib.geometry import CameraModel, CameraModelTolerance
 from he11lib.modes import LGBasis
 from he11lib.synthetic import SyntheticBeamGenerator
 
@@ -10,6 +11,7 @@ W0 = 5e-3
 Z0 = 0.5
 WAVELENGTH = 1.76e-3
 PIXEL_SCALE = 4e-4
+CAMERA_DISTANCE = 5.0
 IMAGE_SHAPE = (61, 61)
 Z_LIST = [0.35, 0.5, 0.65, 0.8]
 
@@ -18,8 +20,21 @@ def make_basis():
     return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
 
 
-def make_generator(basis):
+def make_camera(pixel_scale=PIXEL_SCALE, position=(0.0, 0.0, -CAMERA_DISTANCE), orientation_deg=(0.0, 0.0, 0.0)):
-    return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE)
+    focal_length_px = (CAMERA_DISTANCE + Z0) / pixel_scale
+    return CameraModel(
+        focal_length_px=focal_length_px, position=position, orientation_deg=orientation_deg
+    )
+
+
+def zero_tolerance():
+    return CameraModelTolerance(
+        focal_length_px=0.0, position=(0.0, 0.0, 0.0), orientation_deg=(0.0, 0.0, 0.0)
+    )
+
+
+def make_generator(basis, camera):
+    return SyntheticBeamGenerator(basis=basis, camera=camera)
 
 
 def test_generate_mode_shells_orders_by_2p_plus_abs_l():
@@ -31,13 +46,14 @@ def test_generate_mode_shells_orders_by_2p_plus_abs_l():
 
 def test_fit_recovers_pure_fundamental_mode():
     basis = make_basis()
-    gen = make_generator(basis)
+    camera = make_camera()
+    gen = make_generator(basis, camera)
     planes = gen.generate(
-        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0
+        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=0
     )
 
     fitter = ModalFitter(basis)
-    result = fitter.fit(planes, modes=[(0, 0)])
+    result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())
 
     power_fraction, _ = result.purity[(0, 0)]
     assert power_fraction == pytest.approx(1.0, abs=1e-6)
@@ -48,12 +64,17 @@ def test_fit_recovers_pure_fundamental_mode():
 
 def test_fit_recovers_two_mode_purity_ratio():
     basis = make_basis()
-    gen = make_generator(basis)
+    camera = make_camera()
+    gen = make_generator(basis, camera)
     true_coeffs = {(0, 0): 0.9 + 0j, (1, 0): 0.3 + 0.1j}
-    planes = gen.generate(coefficients=true_coeffs, z_list=Z_LIST, noise_std=1e-4, seed=1)
+    planes = gen.generate(
+        coefficients=true_coeffs, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=1
+    )
 
     fitter = ModalFitter(basis)
-    result = fitter.fit(planes, modes=list(true_coeffs.keys()))
+    result = fitter.fit(
+        planes, modes=list(true_coeffs.keys()), camera=camera, camera_tolerance=zero_tolerance()
+    )
 
     true_total = sum(abs(c) ** 2 for c in true_coeffs.values())
     for mode, c in true_coeffs.items():
@@ -64,49 +85,161 @@ def test_fit_recovers_two_mode_purity_ratio():
 
 def test_fit_recovers_center_offset():
     basis = make_basis()
-    gen = make_generator(basis)
+    camera = make_camera()
+    gen = make_generator(basis, camera)
     true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
     planes = gen.generate(
         coefficients={(0, 0): 1.0 + 0j},
         z_list=Z_LIST,
+        image_shape=IMAGE_SHAPE,
         center=true_center,
         noise_std=1e-4,
         seed=2,
     )
 
     fitter = ModalFitter(basis)
-    result = fitter.fit(planes, modes=[(0, 0)], initial_center=true_center)
+    result = fitter.fit(
+        planes,
+        modes=[(0, 0)],
+        camera=camera,
+        camera_tolerance=zero_tolerance(),
+        initial_center=true_center,
+    )
 
     for cx, cy in result.centers:
         assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE)
         assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE)
 
 
-def test_fit_recovers_unknown_pixel_scale():
+def test_fit_recovers_pointing_angles_independently():
-    # Use a coarser pixel scale so the (much wider, far-field) beam at the
-    # outer z distances still fits within the frame -- otherwise pixel scale
-    # becomes unobservable from clipped images.
     basis = make_basis()
-    local_pixel_scale = 1.5e-3
+    camera = make_camera()
-    gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=local_pixel_scale)
+    gen = make_generator(basis, camera)
-    planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=3)
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j},
+        z_list=Z_LIST,
+        image_shape=IMAGE_SHAPE,
+        pointing_angle_horizontal_deg=0.3,
+        pointing_angle_vertical_deg=-0.15,
+        noise_std=1e-4,
+        seed=6,
+    )
 
-    # hide the known calibration to force the fitter to solve for it
+    fitter = ModalFitter(basis)
-    for plane in planes:
+    result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())
-        plane.pixel_scale = None
+
-        plane.viewing_angle_deg = None
+    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,
+        planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=zero_tolerance()
-        modes=[(0, 0)],
-        initial_pixel_scale=local_pixel_scale * 1.1,
-        initial_viewing_angle_deg=0.0,
     )
 
-    fitted_scales = [result.geometry[f"pixel_scale_{i}"] for i in range(len(planes))]
+    assert result.geometry["focal_length_px"] == wrong_focal_length
-    for scale in fitted_scales:
+
-        assert scale == pytest.approx(local_pixel_scale, rel=0.05)
+
+def test_fit_recovers_offset_camera_field_within_tolerance():
+    # A tolerance>0 field recovers a ground-truth offset from nominal, but
+    # within its band.
+    basis = make_basis()
+    true_camera = make_camera()
+    gen = make_generator(basis, true_camera)
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=8
+    )
+
+    offset = true_camera.focal_length_px * 0.02  # 2% off nominal
+    nominal_camera = CameraModel(
+        focal_length_px=true_camera.focal_length_px + offset,
+        position=true_camera.position,
+        orientation_deg=true_camera.orientation_deg,
+    )
+    tolerance = CameraModelTolerance(
+        focal_length_px=true_camera.focal_length_px * 0.05,  # +/-5% band
+        position=(0.0, 0.0, 0.0),
+        orientation_deg=(0.0, 0.0, 0.0),
+    )
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit(planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tolerance)
+
+    assert result.geometry["focal_length_px"] == pytest.approx(
+        true_camera.focal_length_px, rel=0.02
+    )
+
+
+def test_fit_clips_out_of_band_ground_truth_to_bound():
+    # A ground truth placed outside a deliberately too-tight band is
+    # clipped to the bound rather than escaping it.
+    basis = make_basis()
+    true_camera = make_camera()
+    gen = make_generator(basis, true_camera)
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=9
+    )
+
+    # nominal is 10% off true, but the band only allows +/-1%.
+    nominal_focal_length = true_camera.focal_length_px * 1.10
+    nominal_camera = CameraModel(
+        focal_length_px=nominal_focal_length,
+        position=true_camera.position,
+        orientation_deg=true_camera.orientation_deg,
+    )
+    tight_tolerance = CameraModelTolerance(
+        focal_length_px=nominal_focal_length * 0.01,
+        position=(0.0, 0.0, 0.0),
+        orientation_deg=(0.0, 0.0, 0.0),
+    )
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit(
+        planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tight_tolerance
+    )
+
+    lower_bound = nominal_focal_length - tight_tolerance.focal_length_px
+    assert result.geometry["focal_length_px"] == pytest.approx(lower_bound, rel=1e-3)
+
+
+def test_fit_recovers_offset_z_within_tolerance():
+    basis = make_basis()
+    camera = make_camera()
+    gen = make_generator(basis, camera)
+    true_z_list = [0.35, 0.5, 0.65, 0.8]
+    offsets = {z: 0.01 for z in true_z_list}  # nominal is 1 cm off true
+    planes = gen.generate(
+        coefficients={(0, 0): 1.0 + 0j},
+        z_list=true_z_list,
+        image_shape=IMAGE_SHAPE,
+        nominal_z_offsets=offsets,
+        z_tolerance=0.03,
+        noise_std=1e-4,
+        seed=10,
+    )
+
+    fitter = ModalFitter(basis)
+    result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())
+
+    for i, true_z in enumerate(true_z_list):
+        assert result.geometry[f"z_{i}"] == pytest.approx(true_z, abs=0.005)
 
 
 def test_fit_auto_does_not_add_modes_for_pure_fundamental():