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:
+126
-18
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user