Feature/jax w tilde preload by Jammy2211 · Pull Request #202 · Jammy2211/PyAutoArray

Jammy2211 · 2026-01-16T16:37:11Z

The fast interferometry JAX GPU implementation (#201) uses a preload of the NUFFT in order to compute the curvature_matrix.

This curvature_preload matrix is computed once, before lens modeling, and reused throughout lens modeling. For high number of visibilities and resolution real space mask, this calculation can take minutes or hours on a CPU.

The previous PR did not convert this calculation to JAX or run on GPU.

This pull request makes the W-Tilde curvature preload computation support JAX and GPU, with profiling suggeting at least x100 speed up for high resolution datasets, with the calculation taking under a minute for the highest resolution / visibilities ALMA datasets tested.

It also includes utilities for safely saving and loading precomputed data, which check metadata to ensure the loaded data matches the data being analysed.

…d calculaiton

…e mask meta data checks

Copilot

Pull request overview

This pull request adds JAX/GPU support for the W-Tilde curvature preload computation in interferometer analysis, enabling significant performance improvements (100x+ speedup) for high-resolution datasets. The PR introduces utilities for safely saving and loading precomputed curvature data with metadata validation.

Changes:

Adds JAX implementation of W-Tilde curvature preload computation alongside existing NumPy implementation
Introduces metadata-based save/load utilities to ensure precomputed data is only reused when compatible with current analysis parameters
Updates API to expose use_jax parameter for enabling GPU acceleration

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
`autoarray/__init__.py`	Exports new `load_curvature_preload_if_compatible` utility function
`autoarray/dataset/interferometer/dataset.py`	Adds `use_jax` parameter to `apply_w_tilde` method and updates log message
`autoarray/dataset/interferometer/w_tilde.py`	Implements metadata utilities and save/load methods for curvature preload
`autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py`	Refactors main function into NumPy and JAX implementations with improved documentation
`test_autoarray/dataset/interferometer/test_dataset.py`	Adds test for metadata validation when loading precomputed data

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

Copilot · 2026-01-16T16:42:19Z

autoarray/dataset/interferometer/w_tilde.py

@@ -1,9 +1,208 @@
+import json
+import hashlib
+from dataclasses import dataclass


The dataclass import is unused in this file. It should be removed to keep the imports clean.

Suggested change

from dataclasses import dataclass

Copilot · 2026-01-16T16:42:19Z

autoarray/dataset/interferometer/w_tilde.py

+    rtol, atol
+        Tolerances for pixel scale comparisons (normally exact is fine


The docstring mentions rtol parameter but the function signature only has atol. Either remove rtol from the docstring or add it to the function signature.

Suggested change

rtol, atol

Tolerances for pixel scale comparisons (normally exact is fine

atol

Tolerance for pixel scale comparisons (normally exact is fine

Copilot · 2026-01-16T16:42:20Z

autoarray/dataset/interferometer/w_tilde.py

+    np.ndarray or None
+        The loaded curvature_preload if compatible, otherwise None (unless raise_on_mismatch=True).


The return type documentation is incorrect. The function always raises ValueError if incompatible (line 201), it never returns None. The docstring should state: "The loaded curvature_preload if compatible, otherwise raises ValueError." Also, raise_on_mismatch parameter is mentioned but doesn't exist.

Suggested change

np.ndarray or None

The loaded curvature_preload if compatible, otherwise None (unless raise_on_mismatch=True).

np.ndarray

The loaded curvature_preload if compatible, otherwise raises ValueError.

Copilot · 2026-01-16T16:42:20Z

autoarray/dataset/interferometer/dataset.py

        if curvature_preload is None:

-            logger.info("INTERFEROMETER - Computing W-Tilde... May take a moment.")
+            logger.info("INTERFEROMETER – Computing W-Tilde; runtime scales with visibility count and mask resolution, extreme inputs may exceed hours.")


The character between 'INTERFEROMETER' and 'Computing' appears to be an en-dash (U+2013) rather than a hyphen-minus. For consistency with the rest of the codebase (e.g., line 489 in inversion_interferometer_util.py uses regular hyphen), this should use a standard hyphen-minus character.

Suggested change

logger.info("INTERFEROMETER – Computing W-Tilde; runtime scales with visibility count and mask resolution, extreme inputs may exceed hours.")

logger.info("INTERFEROMETER - Computing W-Tilde; runtime scales with visibility count and mask resolution, extreme inputs may exceed hours.")

Copilot · 2026-01-16T16:42:20Z

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py

+
+    def _compute_all_quadrants(gy, gx, *, chunk_k: int):


The nested function _compute_all_quadrants captures multiple variables from the outer scope (K, n_chunks, idx, ku_x, kv_x, w_x, y_shape, x_shape) which can make the code harder to reason about. Consider passing these as explicit parameters or restructuring the code to make dependencies clearer.

Suggested change

def _compute_all_quadrants(gy, gx, *, chunk_k: int):

quadrant_context = {

"K": K,

"n_chunks": n_chunks,

"idx": idx,

"ku_x": ku_x,

"kv_x": kv_x,

"w_x": w_x,

"y_shape": y_shape,

"x_shape": x_shape,

}

def _compute_all_quadrants(gy, gx, *, chunk_k: int):

K = quadrant_context["K"]

n_chunks = quadrant_context["n_chunks"]

idx = quadrant_context["idx"]

ku_x = quadrant_context["ku_x"]

kv_x = quadrant_context["kv_x"]

w_x = quadrant_context["w_x"]

y_shape = quadrant_context["y_shape"]

x_shape = quadrant_context["x_shape"]

Copilot · 2026-01-16T16:42:20Z

test_autoarray/dataset/interferometer/test_dataset.py

    assert (dataset.uv_wavelengths == 3.0 * np.ones((19, 2))).all()
+
+
+def test__curvature_preload_metadata_from(


The test only validates the error case (incompatible mask) but doesn't verify that the loaded preload data is correct when it is compatible. Consider adding an assertion to check that the loaded curvature_preload matches the original saved data.

Copilot · 2026-01-16T16:42:21Z

test_autoarray/dataset/interferometer/test_dataset.py

+
+    with pytest.raises(ValueError):
+
+        curvature_preload = aa.load_curvature_preload_if_compatible(


Variable curvature_preload is not used.

Suggested change

curvature_preload = aa.load_curvature_preload_if_compatible(

aa.load_curvature_preload_if_compatible(

Copilot · 2026-01-16T16:42:21Z

test_autoarray/dataset/interferometer/test_dataset.py

+        overwrite=True,
+    )
+
+    curvature_preload = aa.load_curvature_preload_if_compatible(


This assignment to 'curvature_preload' is unnecessary as it is redefined before this value is used.

Suggested change

curvature_preload = aa.load_curvature_preload_if_compatible(

aa.load_curvature_preload_if_compatible(

Copilot · 2026-01-16T16:42:21Z

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py

+            if show_memory and show_progress and "_report_memory" in globals():
+                try:
+                    globals()["_report_memory"](acc)
+                except Exception:


'except' clause does nothing but pass and there is no explanatory comment.

Suggested change

except Exception:

except Exception:

# _report_memory is a best-effort diagnostic; ignore any errors it raises.

Jammy2211 added 6 commits January 16, 2026 13:31

successfully implemented JAX version of interferometer w tilde preloa…

565af3a

…d calculaiton

static argnam on chunk k in jax.jit

1ad6572

added functionality to save and load curvature preload with real spac…

d185f9f

…e mask meta data checks

add unit test on loading curvature preload

3cf259f

black

7979d3c

tming of functions including JAX one improved

afb6e90

Jammy2211 requested review from Sketos and Copilot January 16, 2026 16:37

Copilot started reviewing on behalf of Jammy2211 January 16, 2026 16:37 View session

Copilot AI reviewed Jan 16, 2026

View reviewed changes

copilot review

6145a2a

Jammy2211 merged commit e3ce0f7 into main Jan 17, 2026
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feature/jax w tilde preload#202

Feature/jax w tilde preload#202
Jammy2211 merged 7 commits intomainfrom
feature/jax_w_tilde_preload

Jammy2211 commented Jan 16, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Copilot AI Jan 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		rtol, atol
		Tolerances for pixel scale comparisons (normally exact is fine

		np.ndarray or None
		The loaded curvature_preload if compatible, otherwise None (unless raise_on_mismatch=True).

	logger.info("INTERFEROMETER – Computing W-Tilde; runtime scales with visibility count and mask resolution, extreme inputs may exceed hours.")
	logger.info("INTERFEROMETER - Computing W-Tilde; runtime scales with visibility count and mask resolution, extreme inputs may exceed hours.")

-    def _compute_all_quadrants(gy, gx, *, chunk_k: int):
+    quadrant_context = {
+        "K": K,
+        "n_chunks": n_chunks,
+        "idx": idx,
+        "ku_x": ku_x,
+        "kv_x": kv_x,
+        "w_x": w_x,
+        "y_shape": y_shape,
+        "x_shape": x_shape,
+    }
+    def _compute_all_quadrants(gy, gx, *, chunk_k: int):
+        K = quadrant_context["K"]
+        n_chunks = quadrant_context["n_chunks"]
+        idx = quadrant_context["idx"]
+        ku_x = quadrant_context["ku_x"]
+        kv_x = quadrant_context["kv_x"]
+        w_x = quadrant_context["w_x"]
+        y_shape = quadrant_context["y_shape"]
+        x_shape = quadrant_context["x_shape"]

		assert (dataset.uv_wavelengths == 3.0 * np.ones((19, 2))).all()


		def test__curvature_preload_metadata_from(


		with pytest.raises(ValueError):

		curvature_preload = aa.load_curvature_preload_if_compatible(

	curvature_preload = aa.load_curvature_preload_if_compatible(
	aa.load_curvature_preload_if_compatible(

	except Exception:
	except Exception:
	# _report_memory is a best-effort diagnostic; ignore any errors it raises.

Conversation

Jammy2211 commented Jan 16, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants