epgx 0.1.dev0__py3-none-any.whl

epgx/compiler.py

@@ -0,0 +1,347 @@
+"""Sequence compiler: repeated blocks -> lax.scan
+
+Generic: the compiler knows nothing about EPG. The caller provides the lowering
+(item -> Step) and the run context (passed through to the steps). The EPG
+front-end is `lowering.compile_sequence`.
+
+Stages:
+    1. grouping: a (nested) list of items is split into segments:
+        - Unrolled: steps applied one after the other
+        - Scan: consecutive sub-lists with the same pytree structure (repeated
+          blocks), applied with `lax.scan`; leaves that differ between blocks are
+          stacked, leaves shared by all blocks (same object) are constants.
+          Blocks are compiled recursively (nested repeats: nested scans).
+    2. lowering: items (anything but lists) -> steps, by the caller's `lower`
+    3. passes: functions Program -> Program, applied to the program of each
+       level (top level and scan bodies), e.g. merging (cf. `merging`)
+    4. execution: `Program.run(state, context=None, monitor=None)`; `context` is
+       passed to the steps (e.g. sample and scanner); the optional monitor is
+       called (when tracing) with the state and a number of items done: once per
+       top-level unrolled segment and once per iteration of top-level scans
+       (e.g. progress from the compiled program, cf. `progress`)
+
+Grouping only depends on the structure (static), so a program can be compiled
+and run with traced values (under jit, grad...). Acquired values may be
+pytrees (e.g. tuples of arrays, for multiple probes).
+"""
+
+import dataclasses
+import typing
+
+import jax
+import jax.numpy as jnp
+
+MAX_UNROLL = 1000
+
+
+@dataclasses.dataclass(frozen=True)
+class Step:
+    """Lowered item
+
+    Attributes:
+        kernel: function applied as kernel(state, **params(context))
+        params: function context -> dict of kernel parameters
+        acquire: function (state, context) -> value, or None
+        label: description (e.g. operator repr)
+        count: number of items represented by the step (e.g. merged steps: several)
+        source: object the step was lowered from (e.g. operator), for passes
+    """
+
+    kernel: typing.Callable
+    params: typing.Callable
+    acquire: typing.Callable = None
+    label: str = ""
+    count: int = 1
+    source: typing.Any = None
+
+
+@dataclasses.dataclass
+class Unrolled:
+    """Segment of steps applied one after the other"""
+
+    steps: list
+
+    def __repr__(self):
+        return f"Unrolled({len(self.steps)})"
+
+    @property
+    def nunroll(self):
+        return len(self.steps)
+
+    @property
+    def count(self):
+        """number of items"""
+        return sum(step.count for step in self.steps)
+
+    def run(self, state, context=None, monitor=None):
+        chunks = []
+        for step in self.steps:
+            state = step.kernel(state, **step.params(context))
+            if step.acquire is not None:
+                value = step.acquire(state, context)
+                chunks.append(jax.tree.map(lambda leaf: leaf[jnp.newaxis], value))
+        if monitor is not None:
+            monitor(state, self.count)
+        return state, chunks
+
+
+@dataclasses.dataclass
+class Scan:
+    """Segment of repeated blocks applied with lax.scan
+
+    Attributes:
+        treedef: pytree structure of one block
+        leaves: leaves of the first block (constants where `varying` is False)
+        varying: for each leaf, True if it differs between blocks (stacked)
+        stacked: stacked varying leaves (leading axis: block index)
+        length: number of blocks
+        body: program of one block (compiled from the first block)
+        options: compilation options (to compile the body of each block)
+    """
+
+    treedef: typing.Any
+    leaves: list
+    varying: list
+    stacked: list
+    length: int
+    body: "Program"
+    options: dict
+
+    def __repr__(self):
+        return f"Scan({self.length} x {self.body!r})"
+
+    @property
+    def nunroll(self):
+        return self.body.nunroll
+
+    @property
+    def count(self):
+        """number of items"""
+        return self.length * self.body.count
+
+    @property
+    def xs(self):
+        """scanned values (stacked varying leaves), None if all blocks are the same"""
+        return tuple(self.stacked) if self.stacked else None
+
+    def block(self, xs):
+        """block rebuilt from one slice of the stacked leaves"""
+        xs = iter(xs or ())
+        leaves = [next(xs) if v else leaf for v, leaf in zip(self.varying, self.leaves)]
+        return jax.tree.unflatten(self.treedef, leaves)
+
+    def body_at(self, xs):
+        """program of the block with the scanned values `xs` (one slice)"""
+        return _compile(self.block(xs), **self.options)
+
+    def run(self, state, context=None, monitor=None):
+        def body(state, xs):
+            state, chunks = self.body_at(xs)._run(state, context)
+            if monitor is not None:
+                monitor(state, self.body.count)
+            return state, tuple(chunks)
+
+        if self.options.get("checkpoint"):
+            body = jax.checkpoint(body)
+        state, ys = jax.lax.scan(body, state, self.xs, length=self.length)
+        return state, _interleave(ys, self.length)
+
+
+@dataclasses.dataclass
+class Program:
+    """Compiled sequence
+
+    Attributes:
+        segments: list of Unrolled and Scan segments
+    """
+
+    segments: list
+
+    def __repr__(self):
+        return f"Program({', '.join(map(repr, self.segments))})"
+
+    @property
+    def nunroll(self):
+        """number of unrolled steps (traced items), including scan bodies"""
+        return sum(segment.nunroll for segment in self.segments)
+
+    @property
+    def count(self):
+        """number of items (steps standing for several items count them all)"""
+        return sum(segment.count for segment in self.segments)
+
+    @property
+    def nscan(self):
+        """number of scan segments (including nested scans)"""
+        return sum(
+            1 + segment.body.nscan
+            for segment in self.segments
+            if isinstance(segment, Scan)
+        )
+
+    def run(self, state, context=None, monitor=None):
+        """apply the sequence to the state
+
+        Args:
+            state: initial state (pytree)
+            context: passed to the steps' `params` and `acquire` (e.g. sample
+                and scanner)
+            monitor: function (state, count) called after each top-level segment
+                and each iteration of top-level scans (cf. module docstring)
+
+        Returns:
+            final state, acquired values: if all values have the same structure
+            and shapes, a single value (array or pytree of arrays) with leading
+            axis (nadc, ...), else a list of values
+        """
+        state, chunks = self._run(state, context, monitor)
+        return state, _assemble(chunks)
+
+    def _run(self, state, context=None, monitor=None):
+        chunks = []
+        for segment in self.segments:
+            state, new = segment.run(state, context, monitor)
+            chunks.extend(new)
+        return state, chunks
+
+
+def build(sequence, *, lower, passes=(), max_unroll=MAX_UNROLL, checkpoint=False):
+    """compile a (nested) list of items into a Program
+
+    Args:
+        sequence: (nested) list of items; repeated blocks are consecutive
+            sub-lists with the same structure (e.g. `[exc] + [block] * N`)
+        lower: function item -> Step (raises for invalid items)
+        passes: functions Program -> Program, applied to each level (top level
+            and scan bodies), in order
+        max_unroll: maximum number of unrolled steps (error if exceeded)
+        checkpoint: wrap scan bodies with `jax.checkpoint` (memory vs. compute
+            trade-off for reverse-mode differentiation)
+
+    Returns:
+        Program
+    """
+    program = _compile(
+        sequence, lower=lower, passes=tuple(passes), checkpoint=checkpoint
+    )
+    if program.nunroll > max_unroll:
+        raise ValueError(
+            f"Too many unrolled operators ({program.nunroll} > max_unroll={max_unroll}). "
+            "Group repeated blocks into sub-lists with the same structure, "
+            "e.g. `[exc] + [[rf, shift, adc]] * N`, or increase `max_unroll`."
+        )
+    return program
+
+
+def join_unrolled(segments):
+    """merge consecutive Unrolled segments (e.g. for passes)"""
+    joined = []
+    for segment in segments:
+        if (
+            isinstance(segment, Unrolled)
+            and joined
+            and isinstance(joined[-1], Unrolled)
+        ):
+            joined[-1] = Unrolled(joined[-1].steps + segment.steps)
+        else:
+            joined.append(segment)
+    return joined
+
+
+#
+# private
+
+
+def _compile(sequence, **options):
+    """group a (nested) list of items into segments, then apply the passes"""
+    items = list(sequence) if isinstance(sequence, (list, tuple)) else [sequence]
+    segments = []
+    i = 0
+    while i < len(items):
+        item = items[i]
+        if not isinstance(item, (list, tuple)):
+            segments.append(Unrolled([options["lower"](item)]))
+            i += 1
+            continue
+
+        # run of consecutive blocks with the same structure
+        structure = jax.tree.structure(item)
+        j = i + 1
+        while (
+            j < len(items)
+            and isinstance(items[j], (list, tuple))
+            and jax.tree.structure(items[j]) == structure
+        ):
+            j += 1
+        if j - i == 1:
+            # single block: inline
+            segments.extend(_compile(item, **options).segments)
+        else:
+            segments.append(_scan(items[i:j], **options))
+        i = j
+
+    program = Program(join_unrolled(segments))
+    for transform in options.get("passes", ()):
+        program = transform(program)
+    return program
+
+
+def _scan(blocks, **options):
+    """Scan segment from a list of blocks with the same structure"""
+    flat = [jax.tree.flatten(block) for block in blocks]
+    treedef = flat[0][1]
+    leaves = [leaves for leaves, _ in flat]
+    first = leaves[0]
+    varying = [
+        any(other[j] is not first[j] for other in leaves[1:]) for j in range(len(first))
+    ]
+    stacked = []
+    for j, vary in enumerate(varying):
+        if not vary:
+            continue
+        values = [jnp.asarray(other[j]) for other in leaves]
+        shape = jnp.broadcast_shapes(*(value.shape for value in values))
+        stacked.append(jnp.stack([jnp.broadcast_to(value, shape) for value in values]))
+    body = _compile(blocks[0], **options)  # structure of the body (same for all blocks)
+    return Scan(treedef, first, varying, stacked, len(blocks), body, options)
+
+
+def _signature(chunk, skip):
+    """structure and trailing leaf shapes of a chunk (skipping `skip` leading axes)"""
+    leaves, treedef = jax.tree.flatten(chunk)
+    return treedef, tuple(leaf.shape[skip:] for leaf in leaves)
+
+
+def _count(chunk, axis=0):
+    """number of values along a leading axis of a chunk"""
+    return jax.tree.leaves(chunk)[0].shape[axis]
+
+
+def _interleave(ys, length):
+    """scan outputs (one per chunk of the body, leading axes (length, k)) -> chunks"""
+    if not ys:
+        return []
+    if len({_signature(y, 2) for y in ys}) == 1:
+        # same structures and shapes: one chunk of length * k values, in sequence order
+        merged = jax.tree.map(lambda *leaves: jnp.concatenate(leaves, axis=1), *ys)
+        return [jax.tree.map(lambda leaf: leaf.reshape((-1,) + leaf.shape[2:]), merged)]
+    # otherwise: one chunk per value
+    return [
+        jax.tree.map(lambda leaf: leaf[i, j][jnp.newaxis], y)
+        for i in range(length)
+        for y in ys
+        for j in range(_count(y, axis=1))
+    ]
+
+
+def _assemble(chunks):
+    """acquired values: single value if all have the same structure and shapes, else list"""
+    if not chunks:
+        return jnp.zeros((0,))
+    if len({_signature(chunk, 1) for chunk in chunks}) == 1:
+        return jax.tree.map(lambda *leaves: jnp.concatenate(leaves), *chunks)
+    return [
+        jax.tree.map(lambda leaf: leaf[j], chunk)
+        for chunk in chunks
+        for j in range(_count(chunk))
+    ]

