Dispatch density methods in operators

.github/workflows/test-ci.yml

@@ -16,7 +16,7 @@ env:
 jobs:
   test-cpu:
     runs-on: ubuntu-latest
-    if: ${{ !contains(github.event.head_commit.message, "style")}}
+    if: ${{ !contains(github.event.head_commit.message, 'style')}}
     strategy:
       matrix:
         backend: [finufft, pynfft, pynufft, bart, sigpy]
@@ -90,7 +90,7 @@ jobs:
 
   test-gpu:
     runs-on: GPU
-    if: ${{ !contains(github.event.head_commit.message, "style")}}
+    if: ${{ !contains(github.event.head_commit.message, 'style')}}
     strategy:
       matrix:
         backend: [gpunufft, cufinufft]
@@ -107,49 +107,17 @@ jobs:
           source $RUNNER_WORKSPACE/venv/bin/activate
           pip install --upgrade pip wheel
           pip install -e mri-nufft[test]
-          pip install finufft
-
-      - name: Install Experimental cufinufft
-        if: ${{ matrix.backend == 'cufinufft' }}
-        shell: bash
-        run: |
-          cd $RUNNER_WORKSPACE
-          rm -rf finufft
-          git clone https://github.com/flatironinstitute/finufft
-          cd finufft && mkdir build && cd build
-
-          export PATH=/usr/local/cuda-11.8/bin:$PATH
-          export CUDA_PATH=/usr/local/cuda-11.8
-          export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
-          export CUDA_BIN_PATH=/usr/local/cuda-11.8
-
-          cmake -DFINUFFT_USE_CUDA=1 ../ && cmake --build . && cp libcufinufft.so ../python/cufinufft/.
-          # enter venv
-          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
+          pip install finufft
 
-      - name: Install gpuNUFFT
-        if: ${{ matrix.backend == 'gpunufft' }}
+      - name: Install backend
         shell: bash
         run: |
-          cd $RUNNER_WORKSPACE
-          rm -rf gpuNUFFT
-          git clone https://github.com/chaithyagr/gpuNUFFT
-          cd gpuNUFFT
-
-          export PATH=/usr/local/cuda-11.8/bin:$PATH
-          export CUDA_PATH=/usr/local/cuda-11.8
-          export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
-          export CUDA_BIN_PATH=/usr/local/cuda-11.8
-
           source $RUNNER_WORKSPACE/venv/bin/activate
