Files
he11lib/tests/test_data.py
T
Martino Ferrari 03b63ba03a Initial commit: he11lib mode-purity reconstruction library
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>
2026-07-02 21:47:49 +02:00

94 lines
2.8 KiB
Python

import numpy as np
import pytest
from he11lib.data import MeasurementPlane, ReconstructionResult, validate_planes
def test_measurement_plane_stores_fields():
flux = np.ones((4, 4))
plane = MeasurementPlane(flux=flux, z=0.3)
assert plane.z == 0.3
assert np.array_equal(plane.flux, flux)
assert plane.pixel_scale is None
assert plane.viewing_angle_deg is None
assert plane.label is None
def test_measurement_plane_stores_optional_fields():
flux = np.ones((4, 4))
plane = MeasurementPlane(
flux=flux, z=0.4, pixel_scale=0.1, viewing_angle_deg=5.0, label="plane_40cm"
)
assert plane.pixel_scale == 0.1
assert plane.viewing_angle_deg == 5.0
assert plane.label == "plane_40cm"
def test_measurement_plane_rejects_non_2d_flux():
with pytest.raises(ValueError, match="2D"):
MeasurementPlane(flux=np.ones((4, 4, 3)), z=0.3)
def test_measurement_plane_rejects_non_positive_z():
with pytest.raises(ValueError, match="positive"):
MeasurementPlane(flux=np.ones((4, 4)), z=0.0)
with pytest.raises(ValueError, match="positive"):
MeasurementPlane(flux=np.ones((4, 4)), z=-0.1)
def test_validate_planes_rejects_fewer_than_three():
planes = [
MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
MeasurementPlane(flux=np.ones((4, 4)), z=0.4),
]
with pytest.raises(ValueError, match="[Aa]t least 3"):
validate_planes(planes)
def test_validate_planes_rejects_mismatched_shapes():
planes = [
MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
MeasurementPlane(flux=np.ones((5, 5)), z=0.4),
MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
]
with pytest.raises(ValueError, match="same shape"):
validate_planes(planes)
def test_validate_planes_rejects_duplicate_z():
planes = [
MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
]
with pytest.raises(ValueError, match="distinct"):
validate_planes(planes)
def test_validate_planes_accepts_valid_list():
planes = [
MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
MeasurementPlane(flux=np.ones((4, 4)), z=0.4),
MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
]
validate_planes(planes) # should not raise
def test_reconstruction_result_stores_fields():
result = ReconstructionResult(
purity={(0, 0): (1.0, 0.0)},
reconstructed_field=np.ones((4, 4), dtype=complex),
centers=[(0.0, 0.0)],
pointing_angle_deg=0.0,
geometry={"pixel_scale": 0.1, "viewing_angle_deg": 2.0},
residuals=[np.zeros((4, 4))],
coefficient_uncertainty={(0, 0): 0.01},
used_phase_retrieval=False,
)
assert result.purity[(0, 0)] == (1.0, 0.0)
assert result.used_phase_retrieval is False