Martino Ferrari
ef57ec81e4
The optimizer's parameter vector is now built dynamically: LG coefficients, per-plane center, and both pointing angles stay always free; each CameraModel field and each plane's z join the fit (bounded to its +/- tolerance) only when its paired tolerance is nonzero, and are otherwise substituted as fixed constants. Also seeds the first mode's coefficient with a small imaginary offset (1.0 + 0.05j instead of 1.0 + 0j) rather than exactly on the real axis: for a single mode, intensity depends only on |c| (phase is unobservable), so Im=0 sits exactly on that flat/degenerate valley, aligned with a coordinate axis, giving a zero-gradient Jacobian column that destabilized trf's trust-region step and prevented pointing-angle recovery from a default (0, 0) initial guess. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
281 lines
12 KiB
Python
281 lines
12 KiB
Python
"""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 (
|
|
CameraModel,
|
|
CameraModelTolerance,
|
|
GeometryCalibration,
|
|
camera_from_values,
|
|
camera_to_values,
|
|
tolerance_to_values,
|
|
)
|
|
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]],
|
|
camera: CameraModel,
|
|
camera_tolerance: CameraModelTolerance,
|
|
initial_coefficients: dict[tuple[int, int], complex] | None = None,
|
|
initial_center: tuple[float, float] = (0.0, 0.0),
|
|
initial_pointing_deg: tuple[float, float] = (0.0, 0.0),
|
|
) -> ReconstructionResult:
|
|
"""Jointly fit complex coefficients for `modes` plus center/pointing/geometry.
|
|
|
|
Every `CameraModel` field with a nonzero `camera_tolerance` entry,
|
|
and every plane whose `z_tolerance` is nonzero, is refined within
|
|
`[nominal - tolerance, nominal + tolerance]`; zero-tolerance fields
|
|
are held fixed at their nominal value.
|
|
"""
|
|
validate_planes(planes)
|
|
weights = [np.sqrt(self.noise_estimator.weights(p.flux)) for p in planes]
|
|
|
|
camera_nominal = camera_to_values(camera)
|
|
camera_tol = tolerance_to_values(camera_tolerance)
|
|
free_camera_idx = [i for i, t in enumerate(camera_tol) if t > 0]
|
|
|
|
free_z_idx = [i for i, p in enumerate(planes) if p.z_tolerance > 0]
|
|
|
|
n_modes = len(modes)
|
|
n_always_free = 2 * n_modes + 4 # coefficients + center(2) + pointing(2)
|
|
|
|
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.
|
|
# A purely-real seed (Im=0) sits exactly on the flat/
|
|
# degenerate valley of a single mode's phase (for one
|
|
# mode alone, |c|^2 -- not arg(c) -- is all that's
|
|
# observable in intensity), giving a zero-gradient
|
|
# column that destabilizes trf's trust-region step;
|
|
# a small imaginary offset avoids landing on that axis.
|
|
c = 1.0 + 0.05j if i == 0 else 0.1 + 0.05j
|
|
x += [c.real, c.imag]
|
|
x += [initial_center[0], initial_center[1], initial_pointing_deg[0], initial_pointing_deg[1]]
|
|
for i in free_camera_idx:
|
|
x.append(camera_nominal[i])
|
|
for i in free_z_idx:
|
|
x.append(planes[i].z)
|
|
return np.array(x, dtype=float)
|
|
|
|
def pack_bounds() -> tuple[np.ndarray, np.ndarray]:
|
|
lower = [-np.inf] * n_always_free
|
|
upper = [np.inf] * n_always_free
|
|
for i in free_camera_idx:
|
|
lower.append(camera_nominal[i] - camera_tol[i])
|
|
upper.append(camera_nominal[i] + camera_tol[i])
|
|
for i in free_z_idx:
|
|
lower.append(planes[i].z - planes[i].z_tolerance)
|
|
upper.append(planes[i].z + planes[i].z_tolerance)
|
|
return np.array(lower), np.array(upper)
|
|
|
|
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_h_deg, tilt_v_deg = x[offset : offset + 4]
|
|
offset += 4
|
|
|
|
camera_values = list(camera_nominal)
|
|
for i in free_camera_idx:
|
|
camera_values[i] = x[offset]
|
|
offset += 1
|
|
fitted_camera = camera_from_values(camera_values)
|
|
|
|
z_values = [p.z for p in planes]
|
|
for i in free_z_idx:
|
|
z_values[i] = x[offset]
|
|
offset += 1
|
|
|
|
return coeffs, (x0, y0), (tilt_h_deg, tilt_v_deg), fitted_camera, z_values
|
|
|
|
def plane_center(x0: float, y0: float, pointing_deg: tuple[float, float], z: float):
|
|
drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[0]))
|
|
drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[1]))
|
|
return x0 + drift_x, y0 + drift_y
|
|
|
|
def model_flux_for_plane(plane, fitted_camera, z, coeffs, center0, pointing_deg):
|
|
calib = GeometryCalibration(fitted_camera)
|
|
x_grid, y_grid = calib.physical_coordinates(plane.flux.shape, z)
|
|
cx, cy = plane_center(center0[0], center0[1], pointing_deg, z)
|
|
field = self.basis.field_superposition(x_grid - cx, y_grid - cy, z, coeffs)
|
|
return np.abs(field) ** 2
|
|
|
|
def residuals(x: np.ndarray) -> np.ndarray:
|
|
coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(x)
|
|
parts = []
|
|
for i, plane in enumerate(planes):
|
|
model_flux = model_flux_for_plane(
|
|
plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg
|
|
)
|
|
parts.append(((plane.flux - model_flux) * weights[i]).ravel())
|
|
return np.concatenate(parts)
|
|
|
|
x0_vec = pack_initial()
|
|
lower, upper = pack_bounds()
|
|
# 'trf' + x_scale='jac' handles the very different natural
|
|
# magnitudes of these parameters (coefficients ~O(1), focal length
|
|
# ~O(1e3-1e4), angles ~O(1-90), z ~O(0.1-1)); plain 'lm' can
|
|
# terminate prematurely on 'xtol' because its unscaled step-size
|
|
# test is dominated by the largest parameters. 'lm' also doesn't
|
|
# support bounds, which the tolerance mechanism requires.
|
|
opt_result = least_squares(
|
|
residuals, x0_vec, method="trf", x_scale="jac", bounds=(lower, upper), max_nfev=5000
|
|
)
|
|
|
|
coeffs, center0, pointing_deg, fitted_camera, z_values = 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], pointing_deg, z_values[i])
|
|
for i in range(len(planes))
|
|
]
|
|
|
|
geometry: dict[str, float] = dict(zip(
|
|
(
|
|
"focal_length_px", "position_x", "position_y", "position_z",
|
|
"yaw_deg", "pitch_deg", "roll_deg",
|
|
"principal_point_x", "principal_point_y",
|
|
),
|
|
camera_to_values(fitted_camera),
|
|
))
|
|
for i in range(len(planes)):
|
|
geometry[f"z_{i}"] = z_values[i]
|
|
|
|
residual_maps = []
|
|
for i, plane in enumerate(planes):
|
|
model_flux = model_flux_for_plane(
|
|
plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg
|
|
)
|
|
residual_maps.append(plane.flux - model_flux)
|
|
|
|
coefficient_uncertainty = self._estimate_uncertainty(opt_result, modes, coeffs, total_power)
|
|
|
|
reference_idx = min(range(len(planes)), key=lambda i: abs(z_values[i] - self.basis.z0))
|
|
field_at_reference = self._field_on_default_grid(coeffs, z_values[reference_idx])
|
|
|
|
return ReconstructionResult(
|
|
purity=purity,
|
|
reconstructed_field=field_at_reference,
|
|
centers=centers,
|
|
pointing_angle_horizontal_deg=pointing_deg[0],
|
|
pointing_angle_vertical_deg=pointing_deg[1],
|
|
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)
|