Improve sparse mapping matrix in the interferometer sparse inversion (autoarray)

[HRSAstro]

Summary

Thanks for making InversionInterferometerSparse.data_vector sparse (_data_vector_mapper_from, the bincount / .at[].add segment-sum, in 2026.10.18.1). One step on the sparse path still builds the dense mapping matrix L (image pixels × mesh pixels): the reconstructed image. For large images this is still the largest array in the fit.

Where

AbstractInversionInterferometer.mapped_reconstructed_data_dict (inversion/inversion/interferometer/abstract.py), which the sparse inversion inherits:

mapped_reconstructed_data = (
    inversion_util.mapped_reconstructed_data_via_mapping_matrix_from(
        mapping_matrix=linear_obj.mapping_matrix,
        reconstruction=reconstruction,
        xp=self._xp,
    )
)

[Jammy2211]

@HRSAstro you are right, and thank you for spotting it. The dense mapping matrix is still built on the sparse path by the inherited mapped_reconstructed_data_dict, though only after the likelihood: sampling uses the fast chi-squared and never touches it, so the cost lands on model visibilities, galaxy image dicts, the natural model image and the array-free chi-squared, which is fit outputs, visualization and array-free log_likelihood.

We are fixing it now as the mirror of the sparse data-vector segment-sum (scatter-add of vals * reconstruction[cols] over image pixels, static shapes under jit), with parity tests against the dense result on numpy and JAX.

One thing: your post ends at the quoted code block, so if you had a specific replacement in mind it did not make it into the text; if it differs from the above, please add it. Also, where did you hit the memory peak, and at what mask size, mesh size and backend?

[HRSAstro]

Thanks, that matches what I see. The replacement I had in mind is the same scatter-add, so nothing is missing beyond what you describe. For completeness, this is what the end of the post should have said:

rows, cols, vals = self._sparse_triplets_slim_from(mapper)   # as in _data_vector_mapper_from
# NumPy
image = np.bincount(rows, weights=vals * reconstruction[cols], minlength=n_image_slim)
# JAX
image = jnp.zeros(n_image_slim).at[rows].add(vals * reconstruction[cols])

Linear light profiles can keep the dense path, since they only add a few columns.

On where we hit the memory peak: the crash we actually saw was in data_vector on 2026.10.2.1 (the dense L plus the hstack copy), which 2026.10.18.1 already fixes. PyUVimage maps the reconstruction to the image with its own sparse M and never calls mapped_reconstructed_data_dict, so I found the remaining one by reading the code rather than by running out of memory.

The setup where data_vector was killed:

• ALMA data of NGC 4274, 2.6 M visibilities
• 35″ field on a 384 × 384 native grid with a square mask, so all 147,456 pixels unmasked
• 48 × 48 rectangular mesh (2,304 pixels)
• InversionInterferometerSparse, JAX backend on CPU in float64
• L = 2.7 GB, 5.4 GB with the hstack copy
• killed in a 7 GB, 2-core Linux container

With a sparse mapping in its place, the same fit peaked at 0.73 GB. At a 192 × 192 mesh on the same grid, L alone would be 43 GB, so mapped_reconstructed_data_dict would hit the same wall when outputs are written.

[Jammy2211]

@HRSAstro thank you for the snippet, which is exactly what we implemented: InversionInterferometerSparse.mapped_reconstructed_data_dict now scatter-adds vals * reconstruction[cols] over the slim triplets (bincount on numpy, .at[].add under JAX), with linear light profiles keeping the dense path. Merged in PyAutoLabs/PyAutoArray#641 with parity tests against the dense result on Rectangular and Delaunay meshes, sub-gridding, two mappers and mapper-plus-linear-profile cases. It will be in the next release; we will post here when that is on PyPI.