From c5cbf2aeb0cd6f76ac32a36e89e96d2f51c5e955 Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Wed, 13 Dec 2023 12:34:48 +0100
Subject: [PATCH 01/14] feat: refactor density compensation to a module.

---
 src/mrinufft/trajectories/density/__init__.py |  5 ++
 .../{density.py => density/geometry_based.py} | 49 +++++--------------
 .../trajectories/density/nufft_based.py       | 27 ++++++++++
 3 files changed, 44 insertions(+), 37 deletions(-)
 create mode 100644 src/mrinufft/trajectories/density/__init__.py
 rename src/mrinufft/trajectories/{density.py => density/geometry_based.py} (66%)
 create mode 100644 src/mrinufft/trajectories/density/nufft_based.py

diff --git a/src/mrinufft/trajectories/density/__init__.py b/src/mrinufft/trajectories/density/__init__.py
new file mode 100644
index 00000000..179fc1ed
--- /dev/null
+++ b/src/mrinufft/trajectories/density/__init__.py
@@ -0,0 +1,5 @@
+from .geometry_based import voronoi, voronoi_unique
+from .nufft_based import pipe
+
+
+__all__ = ["voronoi", "voronoi_unique", "pipe"]
diff --git a/src/mrinufft/trajectories/density.py b/src/mrinufft/trajectories/density/geometry_based.py
similarity index 66%
rename from src/mrinufft/trajectories/density.py
rename to src/mrinufft/trajectories/density/geometry_based.py
index 64188efe..1137e0d6 100644
--- a/src/mrinufft/trajectories/density.py
+++ b/src/mrinufft/trajectories/density/geometry_based.py
@@ -1,22 +1,15 @@
-"""
-Estimation of the density compensation array methods.
+"""Compute density compensation weights using Voronoi parcellation.
+
+
+Based on the MATLAB implementation in MIRT: https://github.com/JeffFessler/mirt/blob/main/mri/ir_mri_density_comp.m
 
-Those methods are agnostic of the NUFFT operator.
 """
 import numpy as np
 from scipy.spatial import Voronoi
 from mrinufft.operators import proper_trajectory
-from mrinufft.operators.interfaces.cufinufft import pipe as pipe_cufinufft
-from mrinufft.operators.interfaces.tfnufft import pipe as pipe_tfnufft
-from mrinufft.operators.interfaces.gpunufft import pipe as pipe_gpunufft
-
-
-def compute_tetrahedron_volume(A, B, C, D):
-    """Compute the volume of a tetrahedron."""
-    return np.abs(np.dot(np.cross(B - A, C - A), D - A)) / 6.0
 
 
-def vol3d(points):
+def _vol3d(points):
     """Compute the volume of a convex 3D polygon.
 
     Parameters
@@ -35,7 +28,7 @@ def vol3d(points):
     return np.sum(np.abs(np.dot(np.cross(B, C), A.T))) / 6.0
 
 
-def vol2d(points):
+def _vol2d(points):
     """Compute the area of a convex 2D polygon.
 
     Parameters
@@ -57,7 +50,7 @@ def vol2d(points):
     return abs(area) / 2.0
 
 
-def _voronoi(kspace):
+def voronoi_unique(kspace):
     """Estimate  density compensation weight using voronoi parcellation.
 
     This assume unicity of the point in the kspace.
@@ -74,9 +67,9 @@ def _voronoi(kspace):
     """
     M = kspace.shape[0]
     if kspace.shape[1] == 2:
-        vol = vol2d
+        vol = _vol2d
     else:
-        vol = vol3d
+        vol = _vol3d
     wi = np.zeros(M)
     v = Voronoi(kspace)
     for mm in range(M):
@@ -112,32 +105,14 @@ def voronoi(kspace):
         i0f = i0f[0]
         i0[i0f] = False
         wi = np.zeros(len(kspace))
-        wi[~i0] = _voronoi(kspace[~i0])
+        wi[~i0] = voronoi_unique(kspace[~i0])
         i0[i0f] = True
         wi[i0] = wi[i0f] / np.sum(i0)
     else:
-        wi = _voronoi(kspace)
+        wi = voronoi_unique(kspace)
     wi /= np.sum(wi)
     return wi
 
 
-def pipe(kspace_traj, grid_size, backend="cufinufft", **kwargs):
-    """Compute the density compensation weights using the pipe method.
 
