epgx 0.1.dev0__py3-none-any.whl

epgx/kernels/affine.py

@@ -0,0 +1,187 @@
+"""Affine maps on the magnetization states (operator merging)
+
+Several operators act on each state (each phase state k, or each isochromat)
+as an affine map on the vector v = (F+, F-, Z):
+
+    v' = M v + o * eq
+
+with M a 3x3 matrix and o an offset (recovery) multiplied by the equilibrium
+magnetization (density at k=0, or at each isochromat), both per batch element:
+arrays of shape (*shape, 3, 3) and (*shape, 3).
+
+The composition of affine maps is an affine map: consecutive operators between
+barriers (shifts, probes) can be merged into a single map (`apply_affine`).
+
+Builders return (M, o) from the kernels' (dimensionless) parameters:
+    rotation_matrix(alpha, phi)       rotate
+    phase_matrix(phi)                 phase
+    evolution_matrix(rT, rL, r0)      evolve
+    spoil_matrix()                    spoil
+    reset_matrix()                    reset
+
+`AffineAlgebra` gathers these operations for the merging pass (cf. merging).
+"""
+
+import jax
+import jax.numpy as jnp
+
+
+def apply_affine(state, matrix, offset):
+    """apply an affine map to the states: v' = M v + o * eq
+
+    Args:
+        state: State
+        matrix: (*shape, 3, 3) matrix acting on (F+, F-, Z)
+        offset: (*shape, 3) offset, times the equilibrium magnetization
+
+    Only the F+ and Z rows are applied (F- follows by symmetry).
+
+    Returns:
+        new State
+    """
+    matrix, offset = jnp.asarray(matrix), jnp.asarray(offset)
+    state = state.broadcast(matrix.shape[:-2], offset.shape[:-1])
+    coefs = [matrix[..., i, j] for i in (0, 2) for j in range(3)]
+    _, (m00, m01, m02, m20, m21, m22) = state.prepare(*coefs)
+    _, (o0, o2) = state.prepare(offset[..., 0], offset[..., 2], expand=False)
+    F, Fm, Z = state.F, state.mirror(state.F), state.Z
+    F_ = state.add_equilibrium(m00 * F + m01 * Fm + m02 * Z, o0)
+    Z_ = state.add_equilibrium(m20 * F + m21 * Fm + m22 * Z, o2)
+    return state.replace(F=F_, Z=Z_)
+
+
+def compose(*maps):
+    """composition of affine maps, applied in the given order
+
+    Args:
+        maps: (M, o) pairs, first applied first
+
+    Returns:
+        (M, o) of the composed map
+    """
+    matrix, offset = maps[0]
+    for m, o in maps[1:]:
+        # explicit elementwise products: much faster than batched 3x3 matmul
+        matrix = (m[..., :, :, None] * matrix[..., None, :, :]).sum(-2)
+        offset = (m * offset[..., None, :]).sum(-1) + o
+    return matrix, offset
+
+
+def power(affine_map, n):
+    """affine map applied n times (static n), by repeated squaring"""
+    if n < 1:
+        raise ValueError(f"Expected n >= 1, got {n}")
+    result, base = None, affine_map
+    while n:
+        if n & 1:
+            result = base if result is None else compose(result, base)
+        n >>= 1
+        if n:
+            base = compose(base, base)
+    return result
+
+
+def identity(shape=(), dtype=complex):
+    """identity map"""
+    matrix = jnp.broadcast_to(jnp.eye(3, dtype=dtype), tuple(shape) + (3, 3))
+    return matrix, jnp.zeros(tuple(shape) + (3,), dtype=dtype)
+
+
+def fold(stacked):
+    """composition of stacked maps (leading axis), in order
+
+    Sequential composition (measured faster than associative_scan or a pairwise
+    tree on CPU).
+    """
+    matrices, _ = stacked
+    init = identity(matrices.shape[1:-2], matrices.dtype)
+
+    def body(carry, affine_map):
+        return compose(carry, affine_map), None
+
+    composed, _ = jax.lax.scan(body, init, stacked)
+    return composed
+
+
+class AffineAlgebra:
+    """affine maps (M, o) as a map algebra for the merging pass (cf. merging)
+
+    Maps are built by the operators' `affine` attribute (function of the kernel
+    parameters -> map, cf. ops.Operator.affine); they do not depend on the
+    representation of the states (the run context is not used).
+    """
+
+    @staticmethod
+    def mergeable(op):
+        return getattr(op, "affine", None) is not None
+
+    @staticmethod
+    def build(op, params, context):
+        return op.affine(**params)
+
+    @staticmethod
+    def apply(state, value):
+        return apply_affine(state, *value)
+
+    compose = staticmethod(compose)
+    power = staticmethod(power)
+    fold = staticmethod(fold)
+
+
+#
+# builders
+
+
+def rotation_matrix(alpha, phi):
+    """affine map of `rotate`"""
+    alpha, phi = jnp.broadcast_arrays(jnp.asarray(alpha), jnp.asarray(phi))
+    cos2, sin2 = jnp.cos(alpha / 2) ** 2, jnp.sin(alpha / 2) ** 2
+    sin, cos = jnp.sin(alpha), jnp.cos(alpha)
+    e1, e2 = jnp.exp(1j * phi), jnp.exp(2j * phi)
+    rows = [
+        [cos2 + 0j, sin2 * e2, -1j * sin * e1],
+        [sin2 * e2.conj(), cos2 + 0j, 1j * sin * e1.conj()],
+        [-0.5j * sin * e1.conj(), 0.5j * sin * e1, cos + 0j],
+    ]
+    matrix = jnp.stack([jnp.stack(row, axis=-1) for row in rows], axis=-2)
+    return matrix, jnp.zeros(matrix.shape[:-1], dtype=matrix.dtype)
+
+
+def phase_matrix(phi):
+    """affine map of `phase`"""
+    e1 = jnp.exp(1j * jnp.asarray(phi))
+    return _diagonal(e1, e1.conj(), jnp.ones_like(e1))
+
+
+def evolution_matrix(rT=0, rL=0, r0=None):
+    """affine map of `evolve` (exponents broadcast with the batch shape)"""
+    rT, rL = jnp.asarray(rT) + 0j, jnp.asarray(rL)
+    shape = jnp.broadcast_shapes(
+        rT.shape, rL.shape, jnp.shape(r0) if r0 is not None else ()
+    )
+    eT = jnp.broadcast_to(jnp.exp(-rT), shape)
+    matrix, offset = _diagonal(
+        eT.conj(), eT, jnp.broadcast_to(jnp.exp(-rL), shape) + 0j
+    )
+    if r0 is not None:
+        recovery = jnp.broadcast_to(-jnp.expm1(-jnp.asarray(r0)), shape)
+        offset = offset.at[..., 2].set(recovery)
+    return matrix, offset
+
+
+def spoil_matrix():
+    """affine map of `spoil` (F = 0)"""
+    return _diagonal(0j, 0j, 1 + 0j)
+
+
+def reset_matrix():
+    """affine map of `reset` (return to equilibrium)"""
+    matrix = jnp.zeros((3, 3), dtype=complex)
+    return matrix, jnp.array([0, 0, 1], dtype=complex)
+
+
+def _diagonal(d0, d1, d2):
+    d0, d1, d2 = jnp.broadcast_arrays(jnp.asarray(d0), jnp.asarray(d1), jnp.asarray(d2))
+    matrix = jnp.zeros(d0.shape + (3, 3), dtype=jnp.result_type(d0, d1, d2, complex))
+    matrix = matrix.at[..., 0, 0].set(d0).at[..., 1, 1].set(d1).at[..., 2, 2].set(d2)
+    return matrix, jnp.zeros(d0.shape + (3,), dtype=matrix.dtype)

