fixup! Add arc length parameterization

Author: Daval-G

examples/example_traj_projection.py

@@ -0,0 +1,101 @@
+"""
+======================
+Trajectory constraints
+======================
+
+A collection of methods to make trajectories fit hardware constraints.
+
+"""
+
+# %%
+# Hereafter we illustrate different methods to reduce the gradient
+# strengths and slew rates required for the trajectory to match the
+# hardware constraints of MRI machines. A summary table is available
+# below.
+#
+
+# %%
+# .. list-table:: Constraint fitting methods
+#    :header-rows: 1
+#
+#    * - 
+#      - Gradient strength
+#      - Slew rate
+#      - Path preserved
+#      - Density preserved
+#    * - Arc-length parameterization
+#      - Yes
+#      - No
+#      - Yes
+#      - No
+#
+
+# Internal
+import mrinufft as mn
+from mrinufft.trajectories.utils import compute_gradients_and_slew_rates
+from utils import show_trajectory_full
+
+# External
+import numpy as np
+
+
+# %%
+# Script options
+# ==============
+# These options are used in the examples below as default values for all trajectories.
+
+# Acquisition parameters
+resolution = 1e-3  # Resolution in meters
+raster_time = 40e-3  # Raster time in milliseconds
+
+# %%
+
+# Trajectory parameters
+Nc = 16  # Number of shots
+Ns = 3000  # Number of samples per shot
+in_out = False  # Choose between in-out or center-out trajectories
+nb_zigzags = 5  # Number of zigzags for base trajectories
+
+# %%
+
+# Display parameters
+figure_size = 10  # Figure size for trajectory plots
+subfigure_size = 6  # Figure size for subplots
+one_shot = 2 * Nc // 3  # Highlight one shot in particular
+sample_freq = 60  # Frequency of samples to display in the trajectory plots
+
+# %%
+# We will be using a cone trajectory to showcase the different methods as
+# it switches several times between high gradients and slew rates.
+
+original_trajectory = mn.initialize_2D_cones(Nc, Ns, in_out=in_out, nb_zigzags=nb_zigzags)
+
+# %%
+# Arc-length parameterization
+# ===========================
+# Arc-length parameterization is the simplest method to reduce the gradient
+# strength as it resamples the trajectory to have a constant distance between
+# samples. This is technically the lowest gradient strength achievable while
+# preserving the path of the trajectory, but it does not preserve the k-space 
+# density and can lead to high slew rates as shown below.
+
+show_trajectory_full(original_trajectory, one_shot, subfigure_size, sample_freq)
+
+grads, slews = compute_gradients_and_slew_rates(original_trajectory, resolution, raster_time)
+grad_max, slew_max = np.max(grads), np.max(slews)
+print(f"Max gradient: {grad_max:.3f} T/m, Max slew rate: {slew_max:.3f} T/m/ms")
+
+# %%
+#
+
+from mrinufft.trajectories.projection import parameterize_by_arc_length
+
+projected_trajectory = parameterize_by_arc_length(original_trajectory)
+
+# %%
+
+show_trajectory_full(projected_trajectory, one_shot, subfigure_size, sample_freq)
+
+grads, slews = compute_gradients_and_slew_rates(projected_trajectory, resolution, raster_time)
+grad_max, slew_max = np.max(grads), np.max(slews)
+print(f"Max gradient: {grad_max:.3f} T/m, Max slew rate: {slew_max:.3f} T/m/ms")

examples/utils.py

@@ -9,7 +9,12 @@
 import matplotlib.pyplot as plt
 
 # Internal imports
-from mrinufft import display_2D_trajectory, display_3D_trajectory, displayConfig
+from mrinufft import (
+    display_2D_trajectory,
+    display_3D_trajectory,
+    displayConfig,
+    display_gradients_simply,
+)
 from mrinufft.trajectories.utils import KMAX
 
 
@@ -33,6 +38,59 @@ def show_trajectory(trajectory, one_shot, figure_size):
         plt.show()
 
 