-    Parameters
-    ----------
-    kspace_traj: array_like
-        array of shape (M, 2) or (M, 3) containing the coordinates of the points.
-    grid_size: array_like
-        array of shape (2,) or (3,) containing the size of the grid.
-    backend: str
-        backend to use for the computation. Either "cufinufft" or "tensorflow".
-    """
-    if backend == "cufinufft":
-        return pipe_cufinufft(kspace_traj, grid_size, **kwargs)
-    elif backend == "tensorflow":
-        return pipe_tfnufft(kspace_traj, grid_size, **kwargs)
-    elif backend == "gpunufft":
-        return pipe_gpunufft(kspace_traj, grid_size, **kwargs)
-    else:
-        raise ValueError("backend not supported")
+def
diff --git a/src/mrinufft/trajectories/density/nufft_based.py b/src/mrinufft/trajectories/density/nufft_based.py
new file mode 100644
index 00000000..b72214bf
--- /dev/null
+++ b/src/mrinufft/trajectories/density/nufft_based.py
@@ -0,0 +1,27 @@
+"""Density compensation weights using the NUFFT-based methods."""
+
+from .cufinufft import pipe_cufinufft
+from .tfnufft import pipe_tfnufft
+from .gpunufft import pipe_gpunufft
+
+
+def pipe(kspace_traj, grid_size, backend="cufinufft", **kwargs):
+    """Compute the density compensation weights using the pipe method.
+
+    Parameters
+    ----------
+    kspace_traj: array_like
+        array of shape (M, 2) or (M, 3) containing the coordinates of the points.
+    grid_size: array_like
+        array of shape (2,) or (3,) containing the size of the grid.
+    backend: str
+        backend to use for the computation. Either "cufinufft" or "tensorflow".
+    """
+    if backend == "cufinufft":
+        return pipe_cufinufft(kspace_traj, grid_size, **kwargs)
+    elif backend == "tensorflow":
+        return pipe_tfnufft(kspace_traj, grid_size, **kwargs)
+    elif backend == "gpunufft":
+        return pipe_gpunufft(kspace_traj, grid_size, **kwargs)
+    else:
+        raise ValueError("backend not supported")

From 67979ebcfeac6346a527f5c0c54778260c7063ff Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Thu, 14 Dec 2023 11:32:05 +0100
Subject: [PATCH 02/14] feat: add decorator to flatten trajectories.

---
 src/mrinufft/trajectories/density/utils.py | 23 ++++++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 src/mrinufft/trajectories/density/utils.py

diff --git a/src/mrinufft/trajectories/density/utils.py b/src/mrinufft/trajectories/density/utils.py
new file mode 100644
index 00000000..6ab6f31e
--- /dev/null
+++ b/src/mrinufft/trajectories/density/utils.py
@@ -0,0 +1,23 @@
+"""Utilities for density compensation."""
+from mrinufft.operators import proper_trajectory
+from functools import wraps
+
+
+def flat_traj(normalize="unit"):
+    """Decorate function to ensure that the trajectory is flatten before calling."""
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            args = list(args)
+            args[0] = proper_trajectory(args[0], normalize=normalize)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    if callable(normalize):  # call without argument
+        func = normalize
+        normalize = "unit"
+        return decorator(func)
+    else:
+        return decorator

From 5f324a9c8528577084464f7a4c38c2e91516595e Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Thu, 14 Dec 2023 11:32:55 +0100
Subject: [PATCH 03/14] feat: add ultra fast cell counting method.

---
 .../trajectories/density/geometry_based.py    | 57 ++++++++++++++++++-
 1 file changed, 56 insertions(+), 1 deletion(-)

diff --git a/src/mrinufft/trajectories/density/geometry_based.py b/src/mrinufft/trajectories/density/geometry_based.py
index 1137e0d6..0d327e4a 100644
--- a/src/mrinufft/trajectories/density/geometry_based.py
+++ b/src/mrinufft/trajectories/density/geometry_based.py
@@ -114,5 +114,60 @@ def voronoi(kspace):
     return wi
 
 
+@flat_traj
+def cell_count(traj, shape, osf=1.0):
+    """
+    Compute the number of points in each cell of the grid.
+
+    Parameters
+    ----------
+    traj: array_like
+        array of shape (M, 2) or (M, 3) containing the coordinates of the points.
+    shape: tuple
+        shape of the grid.
+    osf: float
+        oversampling factor for the grid. default 1
+
+    Returns
+    -------
+    weights: array_like
+        array of shape (M,) containing the density compensation weights.
 
-def
+    """
+    bins = [np.linspace(-0.5, 0.5, int(osf * s) + 1) for s in shape]
+
+    h, edges = np.histogramdd(traj, bins)
+    if len(shape) == 2:
+        hsum = [np.sum(h, axis=1).astype(int), np.sum(h, axis=0).astype(int)]
+    else:
+        hsum = [
+            np.sum(h, axis=(1, 2)).astype(int),
+            np.sum(h, axis=(0, 2)).astype(int),
+            np.sum(h, axis=(0, 1)).astype(int),
+        ]
+
+    # indices of ascending coordinate in each dimension.
+    locs_sorted = [np.argsort(traj[:, i]) for i in range(len(shape))]
+
+    weights = np.ones(len(traj))
+    set_xyz = [[], [], []]
+    for i in range(len(hsum)):
+        ind = 0
+        for binsize in hsum[i]:
+            s = set(locs_sorted[i][ind : ind + binsize])
+            if s:
+                set_xyz[i].append(s)
+            ind += binsize
+
+    for sx in set_xyz[0]:
+        for sy in set_xyz[1]:
+            sxy = sx.intersection(sy)
+            if not sxy:
+                continue
+            for sz in set_xyz[2]:
+                final = sxy.intersection(sz)
+                if final:
+                    weights[list(final)] = len(final)
+
+    weights /= np.sum(weights)
+    return 1 / weights

From 21c10828078879a819eeb0e8b318ef140deba741 Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Thu, 14 Dec 2023 11:33:28 +0100
Subject: [PATCH 04/14] refactor: use the flat_traj decorator.

---
 .../trajectories/density/geometry_based.py    | 29 ++++++++++---------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/mrinufft/trajectories/density/geometry_based.py b/src/mrinufft/trajectories/density/geometry_based.py
index 0d327e4a..cd880279 100644
--- a/src/mrinufft/trajectories/density/geometry_based.py
+++ b/src/mrinufft/trajectories/density/geometry_based.py
@@ -1,12 +1,12 @@
 """Compute density compensation weights using Voronoi parcellation.
 
-
 Based on the MATLAB implementation in MIRT: https://github.com/JeffFessler/mirt/blob/main/mri/ir_mri_density_comp.m
 
 """
 import numpy as np
 from scipy.spatial import Voronoi
-from mrinufft.operators import proper_trajectory
+
+from .utils import flat_traj
 
 
 def _vol3d(points):
@@ -50,7 +50,8 @@ def _vol2d(points):
     return abs(area) / 2.0
 
 
-def voronoi_unique(kspace):
+@flat_traj
+def voronoi_unique(traj):
     """Estimate  density compensation weight using voronoi parcellation.
 
     This assume unicity of the point in the kspace.
