Simulators API¶

Bases: AcquisitionScheme

A class representing a diffusion MRI acquisition scheme.

This class encapsulates the parameters needed to define a diffusion MRI acquisition, including b-values, gradient directions, and timing parameters.

Attributes:

bvals: Array

bvecs: Array

delta: ArrayLike = field(default_factory=lambda: 0.0106)

Delta: ArrayLike = field(default_factory=lambda: 0.0431)

qvals: ArrayLike

tau: ArrayLike

gradient_strengths: ArrayLike

from_gradient_strengths(gradient_strengths: ArrayLike, gradient_directions: ArrayLike, delta: ArrayLike, Delta: ArrayLike) -> acquisition_scheme

Bases: AcquisitionScheme

Steady-State Free Precession (SSFP) acquisition description.

Parameters are stored in common MR units and converted as properties:

• bvecs: diffusion gradient unit vectors.
• TRs: repetition times [s].
• flipAngles_raw: flip angles [deg] (property flipAngles yields radians).
• diffGradAmps_raw: diffusion gradient amplitudes [T/m] (property diffGradAmps yields G/mm).
• diffGradDur: diffusion gradient durations [s].
• B1: transmit field scaling, unitless.
• T1_raw / T2_raw: relaxation times [ms] (properties T1/T2 yield seconds).
• delta / Delta: diffusion pulse duration and separation [s].

Derived quantities:

• sa, ca: \(\sin(\alpha B_1)\), \(\cos(\alpha B_1)\) with \(\alpha\) in radians.
• E1, E2: relaxation terms \(\exp(-\mathrm{TR}/T_1)\), \(\exp(-\mathrm{TR}/T_2)\).
• qvals: \(q = \gamma \, G \, \delta\) [1/mm] using gyromag_ratio in rad/ms/mT.
• bvals: \(b = q^2 \, 4\pi^2 \, \Delta\).

bvecs: Array

T1_raw: ArrayLike

T2_raw: ArrayLike

B1: ArrayLike

diffGradAmps_raw: ArrayLike

flipAngles_raw: ArrayLike = field(default_factory=lambda: 14.0)

TRs: ArrayLike = field(default_factory=lambda: 0.021)

diffGradDur: ArrayLike = field(default_factory=lambda: 0.01016)

delta: ArrayLike = field(default_factory=lambda: 0.0106)

Delta: ArrayLike = field(default_factory=lambda: 0.008)

gyromag_ratio: ArrayLike = field(default_factory=lambda: 4258 * 2 * jnp.pi)

T1: ArrayLike

T2: ArrayLike

diffGradAmps: ArrayLike

flipAngles: ArrayLike

sa: ArrayLike

ca: ArrayLike

E1: ArrayLike

E2: ArrayLike

logE1: ArrayLike

logE2: ArrayLike

qvals: ArrayLike

bvals: ArrayLike

Bases: ABC

to_theta(*kwargs) -> Array

to_params(theta: ArrayLike) -> Any

from_theta(theta: ArrayLike, **kwargs) -> Compartment

tree_flatten() -> tuple

tree_unflatten(aux_data: Any, children: list) -> Compartment

Bases: Compartment

signal_fn(acq: acquisition_scheme, *args, **kwargs) -> ArrayLike

log_signal_fn(acq: acquisition_scheme, *args, **kwargs) -> ArrayLike

signal(acq: acquisition_scheme, rng=None) -> ArrayLike

log_signal(acq: acquisition_scheme, rng=None) -> ArrayLike

fit(logS: ArrayLike, bvals: ArrayLike, bvecs: ArrayLike) -> Any

to_fod()

Bases: Compartment

noise(signal, rng: key) -> ArrayLike

log_likelihood(signal_pred, signal_obs) -> Array

Bases: SignalCompartment

Composable mixture of signal compartments with optional shared parameters.

A MultiCompartment turns a set of signal compartments :math:S_k(b, \mathbf{g}) and fractions :math:f_k into a single signal

.. math::

S(b, \mathbf{g}) = \sum_k f_k \; S_k(b, \mathbf{g})

Noise compartments can be added on top through noise_types. The class exposes to_theta / from_theta so priors over fractions, shared diffusivities and per-compartment parameters stay Gaussian in optimization space.

