Merge pull request #101 from flatironinstitute/maint/cufinufft-released

Try released finufft 2.2

.github/workflows/main.yml

@@ -29,11 +29,11 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install .[dev]
+        python -m pip install -e .[dev]
 
     - name: Pytest
       run: |
-        pytest tests/ --cov
+        pytest tests/ --cov=pytorch_finufft
 
 
   lint:

.pre-commit-config.yaml

@@ -27,7 +27,7 @@ repos:
         - id: ruff
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.5.1
+    rev: v1.8.0
     hooks:
       - id: mypy
         exclude: ^tests/|^examples/

ci/Jenkinsfile

@@ -16,53 +16,21 @@ pipeline {
       }
       environment {
     HOME = "$WORKSPACE"
-    LIBRARY_PATH = "$WORKSPACE/finufft/build"
-    LD_LIBRARY_PATH = "$WORKSPACE/finufft/build"
       }
       steps {
 
-    // TODO - reconsider install strategy once finufft/cufinufft 2.2 is released
-  checkout scmGit(branches: [[name: '*/master']],
-                  extensions: [cloneOption(noTags: true, reference: '', shallow: true),
-                               [$class: 'RelativeTargetDirectory', relativeTargetDir: 'finufft'],
-                               cleanAfterCheckout()],
-                  userRemoteConfigs: [[url: 'https://github.com/flatironinstitute/finufft']])
-
-    sh '''#!/bin/bash -ex
-      nvidia-smi
-    '''
-    sh '''#!/bin/bash -ex
-      echo $HOME
-      ls
-    '''
-    sh label: "build CUFINUFFT", script: '''#!/bin/bash -ex
-        cd finufft
-        # v100 cuda arch
-        cuda_arch="70"
-        mkdir build
-        cd build
-        cmake .. -DFINUFFT_USE_CUDA=ON \
-                  -DFINUFFT_USE_CPU=OFF \
-                  -DFINUFFT_BUILD_TESTS=OFF \
-                  -DCMAKE_CUDA_ARCHITECTURES="$cuda_arch" \
-                  -DBUILD_TESTING=OFF \
-                  -DCMAKE_BUILD_TYPE=RelWithDebInfo
-        cmake --build . -j 4
-    '''
-
     sh 'python3 -m venv $HOME'
     sh label: "set up virtual environment", script: '''#!/bin/bash -ex
       source $HOME/bin/activate
       python3 -m pip install --upgrade pip
       # we could also move pytorch install inside docker
       python3 -m pip install "torch~=2.1.0" --index-url https://download.pytorch.org/whl/cu118
-      python3 -m pip install finufft/python/cufinufft
 
-      python3 -m pip install -e .[dev] # can make [dev,cuda] once cufinufft released?
+      python3 -m pip install -e .[dev,cuda]
     '''
     sh label: "run tests", script: '''#!/bin/bash -ex
       source $HOME/bin/activate
-      python3 -m pytest -k "cuda" tests/ --cov -v --error-for-skips
+      python3 -m pytest -k "cuda" tests/ --cov=pytorch_finufft -v --error-for-skips
     '''
       }
     }

pyproject.toml

@@ -12,7 +12,7 @@ authors = [
   {name = "Brian Ward", email="[email protected]"},
 ]
 license = { text = "MIT" }
-dependencies = ["finufft>= 2.1", "torch >= 2", "numpy", "scipy"]
+dependencies = ["finufft>= 2.2", "torch >= 2", "numpy", "scipy"]
 classifiers = [
   "Development Status :: 3 - Alpha",
   "Intended Audience :: Developers",
@@ -101,12 +101,3 @@ norecursedirs = [
   ".git",
   "__pycache__",
 ]
-
-[tool.coverage.run]
-source = ["tests"]
-
-[coverage.paths]
-source = "pytorch_finufft"
-
-[coverage.run]
-branch = true

pytorch_finufft/functional.py

@@ -96,7 +96,7 @@ class FinufftType1(torch.autograd.Function):
     MODEORD_DEFAULT = 1  # note: FINUFFT default is 0
 
     @staticmethod
-    def setup_context(  # type: ignore[override]
+    def setup_context(
         ctx: Any,
         inputs: Tuple[
             torch.Tensor, torch.Tensor, Any, Optional[Dict[str, Union[int, float]]]
@@ -142,7 +142,7 @@ def forward(  # type: ignore[override]
         batch_dims = values.shape[:-1]
         finufft_out = nufft_func(
             *points,
-            values.reshape(-1, values.shape[-1]).squeeze(),
+            values.reshape(-1, values.shape[-1]),
             output_shape,
             **finufftkwargs,
         )
@@ -237,10 +237,8 @@ def backward(  # type: ignore[override]
             batched_values = values.reshape(nbatch, 1, values.shape[-1])
 
             ramped_grad_output = (
-                (coord_ramps * batched_grad_output * 1j * _i_sign)
-                .reshape(-1, *shape)
-                .squeeze()
-            )
+                coord_ramps * batched_grad_output * 1j * _i_sign
+            ).reshape(-1, *shape)
 
             backprop_ramp = (
                 nufft_func(*points, ramped_grad_output, isign=_i_sign, **finufftkwargs)
@@ -277,7 +275,7 @@ class FinufftType2(torch.autograd.Function):
     MODEORD_DEFAULT = 1  # note: FINUFFT default is 0
 
     @staticmethod
-    def setup_context(  # type: ignore[override]
+    def setup_context(
         ctx: Any,
         inputs: Tuple[
             torch.Tensor, torch.Tensor, Optional[Dict[str, Union[int, float]]]
@@ -324,7 +322,7 @@ def forward(  # type: ignore[override]
         shape = targets.shape[-ndim:]
         finufft_out = nufft_func(
             *points,
-            targets.reshape(-1, *shape).squeeze(),
+            targets.reshape(-1, *shape),
             **finufftkwargs,
         )
         finufft_out = finufft_out.reshape(*batch_dims, npoints)
@@ -410,10 +408,8 @@ def backward(  # type: ignore[override]
 
             coord_ramps = coordinate_ramps(shape, device)
 
-            ramped_targets = (
-                (coord_ramps * batched_targets * 1j * _i_sign)
-                .reshape(-1, *shape)
-                .squeeze()  # squeeze to work around finufft issue#367
+            ramped_targets = (coord_ramps * batched_targets * 1j * _i_sign).reshape(
+                -1, *shape
             )
 
             backprop_ramp = (