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>
113 lines
3.8 KiB
Python
113 lines
3.8 KiB
Python
from dataclasses import replace
|
|
|
|
import pytest
|
|
|
|
from he11lib.deconvolution import DiffusionDeconvolver
|
|
from he11lib.fitting import ModalFitter
|
|
from he11lib.modes import LGBasis
|
|
from he11lib.reconstruct import BeamReconstructor
|
|
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_reconstruct_recovers_pure_mode_purity():
|
|
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)
|
|
|
|
reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2)
|
|
result = reconstructor.reconstruct(planes)
|
|
|
|
power_fraction, _ = result.purity[(0, 0)]
|
|
assert power_fraction == pytest.approx(1.0, abs=1e-3)
|
|
assert result.used_phase_retrieval is False
|
|
|
|
|
|
def test_reconstruct_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=1
|
|
)
|
|
|
|
reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2)
|
|
result = reconstructor.reconstruct(planes)
|
|
|
|
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_reconstruct_with_deconvolution_corrects_blur():
|
|
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=2)
|
|
|
|
deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=30.0)
|
|
blurred_planes = [
|
|
replace(p, flux=deconvolver.blur(p.flux, p.pixel_scale)) for p in planes
|
|
]
|
|
|
|
# Without deconvolution, blur should measurably hurt purity recovery.
|
|
fitter = ModalFitter(basis)
|
|
result_no_deconv = fitter.fit(blurred_planes, modes=[(0, 0), (1, 0), (0, 1)])
|
|
purity_no_deconv, _ = result_no_deconv.purity[(0, 0)]
|
|
|
|
reconstructor = BeamReconstructor(
|
|
w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, deconvolver=deconvolver
|
|
)
|
|
result = reconstructor.reconstruct(blurred_planes)
|
|
purity_with_deconv, _ = result.purity[(0, 0)]
|
|
|
|
assert purity_with_deconv > purity_no_deconv
|
|
assert purity_with_deconv > 0.9
|
|
|
|
|
|
def test_reconstruct_forces_phase_retrieval_fallback():
|
|
basis = make_basis()
|
|
gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4)
|
|
z_list = [0.47, 0.5, 0.53]
|
|
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=3)
|
|
|
|
reconstructor = BeamReconstructor(
|
|
w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, force_phase_retrieval=True
|
|
)
|
|
result = reconstructor.reconstruct(planes)
|
|
|
|
assert result.used_phase_retrieval is True
|
|
power_fraction, _ = result.purity[(0, 0)]
|
|
assert power_fraction > 0.9
|
|
|
|
|
|
def test_reconstruct_falls_back_automatically_on_high_residual():
|
|
basis = make_basis()
|
|
gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4)
|
|
z_list = [0.47, 0.5, 0.53]
|
|
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=4)
|
|
|
|
reconstructor = BeamReconstructor(
|
|
w0=W0,
|
|
z0=Z0,
|
|
wavelength=WAVELENGTH,
|
|
max_order=2,
|
|
phase_retrieval_residual_threshold=1e-8,
|
|
)
|
|
result = reconstructor.reconstruct(planes)
|
|
|
|
assert result.used_phase_retrieval is True
|