add simplified MPRAGE model

mcencini · mcencini · commit d2fac9704857 · 2024-12-05T12:08:59.000+01:00
diff --git a/docs/api.rst b/docs/api.rst
@@ -124,6 +124,7 @@ Iterative
    :nosignatures:  
    
    torchsim.models.FSEModel
+   torchsim.models.MPRAGEModel
    torchsim.models.MP2RAGEModel
    torchsim.models.MPnRAGEModel
    torchsim.models.MRFModel
@@ -150,6 +151,7 @@ Iterative
    :nosignatures:  
    
    torchsim.fse_sim
+  torchsim.mprage_sim
    torchsim.mp2rage_sim
    torchsim.mpnrage_sim
    torchsim.mrf_sim
diff --git a/docs/sg_execution_times.rst b/docs/sg_execution_times.rst
@@ -6,7 +6,7 @@
 
 Computation times
 =================
-**01:07.135** total execution time for 4 files **from all galleries**:
+**00:28.558** total execution time for 4 files **from all galleries**:
 
 .. container::
 
@@ -33,10 +33,10 @@ Computation times
      - Time
      - Mem (MB)
    * - :ref:`sphx_glr_generated_autoexamples_02-synth-data.py` (``../examples/02-synth-data.py``)
-     - 00:42.438
+     - 00:22.610
      - 0.0
    * - :ref:`sphx_glr_generated_autoexamples_03-fitting.py` (``../examples/03-fitting.py``)
-     - 00:24.697
+     - 00:05.949
      - 0.0
    * - :ref:`sphx_glr_generated_autoexamples_01-simple-simulator.py` (``../examples/01-simple-simulator.py``)
      - 00:00.000
diff --git a/examples/02-synth-data.py b/examples/02-synth-data.py
@@ -28,20 +28,15 @@
 
 warnings.filterwarnings("ignore")
 
+import os
 import torchio as tio
 
-try:
-    ixi_dataset = tio.datasets.IXI(
-        "/home/mcencini/ixi/",
-        modalities=("PD", "T1", "T2"),
-        download=False,
-    )
-except:
-    ixi_dataset = tio.datasets.IXI(
-        "$HOME/.ixi/",
-        modalities=("PD", "T1", "T2"),
-        download=True,
-    )
+path = os.path.realpath("data")
+ixi_dataset = tio.datasets.IXI(
+    path,
+    modalities=("PD", "T2"),
+    download=False,
+)
 
 # get subject 0
 sample_subject = ixi_dataset[0]
diff --git a/examples/03-fitting.py b/examples/03-fitting.py
@@ -25,23 +25,17 @@
 
 warnings.filterwarnings("ignore")
 
+import os
 import numpy as np
 import torchio as tio
 import torchsim
 
-
-try:
-    ixi_dataset = tio.datasets.IXI(
-        "/home/mcencini/ixi/",
-        modalities=("PD", "T1", "T2"),
-        download=False,
-    )
-except:
-    ixi_dataset = tio.datasets.IXI(
-        "$HOME/.ixi/",
-        modalities=("PD", "T1", "T2"),
-        download=True,
-    )
+path = os.path.realpath("data")
+ixi_dataset = tio.datasets.IXI(
+    path,
+    modalities=("PD", "T2"),
+    download=False,
+)
 
 # get subject 0
 sample_subject = ixi_dataset[0]
diff --git a/examples/data/PD/IXI002-Guys-0828-PD.nii.gz b/examples/data/PD/IXI002-Guys-0828-PD.nii.gz
diff --git a/examples/data/PD/IXI002-Guys-0828-PD.nii.gz:Zone.Identifier b/examples/data/PD/IXI002-Guys-0828-PD.nii.gz:Zone.Identifier
diff --git a/examples/data/T2/IXI002-Guys-0828-T2.nii.gz b/examples/data/T2/IXI002-Guys-0828-T2.nii.gz
diff --git a/examples/data/T2/IXI002-Guys-0828-T2.nii.gz:Zone.Identifier b/examples/data/T2/IXI002-Guys-0828-T2.nii.gz:Zone.Identifier
diff --git a/src/torchsim/_functional/__init__.py b/src/torchsim/_functional/__init__.py
@@ -17,6 +17,11 @@
 __all__.append("mp2rage_sim")
 
 
