Skip to content

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

New issue
New issue
Closed
Closed
fix(inversion): make AbstractMeshGeometry picklable (xp module → _use_jax bool + property)#320

Description

@Jammy2211
Jammy2211
opened on May 16, 2026

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)

Activity

  1. Jammy2211
    closed this as completedin aee5e94on May 16, 2026

  2. Jammy2211 commented on May 16, 2026

    @Jammy2211
    Jammy2211
    on May 16, 2026
    CollaboratorAuthor

    Shipped — 2026-05-16

    PR

    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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions