Merge branch 'master' into pulseq

Author: paquiteau

examples/GPU/example_3d_trajectory_display.py

@@ -14,14 +14,22 @@
 import os
 from mrinufft.trajectories.display3D import get_gridded_trajectory
 import mrinufft.trajectories.trajectory3D as mtt
-from mrinufft.trajectories.utils import Gammas
+from mrinufft.trajectories.utils import Gammas, Acquisition
 import matplotlib.pyplot as plt
 import numpy as np
 
 
 BACKEND = os.environ.get("MRINUFFT_BACKEND", "gpunufft")
 
 
+# %%
+# Acquisition parameters
+# ======================
+# Here we use acquisition defaults for the  trajectory gridding.
+
+acq = Acquisition.default
+
+
 # %%
 # Helper function to Displaying 3D Gridded Trajectories
 # =====================================================
@@ -57,9 +65,8 @@ def create_grid(grid_type, trajectories, traj_params, **kwargs):
     for i, (name, traj) in enumerate(trajectories.items()):
         grid = get_gridded_trajectory(
             traj,
-            traj_params["img_size"],
+            acq,
             grid_type=grid_type,
-            traj_params=traj_params,
             backend=BACKEND,
             osf=2,
             **kwargs,

examples/operators/example_offresonance.py

@@ -85,10 +85,10 @@
 
 from mrinufft import initialize_2D_spiral
 from mrinufft.density import voronoi
-from mrinufft.trajectories.utils import DEFAULT_RASTER_TIME
+from mrinufft.trajectories.utils import Acquisition
 
 samples = initialize_2D_spiral(Nc=48, Ns=600, nb_revolutions=10)
-t_read = np.arange(samples.shape[1]) * DEFAULT_RASTER_TIME * 1e-3
+t_read = np.arange(samples.shape[1]) * Acquisition.default.raster_time
 t_read = np.repeat(t_read[None, ...], samples.shape[0], axis=0)
 density = voronoi(samples)
 

examples/trajectories/example_3D_trajectories.py

@@ -30,7 +30,6 @@
 
 # Internal
 import mrinufft as mn
-from mrinufft import display_2D_trajectory, display_3D_trajectory
 
 # %%
 # Script options

examples/trajectories/example_display_config.py

@@ -15,6 +15,7 @@
 
 from mrinufft import display_2D_trajectory, display_3D_trajectory, displayConfig
 from mrinufft.trajectories import conify, initialize_2D_spiral
+from mrinufft.trajectories.utils import Acquisition, Hardware
 
 # %%
 # Script options
@@ -118,3 +119,33 @@ def show_traj(traj, name, values, **kwargs):
     ["tab:blue", "tab:orange", "tab:red"],
     show_constraints=True,
 )
+
+# You can also change the values of gmax and smax in order to see how the constraint
+# violations change.
+#
+acqs = [
+    Acquisition(
+        fov=(0.256, 0.256, 0.256),
+        img_size=(256, 256, 256),
+        hardware=Hardware(gmax=0.04, smax=50),
+    ),  # limiting slew rate to 50 T/m/s
+    Acquisition(
+        fov=(0.256, 0.256, 0.256),
+        img_size=(256, 256, 256),
+        hardware=Hardware(gmax=0.04, smax=100),
+    ),  # limiting slew rate to 100 T/m/s
+    Acquisition(
+        fov=(0.256, 0.256, 0.256),
+        img_size=(256, 256, 256),
+        hardware=Hardware(gmax=0.04, smax=200),
+    ),
+]  # limiting slew rate to 200 T/m/s
+
+# you can use Acquisition as a Context Manager, like display config.
+# Or pass it to the display function as well.
+for acq in acqs:
+    with acq:
+        display_3D_trajectory(traj, show_constraints=True)
+
+    # equivalent to
+    # display_3D_trajectory(traj, show_constraints=True, acq=acq)

examples/trajectories/utils.py

@@ -16,15 +16,15 @@
 def show_trajectory(trajectory, one_shot, figure_size):
     if trajectory.shape[-1] == 2:
         ax = display_2D_trajectory(
-            trajectory, size=figure_size, one_shot=one_shot % trajectory.shape[0]
+            trajectory, figsize=figure_size, one_shot=one_shot % trajectory.shape[0]
         )
         ax.set_aspect("equal")
         plt.tight_layout()
         plt.show()
     else:
         ax = display_3D_trajectory(
             trajectory,
-            size=figure_size,
+            figsize=figure_size,
             one_shot=one_shot % trajectory.shape[0],
             per_plane=False,
         )