@@ -65,20 +66,20 @@ def voronoi_unique(kspace):
     wi: array_like
         array of shape (M,) containing the density compensation weights.
     """
-    M = kspace.shape[0]
-    if kspace.shape[1] == 2:
+    M = traj.shape[0]
+    if traj.shape[1] == 2:
         vol = _vol2d
     else:
         vol = _vol3d
     wi = np.zeros(M)
-    v = Voronoi(kspace)
+    v = Voronoi(traj)
     for mm in range(M):
         idx_vertices = v.regions[v.point_region[mm]]
         if np.all([i != -1 for i in idx_vertices]):
             wi[mm] = vol(v.vertices[idx_vertices])
     # For edge point (infinite voronoi cells) we extrapolate from neighbours
     # Initial implementation in Jeff Fessler's MIRT
-    rho = np.sum(kspace**2, axis=1)
+    rho = np.sum(traj**2, axis=1)
     igood = rho > 0.6 * np.max(rho)
     if len(igood) < 10:
         print("dubious extrapolation with", len(igood), "points")
@@ -87,29 +88,29 @@ def voronoi_unique(kspace):
     return wi
 
 
-def voronoi(kspace):
+@flat_traj
+def voronoi(traj):
     """Estimate  density compensation weight using voronoi parcellation.
 
     In case of multiple point in the center of kspace, the weight is split evenly.
 
     Parameters
     ----------
-    kspace: array_like
+    traj: array_like
         array of shape (M, 2) or (M, 3) containing the coordinates of the points.
     """
     # deduplication only works for the 0,0 coordinate !!
-    kspace = proper_trajectory(kspace)
-    i0 = np.sum(np.abs(kspace), axis=1) == 0
+    i0 = np.sum(np.abs(traj), axis=1) == 0
     if np.any(i0):
         i0f = np.where(i0)
         i0f = i0f[0]
         i0[i0f] = False
-        wi = np.zeros(len(kspace))
-        wi[~i0] = voronoi_unique(kspace[~i0])
+        wi = np.zeros(len(traj))
+        wi[~i0] = voronoi_unique(traj[~i0])
         i0[i0f] = True
         wi[i0] = wi[i0f] / np.sum(i0)
     else:
-        wi = voronoi_unique(kspace)
+        wi = voronoi_unique(traj)
     wi /= np.sum(wi)
     return wi
 

From 7481b63a0af68b8c649cc4878b7b6716b33f9125 Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Thu, 14 Dec 2023 12:12:35 +0100
Subject: [PATCH 05/14] refactor: make pipe a classmethod.

---
 .../operators/interfaces/cufinufft.py         | 78 +++++++++----------
 src/mrinufft/operators/interfaces/gpunufft.py | 54 ++++++-------
 src/mrinufft/operators/interfaces/tfnufft.py  | 36 ++++-----
 .../trajectories/density/nufft_based.py       | 16 ++--
 4 files changed, 89 insertions(+), 95 deletions(-)

diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py
index 8d804b70..3a86f29a 100644
--- a/src/mrinufft/operators/interfaces/cufinufft.py
+++ b/src/mrinufft/operators/interfaces/cufinufft.py
@@ -264,7 +264,7 @@ def __init__(
 
         # density compensation support
         if density is True:
-            self.density = pipe(self.samples, shape)
+            self.density = self.pipe(self.samples, shape)
         elif is_host_array(density) or is_cuda_array(density):
             self.density = cp.array(density)
         else:
@@ -800,6 +800,44 @@ def __repr__(self):
             ")"
         )
 
+    @classmethod
+    def pipe(kspace, grid_shape, num_iter=10, tol=2e-7):
+        """Estimate density compensation weight using the Pipe method.
+
+        Parameters
+        ----------
+        kspace: array_like
+            array of shape (M, 2) or (M, 3) containing the coordinates of the points.
+        grid_shape: tuple
+            shape of the image grid.
+        num_iter: int, optional
+            number of iterations.
+        tol: float, optional
+            tolerance of the density estimation.
+
+        Returns
+        -------
+        density: array_like
+            array of shape (M,) containing the density compensation weights.
+        """
+        if not CUFINUFFT_AVAILABLE:
+            raise ImportError(
+                "cuFINUFFT library not available to do Pipe density estimation"
+            )
+        import cupy as cp
+
+        kspace = proper_trajectory(kspace, normalize="pi").astype(np.float32)
+        if is_host_array(kspace):
+            kspace = cp.array(kspace, order="F")
+        image = cp.empty(grid_shape, dtype=np.complex64)
+        density = cp.ones(kspace.shape[0], dtype=np.complex64)
+        update = cp.empty_like(density)
+        for _ in range(num_iter):
+            _do_spread_interp(kspace, density, image, type=1, tol=tol)
+            _do_spread_interp(kspace, update, image, type=2, tol=tol)
+            update_density(density, update)
+        return density.real
+
 
 def _get_samples_ptr(samples):
     x, y, z = None, None, None
@@ -831,41 +869,3 @@ def _do_spread_interp(samples, c, f, tol=1e-4, type=1):
         opts,
         tol,
     )
-
-
-def pipe(kspace, grid_shape, num_iter=10, tol=2e-7):
-    """Estimate density compensation weight using the Pipe method.
-
-    Parameters
-    ----------
-    kspace: array_like
-        array of shape (M, 2) or (M, 3) containing the coordinates of the points.
-    grid_shape: tuple
-        shape of the image grid.
-    num_iter: int, optional
-        number of iterations.
-    tol: float, optional
-        tolerance of the density estimation.
-
-    Returns
-    -------
-    density: array_like
-        array of shape (M,) containing the density compensation weights.
-    """
-    if not CUFINUFFT_AVAILABLE:
-        raise ImportError(
-            "cuFINUFFT library not available to do Pipe density estimation"
-        )
-    import cupy as cp
-
-    kspace = proper_trajectory(kspace, normalize="pi").astype(np.float32)
-    if is_host_array(kspace):
-        kspace = cp.array(kspace, order="F")
-    image = cp.empty(grid_shape, dtype=np.complex64)
-    density = cp.ones(kspace.shape[0], dtype=np.complex64)
-    update = cp.empty_like(density)
-    for _ in range(num_iter):
-        _do_spread_interp(kspace, density, image, type=1, tol=tol)
-        _do_spread_interp(kspace, update, image, type=2, tol=tol)
-        update_density(density, update)
-    return density.real
diff --git a/src/mrinufft/operators/interfaces/gpunufft.py b/src/mrinufft/operators/interfaces/gpunufft.py
index f1cee6ab..dac2d60a 100644
--- a/src/mrinufft/operators/interfaces/gpunufft.py
+++ b/src/mrinufft/operators/interfaces/gpunufft.py
@@ -228,7 +228,7 @@ def __init__(
         self.n_coils = n_coils
         self.smaps = smaps
         if density is True:
-            self.density = pipe(self.samples, shape)
+            self.density = self.pipe(self.samples, shape)
         elif isinstance(density, np.ndarray):
             self.density = density
         else:
@@ -279,31 +279,31 @@ def uses_sense(self):
         """Return True if the Fourier Operator uses the SENSE method."""
         return self.impl.uses_sense
 
+    @classmethod
+    def pipe(cls, kspace_loc, volume_shape, num_iterations=10):
+        """Compute the density compensation weights for a given set of kspace locations.
 
