Add arc length parameterization (#236)

* Add arc length parameterization

* Clarify annular ring description, extend example trajectory display function

Author: Daval-G

examples/trajectories/example_traj_projection.py

@@ -0,0 +1,102 @@
+"""
+======================
+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 Acquisition, 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
+acq = Acquisition.default
+
+# %%
+
+# 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, acq)
+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, acq)
+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/trajectories/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,65 @@ def show_trajectory(trajectory, one_shot, figure_size):
         plt.show()
 
 
+def show_trajectory_full(trajectory, one_shot, figure_size, sample_freq=10):
+    # General configuration
+    nb_dim = trajectory.shape[-1]
+    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 nb_dim == 2:
+        ax = display_2D_trajectory(
+            trajectory,
+            figsize=figure_size,
+            one_shot=one_shot,
+            subfigure=subfigs[0],
+        )
+    else:
+        ax = display_3D_trajectory(
+            trajectory,
+            figsize=figure_size,
+            one_shot=one_shot,
+            per_plane=False,
+            subfigure=subfigs[0],
+        )
+    for i in range(trajectory.shape[0]):
+        ax.scatter(*(trajectory[i, ::sample_freq].T), s=15)
+    ax.set_aspect("equal")
+
+    # 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")
+    if nb_dim == 2:
+        subfigs[2].axes[2].set_ylabel("|S|")
+    else:
+        subfigs[2].axes[2].set_ylabel("Sz")
+        subfigs[2].axes[3].set_ylabel("|S|")
+    plt.show()
+
+
 def show_trajectories(
     function, arguments, one_shot, subfig_size, dim="3D", axes=(0, 1)
 ):

src/mrinufft/__init__.py

@@ -65,6 +65,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

@@ -3,7 +3,13 @@
 See also the trajectories examples: :ref:`sphx_glr_generated_autoexamples_trajectories`
 """
 
-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_eccentric,
@@ -13,6 +19,7 @@
     initialize_3D_random_walk,
     initialize_3D_travelling_salesman,
 )
+from .projection import parameterize_by_arc_length
 from .sampling import (
     create_chauffert_density,
     create_cutoff_decay_density,
@@ -122,4 +129,8 @@
     "displayConfig",
     "display_2D_trajectory",
     "display_3D_trajectory",
+    "display_gradients",
+    "display_gradients_simply",
+    # projection
+    "parameterize_by_arc_length",
 ]