Martino Ferrari
03b63ba03a
Full implementation of Laguerre-Gauss modal reconstruction for gyrotron beam diagnostics, per the approved design spec, plus tests, docs, and a runnable end-to-end example. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from he11lib.modes import LGBasis
|
|
from he11lib.synthetic import SyntheticBeamGenerator
|
|
|
|
|
|
W0 = 5e-3
|
|
Z0 = 0.5
|
|
WAVELENGTH = 1.76e-3
|
|
PIXEL_SCALE = 2e-4 # 0.2 mm/px
|
|
IMAGE_SHAPE = (161, 161) # odd so there's a well-defined center pixel
|
|
|
|
|
|
def make_generator():
|
|
basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
|
|
return SyntheticBeamGenerator(
|
|
basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
assert [p.z for p in planes] == z_list
|
|
assert all(p.flux.shape == IMAGE_SHAPE for p in planes)
|
|
|
|
|
|
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))
|
|
flux = planes[0].flux
|
|
|
|
peak_idx = np.unravel_index(np.argmax(flux), flux.shape)
|
|
center_idx = (IMAGE_SHAPE[0] // 2, IMAGE_SHAPE[1] // 2)
|
|
assert peak_idx == center_idx
|
|
|
|
|
|
def test_generate_applies_center_offset():
|
|
gen = make_generator()
|
|
offset_m = 20 * PIXEL_SCALE # 20 pixels
|
|
planes = gen.generate(
|
|
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(offset_m, 0.0)
|
|
)
|
|
flux = planes[0].flux
|
|
|
|
peak_idx = np.unravel_index(np.argmax(flux), flux.shape)
|
|
center_row = IMAGE_SHAPE[0] // 2
|
|
center_col = IMAGE_SHAPE[1] // 2
|
|
assert peak_idx[0] == center_row
|
|
assert peak_idx[1] == pytest.approx(center_col + 20, abs=1)
|
|
|
|
|
|
def test_generate_applies_pointing_angle_as_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,
|
|
center=(0.0, 0.0),
|
|
pointing_angle_deg=pointing_angle_deg,
|
|
)
|
|
|
|
peaks_col = []
|
|
for plane in planes:
|
|
peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape)
|
|
peaks_col.append(peak_idx[1])
|
|
|
|
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)
|
|
|
|
|
|
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
|
|
)
|
|
planes_b = gen.generate(
|
|
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42
|
|
)
|
|
np.testing.assert_array_equal(planes_a[0].flux, planes_b[0].flux)
|
|
|
|
|
|
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
|
|
)
|
|
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():
|
|
gen = make_generator()
|
|
planes_straight = gen.generate(
|
|
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=0.0
|
|
)
|
|
planes_tilted = gen.generate(
|
|
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=60.0
|
|
)
|
|
|
|
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)
|