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>
This commit is contained in:
Martino Ferrari
2026-07-03 12:26:18 +02:00
parent 1d399ccc9b
commit 57dc7d743e
2 changed files with 145 additions and 28 deletions
+19 -10
View File
@@ -9,6 +9,7 @@ import numpy as np
from .data import MeasurementPlane, ReconstructionResult, validate_planes from .data import MeasurementPlane, ReconstructionResult, validate_planes
from .deconvolution import DiffusionDeconvolver from .deconvolution import DiffusionDeconvolver
from .fitting import ModalFitter, generate_mode_shells from .fitting import ModalFitter, generate_mode_shells
from .geometry import CameraModel, CameraModelTolerance, GeometryCalibration
from .modes import LGBasis from .modes import LGBasis
from .noise import NoiseEstimator from .noise import NoiseEstimator
from .phase_retrieval import PhaseRetriever from .phase_retrieval import PhaseRetriever
@@ -25,12 +26,15 @@ class BeamReconstructor:
Parameters Parameters
---------- ----------
w0, z0, wavelength : known reference beam parameters (see `LGBasis`). 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 max_order : cap on automatic candidate-mode-set growth (see
`ModalFitter.fit_auto`), and also the mode set used to project the `ModalFitter.fit_auto`), and also the mode set used to project the
phase-retrieval fallback's recovered field onto the LG basis. phase-retrieval fallback's recovered field onto the LG basis.
noise_estimator : shared noise model; defaults to `NoiseEstimator()`. noise_estimator : shared noise model; defaults to `NoiseEstimator()`.
deconvolver : if given, each plane's flux is deblurred (its deconvolver : if given, each plane's flux is deblurred (using
`pixel_scale` must be known) before fitting. `GeometryCalibration(camera).effective_pixel_scale`) before fitting.
force_phase_retrieval : if True, always run the phase-retrieval fallback force_phase_retrieval : if True, always run the phase-retrieval fallback
instead of the modal fit. instead of the modal fit.
phase_retrieval_residual_threshold : if set (and `force_phase_retrieval` phase_retrieval_residual_threshold : if set (and `force_phase_retrieval`
@@ -43,6 +47,8 @@ class BeamReconstructor:
w0: float, w0: float,
z0: float, z0: float,
wavelength: float, wavelength: float,
camera: CameraModel,
camera_tolerance: CameraModelTolerance,
max_order: int = 4, max_order: int = 4,
noise_estimator: NoiseEstimator | None = None, noise_estimator: NoiseEstimator | None = None,
deconvolver: DiffusionDeconvolver | None = None, deconvolver: DiffusionDeconvolver | None = None,
@@ -51,6 +57,8 @@ class BeamReconstructor:
): ):
self.basis = LGBasis(w0=w0, z0=z0, wavelength=wavelength) self.basis = LGBasis(w0=w0, z0=z0, wavelength=wavelength)
self.wavelength = wavelength self.wavelength = wavelength
self.camera = camera
self.camera_tolerance = camera_tolerance
self.max_order = max_order self.max_order = max_order
self.noise_estimator = noise_estimator or NoiseEstimator() self.noise_estimator = noise_estimator or NoiseEstimator()
self.deconvolver = deconvolver self.deconvolver = deconvolver
@@ -63,7 +71,9 @@ class BeamReconstructor:
planes = self._deconvolve(planes) planes = self._deconvolve(planes)
fitter = ModalFitter(self.basis, self.noise_estimator) 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): if self.force_phase_retrieval or self._residual_too_high(result, planes):
result = self._phase_retrieval_fallback(planes) result = self._phase_retrieval_fallback(planes)
@@ -73,13 +83,11 @@ class BeamReconstructor:
def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]: def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]:
if self.deconvolver is None: if self.deconvolver is None:
return planes return planes
calib = GeometryCalibration(self.camera)
deblurred = [] deblurred = []
for plane in planes: for plane in planes:
if plane.pixel_scale is None: pixel_scale = calib.effective_pixel_scale(plane.flux.shape, plane.z)
raise ValueError( flux = self.deconvolver.deconvolve(plane.flux, pixel_scale)
"Deconvolution requires a known pixel_scale on every MeasurementPlane."
)
flux = self.deconvolver.deconvolve(plane.flux, plane.pixel_scale)
deblurred.append(replace(plane, flux=flux)) deblurred.append(replace(plane, flux=flux))
return deblurred return deblurred
@@ -101,7 +109,7 @@ class BeamReconstructor:
self, planes: list[MeasurementPlane] self, planes: list[MeasurementPlane]
) -> ReconstructionResult: ) -> ReconstructionResult:
retriever = PhaseRetriever(self.wavelength) 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] 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]) dx = float(pr_result.x[0, 1] - pr_result.x[0, 0])
@@ -116,7 +124,8 @@ class BeamReconstructor:
purity=purity, purity=purity,
reconstructed_field=pr_result.field, reconstructed_field=pr_result.field,
centers=[pr_result.center for _ in planes], 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={}, geometry={},
residuals=[], residuals=[],
coefficient_uncertainty={mode: float("nan") for mode in modes}, coefficient_uncertainty={mode: float("nan") for mode in modes},
+126 -18
View File
@@ -4,6 +4,7 @@ import pytest
from he11lib.deconvolution import DiffusionDeconvolver from he11lib.deconvolution import DiffusionDeconvolver
from he11lib.fitting import ModalFitter from he11lib.fitting import ModalFitter
from he11lib.geometry import CameraModel, CameraModelTolerance, GeometryCalibration
from he11lib.modes import LGBasis from he11lib.modes import LGBasis
from he11lib.reconstruct import BeamReconstructor from he11lib.reconstruct import BeamReconstructor
from he11lib.synthetic import SyntheticBeamGenerator from he11lib.synthetic import SyntheticBeamGenerator
@@ -12,6 +13,7 @@ W0 = 5e-3
Z0 = 0.5 Z0 = 0.5
WAVELENGTH = 1.76e-3 WAVELENGTH = 1.76e-3
PIXEL_SCALE = 4e-4 PIXEL_SCALE = 4e-4
CAMERA_DISTANCE = 5.0
IMAGE_SHAPE = (61, 61) IMAGE_SHAPE = (61, 61)
Z_LIST = [0.35, 0.5, 0.65, 0.8] 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) return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
def make_generator(basis): def make_camera():
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=(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(): def test_reconstruct_recovers_pure_mode_purity():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0) 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) result = reconstructor.reconstruct(planes)
power_fraction, _ = result.purity[(0, 0)] power_fraction, _ = result.purity[(0, 0)]
@@ -39,13 +61,21 @@ def test_reconstruct_recovers_pure_mode_purity():
def test_reconstruct_recovers_center_offset(): def test_reconstruct_recovers_center_offset():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
gen = make_generator(basis, camera)
true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE) true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
planes = gen.generate( 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) result = reconstructor.reconstruct(planes)
for cx, cy in result.centers: for cx, cy in result.centers:
@@ -55,21 +85,34 @@ def test_reconstruct_recovers_center_offset():
def test_reconstruct_with_deconvolution_corrects_blur(): def test_reconstruct_with_deconvolution_corrects_blur():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=2) 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) deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=30.0)
calib = GeometryCalibration(camera)
blurred_planes = [ 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. # Without deconvolution, blur should measurably hurt purity recovery.
fitter = ModalFitter(basis) 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)] purity_no_deconv, _ = result_no_deconv.purity[(0, 0)]
reconstructor = BeamReconstructor( 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) result = reconstructor.reconstruct(blurred_planes)
purity_with_deconv, _ = result.purity[(0, 0)] 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(): def test_reconstruct_forces_phase_retrieval_fallback():
basis = make_basis() 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] 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( 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) result = reconstructor.reconstruct(planes)
@@ -96,17 +148,73 @@ def test_reconstruct_forces_phase_retrieval_fallback():
def test_reconstruct_falls_back_automatically_on_high_residual(): def test_reconstruct_falls_back_automatically_on_high_residual():
basis = make_basis() 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] 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( reconstructor = BeamReconstructor(
w0=W0, w0=W0,
z0=Z0, z0=Z0,
wavelength=WAVELENGTH, wavelength=WAVELENGTH,
camera=camera,
camera_tolerance=zero_tolerance(),
max_order=2, max_order=2,
phase_retrieval_residual_threshold=1e-8, phase_retrieval_residual_threshold=1e-8,
) )
result = reconstructor.reconstruct(planes) result = reconstructor.reconstruct(planes)
assert result.used_phase_retrieval is True 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)