Fix MaxRecursionDepth error, Memory leaks and move to CUDA12 (#174)

* Added * Fix * Remove bymistake add * Fix * Fixed lint * Lint * Added refbackend * Fix NDFT * feat: use finufft as ref backend. * feat(tests): move ndft vs nufft tests to own file. * All * as * Move to density_method * Move to cuda 12 * docs_build * docs_build added support for binder * docs_build minor check change * docs_build minor check change * Update src/mrinufft/extras/smaps.py Co-authored-by: Pierre-Antoine Comby <pierre-antoine.comby@crans.org> * make it private * Fix gpunufft * testing docs_build * docs_build tryu to see meesage * docs_build again * docs_build again * event_name * event * try last message * try last message * try last message * try last message * docs_build all setup * docs_build all set * docs_build all set * Go back * Clean up docs_build style --------- Co-authored-by: chaithyagr <chaithyagr@gitlab.com> Co-authored-by: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
mind-inria · Aug 8, 2024 · 87ffecf · 87ffecf
1 parent 60e32d2
commit 87ffecf
Show file tree

Hide file tree

Showing 10 changed files with 55 additions and 31 deletions.
diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml
@@ -118,23 +118,22 @@ jobs:
           source $RUNNER_WORKSPACE/venv/bin/activate
           pip install --upgrade pip wheel
           pip install -e mri-nufft[test]
-          pip install cupy-cuda11x
-          pip install finufft "numpy<2.0"
+          pip install cupy-cuda12x finufft "numpy<2.0"
 
-      - name: Install torch with CUDA 11.8
+      - name: Install torch with CUDA 12.1
         shell: bash
         if: ${{ matrix.backend != 'tensorflow'}}
         run: |
           source $RUNNER_WORKSPACE/venv/bin/activate
-          pip install torch --index-url https://download.pytorch.org/whl/cu118
+          pip install torch
 
       - name: Install backend
         shell: bash
         run: |
           source $RUNNER_WORKSPACE/venv/bin/activate
-          export CUDA_BIN_PATH=/usr/local/cuda-11.8/
-          export PATH=/usr/local/cuda-11.8/bin/:${PATH}
-          export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64/:${LD_LIBRARY_PATH}
+          export CUDA_BIN_PATH=/usr/local/cuda-12.1/
+          export PATH=/usr/local/cuda-12.1/bin/:${PATH}
+          export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH}
           if [[ ${{ matrix.backend }} == "torchkbnufft-gpu" ]]; then
             pip install torchkbnufft
           elif [[ ${{ matrix.backend }} == "tensorflow" ]]; then
@@ -195,11 +194,10 @@ jobs:
         
       - name: Install GPU related interfaces
         run: |
-          export CUDA_BIN_PATH=/usr/local/cuda-11.8/
-          export PATH=/usr/local/cuda-11.8/bin/:${PATH}
-          export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64/:${LD_LIBRARY_PATH}
-          pip install cupy-cuda11x
-          pip install torch --index-url https://download.pytorch.org/whl/cu118
+          export CUDA_BIN_PATH=/usr/local/cuda-12.1/
+          export PATH=/usr/local/cuda-12.1/bin/:${PATH}
+          export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH}
+          pip install cupy-cuda12x torch
           python -m pip install gpuNUFFT cufinufft
 
       - name: Run examples
@@ -282,7 +280,7 @@ jobs:
   BuildDocs:
     name: Build API Documentation
     runs-on: gpu
-    if: ${{ (contains(github.event.head_commit.message, 'docs_build')) || (github.ref == 'refs/heads/master')}}
+    if: ${{ contains(github.event.head_commit.message, 'docs_build') || github.ref == 'refs/heads/master' }}
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -306,11 +304,10 @@ jobs:
       
       - name: Install GPU related interfaces
         run: |
-          export CUDA_BIN_PATH=/usr/local/cuda-11.8/
-          export PATH=/usr/local/cuda-11.8/bin/:${PATH}
-          export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64/:${LD_LIBRARY_PATH}
-          pip install cupy-cuda11x
-          pip install torch --index-url https://download.pytorch.org/whl/cu118
+          export CUDA_BIN_PATH=/usr/local/cuda-12.1/
+          export PATH=/usr/local/cuda-12.1/bin/:${PATH}
+          export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64/:${LD_LIBRARY_PATH}
+          pip install cupy-cuda12x torch
           python -m pip install gpuNUFFT cufinufft
       
       - name: Build API documentation