-          python setup.py install
-          pip install cupy-cuda11x
+          export CUDA_BIN_PATH=/usr/local/cuda-11.8/
+          export PATH=/usr/local/cuda-11.8/bin/${PATH:+:${PATH}}
+          export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib/{LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+          pip install ${{ matrix.backend }}
 
       - name: Run Tests
         shell: bash
@@ -178,7 +146,7 @@ jobs:
 
   test-examples:
     runs-on: ubuntu-latest
-    needs: linter-check
+    if: ${{ !contains(github.event.head_commit.message, 'style')}}
 
     steps:
       - uses: actions/checkout@v3

examples/example_density.py

@@ -0,0 +1,150 @@
+# %%
+"""
+=============================
+Density Compensation Routines
+=============================
+
+Examples of differents density compensation methods.
+
+Density compensation depends on the sampling trajectory,and is apply before the
+adjoint operation to act as preconditioner, and should make the lipschitz constant
+of the operator roughly equal to 1.
+
+"""
+import brainweb_dl as bwdl
+
+import matplotlib.pyplot as plt
+import numpy as np
+from mrinufft import get_density, get_operator, check_backend
+from mrinufft.trajectories import initialize_2D_radial
+from mrinufft.trajectories.display import display_2D_trajectory
+
+
+# %%
+# Create sample data
+# ------------------
+
+mri_2D = bwdl.get_mri(4, "T1")[80, ...].astype(np.float32)
+
+print(mri_2D.shape)
+
+traj = initialize_2D_radial(192, 192).astype(np.float32)
+
+nufft = get_operator("finufft")(traj, mri_2D.shape, density=False)
+kspace = nufft.op(mri_2D)
+adjoint = nufft.adj_op(kspace)
+
+fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+axs[0].imshow(abs(mri_2D))
+display_2D_trajectory(traj, subfigure=axs[1])
+axs[2].imshow(abs(adjoint))
+
+# %%
+# As you can see, the radial sampling pattern as a strong concentration of sampling point in the center, resulting in a  low-frequency biased adjoint reconstruction.
+
+# %%
+# Geometry based methods
+# ======================
+#
+# Voronoi
+# -------
+#
+# Voronoi Parcellation attribute a weights to each k-space coordinate, inversely
+# proportional to its voronoi cell area.
+
+
+# .. warning::
+#    The current implementation of voronoi parcellation is CPU only, and is thus
+#    **very** slow in 3D ( > 1h).
+
+# %%
+voronoi_weights = get_density("voronoi", traj)
+
+nufft_voronoi = get_operator("finufft")(
+    traj, shape=mri_2D.shape, density=voronoi_weights
+)
+adjoint_voronoi = nufft_voronoi.adj_op(kspace)
+fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+axs[0].imshow(abs(mri_2D))
+axs[0].set_title("Ground Truth")
+axs[1].imshow(abs(adjoint))
+axs[1].set_title("no density compensation")
+axs[2].imshow(abs(adjoint_voronoi))
+axs[2].set_title("Voronoi density compensation")
+
+
+# %%
+# Cell Counting
+# -------------
+#
+# Cell Counting attributes weights based on the number of trajectory point lying in a same k-space nyquist voxel.
+# This can be viewed as an approximation to the voronoi neth
+
+# .. note::
+#    Cell counting is faster than voronoi (especially in 3D), but is less precise.
+
+# The size of the niquist voxel can be tweak by using the osf parameter. Typically as the NUFFT (and by default in MRI-NUFFT) is performed at an OSF of 2
+
+
+# %%
+cell_count_weights = get_density("cell_count", traj, shape=mri_2D.shape, osf=2.0)
+
+nufft_cell_count = get_operator("finufft")(
+    traj, shape=mri_2D.shape, density=cell_count_weights, upsampfac=2.0
+)
+adjoint_cell_count = nufft_cell_count.adj_op(kspace)
+fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+axs[0].imshow(abs(mri_2D))
+axs[0].set_title("Ground Truth")
+axs[1].imshow(abs(adjoint))
+axs[1].set_title("no density compensation")
+axs[2].imshow(abs(adjoint_cell_count))
+axs[2].set_title("cell_count density compensation")
+
+# %%
+# Manual Density Estimation
+# -------------------------
+#
+# For some analytical trajectory it is also possible to determine the density compensation vector directly.
+# In radial trajectory for instance, a sample's weight can be determined from its distance to the center.
+
+
+# %%
+flat_traj = traj.reshape(-1, 2)
+weights = np.sqrt(np.sum(flat_traj**2, axis=1))
+nufft = get_operator("finufft")(traj, shape=mri_2D.shape, density=weights)
+adjoint_manual = nufft.adj_op(kspace)
+fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+axs[0].imshow(abs(mri_2D))
+axs[0].set_title("Ground Truth")
+axs[1].imshow(abs(adjoint))
+axs[1].set_title("no density compensation")
+axs[2].imshow(abs(adjoint_manual))
+axs[2].set_title("manual density compensation")
+
+# %%
+# Operator-based method
+# =====================
+#
+# Pipe's Method
+# -------------
+# Pipe's method is an iterative scheme, that use the interpolation and spreading kernel operator for computing the density compensation.
+#
+# .. warning::
+#    If this method is widely used in the literature, there exists no convergence guarantees for it.
+
+# .. note::
+#    The Pipe method is currently only implemented for gpuNUFFT.
+
+# %%
+if check_backend("gpunufft"):
+    flat_traj = traj.reshape(-1, 2)
+    nufft = get_operator("gpunufft")(traj, shape=mri_2D.shape, density=False)
+    adjoint_manual = nufft.adj_op(kspace)
+    fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+    axs[0].imshow(abs(mri_2D))
+    axs[0].set_title("Ground Truth")
+    axs[1].imshow(abs(adjoint))
+    axs[1].set_title("no density compensation")
+    axs[2].imshow(abs(adjoint_manual))
+    axs[2].set_title("manual density compensation")

src/mrinufft/__init__.py

@@ -35,7 +35,7 @@
     display_3D_trajectory,
 )
 
