fix(inversion): make AbstractMeshGeometry picklable (xp module → _use_jax bool + property)

[Jammy2211]

Overview

Carve-out from PyAutoFit #1279 (Phase 4 of the JAX visualization roadmap). A picklability spike found that FitImaging cannot be pickled today because AbstractMeshGeometry.__init__ stores self._xp = xp — the literal numpy or jax.numpy module — and Python can't pickle modules. This blocks sending a FitImaging over an mp.Process+Queue or ProcessPoolExecutor boundary, which is the production design for Phase 4 subprocess visualization.

Fix: replace the module-attribute pattern with the _use_jax: bool + _xp property pattern already used in Analysis._xp (PyAutoFit) and AbstractMaker._xp (PyAutoArray decorators per CLAUDE.md). Eliminates the pickle barrier at the API level instead of working around it with __getstate__/__setstate__ hooks.

Spike evidence (PyAutoFit #1279):

• Before fix: TypeError: cannot pickle 'module' object at inversion.linear_obj_list[0].interpolator.mesh_geometry._xp
• After fix (proven via monkey-patch): FitImaging round-trips through pickle.dumps/loads with log_likelihood delta 0.00e+00 on both numpy and JAX backends. Pickle size ~4.6 MB.

Plan

• Change AbstractMeshGeometry.__init__ to store self._use_jax = xp is not np instead of self._xp = xp
• Add _xp as a @property deriving from self._use_jax (returns np or jnp on demand)
• All existing call sites read self._xp (no writes) — they continue to work transparently via the property
• Add a unit test that pickles a MeshGeometryRectangular instance and verifies _xp round-trips to the correct module

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary, library, only repo touched)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean (canonical) — knn-barycentric task holds it via worktree, file-disjoint

Suggested branch: feature/mesh-geometry-picklable
Worktree root: ~/Code/PyAutoLabs-wt/mesh-geometry-picklable/

Conflict resolution: File-level coexistence on PyAutoArray with knn-barycentric (#317). KNN edits inversion/mesh/interpolator/, inversion/mesh/mesh/, and plot/* — no overlap with our scope (inversion/mesh/mesh_geometry/abstract.py). Same precedent as ag-interferometer-kwargs.

Implementation Steps

1. autoarray/inversion/mesh/mesh_geometry/abstract.py:
   • In __init__: replace self._xp = xp with self._use_jax = xp is not np
   • Add @property method _xp that returns jax.numpy if _use_jax else numpy
   • Import of jax.numpy stays lazy (only when _use_jax=True and _xp is accessed) — matches Analysis._xp pattern in PyAutoFit
2. test_autoarray/inversion/pixelization/mesh_geometry/test_picklability.py (new):
   • test_picklable_numpy_backend — construct via __new__ + _use_jax=False, pickle round-trip, assert restored._xp is np
   • test_picklable_jax_backend — same with _use_jax=True, assert restored._xp is jnp (pytest.importorskip('jax'))
3. Run the existing mesh_geometry tests + the new pickle tests + any pixelization tests that touch _xp to verify no regression

Key Files

• PyAutoArray/autoarray/inversion/mesh/mesh_geometry/abstract.py (line 23 + new property)
• PyAutoArray/test_autoarray/inversion/pixelization/mesh_geometry/test_picklability.py (new)

Out of scope

• __getstate__/__setstate__ workaround (Option A in the parent issue's spike) — rejected in favour of the cleaner property pattern (Option B)
• Subprocess visualization itself — that's the parent task #1279
• FitInterferometer NUFFT picklability — separate spike under #1279

Parent

PyAutoFit #1279 — Phase 4 subprocess visualization feasibility (carve-out of the picklability prereq)

[Jammy2211]

Shipped — 2026-05-16

PR

• fix(inversion): make AbstractMeshGeometry picklable (xp module → _use_jax bool + property) #321

Summary

AbstractMeshGeometry now stores _use_jax: bool instead of _xp as a module reference, with _xp exposed as a property derived from the flag. FitImaging is now picklable end-to-end (4.6 MB / fit, bit-exact log_likelihood round-trip on both numpy and JAX backends). Unblocks PyAutoFit #1279 (Phase 4 subprocess visualization).

Session Notes

• Q2 of the Phase 4 feasibility task was deliberately tackled first because picklability constrains the IPC choice for Q1. Outcome: FitImaging pickles cleanly with a single-class fix — strong positive signal for mp.Process+Queue / ProcessPoolExecutor over disk-staged subprocess. Q1 measurement now focuses on which Python-native IPC is fastest under JAX, not whether they can carry a FitImaging.
• Spike found the breakage was ONE attribute on ONE class — not the broader picklability concern feared. Mappers, linear operators, computed inversion matrices, the tracer / dataset / mask graph all pickled fine out-of-box.
• Chose Option B (property-from-bool) over Option A (__getstate__/__setstate__ hooks) for consistency with Analysis._xp (PyAutoFit Phase 2) and AbstractMaker._xp (PyAutoArray decorators per CLAUDE.md). Same pattern in three places now.
• File-disjoint coexistence with knn-barycentric (feat: KNN-barycentric interpolator — drop-in JAX-native Delaunay alternative #317) worked cleanly: both tasks in inversion/mesh/ but on different files. Confirms the soft-conflict-guard / file-level-safety pattern from ag-interferometer-kwargs precedent.
• Spike scripts and q2_findings.md remain in the viz-subprocess-feasibility worktree (~/Code/PyAutoLabs-wt/viz-subprocess-feasibility/) as in-progress notes for Q1/Q3/Q4 work.