fix: drop deleted TransformerNUFFTPyNUFFT from jax_profiling dataset setup

jax_profiling/dataset_setup/interferometer.py

@@ -62,9 +62,14 @@
     },
     "alma_high_res": {
         # ALMA-style 5000-vis coverage at 0.025"/px on a 512x512 real-space
-        # grid. Simulator uses pynufft because (a) DFT would need a dense
-        # (n_vis x n_real_space) matrix = ~20GB and OOM on a 15GB laptop,
-        # and (b) nufftax requires Python >= 3.12 (PyAutoGPU venv is 3.10).
+        # grid. Simulator uses the nufftax-backed NUFFT because a DFT would
+        # need a dense (n_vis x n_real_space) matrix = ~20GB and OOM on a 15GB
+        # laptop: 5000 vis x 512x512 = 1.31e9, far above the ~1e7 n_vis*n_pix
+        # crossover past which the NUFFT is the only feasible path.
+        # This was pynufft until 2026-08-23; TransformerNUFFTPyNUFFT was
+        # deleted by PyAutoArray#475. The old "nufftax needs Python >= 3.12
+        # but the PyAutoGPU venv is 3.10" caveat is moot — the whole stack now
+        # declares requires-python >= 3.12, so a 3.10 venv cannot run it at all.
         # Likelihood scripts can still use TransformerDFT downstream — they
         # only read uv_wavelengths from this file.
         "n_visibilities": 5000,
@@ -73,7 +78,7 @@
         "shape_native": (512, 512),
         "noise_sigma": 100.0,
         "seed": 1,
-        "transformer_class": "nufft_pynufft",
+        "transformer_class": "nufft",
     },
     "hannah": {
         "n_visibilities": 16984,
@@ -137,7 +142,6 @@ def simulate(instrument: str):
     transformer_class = {
         "dft": al.TransformerDFT,
         "nufft": al.TransformerNUFFT,
-        "nufft_pynufft": al.TransformerNUFFTPyNUFFT,
     }[transformer_choice]
 
     simulator = al.SimulatorInterferometer(

jax_profiling/jit/datacube/delaunay.py

@@ -231,8 +231,12 @@ def jit_profile(func, label, *args, n_repeats=10):
             real_space_mask=real_space_mask,
             transformer_class=al.TransformerDFT,
             # DFT is intentional even at ALMA-scale visibility counts — profiling
-            # the JAX-traceable path is the goal, NUFFT (pynufft) is not yet
-            # JIT-friendly.
+            # the JAX-traceable path is the goal. NOTE (2026-08-23): the old
+            # reason given here, "NUFFT (pynufft) is not yet JIT-friendly", is
+            # no longer true — the pynufft backend is deleted and the
+            # nufftax-backed TransformerNUFFT jits fine via xp=jnp. DFT is left
+            # in place so existing profiling baselines stay comparable; revisit
+            # if NUFFT-path timings are wanted.
             raise_error_dft_visibilities_limit=False,
         ).apply_sparse_operator(use_jax=True, show_progress=False)
         for _ in range(n_channels)

jax_profiling/jit/interferometer/delaunay.py

@@ -64,8 +64,11 @@
 
 - ``dataset.transformer.transform_mapping_matrix`` is JIT-friendly for
   ``TransformerDFT`` (a single matrix multiply) and the default SMA preset
-  uses it. ``TransformerNUFFT`` (pynufft-based) is not JIT-friendly; if you
-  swap the transformer the step-8 timing will fall back to eager-only.
+  uses it. ``TransformerNUFFT`` (nufftax-backed) is JIT-friendly *too*, but
+  only when called with ``xp=jnp``: ``xp`` is a per-call argument that
+  defaults to ``np``, and the numpy path raises
+  ``TracerArrayConversionError`` under ``jax.jit``. Verified 2026-08-23 —
+  the two paths agree to 5e-16 relative.
 - The visibility-space χ² in step 13 separates the complex visibilities and
   noise into real/imag components inside the JIT body (matching the
   ``pixelization/likelihood_function.py`` reference). Complex-valued JIT
@@ -214,8 +217,12 @@ def jit_profile(func, label, *args, n_repeats=10):
         real_space_mask=real_space_mask,
         transformer_class=al.TransformerDFT,
         # DFT is intentional even at ALMA-scale visibility counts — profiling
-        # the JAX-traceable path is the goal, NUFFT (pynufft) is not yet
-        # JIT-friendly.
+        # the JAX-traceable path is the goal. NOTE (2026-08-23): the old
+        # reason given here, "NUFFT (pynufft) is not yet JIT-friendly", is no
+        # longer true — the pynufft backend is deleted and the nufftax-backed
+        # TransformerNUFFT jits fine via xp=jnp. DFT is left in place here so
+        # existing profiling baselines stay comparable, not because NUFFT
+        # cannot be traced; revisit if NUFFT-path timings are wanted.
         raise_error_dft_visibilities_limit=False,
     )
 
@@ -553,8 +560,9 @@ def ray_trace_mesh_raw(mesh_raw):
 # For ``TransformerDFT`` this is a single complex matrix multiply
 # ``D @ M`` where D is the discrete Fourier matrix (n_vis × n_image) and
 # M is the real mapping matrix (n_image × source_pixels). JIT-friendly.
-# For ``TransformerNUFFT`` (pynufft-based, ALMA-scale) this step would fall
-# back to eager — flag if you swap the transformer.
+# For ``TransformerNUFFT`` (nufftax-backed, ALMA-scale) this step also jits,
+# provided ``transform_mapping_matrix`` is called with ``xp=jnp``; with the
+# default ``xp=np`` it raises ``TracerArrayConversionError`` under jit.
 
 print("\n--- Step 8: Transformed mapping matrix (NUFFT) ---")