References¶

• Stejskal & Tanner, 1965. Spin diffusion measurements: spin echoes in the presence of a time‐dependent field gradient.
• Behrens et al., 2003. Characterization and propagation of uncertainty in diffusion-weighted MR imaging.

num_compartments()

get_all_params()

signal_fn(acq: acquisition_scheme, model_compartments, noise_compartments, model_fractions, model_mask, shared_parameter, rng=None)

to_theta(model_fractions: ArrayLike, model_compartments: list, noise_compartments: list, model_mask: ArrayLike | None = None, shared_parameter: SharedParameterState | None = None)

Bases: SignalCompartment

Isotropic Gaussian diffusion (free water).

Signal follows the Stejskal–Tanner form

.. math:: S(b) = \exp(-b\,\lambda),

with scalar diffusivity :math:\lambda \in [0, 10^{-2}]\,\mathrm{mm}^2/\mathrm{s} (Stejskal & Tanner, 1965).

Initialize the Ball model with a lambda value.

theta_dim: int = 1

lam_min: float = 0.0

lam_max: float = 0.01

lam = lam

log_signal_fn(acq: acquisition_scheme, lam: float, rng=None) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

to_theta(lam: ArrayLike) -> Array

fit(logS: ArrayLike, bvals: ArrayLike, bvecs: Array) -> tuple

to_fod()

Bases: Ball

The StaticBall model is a Ball with a fixed lambda value. The lam parameter is shared from a global parameter state as a class attribute.

Initialize the StaticBall model. The lam parameter is a class attribute and not passed to the constructor.

theta_dim: int = 0

lam: float = None

to_theta() -> ArrayLike

to_params(theta: ArrayLike) -> tuple

from_global_params(params: ArrayLike, idx: list[int]) -> StaticBall

log_signal_fn(acq: acquisition_scheme, rng=None) -> ArrayLike

Bases: Ball

The MultiShellBall model is a Ball with multiple lambda values. The lam parameter is shared from a global parameter state as a class attribute.

Initialize the MultiShellBall model. The lam parameter is a class attribute and not passed to the constructor.

theta_dim: int = 2

lam_std_min: float = 1e-07

lam_std_max: float = 0.01

lam: float = lam

lam_std: float = lam_std

log_signal_fn(acq: acquisition_scheme, lam: float, lam_std: float, rng=None) -> ArrayLike

to_theta(lam: float, lam_std: float) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

Bases: MultiShellBall

The MultiShellBall model is a Ball with multiple lambda values. The lam parameter is shared from a global parameter state as a class attribute.

theta_dim: int = 0

lam: float = None

lam_std: float = None

from_global_params(params: ArrayLike, idx: list[int]) -> StaticBall

to_theta() -> ArrayLike

to_params(theta: ArrayLike) -> tuple

log_signal_fn(acq: acquisition_scheme, rng=None) -> ArrayLike

Bases: SignalCompartment

The SSFPBall model is a Ball with a fixed lambda value. The lam parameter is shared from a global parameter state as a class attribute.

theta_dim: int = 1

lam_min: float = 0.0

lam_max: float = 0.01

lam = lam

signal_fn(acq: ssfp_acquisition_scheme, lam: float, rng=None) -> ArrayLike

log_signal_fn(acq: ssfp_acquisition_scheme, lam: float, rng=None) -> ArrayLike

to_theta(lam: float) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

Bases: SSFPBall

The SSFPStaticBall model is a SSFPBall with a fixed lambda value. The lam parameter is shared from a global parameter state as a class attribute.

theta_dim: int = 0

lam: float = None

from_global_params(params: ArrayLike, idx: list[int]) -> StaticBall

to_theta() -> ArrayLike

to_params(theta: ArrayLike) -> tuple

signal_fn(acq: ssfp_acquisition_scheme, rng=None) -> ArrayLike

log_signal_fn(acq: ssfp_acquisition_scheme, rng=None) -> ArrayLike

Bases: SignalCompartment

The Stejskal Tanner signal approximation of a sphere model. It assumes that pulse length is infinitessimally small and diffusion time large enough so that the diffusion is completely restricted. Only depends on q-value.

Parameters¶