+def show_trajectory_full(trajectory, one_shot, figure_size, sample_freq=10):
+    # General configuration
+    fig = plt.figure(figsize=(3 * figure_size, figure_size))
+    subfigs = fig.subfigures(1, 3, wspace=0)
+
+    # Trajectory display
+    subfigs[0].suptitle("Trajectory", fontsize=displayConfig.fontsize, x=0.5, y=0.98)
+    if trajectory.shape[-1] == 2:
+        ax = display_2D_trajectory(
+            trajectory,
+            size=figure_size,
+            one_shot=one_shot,
+            subfigure=subfigs[0],
+        )
+    else:
+        ax = display_3D_trajectory(
+            trajectory,
+            size=figure_size,
+            one_shot=one_shot,
+            per_plane=False,
+            subfigure=subfigs[0],
+        )
+    ax.set_aspect("equal")
+    for i in range(trajectory.shape[0]):
+        ax.scatter(trajectory[i, ::sample_freq, 0], trajectory[i, ::sample_freq, 1], s=15)
+
+    # Gradient display
+    subfigs[1].suptitle("Gradients", fontsize=displayConfig.fontsize, x=0.5, y=0.98)
+    display_gradients_simply(
+        trajectory,
+        shot_ids=[one_shot],
+        figsize=figure_size,
+        subfigure=subfigs[1],
+        uni_gradient="k",
+        uni_signal="gray",
+    )
+
+    # Slew rates display
+    subfigs[2].suptitle("Slew rates", fontsize=displayConfig.fontsize, x=0.5, y=0.98)
+    display_gradients_simply(
+        np.diff(trajectory, axis=1),
+        shot_ids=[one_shot],
+        figsize=figure_size,
+        subfigure=subfigs[2],
+        uni_gradient="k",
+        uni_signal="gray",
+    )
+    subfigs[2].axes[0].set_ylabel("Sx")
+    subfigs[2].axes[1].set_ylabel("Sy")
+    subfigs[2].axes[2].set_ylabel("|S|")
+    plt.show()
+
+
 def show_trajectories(
     function, arguments, one_shot, subfig_size, dim="3D", axes=(0, 1)
 ):

src/mrinufft/__init__.py

@@ -64,6 +64,8 @@
     displayConfig,
     display_2D_trajectory,
     display_3D_trajectory,
+    display_gradients,
+    display_gradients_simply,
 )
 
 from .density import voronoi, cell_count, pipe, get_density
@@ -130,6 +132,8 @@
     "displayConfig",
     "display_2D_trajectory",
     "display_3D_trajectory",
+    "display_gradients",
+    "display_gradients_simply",
 ]
 
 from importlib.metadata import version, PackageNotFoundError

src/mrinufft/trajectories/__init__.py

@@ -1,14 +1,20 @@
 """Collection of trajectories and tools used for non-Cartesian MRI."""
 
-from .display import display_2D_trajectory, display_3D_trajectory, displayConfig
+from .display import (
+    display_2D_trajectory,
+    display_3D_trajectory,
+    displayConfig,
+    display_gradients,
+    display_gradients_simply,
+)
 from .gradients import patch_center_anomaly
 from .inits import (
     initialize_2D_random_walk,
     initialize_2D_travelling_salesman,
     initialize_3D_random_walk,
     initialize_3D_travelling_salesman,
 )
-from .projection import fit_arc_length
+from .projection import parameterize_by_arc_length
 from .sampling import (
     create_chauffert_density,
     create_cutoff_decay_density,
@@ -117,6 +123,8 @@
     "displayConfig",
     "display_2D_trajectory",
     "display_3D_trajectory",
+    "display_gradients",
+    "display_gradients_simply",
     # projection
-    "fit_arc_length",
+    "parameterize_by_arc_length",
 ]

src/mrinufft/trajectories/display.py