+from ._mprage import mprage_sim  # noqa
+
+__all__.append("mprage_sim")
+
+
 from ._fse import fse_sim  # noqa
 
 __all__.append("fse_sim")
diff --git a/src/torchsim/_functional/_mprage.py b/src/torchsim/_functional/_mprage.py
@@ -0,0 +1,72 @@
+"""MPRAGE simulator."""
+
+__all__ = ["mprage_sim"]
+
+import numpy.typing as npt
+import torch
+
+from ..models.mprage import MPRAGEModel
+
+
+def mprage_sim(
+    TI: float,
+    flip: float,
+    TRspgr: float,
+    nshots: int | npt.ArrayLike,
+    T1: float | npt.ArrayLike,
+    diff: str | tuple[str] = None,
+    inv_efficiency: float | npt.ArrayLike = 1.0,
+    M0: float | npt.ArrayLike = 1.0,
+    chunk_size: int = None,
+    device: str | torch.device = None,
+):
+    """
+    MPRAGE simulator wrapper.
+
+    Parameters
+    ----------
+    TI : float
+        Inversion time (s) in milliseconds.
+    flip : float | npt.ArrayLike
+        Flip angle train in degrees of shape ``(2,)``.
+        If scalar, assume same angle for both blocks.
+    TRspgr : float
+        Repetition time in milliseconds for each SPGR readout.
+    nshots : int | npt.ArrayLike
+        Number of SPGR readout within the inversion block of shape ``(npre, npost)``
+        If scalar, assume ``npre == npost == 0.5 * nshots``. Usually, this
+        is the number of slice encoding lines ``(nshots = nz / Rz)``,
+        i.e., the number of slices divided by the total acceleration factor along ``z``.
+    T1 : float | npt.ArrayLike
+        Longitudinal relaxation time in milliseconds.
+    diff : str | tuple[str], optional
+        Arguments to get the signal derivative with respect to.
+        The default is ``None`` (no differentation).
+    inv_efficiency : float | npt.ArrayLike, optional
+        Inversion efficiency map, default is ``1.0``.
+    M0 : float or array-like, optional
+        Proton density scaling factor, default is ``1.0``.
+    TI : float | npt.ArrayLike, optional
+        Inversion time in milliseconds.
+        The default is ``0.0``.
+    chunk_size : int, optional
+        Number of atoms to be simulated in parallel.
+        The default is ``None``.
+    device : str | torch.device, optional
+        Computational device for simulation.
+        The default is ``None`` (infer from input).
+
+    Returns
+    -------
+    sig : npt.ArrayLike
+        Signal evolution of shape ``(...,)``.
+    jac : npt.ArrayLike
+        Derivatives of signal wrt ``diff`` parameters,
+        of shape ``(..., len(diff))``.
+        Not returned if ``diff`` is ``None``.
+
+    """
+    model = MPRAGEModel(diff, chunk_size, device)
+    model.set_properties(T1, M0, inv_efficiency)
+    model.set_sequence(TI, flip, TRspgr, nshots)
+    return model()
diff --git a/src/torchsim/models/__init__.py b/src/torchsim/models/__init__.py
@@ -30,3 +30,8 @@
 from .mp2rage import MP2RAGEModel  # noqa
 
 __all__.append("MP2RAGEModel")
+
+
+from .mprage import MPRAGEModel  # noqa
+
+__all__.append("MPRAGEModel")
diff --git a/src/torchsim/models/fse.py b/src/torchsim/models/fse.py
@@ -15,7 +15,7 @@ class FSEModel(AbstractModel):
     """
     Fast Spin Echo (FSE) Model.
 
-    This class simulates fast spin echo (FSE) signals based on
+    This class models fast spin echo (FSE) signals based on
     tissue properties, pulse sequence parameters, and experimental conditions. It
     uses Extended Phase Graph (EPG) formalism to compute the magnetization evolution
     over time.
diff --git a/src/torchsim/models/mp2rage.py b/src/torchsim/models/mp2rage.py
@@ -13,7 +13,7 @@ class MP2RAGEModel(AbstractModel):
     """
     Magnetization Prepared (2) RApid Gradient Echo (MPnRAGE) Model.
 