-def pipe(kspace_loc, volume_shape, num_iterations=10):
-    """Compute the density compensation weights for a given set of kspace locations.
-
-    Parameters
-    ----------
-    kspace_loc: np.ndarray
-        the kspace locations
-    volume_shape: np.ndarray
-        the volume shape
-    num_iterations: int default 10
-        the number of iterations for density estimation
-    """
-    if GPUNUFFT_AVAILABLE is False:
-        raise ValueError(
-            "gpuNUFFT is not available, cannot " "estimate the density compensation"
-        )
-    grid_op = MRIGpuNUFFT(
-        samples=kspace_loc,
-        shape=volume_shape,
-        osf=1,
-    )
-    density_comp = np.ones(kspace_loc.shape[0])
-    for _ in range(num_iterations):
-        density_comp = density_comp / np.abs(
-            grid_op.op(grid_op.adj_op(density_comp, True), True)
+        Parameters
+        ----------
+        kspace_loc: np.ndarray
+            the kspace locations
+        volume_shape: np.ndarray
+            the volume shape
+        num_iterations: int default 10
+            the number of iterations for density estimation
+        """
+        if GPUNUFFT_AVAILABLE is False:
+            raise ValueError(
+                "gpuNUFFT is not available, cannot " "estimate the density compensation"
+            )
+        grid_op = MRIGpuNUFFT(
+            samples=kspace_loc,
+            shape=volume_shape,
+            osf=1,
         )
-    return density_comp
+        density_comp = np.ones(kspace_loc.shape[0])
+        for _ in range(num_iterations):
+            density_comp = density_comp / np.abs(
+                grid_op.op(grid_op.adj_op(density_comp, True), True)
+            )
+        return density_comp
diff --git a/src/mrinufft/operators/interfaces/tfnufft.py b/src/mrinufft/operators/interfaces/tfnufft.py
index 5252c7ed..06fb3088 100644
--- a/src/mrinufft/operators/interfaces/tfnufft.py
+++ b/src/mrinufft/operators/interfaces/tfnufft.py
@@ -135,23 +135,23 @@ def data_consistency(self, data, obs_data):
         """
         return self.adj_op(self.op(data) - obs_data)
 
+    @classmethod
+    def pipe(samples, shape, n_iter=15):
+        """Estimate the density compensation using the pipe method.
 
-def pipe(samples, shape, n_iter=15):
-    """Estimate the density compensation using the pipe method.
-
-    Parameters
-    ----------
-    samples: Tensor
-        The samples location of shape ``Nsamples x N_dimensions``.
-        It should be C-contiguous.
-    shape: tuple
-        Shape of the image space.
-    n_iter: int
-        Number of iterations.
+        Parameters
+        ----------
+        samples: Tensor
+            The samples location of shape ``Nsamples x N_dimensions``.
+            It should be C-contiguous.
+        shape: tuple
+            Shape of the image space.
+        n_iter: int
+            Number of iterations.
 
-    Returns
-    -------
-    Tensor
-        The estimated density compensation.
-    """
-    return tfmri.estimate_density(samples, shape, method="pipe", max_iter=n_iter)
+        Returns
+        -------
+        Tensor
+            The estimated density compensation.
+        """
+        return tfmri.estimate_density(samples, shape, method="pipe", max_iter=n_iter)
diff --git a/src/mrinufft/trajectories/density/nufft_based.py b/src/mrinufft/trajectories/density/nufft_based.py
index b72214bf..cb48684f 100644
--- a/src/mrinufft/trajectories/density/nufft_based.py
+++ b/src/mrinufft/trajectories/density/nufft_based.py
@@ -1,8 +1,6 @@
 """Density compensation weights using the NUFFT-based methods."""
 
-from .cufinufft import pipe_cufinufft
-from .tfnufft import pipe_tfnufft
-from .gpunufft import pipe_gpunufft
+from mrinufft import get_operator
 
 
 def pipe(kspace_traj, grid_size, backend="cufinufft", **kwargs):
@@ -17,11 +15,7 @@ def pipe(kspace_traj, grid_size, backend="cufinufft", **kwargs):
     backend: str
         backend to use for the computation. Either "cufinufft" or "tensorflow".
     """
-    if backend == "cufinufft":
-        return pipe_cufinufft(kspace_traj, grid_size, **kwargs)
-    elif backend == "tensorflow":
-        return pipe_tfnufft(kspace_traj, grid_size, **kwargs)
-    elif backend == "gpunufft":
-        return pipe_gpunufft(kspace_traj, grid_size, **kwargs)
-    else:
-        raise ValueError("backend not supported")
+    nufft_class = get_operator(backend)
+    if hasattr(nufft_class, "pipe"):
+        return nufft_class.pipe(kspace_traj, grid_size, **kwargs)
+    raise ValueError("backend does not have pipe iterations method.")

From ab99858dcfa6e0c0bc6d3f33123c6f322fafd35d Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Thu, 14 Dec 2023 12:13:11 +0100
Subject: [PATCH 06/14] fix: special case for 2d trajectories.

---
 src/mrinufft/trajectories/density/__init__.py       | 4 ++--
 src/mrinufft/trajectories/density/geometry_based.py | 9 ++++++---
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/mrinufft/trajectories/density/__init__.py b/src/mrinufft/trajectories/density/__init__.py
index 179fc1ed..6b25b125 100644
--- a/src/mrinufft/trajectories/density/__init__.py
+++ b/src/mrinufft/trajectories/density/__init__.py
@@ -1,5 +1,5 @@
-from .geometry_based import voronoi, voronoi_unique
+from .geometry_based import voronoi, voronoi_unique, cell_count
 from .nufft_based import pipe
 
 
-__all__ = ["voronoi", "voronoi_unique", "pipe"]
+__all__ = ["voronoi", "voronoi_unique", "pipe", "cell_count"]
diff --git a/src/mrinufft/trajectories/density/geometry_based.py b/src/mrinufft/trajectories/density/geometry_based.py
index cd880279..2df503b1 100644
--- a/src/mrinufft/trajectories/density/geometry_based.py
+++ b/src/mrinufft/trajectories/density/geometry_based.py
@@ -165,10 +165,13 @@ def cell_count(traj, shape, osf=1.0):
             sxy = sx.intersection(sy)
             if not sxy:
                 continue
+            if len(shape) == 2:
+                weights[list(sxy)] = len(sxy)
+                continue
             for sz in set_xyz[2]:
-                final = sxy.intersection(sz)
-                if final:
-                    weights[list(final)] = len(final)
+                sxyz = sxy.intersection(sz)
+                if sxyz:
+                    weights[list(sxyz)] = len(sxyz)
 
     weights /= np.sum(weights)
     return 1 / weights

From 9d4e060d15a679d07d75f167e9e6cf72b6f3c113 Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Fri, 15 Dec 2023 10:03:58 +0100
Subject: [PATCH 07/14] feat: enforce common api for density compensation
 functions.

---
 src/mrinufft/trajectories/density/nufft_based.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/mrinufft/trajectories/density/nufft_based.py b/src/mrinufft/trajectories/density/nufft_based.py
index cb48684f..6e41ca43 100644
--- a/src/mrinufft/trajectories/density/nufft_based.py
+++ b/src/mrinufft/trajectories/density/nufft_based.py
@@ -1,21 +1,23 @@
 """Density compensation weights using the NUFFT-based methods."""
 
 from mrinufft import get_operator
+from .utils import flat_traj
 
 
-def pipe(kspace_traj, grid_size, backend="cufinufft", **kwargs):
+@flat_traj
+def pipe(traj, shape, backend="cufinufft", **kwargs):
     """Compute the density compensation weights using the pipe method.
 
     Parameters
     ----------
-    kspace_traj: array_like
+    traj: array_like
         array of shape (M, 2) or (M, 3) containing the coordinates of the points.
-    grid_size: array_like
+    shape: array_like
         array of shape (2,) or (3,) containing the size of the grid.
     backend: str
         backend to use for the computation. Either "cufinufft" or "tensorflow".
     """
     nufft_class = get_operator(backend)
     if hasattr(nufft_class, "pipe"):
-        return nufft_class.pipe(kspace_traj, grid_size, **kwargs)
+        return nufft_class.pipe(traj, shape, **kwargs)
     raise ValueError("backend does not have pipe iterations method.")

From 6ab996eec535d0a8b41ef247c3634923f2be5790 Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Fri, 15 Dec 2023 10:40:45 +0100
Subject: [PATCH 08/14] feat: move density module up.

---
 src/mrinufft/__init__.py                      |  5 ++++
 .../{trajectories => }/density/__init__.py    |  0
 .../density/geometry_based.py                 | 27 ++++++++++---------
 .../{trajectories => }/density/nufft_based.py |  0
 .../{trajectories => }/density/utils.py       | 14 +++++++++-
 5 files changed, 33 insertions(+), 13 deletions(-)
 rename src/mrinufft/{trajectories => }/density/__init__.py (100%)
 rename src/mrinufft/{trajectories => }/density/geometry_based.py (88%)
 rename src/mrinufft/{trajectories => }/density/nufft_based.py (100%)
 rename src/mrinufft/{trajectories => }/density/utils.py (71%)

diff --git a/src/mrinufft/__init__.py b/src/mrinufft/__init__.py
index 69503718..864c7b14 100644
--- a/src/mrinufft/__init__.py
+++ b/src/mrinufft/__init__.py
@@ -35,6 +35,8 @@
     display_3D_trajectory,
 )
 
+from .density import voronoi, cell_count, pipe
+
 __all__ = [
     "get_operator",
     "check_backend",
@@ -60,6 +62,9 @@
     "displayConfig",
     "display_2D_trajectory",
     "display_3D_trajectory",
+    "voronoi",
+    "cell_count",
+    "pipe",
 ]
 
 
diff --git a/src/mrinufft/trajectories/density/__init__.py b/src/mrinufft/density/__init__.py
similarity index 100%
rename from src/mrinufft/trajectories/density/__init__.py
rename to src/mrinufft/density/__init__.py
diff --git a/src/mrinufft/trajectories/density/geometry_based.py b/src/mrinufft/density/geometry_based.py
similarity index 88%
rename from src/mrinufft/trajectories/density/geometry_based.py
rename to src/mrinufft/density/geometry_based.py
index 2df503b1..dcc43999 100644
--- a/src/mrinufft/trajectories/density/geometry_based.py
+++ b/src/mrinufft/density/geometry_based.py
@@ -1,12 +1,8 @@
-"""Compute density compensation weights using Voronoi parcellation.
-
-Based on the MATLAB implementation in MIRT: https://github.com/JeffFessler/mirt/blob/main/mri/ir_mri_density_comp.m
-
-"""
+"""Compute density compensation weights using geometry-based methods."""
 import numpy as np
 from scipy.spatial import Voronoi
 
-from .utils import flat_traj
+from .utils import flat_traj, normalize_weights
 
 
 def _vol3d(points):
@@ -51,7 +47,7 @@ def _vol2d(points):
 
 
 @flat_traj
-def voronoi_unique(traj):
+def voronoi_unique(traj, *args, **kwargs):
     """Estimate  density compensation weight using voronoi parcellation.
 
     This assume unicity of the point in the kspace.
@@ -60,6 +56,8 @@ def voronoi_unique(traj):
     ----------
     kspace: array_like
         array of shape (M, 2) or (M, 3) containing the coordinates of the points.
+    *args, **kwargs:
+        Dummy arguments to be compatible with other methods.
 
     Returns
     -------
@@ -89,7 +87,7 @@ def voronoi_unique(traj):
 
 
 @flat_traj
-def voronoi(traj):
+def voronoi(traj, *args, **kwargs):
     """Estimate  density compensation weight using voronoi parcellation.
 
     In case of multiple point in the center of kspace, the weight is split evenly.
@@ -98,6 +96,13 @@ def voronoi(traj):
     ----------
     traj: array_like
         array of shape (M, 2) or (M, 3) containing the coordinates of the points.
+
+    *args, **kwargs:
+        Dummy arguments to be compatible with other methods.
+
+    References
+    ----------
+    Based on the MATLAB implementation in MIRT: https://github.com/JeffFessler/mirt/blob/main/mri/ir_mri_density_comp.m
     """
     # deduplication only works for the 0,0 coordinate !!
     i0 = np.sum(np.abs(traj), axis=1) == 0
@@ -111,8 +116,7 @@ def voronoi(traj):
         wi[i0] = wi[i0f] / np.sum(i0)
     else:
         wi = voronoi_unique(traj)
-    wi /= np.sum(wi)
-    return wi
+    return normalize_weights(wi)
 
 
 @flat_traj
@@ -173,5 +177,4 @@ def cell_count(traj, shape, osf=1.0):
                 if sxyz:
                     weights[list(sxyz)] = len(sxyz)
 
-    weights /= np.sum(weights)
-    return 1 / weights
+    return normalize_weights(weights)
diff --git a/src/mrinufft/trajectories/density/nufft_based.py b/src/mrinufft/density/nufft_based.py
similarity index 100%
rename from src/mrinufft/trajectories/density/nufft_based.py
rename to src/mrinufft/density/nufft_based.py
diff --git a/src/mrinufft/trajectories/density/utils.py b/src/mrinufft/density/utils.py
similarity index 71%
rename from src/mrinufft/trajectories/density/utils.py
rename to src/mrinufft/density/utils.py
index 6ab6f31e..5559410d 100644
--- a/src/mrinufft/trajectories/density/utils.py
+++ b/src/mrinufft/density/utils.py
@@ -1,7 +1,10 @@
 """Utilities for density compensation."""
-from mrinufft.operators import proper_trajectory
 from functools import wraps
 
+import numpy as np
+
+from mrinufft.operators import proper_trajectory
+
 
 def flat_traj(normalize="unit"):
     """Decorate function to ensure that the trajectory is flatten before calling."""
@@ -21,3 +24,12 @@ def wrapper(*args, **kwargs):
         return decorator(func)
     else:
         return decorator
+
+
+def normalize_weights(weights):
+    """Normalize samples weights to reflect their importance.
+
+    Higher weights have lower importance.
+    """
+    inv_weights = np.sum(weights) / weights
+    return inv_weights / (np.sum(inv_weights))

From c4d336979c81e534d5d9d651aff7c92c3aaee815 Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Fri, 15 Dec 2023 10:44:31 +0100
Subject: [PATCH 09/14] feat(density): add test.

---
 tests/test_density.py | 71 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)
 create mode 100644 tests/test_density.py