@@ -49,15 +49,15 @@ def show_trajectories(
         if dim == "3D" and traj.shape[-1] == 3:
             ax = display_3D_trajectory(
                 traj,
-                size=subfig_size,
+                figsize=subfig_size,
                 one_shot=one_shot % traj.shape[0],
                 subfigure=subfig,
                 per_plane=False,
             )
         else:
             ax = display_2D_trajectory(
                 traj[..., axes],
-                size=subfig_size,
+                figsize=subfig_size,
                 one_shot=one_shot % traj.shape[0],
                 subfigure=subfig,
             )

pyproject.toml

@@ -4,21 +4,21 @@ description = "MRI Non-Cartesian Fourier Operators with multiple computation bac
 authors = [{name="Pierre-antoine Comby", email="[email protected]"}]
 
 readme = "README.md"
-dependencies = ["numpy", "scipy", "matplotlib", "tqdm", "joblib"]
-requires-python = ">=3.9"
+dependencies = ["numpy>=2.2", "scipy>=1.13", "matplotlib>=3.9", "tqdm", "joblib>=1.4"]
+requires-python = ">=3.10"
 
 dynamic = ["version"]
 
 [project.optional-dependencies]
 
-gpunufft = ["gpuNUFFT>=0.9.0", "cupy-cuda12x"]
+gpunufft = ["gpuNUFFT>=0.10.1", "cupy-cuda12x"]
 
 torchkbnufft = ["torchkbnufft", "cupy-cuda12x"]
 torchkbnufft-cpu = ["torchkbnufft", "cupy-cuda12x"]
 torchkbnufft-gpu = ["torchkbnufft", "cupy-cuda12x"]
 
 cufinufft = ["cufinufft>=2.4.0", "cupy-cuda12x"]
-tensorflow = ["tensorflow-mri==0.21.0", "tensorflow-probability==0.17.0", "tensorflow-io==0.27.0", "matplotlib==3.7"]
+tensorflow = ["tensorflow-mri>=0.22.0"]
 finufft = ["finufft>=2.4.0"]
 sigpy = ["sigpy"]
 pynfft = ["pynfft3"]

src/mrinufft/extras/smaps.py

@@ -4,17 +4,21 @@
 
 from mrinufft.density.utils import flat_traj
 from mrinufft._utils import get_array_module
+from mrinufft._array_compat import with_numpy_cupy
 from .utils import register_smaps
 import numpy as np
+from numpy.typing import NDArray
+
+from collections.abc import Callable
 
 
 def _extract_kspace_center(
-    kspace_data,
-    kspace_loc,
-    threshold=None,
-    density=None,
-    window_fun="ellipse",
-):
+    kspace_data: NDArray,
+    kspace_loc: NDArray,
+    threshold: float | tuple[float, ...] = None,
+    density: NDArray | None = None,
+    window_fun: str | Callable[[NDArray], NDArray] = "ellipse",
+) -> tuple[NDArray, NDArray, NDArray | None]:
     r"""Extract k-space center and corresponding sampling locations.
 
     The extracted center of the k-space, i.e. both the kspace locations and
@@ -81,7 +85,7 @@ def _extract_kspace_center(
         return data_thresholded, center_locations, dc
     else:
         if callable(window_fun):
-            window = window_fun(center_locations)
+            window = window_fun(kspace_loc)
         else:
             if window_fun in ["hann", "hanning", "hamming"]:
                 radius = xp.linalg.norm(kspace_loc, axis=1)
@@ -99,16 +103,16 @@ def _extract_kspace_center(
 @register_smaps
 @flat_traj
 def low_frequency(
-    traj,
-    shape,
-    kspace_data,
-    backend,
+    traj: NDArray,
+    shape: tuple[int, ...],
+    kspace_data: NDArray,
+    backend: str,
     threshold: float | tuple[float, ...] = 0.1,
-    density=None,
-    window_fun: str = "ellipse",
+    density: NDArray | None = None,
+    window_fun: str | Callable[[NDArray], NDArray] = "ellipse",
     blurr_factor: int | float | tuple[float, ...] = 0.0,
     mask: bool = False,
-):
+) -> tuple[NDArray, NDArray]:
     """
     Calculate low-frequency sensitivity maps.
 
@@ -190,3 +194,61 @@ def low_frequency(
         SOS = np.linalg.norm(Smaps, axis=0) + 1e-10
     Smaps = Smaps / SOS
     return Smaps, SOS
+
+
+@with_numpy_cupy
+def coil_compression(
+    kspace_data: NDArray,
+    K: int | float,
+    traj: NDArray | None = None,
+    krad_thresh: float | None = None,
+) -> NDArray:
+    """
+    Coil compression using principal component analysis on k-space data.
+
+    Parameters
+    ----------
+    kspace_data : NDArray
+        Multi-coil k-space data. Shape: (n_coils, n_samples).
+    K : int or float
+        Number of virtual coils to retain (if int), or energy threshold (if
+        float between 0 and 1).
+    traj : NDArray, optional
+        Sampling trajectory. Shape: (n_samples, n_dims).
+    krad_thresh : float, optional
+        Relative k-space radius (as a fraction of maximum) to use for selecting
+        the calibration region for principal component analysis. If None, use
+        all k-space samples.
+
+    Returns
+    -------
+    NDArray
+        Coil-compressed data. Shape: (K, n_samples) if K is int, number of
+        retained components otherwise.
+    """
+    xp = get_array_module(kspace_data)
+
+    if krad_thresh is not None and traj is not None:
+        traj_rad = xp.sqrt(xp.sum(traj**2, axis=-1))
+        center_data = kspace_data[:, traj_rad < krad_thresh * xp.max(traj)]
+    elif krad_thresh is None:
+        center_data = kspace_data
+    else:
+        raise ValueError("traj and krad_thresh must be specified.")
+
+    # Compute the covar matrix of selected data
+    cov = center_data @ center_data.T.conj()
+    w, v = xp.linalg.eigh(cov)
+    # sort eigenvalues largest to smallest
+    si = xp.argsort(w)[::-1]
+    w_sorted = w[si]
+    v_sorted = v[si]
+    if isinstance(K, float):
+        # retain enough components to reach energy K
+        w_cumsum = xp.cumsum(w_sorted)  # from largest to smallest
+        total_energy = xp.sum(w_sorted)
+        K = int(xp.searchsorted(w_cumsum / total_energy, K, side="left") + 1)
+        K = min(K, w_sorted.size)
+    V = v_sorted[:K]  # use top K component
+    compress_data = V @ kspace_data
+    return compress_data