refactor: unify FitEllipse perimeter sampling; add JAX support

Replace the 300-iteration mask-rejection loop in
FitEllipse.points_from_major_axis_from with a single-pass algorithm
that traces under jax.jit: generate N perimeter angles, evaluate the
mask, mark masked positions with NaN via xp.where. Downstream
nansum / nanmean reductions in residual_map / chi_squared already
handle NaN positions correctly, so the reduction layer is unchanged.

Add use_jax: bool = False to FitEllipse.__init__ and an _xp property
dispatching to jax.numpy or numpy. The 8 reduction properties
(data_interp, noise_map_interp, residual_map, etc.) thread self._xp
through their inner arithmetic and guard the aa.ArrayIrregular
return-wrap behind `if xp is np:` — under JAX they return raw
jax.Array instead.

Behaviour changes for masked datasets (no impact on unmasked):

- The dynamic regen-with-n_i mechanism that maintained exactly N
  unmasked points is gone. chi_squared now sums over the perimeter
  positions that aren't masked (fewer terms when some are masked).
- ValueError for geometrically impossible masks is gone; output is
  all-NaN, chi_squared is NaN, the non-linear search treats this as
  -inf likelihood (loud rejection at the search level).

Net diff: fit_ellipse.py -71 lines. JAX-numpy chi_squared agreement
on the prompt-2 dataset: rel diff 1.5e-14 (machine epsilon).

Step 6 of 7 in z_features/ellipse_fitting_jax.md. Issue PyAutoGalaxy#409.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 2 people