Density compensation (mind-inria#66)

* feat: refactor density compensation to a module.

* feat: add decorator to flatten trajectories.

* feat: add ultra fast cell counting method.

* refactor: use the flat_traj decorator.

* refactor: make pipe a classmethod.

* fix: special case for 2d trajectories.

* feat: enforce common api for density compensation functions.

* feat: move density module up.

* feat(density): add test.

* feat(cufinufft)!: drop pipe implementation and follow upstream.

* fix: update module level for density.

* fix: cleanup

* feat(ci): Install most recent cufinufft

* fix(ci): copy libcufinufft.so at the right place.

.github/workflows/test-ci.yml

@@ -119,6 +119,7 @@ jobs:
     strategy:
       matrix:
         backend: [gpunufft, cufinufft]
+
     steps:
       - uses: actions/checkout@v3
 
@@ -133,13 +134,13 @@ jobs:
           pip install -e mri-nufft[test]
           pip install finufft
 
-      - name: Install Cufinufft
+      - name: Install Experimental cufinufft
         if: ${{ matrix.backend == 'cufinufft' }}
         shell: bash
         run: |
           cd $RUNNER_WORKSPACE
           rm -rf finufft
-          git clone https://github.com/chaithyagr/finufft --branch chaithyagr/issue306
+          git clone https://github.com/flatironinstitute/finufft
           cd finufft && mkdir build && cd build
 
           export PATH=/usr/local/cuda-11.8/bin:$PATH

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)

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)

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)

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",
 ]
 
 

src/mrinufft/density/__init__.py

@@ -0,0 +1,6 @@
+"""Density compensation methods."""
+from .geometry_based import voronoi, voronoi_unique, cell_count
+from .nufft_based import pipe
+
+
+__all__ = ["voronoi", "voronoi_unique", "pipe", "cell_count"]

src/mrinufft/trajectories/density.py → src/mrinufft/density/geometry_based.py

@@ -1,22 +1,11 @@
-"""
-Estimation of the density compensation array methods.
-
-Those methods are agnostic of the NUFFT operator.
-"""
+"""Compute density compensation weights using geometry-based methods."""
 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
+from .utils import flat_traj, normalize_weights
 
 
-def vol3d(points):
+def _vol3d(points):
     """Compute the volume of a convex 3D polygon.
 
     Parameters
@@ -35,7 +24,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 +46,8 @@ def vol2d(points):
     return abs(area) / 2.0
 
 
-def _voronoi(kspace):
+@flat_traj
+def voronoi_unique(traj, *args, **kwargs):
     """Estimate  density compensation weight using voronoi parcellation.
 
     This assume unicity of the point in the kspace.
@@ -66,26 +56,28 @@ def _voronoi(kspace):
     ----------
     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
     -------
     wi: array_like
         array of shape (M,) containing the density compensation weights.
     """
-    M = kspace.shape[0]
-    if kspace.shape[1] == 2:
-        vol = vol2d
+    M = traj.shape[0]
+    if traj.shape[1] == 2:
+        vol = _vol2d
     else:
-        vol = vol3d
+        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")
@@ -94,50 +86,95 @@ def _voronoi(kspace):
     return wi
 
 
-def voronoi(kspace):
+@flat_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.
 
     Parameters
     ----------
-    kspace: array_like
+    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 !!
-    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(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(kspace)
-    wi /= np.sum(wi)
-    return wi
+        wi = voronoi_unique(traj)
+    return normalize_weights(wi)
 
 
-def pipe(kspace_traj, grid_size, backend="cufinufft", **kwargs):
-    """Compute the density compensation weights using the pipe method.
+@flat_traj
+def cell_count(traj, shape, osf=1.0):
+    """
+    Compute the number of points in each cell of the grid.
 
     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
-        array of shape (2,) or (3,) containing the size of the grid.
-    backend: str
-        backend to use for the computation. Either "cufinufft" or "tensorflow".
+    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.
+
     """
-    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)
+    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:
-        raise ValueError("backend not supported")
+        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
+            if len(shape) == 2:
+                weights[list(sxy)] = len(sxy)
+                continue
+            for sz in set_xyz[2]:
+                sxyz = sxy.intersection(sz)
+                if sxyz:
+                    weights[list(sxyz)] = len(sxyz)
+
+    return normalize_weights(weights)

src/mrinufft/density/nufft_based.py

@@ -0,0 +1,23 @@
+"""Density compensation weights using the NUFFT-based methods."""
+
+from mrinufft import get_operator
+from .utils import flat_traj
+
+
+@flat_traj
+def pipe(traj, shape, backend="cufinufft", **kwargs):
+    """Compute the density compensation weights using the pipe method.
+
+    Parameters
+    ----------
+    traj: array_like
+        array of shape (M, 2) or (M, 3) containing the coordinates of the points.
+    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(traj, shape, **kwargs)
+    raise ValueError("backend does not have pipe iterations method.")

src/mrinufft/density/utils.py

@@ -0,0 +1,35 @@
+"""Utilities for density compensation."""
+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."""
+
+    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
+
+
+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))