epgx/kernels/evolution.py

@@ -0,0 +1,40 @@
+"""Evolution kernels (relaxation, precession, recovery)
+
+Dimensionless exponents (cf. the `R(rT, rL, r0)` operator):
+    F-  -> exp(-rT) F-, i.e. F+ -> exp(-conj(rT)) F+
+    Z   -> exp(-rL) Z + (1 - exp(-r0)) * equilibrium
+
+For E(tau, T1, T2, g): rT = tau * (1/T2 + 2i pi g), rL = r0 = tau / T1.
+
+Exponents may vary along the core axes (`core=True`), e.g. a shift on
+isochromats is a precession depending on their positions: rT = 2i pi k x.
+"""
+
+import jax.numpy as jnp
+
+
+def evolve(state, rT=0, rL=0, r0=None, *, core=False):
+    """relaxation / precession / recovery
+
+    Args:
+        state: State
+        rT: transverse exponent (complex: relaxation + precession)
+        rL: longitudinal exponent (real)
+        r0: recovery exponent (real), None: no recovery
+        core: if False, exponents are broadcast with the batch shape; if True,
+            with the batch and core axes (`(*shape, *core)`)
+
+    Returns:
+        new State
+    """
+    state, (rT, rL, r0) = state.prepare(rT, rL, r0, expand=False, core=core)
+    F = state.expand_param(jnp.exp(-jnp.conj(rT)), core=core) * state.F
+    Z = state.expand_param(jnp.exp(-rL), core=core) * state.Z
+    if r0 is not None:
+        # 1 - exp(-r0), accurate for small r0
+        recovery = -jnp.expm1(-r0)
+        if core:
+            Z = Z + recovery * state.add_equilibrium(jnp.zeros_like(Z))
+        else:
+            Z = state.add_equilibrium(Z, recovery)
+    return state.replace(F=F, Z=Z)

