From 1d399ccc9b9544055607d0af449bfdf6f77d72e5 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Fri, 3 Jul 2026 12:19:34 +0200 Subject: [PATCH] Thread CameraModel through PhaseRetriever.retrieve retrieve() now takes a CameraModel directly instead of the removed pixel_scale/viewing_angle_deg override kwargs, matching the rest of the pipeline's shared-camera convention. Co-Authored-By: Claude Sonnet 5 --- he11lib/phase_retrieval.py | 17 ++++++++--------- tests/test_phase_retrieval.py | 36 +++++++++++++++++++++++++++-------- 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/he11lib/phase_retrieval.py b/he11lib/phase_retrieval.py index daa9c87..7fd4e16 100644 --- a/he11lib/phase_retrieval.py +++ b/he11lib/phase_retrieval.py @@ -15,7 +15,7 @@ from dataclasses import dataclass import numpy as np from .data import MeasurementPlane, validate_planes -from .geometry import GeometryCalibration +from .geometry import CameraModel, GeometryCalibration def propagate_angular_spectrum( @@ -79,22 +79,21 @@ class PhaseRetriever: def retrieve( self, planes: list[MeasurementPlane], - pixel_scale: float | None = None, - viewing_angle_deg: float | None = None, + camera: CameraModel, max_iterations: int = 200, ) -> PhaseRetrievalResult: """Run Gerchberg-Saxton phase retrieval across the given planes. - Planes must share the same known (or overridden) pixel_scale and - viewing_angle_deg, since all planes are propagated on one common - physical grid. + All planes are propagated on one common physical grid, derived + from `camera` at the smallest-z plane's depth (an existing + approximation: the shared grid is only exact at that one z, since + other planes may sit at a slightly different true depth under true + perspective projection). """ validate_planes(planes) ordered = sorted(planes, key=lambda p: p.z) - x, y = GeometryCalibration(ordered[0]).physical_coordinates( - pixel_scale=pixel_scale, viewing_angle_deg=viewing_angle_deg - ) + x, y = GeometryCalibration(camera).physical_coordinates(ordered[0].flux.shape, ordered[0].z) dx = float(x[0, 1] - x[0, 0]) amplitudes = [np.sqrt(np.clip(p.flux, 0, None)) for p in ordered] diff --git a/tests/test_phase_retrieval.py b/tests/test_phase_retrieval.py index ad66a5d..5808be7 100644 --- a/tests/test_phase_retrieval.py +++ b/tests/test_phase_retrieval.py @@ -1,6 +1,7 @@ import numpy as np import pytest +from he11lib.geometry import CameraModel from he11lib.modes import LGBasis from he11lib.phase_retrieval import PhaseRetriever, propagate_angular_spectrum from he11lib.synthetic import SyntheticBeamGenerator @@ -9,6 +10,7 @@ W0 = 5e-3 Z0 = 0.5 WAVELENGTH = 1.76e-3 PIXEL_SCALE = 3e-4 +CAMERA_DISTANCE = 5.0 IMAGE_SHAPE = (121, 121) @@ -16,6 +18,15 @@ def make_basis(): return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) +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 make_grid(): coords = (np.arange(IMAGE_SHAPE[0]) - IMAGE_SHAPE[0] // 2) * PIXEL_SCALE x, y = np.meshgrid(coords, coords) @@ -42,7 +53,6 @@ def test_propagate_matches_lgbasis_analytic_evolution(): propagated = propagate_angular_spectrum(field_at_waist, PIXEL_SCALE, dz=dz, wavelength=WAVELENGTH) analytic = basis.field(x, y, Z0 + dz, p=0, l=0) - # compare intensity profiles (phase reference/global constant may differ) np.testing.assert_allclose( np.abs(propagated) ** 2, np.abs(analytic) ** 2, atol=1e-2 * np.max(np.abs(analytic) ** 2) ) @@ -53,15 +63,19 @@ def test_retrieve_recovers_pure_mode_purity(): # within the frame -- otherwise FFT wraparound/clipping at the edges # degrades angular-spectrum propagation accuracy. basis = make_basis() - gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) + 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=0) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE, noise_std=1e-5, seed=0 + ) retriever = PhaseRetriever(wavelength=WAVELENGTH) - result = retriever.retrieve(planes, viewing_angle_deg=0.0, max_iterations=100) + result = retriever.retrieve(planes, camera, max_iterations=100) + dx = float(result.x[0, 1] - result.x[0, 0]) coeffs = basis.project( - result.field, result.x, result.y, PIXEL_SCALE, result.z, modes=[(0, 0), (1, 0), (0, 1)] + result.field, result.x, result.y, dx, result.z, modes=[(0, 0), (1, 0), (0, 1)] ) total_power = sum(abs(c) ** 2 for c in coeffs.values()) purity_00 = abs(coeffs[(0, 0)]) ** 2 / total_power @@ -70,15 +84,21 @@ def test_retrieve_recovers_pure_mode_purity(): def test_retrieve_estimates_beam_center(): basis = make_basis() - gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) + camera = make_camera() + gen = SyntheticBeamGenerator(basis=basis, camera=camera) z_list = [0.47, 0.5, 0.53] true_center = (15 * PIXEL_SCALE, -8 * PIXEL_SCALE) planes = gen.generate( - coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, center=true_center, noise_std=1e-5, seed=1 + coefficients={(0, 0): 1.0 + 0j}, + z_list=z_list, + image_shape=IMAGE_SHAPE, + center=true_center, + noise_std=1e-5, + seed=1, ) retriever = PhaseRetriever(wavelength=WAVELENGTH) - result = retriever.retrieve(planes, viewing_angle_deg=0.0, max_iterations=100) + result = retriever.retrieve(planes, camera, max_iterations=100) assert result.center[0] == pytest.approx(true_center[0], abs=3 * PIXEL_SCALE) assert result.center[1] == pytest.approx(true_center[1], abs=3 * PIXEL_SCALE)