diff --git a/he11lib/fitting.py b/he11lib/fitting.py index 4c3184f..37bb6f7 100644 --- a/he11lib/fitting.py +++ b/he11lib/fitting.py @@ -8,7 +8,14 @@ import numpy as np from scipy.optimize import least_squares from .data import MeasurementPlane, ReconstructionResult, validate_planes -from .geometry import GeometryCalibration +from .geometry import ( + CameraModel, + CameraModelTolerance, + GeometryCalibration, + camera_from_values, + camera_to_values, + tolerance_to_values, +) from .modes import LGBasis from .noise import NoiseEstimator @@ -35,19 +42,31 @@ class ModalFitter: self, planes: list[MeasurementPlane], modes: list[tuple[int, int]], + camera: CameraModel, + camera_tolerance: CameraModelTolerance, initial_coefficients: dict[tuple[int, int], complex] | None = None, initial_center: tuple[float, float] = (0.0, 0.0), - initial_tilt_deg: tuple[float, float] = (0.0, 0.0), - initial_pixel_scale: float | None = None, - initial_viewing_angle_deg: float = 0.0, + initial_pointing_deg: tuple[float, float] = (0.0, 0.0), ) -> ReconstructionResult: - """Jointly fit complex coefficients for `modes` plus center/tilt/geometry.""" - validate_planes(planes) + """Jointly fit complex coefficients for `modes` plus center/pointing/geometry. - unknown_scale_idx = [i for i, p in enumerate(planes) if p.pixel_scale is None] - unknown_angle_idx = [i for i, p in enumerate(planes) if p.viewing_angle_deg is None] + Every `CameraModel` field with a nonzero `camera_tolerance` entry, + and every plane whose `z_tolerance` is nonzero, is refined within + `[nominal - tolerance, nominal + tolerance]`; zero-tolerance fields + are held fixed at their nominal value. + """ + validate_planes(planes) weights = [np.sqrt(self.noise_estimator.weights(p.flux)) for p in planes] + camera_nominal = camera_to_values(camera) + camera_tol = tolerance_to_values(camera_tolerance) + free_camera_idx = [i for i, t in enumerate(camera_tol) if t > 0] + + free_z_idx = [i for i, p in enumerate(planes) if p.z_tolerance > 0] + + n_modes = len(modes) + n_always_free = 2 * n_modes + 4 # coefficients + center(2) + pointing(2) + def pack_initial() -> np.ndarray: x: list[float] = [] for i, mode in enumerate(modes): @@ -56,98 +75,126 @@ class ModalFitter: # Nonzero seed for every mode: starting a coefficient at # exactly 0+0j sits at a flat/degenerate point for the # optimizer and can prevent it from ever leaving zero. - c = 1.0 + 0j if i == 0 else 0.1 + 0.05j + # A purely-real seed (Im=0) sits exactly on the flat/ + # degenerate valley of a single mode's phase (for one + # mode alone, |c|^2 -- not arg(c) -- is all that's + # observable in intensity), giving a zero-gradient + # column that destabilizes trf's trust-region step; + # a small imaginary offset avoids landing on that axis. + c = 1.0 + 0.05j if i == 0 else 0.1 + 0.05j x += [c.real, c.imag] - x += [initial_center[0], initial_center[1], initial_tilt_deg[0], initial_tilt_deg[1]] - for _ in unknown_scale_idx: - x.append(initial_pixel_scale if initial_pixel_scale is not None else 1e-4) - for _ in unknown_angle_idx: - x.append(initial_viewing_angle_deg) + x += [initial_center[0], initial_center[1], initial_pointing_deg[0], initial_pointing_deg[1]] + for i in free_camera_idx: + x.append(camera_nominal[i]) + for i in free_z_idx: + x.append(planes[i].z) return np.array(x, dtype=float) - n_modes = len(modes) + def pack_bounds() -> tuple[np.ndarray, np.ndarray]: + lower = [-np.inf] * n_always_free + upper = [np.inf] * n_always_free + for i in free_camera_idx: + lower.append(camera_nominal[i] - camera_tol[i]) + upper.append(camera_nominal[i] + camera_tol[i]) + for i in free_z_idx: + lower.append(planes[i].z - planes[i].z_tolerance) + upper.append(planes[i].z + planes[i].z_tolerance) + return np.array(lower), np.array(upper) def unpack(x: np.ndarray): coeffs = {mode: complex(x[2 * i], x[2 * i + 1]) for i, mode in enumerate(modes)} offset = 2 * n_modes - x0, y0, tilt_x_deg, tilt_y_deg = x[offset : offset + 4] + x0, y0, tilt_h_deg, tilt_v_deg = x[offset : offset + 4] offset += 4 - scales = {} - for idx in unknown_scale_idx: - scales[idx] = x[offset] - offset += 1 - angles = {} - for idx in unknown_angle_idx: - angles[idx] = x[offset] - offset += 1 - return coeffs, (x0, y0), (tilt_x_deg, tilt_y_deg), scales, angles - def plane_center(x0: float, y0: float, tilt_deg: tuple[float, float], z: float): - drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[0])) - drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[1])) + camera_values = list(camera_nominal) + for i in free_camera_idx: + camera_values[i] = x[offset] + offset += 1 + fitted_camera = camera_from_values(camera_values) + + z_values = [p.z for p in planes] + for i in free_z_idx: + z_values[i] = x[offset] + offset += 1 + + return coeffs, (x0, y0), (tilt_h_deg, tilt_v_deg), fitted_camera, z_values + + def plane_center(x0: float, y0: float, pointing_deg: tuple[float, float], z: float): + drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[0])) + drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[1])) return x0 + drift_x, y0 + drift_y - def model_flux_for_plane(i: int, plane: MeasurementPlane, coeffs, center0, tilt_deg, scales, angles): - scale = plane.pixel_scale if plane.pixel_scale is not None else scales[i] - angle = plane.viewing_angle_deg if plane.viewing_angle_deg is not None else angles[i] - calib = GeometryCalibration(plane) - x_grid, y_grid = calib.physical_coordinates(pixel_scale=scale, viewing_angle_deg=angle) - cx, cy = plane_center(center0[0], center0[1], tilt_deg, plane.z) - field = self.basis.field_superposition(x_grid - cx, y_grid - cy, plane.z, coeffs) + def model_flux_for_plane(plane, fitted_camera, z, coeffs, center0, pointing_deg): + calib = GeometryCalibration(fitted_camera) + x_grid, y_grid = calib.physical_coordinates(plane.flux.shape, z) + cx, cy = plane_center(center0[0], center0[1], pointing_deg, z) + field = self.basis.field_superposition(x_grid - cx, y_grid - cy, z, coeffs) return np.abs(field) ** 2 def residuals(x: np.ndarray) -> np.ndarray: - coeffs, center0, tilt_deg, scales, angles = unpack(x) + coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(x) parts = [] for i, plane in enumerate(planes): - model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles) + model_flux = model_flux_for_plane( + plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg + ) parts.append(((plane.flux - model_flux) * weights[i]).ravel()) return np.concatenate(parts) x0_vec = pack_initial() - # 'trf' + x_scale='jac' handles the very different natural magnitudes - # of these parameters (coefficients ~O(1), pixel_scale ~O(1e-3), - # angles ~O(1-90)); plain 'lm' can terminate prematurely on 'xtol' - # because its unscaled step-size test is dominated by the largest - # parameters. + lower, upper = pack_bounds() + # 'trf' + x_scale='jac' handles the very different natural + # magnitudes of these parameters (coefficients ~O(1), focal length + # ~O(1e3-1e4), angles ~O(1-90), z ~O(0.1-1)); plain 'lm' can + # terminate prematurely on 'xtol' because its unscaled step-size + # test is dominated by the largest parameters. 'lm' also doesn't + # support bounds, which the tolerance mechanism requires. opt_result = least_squares( - residuals, x0_vec, method="trf", x_scale="jac", max_nfev=5000 + residuals, x0_vec, method="trf", x_scale="jac", bounds=(lower, upper), max_nfev=5000 ) - coeffs, center0, tilt_deg, scales, angles = unpack(opt_result.x) + coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(opt_result.x) total_power = sum(abs(c) ** 2 for c in coeffs.values()) if total_power == 0: total_power = 1.0 purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()} - centers = [plane_center(center0[0], center0[1], tilt_deg, p.z) for p in planes] - pointing_angle_deg = float(np.hypot(tilt_deg[0], tilt_deg[1])) + centers = [ + plane_center(center0[0], center0[1], pointing_deg, z_values[i]) + for i in range(len(planes)) + ] - geometry: dict[str, float] = {} + geometry: dict[str, float] = dict(zip( + ( + "focal_length_px", "position_x", "position_y", "position_z", + "yaw_deg", "pitch_deg", "roll_deg", + "principal_point_x", "principal_point_y", + ), + camera_to_values(fitted_camera), + )) for i in range(len(planes)): - geometry[f"pixel_scale_{i}"] = ( - planes[i].pixel_scale if planes[i].pixel_scale is not None else scales[i] - ) - geometry[f"viewing_angle_deg_{i}"] = ( - planes[i].viewing_angle_deg if planes[i].viewing_angle_deg is not None else angles[i] - ) + geometry[f"z_{i}"] = z_values[i] residual_maps = [] for i, plane in enumerate(planes): - model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles) + model_flux = model_flux_for_plane( + plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg + ) residual_maps.append(plane.flux - model_flux) coefficient_uncertainty = self._estimate_uncertainty(opt_result, modes, coeffs, total_power) - reference_z = min(planes, key=lambda p: abs(p.z - self.basis.z0)).z - field_at_reference = self._field_on_default_grid(coeffs, reference_z) + reference_idx = min(range(len(planes)), key=lambda i: abs(z_values[i] - self.basis.z0)) + field_at_reference = self._field_on_default_grid(coeffs, z_values[reference_idx]) return ReconstructionResult( purity=purity, reconstructed_field=field_at_reference, centers=centers, - pointing_angle_deg=pointing_angle_deg, + pointing_angle_horizontal_deg=pointing_deg[0], + pointing_angle_vertical_deg=pointing_deg[1], geometry=geometry, residuals=residual_maps, coefficient_uncertainty=coefficient_uncertainty, diff --git a/tests/test_fitting.py b/tests/test_fitting.py index ff0c170..cd58cdb 100644 --- a/tests/test_fitting.py +++ b/tests/test_fitting.py @@ -3,6 +3,7 @@ import pytest from he11lib.data import validate_planes from he11lib.fitting import ModalFitter, generate_mode_shells +from he11lib.geometry import CameraModel, CameraModelTolerance from he11lib.modes import LGBasis from he11lib.synthetic import SyntheticBeamGenerator @@ -10,6 +11,7 @@ W0 = 5e-3 Z0 = 0.5 WAVELENGTH = 1.76e-3 PIXEL_SCALE = 4e-4 +CAMERA_DISTANCE = 5.0 IMAGE_SHAPE = (61, 61) Z_LIST = [0.35, 0.5, 0.65, 0.8] @@ -18,8 +20,21 @@ def make_basis(): return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) -def make_generator(basis): - return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) +def make_camera(pixel_scale=PIXEL_SCALE, position=(0.0, 0.0, -CAMERA_DISTANCE), orientation_deg=(0.0, 0.0, 0.0)): + focal_length_px = (CAMERA_DISTANCE + Z0) / pixel_scale + return CameraModel( + focal_length_px=focal_length_px, position=position, orientation_deg=orientation_deg + ) + + +def zero_tolerance(): + return CameraModelTolerance( + focal_length_px=0.0, position=(0.0, 0.0, 0.0), orientation_deg=(0.0, 0.0, 0.0) + ) + + +def make_generator(basis, camera): + return SyntheticBeamGenerator(basis=basis, camera=camera) def test_generate_mode_shells_orders_by_2p_plus_abs_l(): @@ -31,13 +46,14 @@ def test_generate_mode_shells_orders_by_2p_plus_abs_l(): def test_fit_recovers_pure_fundamental_mode(): basis = make_basis() - gen = make_generator(basis) + camera = make_camera() + gen = make_generator(basis, camera) planes = gen.generate( - coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0 + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=0 ) fitter = ModalFitter(basis) - result = fitter.fit(planes, modes=[(0, 0)]) + result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance()) power_fraction, _ = result.purity[(0, 0)] assert power_fraction == pytest.approx(1.0, abs=1e-6) @@ -48,12 +64,17 @@ def test_fit_recovers_pure_fundamental_mode(): def test_fit_recovers_two_mode_purity_ratio(): basis = make_basis() - gen = make_generator(basis) + camera = make_camera() + gen = make_generator(basis, camera) true_coeffs = {(0, 0): 0.9 + 0j, (1, 0): 0.3 + 0.1j} - planes = gen.generate(coefficients=true_coeffs, z_list=Z_LIST, noise_std=1e-4, seed=1) + planes = gen.generate( + coefficients=true_coeffs, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=1 + ) fitter = ModalFitter(basis) - result = fitter.fit(planes, modes=list(true_coeffs.keys())) + result = fitter.fit( + planes, modes=list(true_coeffs.keys()), camera=camera, camera_tolerance=zero_tolerance() + ) true_total = sum(abs(c) ** 2 for c in true_coeffs.values()) for mode, c in true_coeffs.items(): @@ -64,49 +85,161 @@ def test_fit_recovers_two_mode_purity_ratio(): def test_fit_recovers_center_offset(): basis = make_basis() - gen = make_generator(basis) + camera = make_camera() + gen = make_generator(basis, camera) true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE) planes = gen.generate( coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, + image_shape=IMAGE_SHAPE, center=true_center, noise_std=1e-4, seed=2, ) fitter = ModalFitter(basis) - result = fitter.fit(planes, modes=[(0, 0)], initial_center=true_center) + result = fitter.fit( + planes, + modes=[(0, 0)], + camera=camera, + camera_tolerance=zero_tolerance(), + initial_center=true_center, + ) for cx, cy in result.centers: assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE) assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE) -def test_fit_recovers_unknown_pixel_scale(): - # Use a coarser pixel scale so the (much wider, far-field) beam at the - # outer z distances still fits within the frame -- otherwise pixel scale - # becomes unobservable from clipped images. +def test_fit_recovers_pointing_angles_independently(): basis = make_basis() - local_pixel_scale = 1.5e-3 - gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=local_pixel_scale) - planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=3) + camera = make_camera() + gen = make_generator(basis, camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, + z_list=Z_LIST, + image_shape=IMAGE_SHAPE, + pointing_angle_horizontal_deg=0.3, + pointing_angle_vertical_deg=-0.15, + noise_std=1e-4, + seed=6, + ) - # hide the known calibration to force the fitter to solve for it - for plane in planes: - plane.pixel_scale = None - plane.viewing_angle_deg = None + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance()) + + assert result.pointing_angle_horizontal_deg == pytest.approx(0.3, abs=0.05) + assert result.pointing_angle_vertical_deg == pytest.approx(-0.15, abs=0.05) + + +def test_fit_holds_zero_tolerance_camera_field_fixed_at_wrong_nominal(): + # A tolerance=0 field must stay exactly at its (deliberately wrong) + # nominal value rather than being corrected. + basis = make_basis() + true_camera = make_camera() + gen = make_generator(basis, true_camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=7 + ) + + wrong_focal_length = true_camera.focal_length_px * 1.2 + nominal_camera = CameraModel( + focal_length_px=wrong_focal_length, + position=true_camera.position, + orientation_deg=true_camera.orientation_deg, + ) fitter = ModalFitter(basis) result = fitter.fit( - planes, - modes=[(0, 0)], - initial_pixel_scale=local_pixel_scale * 1.1, - initial_viewing_angle_deg=0.0, + planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=zero_tolerance() ) - fitted_scales = [result.geometry[f"pixel_scale_{i}"] for i in range(len(planes))] - for scale in fitted_scales: - assert scale == pytest.approx(local_pixel_scale, rel=0.05) + assert result.geometry["focal_length_px"] == wrong_focal_length + + +def test_fit_recovers_offset_camera_field_within_tolerance(): + # A tolerance>0 field recovers a ground-truth offset from nominal, but + # within its band. + basis = make_basis() + true_camera = make_camera() + gen = make_generator(basis, true_camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=8 + ) + + offset = true_camera.focal_length_px * 0.02 # 2% off nominal + nominal_camera = CameraModel( + focal_length_px=true_camera.focal_length_px + offset, + position=true_camera.position, + orientation_deg=true_camera.orientation_deg, + ) + tolerance = CameraModelTolerance( + focal_length_px=true_camera.focal_length_px * 0.05, # +/-5% band + position=(0.0, 0.0, 0.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tolerance) + + assert result.geometry["focal_length_px"] == pytest.approx( + true_camera.focal_length_px, rel=0.02 + ) + + +def test_fit_clips_out_of_band_ground_truth_to_bound(): + # A ground truth placed outside a deliberately too-tight band is + # clipped to the bound rather than escaping it. + basis = make_basis() + true_camera = make_camera() + gen = make_generator(basis, true_camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=9 + ) + + # nominal is 10% off true, but the band only allows +/-1%. + nominal_focal_length = true_camera.focal_length_px * 1.10 + nominal_camera = CameraModel( + focal_length_px=nominal_focal_length, + position=true_camera.position, + orientation_deg=true_camera.orientation_deg, + ) + tight_tolerance = CameraModelTolerance( + focal_length_px=nominal_focal_length * 0.01, + position=(0.0, 0.0, 0.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + + fitter = ModalFitter(basis) + result = fitter.fit( + planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tight_tolerance + ) + + lower_bound = nominal_focal_length - tight_tolerance.focal_length_px + assert result.geometry["focal_length_px"] == pytest.approx(lower_bound, rel=1e-3) + + +def test_fit_recovers_offset_z_within_tolerance(): + basis = make_basis() + camera = make_camera() + gen = make_generator(basis, camera) + true_z_list = [0.35, 0.5, 0.65, 0.8] + offsets = {z: 0.01 for z in true_z_list} # nominal is 1 cm off true + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, + z_list=true_z_list, + image_shape=IMAGE_SHAPE, + nominal_z_offsets=offsets, + z_tolerance=0.03, + noise_std=1e-4, + seed=10, + ) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance()) + + for i, true_z in enumerate(true_z_list): + assert result.geometry[f"z_{i}"] == pytest.approx(true_z, abs=0.005) def test_fit_auto_does_not_add_modes_for_pure_fundamental():