epgx/kernels/operator.py

@@ -0,0 +1,32 @@
+"""Spoiler and reset kernels"""
+
+import jax.numpy as jnp
+
+
+def spoil(state):
+    """perfect spoiler: set transverse states to 0"""
+    return state.replace(F=jnp.zeros_like(state.F))
+
+
+def spoil_axis(state, direction):
+    """keep the states of index 0 along a dephasing direction (F and Z)
+
+    States dephased along the direction are removed: on the lattice, all
+    states with a nonzero index along its axis; on samples (isochromats), the
+    magnetization is replaced by its mean along the axis.
+    """
+    axis = state.core_axis(direction)
+    if state.axis_kind(state.axes[axis - state.ndim]) == "sample":
+        F = jnp.broadcast_to(state.F.mean(axis=axis, keepdims=True), state.F.shape)
+        Z = jnp.broadcast_to(state.Z.mean(axis=axis, keepdims=True), state.Z.shape)
+        return state.replace(F=F, Z=Z)
+    size = state.F.shape[axis]
+    mask = (jnp.arange(size) == size // 2).reshape(
+        (size,) + (1,) * (state.F.ndim - axis - 1)
+    )
+    return state.replace(F=state.F * mask, Z=state.Z * mask)
+
+
+def reset(state):
+    """return to equilibrium (capacity and batch shape are kept)"""
+    return state.equilibrium

epgx/kernels/probe.py