diff --git a/tests/test_density.py b/tests/test_density.py
new file mode 100644
index 00000000..a0fc72cc
--- /dev/null
+++ b/tests/test_density.py
@@ -0,0 +1,71 @@
+"""Test the density compensation methods."""
+
+
+import numpy as np
+import numpy.testing as npt
+from pytest_cases import fixture, parametrize, parametrize_with_cases
+
+from case_trajectories import CasesTrajectories
+from helpers import assert_correlate
+from mrinufft.density import cell_count, voronoi
+from mrinufft.density.utils import normalize_weights
+from mrinufft.operators import proper_trajectory
+
+
+def slow_cell_count2D(traj, shape, osf):
+    """Perform the cell count but it is slow."""
+    traj = proper_trajectory(traj, normalize="unit")
+    bins = [np.linspace(-0.5, 0.5, int(osf * s) + 1) for s in shape]
+
+    h, edges = np.histogramdd(
+        traj,
+        bins,
+    )
+
+    weights = np.ones(len(traj))
+
+    bx = bins[0]
+    by = bins[1]
+    for i, (bxmin, bxmax) in enumerate(zip(bx[:-1], bx[1:])):
+        for j, (bymin, bymax) in enumerate(zip(by[:-1], by[1:])):
+            weights[
+                (bxmin <= traj[:, 0])
+                & (traj[:, 0] <= bxmax)
+                & (bymin <= traj[:, 1])
+                & (traj[:, 1] <= bymax)
+            ] = h[i, j]
+
+    return normalize_weights(weights)
+
+
+@fixture(scope="module")
+def radial_distance():
+    """Compute the radial distance of a trajectory."""
+    traj, shape = CasesTrajectories().case_radial2D()
+
+    proper_traj = proper_trajectory(traj, normalize="unit")
+    weights = 2 * np.pi * np.sqrt(proper_traj[:, 0] ** 2 + proper_traj[:, 1] ** 2)
+
+    return normalize_weights(weights)
+
+
+@parametrize("osf", [1, 1.25, 2])
+@parametrize_with_cases("traj, shape", cases=[CasesTrajectories.case_radial2D])
+def test_cell_count2D(traj, shape, osf):
+    """Test the cell count method."""
+    count_ref = slow_cell_count2D(traj, shape, osf)
+    count_real = cell_count(traj, shape, osf)
+    npt.assert_allclose(count_real, count_ref, atol=1e-5)
+
+
+@parametrize_with_cases("traj, shape", cases=[CasesTrajectories.case_radial2D])
+def test_voronoi(traj, shape, radial_distance):
+    """Test the voronoi method."""
+    result = voronoi(traj)
+
+    assert_correlate(result, radial_distance, slope=2 * np.pi)
+
+
+def test_pipe():
+    """Test the pipe method."""
+    pass

