Display 3D trajectories, with a grid interface

examples/GPU/example_3d_trajectory_display.py

@@ -0,0 +1,111 @@
+"""
+======================
+3D Trajectory display
+======================
+
+In this example, we show some tools available to display 3D trajectories.
+It can be used to understand the k-space sampling patterns, visualize the trajectories, see the sampling times, gradient strengths, slew rates etc.
+Another key feature is to display the sampling density in k-space, for example to check for k-space holes or irregularities in the learning-based trajectories that would lead to artifacts in the images.
+"""
+
+# %%
+# Imports
+from mrinufft.trajectories.display3D import get_gridded_trajectory
+import mrinufft.trajectories.trajectory3D as mtt
+from mrinufft.trajectories.utils import Gammas
+import matplotlib.pyplot as plt
+from matplotlib import gridspec
+import numpy as np
+
+
+# %%
+# Utility function to plot mid-plane slices for 3D volumes
+def plot_slices(axs, volume, title=""):
+    def set_labels(ax, axis_num=None):
+        ax.set_xticks([0, 32, 64])
+        ax.set_yticks([0, 32, 64])
+        ax.set_xticklabels([r"$-\pi$", 0, r"$\pi$"])
+        ax.set_yticklabels([r"$-\pi$", 0, r"$\pi$"])
+        if axis_num is not None:
+            ax.set_xlabel(r"$k_" + "zxy"[axis_num] + r"$")
+            ax.set_ylabel(r"$k_" + "yzx"[axis_num] + r"$")
+
+    for i in range(3):
+        volume = np.rollaxis(volume, i, 0)
+        axs[i].imshow(volume[volume.shape[0] // 2])
+        axs[i].set_title(
+            ((title + f"\n") if i == 0 else "") + r"$k_{" + "xyz"[i] + r"}=0$"
+        )
+        set_labels(axs[i], i)
+
+
+def create_grid(grid_type, trajectories, **kwargs):
+    fig, axs = plt.subplots(3, 3, figsize=(10, 10))
+    for i, (name, traj) in enumerate(trajectories.items()):
+        grid = get_gridded_trajectory(
+            traj,
+            traj_params["img_size"],
+            grid_type=grid_type,
+            traj_params=traj_params,
+            **kwargs,
+        )
+        plot_slices(axs[:, i], grid, title=name)
+    plt.tight_layout()
+
+
+# %%
+# Create a bunch of sample trajectories
+trajectories = {
+    "Radial": mtt.initialize_3D_phyllotaxis_radial(64 * 8, 64),
+    "FLORETs": mtt.initialize_3D_floret(64 * 8, 64, nb_revolutions=2),
+    "Seiffert Spirals": mtt.initialize_3D_seiffert_spiral(64 * 8, 64),
+}
+traj_params = {
+    "FOV": (0.23, 0.23, 0.23),
+    "img_size": (64, 64, 64),
+    "gamma": Gammas.HYDROGEN,
+}
+
+# %%
+# Display the density of the trajectories, along the 3 mid-planes. For this, make `grid_type="density"`.
+create_grid("density", trajectories)
+plt.suptitle("Sampling Density")
+plt.show()
+
+
+# %%
+# Display the sampling times over the trajectories. For this, make `grid_type="time"`.
+# It helps to check the sampling times over the k-space trajectories, which can be responsible for excessive off-resonance artifacts.
+create_grid("time", trajectories)
+plt.suptitle("Sampling Time")
+plt.show()
+
+# %%
+# Display the inversion time of the trajectories. For this, make `grid_type="inversion"`.
+# This helps in obtaining the inversion time when particular region of k-space is sampled, assuming the trajectories are time ordered.
+# This helps understand any issues for imaging involving inversion recovery.
+# The argument `turbo_factor` can be used to tell what is the number of echoes between 2 inversion pulses.
+create_grid("inversion", trajectories, turbo_factor=64)
+plt.suptitle("Inversion Time")
+plt.show()
+# %%
+# Display the k-space holes in the trajectories. For this, make `grid_type="holes"`.
+# This helps in understanding the k-space holes, and can help debug artifacts in reconstructions.
+create_grid("holes", trajectories, threshold=1e-2)
+plt.suptitle("K-space Holes")
+plt.show()
+# %%
+# Display the gradient strength of the trajectories. For this, make `grid_type="gradients"`.
+# This helps in understanding the gradient strength applied at specific k-space region.
+# This can also be used as a surrogate to k-space "velocity", i.e. how fast does trajectory pass through a given region in k-space
+create_grid("gradients", trajectories)
+plt.suptitle("Gradient Strength")
+plt.show()
+
+# %%
+# Display the slew rates of the trajectories. For this, make `grid_type="slew"`.
+# This helps in understanding the slew rates applied at specific k-space region.
+# This can also be used as a surrogate to k-space "acceleration", i.e. how fast does trajectory change in a given region in k-space
+create_grid("slew", trajectories)
+plt.suptitle("Slew Rates")
+plt.show()

src/mrinufft/trajectories/display3D.py

@@ -0,0 +1,140 @@
+"""Utils for displaying 3D trajectories."""
+
+from mrinufft import get_operator, get_density
+from mrinufft.trajectories.utils import (
+    convert_trajectory_to_gradients,
+    convert_gradients_to_slew_rates,
+    KMAX,
+    DEFAULT_RASTER_TIME,
+)
+import numpy as np
+from typing import Tuple
+from mrinufft.io import read_trajectory
+
+
+def get_gridded_trajectory(
+    shots: np.ndarray,
+    shape: Tuple,
+    grid_type: str = "density",
+    osf: int = 1,
+    backend: str = "gpunufft",
+    traj_params: dict = None,
+    turbo_factor: int = 176,
+    elliptical_samp: bool = True,
+    threshold: float = 1e-3,
+):
+    """
+    Compute the gridded trajectory for MRI reconstruction.
+
+    This function helps in gridding a k-space sampling trajectory to a desired shape.
+    The gridding process can be carried out to reflect the sampling density,
+    sampling time, inversion time, k-space holes, gradient strengths, or slew rates.
+    Please check `grid_type` parameter to know the benefits of each type of gridding.
+
+    Parameters
+    ----------
+    shots : ndarray
+        The input array of shape (N, M, D), where N is the number of shots and M is the
+        number of samples per shot and D is the dimension of the trajectory (usually 3)
+    shape : tuple
+        The desired shape of the gridded trajectory.
+    grid_type : str, optional
+        The type of gridded trajectory to compute. Default is "density".
+        It can be one of the following:
+        "density" : Get the sampling density in closest number of samples per voxel.
+            Helps understand suboptimal sampling, by showcasing regions with strong
+            oversampling.
+        "time" : Showcases when the k-space data is acquired in time.
+            This is helpful to view and understand off-resonance effects.
+            Generally, lower off-resonance effects occur when the sampling trajectory
+            has smoother k-space sampling time over the k-space.
+        "inversion" : Relative inversion time at the sampling location. Needs
+            turbo_factor to be set. This is useful for analyzing the exact inversion
+            time when the k-space is acquired, for sequences like MP(2)RAGE.
+        "holes": Show the k-space holes within a ellipsoid of the k-space.
+        "gradients": Show the gradient strengths of the k-space trajectory.
+        "slew": Show the slew rate of the k-space trajectory.
+    osf : int, optional
+        The oversampling factor for the gridded trajectory. Default is 1.
+    backend : str, optional
+        The backend to use for gridding. Default is "gpunufft".
+        Note that "gpunufft" is anyway used to get the `pipe` density internally.
+    traj_params : dict, optional
+        The trajectory parameters. Default is None.
+        This is only needed when `grid_type` is "gradients" or "slew".
+        The parameters needed include `img_size`, `FOV`, and `gamma` of the sequence.
+        Generally these values are stored in the header of the trajectory file.
+    turbo_factor : int, optional
+        The turbo factor when sampling is with inversion. Default is 176, which is
+        the default turbo factor for MPRAGE acquisitions at 1mm whole
+        brain acquisitions.
+    elliptical_samp : bool, optional
+        Whether to use elliptical sampling. Default is True.
+        This is useful while analyzing the k-space holes, especially if the k-space
+        trajectory is expected to be elliptical sampling of k-space
+        (i.e. ellipsoid over cuboid).
+    threshold: float, optional
+        The threshold for the k-space holes. Default is 1e-3.
+        This value is set heuristically to visualize the k-space hole.
+
+    Returns
+    -------
+    ndarray
+        The gridded trajectory of shape `shape`.
+    """
+    samples = shots.reshape(-1, shots.shape[-1])
+    dcomp = get_density("pipe")(samples, shape)
+    grid_op = get_operator(backend)(
+        samples, [sh * osf for sh in shape], density=dcomp, upsampfac=1
+    )
+    gridded_ones = grid_op.raw_op.adj_op(np.ones(samples.shape[0]), None, True)
+    if grid_type == "density":
+        return np.abs(gridded_ones).squeeze()
+    elif grid_type == "time":
+        data = grid_op.raw_op.adj_op(
+            np.tile(np.linspace(1, 10, shots.shape[1]), (shots.shape[0],)),
+            None,
+            True,
+        )
+    elif grid_type == "inversion":
+        data = grid_op.raw_op.adj_op(
+            np.repeat(
+                np.linspace(1, 10, turbo_factor), samples.shape[0] // turbo_factor + 1
+            )[: samples.shape[0]],
+            None,
+            True,
+        )
+    elif grid_type == "holes":
+        data = np.abs(gridded_ones).squeeze() < threshold
+        if elliptical_samp:
+            # If the trajectory uses elliptical sampling, ignore the k-space holes
+            # outside the ellipsoid.
+            data[
+                np.linalg.norm(
+                    np.meshgrid(
+                        *[np.linspace(-1, 1, sh) for sh in shape], indexing="ij"
+                    ),
+                    axis=0,
+                )
+                > 1
+            ] = 0
+    elif grid_type in ["gradients", "slew"]:
+        gradients, initial_position = convert_trajectory_to_gradients(
+            shots,
+            norm_factor=KMAX,
+            resolution=np.asarray(traj_params["FOV"])
+            / np.asarray(traj_params["img_size"]),
+            raster_time=DEFAULT_RASTER_TIME,
+            gamma=traj_params["gamma"],
+        )
+        if grid_type == "gradients":
+            data = np.hstack(
+                [gradients, np.zeros((gradients.shape[0], 1, gradients.shape[2]))]
+            )
+        else:
+            slews, _ = convert_gradients_to_slew_rates(gradients, DEFAULT_RASTER_TIME)
+            data = np.hstack([slews, np.zeros((slews.shape[0], 2, slews.shape[2]))])
+        data = grid_op.raw_op.adj_op(
+            np.linalg.norm(data, axis=-1).flatten(), None, True
+        )
+    return np.squeeze(np.abs(data))