@@ -0,0 +1,107 @@
+"""Probe kernels (signal acquisition)
+
+    readout: quantity probed from the state (F0, Z0, F, Fminus, Z, states)
+    reduce: weighted sum over batch axes (static indices)
+    demodulate: removal of the receiver phase (scanner convention)
+    adc: readout -> weights -> reduce -> demodulate
+
+Axis indices refer to the batch axes; values may have `ncore` trailing axes
+(e.g. core axes of F and Z), which are never reduced. Readouts always keep the
+batch axes first, so `ncore` is derived from the shapes (layout independent).
+Weights and phases are broadcast with the batch axes (right-aligned).
+
+ADC phase follows the scanner convention: the ADC phase equals the phase of the
+RF pulse to compensate, and the signal is multiplied by exp(-i phase).
+"""
+
+import jax.numpy as jnp
+
+NAX = jnp.newaxis
+
+ATTRIBUTES = ["F0", "Z0", "F", "Fminus", "Z", "states"]
+
+
+def readout(state, attr="F0"):
+    """quantity probed from the state
+
+    Args:
+        state: State
+        attr: "F0", "Z0" (batch shape), "F", "Fminus", "Z" (batch and core axes),
+            "states" ([F+, F-, Z] layout, 1d lattice only)
+    """
+    if attr not in ATTRIBUTES:
+        raise ValueError(f"Unknown attribute: {attr}, expected one of {ATTRIBUTES}")
+    return getattr(state, attr)
+
+
+def reduce(value, axes=None, weights=None, *, ncore=0):
+    """weighted sum over batch axes
+
+    Args:
+        value: array of shape (*shape, *core) with len(core) == ncore
+        axes: None (no sum, or weights' axes if weights is set), True (all batch axes),
+            int or tuple of int (batch axes, negative: from the last batch axis)
+        weights: weights broadcast with the batch axes (right-aligned)
+        ncore: number of trailing non-batch axes in value
+
+    Returns:
+        array
+    """
+    if weights is not None:
+        weights = _cast(weights, value)
+        value = value * weights[(...,) + (NAX,) * ncore]
+        if axes is None:
+            # sum along the weights' axes (trailing batch axes)
+            axes = tuple(range(-weights.ndim, 0))
+    if axes is None or axes is False:
+        return value
+    ndim = value.ndim - ncore
+    if axes is True:
+        axes = tuple(range(ndim))
+    axes = (axes,) if isinstance(axes, int) else tuple(axes)
+    if not all(-ndim <= ax < ndim for ax in axes):
+        raise ValueError(f"Invalid axes {axes} for {ndim} batch dimension(s)")
+    return value.sum(axis=tuple(ax % ndim for ax in axes))
+
+
+def demodulate(value, phase, *, ncore=0):
+    """remove receiver phase: value * exp(-i phase)
+
+    Args:
+        value: array of shape (*shape, *core) with len(core) == ncore
+        phase: phase (rad), broadcast with the batch axes (right-aligned)
+        ncore: number of trailing non-batch axes in value
+    """
+    phase = _cast(phase, value, real=True)
+    return value * jnp.exp(-1j * phase)[(...,) + (NAX,) * ncore]
+
+
+def adc(state, attr="F0", *, phase=None, weights=None, axes=None):
+    """acquire signal: readout -> weights -> reduce -> demodulate
+
+    Args:
+        state: State
+        attr: probed attribute (cf. `readout`)
+        phase: ADC phase (rad), applied after reduction (scanner convention)
+        weights, axes: cf. `reduce`
+    """
+    value = readout(state, attr)
+    # readouts keep the batch axes first: any trailing axis is a non-batch axis
+    ncore = value.ndim - state.ndim
+    value = reduce(value, axes, weights, ncore=ncore)
+    if phase is not None:
+        value = demodulate(value, phase, ncore=ncore)
+    return value
+
+
+#
+# functions
+
+
+def _cast(param, value, real=False):
+    """cast parameter to the precision of value (keep real/complex kind)"""
+    dtype = jnp.result_type(value)
+    rdtype = jnp.finfo(dtype).dtype if jnp.issubdtype(dtype, jnp.inexact) else dtype
+    if real or not jnp.iscomplexobj(param):
+        return jnp.asarray(param, dtype=rdtype)
+    return jnp.asarray(param, dtype=dtype)

epgx/kernels/shift.py

