feat(weak): FitWeak class + plotters (#524)

[Jammy2211]

Summary

Step 3 of the weak-lensing series. Adds a FitWeak class that compares a model shear field (derived from a Tracer's mass profiles) against an observed WeakDataset, plus a matching set of four aplt plotters for visualising the fit.

The model shear comes from LensCalc.from_tracer(tracer).shear_yx_2d_via_hessian_from(grid=...) — the exact primitive SimulatorShearYX uses. A noise-free round-trip is therefore bit-exact: residual_map ≈ 0 for the truth tracer.

Each background galaxy contributes two independent measurements (γ₁ and γ₂ share per-galaxy noise but are independent Gaussian draws), so chi_squared sums over N × 2 elements and noise_normalization carries a factor of 2 to match.

API Changes

Pure additions. One new class (autolens.weak.FitWeak) and four new module-level plot helpers (plot_data_vs_model, plot_residuals, plot_chi_squared_map, subplot_fit_weak), all re-exported into autolens.plot and accessible at the top level as al.FitWeak. See full details below.

Test Plan

• pytest test_autolens/weak/ -v → 26 passed (10 new + 16 existing dataset/simulator/plotter)
• python -c "import autolens as al; al.FitWeak" resolves without ImportError
• python -c "import autolens.plot as aplt; aplt.subplot_fit_weak" resolves without ImportError
• grep "^import jax\|^from jax" autolens/weak/fit.py autolens/weak/plot/fit_weak_plots.py → empty (numpy-only)

Full API Changes (for automation & release notes)

Added

New class autolens.weak.fit.FitWeak (also reachable as al.FitWeak):

• FitWeak(dataset: WeakDataset, tracer: Tracer)
• @cached_property model_shear -> ShearYX2DIrregular
• @property residual_map -> np.ndarray (shape (N, 2))
• @property normalized_residual_map -> np.ndarray (shape (N, 2); residual / σ broadcast)
• @property chi_squared_map -> np.ndarray (shape (N, 2))
• @property chi_squared -> float (sum over N*2 elements)
• @property noise_normalization -> float (2 * Σ log(2π σ²) — factor of 2 reflects independent γ₁/γ₂)
• @property log_likelihood -> float (-0.5 * (chi_squared + noise_normalization))
• @property figure_of_merit -> float (same as log_likelihood — no inversion / no log_evidence)

New module autolens.weak.plot.fit_weak_plots, all re-exported into autolens.plot (the aplt namespace):

• plot_data_vs_model(fit, ax=None, title="Data vs Model", output_path=None, output_filename="data_vs_model", output_format=None) — overlays the data shear field (black) and the model shear field (red, alpha 0.6) as headless quivers on a single axes.
• plot_residuals(fit, ax=None, title="Residuals", output_path=None, output_filename="residuals", output_format=None) — headless-quiver of the residual map, colour-coded by |residual| with the RdBu_r colormap.
• plot_chi_squared_map(fit, ax=None, title=r"$\chi^2$ Map", output_path=None, output_filename="chi_squared_map", output_format=None) — scalar-coloured scatter of per-galaxy chi-squared (sum across γ₁/γ₂) via autoarray.plot.grid.plot_grid with colormap="magma".
• subplot_fit_weak(fit, output_path=None, output_filename="subplot_fit_weak", output_format=None, title_prefix=None) — 2×2 mosaic: data | model | data-vs-model overlay | chi-squared.

Migration

None required — additive change.

Design notes

• Standalone class — not inheriting from aa.AbstractFit. That base is shaped for "data + noise_map + mask" pixel-grid fits and doesn't fit shear catalogues cleanly. FitPoint (in autolens.point) follows the same standalone pattern.
• Storage convention preserved. Plotters access the shear field through .ellipticities / .phis only. The fit accesses raw (N, 2) storage but only via element-wise arithmetic (subtraction, division) which is component-symmetric — never raw [:, 0] / [:, 1] indexing.
• (N, 2) chi-squared shape. Each galaxy contributes two independent measurements. Reflected in chi_squared (sums all 2N elements) and noise_normalization (factor of 2 over the (N,) noise array).
• Pure numpy. No JAX imports — JAX adoption is a separate ratchet for the whole weak-lensing module.

🤖 Generated with Claude Code

[Jammy2211]

Workspace follow-up PR: PyAutoLabs/autolens_workspace#188