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 <noreply@anthropic.com>
This commit is contained in:
@@ -15,7 +15,7 @@ from dataclasses import dataclass
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
from .data import MeasurementPlane, validate_planes
|
from .data import MeasurementPlane, validate_planes
|
||||||
from .geometry import GeometryCalibration
|
from .geometry import CameraModel, GeometryCalibration
|
||||||
|
|
||||||
|
|
||||||
def propagate_angular_spectrum(
|
def propagate_angular_spectrum(
|
||||||
@@ -79,22 +79,21 @@ class PhaseRetriever:
|
|||||||
def retrieve(
|
def retrieve(
|
||||||
self,
|
self,
|
||||||
planes: list[MeasurementPlane],
|
planes: list[MeasurementPlane],
|
||||||
pixel_scale: float | None = None,
|
camera: CameraModel,
|
||||||
viewing_angle_deg: float | None = None,
|
|
||||||
max_iterations: int = 200,
|
max_iterations: int = 200,
|
||||||
) -> PhaseRetrievalResult:
|
) -> PhaseRetrievalResult:
|
||||||
"""Run Gerchberg-Saxton phase retrieval across the given planes.
|
"""Run Gerchberg-Saxton phase retrieval across the given planes.
|
||||||
|
|
||||||
Planes must share the same known (or overridden) pixel_scale and
|
All planes are propagated on one common physical grid, derived
|
||||||
viewing_angle_deg, since all planes are propagated on one common
|
from `camera` at the smallest-z plane's depth (an existing
|
||||||
physical grid.
|
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)
|
validate_planes(planes)
|
||||||
ordered = sorted(planes, key=lambda p: p.z)
|
ordered = sorted(planes, key=lambda p: p.z)
|
||||||
|
|
||||||
x, y = GeometryCalibration(ordered[0]).physical_coordinates(
|
x, y = GeometryCalibration(camera).physical_coordinates(ordered[0].flux.shape, ordered[0].z)
|
||||||
pixel_scale=pixel_scale, viewing_angle_deg=viewing_angle_deg
|
|
||||||
)
|
|
||||||
dx = float(x[0, 1] - x[0, 0])
|
dx = float(x[0, 1] - x[0, 0])
|
||||||
|
|
||||||
amplitudes = [np.sqrt(np.clip(p.flux, 0, None)) for p in ordered]
|
amplitudes = [np.sqrt(np.clip(p.flux, 0, None)) for p in ordered]
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
from he11lib.geometry import CameraModel
|
||||||
from he11lib.modes import LGBasis
|
from he11lib.modes import LGBasis
|
||||||
from he11lib.phase_retrieval import PhaseRetriever, propagate_angular_spectrum
|
from he11lib.phase_retrieval import PhaseRetriever, propagate_angular_spectrum
|
||||||
from he11lib.synthetic import SyntheticBeamGenerator
|
from he11lib.synthetic import SyntheticBeamGenerator
|
||||||
@@ -9,6 +10,7 @@ W0 = 5e-3
|
|||||||
Z0 = 0.5
|
Z0 = 0.5
|
||||||
WAVELENGTH = 1.76e-3
|
WAVELENGTH = 1.76e-3
|
||||||
PIXEL_SCALE = 3e-4
|
PIXEL_SCALE = 3e-4
|
||||||
|
CAMERA_DISTANCE = 5.0
|
||||||
IMAGE_SHAPE = (121, 121)
|
IMAGE_SHAPE = (121, 121)
|
||||||
|
|
||||||
|
|
||||||
@@ -16,6 +18,15 @@ def make_basis():
|
|||||||
return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
|
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():
|
def make_grid():
|
||||||
coords = (np.arange(IMAGE_SHAPE[0]) - IMAGE_SHAPE[0] // 2) * PIXEL_SCALE
|
coords = (np.arange(IMAGE_SHAPE[0]) - IMAGE_SHAPE[0] // 2) * PIXEL_SCALE
|
||||||
x, y = np.meshgrid(coords, coords)
|
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)
|
propagated = propagate_angular_spectrum(field_at_waist, PIXEL_SCALE, dz=dz, wavelength=WAVELENGTH)
|
||||||
analytic = basis.field(x, y, Z0 + dz, p=0, l=0)
|
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.testing.assert_allclose(
|
||||||
np.abs(propagated) ** 2, np.abs(analytic) ** 2, atol=1e-2 * np.max(np.abs(analytic) ** 2)
|
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
|
# within the frame -- otherwise FFT wraparound/clipping at the edges
|
||||||
# degrades angular-spectrum propagation accuracy.
|
# degrades angular-spectrum propagation accuracy.
|
||||||
basis = make_basis()
|
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]
|
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)
|
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(
|
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())
|
total_power = sum(abs(c) ** 2 for c in coeffs.values())
|
||||||
purity_00 = abs(coeffs[(0, 0)]) ** 2 / total_power
|
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():
|
def test_retrieve_estimates_beam_center():
|
||||||
basis = make_basis()
|
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]
|
z_list = [0.47, 0.5, 0.53]
|
||||||
true_center = (15 * PIXEL_SCALE, -8 * PIXEL_SCALE)
|
true_center = (15 * PIXEL_SCALE, -8 * PIXEL_SCALE)
|
||||||
planes = gen.generate(
|
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)
|
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[0] == pytest.approx(true_center[0], abs=3 * PIXEL_SCALE)
|
||||||
assert result.center[1] == pytest.approx(true_center[1], abs=3 * PIXEL_SCALE)
|
assert result.center[1] == pytest.approx(true_center[1], abs=3 * PIXEL_SCALE)
|
||||||
|
|||||||
Reference in New Issue
Block a user