diameter : float, sphere diameter in meters.

References¶

.. [1] Balinov, Balin, et al. "The NMR self-diffusion method applied to restricted diffusion. Simulation of echo attenuation from molecules in spheres and between planes." Journal of Magnetic Resonance, Series A 104.1 (1993): 17-25.

theta_dim = 1

radius_mean = math.log(0.01)

radius_scale = 0.5

radius = radius

log_signal_fn(acq, radius: float, rng=None)

to_theta(radius: float) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

Bases: SignalCompartment

Axially symmetric Gaussian tensor (:math:\lambda_{\parallel}, \lambda_{\perp}).

Signal:

.. math:: S(b, \mathbf{g}) = \exp\Big(-b \big[\lambda_{\perp} + (\lambda_{\parallel}-\lambda_{\perp})(\mathbf{g}\cdot\boldsymbol{\mu})^2\big]\Big),

often used for extra-axonal water (Basser et al., 1994).

Initialize the Zeppelin model with orientation and diffusivity parameters.

theta_dim = 4

min_lam = 0.0

max_lam = 0.01

mu = mu

lambda_par = lambda_par

lambda_perp = lambda_perp

log_signal_fn(acq: acquisition_scheme, mu: ArrayLike, lambda_par: float, lambda_perp: float, rng=None) -> ArrayLike

to_theta(mu: ArrayLike, lambda_par: float, lambda_perp: float) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

fit(logS: ArrayLike, acq) -> tuple

to_fod()

Bases: Zeppelin

The StaticZeppelin model is a Zeppelin with a fixed lambda value.

theta_dim: int = 3

lambda_par: float = None

lambda_perp: float = None

mu = mu

to_theta(mu: ArrayLike) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

Bases: SignalCompartment

The Stejskal-Tanner approximation of a finite-radius cylinder (Callaghan, 1991). Assumes rectangular pulses and :math:\Delta \gg \delta.

NOTE: This model assumes that the gradient pulses are rectangular and that the diffusion time is much larger than the pulse duration and available in the acquisition scheme.

theta_dim = 4

lam_par_max: float = 0.01

radius_mean = math.log(0.01)

radius_scale = 0.5

mu = mu

lam_par = lam_par

radius = radius

log_signal_fn(acq, mu: ArrayLike, lam_par: float, radius: float, rng=None)

to_theta(mu: ArrayLike, lam_par: float, radius: float) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

Bases: ConvolvedSignalCompartment

Watson distribution with zeppelin kernel.

fod_type = Watson

signal_kernel_type = ZeppelinKernel

Bases: ConvolvedSignalCompartment

Bingham distribution with zeppelin kernel.

fod_type = Bingham

signal_kernel_type = ZeppelinKernel

Bases: SignalCompartment

Zero-radius cylinder aligned with unit vector :math:\boldsymbol{\mu}.

Signal:

.. math:: S(b, \mathbf{g}) = \exp\big(-b\,\lambda_{\parallel} (\mathbf{g}\cdot\boldsymbol{\mu})^2\big),

capturing fully anisotropic diffusion along the fiber (Behrens et al., 2003).

Initialize the Stick model with a lambda value and eigenvector.

theta_dim = 3

min_lam = 0.0

max_lam = 0.01

mu = mu

lam_par = lam_par

log_signal_fn(acq: acquisition_scheme, mu: ArrayLike, lam_par: float, rng=None) -> ArrayLike

to_theta(mu: ArrayLike, lam_par: float) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

to_fod()

fit(logS: ArrayLike, acq: acquisition_scheme) -> tuple

Bases: Stick

The StaticStick model is a Stick with a fixed lambda value. The lam_par parameter is shared from a global parameter state as a class attribute, while the mu parameter remains learnable.

Initialize the StaticStick model with a mu value. The lam_par parameter is a class attribute and not passed to the constructor.

theta_dim: int = 2

lam_par: float = None

mu = mu

to_theta(mu: ArrayLike) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

from_global_params(params: ArrayLike, idx: list[int]) -> type

log_signal_fn(acq: acquisition_scheme, mu: ArrayLike, rng=None) -> ArrayLike

Bases: Stick

