fix: use_jax no longer demotes the NUFFT preload to the JAX brute force (#539)

Every existing workspace `apply_sparse_operator(use_jax=True)` call was mapped
onto `method="jax"`, silently putting the production path on the brute force;
`use_jax` is now only honoured when a brute-force method is already selected.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF

Author: 2 people

autoarray/dataset/interferometer/dataset.py

@@ -251,10 +251,10 @@ def apply_sparse_operator(
 
         The default builder (`method="nufft"`) computes the precision operator as a type-1 NUFFT, so
         it costs `O(N_vis * nspread^2 + M log M)` for `M = 4 * Ny * Nx` — seconds even at a million
-        visibilities. The brute-force builders (`method="numpy"` / `"jax"`, and the `use_jax` kwarg)
-        are `O(N_vis * N_pix)` and can take minutes to hours; they are kept as the reference the
-        NUFFT builder is pinned against. Either way the result can be cached to disk and reloaded
-        via `nufft_precision_operator=`.
+        visibilities. The brute-force builders (`method="numpy"` / `"jax"`) are `O(N_vis * N_pix)`
+        and can take minutes to hours; they are kept as the reference the NUFFT builder is pinned
+        against. Either way the result can be cached to disk and reloaded via
+        `nufft_precision_operator=`.
 
         Parameters
         ----------
@@ -284,15 +284,21 @@ def apply_sparse_operator(
         show_memory
             If `True`, memory usage statistics are printed while computing the NUFFT precision matrix.
         use_jax
-            If `True`, JAX is used to accelerate the NUFFT precision matrix computation.
+            Only honoured when a brute-force builder is selected: `method="numpy"` with
+            `use_jax=True` runs the JAX brute force (equivalent to `method="jax"`). Under the
+            default `method="nufft"` it is ignored, because the NUFFT already runs on JAX --
+            so an existing `use_jax=True` call keeps the fast path rather than being demoted
+            to the `O(N_vis * N_pix)` brute force.
 
             `PYAUTO_DISABLE_JAX=1` overrides this to `False`. That variable is a
             harness-level switch, not a preference: it is the documented way to force the
             NumPy path (the workspace `start_here` guides name it beside `use_jax=False`),
             and the smoke profiles set it so a fast run does not pay a JIT compile. An
             explicit `use_jax=True` in a script -- which is the right thing for a script
             demonstrating the production path to say -- must therefore not defeat it, or
-            the harness pays 2.3-3.2 s of compile for a backend it asked to disable.
+            the harness pays 2.3-3.2 s of compile for a backend it asked to disable. (The
+            same switch demotes the `"nufft"` builder to the NumPy brute force inside
+            `nufft_precision_operator_from`.)
 
         Precondition
         ------------
@@ -317,6 +323,9 @@ def apply_sparse_operator(
             If any visibility has unequal real and imaginary noise sigma.
         """
 
+        # `use_jax` now only selects between the two brute forces (the `"nufft"` builder
+        # runs on JAX whatever it says), so this clears it before it can upgrade
+        # `method="numpy"` to the JAX brute force under the kill switch.
         if disable_jax():
             use_jax = False
 
@@ -357,7 +366,7 @@ def apply_sparse_operator(
             n_vis = self.uv_wavelengths.shape[0]
             n_pix = self.real_space_mask.pixels_in_mask
 
-            if method != "nufft" or use_jax:
+            if method != "nufft":
                 logger.info(
                     f"INTERFEROMETER - The precision operator is being built by a brute-force "
                     f"builder, which is O(N_vis x N_pix) = O({n_vis * n_pix:.1e}) and can take "
@@ -432,9 +441,9 @@ def psf_precision_operator_from(
 
         The default builder (`method="nufft"`) computes this as a type-1 (adjoint) NUFFT, which is
         `O(N_vis * nspread^2 + M log M)` for `M = 4 * Ny * Nx` — seconds even at a million
-        visibilities. The brute-force builders (`method="numpy"` / `"jax"`, and the `use_jax`
-        kwarg) are `O(N_vis * N_pix)` and can take minutes to hours on a CPU for a
-        high-resolution mask; they are kept as the reference the NUFFT builder is pinned against.
+        visibilities. The brute-force builders (`method="numpy"` / `"jax"`) are
+        `O(N_vis * N_pix)` and can take minutes to hours on a CPU for a high-resolution mask;
+        they are kept as the reference the NUFFT builder is pinned against.
         The result can still be saved to disk and reloaded rather than recomputed on each run —
         use `apply_sparse_operator(nufft_precision_operator=...)` to attach a cached result.
 
@@ -448,7 +457,9 @@ def psf_precision_operator_from(
         show_memory
             If `True`, memory usage statistics are printed during computation.
         use_jax
-            If `True`, the JAX brute-force builder is used (equivalent to `method="jax"`).
+            Only honoured when a brute-force builder is selected: `method="numpy"` with
+            `use_jax=True` runs the JAX brute force (equivalent to `method="jax"`). It is
+            ignored under the default `method="nufft"`, which already runs on JAX.
         method
             Which builder computes the operator: `"nufft"` (default), `"numpy"` or `"jax"`.
         eps

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py

@@ -122,15 +122,15 @@ def nufft_precision_operator_from(
      matrix construction without performing a NUFFT per source pixel.
 
      -------------------------------------------------------------------------------
-     Backend behaviour
+     Backend behaviour (the two brute-force builders)
      -------------------------------------------------------------------------------
-     - NumPy backend (use_jax=False, default):
+     - NumPy backend (`method="numpy"`):
          * CPU execution
          * Explicit Python loop over visibility chunks
          * Supports progress bars and optional memory reporting
          * Numerically closest to the original reference implementation
 
-     - JAX backend (use_jax=True):
+     - JAX backend (`method="jax"`, or `method="numpy"` with `use_jax=True`):
          * JIT-compilable and GPU/TPU capable
          * Uses fixed-size chunking and lax.fori_loop
          * No Python-side loops during execution
@@ -240,7 +240,13 @@ def nufft_precision_operator_from(
       `O(N_pix*K)` reference builder. Kept as the reference the NUFFT builder is
       pinned against, and used as the fallback below.
     - `"jax"` -- `nufft_precision_operator_via_jax_from`, the same brute force on
-      JAX. `use_jax=True` is kept for backwards compatibility and maps to this.
+      JAX. `method="jax"` is the explicit way to ask for it.
+
+    `use_jax` is only honoured when a brute-force method is selected: it upgrades
+    `method="numpy"` to `method="jax"` and is otherwise ignored. Under the default
+    `method="nufft"` it does nothing, because the NUFFT already runs on JAX -- an
+    existing `use_jax=True` caller therefore keeps the fast path rather than being
+    demoted to the `O(N_pix*K)` brute force.
 
     Two fallbacks to `"numpy"` are taken, both logged loudly (never silently),
     because they cost `O(N_pix*K)` where the NUFFT is `O(K*nspread^2 + M log M)`:
@@ -273,8 +279,12 @@ def nufft_precision_operator_from(
         optimisation), used by `method="nufft"` only. `None` is one shot.
     chunk_k
         The visibility chunk size of the two brute-force builders.
+    use_jax
+        Only honoured when a brute-force method is selected: `method="numpy"` with
+        `use_jax=True` runs the JAX brute force (equivalent to `method="jax"`). It is
+        ignored under the default `method="nufft"`, which already runs on JAX.
     """
-    if use_jax:
+    if method == "numpy" and use_jax:
         method = "jax"
 
     if method not in ("nufft", "numpy", "jax"):

test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py

@@ -577,16 +577,33 @@ def test__nufft_precision_operator_from__method_routes_to_each_builder():
         operator_via_jax,
     )
 
-    # `use_jax=True` is kept for backwards compatibility and maps onto `method="jax"`.
+    # `use_jax=True` only upgrades a brute-force method, so `method="numpy"` with it set is
+    # the JAX brute force.
     np.testing.assert_array_equal(
         np.asarray(
             aa.util.inversion_interferometer.nufft_precision_operator_from(
-                use_jax=True, **inputs
+                method="numpy", use_jax=True, **inputs
             )
         ),
         operator_via_jax,
     )
 
+    # Under the default `method="nufft"` it is ignored -- the NUFFT already runs on JAX, so
+    # honouring it there would demote every existing `use_jax=True` caller (the workspace
+    # `apply_sparse_operator(use_jax=True)` calls) from seconds to the O(N_pix * K) brute force.
+    operator_use_jax = np.asarray(
+        aa.util.inversion_interferometer.nufft_precision_operator_from(
+            use_jax=True, **inputs
+        )
+    )
+
+    np.testing.assert_array_equal(operator_use_jax, operator_via_nufft)
+    _assert_matches_brute_force(operator_use_jax, operator_via_np)
+
+    # The control: the NUFFT array is not the JAX brute-force array, so the assertion above is
+    # testing the routing and not two builders that happen to agree bitwise.
+    assert not np.array_equal(operator_use_jax, operator_via_jax)
+
     # The default is the NUFFT builder, and it agrees with the brute force.
     operator_default = np.asarray(
         aa.util.inversion_interferometer.nufft_precision_operator_from(**inputs)
@@ -616,10 +633,17 @@ def test__nufft_precision_operator_from__disable_jax_falls_back_to_the_numpy_bui
 
     # `PYAUTO_DISABLE_JAX=1` is the harness-level kill switch. Both the default NUFFT builder and
     # the `"jax"` brute force run on JAX, so both must fall back -- and loudly, because the NumPy
-    # brute force is O(N_pix * K) where the NUFFT is O(K * nspread^2 + M log M).
+    # brute force is O(N_pix * K) where the NUFFT is O(K * nspread^2 + M log M). `use_jax=True`
+    # falls back either way: ignored under the default, and demoted again when it upgrades
+    # `method="numpy"` to the JAX brute force.
     monkeypatch.setenv("PYAUTO_DISABLE_JAX", "1")
 
-    for kwargs in ({}, {"method": "jax"}, {"use_jax": True}):
+    for kwargs in (
+        {},
+        {"method": "jax"},
+        {"use_jax": True},
+        {"method": "numpy", "use_jax": True},
+    ):
         caplog.clear()
 
         with caplog.at_level("WARNING"):