From 90bd238356203233ab766976306865999a247ad6 Mon Sep 17 00:00:00 2001
From: James Nightingale <james.w.nightingale@durham.ac.uk>
Date: Sun, 25 Oct 2026 09:05:54 -0700
Subject: [PATCH] fix: install finufft in Array-based Colab bootstraps (#204)

---
 autonerves/setup_colab.py           |  5 ++++-
 test_autonerves/test_setup_colab.py | 30 +++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/autonerves/setup_colab.py b/autonerves/setup_colab.py
index 139b139..b09374d 100644
--- a/autonerves/setup_colab.py
+++ b/autonerves/setup_colab.py
@@ -75,7 +75,10 @@
 ]
 
 _AUTOFIT_STACK = ["autonerves", "autofit"]
-_AUTOGALAXY_STACK = _AUTOFIT_STACK + ["autoarray", "autogalaxy"]
+# AutoArray's CPU NUFFT backend is not supplied by Colab. The bootstrap uses
+# --no-deps, so array-based projects must install it explicitly; the floor
+# matches AutoArray's Linux requirement.
+_AUTOGALAXY_STACK = _AUTOFIT_STACK + ["autoarray", "autogalaxy", "finufft>=2.4.0"]
 _AUTOLENS_STACK = _AUTOGALAXY_STACK + ["autolens"]
 
 # One entry per Colab-enabled notebook repository. ``top_package`` is the
diff --git a/test_autonerves/test_setup_colab.py b/test_autonerves/test_setup_colab.py
index 2c4a004..d35344c 100644
--- a/test_autonerves/test_setup_colab.py
+++ b/test_autonerves/test_setup_colab.py
@@ -108,6 +108,36 @@ def test_empty_device_list(self, fake_jax):
 
 
 class TestRegistry:
+    @pytest.mark.parametrize("project", tuple(setup_colab._PROJECTS))
+    def test_colab_installs_cpu_nufft_for_array_stacks(self, project, monkeypatch, tmp_path):
+        """Exercise the no-deps install command, where a missing package stays missing."""
+        from autonerves import conf
+        from packaging.requirements import Requirement
+
+        google = types.ModuleType("google")
+        google.__path__ = []
+        google.colab = types.ModuleType("google.colab")
+        monkeypatch.setitem(sys.modules, "google", google)
+        monkeypatch.setitem(sys.modules, "google.colab", google.colab)
+        monkeypatch.setenv("XLA_FLAGS", "")
+        monkeypatch.chdir(tmp_path)
+        with mock.patch.object(subprocess, "check_call") as install, \
+                mock.patch.object(setup_colab, "check_jax_using_gpu", return_value=False), \
+                mock.patch.object(setup_colab, "_clone_workspace"), \
+                mock.patch.object(conf.instance, "push"):
+            setup_colab.setup(project, workspace_dir=str(tmp_path))
+
+        command = install.call_args.args[0]
+        assert command[:4] == [sys.executable, "-m", "pip", "install"]
+        assert command[-1] == "--no-deps"
+        requirements = {r.name: r for r in map(Requirement, command[4:-1])}
+        if "autoarray" in requirements:
+            assert "finufft" in requirements, project
+            assert requirements["finufft"].specifier.contains("2.4.0")
+            assert not requirements["finufft"].specifier.contains("2.3.0")
+        else:
+            assert "finufft" not in requirements, project
+
     def test_all_projects_have_required_fields(self):
         required = {
             "project_name",