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>

he11lib/fitting.py

@@ -0,0 +1,233 @@
+"""Joint nonlinear least-squares modal fit with automatic mode-set growth."""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+from scipy.optimize import least_squares
+
+from .data import MeasurementPlane, ReconstructionResult, validate_planes
+from .geometry import GeometryCalibration
+from .modes import LGBasis
+from .noise import NoiseEstimator
+
+
+def generate_mode_shells(max_order: int) -> list[list[tuple[int, int]]]:
+    """Group candidate LG_{p,l} modes into shells of increasing order 2p+|l|."""
+    shells: list[list[tuple[int, int]]] = [[] for _ in range(max_order + 1)]
+    for p in range(0, max_order + 1):
+        for l in range(-max_order, max_order + 1):
+            order = 2 * p + abs(l)
+            if order <= max_order:
+                shells[order].append((p, l))
+    return shells
+
+
+class ModalFitter:
+    """Fits LG mode coefficients, beam center/pointing, and geometry to measured planes."""
+
+    def __init__(self, basis: LGBasis, noise_estimator: NoiseEstimator | None = None):
+        self.basis = basis
+        self.noise_estimator = noise_estimator or NoiseEstimator()
+
+    def fit(
+        self,
+        planes: list[MeasurementPlane],
+        modes: list[tuple[int, int]],
+        initial_coefficients: dict[tuple[int, int], complex] | None = None,
+        initial_center: tuple[float, float] = (0.0, 0.0),
+        initial_tilt_deg: tuple[float, float] = (0.0, 0.0),
+        initial_pixel_scale: float | None = None,
+        initial_viewing_angle_deg: float = 0.0,
+    ) -> ReconstructionResult:
+        """Jointly fit complex coefficients for `modes` plus center/tilt/geometry."""
+        validate_planes(planes)
+
+        unknown_scale_idx = [i for i, p in enumerate(planes) if p.pixel_scale is None]
+        unknown_angle_idx = [i for i, p in enumerate(planes) if p.viewing_angle_deg is None]
+        weights = [np.sqrt(self.noise_estimator.weights(p.flux)) for p in planes]
+
+        def pack_initial() -> np.ndarray:
+            x: list[float] = []
+            for i, mode in enumerate(modes):
+                c = (initial_coefficients or {}).get(mode)
+                if c is None:
+                    # Nonzero seed for every mode: starting a coefficient at
+                    # exactly 0+0j sits at a flat/degenerate point for the
+                    # optimizer and can prevent it from ever leaving zero.
+                    c = 1.0 + 0j if i == 0 else 0.1 + 0.05j
+                x += [c.real, c.imag]
+            x += [initial_center[0], initial_center[1], initial_tilt_deg[0], initial_tilt_deg[1]]
+            for _ in unknown_scale_idx:
+                x.append(initial_pixel_scale if initial_pixel_scale is not None else 1e-4)
+            for _ in unknown_angle_idx:
+                x.append(initial_viewing_angle_deg)
+            return np.array(x, dtype=float)
+
+        n_modes = len(modes)
+
+        def unpack(x: np.ndarray):
+            coeffs = {mode: complex(x[2 * i], x[2 * i + 1]) for i, mode in enumerate(modes)}
+            offset = 2 * n_modes
+            x0, y0, tilt_x_deg, tilt_y_deg = x[offset : offset + 4]
+            offset += 4
+            scales = {}
+            for idx in unknown_scale_idx:
+                scales[idx] = x[offset]
+                offset += 1
+            angles = {}
+            for idx in unknown_angle_idx:
+                angles[idx] = x[offset]
+                offset += 1
+            return coeffs, (x0, y0), (tilt_x_deg, tilt_y_deg), scales, angles
+
+        def plane_center(x0: float, y0: float, tilt_deg: tuple[float, float], z: float):
+            drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[0]))
+            drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[1]))
+            return x0 + drift_x, y0 + drift_y
+
+        def model_flux_for_plane(i: int, plane: MeasurementPlane, coeffs, center0, tilt_deg, scales, angles):
+            scale = plane.pixel_scale if plane.pixel_scale is not None else scales[i]
+            angle = plane.viewing_angle_deg if plane.viewing_angle_deg is not None else angles[i]
+            calib = GeometryCalibration(plane)
+            x_grid, y_grid = calib.physical_coordinates(pixel_scale=scale, viewing_angle_deg=angle)
+            cx, cy = plane_center(center0[0], center0[1], tilt_deg, plane.z)
+            field = self.basis.field_superposition(x_grid - cx, y_grid - cy, plane.z, coeffs)
+            return np.abs(field) ** 2
+
+        def residuals(x: np.ndarray) -> np.ndarray:
+            coeffs, center0, tilt_deg, scales, angles = unpack(x)
+            parts = []
+            for i, plane in enumerate(planes):
+                model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles)
+                parts.append(((plane.flux - model_flux) * weights[i]).ravel())
+            return np.concatenate(parts)
+
+        x0_vec = pack_initial()
+        # 'trf' + x_scale='jac' handles the very different natural magnitudes
+        # of these parameters (coefficients ~O(1), pixel_scale ~O(1e-3),
+        # angles ~O(1-90)); plain 'lm' can terminate prematurely on 'xtol'
+        # because its unscaled step-size test is dominated by the largest
+        # parameters.
+        opt_result = least_squares(
+            residuals, x0_vec, method="trf", x_scale="jac", max_nfev=5000
+        )
+
+        coeffs, center0, tilt_deg, scales, angles = unpack(opt_result.x)
+
+        total_power = sum(abs(c) ** 2 for c in coeffs.values())
+        if total_power == 0:
+            total_power = 1.0
+        purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()}
+
+        centers = [plane_center(center0[0], center0[1], tilt_deg, p.z) for p in planes]
+        pointing_angle_deg = float(np.hypot(tilt_deg[0], tilt_deg[1]))
+
+        geometry: dict[str, float] = {}
+        for i in range(len(planes)):
+            geometry[f"pixel_scale_{i}"] = (
+                planes[i].pixel_scale if planes[i].pixel_scale is not None else scales[i]
+            )
+            geometry[f"viewing_angle_deg_{i}"] = (
+                planes[i].viewing_angle_deg if planes[i].viewing_angle_deg is not None else angles[i]
+            )
+
+        residual_maps = []
+        for i, plane in enumerate(planes):
+            model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles)
+            residual_maps.append(plane.flux - model_flux)
+
+        coefficient_uncertainty = self._estimate_uncertainty(opt_result, modes, coeffs, total_power)
+
+        reference_z = min(planes, key=lambda p: abs(p.z - self.basis.z0)).z
+        field_at_reference = self._field_on_default_grid(coeffs, reference_z)
+
+        return ReconstructionResult(
+            purity=purity,
+            reconstructed_field=field_at_reference,
+            centers=centers,
+            pointing_angle_deg=pointing_angle_deg,
+            geometry=geometry,
+            residuals=residual_maps,
+            coefficient_uncertainty=coefficient_uncertainty,
+            used_phase_retrieval=False,
+        )
+
+    def fit_auto(
+        self,
+        planes: list[MeasurementPlane],
+        max_order: int = 4,
+        bic_improvement_threshold: float = 10.0,
+    ) -> ReconstructionResult:
+        """Fit with automatic mode-set growth, capped at `max_order`."""
+        validate_planes(planes)
+        shells = generate_mode_shells(max_order)
+
+        current_modes = list(shells[0])
+        best_result = self.fit(planes, current_modes)
+        best_bic = self._bic(planes, best_result, current_modes)
+
+        grew_until_cap = True
+        for shell in shells[1:]:
+            trial_modes = current_modes + shell
+            warm_start = self._warm_start_coefficients(best_result, current_modes)
+            trial_result = self.fit(planes, trial_modes, initial_coefficients=warm_start)
+            trial_bic = self._bic(planes, trial_result, trial_modes)
+
+            if trial_bic < best_bic - bic_improvement_threshold:
+                current_modes = trial_modes
+                best_result = trial_result
+                best_bic = trial_bic
+            else:
+                grew_until_cap = False
+                break
+
+        if grew_until_cap and len(shells) > 1:
+            warnings.warn(
+                "Automatic mode-set growth hit the configured max_order cap "
+                f"({max_order}) while still improving the fit; consider raising max_order.",
+                stacklevel=2,
+            )
+
+        return best_result
+
+    def _warm_start_coefficients(
+        self, previous_result: ReconstructionResult, previous_modes: list[tuple[int, int]]
+    ) -> dict[tuple[int, int], complex]:
+        """Reconstruct approximate complex coefficients from a previous fit's purity."""
+        coeffs = {}
+        for mode in previous_modes:
+            fraction, phase = previous_result.purity[mode]
+            amplitude = np.sqrt(max(fraction, 0.0))
+            coeffs[mode] = amplitude * np.exp(1j * phase)
+        return coeffs
+
+    def _bic(self, planes: list[MeasurementPlane], result: ReconstructionResult, modes: list[tuple[int, int]]) -> float:
+        chi2 = sum(np.sum((r * np.sqrt(self.noise_estimator.weights(p.flux))) ** 2) for r, p in zip(result.residuals, planes))
+        n_data = sum(p.flux.size for p in planes)
+        n_params = 2 * len(modes) + 4
+        return float(chi2 + n_params * np.log(n_data))
+
+    def _estimate_uncertainty(self, opt_result, modes, coeffs, total_power):
+        try:
+            jac = opt_result.jac
+            cov = np.linalg.pinv(jac.T @ jac)
+        except np.linalg.LinAlgError:
+            return {mode: float("nan") for mode in modes}
+
+        uncertainty = {}
+        for i, mode in enumerate(modes):
+            var_re = cov[2 * i, 2 * i]
+            var_im = cov[2 * i + 1, 2 * i + 1]
+            c = coeffs[mode]
+            sigma_c = np.sqrt(max(var_re, 0) + max(var_im, 0))
+            uncertainty[mode] = float(2 * abs(c) * sigma_c / total_power)
+        return uncertainty
+
+    def _field_on_default_grid(self, coeffs, z: float, n: int = 128, half_width_in_w: float = 6.0):
+        w_z = self.basis.beam_radius(z)
+        extent = half_width_in_w * w_z
+        coords = np.linspace(-extent, extent, n)
+        x, y = np.meshgrid(coords, coords)
+        return self.basis.field_superposition(x, y, z, coeffs)