Rewrite SyntheticBeamGenerator around CameraModel and 2D beam pointing

Renders each plane via true pinhole projection through a shared
CameraModel instead of the old cosine-compression formula, adds
independent horizontal/vertical pointing drift, and supports generating
a deliberately-offset nominal z (vs. true z) per plane for tolerance-
recovery testing in later tasks.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Martino Ferrari
2026-07-03 11:43:20 +02:00
parent dffca62f81
commit c2ae95e01f
2 changed files with 93 additions and 74 deletions
+30 -34
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import numpy as np import numpy as np
from .data import MeasurementPlane from .data import MeasurementPlane
from .geometry import CameraModel, GeometryCalibration
from .modes import LGBasis from .modes import LGBasis
@@ -19,65 +20,60 @@ class SyntheticBeamGenerator:
Parameters Parameters
---------- ----------
basis : LGBasis defining the reference w0, z0, wavelength. basis : LGBasis defining the reference w0, z0, wavelength.
image_shape : (rows, cols) pixel shape of generated images. camera : ground-truth CameraModel (position/orientation/intrinsics) used
pixel_scale : physical size of one pixel, in meters, along the to render each plane via true perspective projection.
non-tilted (y) axis. The tilt/projection axis is assumed to be x.
""" """
def __init__(self, basis: LGBasis, image_shape: tuple[int, int], pixel_scale: float): def __init__(self, basis: LGBasis, camera: CameraModel):
self.basis = basis self.basis = basis
self.image_shape = image_shape self.camera = camera
self.pixel_scale = pixel_scale self.calibration = GeometryCalibration(camera)
def _pixel_grid(self, center: tuple[float, float], viewing_angle_deg: float):
rows, cols = self.image_shape
row_idx = np.arange(rows) - rows // 2
col_idx = np.arange(cols) - cols // 2
col_grid, row_grid = np.meshgrid(col_idx, row_idx)
cos_angle = np.cos(np.deg2rad(viewing_angle_deg))
x = col_grid * self.pixel_scale / cos_angle - center[0]
y = row_grid * self.pixel_scale - center[1]
return x, y
def generate( def generate(
self, self,
coefficients: dict[tuple[int, int], complex], coefficients: dict[tuple[int, int], complex],
z_list: list[float], z_list: list[float],
image_shape: tuple[int, int],
*, *,
center: tuple[float, float] = (0.0, 0.0), center: tuple[float, float] = (0.0, 0.0),
pointing_angle_deg: float = 0.0, pointing_angle_horizontal_deg: float = 0.0,
viewing_angle_deg: float = 0.0, pointing_angle_vertical_deg: float = 0.0,
z_tolerance: float = 0.0,
nominal_z_offsets: dict[float, float] | None = None,
noise_std: float = 0.0, noise_std: float = 0.0,
seed: int | None = None, seed: int | None = None,
) -> list[MeasurementPlane]: ) -> list[MeasurementPlane]:
"""Generate one MeasurementPlane per requested z distance. """Generate one MeasurementPlane per requested (true) z distance.
The beam transverse center drifts linearly with z according to The beam transverse center drifts linearly with z according to the
pointing_angle_deg (tilt of the beam axis along x), starting from two pointing angles, starting from `center` at the basis's
`center` at the basis's reference z0. reference z0. `nominal_z_offsets`, if given, maps a true z (as
given in z_list) to an offset applied to the *nominal* z stored on
the resulting MeasurementPlane -- letting tests verify a fit
recovers the true z despite a deliberately-offset nominal input.
Every resulting plane shares `z_tolerance`.
""" """
rng = np.random.default_rng(seed) rng = np.random.default_rng(seed)
tilt_rad = np.deg2rad(pointing_angle_deg) tilt_h_rad = np.deg2rad(pointing_angle_horizontal_deg)
tilt_v_rad = np.deg2rad(pointing_angle_vertical_deg)
offsets = nominal_z_offsets or {}
planes = [] planes = []
for z in z_list: for z in z_list:
drift_x = (z - self.basis.z0) * np.tan(tilt_rad) drift_x = (z - self.basis.z0) * np.tan(tilt_h_rad)
plane_center = (center[0] + drift_x, center[1]) drift_y = (z - self.basis.z0) * np.tan(tilt_v_rad)
cx = center[0] + drift_x
cy = center[1] + drift_y
x, y = self._pixel_grid(plane_center, viewing_angle_deg) x, y = self.calibration.physical_coordinates(image_shape, z)
field = self.basis.field_superposition(x, y, z, coefficients) field = self.basis.field_superposition(x - cx, y - cy, z, coefficients)
flux = np.abs(field) ** 2 flux = np.abs(field) ** 2
if noise_std > 0: if noise_std > 0:
flux = flux + rng.normal(0.0, noise_std, size=flux.shape) flux = flux + rng.normal(0.0, noise_std, size=flux.shape)
nominal_z = z + offsets.get(z, 0.0)
planes.append( planes.append(
MeasurementPlane( MeasurementPlane(flux=flux, z=nominal_z, z_tolerance=z_tolerance)
flux=flux,
z=z,
pixel_scale=self.pixel_scale,
viewing_angle_deg=viewing_angle_deg,
)
) )
return planes return planes
+63 -40
View File
@@ -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.synthetic import SyntheticBeamGenerator from he11lib.synthetic import SyntheticBeamGenerator
@@ -8,21 +9,29 @@ from he11lib.synthetic import SyntheticBeamGenerator
W0 = 5e-3 W0 = 5e-3
Z0 = 0.5 Z0 = 0.5
WAVELENGTH = 1.76e-3 WAVELENGTH = 1.76e-3
PIXEL_SCALE = 2e-4 # 0.2 mm/px PIXEL_SCALE = 2e-4 # 0.2 mm/px, achieved at z=Z0
CAMERA_DISTANCE = 5.0 # camera stands 5 m upstream of the output window
IMAGE_SHAPE = (161, 161) # odd so there's a well-defined center pixel IMAGE_SHAPE = (161, 161) # odd so there's a well-defined center pixel
def make_camera(pixel_scale=PIXEL_SCALE, z0=Z0, camera_distance=CAMERA_DISTANCE):
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_generator(): def make_generator():
basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
return SyntheticBeamGenerator( return SyntheticBeamGenerator(basis=basis, camera=make_camera())
basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE
)
def test_generate_returns_planes_with_requested_z(): def test_generate_returns_planes_with_requested_z():
gen = make_generator() gen = make_generator()
z_list = [0.3, 0.4, 0.5] z_list = [0.3, 0.4, 0.5]
planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=z_list) planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE)
assert [p.z for p in planes] == z_list assert [p.z for p in planes] == z_list
assert all(p.flux.shape == IMAGE_SHAPE for p in planes) assert all(p.flux.shape == IMAGE_SHAPE for p in planes)
@@ -30,7 +39,9 @@ def test_generate_returns_planes_with_requested_z():
def test_generate_pure_mode_peak_at_image_center_when_centered(): def test_generate_pure_mode_peak_at_image_center_when_centered():
gen = make_generator() gen = make_generator()
planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(0.0, 0.0)) planes = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, center=(0.0, 0.0)
)
flux = planes[0].flux flux = planes[0].flux
peak_idx = np.unravel_index(np.argmax(flux), flux.shape) peak_idx = np.unravel_index(np.argmax(flux), flux.shape)
@@ -40,9 +51,9 @@ def test_generate_pure_mode_peak_at_image_center_when_centered():
def test_generate_applies_center_offset(): def test_generate_applies_center_offset():
gen = make_generator() gen = make_generator()
offset_m = 20 * PIXEL_SCALE # 20 pixels offset_m = 20 * PIXEL_SCALE # ~20 pixels at z=Z0
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(offset_m, 0.0) coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, center=(offset_m, 0.0)
) )
flux = planes[0].flux flux = planes[0].flux
@@ -53,35 +64,41 @@ def test_generate_applies_center_offset():
assert peak_idx[1] == pytest.approx(center_col + 20, abs=1) assert peak_idx[1] == pytest.approx(center_col + 20, abs=1)
def test_generate_applies_pointing_angle_as_linear_drift(): def test_generate_applies_pointing_angles_as_2d_linear_drift():
gen = make_generator() gen = make_generator()
pointing_angle_deg = 1.0 # small tilt
z_list = [Z0, Z0 + 0.2] z_list = [Z0, Z0 + 0.2]
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1 + 0j}, coefficients={(0, 0): 1 + 0j},
z_list=z_list, z_list=z_list,
image_shape=IMAGE_SHAPE,
center=(0.0, 0.0), center=(0.0, 0.0),
pointing_angle_deg=pointing_angle_deg, pointing_angle_horizontal_deg=1.0,
pointing_angle_vertical_deg=0.5,
) )
peaks_col = [] peaks = []
for plane in planes: for plane in planes:
peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape) peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape)
peaks_col.append(peak_idx[1]) peaks.append(peak_idx)
expected_shift_m = 0.2 * np.tan(np.deg2rad(pointing_angle_deg)) expected_shift_x_m = 0.2 * np.tan(np.deg2rad(1.0))
expected_shift_px = expected_shift_m / PIXEL_SCALE expected_shift_y_m = 0.2 * np.tan(np.deg2rad(0.5))
actual_shift_px = peaks_col[1] - peaks_col[0] expected_shift_col_px = expected_shift_x_m / PIXEL_SCALE
assert actual_shift_px == pytest.approx(expected_shift_px, abs=1) expected_shift_row_px = expected_shift_y_m / PIXEL_SCALE
actual_shift_col_px = peaks[1][1] - peaks[0][1]
actual_shift_row_px = peaks[1][0] - peaks[0][0]
assert actual_shift_col_px == pytest.approx(expected_shift_col_px, abs=1)
assert actual_shift_row_px == pytest.approx(expected_shift_row_px, abs=1)
def test_generate_noise_is_reproducible_with_seed(): def test_generate_noise_is_reproducible_with_seed():
gen = make_generator() gen = make_generator()
planes_a = gen.generate( planes_a = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42 coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.01, seed=42
) )
planes_b = gen.generate( planes_b = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42 coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.01, seed=42
) )
np.testing.assert_array_equal(planes_a[0].flux, planes_b[0].flux) np.testing.assert_array_equal(planes_a[0].flux, planes_b[0].flux)
@@ -90,34 +107,40 @@ def test_generate_noise_std_matches_requested_level():
gen = make_generator() gen = make_generator()
noise_std = 0.02 noise_std = 0.02
planes_noisy = gen.generate( planes_noisy = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=noise_std, seed=1 coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=noise_std, seed=1
)
planes_clean = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.0
) )
planes_clean = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.0)
diff = planes_noisy[0].flux - planes_clean[0].flux diff = planes_noisy[0].flux - planes_clean[0].flux
assert np.std(diff) == pytest.approx(noise_std, rel=0.15) assert np.std(diff) == pytest.approx(noise_std, rel=0.15)
def test_generate_viewing_angle_compresses_tilt_axis(): def test_generate_applies_z_tolerance_to_every_plane():
gen = make_generator() gen = make_generator()
planes_straight = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=0.0 coefficients={(0, 0): 1 + 0j},
z_list=[0.3, 0.4, 0.5],
image_shape=IMAGE_SHAPE,
z_tolerance=0.02,
) )
planes_tilted = gen.generate( assert all(p.z_tolerance == 0.02 for p in planes)
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=60.0
def test_generate_applies_nominal_z_offset_independent_of_true_z():
gen = make_generator()
true_z_list = [0.3, 0.4, 0.5]
offsets = {0.3: 0.01, 0.4: -0.005, 0.5: 0.0}
planes = gen.generate(
coefficients={(0, 0): 1 + 0j},
z_list=true_z_list,
image_shape=IMAGE_SHAPE,
nominal_z_offsets=offsets,
) )
def width_along_axis(flux, axis): nominal_zs = [p.z for p in planes]
profile = flux[flux.shape[0] // 2, :] if axis == 1 else flux[:, flux.shape[1] // 2] assert nominal_zs == pytest.approx([0.31, 0.395, 0.5])
half_max = profile.max() / 2 # The flux is still rendered at each plane's *true* z (0.3, 0.4, 0.5),
above = np.where(profile >= half_max)[0] # not its offset nominal z -- verified indirectly in Task 7's
return above[-1] - above[0] # end-to-end tolerance-recovery test.
width_straight_x = width_along_axis(planes_straight[0].flux, axis=1)
width_tilted_x = width_along_axis(planes_tilted[0].flux, axis=1)
width_straight_y = width_along_axis(planes_straight[0].flux, axis=0)
width_tilted_y = width_along_axis(planes_tilted[0].flux, axis=0)
# tilt compresses the viewed beam along the tilt (x) axis, y unaffected
assert width_tilted_x < width_straight_x
assert width_tilted_y == pytest.approx(width_straight_y, abs=1)