From 57dc7d743ebc506361f34be3c83ed2898f5bc5d2 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 12:26:18 +0200 Subject: [PATCH] Wire CameraModel/CameraModelTolerance through BeamReconstructor BeamReconstructor now requires camera/camera_tolerance, threading them into ModalFitter.fit_auto and PhaseRetriever.retrieve, and uses GeometryCalibration.effective_pixel_scale for deconvolution instead of the removed MeasurementPlane.pixel_scale. Co-Authored-By: Claude Sonnet 5 --- he11lib/reconstruct.py | 29 +++++--- tests/test_reconstruct.py | 144 +++++++++++++++++++++++++++++++++----- 2 files changed, 145 insertions(+), 28 deletions(-) diff --git a/he11lib/reconstruct.py b/he11lib/reconstruct.py index 0d22b4e..be4ab9a 100644 --- a/he11lib/reconstruct.py +++ b/he11lib/reconstruct.py @@ -9,6 +9,7 @@ import numpy as np from .data import MeasurementPlane, ReconstructionResult, validate_planes from .deconvolution import DiffusionDeconvolver from .fitting import ModalFitter, generate_mode_shells +from .geometry import CameraModel, CameraModelTolerance, GeometryCalibration from .modes import LGBasis from .noise import NoiseEstimator from .phase_retrieval import PhaseRetriever @@ -25,12 +26,15 @@ class BeamReconstructor: Parameters ---------- w0, z0, wavelength : known reference beam parameters (see `LGBasis`). + camera : nominal shared CameraModel (position/orientation/intrinsics). + camera_tolerance : per-field +/- refinement bound for `camera`; a + zero-tolerance field is held fixed at its nominal value. max_order : cap on automatic candidate-mode-set growth (see `ModalFitter.fit_auto`), and also the mode set used to project the phase-retrieval fallback's recovered field onto the LG basis. noise_estimator : shared noise model; defaults to `NoiseEstimator()`. - deconvolver : if given, each plane's flux is deblurred (its - `pixel_scale` must be known) before fitting. + deconvolver : if given, each plane's flux is deblurred (using + `GeometryCalibration(camera).effective_pixel_scale`) before fitting. force_phase_retrieval : if True, always run the phase-retrieval fallback instead of the modal fit. phase_retrieval_residual_threshold : if set (and `force_phase_retrieval` @@ -43,6 +47,8 @@ class BeamReconstructor: w0: float, z0: float, wavelength: float, + camera: CameraModel, + camera_tolerance: CameraModelTolerance, max_order: int = 4, noise_estimator: NoiseEstimator | None = None, deconvolver: DiffusionDeconvolver | None = None, @@ -51,6 +57,8 @@ class BeamReconstructor: ): self.basis = LGBasis(w0=w0, z0=z0, wavelength=wavelength) self.wavelength = wavelength + self.camera = camera + self.camera_tolerance = camera_tolerance self.max_order = max_order self.noise_estimator = noise_estimator or NoiseEstimator() self.deconvolver = deconvolver @@ -63,7 +71,9 @@ class BeamReconstructor: planes = self._deconvolve(planes) fitter = ModalFitter(self.basis, self.noise_estimator) - result = fitter.fit_auto(planes, max_order=self.max_order) + result = fitter.fit_auto( + planes, self.camera, self.camera_tolerance, max_order=self.max_order + ) if self.force_phase_retrieval or self._residual_too_high(result, planes): result = self._phase_retrieval_fallback(planes) @@ -73,13 +83,11 @@ class BeamReconstructor: def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]: if self.deconvolver is None: return planes + calib = GeometryCalibration(self.camera) deblurred = [] for plane in planes: - if plane.pixel_scale is None: - raise ValueError( - "Deconvolution requires a known pixel_scale on every MeasurementPlane." - ) - flux = self.deconvolver.deconvolve(plane.flux, plane.pixel_scale) + pixel_scale = calib.effective_pixel_scale(plane.flux.shape, plane.z) + flux = self.deconvolver.deconvolve(plane.flux, pixel_scale) deblurred.append(replace(plane, flux=flux)) return deblurred @@ -101,7 +109,7 @@ class BeamReconstructor: self, planes: list[MeasurementPlane] ) -> ReconstructionResult: retriever = PhaseRetriever(self.wavelength) - pr_result = retriever.retrieve(planes) + pr_result = retriever.retrieve(planes, self.camera) modes = [mode for shell in generate_mode_shells(self.max_order) for mode in shell] dx = float(pr_result.x[0, 1] - pr_result.x[0, 0]) @@ -116,7 +124,8 @@ class BeamReconstructor: purity=purity, reconstructed_field=pr_result.field, centers=[pr_result.center for _ in planes], - pointing_angle_deg=float("nan"), + pointing_angle_horizontal_deg=float("nan"), + pointing_angle_vertical_deg=float("nan"), geometry={}, residuals=[], coefficient_uncertainty={mode: float("nan") for mode in modes}, diff --git a/tests/test_reconstruct.py b/tests/test_reconstruct.py index 2c20c91..86ea393 100644 --- a/tests/test_reconstruct.py +++ b/tests/test_reconstruct.py @@ -4,6 +4,7 @@ import pytest from he11lib.deconvolution import DiffusionDeconvolver from he11lib.fitting import ModalFitter +from he11lib.geometry import CameraModel, CameraModelTolerance, GeometryCalibration from he11lib.modes import LGBasis from he11lib.reconstruct import BeamReconstructor from he11lib.synthetic import SyntheticBeamGenerator @@ -12,6 +13,7 @@ W0 = 5e-3 Z0 = 0.5 WAVELENGTH = 1.76e-3 PIXEL_SCALE = 4e-4 +CAMERA_DISTANCE = 5.0 IMAGE_SHAPE = (61, 61) Z_LIST = [0.35, 0.5, 0.65, 0.8] @@ -20,16 +22,36 @@ def make_basis(): return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) -def make_generator(basis): - return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) +def make_camera(): + focal_length_px = (CAMERA_DISTANCE + Z0) / PIXEL_SCALE + return CameraModel( + focal_length_px=focal_length_px, + position=(0.0, 0.0, -CAMERA_DISTANCE), + orientation_deg=(0.0, 0.0, 0.0), + ) + + +def zero_tolerance(): + return CameraModelTolerance( + focal_length_px=0.0, position=(0.0, 0.0, 0.0), orientation_deg=(0.0, 0.0, 0.0) + ) + + +def make_generator(basis, camera): + return SyntheticBeamGenerator(basis=basis, camera=camera) def test_reconstruct_recovers_pure_mode_purity(): basis = make_basis() - gen = make_generator(basis) - planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0) + camera = make_camera() + gen = make_generator(basis, camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=0 + ) - reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2) + reconstructor = BeamReconstructor( + w0=W0, z0=Z0, wavelength=WAVELENGTH, camera=camera, camera_tolerance=zero_tolerance(), max_order=2 + ) result = reconstructor.reconstruct(planes) power_fraction, _ = result.purity[(0, 0)] @@ -39,13 +61,21 @@ def test_reconstruct_recovers_pure_mode_purity(): def test_reconstruct_recovers_center_offset(): basis = make_basis() - gen = make_generator(basis) + camera = make_camera() + gen = make_generator(basis, camera) true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE) planes = gen.generate( - coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, center=true_center, noise_std=1e-4, seed=1 + coefficients={(0, 0): 1.0 + 0j}, + z_list=Z_LIST, + image_shape=IMAGE_SHAPE, + center=true_center, + noise_std=1e-4, + seed=1, ) - reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2) + reconstructor = BeamReconstructor( + w0=W0, z0=Z0, wavelength=WAVELENGTH, camera=camera, camera_tolerance=zero_tolerance(), max_order=2 + ) result = reconstructor.reconstruct(planes) for cx, cy in result.centers: @@ -55,21 +85,34 @@ def test_reconstruct_recovers_center_offset(): def test_reconstruct_with_deconvolution_corrects_blur(): basis = make_basis() - gen = make_generator(basis) - planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=2) + camera = make_camera() + gen = make_generator(basis, camera) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=2 + ) deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=30.0) + calib = GeometryCalibration(camera) blurred_planes = [ - replace(p, flux=deconvolver.blur(p.flux, p.pixel_scale)) for p in planes + replace(p, flux=deconvolver.blur(p.flux, calib.effective_pixel_scale(p.flux.shape, p.z))) + for p in planes ] # Without deconvolution, blur should measurably hurt purity recovery. fitter = ModalFitter(basis) - result_no_deconv = fitter.fit(blurred_planes, modes=[(0, 0), (1, 0), (0, 1)]) + result_no_deconv = fitter.fit( + blurred_planes, modes=[(0, 0), (1, 0), (0, 1)], camera=camera, camera_tolerance=zero_tolerance() + ) purity_no_deconv, _ = result_no_deconv.purity[(0, 0)] reconstructor = BeamReconstructor( - w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, deconvolver=deconvolver + w0=W0, + z0=Z0, + wavelength=WAVELENGTH, + camera=camera, + camera_tolerance=zero_tolerance(), + max_order=2, + deconvolver=deconvolver, ) result = reconstructor.reconstruct(blurred_planes) purity_with_deconv, _ = result.purity[(0, 0)] @@ -80,12 +123,21 @@ def test_reconstruct_with_deconvolution_corrects_blur(): def test_reconstruct_forces_phase_retrieval_fallback(): basis = make_basis() - gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4) + camera = make_camera() + gen = SyntheticBeamGenerator(basis=basis, camera=camera) z_list = [0.47, 0.5, 0.53] - planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=3) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE, noise_std=1e-5, seed=3 + ) reconstructor = BeamReconstructor( - w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, force_phase_retrieval=True + w0=W0, + z0=Z0, + wavelength=WAVELENGTH, + camera=camera, + camera_tolerance=zero_tolerance(), + max_order=2, + force_phase_retrieval=True, ) result = reconstructor.reconstruct(planes) @@ -96,17 +148,73 @@ def test_reconstruct_forces_phase_retrieval_fallback(): def test_reconstruct_falls_back_automatically_on_high_residual(): basis = make_basis() - gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4) + camera = make_camera() + gen = SyntheticBeamGenerator(basis=basis, camera=camera) z_list = [0.47, 0.5, 0.53] - planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=4) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE, noise_std=1e-5, seed=4 + ) reconstructor = BeamReconstructor( w0=W0, z0=Z0, wavelength=WAVELENGTH, + camera=camera, + camera_tolerance=zero_tolerance(), max_order=2, phase_retrieval_residual_threshold=1e-8, ) result = reconstructor.reconstruct(planes) assert result.used_phase_retrieval is True + + +def test_reconstruct_recovers_camera_and_z_offset_from_nominal(): + # End-to-end: ground truth is offset from the nominal camera/z inputs + # (within their tolerances), simulating realistic calibration error. + basis = make_basis() + true_camera = make_camera() + gen = make_generator(basis, true_camera) + true_z_list = Z_LIST + z_offsets = {z: 0.01 for z in true_z_list} + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, + z_list=true_z_list, + image_shape=IMAGE_SHAPE, + nominal_z_offsets=z_offsets, + z_tolerance=0.03, + pointing_angle_horizontal_deg=0.2, + pointing_angle_vertical_deg=-0.1, + noise_std=1e-4, + seed=13, + ) + + nominal_focal_offset = true_camera.focal_length_px * 0.03 + nominal_camera = CameraModel( + focal_length_px=true_camera.focal_length_px + nominal_focal_offset, + position=true_camera.position, + orientation_deg=true_camera.orientation_deg, + ) + tolerance = CameraModelTolerance( + focal_length_px=true_camera.focal_length_px * 0.1, + position=(0.0, 0.0, 0.0), + orientation_deg=(0.0, 0.0, 0.0), + ) + + reconstructor = BeamReconstructor( + w0=W0, + z0=Z0, + wavelength=WAVELENGTH, + camera=nominal_camera, + camera_tolerance=tolerance, + max_order=1, + ) + result = reconstructor.reconstruct(planes) + + power_fraction, _ = result.purity[(0, 0)] + assert power_fraction > 0.95 + assert result.pointing_angle_horizontal_deg == pytest.approx(0.2, abs=0.1) + assert result.pointing_angle_vertical_deg == pytest.approx(-0.1, abs=0.1) + assert result.geometry["focal_length_px"] == pytest.approx(true_camera.focal_length_px, rel=0.03) + for i, true_z in enumerate(true_z_list): + assert result.geometry[f"z_{i}"] == pytest.approx(true_z, abs=0.005)