-    This class Magnetization Prepared RApid Gradient Echo with 2 volumes per segment
+    This class models Magnetization Prepared RApid Gradient Echo with 2 volumes per segment
     (MP2RAGE) signals based on tissue properties, pulse sequence parameters,
     and experimental conditions. It uses Extended Phase Graph (EPG) formalism
     to compute the magnetization evolution over time.
@@ -28,9 +28,8 @@ class MP2RAGEModel(AbstractModel):
     set_sequence(nshots, flip, TR, TI=0.0)
         Configures the pulse sequence parameters for the simulation.
 
-    _engine(T1, TI, flip, TRspgr, TRmp2rage, nshots, M0=1.0, inv_efficiency=1.0,
-            slice_prof=1.0)
-        Computes the MPnRAGE signal for given tissue properties and sequence parameters.
+    _engine(T1, TI, flip, TRspgr, TRmp2rage, nshots, M0=1.0, inv_efficiency=1.0)
+        Computes the MP2RAGE signal for given tissue properties and sequence parameters.
 
     Examples
     --------
diff --git a/src/torchsim/models/mpnrage.py b/src/torchsim/models/mpnrage.py
@@ -15,7 +15,7 @@ class MPnRAGEModel(AbstractModel):
     """
     Magnetization Prepared (n) RApid Gradient Echo (MPnRAGE) Model.
 
-    This class Magnetization Prepared RApid Gradient Echo with n volumes per segment
+    This class models Magnetization Prepared RApid Gradient Echo with n volumes per segment
     (MPnRAGE) signals based on tissue properties, pulse sequence parameters,
     and experimental conditions. It uses Extended Phase Graph (EPG) formalism
     to compute the magnetization evolution over time.
diff --git a/src/torchsim/models/mprage.py b/src/torchsim/models/mprage.py
@@ -0,0 +1,155 @@
+"""MPRAGE sub-routines."""
+
+__all__ = ["MPRAGEModel"]
+
+from ..base import AbstractModel
+from ..base import autocast
+
+import numpy.typing as npt
+import torch
+
+from .. import epg
+
+
+class MPRAGEModel(AbstractModel):
+    """
+    Magnetization Prepared RApid Gradient Echo (MPnRAGE) Model.
+
+    This class models Magnetization Prepared RApid Gradient Echo (MPRAGE) signals
+    based on tissue properties, pulse sequence parameters, and experimental conditions.
+    It uses Extended Phase Graph (EPG) formalism to compute the magnetization evolution over time.
+
+    Assume that signal is sampled at center of k-space only.
+
+    Methods
+    -------
+    set_properties(T1, M0=1.0, inv_efficiency=1.0)
+        Sets tissue relaxation properties and experimental conditions.
+
+    set_sequence(nshots, flip, TR, TI=0.0)
+        Configures the pulse sequence parameters for the simulation.
+
+    _engine(T1, TI, flip, TRspgr, nshots, M0=1.0, inv_efficiency=1.0)
+        Computes the MPRAGE signal for given tissue properties and sequence parameters.
+
+    Examples
+    --------
+    .. exec::
+
+        from torchsim.models import MPRAGEModel
+
+        model = MPRAGEModel()
+        model.set_properties(T1=(200, 1000), inv_efficiency=0.95)
+        model.set_sequence(TI=500.0, flip=5.0, TRspgr=5.0, nshots=128)
+        signal = model()
+
+    """
+
+    @autocast
+    def set_properties(
+        self,
+        T1: float | npt.ArrayLike,
+        M0: float | npt.ArrayLike = 1.0,
+        inv_efficiency: float | npt.ArrayLike = 1.0,
+    ):
+        """
+        Set tissue and system-specific properties for the MRF model.
+
+        Parameters
+        ----------
+        T1 : float | npt.ArrayLike
+            Longitudinal relaxation time in milliseconds.
+        M0 : float or array-like, optional
+            Proton density scaling factor, default is ``1.0``.
+        inv_efficiency : float | npt.ArrayLike, optional
+            Inversion efficiency map, default is ``1.0``.
+
+        """
+        self.properties.T1 = T1
+        self.properties.M0 = M0
+        self.properties.inv_efficiency = inv_efficiency
+
+    @autocast
+    def set_sequence(
+        self,
+        TI: float,
+        flip: float,
+        TRspgr: float,
+        nshots: int | npt.ArrayLike,
+    ):
+        """
+        Set sequence parameters for the SPGR model.
+
+        Parameters
+        ----------
+        TI : float
+            Inversion time in milliseconds of shape ``(2,)``.
+        flip : float | npt.ArrayLike
+            Flip angle train in degrees.
+        TRspgr : float
+            Repetition time in milliseconds for each SPGR readout.
+        TRmprage : float
+            Repetition time in milliseconds for the whole inversion block.
+        nshots : int | npt.ArrayLike
+            Number of SPGR readout within the inversion block of shape ``(npre, npost)``
+            If scalar, assume ``npre == npost == 0.5 * nshots``. Usually, this
+            is the number of slice encoding lines ``(nshots = nz / Rz)``,
+            i.e., the number of slices divided by the total acceleration factor along ``z``.
+
+        """
+        self.sequence.nshots = nshots
+        self.sequence.TI = TI * 1e-3  # ms -> s
+        self.sequence.flip = torch.pi * flip / 180.0
+        self.sequence.TRspgr = TRspgr * 1e-3  # ms -> s
+        if nshots.numel() == 1:
+            nshots = torch.repeat_interleave(nshots // 2, 2)
+        self.sequence.nshots = nshots
+
+    @staticmethod
+    def _engine(
+        T1: float | npt.ArrayLike,
+        TI: npt.ArrayLike,
+        flip: float | npt.ArrayLike,
+        TRspgr: float,
+        TRmprage: float,
+        nshots: int | npt.ArrayLike,
+        M0: float | npt.ArrayLike = 1.0,
+        inv_efficiency: float | npt.ArrayLike = 1.0,
+    ):
+        R1 = 1e3 / T1
+
+        # Calculate number of shots and time before DC sampling
+        nshots_bef = nshots[0]
+        time_bef = nshots_bef * TRspgr
+
+        # Prepare EPG states matrix
+        states = epg.states_matrix(
+            device=R1.device,
+            nstates=1,
+        )
+        # Prepare excitation pulse
+        RF = epg.rf_pulse_op(flip)
+
+        # Prepare relaxation operator for preparation pulse
+        E1inv, rE1inv = epg.longitudinal_relaxation_op(R1, TI - time_bef)
+
+        # Prepare relaxation operator for sequence loop
+        E1, rE1 = epg.longitudinal_relaxation_op(R1, TRspgr)
+
+        # Apply inversion
+        states = epg.adiabatic_inversion(states, inv_efficiency)
+        states = epg.longitudinal_relaxation(states, E1inv, rE1inv)
+        states = epg.spoil(states)
+
+        # Scan loop
+        for p in range(nshots_bef):
+
+            # Apply RF pulse
+            states = epg.rf_pulse(states, RF)
+
+            # Evolve
+            states = epg.longitudinal_relaxation(states, E1, rE1)
+            states = epg.spoil(states)
+
+        # Record signal
+        return M0 * 1j * epg.get_signal(states)
diff --git a/src/torchsim/models/mrf.py b/src/torchsim/models/mrf.py
@@ -15,7 +15,7 @@ class MRFModel(AbstractModel):
     """
     SSFP Magnetic Resonance Fingerprinting (MRF) Model.
 
-    This class simulates steady-state free precession (SSFP) MRF signals based on
+    This class models steady-state free precession (SSFP) MRF signals based on
     tissue properties, pulse sequence parameters, and experimental conditions. It
     uses Extended Phase Graph (EPG) formalism to compute the magnetization evolution
     over time.
diff --git a/tests/functionals/test_mprage_func.py b/tests/functionals/test_mprage_func.py