diff --git a/docs/binder/apt.txt b/docs/binder/apt.txt
@@ -0,0 +1,3 @@
+cmake
+libnfft3-dev
+libgl1-mesa-glx
diff --git a/docs/binder/environment.yml b/docs/binder/environment.yml
@@ -0,0 +1,10 @@
+name: nufft
+channels:
+  - conda-forge
+dependencies:
+  - python=3.10
+  - pip>=19.1.1
+  - jupyter>=1.0.0
+  - pip:
+      - mri-nufft[finufft]
+      - pynfft2
diff --git a/docs/conf.py b/docs/conf.py
@@ -79,6 +79,18 @@
     "filename_pattern": "/example_",
     "ignore_pattern": r"(__init__|conftest|utils).py",
     "nested_sections": True,
+    "binder": {
+        "org": "mind-inria",
+        "repo": "mri-nufft",
+        "branch": "master",
+        "binderhub_url": "https://mybinder.org",
+        "dependencies": [
+            "./binder/apt.txt",
+            "./binder/environment.yml",
+        ],
+        "notebooks_dir": "examples",
+        "use_jupyter_lab": True,
+    },
     "parallel": 3,
 }
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -11,9 +11,9 @@ dynamic = ["version"]
 
 [project.optional-dependencies]
 
-gpunufft = ["gpuNUFFT>=0.8.1", "cupy-cuda11x"]
-torchkbnufft = ["torchkbnufft", "cupy-cuda11x"]
-cufinufft = ["cufinufft", "cupy-cuda11x"]
+gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"]
+torchkbnufft = ["torchkbnufft", "cupy-cuda12x"]
+cufinufft = ["cufinufft", "cupy-cuda12x"]
 finufft = ["finufft"]
 pynfft = ["pynfft2>=1.4.3", "numpy>=2.0.0"]
 pynufft = ["pynufft"]

diff --git a/src/mrinufft/extras/smaps.py b/src/mrinufft/extras/smaps.py
@@ -106,7 +106,7 @@ def low_frequency(
     threshold: float | tuple[float, ...] = 0.1,
     density=None,
     window_fun: str = "ellipse",
-    blurr_factor: float | tuple[float, ...] = 0.0,
+    blurr_factor: int | float | tuple[float, ...] = 0.0,
     mask: bool = False,
 ):
     """
@@ -171,7 +171,7 @@ def low_frequency(
         Smaps = Smaps * convex_hull
     # Smooth out the sensitivity maps
     if np.sum(blurr_factor) > 0:
-        if isinstance(blurr_factor, float):
+        if isinstance(blurr_factor, (float, int)):
             blurr_factor = (blurr_factor,) * SOS.ndim
         Smaps = gaussian(np.abs(Smaps), sigma=(0,) + blurr_factor) * np.exp(
             1j * np.angle(Smaps)

diff --git a/src/mrinufft/operators/base.py b/src/mrinufft/operators/base.py
@@ -255,7 +255,7 @@ class FourierOperatorBase(ABC):
 
     interfaces: dict[str, tuple] = {}
     autograd_available = False
-    density_method = None
+    _density_method = None
     _grad_wrt_data = False
     _grad_wrt_traj = False
 
@@ -423,11 +423,12 @@ def compute_density(self, method=None):
             method = get_density(method)
         if not callable(method):
             raise ValueError(f"Unknown density method: {method}")
-        self.density_method = lambda samples, shape: method(
-            samples,
-            shape,
-            **kwargs,
-        )
+        if self._density_method is None:
+            self._density_method = lambda samples, shape: method(
+                samples,
+                shape,
+                **kwargs,
+            )
         self.density = method(self.samples, self.shape, **kwargs)
 
     def get_lipschitz_cst(self, max_iter=10, **kwargs):

diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py
@@ -278,6 +278,7 @@ def samples(self, samples):
             if typ == "grad" and not self._grad_wrt_traj:
                 continue
             self.raw_op._set_pts(typ, samples)
+        self.compute_density(self._density_method)
 
     @with_numpy_cupy
     @nvtx_mark()

diff --git a/src/mrinufft/operators/interfaces/finufft.py b/src/mrinufft/operators/interfaces/finufft.py
@@ -149,7 +149,7 @@ def samples(self, new_samples):
             if typ == "grad" and not self._grad_wrt_traj:
                 continue
             self.raw_op._set_pts(typ, new_samples)
-        self.compute_density(self.density_method)
+        self.compute_density(self._density_method)
 
     def _make_plan_grad(self, **kwargs):
         self.raw_op.grad_plan = Plan(

diff --git a/src/mrinufft/operators/interfaces/gpunufft.py b/src/mrinufft/operators/interfaces/gpunufft.py
@@ -542,7 +542,7 @@ def samples(self, samples):
         )
         # TODO: gpuNUFFT needs to sort the points twice in this case.
         # It could help to have access to directly dorted arrays from gpuNUFFT.
-        self.compute_density(self.density_method)
+        self.compute_density(self._density_method)
         self.raw_op.set_pts(
             self._samples,
             density=self.density,