From 47b1a3966236749abf3fd29134b737982f97d7bb Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Fri, 15 Dec 2023 10:44:48 +0100
Subject: [PATCH 10/14] feat(cufinufft)!: drop pipe implementation and follow
 upstream.

---
 .github/workflows/test-ci.yml                 |   2 +-
 .../operators/interfaces/cufinufft.py         | 121 ------------------
 2 files changed, 1 insertion(+), 122 deletions(-)

diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml
index 0cd91e4d..485bae3c 100644
--- a/.github/workflows/test-ci.yml
+++ b/.github/workflows/test-ci.yml
@@ -139,7 +139,7 @@ jobs:
         run: |
           cd $RUNNER_WORKSPACE
           rm -rf finufft
-          git clone https://github.com/chaithyagr/finufft --branch chaithyagr/issue306
+          git clone https://github.com/flatiron/finufft
           cd finufft && mkdir build && cd build
 
           export PATH=/usr/local/cuda-11.8/bin:$PATH
diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py
index 3a86f29a..c9c3488a 100644
--- a/src/mrinufft/operators/interfaces/cufinufft.py
+++ b/src/mrinufft/operators/interfaces/cufinufft.py
@@ -19,7 +19,6 @@
 try:
     import cupy as cp
     from cufinufft._plan import Plan