-from .density import voronoi, cell_count, pipe
+from .density import voronoi, cell_count, pipe, get_density
 
 __all__ = [
     "get_operator",
@@ -65,6 +65,7 @@
     "voronoi",
     "cell_count",
     "pipe",
+    "get_density",
 ]
 
 

src/mrinufft/_utils.py

@@ -0,0 +1,82 @@
+"""General utility functions for MRI-NUFFT."""
+
+import warnings
+from collections import defaultdict
+from functools import wraps
+import numpy as np
+
+
+def proper_trajectory(trajectory, normalize="pi"):
+    """Normalize the trajectory to be used by NUFFT operators.
+
+    Parameters
+    ----------
+    trajectory: np.ndarray
+        The trajectory to normalize, it might be of shape (Nc, Ns, dim) of (Ns, dim)
+
+    normalize: str
+        if "pi" trajectory will be rescaled in [-pi, pi], if it was in [-0.5, 0.5]
+        if "unit" trajectory will be rescaled in [-0.5, 0.5] if it was not [-0.5, 0.5]
+
+    Returns
+    -------
+    new_traj: np.ndarray
+        The normalized trajectory of shape (Nc * Ns, dim) or (Ns, dim) in -pi, pi
+    """
+    # flatten to a list of point
+    try:
+        new_traj = np.asarray(trajectory).copy()
+    except Exception as e:
+        raise ValueError(
+            "trajectory should be array_like, with the last dimension being coordinates"
+        ) from e
+    new_traj = new_traj.reshape(-1, trajectory.shape[-1])
+
+    if normalize == "pi" and np.max(abs(new_traj)) - 1e-4 < 0.5:
+        warnings.warn(
+            "Samples will be rescaled to [-pi, pi), assuming they were in [-0.5, 0.5)"
+        )
+        new_traj *= 2 * np.pi
+    elif normalize == "unit" and np.max(abs(new_traj)) - 1e-4 > 0.5:
+        warnings.warn(
+            "Samples will be rescaled to [-0.5, 0.5), assuming they were in [-pi, pi)"
+        )
+        new_traj /= 2 * np.pi
+    if normalize == "unit" and np.max(new_traj) >= 0.5:
+        new_traj = (new_traj + 0.5) % 1 - 0.5
+    return new_traj
+
+
+class MethodRegister:
+    """
+    A Decorator to register methods of the same type in dictionnaries.
+
+    Parameters
+    ----------
+    name: str
+        The  register
+    """
+
+    registry = defaultdict(dict)
+
+    def __init__(self, register_name):
+        self.register_name = register_name
+
+    def __call__(self, method_name=None):
+        """Register the function in the registry."""
+
+        def decorator(func):
+            self.registry[self.register_name][method_name] = func
+
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return wrapper
+
+        if callable(method_name):
+            func = method_name
+            method_name = func.__name__
+            return decorator(func)
+        else:
+            return decorator

src/mrinufft/density/__init__.py

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

src/mrinufft/density/geometry_based.py

@@ -2,7 +2,7 @@
 import numpy as np
 from scipy.spatial import Voronoi
 
-from .utils import flat_traj, normalize_weights
+from .utils import flat_traj, normalize_weights, register_density
 
 
 def _vol3d(points):
@@ -46,6 +46,7 @@ def _vol2d(points):
     return abs(area) / 2.0
 
 
+@register_density
 @flat_traj
 def voronoi_unique(traj, *args, **kwargs):
     """Estimate  density compensation weight using voronoi parcellation.
@@ -86,6 +87,7 @@ def voronoi_unique(traj, *args, **kwargs):
     return wi
 
 
+@register_density
 @flat_traj
 def voronoi(traj, *args, **kwargs):
     """Estimate  density compensation weight using voronoi parcellation.
@@ -116,9 +118,10 @@ def voronoi(traj, *args, **kwargs):
         wi[i0] = wi[i0f] / np.sum(i0)
     else:
         wi = voronoi_unique(traj)
-    return normalize_weights(wi)
+    return 1 / normalize_weights(wi)
 
 
+@register_density
 @flat_traj
 def cell_count(traj, shape, osf=1.0):
     """