epgx 0.1.dev0__py3-none-any.whl

epgx/methods.py

@@ -0,0 +1,329 @@
+"""Simulation methods: configuration of the states
+
+A method configures the representation of the magnetization states (dense
+lattice of phase states, isochromats, later sparse slots) and provides
+defaults. It is a static configuration (hashable), which:
+    - checks that every operator of the sequence can be represented,
+    - builds the initial state (representation, capacity, coordinates),
+    - provides defaults and warnings (e.g. capacity, aliasing).
+
+Operators select their kernels from the state (`State.representation`), so
+they also apply directly to states.
+
+A simulation uses a single method. By default, it is inferred from the
+operators: the simplest representation covering all of them.
+"""
+
+import dataclasses
+import math
+import warnings
+
+import jax.numpy as jnp
+import numpy as np
+
+from . import kernels
+from .ops import S, flatten_sequence, get_dephasing_axes, getnshifts
+
+# names of the dephasing axes of the shift components (keys of `Lattice.nmax`)
+AXIS_NAMES = ("x", "y", "z")
+
+
+@dataclasses.dataclass(frozen=True)
+class Method:
+    """Base class of simulation methods"""
+
+    # representation of the states (cf. kernels.State.representation)
+    representation = None
+    # names of the core axes of the states (e.g. for labelled outputs)
+    core_dims = ()
+
+    def core_axes(self, sequence):
+        """names of the core axes of the states for a sequence"""
+        return self.core_dims
+
+    @property
+    def coords(self):
+        """coordinates of the core axes (e.g. for labelled outputs)"""
+        return {}
+
+    def core_coords(self, sequence):
+        """coordinates of the core axes for a sequence (dict: axis -> values)"""
+        return self.coords
+
+    def check(self, sequence):
+        """raise if an operator of the sequence cannot be represented"""
+
+    def init_state(self, sequence, shape, density=1.0, dtype=None):
+        """initial state (at equilibrium)"""
+        raise NotImplementedError()
+
+
+@dataclasses.dataclass(frozen=True)
+class Lattice(Method):
+    """Dense lattice of phase states (integer shifts)
+
+    One lattice axis per dephasing direction: "k" if the sequence has only
+    scalar shifts along x, else "kx", "ky", "kz" (as many as the longest
+    vector shift) and "k_<name>" for named dephasing axes.
+
+    Args:
+        nmax: maximum number of phase states along each axis; states shifted
+            beyond are dropped. An int (all axes), a tuple (x, y, z), a dict
+            by direction ("x", "y", "z", names of dephasing axes), or None: no
+            truncation (capacity: total number of shifts along each axis)
+            TODO: default value from the sequence and sample (e.g. relaxation)
+        warn: if nmax is None, warn when the capacity exceeds this value
+    """
+
+    nmax: object = None
+    warn: int = 256
+
+    representation = "lattice"
+    core_dims = ("k",)
+
+    def __post_init__(self):
+        # hashable (static argument of the compiled program)
+        if isinstance(self.nmax, dict):
+            object.__setattr__(self, "nmax", tuple(sorted(self.nmax.items())))
+        elif isinstance(self.nmax, list):
+            object.__setattr__(self, "nmax", tuple(self.nmax))
+
+    def core_axes(self, sequence):
+        return _core_axes(sequence, self.representation)
+
+    def check(self, sequence):
+        for op in _shifts(sequence):
+            if not op.integer:
+                raise ValueError(
+                    f"Non-integer shift on the lattice: {op}, "
+                    "use another method, e.g. `method=Isochromats(n=...)`"
+                )
+        get_dephasing_axes(sequence)  # consistent axes
+        self._limits(_directions(sequence))  # nmax consistent with the axes
+
+    def nstates(self, sequence):
+        """state capacity along each axis: min(number of shifts, nmax)"""
+        nshifts = getnshifts(sequence) or {"x": 0}
+        limits = self._limits(tuple(nshifts))
+        nstates = []
+        for nshift, limit in zip(nshifts.values(), limits):
+            if limit is not None:
+                nstates.append(min(nshift, limit))
+                continue
+            elif nshift > self.warn:
+                warnings.warn(
+                    f"Large phase-state capacity ({nshift} states, without "
+                    "truncation): set a maximum number of states, e.g. "
+                    "`method=Lattice(nmax=...)`",
+                    stacklevel=4,
+                )
+            nstates.append(nshift)
+        return tuple(nstates)
+
+    def nstate(self, sequence):
+        """state capacity of a single-axis lattice: min(number of shifts, nmax)"""
+        nstates = self.nstates(sequence)
+        if len(nstates) != 1:
+            raise ValueError("`nstate` requires a single lattice axis: use `nstates`")
+        return nstates[0]
+
+    def init_state(self, sequence, shape, density=1.0, dtype=None):
+        axes = self.core_axes(sequence)
+        nstates = self.nstates(sequence)
+        return kernels.State.at_equilibrium(
+            shape, nstates, axes=axes, density=density, dtype=dtype
+        )
+
+    def _limits(self, directions):
+        """nmax along each direction (None: no truncation)"""
+        return _per_direction(self.nmax, directions, "nmax", "lattice axes")
+
+
+@dataclasses.dataclass(frozen=True)
+class Isochromats(Method):
+    """Isochromats along the dephasing axes (reference simulator)
+
+    Each isochromat is the magnetization at a position x within the voxel
+    (unit: the period of a unit shift along each axis). A shift S(k) is a
+    precession depending on the position: M+ -> exp(2i pi k.x) M+, exact for
+    any k (integer or not). The voxel signal (F0) is the mean over isochromats.
+    The positions are stored in the state (static coordinates of the sample
+    axes), so that operators select their kernel from the state.
+
+    Sample axes: "iso" if the sequence has only scalar shifts along x, else
+    "iso_x", "iso_y", "iso_z" (a regular grid, as many axes as the longest
+    vector shift) and "iso_<name>" for named dephasing axes.
+
+    Args:
+        n: number of isochromats on a regular grid over one period along each
+            axis, x = (j + 1/2) / n - 1/2 (j = 0..n-1): an int (all axes), a
+            tuple (x, y, z) or a dict by direction ("x", "y", "z", names of
+            dephasing axes)
+        positions: explicit positions along a single axis (e.g. random
+            samples), instead of n
+        warn: warn if a regular grid is too coarse for the sequence (aliasing)
+
+    On a regular grid, n > nshift along each axis (the total number of shifts)
+    gives the same signal as the lattice method without truncation
+    (n >= 2 nshift + 1 for all the phase states).
+    """
+
+    n: object = None
+    positions: tuple = None
+    warn: bool = True
+
+    representation = "isochromats"
+    core_dims = ("iso",)
+
+    def __post_init__(self):
+        n = self.n
+        if self.positions is not None:
+            positions = np.ravel(np.asarray(self.positions, dtype=float))
+            if n is not None and n != positions.size:
+                raise ValueError(f"Inconsistent n={n} and {positions.size} positions")
+            object.__setattr__(self, "positions", tuple(positions.tolist()))
+            object.__setattr__(self, "n", positions.size)
+            return
+        if isinstance(n, dict):
+            n = tuple(sorted(n.items()))  # hashable
+        elif isinstance(n, list):
+            n = tuple(n)
+        counts = [v for _, v in n] if _is_items(n) else n
+        counts = [counts] if isinstance(counts, int) else counts
+        if (
+            n is None
+            or not counts
+            or any(not isinstance(v, int) or v < 1 for v in counts)
+        ):
+            raise ValueError("Isochromats require `n` > 0 (per axis) or `positions`")
+        object.__setattr__(self, "n", n)
+        if isinstance(n, int):  # positions of a single axis (1d)
+            object.__setattr__(self, "positions", tuple(_grid(n).tolist()))
+
+    @property
+    def regular(self):
+        """True if the positions are regular grids"""
+        if self.positions is None:
+            return True
+        return np.allclose(self.positions, _grid(self.n))
+
+    @property
+    def coords(self):
+        """coordinates of a single sample axis (1d)"""
+        if self.positions is None:
+            return {}
+        return {"iso": np.asarray(self.positions)}
+
+    def core_axes(self, sequence):
+        return _core_axes(sequence, self.representation)
+
+    def core_coords(self, sequence):
+        axes = self.core_axes(sequence)
+        if axes == ("iso",) and self.positions is not None:
+            return {"iso": np.asarray(self.positions)}
+        counts = self._counts(_directions(sequence))
+        return {axis: _grid(n) for axis, n in zip(axes, counts)}
+
+    def check(self, sequence):
+        axes = self.core_axes(sequence)
+        if not self.regular and axes != ("iso",):
+            raise NotImplementedError(
+                f"Explicit positions with nd shifts (axes {axes}): use `n`"
+            )
+        get_dephasing_axes(sequence)  # consistent axes
+        self._counts(_directions(sequence))  # n consistent with the axes
+
+    def init_state(self, sequence, shape, density=1.0, dtype=None):
+        coords = self.core_coords(sequence)
+        nshifts = (getnshifts(sequence) or {"x": 0}).values()
+        if self.warn and self.regular:
+            for (axis, x), nshift in zip(coords.items(), nshifts):
+                if len(x) <= nshift:
+                    warnings.warn(
+                        f"Aliasing: {len(x)} isochromats for {nshift:g} shifts along "
+                        f"{axis!r}, use n > {nshift:g} (e.g. n={2 * math.ceil(nshift) + 1})",
+                        stacklevel=3,
+                    )
+        density = jnp.asarray(density)
+        shape = jnp.broadcast_shapes(tuple(shape), density.shape)
+        sizes = tuple(len(x) for x in coords.values())
+        F = jnp.zeros(shape + sizes, dtype=dtype or complex)
+        density = jnp.broadcast_to(density, shape)
+        state = kernels.State(
+            F,
+            F,
+            density,
+            axes=tuple(coords),
+            coords=tuple(
+                (axis, tuple(np.asarray(x).tolist())) for axis, x in coords.items()
+            ),
+            representation=self.representation,
+        )
+        return state.equilibrium
+
+    def _counts(self, directions):
+        """number of isochromats along each direction"""
+        counts = _per_direction(self.n, directions, "n", "sample axes")
+        if any(c is None for c in counts):
+            raise ValueError(f"Isochromats n {self.n} for axes {directions}")
+        return counts
+
+
+def infer_method(sequence):
+    """simplest method covering all operators of the sequence"""
+    for op in _shifts(sequence):
+        if not op.integer:
+            raise ValueError(
+                f"Non-integer shift {op}: set the simulation method explicitly, "
+                "e.g. `method=Isochromats(n=...)`"
+            )
+    return Lattice()
+
+
+def _shifts(sequence):
+    return [op for op in flatten_sequence(sequence) if isinstance(op, S)]
+
+
+def _grid(n):
+    """regular grid of n positions over one period, centred on 0"""
+    return (np.arange(n) + 0.5) / n - 0.5
+
+
+def _is_items(value):
+    """True for a tuple of (name, value) pairs (from a dict)"""
+    return isinstance(value, tuple) and bool(value) and isinstance(value[0], tuple)
+
+
+def _directions(sequence):
+    """dephasing directions of the sequence's shifts (default: "x")"""
+    return tuple(getnshifts(sequence)) or ("x",)
+
+
+def _core_axes(sequence, representation):
+    """names of the core axes storing the directions of the sequence"""
+    nd = any(op.nd for op in _shifts(sequence))
+    name = kernels.state.core_axis_name
+    return tuple(name(d, representation, nd) for d in _directions(sequence))
+
+
+def _per_direction(value, directions, label, kind):
+    """setting along each direction: an int or None (all directions), a tuple
+    (x, y, z), or pairs from a dict by direction (missing: None)"""
+    if value is None or isinstance(value, int):
+        return (value,) * len(directions)
+    elif _is_items(value):
+        given = dict(value)
+        unknown = set(given) - set(directions)
+        if unknown:
+            raise ValueError(
+                f"Unknown axes in {label}: {sorted(unknown)}, axes: {directions}"
+            )
+        return tuple(given.get(d) for d in directions)
+    default = [d for d in directions if d in AXIS_NAMES]
+    if len(value) != len(default):
+        raise ValueError(
+            f"{label} {value} has {len(value)} values for {len(default)} {kind} "
+            "(x, y, z; named axes: use a dict)"
+        )
+    given = dict(zip(default, value))
+    return tuple(given.get(d) for d in directions)