-    from cufinufft._cufinufft import _spread_interpf, NufftOpts, _default_opts
 except ImportError:
     CUFINUFFT_AVAILABLE = False
 
@@ -42,56 +41,6 @@ def _error_check(ier, msg):
         raise RuntimeError(msg)
 
 
-def repr_opts(self):
-    """Get the value of the struct, like a dict."""
-    ret = "Struct(\n"
-    for fieldname, _ in self._fields_:
-        ret += f"{fieldname}: {getattr(self, fieldname)},\n"
-    ret += ")"
-    return ret
-
-
-def str_opts(self):
-    """Get the value of the struct, with their meaning."""
-    ret = "Struct(\n"
-    for fieldname, _ in self._fields_:
-        ret += f"{fieldname}: {getattr(self, fieldname)}"
-        decode = OPTS_FIELD_DECODE.get(fieldname)
-        if decode:
-            ret += f" [{decode[getattr(self, fieldname)]}]"
-        ret += "\n"
-    ret += ")"
-    return ret
-
-
-if CUFINUFFT_AVAILABLE:
-    NufftOpts.__repr__ = lambda self: repr_opts(self)
-    NufftOpts.__str__ = lambda self: str_opts(self)
-
-
-def get_default_opts(nufft_type, dim):
-    """
-    Generate a cufinufft opt struct of the dtype coresponding to plan.
-
-    Parameters
-    ----------
-    finufft_type: int
-        Finufft Type (`1` or `2`)
-    dim: int
-        Number of dimension (1, 2, 3)
-
-    Returns
-    -------
-    nufft_opts structure.
-    """
-    nufft_opts = NufftOpts()
-
-    ier = _default_opts(nufft_type, dim, nufft_opts)
-    _error_check(ier, "Configuration not yet implemented.")
-
-    return nufft_opts
-
-
 class RawCufinufftPlan:
     """Light wrapper around the guru interface of finufft."""
 
@@ -799,73 +748,3 @@ def __repr__(self):
             f"  eps:{self.raw_op.eps:.0e}\n"
             ")"
         )
-
-    @classmethod
-    def pipe(kspace, grid_shape, num_iter=10, tol=2e-7):
-        """Estimate density compensation weight using the Pipe method.
-
-        Parameters
-        ----------
-        kspace: array_like
-            array of shape (M, 2) or (M, 3) containing the coordinates of the points.
-        grid_shape: tuple
-            shape of the image grid.
-        num_iter: int, optional
-            number of iterations.
-        tol: float, optional
-            tolerance of the density estimation.
-
-        Returns
-        -------
-        density: array_like
-            array of shape (M,) containing the density compensation weights.
-        """
-        if not CUFINUFFT_AVAILABLE:
-            raise ImportError(
-                "cuFINUFFT library not available to do Pipe density estimation"
-            )
-        import cupy as cp
-
-        kspace = proper_trajectory(kspace, normalize="pi").astype(np.float32)
-        if is_host_array(kspace):
-            kspace = cp.array(kspace, order="F")
-        image = cp.empty(grid_shape, dtype=np.complex64)
-        density = cp.ones(kspace.shape[0], dtype=np.complex64)
-        update = cp.empty_like(density)
-        for _ in range(num_iter):
-            _do_spread_interp(kspace, density, image, type=1, tol=tol)
-            _do_spread_interp(kspace, update, image, type=2, tol=tol)
-            update_density(density, update)
-        return density.real
-
-
-def _get_samples_ptr(samples):
-    x, y, z = None, None, None
-    x = cp.ascontiguousarray(samples[:, 0])
-    y = cp.ascontiguousarray(samples[:, 1])
-    fpts_axes = [get_ptr(y), get_ptr(x), None]
-    if samples.shape[1] == 3:
-        z = cp.ascontiguousarray(samples[:, 2])
-        fpts_axes.insert(0, get_ptr(z))
-    return fpts_axes[:3], x.size
-
-
-def _convert_shape_to_3D(shape, dim):
-    return shape[::-1] + (1,) * (3 - dim)
-
-
-def _do_spread_interp(samples, c, f, tol=1e-4, type=1):
-    fpts_axes, n_samples = _get_samples_ptr(samples)
-    shape = _convert_shape_to_3D(f.shape, samples.shape[-1])
-    opts = get_default_opts(type, samples.shape[-1])
-    _spread_interpf(
-        type,
-        samples.shape[-1],
-        *shape,
-        n_samples,
-        *fpts_axes,
-        get_ptr(c),
-        get_ptr(f),
-        opts,
-        tol,
-    )

From 1a794dde576d34aa56d05dca0f2a2d21ce8e2928 Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Fri, 15 Dec 2023 11:02:49 +0100
Subject: [PATCH 11/14] fix: update module level for density.

