fix(transformer): tell Python <3.12 users to upgrade, not to pip install nufftax (#409)

The nufftax ModuleNotFoundError said "Install it via the command
`pip install nufftax`". On Python 3.11 that is actively bad advice.

The nufftax releases that work with our JAX stack (0.4.x) declare
requires_python >= 3.12, so pip resolves 0.6.x instead — and both 0.6.x
releases break the interferometer inversion:

  - 0.6.0 uses jax.interpreters.batching.not_mapped, removed in JAX 0.10
  - 0.6.1 cannot handle the rank-4 input transform_mapping_matrix produces
    under vmap (ValueError: too many values to unpack (expected 3))

So a 3.11 user who follows the old message installs a broken nufftax and fails
partway through a fit instead of at import — strictly worse than the original
error. There is currently no nufftax release that both installs on 3.11 and
works.

The message now branches on sys.version_info: below 3.12 it says not to install
nufftax and offers upgrading to 3.12+ or the pynufft backend; 3.12+ keeps the
existing text. Both branches covered by tests.

Only the error text changes — no behaviour change on 3.12+.


Claude-Session: https://claude.ai/code/session_012EZTtyLUAyKWuATyk4mkcm

Co-authored-by: Jammy2211 <JNightingale2211@gmail.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: 3 people

autoarray/operators/transformer.py

@@ -1,5 +1,6 @@
 import copy
 import numpy as np
+import sys
 import warnings
 from typing import Optional, Tuple
 
@@ -41,6 +42,32 @@ def pynufft_exception():
 
 
 def nufftax_exception():
+    # On Python 3.11 `pip install nufftax` is actively bad advice. The releases
+    # that work with our JAX stack (0.4.x) require Python >= 3.12, so pip resolves
+    # 0.6.x instead — and 0.6.0/0.6.1 both break the interferometer inversion
+    # (0.6.0 uses `jax.interpreters.batching.not_mapped`, removed in JAX 0.10;
+    # 0.6.1 cannot handle the rank-4 input `transform_mapping_matrix` produces
+    # under vmap). There is currently no nufftax release that both installs on
+    # 3.11 and works, so send 3.11 users to 3.12+ or to the pynufft backend.
+    if sys.version_info < (3, 12):
+        raise ModuleNotFoundError(
+            "\n--------------------\n"
+            "You are attempting to perform interferometer analysis with the default "
+            f"JAX-native `TransformerNUFFT`, on Python {sys.version_info.major}.{sys.version_info.minor}.\n\n"
+            "The optional library nufftax (https://github.com/GragasLab/nufftax) is not installed, "
+            "and on this Python version it CANNOT be usefully installed:\n\n"
+            "  - nufftax releases that work with PyAutoArray require Python >= 3.12.\n"
+            "  - The releases that do install here (0.6.x) are incompatible with the "
+            "JAX version PyAutoArray requires, and fail partway through a fit.\n\n"
+            "Do NOT `pip install nufftax` on this Python version. Instead either:\n\n"
+            "  1. Upgrade to Python 3.12 or newer (recommended), then "
+            "`pip install 'autoarray[optional]'`; or\n"
+            "  2. Use the legacy pynufft backend, by passing "
+            "`transformer_class=TransformerNUFFTPyNUFFT` and running "
+            "`pip install pynufft==2022.2.2`.\n\n"
+            "----------------------"
+        )
+
     raise ModuleNotFoundError(
         "\n--------------------\n"
         "You are attempting to perform interferometer analysis with the default "

test_autoarray/operators/test_transformer.py

@@ -306,3 +306,50 @@ def f(img_arr):
     assert np.asarray(result) == pytest.approx(
         np.asarray(expected), rel=1.0e-6, abs=1.0e-10
     )
+
+
+def test__nufftax_exception__pre_3_12_tells_user_to_upgrade_not_to_pip_install():
+    """
+    On Python < 3.12 there is no usable nufftax release: the versions that work
+    require >= 3.12, and the ones that do install (0.6.x) break mid-fit. The
+    error must therefore steer the user away from `pip install nufftax`.
+    """
+    import collections
+    from unittest import mock
+    from autoarray.operators import transformer
+
+    version_info = collections.namedtuple(
+        "version_info", "major minor micro releaselevel serial"
+    )
+
+    with mock.patch.object(
+        transformer.sys, "version_info", version_info(3, 11, 0, "final", 0)
+    ):
+        with pytest.raises(ModuleNotFoundError) as exc:
+            transformer.nufftax_exception()
+
+    message = str(exc.value)
+    assert "Python 3.11" in message
+    assert "Do NOT `pip install nufftax`" in message
+    assert "Python 3.12 or newer" in message
+    assert "TransformerNUFFTPyNUFFT" in message
+
+
+def test__nufftax_exception__3_12_and_later_keeps_the_plain_install_instruction():
+    import collections
+    from unittest import mock
+    from autoarray.operators import transformer
+
+    version_info = collections.namedtuple(
+        "version_info", "major minor micro releaselevel serial"
+    )
+
+    with mock.patch.object(
+        transformer.sys, "version_info", version_info(3, 12, 0, "final", 0)
+    ):
+        with pytest.raises(ModuleNotFoundError) as exc:
+            transformer.nufftax_exception()
+
+    message = str(exc.value)
+    assert "Install it via the command `pip install nufftax`" in message
+    assert "Do NOT" not in message