epgx 0.1.dev0__py3-none-any.whl

epgx/functions.py

@@ -0,0 +1,602 @@
+"""Sequence functions
+
+    simulate = resolve (axes) -> method (representation) -> compile -> run
+    trace = simulate with a snapshot of the states after every operator
+
+`simulate` and `trace` are pure and jit-able (except for the "numpy" and
+"xarray" outputs). `simulate` runs a cached compiled program (`jit=True`):
+the sequence is flattened preserving the identity of shared operators
+(`ops.flatten_operators`), and its structure and the static options are the
+cache key, so that new values (e.g. other T2 or flip angles) do not recompile.
+Inside a user's `jax.jit`, build the sequence in the jitted function (or close
+over it) rather than passing it as an argument: JAX's own flattening does not
+preserve identity (correct results, shared operators stacked in scans).
+Sequence helpers (flatten_sequence, getnshift, get_adc_times) are defined in
+`ops.functions` and re-exported here.
+"""
+
+import dataclasses
+import functools
+import operator as _operator
+import warnings
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from . import axes, compiler, lowering, methods
+from .common import pytree, static
+from .kernels import probe as kprobe
+from .ops import (
+    Adc,
+    Operator,
+    Probe,
+    Probes,
+    check_conflicts,
+    flatten_operators,
+    flatten_sequence,
+    get_adc_times,
+    get_dephasing_axes,
+    getnshift,  # noqa: F401 (re-export)
+    unflatten_operators,
+)
+from .progress import Progress
+
+OUTPUTS = ("jax", "numpy", "layout", "times", "xarray")
+GROUPS = (None, "name")
+
+# probe of the full states, inserted after every operator by `trace`
+SNAPSHOT = Probes([Adc("F"), Adc("Z")], name="SNAPSHOT")
+
+
+@dataclasses.dataclass(frozen=True)
+class OutputLayout:
+    """Labels of a simulated output
+
+    Attributes:
+        dims: axis names ("adc" first, then batch and core axes)
+        shape: axis sizes
+        coords: coordinates of labelled axes (dict: name -> values)
+        names: name of the probe of each ADC (if "adc" is the first axis)
+    """
+
+    dims: tuple
+    shape: tuple
+    coords: dict
+    names: tuple = None
+
+    def __repr__(self):
+        axes = ", ".join(f"{d}={s}" for d, s in zip(self.dims, self.shape))
+        return f"OutputLayout({axes})"
+
+
+def simulate(
+    sequence,
+    sample=None,
+    scanner=None,
+    *,
+    probe=None,
+    output="jax",
+    group=None,
+    adc_time=False,
+    method=None,
+    init=None,
+    dtype=None,
+    dims=None,
+    alignment=None,
+    max_unroll=compiler.MAX_UNROLL,
+    checkpoint=False,
+    merge=True,
+    jit=True,
+    disp=False,
+):
+    """simulate a sequence
+
+    Args:
+        sequence: (nested) list of operators (repeated blocks: sub-lists with
+            the same structure, cf. compiler)
+        sample: Sample (tissue properties)
+        scanner: Scanner (hardware properties)
+        probe: replaces the acquisition of all probes (ADC) of the sequence:
+            attribute name (e.g. "Z0", "F"), Probe operator, or tuple of these
+            (the acquired value is then a tuple). An `Adc` sets what is
+            measured (attribute, weights, reduction) and each ADC of the
+            sequence keeps its phase (a phase given by both: error)
+        output: name or tuple of names of the returned objects:
+            "jax": acquired values, jax array(s) of shape (nadc, *dims) (default)
+            "numpy": same, as numpy arrays (not under jit)
+            "layout": OutputLayout (dims, shape, coords) of the values (without
+                simulation if no values are requested)
+            "times": ADC times (cf. get_adc_times)
+            "xarray": xarray.DataArray (requires xarray, not under jit)
+        group: arrangement of every returned object:
+            None: in sequence order (default)
+            "name": dict by probe name, in order of first appearance (e.g.
+                DESS with `Adc(name="fid")` and `Adc(name="echo")`); each
+                group is arranged as a whole output
+        adc_time: if True, also return the ADC times first, i.e.
+            output=("times", output)
+        method: representation of the states (e.g. `Lattice(nmax=...)`),
+            default: inferred from the sequence
+        init: initial State (default: equilibrium with density sample.PD)
+        dtype: complex dtype of the states (default: JAX's default)
+        dims, alignment: order of the axes, alignment of unnamed arrays (cf.
+            axes.resolve)
+        max_unroll, checkpoint: cf. lowering.compile_sequence
+        merge: merge consecutive affine operators (RF pulses, relaxation...)
+            into single maps, cf. lowering.compile_sequence (results are
+            identical up to rounding)
+        jit: run a cached compiled program (compiled again only if the
+            structure of the sequence or the static options change); False:
+            eager execution (e.g. for debugging)
+        disp: display a progress bar (cf. epgx.progress; compilation not shown)
+
+    Returns:
+        object(s) selected by `output`. If the acquired values do not all have
+        the same shape (e.g. ADCs with different reductions), they are returned
+        as a list (one value per ADC); naming the probes and using
+        `group="name"` keeps values of the same shape together.
+    """
+    names = (output,) if isinstance(output, str) else tuple(output)
+    if adc_time:
+        names = ("times",) + names
+    if not set(names) <= set(OUTPUTS):
+        raise ValueError(f"Unknown output(s): {set(names) - set(OUTPUTS)}")
+    if group not in GROUPS:
+        raise ValueError(f"Unknown group: {group}, expected one of {GROUPS}")
+
+    # probes
+    sequence = replace_probes(sequence, probe)
+    probes = [op for op in flatten_sequence(sequence) if isinstance(op, Probe)]
+    if not probes:
+        raise ValueError("Cannot simulate a sequence without any probe (ADC)")
+    groups = _group_by_name(probes) if group == "name" else {None: None}
+
+    check_sample(sequence, sample)
+    get_dephasing_axes(sequence)  # consistent axes (concrete scales, before jit)
+
+    # static options (cache key of the compiled program)
+    method = method or methods.infer_method(sequence)
+    options = (
+        ("method", method),
+        ("dtype", dtype),
+        ("dims", None if dims is None else tuple(dims)),
+        ("alignment", alignment),
+        ("max_unroll", max_unroll),
+        ("checkpoint", checkpoint),
+        ("merge", merge),
+        ("disp", disp),
+    )
+    leaves, skeleton = flatten_operators(sequence)
+    execute = functools.partial(
+        _execute_jit if jit else _execute, skeleton=skeleton, options=options
+    )
+    if set(names) & {"jax", "numpy", "xarray"}:
+        values, times = execute(leaves, sample, scanner, init)
+    else:
+        # labels and times only: shapes of the values, without simulation
+        values, _ = jax.eval_shape(execute, leaves, sample, scanner, init)
+        times = get_adc_times(sequence) if "times" in names else None
+
+    # labels (computed outside of the compiled program: static)
+    if set(names) & {"layout", "xarray"}:
+        coords = axes.get_coords(sequence, sample, scanner)
+        if getattr(sample, "compartments", None):
+            coords["compartment"] = np.array(sample.compartments)
+        coords.update(method.core_coords(sequence))
+        core = method.core_axes(sequence)
+        layout, (resolved, _, _) = axes.resolve(
+            sequence, sample, scanner, dims=dims, alignment=alignment
+        )
+        resolved = [op for op in flatten_sequence(resolved) if isinstance(op, Probe)]
+
+    # outputs, by group (positions of the ADCs: static)
+    results = [{} for _ in names]
+    for key, indices in groups.items():
+        values_, times_, probes_ = values, times, probes
+        if indices is not None:
+            values_ = _take(values, indices)
+            times_ = None if times is None else times[np.asarray(indices)]
+        if set(names) & {"layout", "xarray"}:
+            probes_ = resolved if indices is None else [resolved[i] for i in indices]
+            labels = _labels(values_, probes_, layout, core, coords)
+        for result, name in zip(results, names):
+            if name == "jax":
+                result[key] = values_
+            elif name == "numpy":
+                result[key] = _to_numpy(values_)
+            elif name == "times":
+                result[key] = times_
+            elif name == "layout":
+                result[key] = labels
+            elif name == "xarray":
+                result[key] = _to_xarray(values_, labels, times_)
+    results = [result[None] if group is None else result for result in results]
+    return results[0] if len(results) == 1 else tuple(results)
+
+
+def _execute(leaves, sample, scanner, init, *, skeleton, options):
+    """run a flattened sequence (cf. `simulate`)
+
+    Returns:
+        acquired values, ADC times
+    """
+    options = dict(options)
+    sequence = unflatten_operators(skeleton, leaves)
+    layout, (sequence, sample, scanner) = axes.resolve(
+        sequence, sample, scanner, dims=options["dims"], alignment=options["alignment"]
+    )
+
+    # method and initial state
+    method = options["method"]
+    method.check(sequence)
+    if init is None:
+        density = 1.0 if sample is None else sample.PD
+        state = method.init_state(
+            sequence, layout.shape, density=density, dtype=options["dtype"]
+        )
+    else:
+        state = init.broadcast(layout.shape)
+
+    # compile and run
+    program = lowering.compile_sequence(
+        sequence,
+        max_unroll=options["max_unroll"],
+        checkpoint=options["checkpoint"],
+        merge=options["merge"],
+    )
+    monitor = Progress(program.count).monitor if options["disp"] else None
+    context = lowering.Context(sample, scanner, method)
+    _, values = program.run(state, context, monitor=monitor)
+    return values, get_adc_times(sequence)
+
+
+# compiled once per structure (skeleton) and static options
+_execute_jit = jax.jit(_execute, static_argnames=("skeleton", "options"))
+
+
+@pytree
+class Trace:
+    """States after every operator of a sequence (cf. `trace`)
+
+    Attributes:
+        F, Z: states before the first operator (row 0) and after each operator
+            (row i + 1), shape (nop + 1, *core)
+        times: time of each row (ms): 0, then end time of each operator
+        coords: coordinates of the core axis (phase states k or isochromat
+            positions); nd lattice: tuple of coordinates, one per axis
+        operators: flat tuple of the operators (parameters selected as the states)
+        axis: (static) name of the core axis; nd lattice: tuple of names
+            (e.g. ("kx", "ky", "kz"))
+        representation: (static) representation of the states
+    """
+
+    F: jax.Array = None
+    Z: jax.Array = None
+    times: jax.Array = None
+    coords: jax.Array = None
+    operators: tuple = ()
+    axis: str = static("k")
+    representation: str = static("lattice")
+
+    def __init__(self, F, Z, times, coords, operators, *, axis, representation):
+        for name, value in locals().items():
+            if name != "self":
+                object.__setattr__(self, name, value)
+
+    def __repr__(self):
+        names = (self.axis,) if isinstance(self.axis, str) else self.axis
+        sizes = self.F.shape[len(self.F.shape) - len(names) :]
+        core = ", ".join(f"{n}={s}" for n, s in zip(names, sizes))
+        return f"Trace({len(self.operators)} operators, {core}, {self.representation})"
+
+    @property
+    def nop(self):
+        """number of operators"""
+        return len(self.operators)
+
+
+def trace(
+    sequence,
+    sample=None,
+    scanner=None,
+    *,
+    select=None,
+    method=None,
+    init=None,
+    dtype=None,
+    dims=None,
+    alignment=None,
+    max_unroll=compiler.MAX_UNROLL,
+    jit=True,
+    disp=False,
+):
+    """simulate a sequence, keeping the states after every operator
+
+    The states are acquired by a probe (`SNAPSHOT`) inserted after every
+    operator, in the same nested lists (repeated blocks are still scanned).
+    The probes of the sequence are replaced by snapshots.
+
+    Args:
+        sequence, sample, scanner: cf. `simulate`
+        select: element of the batch axes, as a dict {axis: index} (axis name,
+            or int i for the unnamed axis dim_i); other axes: index 0
+        method, init, dtype, dims, alignment, max_unroll, jit, disp: cf. `simulate`
+
+    Returns:
+        Trace (states and operators of the selected element)
+    """
+    method = method or methods.infer_method(sequence)
+    core = method.core_axes(sequence)
+    traced = [SNAPSHOT] + _insert_snapshots(sequence, {})
+    F, Z = simulate(
+        traced,
+        sample,
+        scanner,
+        method=method,
+        init=init,
+        dtype=dtype,
+        dims=dims,
+        alignment=alignment,
+        max_unroll=max_unroll,
+        merge=False,
+        jit=jit,
+        disp=disp,
+    )
+
+    # selected element of the batch axes
+    batch, (resolved, _, _) = axes.resolve(
+        sequence, sample, scanner, dims=dims, alignment=alignment
+    )
+    index = _select_index(select, batch)
+    F = F[(slice(None),) + index]
+    Z = Z[(slice(None),) + index]
+    operators = tuple(
+        jax.tree.map(lambda x: _select_leaf(x, index), op)
+        for op in flatten_sequence(resolved)
+    )
+    durations = jnp.stack([jnp.asarray(op.duration, dtype=float) for op in operators])
+    times = jnp.concatenate([jnp.zeros(1, durations.dtype), jnp.cumsum(durations)])
+
+    # coordinates of the core axes (nd lattice: one per axis)
+    if method.representation == "lattice":
+        sizes = F.shape[F.ndim - len(core) :]
+        coords = tuple(jnp.arange(-(n // 2), n // 2 + 1) for n in sizes)
+    else:
+        core_coords = method.core_coords(sequence)
+        coords = tuple(jnp.asarray(core_coords[axis]) for axis in core)
+    axis, coords = (core[0], coords[0]) if len(core) == 1 else (core, coords)
+    return Trace(
+        F,
+        Z,
+        times,
+        coords,
+        operators,
+        axis=axis,
+        representation=method.representation,
+    )
+
+
+def check_sample(sequence, sample):
+    """check the sample against the operators of the sequence
+
+    Raises if a property is given both by an operator and by the sample (e.g.
+    `E(tau, T1, T2)` with a sample setting T2). Warns if properties set in the
+    sample are read by no operator (e.g. a sample with T1/T2 and a sequence of
+    `R` operators, which ignore the sample). PD is always used (equilibrium).
+    """
+    if sample is None:
+        return
+    operators = flatten_sequence(sequence)
+    for op in operators:
+        check_conflicts(op, sample)
+    fields = {name for name in ("T1", "T2", "df") if getattr(sample, name) is not None}
+    used = set()
+    for op in operators:
+        used.update(op.sample_fields())
+    unused = sorted(fields - used)
+    if unused:
+        warnings.warn(
+            f"Sample properties not used by any operator of the sequence: {unused} "
+            "(they are read by E and P)",
+            stacklevel=3,
+        )
+
+
+def replace_probes(sequence, probe):
+    """replace the probes of a sequence (keeping their durations and names)
+
+    Args:
+        probe: attribute name (str), Probe, or tuple of these (-> Probes)
+    """
+    if probe is None:
+        return sequence
+    multiple = isinstance(probe, tuple)
+    probes = probe if multiple else (probe,)
+    for item in probes:
+        if isinstance(item, str) and item not in kprobe.ATTRIBUTES:
+            raise ValueError(
+                f"Invalid probe: {item}, expected one of {kprobe.ATTRIBUTES}"
+            )
+        elif not isinstance(item, (str, Probe)):
+            raise TypeError(f"Invalid probe: {item}")
+
+    cache = {}  # keep the identity of shared probes (cf. compiler)
+
+    def new_probe(op, item):
+        if isinstance(item, Adc) and isinstance(op, Adc):
+            # the probe sets what is measured, the sequence ADC keeps its phase
+            if item.phase is not None and op.phase is not None:
+                raise ValueError(
+                    f"ADC phase given both by the probe {item!r} and by the ADC "
+                    f"{op!r} of the sequence: remove it from the probe"
+                )
+            phase = op.phase if item.phase is None else item.phase
+            return item.replace(phase=phase, duration=op.duration)
+        elif isinstance(item, Probe):
+            return item.replace(duration=op.duration)
+        elif isinstance(op, Adc):
+            # same ADC (phase, weights, reduction) with another attribute
+            return op.replace(attr=item)
+        return Adc(item, duration=op.duration)
+
+    def replace(op):
+        if not isinstance(op, Probe):
+            return op
+        if id(op) not in cache:
+            new = [new_probe(op, item) for item in probes]
+            if multiple:
+                new = Probes(new, duration=op.duration, name=op.name)
+            else:
+                # the name of the sequence probe is kept (cf. group="name")
+                new = new[0] if new[0].name else new[0].replace(name=op.name)
+            cache[id(op)] = new
+        return cache[id(op)]
+
+    return jax.tree.map(replace, sequence, is_leaf=lambda x: isinstance(x, Operator))
+
+
+#
+# private
+
+
+def _insert_snapshots(sequence, cache):
+    """insert SNAPSHOT after every operator, replace probes (cf. `trace`)"""
+    if isinstance(sequence, Operator):
+        sequence = [sequence]
+    result = []
+    for item in sequence:
+        if isinstance(item, (list, tuple)):
+            result.append(_insert_snapshots(item, cache))
+        elif isinstance(item, Probe):
+            # keep the identity of shared probes (cf. compiler)
+            if id(item) not in cache:
+                cache[id(item)] = SNAPSHOT.replace(duration=item.duration)
+            result.append(cache[id(item)])
+        elif isinstance(item, Operator):
+            result.extend([item, SNAPSHOT])
+        else:
+            raise ValueError(f"Invalid operator: {item}")
+    return result
+
+
+def _select_index(select, layout):
+    """index of the selected element of the batch axes (cf. `trace`)"""
+    select = dict(select or {})
+    index = [0] * len(layout.dims)
+    for axis, value in select.items():
+        name = f"dim_{axis}" if isinstance(axis, int) else axis
+        if name not in layout.dims:
+            raise ValueError(f"Unknown axis: {axis!r}, expected one of {layout.dims}")
+        pos = layout.dims.index(name)
+        size = layout.shape[pos]
+        value = _operator.index(value)
+        if not -size <= value < size:
+            raise IndexError(f"Index {value} out of range for axis {name!r} ({size})")
+        index[pos] = value % size
+    return tuple(index)
+
+
+def _select_leaf(value, index):
+    """selected element of a resolved leaf (scalar, or broadcastable array)"""
+    if jnp.ndim(value) < len(index):
+        return value  # scalar
+    shape = jnp.shape(value)
+    return value[tuple(i if s > 1 else 0 for i, s in zip(index, shape))]
+
+
+def _probe_dims(probe, layout, core):
+    """axis names of the value acquired by a (resolved) probe"""
+    if isinstance(probe, Probes):
+        return tuple(_probe_dims(p, layout, core) for p in probe.probes)
+    elif not isinstance(probe, Adc):
+        return None  # unknown probe
+    dims = layout.reduced_dims(probe.reduce)
+    if probe.attr in ("F", "Fminus", "Z"):
+        dims += core
+    elif probe.attr == "states":
+        dims += core + ("component",)
+    return dims
+
+
+def _labels(values, probes, layout, core, coords):
+    """OutputLayout(s) of the values acquired by (resolved) probes"""
+    specs = [_probe_dims(probe, layout, core) for probe in probes]
+    names = tuple(probe.name for probe in probes)
+
+    def label(dims, value, adc=True):
+        if dims is None:
+            return None
+        dims = (("adc",) if adc else ()) + dims
+        coords_ = {d: coords[d] for d in dims if d in coords}
+        return OutputLayout(dims, tuple(value.shape), coords_, names if adc else None)
+
+    def label_value(spec, value, adc=True):
+        if isinstance(value, tuple):
+            return tuple(label_value(s, v, adc) for s, v in zip(spec, value))
+        return label(spec, value, adc)
+
+    if isinstance(values, list):
+        return [
+            label_value(spec, value, adc=False) for spec, value in zip(specs, values)
+        ]
+    return label_value(specs[0], values)
+
+
+def _group_by_name(probes):
+    """positions of the probes by name, in order of first appearance"""
+    unnamed = [probe for probe in probes if not probe.name]
+    if unnamed:
+        raise ValueError(
+            f'group="name" requires named probes, got: {unnamed[0]!r} '
+            '(e.g. Adc(name="echo"), the ADC instance is named "ADC")'
+        )
+    groups = {}
+    for i, probe in enumerate(probes):
+        groups.setdefault(probe.name, []).append(i)
+    return groups
+
+
+def _take(values, indices):
+    """acquired values of the ADCs at given positions (as a whole output)"""
+    if all(isinstance(leaf, jax.ShapeDtypeStruct) for leaf in jax.tree.leaves(values)):
+        # shapes only (labels without simulation)
+        return jax.eval_shape(functools.partial(_take, indices=indices), values)
+    if isinstance(values, list):
+        chunks = [jax.tree.map(lambda leaf: leaf[None], values[i]) for i in indices]
+        return compiler._assemble(chunks)
+    index = np.asarray(indices)
+    return jax.tree.map(lambda leaf: leaf[index], values)
+
+
+def _to_numpy(values):
+    def convert(leaf):
+        if isinstance(leaf, jax.core.Tracer):
+            raise ValueError('output="numpy" is not available under jit/grad')
+        return np.asarray(leaf)
+
+    return jax.tree.map(convert, values)
+
+
+def _to_xarray(values, labels, times):
+    try:
+        import xarray
+    except ImportError as exc:
+        raise ImportError('output="xarray" requires xarray') from exc
+    times = None if times is None else np.asarray(_to_numpy(times))
+
+    def convert(value, label):
+        if isinstance(value, tuple):
+            return tuple(convert(v, lab) for v, lab in zip(value, label))
+        elif label is None:
+            raise ValueError("Cannot label the output of a custom probe")
+        coords = {name: np.asarray(c) for name, c in label.coords.items()}
+        if label.dims and label.dims[0] == "adc" and times is not None:
+            coords["adc"] = times
+        return xarray.DataArray(
+            np.asarray(_to_numpy(value)), dims=label.dims, coords=coords
+        )
+
+    if isinstance(values, list):
+        return [convert(value, label) for value, label in zip(values, labels)]
+    return convert(values, labels)

epgx/kernels/__init__.py

@@ -0,0 +1,9 @@
+"""Layer 1: numerical kernels (pure functions on dimensionless arrays)"""
+
+from .affine import apply_affine
+from .evolution import evolve
+from .operator import reset, spoil, spoil_axis
+from .probe import adc, demodulate, readout, reduce
+from .shift import shift_lattice, shift_samples
+from .state import State, mirror
+from .transition import phase, rotate