@@ -427,7 +427,6 @@ def display_gradients_simply(
     shot_ids: tuple[int, ...] = (0,),
     figsize: float = 5,
     fill_area: bool = True,
-    show_signal: bool = True,
     uni_signal: str | None = "gray",
     uni_gradient: str | None = None,
     subfigure: plt.Figure | None = None,
@@ -447,10 +446,6 @@ def display_gradients_simply(
         Fills the area under the curve for improved visibility and
         representation of the integral, aka trajectory.
         The default is `True`.
-    show_signal : bool, optional
-        Show an additional illustration of the signal as
-        the modulated distance to the center.
-        The default is `True`.
     uni_signal : str or None, optional
         Define whether the signal should be represented by a
         unique color given as argument or just by the default
@@ -471,13 +466,13 @@ def display_gradients_simply(
         Axes of the figure.
     """
     # Setup figure and labels
-    Nd = trajectory.shape[-1]
+    nb_axes = trajectory.shape[-1] + 1
     if subfigure is None:
-        fig = plt.figure(figsize=(figsize, figsize * (Nd + show_signal) / Nd))
+        fig = plt.figure(figsize=(figsize, figsize * nb_axes / (nb_axes - 1)))
     else:
         fig = subfigure
-    axes = fig.subplots(Nd + show_signal, 1)
-    for i, ax in enumerate(axes[:Nd]):
+    axes = fig.subplots(nb_axes, 1)
+    for i, ax in enumerate(axes[:nb_axes - 1]):
         ax.set_ylabel("G{}".format(["x", "y", "z"][i]), fontsize=displayConfig.fontsize)
     axes[-1].set_xlabel("Time", fontsize=displayConfig.fontsize)
 
@@ -489,50 +484,42 @@ def display_gradients_simply(
 
     # Plot the curves for each axis
     gradients = np.diff(trajectory, axis=1)
-    vmax = 1.1 * np.max(np.abs(gradients[shot_ids, ...]))
-    x_axis = np.arange(gradients.shape[1])
+    vmax = 1.1 * np.max(np.linalg.norm(gradients[shot_ids, ...], axis=-1, ord=1))
+    for ax in axes[:-1]:
+        ax.set_ylim((-vmax, vmax))
+    axes[-1].set_ylim(-0.1 * vmax, vmax)
+
+    time_axis = np.arange(gradients.shape[1])
     colors = displayConfig.get_colorlist()
     for j, s_id in enumerate(shot_ids):
-        for i, ax in enumerate(axes[:Nd]):
-            ax.set_ylim((-vmax, vmax))
-            color = (
-                uni_gradient
-                if uni_gradient is not None
-                else colors[j % displayConfig.nb_colors]
-            )
-            ax.plot(x_axis, gradients[s_id, ..., i], color=color)
+        color = (
+            uni_gradient
+            if uni_gradient is not None
+            else colors[j % displayConfig.nb_colors]
+        )
+
+        # Set each axis individually
+        for i, ax in enumerate(axes[:-1]):
+            ax.plot(time_axis, gradients[s_id, ..., i], color=color)
             if fill_area:
                 ax.fill_between(
-                    x_axis,
+                    time_axis,
                     gradients[s_id, ..., i],
                     alpha=displayConfig.alpha,
                     color=color,
                 )
 
-    # Return axes alone
-    if not show_signal:
-        return axes
-
-    # Show signal as modulated distance to center
-    distances = np.linalg.norm(trajectory[shot_ids, 1:-1], axis=-1)
-    distances = np.tile(distances.reshape((len(shot_ids), -1, 1)), (1, 1, 10))
-    signal = 1 - distances.reshape((len(shot_ids), -1)) / np.max(distances)
-    signal = (
-        signal * np.exp(2j * np.pi * figsize / 100 * np.arange(signal.shape[1]))
-    ).real
-    signal = signal * np.abs(signal) ** 3
-
-    colors = displayConfig.get_colorlist()
-    # Show signal for each requested shot
-    axes[-1].set_ylim((-1, 1))
-    axes[-1].set_ylabel("Signal", fontsize=displayConfig.fontsize)
-    for j in range(len(shot_ids)):
-        color = (
-            uni_signal
-            if (uni_signal is not None)
-            else colors[j % displayConfig.nb_colors]
-        )
-        axes[-1].plot(np.arange(signal.shape[1]), signal[j], color=color)
+        # Set the norm axis if requested
+        gradient_norm = np.linalg.norm(gradients[s_id], axis=-1)
+        axes[-1].set_ylabel("|G|", fontsize=displayConfig.fontsize)
+        axes[-1].plot(gradient_norm, color=color)
+        if fill_area:
+            axes[-1].fill_between(
+                time_axis,
+                gradient_norm,
+                alpha=displayConfig.alpha,
+                color=color,
+                )
     return axes
 
 
@@ -541,7 +528,7 @@ def display_gradients(
     shot_ids: tuple[int, ...] = (0,),
     figsize: float = 5,
     fill_area: bool = True,
-    show_signal: bool = True,
+    show_norm: bool = True,
     uni_signal: str | None = "gray",
     uni_gradient: str | None = None,
     subfigure: plt.Figure | plt.Axes | None = None,
@@ -567,7 +554,7 @@ def display_gradients(
         Fills the area under the curve for improved visibility and
         representation of the integral, aka trajectory.
         The default is `True`.
-    show_signal : bool, optional
+    show_norm : bool, optional
         Show an additional illustration of the signal as
         the modulated distance to the center.
         The default is `True`.
@@ -619,7 +606,7 @@ def display_gradients(
         shot_ids,
         figsize,
         fill_area,
-        show_signal,
+        show_norm,
         uni_signal,
         uni_gradient,
         subfigure,
@@ -633,7 +620,7 @@ def display_gradients(
             fontsize=displayConfig.small_fontsize,
         )
     axes[-1].set_xlabel("Time (ms)", fontsize=displayConfig.small_fontsize)
-    if show_signal:
+    if show_norm:
         axes[-1].set_ylabel("Signal (a.u.)", fontsize=displayConfig.small_fontsize)
 
     # Update axis ticks with rescaled values
@@ -642,7 +629,7 @@ def display_gradients(
         if ax == axes[-1]:
             ax.xaxis.set_tick_params(labelbottom=True)
             ticks = ax.get_xticks()
-            scale = (0.1 if (show_signal and ax == axes[-1]) else 1) * raster_time
+            scale = (0.1 if (show_norm and ax == axes[-1]) else 1) * raster_time
             locator = mticker.FixedLocator(ticks)
             formatter = mticker.FixedFormatter(np.around(scale * ticks, 2))
             ax.xaxis.set_major_locator(locator)
@@ -663,7 +650,7 @@ def display_gradients(
         scale = 1e3 * scale  # Convert from T/m to mT/m
         locator = mticker.FixedLocator(ticks)
         formatter = mticker.FixedFormatter(np.around(scale * ticks, 1))
-        if not show_signal or ax != axes[-1]:
+        if not show_norm or ax != axes[-1]:
             ax.yaxis.set_major_locator(locator)
             ax.yaxis.set_major_formatter(formatter)