@@ -0,0 +1,106 @@
+"""Shift kernels (gradient dephasing), named by state representation
+
+    shift_lattice: integer shift on the dense lattice (1d: regular EPG,
+        nd: 3d gradient extension)
+    shift_samples: samples of the magnetization (isochromats), any k: a shift is
+        a precession depending on the position x, M+ -> exp(2i pi k x) M+
+    (future) shift_slots: sparse slots with stored coordinates (integer or
+        float k, accumulated time); may be split by
+        merge rule (exact, grid, moments)
+
+Shifting F+(k) -> F+(k + s) also shifts F-(k) -> F-(k - s), by symmetry
+F-(k) = conj(F+(-k)); Z states are not shifted.
+"""
+
+import jax.numpy as jnp
+import numpy as np
+
+from .evolution import evolve
+
+
+def shift_samples(state, k):
+    """shift of the transverse magnetization of samples (e.g. isochromats)
+
+    A shift is a precession depending on the position x of each sample
+    (state coordinates, unit: period of a unit shift): M+ -> exp(2i pi k.x) M+,
+    with one sample axis per dephasing direction.
+
+    Args:
+        state: State with sample axes (with coordinates)
+        k: dict {direction: number} ("x", "y", "z" or a named dephasing axis;
+            any real values), or a number (along "x") or a tuple (x, y, z)
+
+    Returns:
+        new State
+    """
+    phase = 0.0
+    core = state.F.shape[state.ndim :]
+    for direction, s in shifts_by_direction(k).items():
+        if s == 0:
+            continue
+        axis = state.core_axis(direction) - state.ndim
+        x = state.coordinates(state.axes[axis])
+        shape = [1] * len(core)
+        shape[axis] = x.size
+        phase = phase + s * x.reshape(shape)  # grid of positions (k.x)
+    return evolve(state, 2j * jnp.pi * phase, core=True)
+
+
+def shift_lattice(state, k):
+    """integer shift of the transverse states on the (fixed-capacity) lattice
+
+    States shifted beyond the capacity are dropped.
+
+    Args:
+        state: State with lattice axes
+        k: dict {direction: static int} ("x", "y", "z" or a named dephasing
+            axis), or an int (along "x") or a tuple of int (x, y, z)
+
+    Returns:
+        new State
+    """
+    F = state.F
+    for direction, s in shifts_by_direction(k).items():
+        if not _isint(s):
+            raise TypeError(f"Shifts must be static integers, got: {k}")
+        elif s == 0:
+            continue
+        F = _shift(F, int(s), state.core_axis(direction))
+    return state.replace(F=F)
+
+
+def shifts_by_direction(k):
+    """shift components as a dict {direction: value}
+
+    A number is a shift along "x", a tuple the components along "x", "y", "z".
+    """
+    if isinstance(k, dict):
+        return dict(k)
+    k = tuple(k) if isinstance(k, (tuple, list)) else (k,)
+    if len(k) > 3:
+        raise ValueError(f"Shifts have 1 to 3 components (x, y, z), got: {k}")
+    return dict(zip(("x", "y", "z"), k))
+
+
+#
+# functions
+
+
+def _isint(value):
+    return isinstance(value, (int, np.integer)) and not isinstance(value, bool)
+
+
+def _shift(arr, shift, axis):
+    """shift array along axis with zero fill (out-of-range values are dropped)"""
+    size = arr.shape[axis]
+    if shift == 0:
+        return arr
+    elif abs(shift) >= size:
+        return jnp.zeros_like(arr)
+    index = [slice(None)] * arr.ndim
+    pad = [(0, 0)] * arr.ndim
+    if shift > 0:
+        index[axis], pad[axis] = slice(0, size - shift), (shift, 0)
+    else:
+        index[axis], pad[axis] = slice(-shift, size), (0, -shift)
+    return jnp.pad(arr[tuple(index)], pad)