The MultiShellStick model is a Stick with a fixed lambda value. The lam_par parameter is shared from a global parameter state as a class attribute, while the mu parameter remains learnable.

Initialize the MultiShellStick model with a mu value, lam_par, and lam_par_std.

theta_dim: int = 4

lam_par_std_min: float = 0.0

lam_par_std_max: float = 0.005

mu = mu

lam_par: float = lam_par

lam_par_std: float = lam_par_std

log_signal_fn(acq: acquisition_scheme, mu: ArrayLike, lam_par: float, lam_par_std: float, rng=None) -> ArrayLike

to_theta(mu: ArrayLike, lam_par: float, lam_par_std: float) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

Bases: MultiShellStick

The MultiShellStaticStick model is a Stick with a fixed lambda value. The lam_par parameter is shared from a global parameter state as a class attribute, while the mu parameter remains learnable.

Initialize the MultiShellStaticStick model with a mu value.

theta_dim: int = 2

lam_par: float = None

lam_par_std: float = None

mu = mu

to_theta(mu: ArrayLike) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

from_global_params(params: ArrayLike, idx: list[int]) -> type

log_signal_fn(acq: acquisition_scheme, mu: ArrayLike, rng=None) -> ArrayLike

Bases: Stick

The SSFPStick model is a Stick with a fixed lambda value. The lam_par parameter is shared from a global parameter state as a class attribute, while the mu parameter remains learnable.

theta_dim: int = 3

lam_min: float = 0.0

lam_max: float = 0.01

mu = mu

lam_par: float = lam_par

signal_fn(acq: acquisition_scheme, mu: ArrayLike, lam_par: float, rng=None) -> ArrayLike

log_signal_fn(acq: acquisition_scheme, mu: ArrayLike, lam_par: float, rng=None) -> ArrayLike

Bases: SSFPStick

The SSFPStaticStick model is a SSFPStick with a fixed lambda value. The lam_par parameter is shared from a global parameter state as a class attribute.

theta_dim: int = 2

lam_par: float = None

mu = mu

from_global_params(params: ArrayLike, idx: list[int]) -> type

to_theta(mu: ArrayLike) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

signal_fn(acq: acquisition_scheme, mu: ArrayLike, rng=None) -> ArrayLike

log_signal_fn(acq: acquisition_scheme, mu: ArrayLike, rng=None) -> ArrayLike

Bases: ConvolvedSignalCompartment

Watson distribution with stick kernel.

fod_type = Watson

signal_kernel_type = StickKernel

Bases: ConvolvedSignalCompartment

Bingham distribution with stick kernel.

fod_type = Bingham

signal_kernel_type = StickKernel

Bases: SignalCompartment

The Dot model is a simple model that represents a single point in space. It has no parameters and the signal is constant across all b-values and b-vectors. It represent trapped water molecules in the tissue.

theta_dim: int = 0

log_signal_fn(acq: acquisition_scheme, rng=None) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

to_theta() -> ArrayLike

to_fod()

Bases: SignalCompartment

The Diffusion Tensor Imaging (DTI) model represents the diffusion of water molecules in the brain. It is a simple model that assumes Gaussian diffusion.

Initialize the DTI model with a diffusion tensor D.

theta_dim: int = 6

D_scale = 0.001

min_lam: float = 0.0001

D = D

log_signal_fn(acq: acquisition_scheme, D: ArrayLike, rng=None) -> ArrayLike

to_theta(D: ArrayLike) -> ArrayLike

to_params(theta: ArrayLike) -> tuple

fit(logS: ArrayLike, acq: acquisition_scheme) -> tuple

to_fod()

Bases: ConvolvedSignalCompartment

Watson distribution with NODDI kernel.

fod_type = Watson

signal_kernel_type = NODDIKernel

Bases: ConvolvedSignalCompartment

Bingham distribution with NODDI kernel.

fod_type = Bingham

signal_kernel_type = NODDIKernel

Bases: ConvolvedSignalCompartment

Watson distribution with SANDI kernel.

fod_type = Watson

signal_kernel_type = SimpleSANDIKernel

Bases: ConvolvedSignalCompartment

Bingham distribution with SANDI kernel.

fod_type = Bingham

signal_kernel_type = SimpleSANDIKernel