epgx 0.1.dev0__py3-none-any.whl

epgx/kernels/state.py

@@ -0,0 +1,447 @@
+"""Phase-state container
+
+The state holds the transverse (F) and longitudinal (Z) phase states
+in two arrays of shape `(*shape, *core)`:
+    - `shape`: batch axes (broadcast with operator parameters)
+    - `core`: structural axes of the phase states, named by `State.axes`
+
+Representations (`State.representation`, static), and their core axes:
+    "lattice": phase states (EPG), lattice axes "k" (1d), or "kx", "ky", "kz"
+        (dense nd-lattice, one axis per shift component):
+        k = -nstate..nstate (size 2 * nstate + 1). F holds the F+ states over
+        the whole lattice; F- states are given by symmetry:
+        F-(k) = conj(F+(-k)), Z(-k) = conj(Z(k))
+    "isochromats": samples of the magnetization, sample axis "iso" with
+        coordinates (positions, `State.coords`): F = M+, Z = Mz;
+        F- = conj(F+) at the same sample; the voxel signal is the mean
+
+The representation defines the kind of each core axis (lattice or sample),
+which the helpers rely on; operators select their kernels from it.
+
+The equilibrium is not stored as a full state array: it is
+Z = density at k=0 on lattice axes, and at all samples on sample axes, with
+`density` of shape `shape` (all leaves share the batch axes, so that `vmap`,
+`tree.map` etc. apply uniformly).
+
+Kernels must not index the core axes directly, but use the helpers
+(`expand_param`, `mirror`, `add_equilibrium`, `F0`, `lattice_axes`...), so that
+new representations (nd-lattice, moments, sparse slots...) only require
+extending this module. Operations that are representation-specific by nature
+(e.g. shift) dispatch on `State.representation`.
+"""
+
+import dataclasses
+import functools
+import numbers
+
+import jax
+import jax.numpy as jnp
+
+NAX = jnp.newaxis
+
+# representations: allowed core axes and their kind
+#   lattice axes: phase-state indices (mirrored by conjugation)
+#   sample axes: samples of the magnetization (averaged for the signal)
+# default lattice and sample axes of nd shifts (x, y, z components of k)
+LATTICE_AXES = ("kx", "ky", "kz")
+SAMPLE_AXES = ("iso_x", "iso_y", "iso_z")
+
+# core axes of each representation: fixed names, and prefixes of the axes of
+# named dephasing directions (e.g. "k_diff", "iso_diff"), with their kind
+REPRESENTATIONS = {
+    "lattice": {"k": "lattice", "kx": "lattice", "ky": "lattice", "kz": "lattice"},
+    "isochromats": {
+        "iso": "sample",
+        "iso_x": "sample",
+        "iso_y": "sample",
+        "iso_z": "sample",
+    },
+}
+PREFIXES = {"lattice": ("k_", "lattice"), "isochromats": ("iso_", "sample")}
+
+
+def core_axis_name(direction, representation, nd=True):
+    """name of the core axis storing a dephasing direction
+
+    Args:
+        direction: "x", "y", "z" or the name of a dephasing axis
+        representation: "lattice" or "isochromats"
+        nd: False for a single default axis ("k", "iso")
+    """
+    prefix = PREFIXES[representation][0]
+    if not nd:
+        return prefix[:-1]  # "k", "iso"
+    elif representation == "lattice" and direction in ("x", "y", "z"):
+        return "k" + direction  # "kx", "ky", "kz"
+    return prefix + direction
+
+
+@functools.partial(
+    jax.tree_util.register_dataclass,
+    data_fields=["F", "Z", "density"],
+    meta_fields=["axes", "coords", "representation"],
+)
+@dataclasses.dataclass(frozen=True)
+class State:
+    """Phase states (pytree)
+
+    Attributes:
+        F: transverse states F+ (*shape, *core)
+        Z: longitudinal states (*shape, *core)
+        density: proton density, i.e. equilibrium Z(k=0) or Z of each sample (*shape)
+        axes: (static) names of the core axes
+        coords: (static) coordinates of the sample axes, as a tuple of
+            (axis name, tuple of values) pairs, e.g. isochromat positions;
+            None for lattice axes (coordinates given by the indices)
+        representation: (static) representation of the states ("lattice",
+            "isochromats"), defining the kind of the core axes
+
+    Use the constructors (`at_equilibrium`, `from_states`), which broadcast `density`.
+    """
+
+    F: jax.Array
+    Z: jax.Array
+    density: jax.Array
+    axes: tuple = ("k",)
+    coords: tuple = None
+    representation: str = "lattice"
+
+    # constructors
+
+    @classmethod
+    def at_equilibrium(cls, shape=(), nstate=0, *, axes=None, density=1, dtype=None):
+        """Init state at equilibrium (Z0 = density)
+
+        Args:
+            shape: batch shape (broadcast with density's shape)
+            nstate: number of phase states (capacity: 2 * nstate + 1), int or
+                tuple of int (one per lattice axis)
+            axes: names of the lattice axes (default: "k" for an int nstate,
+                "kx", "ky", "kz" for a tuple)
+            density: proton density (scalar or array)
+            dtype: complex dtype (default: JAX's default complex type)
+        """
+        dtype = dtype or complex
+        scalar = isinstance(nstate, numbers.Integral)
+        nstates = (int(nstate),) if scalar else tuple(int(n) for n in nstate)
+        if axes is None:
+            axes = ("k",) if scalar else LATTICE_AXES[: len(nstates)]
+        if len(axes) != len(nstates):
+            raise ValueError(f"Expected one nstate per axis {axes}, got {nstate}")
+        density = jnp.asarray(density)
+        shape = jnp.broadcast_shapes(tuple(shape), density.shape)
+        density = jnp.broadcast_to(density, shape)
+        F = jnp.zeros(shape + tuple(2 * n + 1 for n in nstates), dtype=dtype)
+        return cls(F, F, density, axes=tuple(axes)).equilibrium
+
+    @classmethod
+    def from_states(cls, states, *, density=1):
+        """Init state from a [F+, F-, Z] states array (..., 2 * nstate + 1, 3)
+
+        The F- column (states[..., 1]) is dropped: it is assumed to match F+.
+
+        Args:
+            states: [F+, F-, Z] states array
+            density: proton density (broadcastable to the batch shape)
+        """
+        states = jnp.asarray(states)
+        if states.ndim < 2 or states.shape[-1] != 3:
+            raise ValueError(
+                f"Expected states of shape (..., nk, 3), got {states.shape}"
+            )
+        elif states.shape[-2] % 2 != 1:
+            raise ValueError(f"The number of states must be odd: {states.shape[-2]}")
+        if not jnp.iscomplexobj(states):
+            states = states.astype(complex)
+        density = jnp.asarray(density)
+        shape = states.shape[:-2]
+        if jnp.broadcast_shapes(shape, density.shape) != shape:
+            raise ValueError(
+                f"Density shape {density.shape} does not broadcast to {shape}"
+            )
+        density = jnp.broadcast_to(density, shape)
+        return cls(states[..., 0], states[..., 2], density)
+
+    # properties
+
+    @property
+    def naxes(self):
+        """number of core axes"""
+        return len(self.axes)
+
+    @property
+    def shape(self):
+        """batch shape"""
+        return self.F.shape[: self.F.ndim - self.naxes]
+
+    @property
+    def ndim(self):
+        """number of batch dimensions"""
+        return len(self.shape)
+
+    @property
+    def lattice_axes(self):
+        """positions of the lattice axes in F and Z (raise if unknown axis)"""
+        return tuple(
+            self.ndim + i
+            for i, ax in enumerate(self.axes)
+            if self.axis_kind(ax) == "lattice"
+        )
+
+    @property
+    def sample_axes(self):
+        """positions of the sample axes in F and Z (raise if unknown axis)"""
+        return tuple(
+            self.ndim + i
+            for i, ax in enumerate(self.axes)
+            if self.axis_kind(ax) == "sample"
+        )
+
+    def axis_kind(self, axis):
+        """kind of a core axis in the representation ("lattice" or "sample")"""
+        return axis_kind(axis, self.representation)
+
+    def core_axis(self, direction):
+        """position (in F and Z) of the core axis storing a dephasing direction
+
+        Args:
+            direction: "x", "y", "z" or the name of a dephasing axis ("x" is
+                also the single default axis "k" or "iso")
+        """
+        names = [core_axis_name(direction, self.representation)]
+        if direction == "x":
+            names.append(core_axis_name(direction, self.representation, nd=False))
+        for name in names:
+            if name in self.axes:
+                return self.ndim + self.axes.index(name)
+        raise ValueError(
+            f"No core axis for the dephasing direction {direction!r} "
+            f"(axes: {self.axes})"
+        )
+
+    def coordinates(self, axis):
+        """coordinates of a sample axis (array of the core axis size)"""
+        coords = dict(self.coords or ())
+        if axis not in coords:
+            raise ValueError(f"No coordinates for axis {axis!r} (axes: {self.axes})")
+        return jnp.asarray(coords[axis])
+
+    @property
+    def nstates(self):
+        """number of phase states along each lattice axis"""
+        return tuple((self.F.shape[axis] - 1) // 2 for axis in self.lattice_axes)
+
+    @property
+    def nstate(self):
+        """number of phase states (single lattice axis only)"""
+        nstates = self.nstates
+        if len(nstates) != 1:
+            raise ValueError(
+                f"`nstate` requires a single lattice axis, got {self.axes}: use `nstates`"
+            )
+        return nstates[0]
+
+    @property
+    def dtype(self):
+        return self.F.dtype
+
+    @property
+    def i0(self):
+        """index of the k=0 state along the lattice axes"""
+        return tuple((self.F.shape[axis] - 1) // 2 for axis in self.lattice_axes)
+
+    @property
+    def F0(self):
+        """voxel signal: F+(k=0) on lattice axes, mean over sample axes"""
+        return self._voxel(self.F)
+
+    @property
+    def Z0(self):
+        """longitudinal magnetization: Z(k=0) on lattice axes, mean over sample axes"""
+        return self._voxel(self.Z)
+
+    @property
+    def equilibrium(self):
+        """equilibrium state with same layout and shape (F = 0, Z = density)"""
+        zeros = jnp.zeros_like(self.F)
+        return self.replace(F=zeros, Z=self.add_equilibrium(zeros))
+
+    @property
+    def Fminus(self):
+        """transverse states F-"""
+        return self.mirror(self.F)
+
+    @property
+    def states(self):
+        """[F+, F-, Z] states array (..., 2 * nstate + 1, 3)"""
+        if self.axes != ("k",):
+            raise NotImplementedError(
+                f"No [F+, F-, Z] states layout for axes: {self.axes}"
+            )
+        return jnp.stack([self.F, self.Fminus, self.Z], axis=-1)
+
+    # helpers
+
+    def expand_param(self, value, core=False):
+        """expand parameter to broadcast with the state arrays
+
+        Args:
+            value: parameter with batch axes (core=False), or with batch and core
+                axes (core=True, returned as is)
+        """
+        value = jnp.asarray(value)
+        if core:
+            return value
+        return value[(...,) + (NAX,) * self.naxes]
+
+    def prepare(self, *params, expand=True, core=False):
+        """prepare state and kernel parameters
+
+        Parameters are cast to the state's precision (real parameters to the
+        real dtype, complex parameters to the complex dtype), the state is
+        broadcast to the common batch shape (right-aligned), and parameters
+        are expanded to broadcast with the state arrays. None is passed through.
+
+        Args:
+            *params: kernel parameters (broadcastable to the batch shape) or None
+            expand: if False, parameters keep their batch shape
+            core: if True, parameters have the core axes (`(*shape, *core)`)
+
+        Returns:
+            state, list of parameters
+        """
+        rdtype = jnp.finfo(self.dtype).dtype
+        params = [
+            None
+            if param is None
+            else jnp.asarray(
+                param, dtype=self.dtype if jnp.iscomplexobj(param) else rdtype
+            )
+            for param in params
+        ]
+        ncore = self.naxes if core else 0
+        shapes = [p.shape[: p.ndim - ncore] for p in params if p is not None]
+        state = self.broadcast(*shapes)
+        if expand:
+            params = [
+                None if param is None else state.expand_param(param, core=core)
+                for param in params
+            ]
+        return state, params
+
+    def broadcast(self, *shapes):
+        """broadcast state (F, Z and density) to a common batch shape"""
+        shape = jnp.broadcast_shapes(self.shape, *shapes)
+        if shape == self.shape:
+            return self
+        core = self.F.shape[self.ndim :]
+        return self.replace(
+            F=jnp.broadcast_to(self.F, shape + core),
+            Z=jnp.broadcast_to(self.Z, shape + core),
+            density=jnp.broadcast_to(self.density, shape),
+        )
+
+    def mirror(self, arr):
+        """return conj(arr(-k)) for an array of shape (*shape, *core)"""
+        return mirror(arr, self.axes, self.representation)
+
+    def add_equilibrium(self, arr, scale=1.0):
+        """add the equilibrium magnetization (times scale) to a Z-like array
+
+        Args:
+            arr: array of shape (*shape, *core)
+            scale: factor broadcastable to the batch shape
+        """
+        value = jnp.asarray(scale) * self.density
+        value = value[(...,) + (NAX,) * len(self.sample_axes)]
+        return arr.at[(..., *self._equilibrium_index)].add(value)
+
+    def replace(self, **changes):
+        """return new state with given fields replaced"""
+        return dataclasses.replace(self, **changes)
+
+    def resize(self, nstate):
+        """pad (with zeros) or crop phase states to nstate
+
+        Args:
+            nstate: int (all lattice axes) or tuple of int (one per lattice axis)
+        """
+        nstates = self.nstates
+        if isinstance(nstate, int):
+            nstate = (nstate,) * len(nstates)
+        elif len(nstate) != len(nstates):
+            raise ValueError(f"Expected {len(nstates)} values for nstate: {nstate}")
+        F, Z = self.F, self.Z
+        for axis, new, old in zip(self.lattice_axes, nstate, nstates):
+            if new != old:
+                F = _resize(F, new - old, axis)
+                Z = _resize(Z, new - old, axis)
+        if F is self.F:
+            return self
+        return self.replace(F=F, Z=Z)
+
+    # private
+
+    @property
+    def _equilibrium_index(self):
+        """core index of the equilibrium: k=0 on lattice axes, all samples"""
+        index = []
+        for i, ax in enumerate(self.axes):
+            if self.axis_kind(ax) == "lattice":
+                index.append((self.F.shape[self.ndim + i] - 1) // 2)
+            else:
+                index.append(slice(None))
+        return tuple(index)
+
+    def _voxel(self, arr):
+        """value at k=0 on lattice axes, mean over sample axes"""
+        value = arr[(..., *self._equilibrium_index)]
+        nsample = len(self.sample_axes)
+        if nsample == 0:
+            return value
+        return value.mean(axis=tuple(range(-nsample, 0)))
+
+
+#
+# functions
+
+
+def mirror(arr, axes, representation="lattice"):
+    """return conj(arr(-k)) (reverse lattice axes and conjugate)
+
+    Sample axes are not reversed: F- = conj(F+) at the same sample.
+    """
+    offset = arr.ndim - len(axes)
+    flip = tuple(
+        offset + i
+        for i, ax in enumerate(axes)
+        if axis_kind(ax, representation) == "lattice"
+    )
+    return jnp.flip(arr, axis=flip).conj()
+
+
+def axis_kind(axis, representation="lattice"):
+    """kind of a core axis in a representation: "lattice" or "sample" (raise if invalid)"""
+    if representation not in REPRESENTATIONS:
+        raise ValueError(f"Unknown representation: {representation}")
+    kinds = REPRESENTATIONS[representation]
+    prefix, kind = PREFIXES[representation]
+    if axis in kinds:
+        return kinds[axis]
+    elif axis.startswith(prefix) and axis[len(prefix) :].isidentifier():
+        return kind  # axis of a named dephasing direction
+    raise ValueError(
+        f"Invalid core axis {axis!r} for representation {representation!r} "
+        f"(allowed: {tuple(kinds)}, or {prefix}<name>)"
+    )
+
+
+def _resize(arr, diff, axis):
+    """symmetric pad (diff > 0) or crop (diff < 0) along axis"""
+    if diff > 0:
+        pad = [(0, 0)] * arr.ndim
+        pad[axis] = (diff, diff)
+        return jnp.pad(arr, pad)
+    index = [slice(None)] * arr.ndim
+    index[axis] = slice(-diff, arr.shape[axis] + diff)
+    return arr[tuple(index)]

epgx/kernels/transition.py

@@ -0,0 +1,48 @@
+"""Transition kernels (instantaneous RF pulses)
+
+Rotation T(alpha, phi) = Rz(phi) Rx(alpha) Rz(-phi), applied to each
+[F+, F-, Z] phase state. Only the F+ and Z rows are
+computed: F- follows from the symmetry F-(k) = conj(F+(-k)).
+
+Angles are in radians.
+"""
+
+import jax.numpy as jnp
+
+
+def rotate(state, alpha, phi):
+    """RF rotation (instantaneous pulse)
+
+    Args:
+        state: State
+        alpha: flip angle (rad), broadcast with the batch shape
+        phi: phase (rad), broadcast with the batch shape
+
+    Returns:
+        new State
+    """
+    state, (alpha, phi) = state.prepare(alpha, phi)
+    F, Fm, Z = state.F, state.mirror(state.F), state.Z
+
+    cos2 = jnp.cos(alpha / 2) ** 2
+    sin2 = jnp.sin(alpha / 2) ** 2
+    cos, sin = jnp.cos(alpha), jnp.sin(alpha)
+    eiphi = jnp.exp(1j * phi)
+
+    F_ = cos2 * F + sin2 * eiphi**2 * Fm - 1j * sin * eiphi * Z
+    Z_ = 0.5j * sin * (eiphi * Fm - eiphi.conj() * F) + cos * Z
+    return state.replace(F=F_, Z=Z_)
+
+
+def phase(state, phi):
+    """phase offset (z-axis rotation): F+ -> exp(i phi) F+
+
+    Args:
+        state: State
+        phi: phase (rad), broadcast with the batch shape
+
+    Returns:
+        new State
+    """
+    state, (phi,) = state.prepare(phi)
+    return state.replace(F=jnp.exp(1j * phi) * state.F)

epgx/library/__init__.py

@@ -0,0 +1,3 @@
+"""Reference resources: commonly used sequences (later: pulses, other resources)"""
+
+from . import sequences