Acquisition Configuration (#285)

* feat!: acquisition Config.

This is a breaking change!

* fmt: black

* feat(test): use acquisition fixture for tests.

* feat: docstrings

* refactor!: use acq in display functions as well

* fmt:black

* fix: correct field name.

* feat: move adc setup to acquisition level.

* fix: example use the acquisition object

* feat: example use acquisition.

* fix: style

* use correct units

* fix: enum meta has evolved in python 3.12

* feat: use better default for grad raster time.

* feat: more SI units and acq use

* fix: don't modify API for write/read_trajectory

* fmt

* fix examples

* update

* fix: acq and display

* feat: add context manager for the acquisition

* update context manager

* fmt: PEP604 type union

fixes UP045

* [docs] update docstring for acquisition config.

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,
             )

src/mrinufft/io/nsp.py

@@ -1,3 +1,4 @@
+# type: ignore
 """Read/Write trajectory for Neurospin sequences."""
 
 from __future__ import annotations
@@ -8,17 +9,19 @@
 from datetime import datetime
 
 import numpy as np
-from typing import Optional
 
 from mrinufft.trajectories.utils import (
-    DEFAULT_GMAX,
-    DEFAULT_RASTER_TIME,
-    DEFAULT_SMAX,
-    KMAX,
+    Acquisition,
+    Hardware,
     Gammas,
+    SI,
     check_hardware_constraints,
     convert_gradients_to_slew_rates,
     convert_trajectory_to_gradients,
+    DEFAULT_GMAX,
+    DEFAULT_SMAX,
+    KMAX,
+    DEFAULT_RASTER_TIME,
 )
 from mrinufft.trajectories.tools import get_gradient_amplitudes_to_travel_for_set_time
 
@@ -47,7 +50,7 @@ def write_gradients(
     Parameters
     ----------
     gradients : np.ndarray
-        Gradients. Shape (num_shots, num_samples_per_shot, dimension).
+        Gradients. Shape (num_shots, num_samples_per_shot, dimension). in T/m.
     initial_positions : np.ndarray
         Initial positions. Shape (num_shots, dimension).
     grad_filename : str
@@ -211,7 +214,7 @@ def write_trajectory(
     img_size: tuple[int, ...],
     grad_filename: str,
     norm_factor: float = KMAX,
-    gamma: float = Gammas.HYDROGEN,
+    gamma: float = Gammas.HYDROGEN / 1e3,
     raster_time: float = DEFAULT_RASTER_TIME,
     check_constraints: bool = True,
     TE_pos: float = 0.5,
@@ -270,22 +273,20 @@ def write_trajectory(
         These are arguments passed to write_gradients function above.
     """
     # Convert normalized trajectory to gradients
+    acq = Acquisition(
+        fov=FOV,
+        img_size=img_size,
+        gamma=gamma,
+        norm_factor=norm_factor,
+        hardware=Hardware(gmax=gmax, smax=smax, grad_raster_time=raster_time),
+    )
     gradients, initial_positions, final_positions = convert_trajectory_to_gradients(
         trajectory,
-        norm_factor=norm_factor,
-        resolution=np.asarray(FOV) / np.asarray(img_size),
-        raster_time=raster_time,
-        gamma=gamma,
+        acq=acq,
         get_final_positions=True,
     )
     Ns_to_skip_at_start = 0
     Ns_to_skip_at_end = 0
-    scan_consts = {
-        "gamma": gamma,
-        "gmax": gmax,
-        "smax": smax,
-        "raster_time": raster_time,
-    }
     if pregrad == "prephase":
         if version < 5.1:
             raise ValueError(
@@ -295,7 +296,7 @@ def write_trajectory(
         start_gradients = get_gradient_amplitudes_to_travel_for_set_time(
             kspace_end_loc=initial_positions,
             end_gradients=gradients[:, 0],
-            **scan_consts,
+            acq=acq,
         )
         initial_positions = np.zeros_like(initial_positions)
         gradients = np.hstack([start_gradients, gradients])
@@ -314,18 +315,17 @@ def write_trajectory(
             kspace_end_loc=edge_locations,
             start_gradients=gradients[:, -1],
             kspace_start_loc=final_positions,
-            **scan_consts,
+            acq=acq,
         )
         gradients = np.hstack([gradients, end_gradients])
         Ns_to_skip_at_end = end_gradients.shape[1]
     # Check constraints if requested
     if check_constraints:
-        slewrates, _ = convert_gradients_to_slew_rates(gradients, raster_time)
+        slewrates, _ = convert_gradients_to_slew_rates(gradients, acq)
         valid, maxG, maxS = check_hardware_constraints(
             gradients=gradients,
             slewrates=slewrates,
-            gmax=gmax,
-            smax=smax,
+            acq=acq,
         )
         if not valid:
             warnings.warn(