---
 README.rst                  | 2 +-
 examples/example_readme.py  | 2 +-
 examples/example_stacked.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/README.rst b/README.rst
index cb233058..195824d7 100644
--- a/README.rst
+++ b/README.rst
@@ -34,7 +34,7 @@ Usage
       import numpy as np
       import mrinufft
       from mrinufft.trajectories import display
-      from mrinufft.trajectories.density import voronoi
+      from mrinufft.density import voronoi
 
       # Create a 2D Radial trajectory for demo
       samples_loc = mrinufft.initialize_2D_radial(Nc=100, Ns=500)
diff --git a/examples/example_readme.py b/examples/example_readme.py
index e4b32d25..1b9b58d8 100644
--- a/examples/example_readme.py
+++ b/examples/example_readme.py
@@ -10,7 +10,7 @@
 import numpy as np
 import mrinufft
 from mrinufft.trajectories import display
-from mrinufft.trajectories.density import voronoi
+from mrinufft.density import voronoi
 
 # Create a 2D Radial trajectory for demo
 samples_loc = mrinufft.initialize_2D_radial(Nc=100, Ns=500)
diff --git a/examples/example_stacked.py b/examples/example_stacked.py
index 19399a81..d7f5f746 100644
--- a/examples/example_stacked.py
+++ b/examples/example_stacked.py
@@ -37,7 +37,7 @@
 # ----------------------------
 
 from mrinufft import initialize_2D_spiral
-from mrinufft.trajectories.density import voronoi
+from mrinufft.density import voronoi
 
 samples = initialize_2D_spiral(Nc=16, Ns=500, nb_revolutions=10)
 density = voronoi(samples)

From 0eb2293a65925033b44065f43cb35146522014dd Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Fri, 15 Dec 2023 11:05:18 +0100
Subject: [PATCH 12/14] fix: cleanup

---
 .github/workflows/test-ci.yml                  | 2 +-
 src/mrinufft/density/__init__.py               | 1 +
 src/mrinufft/operators/interfaces/cufinufft.py | 1 -
 3 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml
index 485bae3c..f0c24fd2 100644
--- a/.github/workflows/test-ci.yml
+++ b/.github/workflows/test-ci.yml
@@ -139,7 +139,7 @@ jobs:
         run: |
           cd $RUNNER_WORKSPACE
           rm -rf finufft
-          git clone https://github.com/flatiron/finufft
+          git clone https://github.com/flatironinstitute/finufft
           cd finufft && mkdir build && cd build
 
           export PATH=/usr/local/cuda-11.8/bin:$PATH
diff --git a/src/mrinufft/density/__init__.py b/src/mrinufft/density/__init__.py
index 6b25b125..2a6ec4d0 100644
--- a/src/mrinufft/density/__init__.py
+++ b/src/mrinufft/density/__init__.py
@@ -1,3 +1,4 @@
+"""Density compensation methods."""
 from .geometry_based import voronoi, voronoi_unique, cell_count
 from .nufft_based import pipe
 
diff --git a/src/mrinufft/operators/interfaces/cufinufft.py b/src/mrinufft/operators/interfaces/cufinufft.py
index c9c3488a..60cf1cec 100644
--- a/src/mrinufft/operators/interfaces/cufinufft.py
+++ b/src/mrinufft/operators/interfaces/cufinufft.py
@@ -13,7 +13,6 @@
     pin_memory,
     sizeof_fmt,
 )
-from ._cupy_kernels import update_density
 
 CUFINUFFT_AVAILABLE = CUPY_AVAILABLE
 try:

From 784a45c77023fa4b90c5c664b6ee9c09912eb97f Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Fri, 15 Dec 2023 13:03:57 +0100
Subject: [PATCH 13/14] feat(ci): Install most recent cufinufft

---
 .github/workflows/test-ci.yml | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml
index f0c24fd2..02eec549 100644
--- a/.github/workflows/test-ci.yml
+++ b/.github/workflows/test-ci.yml
@@ -119,6 +119,7 @@ jobs:
     strategy:
       matrix:
         backend: [gpunufft, cufinufft]
+
     steps:
       - uses: actions/checkout@v3
 
@@ -133,7 +134,7 @@ jobs:
           pip install -e mri-nufft[test]
           pip install finufft
 
-      - name: Install Cufinufft
+      - name: Install Experimental cufinufft
         if: ${{ matrix.backend == 'cufinufft' }}
         shell: bash
         run: |
@@ -152,10 +153,7 @@ jobs:
           source $RUNNER_WORKSPACE/venv/bin/activate
           pip install cupy-cuda11x
           cd $RUNNER_WORKSPACE/finufft/python/cufinufft
-          python setup.py develop
-          # FIXME: This is hardcoded
-          cp libcufinufft.so  cufinufftc.cpython-310-x86_64-linux-gnu.so
-          cd $RUNNER_WORKSPACE
+          python setup.py install
 
       - name: Install gpuNUFFT
         if: ${{ matrix.backend == 'gpunufft' }}

From 50848d89c6eea22e4fb82609f1da774f0b1c07d3 Mon Sep 17 00:00:00 2001
From: Pierre-antoine Comby <pierre-antoine.comby@crans.org>
Date: Fri, 15 Dec 2023 14:18:54 +0100
Subject: [PATCH 14/14] fix(ci): copy libcufinufft.so at the right place.

---
 .github/workflows/test-ci.yml | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test-ci.yml b/.github/workflows/test-ci.yml
index 02eec549..73839fcb 100644
--- a/.github/workflows/test-ci.yml
+++ b/.github/workflows/test-ci.yml
@@ -153,7 +153,10 @@ jobs:
           source $RUNNER_WORKSPACE/venv/bin/activate
           pip install cupy-cuda11x
           cd $RUNNER_WORKSPACE/finufft/python/cufinufft
-          python setup.py install
+          python setup.py develop
+          # FIXME: This is hardcoded
+          cp libcufinufft.so  cufinufftc.cpython-310-x86_64-linux-gnu.so
+          cd $RUNNER_WORKSPACE
 
       - name: Install gpuNUFFT
         if: ${{ matrix.backend == 'gpunufft' }}