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>
This commit is contained in:
@@ -0,0 +1,64 @@
|
||||
import matplotlib.figure
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from he11lib.data import MeasurementPlane, ReconstructionResult
|
||||
from he11lib.plotting import plot_center_trace, plot_mode_purity, plot_residuals
|
||||
|
||||
|
||||
def make_result(**overrides):
|
||||
defaults = dict(
|
||||
purity={(0, 0): (0.9, 0.1), (1, 0): (0.1, -0.2)},
|
||||
reconstructed_field=np.zeros((5, 5), dtype=complex),
|
||||
centers=[(0.0, 0.0), (1e-4, -1e-4), (2e-4, -2e-4)],
|
||||
pointing_angle_deg=0.5,
|
||||
residuals=[np.ones((5, 5)), np.ones((5, 5)) * 2, np.ones((5, 5)) * 3],
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return ReconstructionResult(**defaults)
|
||||
|
||||
|
||||
def make_planes():
|
||||
return [
|
||||
MeasurementPlane(flux=np.zeros((5, 5)), z=0.3),
|
||||
MeasurementPlane(flux=np.zeros((5, 5)), z=0.5),
|
||||
MeasurementPlane(flux=np.zeros((5, 5)), z=0.7),
|
||||
]
|
||||
|
||||
|
||||
def test_plot_mode_purity_draws_one_bar_per_mode():
|
||||
result = make_result()
|
||||
fig = plot_mode_purity(result)
|
||||
|
||||
assert isinstance(fig, matplotlib.figure.Figure)
|
||||
ax = fig.axes[0]
|
||||
assert len(ax.patches) == len(result.purity)
|
||||
|
||||
|
||||
def test_plot_center_trace_plots_one_point_per_plane():
|
||||
planes = make_planes()
|
||||
result = make_result()
|
||||
fig = plot_center_trace(planes, result)
|
||||
|
||||
assert isinstance(fig, matplotlib.figure.Figure)
|
||||
ax = fig.axes[0]
|
||||
line = ax.lines[0]
|
||||
assert len(line.get_xdata()) == len(planes)
|
||||
|
||||
|
||||
def test_plot_residuals_draws_one_axes_per_plane():
|
||||
planes = make_planes()
|
||||
result = make_result()
|
||||
fig = plot_residuals(planes, result)
|
||||
|
||||
assert isinstance(fig, matplotlib.figure.Figure)
|
||||
image_axes = [ax for ax in fig.axes if ax.images]
|
||||
assert len(image_axes) == len(planes)
|
||||
|
||||
|
||||
def test_plot_residuals_raises_when_phase_retrieval_used_without_residuals():
|
||||
planes = make_planes()
|
||||
result = make_result(residuals=[], used_phase_retrieval=True)
|
||||
|
||||
with pytest.raises(ValueError, match="residuals"):
|
||||
plot_residuals(planes, result)
|
||||
Reference in New Issue
Block a user