Skip to content

feat: sparse mapped reconstructed image on the interferometer sparse inversion path #639

Closed

Description

@Jammy2211

Overview

Community proposal from @HRSAstro (PyAutoLabs/.github#40, accepted 2026-10-21). After 2026.10.18.1 made InversionInterferometerSparse.data_vector sparse, the inherited mapped_reconstructed_data_dict still builds the dense image-pixels × mesh-pixels mapping matrix L for fit outputs (model visibilities, galaxy image dicts, model_image_natural, array-free chi-squared). On a 384×384 unmasked grid with a 48² mesh L is 2.7 GB; at 192² it would be 43 GB. This task reconstructs the mapper image by scatter-add over the existing sparse triplets instead.

Plan

  • Override mapped_reconstructed_data_dict on InversionInterferometerSparse so mapper blocks are reconstructed by a scatter-add over the slim sparse triplets, never building the dense mapping matrix.
  • Non-mapper linear objects (linear light profiles) keep the dense mapping_matrix @ reconstruction path.
  • Parity tests against the dense result on numpy and under JAX jit: Rectangular and Delaunay meshes, sub-size > 1, two mappers, mapper plus linear functions.
  • No PyAutoGalaxy source change expected; memory before/after in autolens_profiling only if cheap.
  • Tier: glance — merge mode: auto-merge on green if the Witness passes, post-merge summary
Detailed implementation plan

Affected Repositories

  • PyAutoArray (primary)
  • PyAutoGalaxy (consumer, verify only)
  • autolens_profiling (optional memory note)

Branch Survey

Repository Current Branch Dirty?
array/PyAutoArray main clean
galaxy/PyAutoGalaxy main clean
lens/autolens_profiling main clean

Suggested branch: feature/sparse-mapped-reconstructed-image

Implementation Steps

  1. autoarray/inversion/inversion/interferometer/sparse.py: add mapped_reconstructed_data_dict mirroring _data_vector_mapper_from (line 174). For each mapper in source_quantity_dict_from(self.reconstruction): rows, cols, vals = self._sparse_triplets_slim_from(mapper); contrib = vals * reconstruction[cols]; numpy np.bincount(rows, weights=contrib, minlength=n_slim), JAX xp.zeros(n_slim, dtype).at[rows].add(contrib); wrap as Array2D on self.mask exactly as the inherited abstract.py:143-192 does. Other linear objects fall back to mapping_matrix @ reconstruction.
  2. sparse.py:517 mapped_reconstructed_operated_data_dict already consumes the dict; model visibilities inherit the saving unchanged.
  3. Tests in test_autoarray/inversion/inversion/interferometer/test_interferometer.py: build the sparse inversion and the dense mapping inversion on the same fixtures; assert dict values agree to 1e-10 (numpy) and ~1e-12 under jit (dense L is fp32 under JAX, so no bit-identity). Cases: Rectangular, Delaunay, sub_size > 1, two mappers, mapper + linear light profile.
  4. Read-only check of imaging/sparse.py:526 for the same negative-index masking gap; report as a side finding, do not fix here.
  5. PyAutoGalaxy: run the interferometer fit tests against the branch to confirm the consumers (abstract_fit.py:186, fit_interferometer.py:298,626) are unchanged.

Key Files

  • autoarray/inversion/inversion/interferometer/sparse.py — new override; _sparse_triplets_slim_from, _data_vector_mapper_from as the template
  • autoarray/inversion/inversion/interferometer/abstract.py — inherited dense implementation to mirror
  • test_autoarray/inversion/inversion/interferometer/test_interferometer.py — parity tests

Original Prompt

Click to expand starting prompt

Sparse mapped reconstructed image on the interferometer sparse inversion path…

Type: feature
Target: PyAutoArray
Repos:

  • PyAutoArray
  • PyAutoGalaxy
  • autolens_profiling
    Difficulty: medium
    Autonomy: safe
    Priority: high
    Memory: reading-queue.md; wiki/lensing/sources/dark-matter-substructure.md; wiki/galaxies/sources/cosmos-survey.md
    Status: formalised
    Consequence: glance
    Witness: a test asserting the sparse mapped_reconstructed_data_dict equals the dense mapping_matrix @ reconstruction result to 1e-10 on numpy and under JAX jit, for Rectangular and Delaunay meshes, sub-size > 1, two mappers and mapper-plus-linear-functions
    Review-minutes: 3
    Unattended: ready

Sparse mapped reconstructed image on the interferometer sparse inversion path (never build the dense mapping matrix L)

Community proposal accepted 2026-10-21: PyAutoLabs/.github#40 (@HRSAstro, Ideas & Proposals, 2026-10-20). user-facing: send the contributor an update when the fix merges and when it is released. Note the discussion body is truncated at its code block, so the proposer's own replacement is unknown; we are implementing our own.

Problem

After 2026.10.18.1 made InversionInterferometerSparse.data_vector sparse (_data_vector_mapper_from, segment-sum), the inherited AbstractInversionInterferometer.mapped_reconstructed_data_dict (autoarray/inversion/inversion/interferometer/abstract.py:179-180) still calls mapped_reconstructed_data_via_mapping_matrix_from(mapping_matrix=linear_obj.mapping_matrix, ...), which builds the dense (image pixels × mesh pixels) matrix L via the cached Mapper.mapping_matrix. Sampling does not hit it (log_evidence uses fast_chi_squared), but it is built for model visibilities (sparse.py:555 mapped_reconstructed_operated_data_dict), galaxy image dicts (autogalaxy/abstract_fit.py:186), model_image_natural (autogalaxy/interferometer/fit_interferometer.py:626) and the array-free sparse_chi_squared_from (fit_interferometer.py:298), so fit outputs, visualization and array-free log_likelihood pay the n_image × n_mesh memory.

Fix

Override mapped_reconstructed_data_dict in InversionInterferometerSparse (and check the numba sibling) as the mirror of _data_vector_mapper_from (sparse.py:174-196): rows, cols, vals from _sparse_triplets_slim_from(mapper) (which already zeroes invalid slots); contrib = vals * reconstruction[cols]; image = bincount(rows, contrib, minlength=n_pixels) on numpy, .at[rows].add / jax.ops.segment_sum(num_segments=n) under JAX. Non-mapper linear objects keep mapping_matrix @ reconstruction. Do not reuse the imaging helper (imaging/inversion_imaging_util.py:280): it imports JAX unconditionally and takes raw triplets without the negative-index masking (check imaging/sparse.py:526 for the same latent bug as a side finding).

Scope notes

  • Mixed precision: dense L is fp32 under JAX while this path runs fp64 (same note as perf: interferometer sparse data vector without the dense mapping matrix #630); expect ~1e-12 agreement, not bit-identity.
  • mapper.mapping_matrix remains lazily available for plotting; mapped_to_source_from and image regions are numpy-only follow-ups.
  • Record memory before/after on one large real_space_mask in autolens_profiling if cheap.

Activity

  1. Jammy2211 commented on Oct 21, 2026

    @Jammy2211
    CollaboratorAuthor

    Library PR created: #641 (branch feature/sparse-mapped-reconstructed-image, commit 422f60a). Workspace impact: none (internal change). Heart YELLOW acknowledged by the human 2026-10-21 (manifest drift 1 mismatch; test run status unknown; no release rehearsal). Tier glance: auto-merge on green if the Witness (sparse-vs-dense parity tests) passes.

  2. Jammy2211 commented on Oct 21, 2026

    @Jammy2211
    CollaboratorAuthor

    Shipped: #641 merged (7ecc6ef, glance tier, auto-merged on green: unittest 3.12 / 3.13 / nojax all pass; Witness parity tests in-run). InversionInterferometerSparse.mapped_reconstructed_data_dict now scatter-adds over the slim sparse triplets; the dense mapping matrix is no longer built on the sparse interferometer path. Pending release (not yet on PyPI). Side finding for a follow-up issue: imaging _mapped_reconstructed_data_dict_from uses unmasked sparse_triplets_curvature and its helper imports JAX unconditionally.

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