cufinufft v2.2.0 migration

.github/workflows/workflow.yml

@@ -128,7 +128,7 @@ jobs:
     - uses: actions/checkout@v3
     - name: Install dependencies
       run: |
-        pip install -e ".[dev,gpu_12x]"
+        pip install -e ".[dev,gpu-12x]"
     - name: Customize config
       run: |
         echo "Setup tmp dirs and chmod so others can cleanup."

docs/source/installation.rst

@@ -141,24 +141,26 @@ driver version, run ``nvidia-smi`` on the intended system.
    * - CUDA Version
      - ASPIRE Extension
    * - 10.2
-     - gpu_102
+     - gpu-102
    * - 11.0
-     - gpu_110
+     - gpu-110
    * - 11.1
-     - gpu_111
+     - gpu-111
    * - >=11.2
-     - gpu_11x
+     - gpu-11x
+   * - >=12
+     - gpu-12x
 
-For example, if you have CUDA 11.7 installed on your system,
+For example, if you have CUDA 12.3 installed on your system,
 the command below would install GPU packages required for ASPIRE.
 
 ::
 
     # From a local git repo
-    pip install -e ".[gpu_11x]"
+    pip install -e ".[gpu-12x]"
 
     # From PyPI
-    pip install "aspire[gpu_11x]"
+    pip install "aspire[gpu-12x]"
 
 
 By default if the required GPU extensions are correctly installed,

pyproject.toml

@@ -63,11 +63,11 @@ dependencies = [
 "Source" = "https://github.com/ComputationalCryoEM/ASPIRE-Python"
 
 [project.optional-dependencies]
-gpu_102 = ["pycuda", "cupy-cuda102", "cufinufft==1.3"]
-gpu_110 = ["pycuda", "cupy-cuda110", "cufinufft==1.3"]
-gpu_111 = ["pycuda", "cupy-cuda111", "cufinufft==1.3"]
-gpu_11x = ["pycuda", "cupy-cuda11x", "cufinufft==1.3"]
-gpu_12x = ["pycuda", "cupy-cuda12x", "cufinufft==1.3"]
+gpu-102 = ["pycuda", "cupy-cuda102", "cufinufft==1.3"]
+gpu-110 = ["pycuda", "cupy-cuda110", "cufinufft==1.3"]
+gpu-111 = ["pycuda", "cupy-cuda111", "cufinufft==1.3"]
+gpu-11x = ["pycuda", "cupy-cuda11x", "cufinufft==1.3"]
+gpu-12x = ["pycuda", "cupy-cuda12x", "cufinufft==2.2.0"]
 dev = [
     "black",
     "bumpversion",

src/aspire/nufft/cufinufft.py

@@ -4,7 +4,7 @@
 import pycuda.autoinit  # noqa: F401
 import pycuda.driver as cuda  # noqa: F401
 import pycuda.gpuarray as gpuarray  # noqa: F401
-from cufinufft import cufinufft
+from cufinufft import Plan as cufPlan
 
 from aspire.nufft import Plan
 from aspire.utils import complex_type
@@ -60,8 +60,8 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs):
         self.num_pts = fourier_pts.shape[1]
         self.epsilon = max(epsilon, np.finfo(self.dtype).eps)
 
-        self._transform_plan = cufinufft(
-            2, self.sz, self.ntransforms, self.epsilon, -1, dtype=self.dtype
+        self._transform_plan = cufPlan(
+            2, self.sz, self.ntransforms, self.epsilon, -1, dtype=self.complex_dtype
         )
 
         self.adjoint_opts = dict()
@@ -73,22 +73,22 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs):
             )
             self.adjoint_opts["gpu_method"] = 1
 
-        self._adjoint_plan = cufinufft(
+        self._adjoint_plan = cufPlan(
             1,
             self.sz,
             self.ntransforms,
             self.epsilon,
             1,
-            dtype=self.dtype,
+            dtype=self.complex_dtype,
             **self.adjoint_opts,
         )
 
         # Note, I store self.fourier_pts_gpu so the GPUArrray life
         #   is tied to instance, instead of this method.
         self.fourier_pts_gpu = gpuarray.to_gpu(self.fourier_pts)
 
-        self._transform_plan.set_pts(*self.fourier_pts_gpu)
-        self._adjoint_plan.set_pts(*self.fourier_pts_gpu)
+        self._transform_plan.setpts(*self.fourier_pts_gpu)
+        self._adjoint_plan.setpts(*self.fourier_pts_gpu)
 
     def transform(self, signal):
         """
@@ -140,7 +140,7 @@ def transform(self, signal):
 
         result_gpu = gpuarray.GPUArray(res_shape, dtype=self.complex_dtype)
 
-        self._transform_plan.execute(result_gpu, signal_gpu)
+        self._transform_plan.execute(signal_gpu, out=result_gpu)
 
         result = result_gpu.get()
         # ASPIRE-Python/703
@@ -187,7 +187,7 @@ def adjoint(self, signal):
 
         result_gpu = gpuarray.GPUArray(res_shape, dtype=self.complex_dtype)
 
-        self._adjoint_plan.execute(signal_gpu, result_gpu)
+        self._adjoint_plan.execute(signal_gpu, out=result_gpu)
 
         result = result_gpu.get()
         # ASPIRE-Python/703