epgx/ops/__init__.py

@@ -0,0 +1,30 @@
+"""Layer 2: operators (pytrees in user units)"""
+
+from .dephasing import DephasingAxis, kAxis, tAxis
+from .evolution import E, P, R
+from .functions import (
+    Skeleton,
+    flatten_operators,
+    flatten_sequence,
+    get_adc_times,
+    get_dephasing_axes,
+    getnshift,
+    getnshifts,
+    unflatten_operators,
+)
+from .operator import (
+    NULL,
+    RESET,
+    SPOILER,
+    ZERO_DURATION,
+    EmptyOperator,
+    Operator,
+    Reset,
+    Spoiler,
+    Wait,
+    check_conflicts,
+    operator,
+)
+from .probe import ADC, Adc, Probe, Probes
+from .shift import S
+from .transition import Phi, T

epgx/ops/dephasing.py

@@ -0,0 +1,71 @@
+"""Dephasing axes: directions along which shifts dephase the magnetization
+
+The default space axes "x", "y", "z" exist implicitly: a vector shift
+`S((kx, ky, kz))` acts along them (a scalar shift along "x"). An independent
+axis, e.g. for a gradient unrelated to the others, is declared with its
+physical scale and passed to the operators (`S(1, axis=diff)`):
+
+    diff = kAxis("diff", kvalue=q)      # space axis, rad/mm per unit shift
+    time = tAxis("t", tvalue=tau)       # time axis, ms per unit shift (not used yet)
+
+Rule: the same axis interacts, different axes do not (no cross terms). The
+simulation collects the axes from the operators (two different axes with the
+same name: error). The representation of the states (method) decides how a
+direction is stored: a lattice axis ("kx", "k_diff"), a sample axis of
+isochromats ("iso_x", "iso_diff")...
+
+The scales are carried for the operators that will need physical units
+(diffusion, imaging); they are not used by the shifts.
+"""
+
+import jax
+
+from ..common import pytree, repr_value, static
+
+# default space axes of the shift components
+SPACE_AXES = ("x", "y", "z")
+
+
+@pytree
+class DephasingAxis:
+    """Dephasing axis with its physical scale (one of kvalue, tvalue)
+
+    Args:
+        name: axis name ("x", "y", "z": the default space axes, e.g. to set
+            their scale)
+        kvalue: space axis: wavenumber of a unit shift (rad/mm)
+        tvalue: time axis: duration of a unit shift (ms)
+    """
+
+    kvalue: jax.Array = None
+    tvalue: jax.Array = None
+    name: str = static()
+
+    def __init__(self, name, *, kvalue=None, tvalue=None):
+        if not isinstance(name, str) or not name.isidentifier():
+            raise ValueError(f"Invalid dephasing axis name: {name!r}")
+        elif (kvalue is None) == (tvalue is None):
+            raise ValueError(f"Dephasing axis {name!r}: give one of kvalue, tvalue")
+        object.__setattr__(self, "name", name)
+        object.__setattr__(self, "kvalue", kvalue)
+        object.__setattr__(self, "tvalue", tvalue)
+
+    @property
+    def kind(self):
+        """kind of the axis: "space" or "time" """
+        return "space" if self.kvalue is not None else "time"
+
+    def __repr__(self):
+        key = "kvalue" if self.kind == "space" else "tvalue"
+        value = self.kvalue if self.kind == "space" else self.tvalue
+        return f"DephasingAxis({self.name!r}, {key}={repr_value(value)})"
+
+
+def kAxis(name, kvalue):  # noqa: N802 (constructor in physics notation)
+    """space dephasing axis (cf. DephasingAxis)"""
+    return DephasingAxis(name, kvalue=kvalue)
+
+
+def tAxis(name, tvalue):  # noqa: N802 (constructor in physics notation)
+    """time dephasing axis (cf. DephasingAxis)"""
+    return DephasingAxis(name, tvalue=tvalue)