epgx/diff.py

@@ -0,0 +1,242 @@
+"""Derivatives of functions of named values (Jacobian, cross Hessian)
+
+JAX differentiates `simulate` directly (`jax.jvp`, `jax.jacfwd`, `jax.grad`...).
+This module only handles what plain JAX makes error-prone:
+
+    - batch axes: `jax.jacfwd` w.r.t. a parameter with one value per batch
+      element (e.g. 1000 T2 values) builds a (..., 1000, 1000) matrix, zero off
+      the diagonal, at the cost of 1000 passes. Since each element depends only
+      on its own value, one pass with a tangent of ones gives all the
+      derivatives;
+    - Jacobian and (cross) Hessian in one computation (nested JVPs);
+    - coupled axes: elements interacting along an axis (e.g. compartments
+      mixed by exchange) need one derivative per value along that axis.
+
+Kinds of values (one entry of `values`):
+    Axis: one value per element of its batch axes -> one column (tangent of
+        ones); along axes listed in `coupled`: one column per value
+    scalar: one parameter -> one column
+    array (not an Axis): a vector of parameters (e.g. flip angles) -> one
+        column per value
+
+Example:
+    def signal(p):
+        seq = [T(90, 90)] + [[S(1), E(5), T(a, 0), S(1), E(5), ADC] for a in p["alpha"]]
+        return simulate(seq, Sample(T1=p["T1"], T2=p["T2"]))
+
+    values = {"T1": 1000.0, "T2": Axis("T2", [30.0, 40.0]), "alpha": alphas}
+    d = derivatives(signal, values, jacobian=["T1", "T2"], hessian=(["T2"], ["alpha"]))
+    d.jacobian  # (nadc, 2, 2): last axis: T1, T2
+    d.hessian   # (nadc, 2, 1, nalpha)
+"""
+
+import itertools
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from .axes import is_axis
+from .common import pytree, static
+
+
+@pytree
+class Derivatives:
+    """Value and derivatives of a function (cf. `derivatives`)
+
+    Attributes:
+        value: function value
+        jacobian: (*shape, ncols) or None, derivatives along the last axis
+        hessian: (*shape, nrows, ncols) or None
+        columns: (static) labels of the Jacobian's columns, e.g. ("T2",) or
+            ("alpha", 3) (parameter name, then index for vectors/coupled axes)
+        rows, cols: (static) labels of the Hessian's rows and columns
+    """
+
+    value: jax.Array = None
+    jacobian: jax.Array = None
+    hessian: jax.Array = None
+    columns: tuple = static(())
+    rows: tuple = static(())
+    cols: tuple = static(())
+
+    def __init__(
+        self, value, jacobian=None, hessian=None, *, columns=(), rows=(), cols=()
+    ):
+        for name, item in [
+            ("value", value),
+            ("jacobian", jacobian),
+            ("hessian", hessian),
+            ("columns", tuple(columns)),
+            ("rows", tuple(rows)),
+            ("cols", tuple(cols)),
+        ]:
+            object.__setattr__(self, name, item)
+
+
+def derivatives(fn, values, *, jacobian=(), hessian=None, coupled=()):
+    """value, Jacobian and cross Hessian of fn(values), in one computation
+
+    Args:
+        fn: function dict -> array (or pytree of arrays)
+        values: dict name -> value (Axis, scalar or array, cf. module docstring)
+        jacobian: names of the Jacobian's parameters
+        hessian: (rows, cols): names of the Hessian's rows and columns
+            (rectangular, not assumed symmetric), or None
+        coupled: names of batch axes along which elements interact
+
+    Returns:
+        Derivatives (value, jacobian, hessian and their labels)
+    """
+    jacobian = list(jacobian)
+    rows, cols = (None, None) if hessian is None else map(list, hessian)
+    names = list(dict.fromkeys(jacobian + (rows or []) + (cols or [])))
+    for name in names:
+        if name not in values:
+            raise ValueError(
+                f"Unknown parameter: {name!r}, expected one of {list(values)}"
+            )
+    params = {name: _as_float(values[name]) for name in names}
+
+    def f(params):
+        return fn({**values, **params})
+
+    if hessian is None:
+        basis, columns = _basis(params, jacobian, coupled)
+        value, jac = _jacobian(f, params, basis)
+        return Derivatives(value, _last(jac), columns=columns)
+
+    # nested JVPs: rows inside, columns outside -> value, J_rows, J_cols, H
+    U, row_labels = _basis(params, rows, coupled)
+    V, col_labels = _basis(params, cols, coupled)
+
+    def per_col(v):
+        def per_row(u):
+            def inner(x):
+                return jax.jvp(f, (x,), (u,))
+
+            (value, ju), (jv, huv) = jax.jvp(inner, (params,), (v,))
+            return value, ju, jv, huv
+
+        return jax.vmap(per_row, out_axes=(None, 0, None, 0))(U)
+
+    value, ju, jv, huv = jax.vmap(per_col, out_axes=(None, None, 0, 0))(V)
+    # huv: (ncols, nrows, *shape) -> (*shape, nrows, ncols)
+    hess = jax.tree.map(lambda h: jnp.moveaxis(h, (0, 1), (-1, -2)), huv)
+
+    # Jacobian: columns from J_rows / J_cols, others computed separately
+    jac, columns = None, ()
+    if jacobian:
+        known = {}
+        for labels, jacs in [(row_labels, ju), (col_labels, jv)]:
+            for i, label in enumerate(labels):
+                known.setdefault(label, jax.tree.map(lambda j, i=i: j[i], jacs))
+        missing = [name for name in jacobian if not _has(known, name)]
+        if missing:
+            basis, labels = _basis(params, missing, coupled)
+            _, extra = _jacobian(f, params, basis)
+            for i, label in enumerate(labels):
+                known[label] = jax.tree.map(lambda j, i=i: j[i], extra)
+        _, columns = _basis(params, jacobian, coupled)
+        jac = jax.tree.map(
+            lambda *js: jnp.stack(js, axis=-1), *[known[c] for c in columns]
+        )
+    return Derivatives(
+        value, jac, hess, columns=columns, rows=row_labels, cols=col_labels
+    )
+
+
+def jacobian(fn, values, wrt, *, coupled=()):
+    """value and Jacobian (*shape, ncols) of fn(values) (cf. `derivatives`)"""
+    d = derivatives(fn, values, jacobian=wrt, coupled=coupled)
+    return d.value, d.jacobian
+
+
+def hessian(fn, values, rows, cols=None, *, coupled=()):
+    """value, Jacobian (rows) and Hessian (*shape, nrows, ncols) (cf. `derivatives`)"""
+    cols = rows if cols is None else cols
+    d = derivatives(fn, values, jacobian=rows, hessian=(rows, cols), coupled=coupled)
+    return d.value, d.jacobian, d.hessian
+
+
+#
+# private
+
+
+def _as_float(value):
+    """differentiable copy of a value (float dtype)"""
+
+    def convert(leaf):
+        leaf = jnp.asarray(leaf)
+        if jnp.issubdtype(leaf.dtype, jnp.inexact):
+            return leaf
+        return leaf.astype(jnp.result_type(float))
+
+    return jax.tree.map(convert, value)
+
+
+def _tangents(name, value, coupled):
+    """tangents of one parameter: (ntangent, *shape) array, labels"""
+    if is_axis(value):
+        shape, dims = value.shape, value.names
+        cdims = [i for i, dim in enumerate(dims) if dim in coupled]
+    else:
+        shape = jnp.shape(value)
+        cdims = list(range(len(shape)))  # vector of parameters
+    csizes = [shape[i] for i in cdims]
+    tangents, labels = [], []
+    for index in itertools.product(*(range(n) for n in csizes)):
+        tangent = np.zeros(shape) if cdims else np.ones(shape)
+        if cdims:
+            where = [slice(None)] * len(shape)
+            for i, k in zip(cdims, index):
+                where[i] = k
+            tangent[tuple(where)] = 1
+        tangents.append(tangent)
+        if not index:
+            labels.append((name,))
+        elif len(index) == 1:
+            labels.append((name, index[0]))
+        else:
+            labels.append((name, index))
+    return np.stack(tangents), labels
+
+
+def _basis(params, names, coupled):
+    """stacked tangents (leading axis) for the parameters `names`, labels"""
+    # tangents in the requested order
+    offsets, labels, start = {}, [], 0
+    for name in names:
+        tangents, name_labels = _tangents(name, params[name], coupled)
+        offsets[name] = (start, tangents)
+        labels += name_labels
+        start += len(tangents)
+    # one stacked array per parameter (zeros outside its own tangents)
+    stacked = {}
+    for name, value in params.items():
+        leaf = value.values if is_axis(value) else jnp.asarray(value)
+        array = np.zeros((start,) + tuple(jnp.shape(leaf)))
+        if name in offsets:
+            offset, tangents = offsets[name]
+            array[offset : offset + len(tangents)] = tangents
+        array = jnp.asarray(array, dtype=leaf.dtype)
+        stacked[name] = value.replace(values=array) if is_axis(value) else array
+    return stacked, tuple(labels)
+
+
+def _jacobian(f, params, basis):
+    """value and Jacobian (ntangent, *shape) for stacked tangents"""
+
+    def column(t):
+        return jax.jvp(f, (params,), (t,))
+
+    return jax.vmap(column, out_axes=(None, 0))(basis)
+
+
+def _last(jac):
+    """leading tangent axis -> last axis"""
+    return jax.tree.map(lambda j: jnp.moveaxis(j, 0, -1), jac)
+
+
+def _has(known, name):
+    return any(label[0] == name for label in known)