martino/he11lib
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rewrite SyntheticBeamGenerator around CameraModel and 2D beam pointing

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_synthetic.py
+63 -40
View File
@@ -1,6 +1,7 @@
import numpy as np
import pytest
from he11lib.geometry import CameraModel
from he11lib.modes import LGBasis
from he11lib.synthetic import SyntheticBeamGenerator
@@ -8,21 +9,29 @@ from he11lib.synthetic import SyntheticBeamGenerator
W0 = 5e-3
Z0 = 0.5
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
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():
basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
return SyntheticBeamGenerator(
basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE
)
return SyntheticBeamGenerator(basis=basis, camera=make_camera())
def test_generate_returns_planes_with_requested_z():
gen = make_generator()
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 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():
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
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():
gen = make_generator()
offset_m = 20 * PIXEL_SCALE # 20 pixels
offset_m = 20 * PIXEL_SCALE # ~20 pixels at z=Z0
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
@@ -53,35 +64,41 @@ def test_generate_applies_center_offset():
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()
pointing_angle_deg = 1.0 # small tilt
z_list = [Z0, Z0 + 0.2]
planes = gen.generate(
coefficients={(0, 0): 1 + 0j},
z_list=z_list,
image_shape=IMAGE_SHAPE,
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:
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_px = expected_shift_m / PIXEL_SCALE
actual_shift_px = peaks_col[1] - peaks_col[0]
assert actual_shift_px == pytest.approx(expected_shift_px, abs=1)
expected_shift_x_m = 0.2 * np.tan(np.deg2rad(1.0))
expected_shift_y_m = 0.2 * np.tan(np.deg2rad(0.5))
expected_shift_col_px = expected_shift_x_m / PIXEL_SCALE
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():
gen = make_generator()
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(
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)
@@ -90,34 +107,40 @@ def test_generate_noise_std_matches_requested_level():
gen = make_generator()
noise_std = 0.02
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
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()
planes_straight = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=0.0
planes = gen.generate(
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(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=60.0
assert all(p.z_tolerance == 0.02 for p in planes)
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):
profile = flux[flux.shape[0] // 2, :] if axis == 1 else flux[:, flux.shape[1] // 2]
half_max = profile.max() / 2
above = np.where(profile >= half_max)[0]
return above[-1] - above[0]
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)
nominal_zs = [p.z for p in planes]
assert nominal_zs == pytest.approx([0.31, 0.395, 0.5])
# The flux is still rendered at each plane's *true* z (0.3, 0.4, 0.5),
# not its offset nominal z -- verified indirectly in Task 7's
# end-to-end tolerance-recovery test.