epgx 0.1.dev0__py3-none-any.whl

epgx/library/sequences.py

@@ -0,0 +1,210 @@
+"""Reference sequences
+
+Commonly used sequences, defined once with `@sequence`, for tests, examples
+and benchmarks where the sequence is not the focus. Conventions:
+
+    - angles in degrees, times in ms; durations are set on the operators, so
+      the ADC times are meaningful (`output="times"`);
+    - relaxation and frequency offsets from the `Sample` (T1, T2, df), B1 and
+      B0 from the `Scanner`;
+    - numeric parameters (angles, timings, trains) are declared variables,
+      structural sizes (numbers of echoes, dummy repetitions) are static;
+    - the ADC phase follows the RF phase (RF-spoiled and phase-cycled
+      sequences are demodulated).
+
+Each sequence gives the signal when called (`mse(esp=10.0, sample=...)`),
+derivatives and statistics by name (`mse.jacobian(["T2"], esp=10.0, ...)`),
+and the layer-2 list with `.build(...)`. A multi-point sequence (e.g. pSPGR)
+is a batch axis of the variables: `spgr(alpha=Axis("point", alphas),
+phi=Axis("point", phis), ...)`.
+
+Preparation modules (`inversion`, `t2prep`) are plain functions returning
+lists of operators, to compose into sequences.
+"""
+
+import jax.numpy as jnp
+
+from .. import ops
+from ..sequences import sequence
+
+#
+# spin echoes
+
+
+@sequence(variables=["esp", "alpha"])
+def mse(esp, alpha=180.0, necho=10):
+    """multi spin echo (CPMG): 90y excitation, `necho` refocusing pulses (x)
+
+    Args:
+        esp: echo spacing (ms)
+        alpha: refocusing flip angle (degrees)
+        necho: (static) number of echoes
+    """
+    shift, relax = ops.S(1), ops.E(esp / 2, duration=True)
+    block = [shift, relax, ops.T(alpha, 0), shift, relax, ops.ADC]
+    return [ops.T(90, 90)] + [block] * necho
+
+
+@sequence(variables=["esp", "alpha[n]"])
+def tse(esp, alpha):
+    """turbo spin echo with a train of refocusing angles (one echo per angle)
+
+    Args:
+        esp: echo spacing (ms)
+        alpha: refocusing flip angles (degrees), one per echo
+    """
+    shift, relax = ops.S(1), ops.E(esp / 2, duration=True)
+    return [ops.T(90, 90)] + [
+        [shift, relax, ops.T(a, 0), shift, relax, ops.ADC] for a in alpha
+    ]
+
+
+@sequence(variables=["alpha", "esp"])
+def hyperecho(alpha, esp=10.0, n=20):
+    """hyperecho: 90y, n pulses (alpha, x), 180x, n pulses (-alpha, x)
+
+    The hyperecho forms at the last ADC (full refocusing without relaxation).
+
+    Args:
+        alpha: flip angle of the trains (degrees)
+        esp: pulse spacing (ms)
+        n: (static) number of pulses of each train
+    """
+    shift, relax = ops.S(1), ops.E(esp, duration=True)
+    first = [[ops.T(alpha, 0), shift, relax, ops.ADC]] * n
+    second = [[ops.T(-alpha, 0), shift, relax, ops.ADC]] * n
+    return (
+        [ops.T(90, 90), shift, relax] + first + [ops.T(180, 0), shift, relax] + second
+    )
+
+
+#
+# gradient echoes
+
+
+@sequence(variables=["alpha", "TR", "TE", "phi"])
+def spgr(alpha, TR, TE, phi=117.0, ndummy=200, nread=1):
+    """RF-spoiled gradient echo (quadratic phase), acquired at steady state
+
+    Phase of pulse n: phi * n * (n + 1) / 2 (increment growing by phi); the
+    ADC phase follows the RF phase. Small `phi` gives partially spoiled
+    sequences (pSPGR).
+
+    Args:
+        alpha: flip angle (degrees)
+        TR, TE: repetition and echo times (ms)
+        phi: phase increment step (degrees)
+        ndummy: (static) number of dummy repetitions
+        nread: (static) number of acquired repetitions (after the dummies)
+    """
+    n = jnp.arange(ndummy + nread)
+    phases = phi * n * (n + 1) / 2
+    rlx1, rlx2 = ops.E(TE, duration=True), ops.E(TR - TE, duration=True)
+    dummies = [[ops.T(alpha, p), rlx1, rlx2, ops.S(1)] for p in phases[:ndummy]]
+    reads = [
+        [ops.T(alpha, p), rlx1, ops.Adc(phase=p), rlx2, ops.S(1)]
+        for p in phases[ndummy:]
+    ]
+    return dummies + reads
+
+
+@sequence(variables=["alpha", "TR"])
+def bssfp(alpha, TR, ndummy=200, nread=1):
+    """balanced SSFP: phase alternation (0, 180), echo at TE = TR / 2
+
+    Off-resonance from the sample (`df`) or the scanner (`B0`); the ADC phase
+    follows the RF phase.
+
+    Args:
+        alpha: flip angle (degrees)
+        TR: repetition time (ms)
+        ndummy: (static) number of dummy repetitions
+        nread: (static) number of acquired repetitions (after the dummies)
+    """
+    half = ops.E(TR / 2, duration=True)
+    blocks = []
+    for n in range(ndummy + nread):
+        phase = 180.0 * (n % 2)
+        adc = [ops.Adc(phase=phase)] if n >= ndummy else []
+        blocks.append([ops.T(alpha, phase), half, *adc, half])
+    return blocks
+
+
+@sequence(variables=["alpha", "TR", "TE"])
+def dess(alpha, TR, TE, ndummy=200, nread=1):
+    """double echo steady state: FID and echo of an unbalanced gradient echo
+
+    The FID is acquired TE after the pulse (probe "fid"), the echo TE before
+    the next pulse (probe "echo"); `group="name"` separates them.
+
+    Args:
+        alpha: flip angle (degrees)
+        TR, TE: repetition time and echo time of the FID (ms)
+        ndummy: (static) number of dummy repetitions
+        nread: (static) number of acquired repetitions (after the dummies)
+    """
+    rf, shift = ops.T(alpha, 0), ops.S(1)
+    rlx1, rlx2 = ops.E(TE, duration=True), ops.E(TR - 2 * TE, duration=True)
+    fid, echo = ops.Adc(name="fid"), ops.Adc(name="echo")
+    dummies = [[rf, rlx1, shift, rlx2, rlx1]] * ndummy
+    reads = [[rf, rlx1, fid, shift, rlx2, echo, rlx1]] * nread
+    return dummies + reads
+
+
+@sequence(variables=["alpha[n]", "TR[n]", "TE", "TI"])
+def mrf_fisp(alpha, TR, TE=3.0, TI=20.0, phi=90.0):
+    """MR fingerprinting (FISP): inversion, then a train of (alpha_i, TR_i)
+
+    Args:
+        alpha: flip angles (degrees), one per repetition
+        TR: repetition times (ms), one per repetition
+        TE: echo time (ms)
+        TI: inversion delay (ms)
+        phi: (static) RF phase (degrees)
+    """
+    rlx1 = ops.E(TE, duration=True)
+    return [ops.T(180, phi), ops.E(TI, duration=True)] + [
+        [ops.T(a, phi), rlx1, ops.ADC, ops.E(tr - TE, duration=True), ops.S(1)]
+        for a, tr in zip(alpha, TR)
+    ]
+
+
+@sequence(variables=["TE", "esp", "alpha"])
+def epi(TE, esp=0.5, alpha=90.0, nlines=64):
+    """single-shot gradient echo EPI, at the signal level (no encoding)
+
+    One ADC per line, every `esp`, the center line at TE (requires TE >=
+    (nlines // 2) * esp). Intra-voxel dephasing (T2*) can be modelled by a
+    frequency distribution of the sample (`df` on an axis), summed with
+    `probe=Adc(reduce=...)`.
+
+    Args:
+        TE: echo time of the center line (ms)
+        esp: line spacing (ms)
+        alpha: flip angle (degrees)
+        nlines: (static) number of lines
+    """
+    first = TE - (nlines // 2) * esp
+    lines = [[ops.E(esp, duration=True), ops.ADC]] * (nlines - 1)
+    return [ops.T(alpha, 90), ops.E(first, duration=True), ops.ADC] + lines
+
+
+#
+# preparation modules (lists of operators)
+
+
+def inversion(TI, phi=0.0, spoil=True):
+    """inversion module: 180 pulse, spoiler (optional), delay TI (ms)"""
+    spoiler = [ops.SPOILER] if spoil else []
+    return [ops.T(180, phi)] + spoiler + [ops.E(TI, duration=True)]
+
+
+def t2prep(TE, nrefocus=4):
+    """T2 preparation: 90x, `nrefocus` 180y pulses over TE (ms), -90x, spoiler
+
+    Stores exp(-TE / T2) of the longitudinal magnetization (on resonance,
+    without T1 recovery).
+    """
+    half = ops.E(TE / nrefocus / 2, duration=True)
+    refocus = [[half, ops.T(180, 90), half]] * nrefocus
+    return [ops.T(90, 0)] + refocus + [ops.T(-90, 0), ops.SPOILER]

epgx/lowering.py

@@ -0,0 +1,95 @@
+"""EPG front-end of the compiler: operators -> steps, compile_sequence
+
+The compiler (cf. compiler) is generic; this module provides what is specific
+to operators:
+    - `lower(op)`: Step applying the operator's kernel with its normalized
+      parameters, from the run context (sample, scanner); probes acquire;
+    - `compile_sequence`: compiler with this lowering and, with `merge=True`,
+      the merging pass over affine maps (cf. merging, kernels.affine).
+
+Programs compiled here run with `program.run(state, Context(sample, scanner,
+method))`; a tuple (sample, scanner) or None is also accepted. The method
+(static) gives map builders access to the representation of the states (e.g.
+the coordinates of a phase grid, for maps of shifts).
+"""
+
+from . import compiler, merging
+from .common import pytree, static
+from .kernels.affine import AffineAlgebra
+from .ops import Operator, Probe
+
+
+@pytree
+class Context:
+    """Run context of a compiled sequence
+
+    Attributes:
+        sample, scanner: Sample and Scanner (or None)
+        method: (static) representation of the states (e.g. `Lattice(...)`)
+    """
+
+    sample: object = None
+    scanner: object = None
+    method: object = static()
+
+    def __init__(self, sample=None, scanner=None, method=None):
+        object.__setattr__(self, "sample", sample)
+        object.__setattr__(self, "scanner", scanner)
+        object.__setattr__(self, "method", method)
+
+
+def lower(op):
+    """Step of an operator (run context: Context, (sample, scanner) or None)"""
+    if not isinstance(op, Operator):
+        raise ValueError(f"Invalid operator: {op}")
+
+    def params(context):
+        return op.normalize(*_unpack(context))
+
+    acquire = None
+    if isinstance(op, Probe):
+
+        def acquire(state, context):
+            return op.acquire(state, *_unpack(context))
+
+    return compiler.Step(op.kernel, params, acquire, repr(op), source=op)
+
+
+def compile_sequence(
+    sequence,
+    *,
+    max_unroll=compiler.MAX_UNROLL,
+    checkpoint=False,
+    merge=False,
+    passes=(),
+):
+    """compile a (nested) list of operators (cf. compiler.build)
+
+    Args:
+        sequence: (nested) list of operators; repeated blocks are consecutive
+            sub-lists with the same structure (e.g. `[exc] + [block] * N`)
+        max_unroll, checkpoint: cf. compiler.build
+        merge: merge consecutive affine operators (RF pulses, relaxation,
+            precession, spoilers) into single maps, including repeated blocks
+        passes: other passes (Program -> Program), applied after merging
+
+    Returns:
+        Program, run with `program.run(state, Context(sample, scanner, method))`
+    """
+    passes = ((merging.merge(AffineAlgebra),) if merge else ()) + tuple(passes)
+    return compiler.build(
+        sequence,
+        lower=lower,
+        passes=passes,
+        max_unroll=max_unroll,
+        checkpoint=checkpoint,
+    )
+
+
+def _unpack(context):
+    """(sample, scanner) from the run context"""
+    if context is None:
+        return None, None
+    elif isinstance(context, Context):
+        return context.sample, context.scanner
+    return context

epgx/merging.py

@@ -0,0 +1,137 @@
+"""Merging pass: consecutive operators composed into single maps
+
+A compiler pass (cf. compiler) parameterized by a map algebra, e.g.
+`kernels.affine.AffineAlgebra` (affine maps on (F+, F-, Z)). An algebra provides:
+
+    mergeable(op): True if the operator has a map in this algebra
+    build(op, params, context): map of an operator from its kernel parameters
+        and the run context (e.g. the method, for maps depending on the
+        representation of the states)
+    apply(state, map): kernel applying a map to the states
+    compose(*maps): composition, first applied first
+    power(map, n): map applied n times (static n)
+    fold(stacked): composition of stacked maps (leading axis), in order
+
+The pass, applied to each level of the program (top level and scan bodies):
+    - reduces repeated blocks made of a single mergeable step (after merging
+      of the body) to one step: `power` for identical blocks, maps of all
+      blocks then `fold` otherwise;
+    - composes runs of consecutive mergeable steps (no acquisition) into one
+      step; other steps (shifts, probes) are barriers.
+
+A merged step is an ordinary step: kernel `algebra.apply`, parameters the
+composed map (computed from the run context when the step runs).
+"""
+
+import dataclasses
+import typing
+
+import jax
+
+from .compiler import Program, Scan, Step, Unrolled, join_unrolled
+
+
+@dataclasses.dataclass(frozen=True)
+class Merged:
+    """source of a merged step: its map as a function of the run context"""
+
+    map: typing.Callable
+
+
+def merge(algebra):
+    """merging pass (Program -> Program) for a map algebra"""
+
+    def merge_pass(program):
+        segments = [_reduce_scan(segment, algebra) for segment in program.segments]
+        segments = [
+            Unrolled(_merge_steps(segment.steps, algebra))
+            if isinstance(segment, Unrolled)
+            else segment
+            for segment in join_unrolled(segments)
+        ]
+        return Program(segments)
+
+    return merge_pass
+
+
+def map_of(step, algebra):
+    """map of a step as a function of the context, or None (not mergeable)"""
+    if isinstance(step.source, Merged):
+        return step.source.map
+    elif step.acquire is not None or not algebra.mergeable(step.source):
+        return None
+    return lambda context: algebra.build(step.source, step.params(context), context)
+
+
+def merged_step(map_fn, algebra, label, count):
+    """step applying the map map_fn(context)"""
+
+    def kernel(state, value):
+        return algebra.apply(state, value)
+
+    def params(context):
+        return {"value": map_fn(context)}
+
+    return Step(kernel, params, None, label, count, Merged(map_fn))
+
+
+#
+# private
+
+
+def _merge_steps(steps, algebra):
+    """compose runs of consecutive mergeable steps"""
+    merged, run = [], []
+
+    def flush():
+        if len(run) > 1:
+            maps = [map_of(step, algebra) for step in run]
+
+            def composed(context, maps=maps):
+                return algebra.compose(*(map_fn(context) for map_fn in maps))
+
+            count = sum(step.count for step in run)
+            merged.append(merged_step(composed, algebra, f"merged({len(run)})", count))
+        else:
+            merged.extend(run)
+        run.clear()
+
+    for step in steps:
+        if map_of(step, algebra) is not None:
+            run.append(step)
+        else:
+            flush()
+            merged.append(step)
+    flush()
+    return merged
+
+
+def _single_step(program):
+    """the only step of a program, or None"""
+    if len(program.segments) != 1 or not isinstance(program.segments[0], Unrolled):
+        return None
+    steps = program.segments[0].steps
+    return steps[0] if len(steps) == 1 else None
+
+
+def _reduce_scan(segment, algebra):
+    """Scan of blocks reduced to a single mergeable step -> one merged step"""
+    if not isinstance(segment, Scan):
+        return segment
+    step = _single_step(segment.body)
+    if step is None or map_of(step, algebra) is None:
+        return segment
+
+    def block_map(xs, context):
+        return map_of(_single_step(segment.body_at(xs)), algebra)(context)
+
+    def scan_map(context):
+        if segment.xs is None:
+            # same block repeated: power (log2(length) compositions)
+            return algebra.power(block_map(None, context), segment.length)
+        # maps of all blocks (vectorized), then composition in order
+        stacked = jax.vmap(lambda xs: block_map(xs, context))(segment.xs)
+        return algebra.fold(stacked)
+
+    step = merged_step(scan_map, algebra, f"merged({segment!r})", segment.count)
+    return Unrolled([step])