Feature/fft jax imaging

[Jammy2211]

This pull request extends the the "w_tilde" formalism of interferometry linear algebra to imaging datasets, enabling efficiently modeling using sparse operators on GPU via JAX. Initial test runs indicate this enables lens modeling of Hubble Space Telescope data using pixelized sources in under 10 minutes, with the formalism enabling small amounts of VRAM use.

It also changes the codebase API to replace the previous use of the term "w_tilde" with a new "sparse operator" formalism for both imaging and interferometer datasets. The changes update class names, method signatures, and internal logic to use the new approach, resulting in improved clarity and maintainability. The most important changes are grouped below:

Refactoring from w_tilde to sparse operator formalism:

• Replaced all references to w_tilde (including WTildeImaging and WTildeInterferometer) with sparse_operator (such as ImagingSparseOperator and InterferometerSparseLinAlg) throughout the codebase, including constructor arguments, attributes, and method names in both imaging and interferometer dataset classes. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13]
• Removed the legacy w_tilde modules and classes, including autoarray/dataset/imaging/w_tilde.py and autoarray/dataset/abstract/w_tilde.py. [1] [2]

API and interface updates:

• Updated the GridsDataset and related classes to use the use_sparse_operator flag instead of use_w_tilde, ensuring consistent configuration for the new formalism. [1] [2] [3]
• Updated import statements and module references throughout the codebase to point to the new sparse operator modules and classes. [1] [2]

These changes modernize the linear algebra approach in the codebase, improve naming consistency, and remove outdated code related to the previous w_tilde formalism.

[Copilot]

Pull request overview

This PR completes the migration from the legacy w_tilde-based inversion formalism to a new sparse_operator-based approach in the lensing pipeline, and updates tests to validate equivalence between the mapping-matrix and sparse-operator paths.

Changes:

• Updated imaging tests to compare inversions built from standard mapping matrices versus the new sparse-operator pipeline, using the new apply_sparse_operator API and naming.
• Refactored FitImaging, FitInterferometer, and TracerToInversion to pass sparse_operator through aa.DatasetInterface instead of the removed w_tilde attributes, ensuring the inversion path is driven entirely by the new formalism.
• Cleaned up legacy w_tilde-related imports in autolens.__init__ and performed minor formatting-only adjustments in tracer_util.py.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
test_autolens/imaging/test_simulate_and_fit_imaging.py	Renames a w_tilde comparison test and switches it to use apply_sparse_operator, asserting curvature and regularization matrices match the mapping-matrix inversion, directly validating the new sparse-operator imaging path.
autolens/lens/tracer_util.py	Reflows the NaN-sanitizing xp.where call for deflections and inlines the time-delay distance expression, preserving behavior while improving readability for JAX-related numeric operations.
autolens/lens/to_inversion.py	Removes the legacy w_tilde handling and forwards self.dataset.sparse_operator into aa.DatasetInterface so that all per-plane inversion datasets use the new sparse-operator field.
autolens/interferometer/fit_interferometer.py	Updates tracer_to_inversion to construct its DatasetInterface with sparse_operator=self.dataset.sparse_operator, aligning the interferometer inversion path with the new sparse-operator API.
autolens/imaging/fit_imaging.py	Similarly updates tracer_to_inversion to pass sparse_operator=self.dataset.sparse_operator when building the inversion dataset, ensuring imaging inversions are configured via the new formalism.
autolens/__init__.py	Drops re-exports of WTildeImaging and load_curvature_preload_if_compatible, removing the public w_tilde surface from the top-level API in favor of the sparse-operator approach.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

autolens/init.py:8

• Removing the top-level WTildeImaging and load_curvature_preload_if_compatible imports without adding equivalents for the new sparse-operator classes changes the public autolens API: code that relied on autolens.WTildeImaging or autolens.load_curvature_preload_if_compatible will now fail. If the new sparse-operator types are meant to replace these, consider re-exporting the new public equivalents here or clearly documenting the new import path so the breaking change is explicit for users.

from autoconf import jax_wrapper
from autoconf.dictable import from_dict, from_json, output_to_json, to_dict
from autoarray import preprocess

from autoarray.dataset.imaging.dataset import Imaging
from autoarray.dataset.interferometer.dataset import (
    Interferometer,
)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.