diff --git a/he11lib/synthetic.py b/he11lib/synthetic.py index 215b34d..b9f2baf 100644 --- a/he11lib/synthetic.py +++ b/he11lib/synthetic.py @@ -10,6 +10,7 @@ from __future__ import annotations import numpy as np from .data import MeasurementPlane +from .geometry import CameraModel, GeometryCalibration from .modes import LGBasis @@ -19,65 +20,60 @@ class SyntheticBeamGenerator: Parameters ---------- basis : LGBasis defining the reference w0, z0, wavelength. - image_shape : (rows, cols) pixel shape of generated images. - pixel_scale : physical size of one pixel, in meters, along the - non-tilted (y) axis. The tilt/projection axis is assumed to be x. + camera : ground-truth CameraModel (position/orientation/intrinsics) used + to render each plane via true perspective projection. """ - def __init__(self, basis: LGBasis, image_shape: tuple[int, int], pixel_scale: float): + def __init__(self, basis: LGBasis, camera: CameraModel): self.basis = basis - self.image_shape = image_shape - self.pixel_scale = pixel_scale - - 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 + self.camera = camera + self.calibration = GeometryCalibration(camera) def generate( self, coefficients: dict[tuple[int, int], complex], z_list: list[float], + image_shape: tuple[int, int], *, center: tuple[float, float] = (0.0, 0.0), - pointing_angle_deg: float = 0.0, - viewing_angle_deg: float = 0.0, + pointing_angle_horizontal_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, seed: int | None = None, ) -> 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 - pointing_angle_deg (tilt of the beam axis along x), starting from - `center` at the basis's reference z0. + The beam transverse center drifts linearly with z according to the + two pointing angles, starting from `center` at the basis's + 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) - 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 = [] for z in z_list: - drift_x = (z - self.basis.z0) * np.tan(tilt_rad) - plane_center = (center[0] + drift_x, center[1]) + drift_x = (z - self.basis.z0) * np.tan(tilt_h_rad) + 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) - field = self.basis.field_superposition(x, y, z, coefficients) + x, y = self.calibration.physical_coordinates(image_shape, z) + field = self.basis.field_superposition(x - cx, y - cy, z, coefficients) flux = np.abs(field) ** 2 if noise_std > 0: flux = flux + rng.normal(0.0, noise_std, size=flux.shape) + nominal_z = z + offsets.get(z, 0.0) planes.append( - MeasurementPlane( - flux=flux, - z=z, - pixel_scale=self.pixel_scale, - viewing_angle_deg=viewing_angle_deg, - ) + MeasurementPlane(flux=flux, z=nominal_z, z_tolerance=z_tolerance) ) return planes diff --git a/tests/test_synthetic.py b/tests/test_synthetic.py index 90e0e9d..00d69d2 100644 --- a/tests/test_synthetic.py +++ b/tests/test_synthetic.py @@ -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.