epgx/ops/evolution.py

@@ -0,0 +1,135 @@
+"""Evolution operators (relaxation, precession)"""
+
+import jax
+import jax.numpy as jnp
+
+from .. import kernels
+from ..kernels import affine as maps
+from .operator import Operator, asparam, check_conflicts, operator, total
+
+
+@operator
+class R(Operator):
+    """Evolution operator with dimensionless exponents (cf. kernels.evolve)
+
+    Args:
+        rT: transverse exponent (complex: relaxation + precession)
+        rL: longitudinal exponent
+        r0: recovery exponent (None: no recovery)
+        cf. Operator for other keyword arguments
+    """
+
+    rT: jax.Array = None
+    rL: jax.Array = None
+    r0: jax.Array = None
+
+    def __init__(self, rT=0, rL=0, *, r0=None, duration=None, name=None):
+        super().__init__(duration=duration, name=name)
+        self._set(rT=asparam(rT), rL=asparam(rL), r0=asparam(r0))
+
+    def normalize(self, sample=None, scanner=None):
+        return {"rT": self.rT, "rL": self.rL, "r0": self.r0}
+
+    kernel = staticmethod(kernels.evolve)
+    affine = staticmethod(maps.evolution_matrix)
+
+
+@operator
+class E(Operator):
+    """Relaxation and precession operator
+
+    Args:
+        tau: evolution time (ms)
+        T1: longitudinal relaxation time (ms), None: from the sample, if any
+        T2: transverse relaxation time (ms), None: from the sample, if any
+            (a value given both here and by the sample raises an error)
+        g: precession rate (kHz), None: no precession
+            (added to the sample frequency offset `df` and the scanner `B0`)
+        duration: operator duration (ms), default: 0, True: tau
+        cf. Operator for other keyword arguments
+    """
+
+    tau: jax.Array = None
+    T1: jax.Array = None
+    T2: jax.Array = None
+    g: jax.Array = None
+
+    def __init__(self, tau, T1=None, T2=None, g=None, *, duration=None, name=None):
+        tau = asparam(tau)
+        duration = tau if duration is True else duration
+        super().__init__(duration=duration, name=name)
+        self._set(tau=tau, T1=asparam(T1), T2=asparam(T2), g=asparam(g))
+
+    def normalize(self, sample=None, scanner=None):
+        # T1/T2 from the operator or the sample (both: error), frequencies add up
+        check_conflicts(self, sample)
+        T1 = self.T1 if self.T1 is not None else getattr(sample, "T1", None)
+        T2 = self.T2 if self.T2 is not None else getattr(sample, "T2", None)
+        g = total(self.g, getattr(sample, "df", None), getattr(scanner, "B0", None))
+        return relaxation_exponents(self.tau, T1, T2, g)
+
+    def sample_fields(self):
+        # T1/T2 from the sample unless explicit, df always added
+        return tuple(f for f in ("T1", "T2") if getattr(self, f) is None) + ("df",)
+
+    def sample_conflicts(self, sample):
+        # T1/T2 given both by the operator and by the sample
+        return tuple(
+            f
+            for f in ("T1", "T2")
+            if getattr(self, f) is not None and getattr(sample, f, None) is not None
+        )
+
+    kernel = staticmethod(kernels.evolve)
+    affine = staticmethod(maps.evolution_matrix)
+
+
+@operator
+class P(Operator):
+    """Precession operator
+
+    Args:
+        tau: evolution time (ms)
+        g: precession rate (kHz)
+        duration: operator duration (ms), default: 0, True: tau
+        cf. Operator for other keyword arguments
+    """
+
+    tau: jax.Array = None
+    g: jax.Array = None
+
+    def __init__(self, tau, g, *, duration=None, name=None):
+        tau = asparam(tau)
+        duration = tau if duration is True else duration
+        super().__init__(duration=duration, name=name)
+        self._set(tau=tau, g=asparam(g))
+
+    def normalize(self, sample=None, scanner=None):
+        # frequencies add up (no relaxation)
+        g = total(self.g, getattr(sample, "df", None), getattr(scanner, "B0", None))
+        return relaxation_exponents(self.tau, None, None, g)
+
+    def sample_fields(self):
+        return ("df",)
+
+    kernel = staticmethod(kernels.evolve)
+    affine = staticmethod(maps.evolution_matrix)
+
+
+#
+# functions
+
+
+def relaxation_exponents(tau, T1=None, T2=None, g=None):
+    """kernel exponents of E(tau, T1, T2, g) (None: no relaxation/precession)
+
+    Returns:
+        dict with rT = tau * (1/T2 + 2i pi g), rL = r0 = tau / T1
+    """
+    rT = 0.0 if T2 is None else tau / T2
+    if g is not None:
+        rT = rT + 2j * jnp.pi * g * tau
+    if T1 is None:
+        return {"rT": rT, "rL": 0.0, "r0": None}
+    rL = tau / T1
+    return {"rT": rT, "rL": rL, "r0": rL}