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>
135 lines
4.3 KiB
Python
135 lines
4.3 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from he11lib.data import validate_planes
|
|
from he11lib.fitting import ModalFitter, generate_mode_shells
|
|
from he11lib.modes import LGBasis
|
|
from he11lib.synthetic import SyntheticBeamGenerator
|
|
|
|
W0 = 5e-3
|
|
Z0 = 0.5
|
|
WAVELENGTH = 1.76e-3
|
|
PIXEL_SCALE = 4e-4
|
|
IMAGE_SHAPE = (61, 61)
|
|
Z_LIST = [0.35, 0.5, 0.65, 0.8]
|
|
|
|
|
|
def make_basis():
|
|
return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
|
|
|
|
|
|
def make_generator(basis):
|
|
return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE)
|
|
|
|
|
|
def test_generate_mode_shells_orders_by_2p_plus_abs_l():
|
|
shells = generate_mode_shells(max_order=2)
|
|
assert shells[0] == [(0, 0)]
|
|
assert set(shells[1]) == {(0, 1), (0, -1)}
|
|
assert set(shells[2]) == {(0, 2), (0, -2), (1, 0)}
|
|
|
|
|
|
def test_fit_recovers_pure_fundamental_mode():
|
|
basis = make_basis()
|
|
gen = make_generator(basis)
|
|
planes = gen.generate(
|
|
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0
|
|
)
|
|
|
|
fitter = ModalFitter(basis)
|
|
result = fitter.fit(planes, modes=[(0, 0)])
|
|
|
|
power_fraction, _ = result.purity[(0, 0)]
|
|
assert power_fraction == pytest.approx(1.0, abs=1e-6)
|
|
for cx, cy in result.centers:
|
|
assert cx == pytest.approx(0.0, abs=2 * PIXEL_SCALE)
|
|
assert cy == pytest.approx(0.0, abs=2 * PIXEL_SCALE)
|
|
|
|
|
|
def test_fit_recovers_two_mode_purity_ratio():
|
|
basis = make_basis()
|
|
gen = make_generator(basis)
|
|
true_coeffs = {(0, 0): 0.9 + 0j, (1, 0): 0.3 + 0.1j}
|
|
planes = gen.generate(coefficients=true_coeffs, z_list=Z_LIST, noise_std=1e-4, seed=1)
|
|
|
|
fitter = ModalFitter(basis)
|
|
result = fitter.fit(planes, modes=list(true_coeffs.keys()))
|
|
|
|
true_total = sum(abs(c) ** 2 for c in true_coeffs.values())
|
|
for mode, c in true_coeffs.items():
|
|
expected_fraction = abs(c) ** 2 / true_total
|
|
recovered_fraction, _ = result.purity[mode]
|
|
assert recovered_fraction == pytest.approx(expected_fraction, abs=0.03)
|
|
|
|
|
|
def test_fit_recovers_center_offset():
|
|
basis = make_basis()
|
|
gen = make_generator(basis)
|
|
true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
|
|
planes = gen.generate(
|
|
coefficients={(0, 0): 1.0 + 0j},
|
|
z_list=Z_LIST,
|
|
center=true_center,
|
|
noise_std=1e-4,
|
|
seed=2,
|
|
)
|
|
|
|
fitter = ModalFitter(basis)
|
|
result = fitter.fit(planes, modes=[(0, 0)], initial_center=true_center)
|
|
|
|
for cx, cy in result.centers:
|
|
assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE)
|
|
assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE)
|
|
|
|
|
|
def test_fit_recovers_unknown_pixel_scale():
|
|
# Use a coarser pixel scale so the (much wider, far-field) beam at the
|
|
# outer z distances still fits within the frame -- otherwise pixel scale
|
|
# becomes unobservable from clipped images.
|
|
basis = make_basis()
|
|
local_pixel_scale = 1.5e-3
|
|
gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=local_pixel_scale)
|
|
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=3)
|
|
|
|
# hide the known calibration to force the fitter to solve for it
|
|
for plane in planes:
|
|
plane.pixel_scale = None
|
|
plane.viewing_angle_deg = None
|
|
|
|
fitter = ModalFitter(basis)
|
|
result = fitter.fit(
|
|
planes,
|
|
modes=[(0, 0)],
|
|
initial_pixel_scale=local_pixel_scale * 1.1,
|
|
initial_viewing_angle_deg=0.0,
|
|
)
|
|
|
|
fitted_scales = [result.geometry[f"pixel_scale_{i}"] for i in range(len(planes))]
|
|
for scale in fitted_scales:
|
|
assert scale == pytest.approx(local_pixel_scale, rel=0.05)
|
|
|
|
|
|
def test_fit_auto_does_not_add_modes_for_pure_fundamental():
|
|
basis = make_basis()
|
|
gen = make_generator(basis)
|
|
planes = gen.generate(
|
|
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=4
|
|
)
|
|
|
|
fitter = ModalFitter(basis)
|
|
result = fitter.fit_auto(planes, max_order=2)
|
|
|
|
assert set(result.purity.keys()) == {(0, 0)}
|
|
|
|
|
|
def test_fit_auto_grows_to_include_second_mode():
|
|
basis = make_basis()
|
|
gen = make_generator(basis)
|
|
true_coeffs = {(0, 0): 0.9 + 0j, (0, 1): 0.4 + 0j}
|
|
planes = gen.generate(coefficients=true_coeffs, z_list=Z_LIST, noise_std=1e-4, seed=5)
|
|
|
|
fitter = ModalFitter(basis)
|
|
result = fitter.fit_auto(planes, max_order=2)
|
|
|
|
assert (0, 1) in result.purity or (0, -1) in result.purity
|