epgx 0.1.dev0__py3-none-any.whl

epgx/ops/functions.py

@@ -0,0 +1,158 @@
+"""Sequence functions"""
+
+import dataclasses
+
+import jax
+import jax.numpy as jnp
+
+from .operator import Operator
+from .probe import Probe
+
+
+def flatten_sequence(sequence):
+    """return a flat list of operators from a (nested) list of operators"""
+    if isinstance(sequence, Operator):
+        return [sequence]
+    flat = []
+    for item in sequence:
+        if isinstance(item, (list, tuple)):
+            flat.extend(flatten_sequence(item))
+        elif isinstance(item, Operator):
+            flat.append(item)
+        else:
+            raise ValueError(f"Invalid operator: {item}")
+    return flat
+
+
+def getnshift(sequence):
+    """number of phase-state shifts of the sequence
+
+    Upper bound of the state capacity needed on the lattice (sum of |k|).
+    """
+    return sum(op.nshift for op in flatten_sequence(sequence))
+
+
+def getnshifts(sequence):
+    """number of phase-state shifts of the sequence along each dephasing direction
+
+    Upper bound of the state capacity needed along each direction (sum of
+    |k|), as a dict {direction: count}: "x", "y", "z" first (as many as the
+    longest vector shift), then the named axes in order of appearance.
+    """
+    counts = {}
+    for op in flatten_sequence(sequence):
+        for direction, s in getattr(op, "shifts", {}).items():
+            counts[direction] = counts.get(direction, 0) + abs(s)
+    order = [d for d in ("x", "y", "z") if d in counts]
+    order += [d for d in counts if d not in ("x", "y", "z")]
+    return {d: counts[d] for d in order}
+
+
+def get_dephasing_axes(sequence):
+    """dephasing axes declared by the operators, by name (cf. DephasingAxis)
+
+    Two different axes with the same name raise an error (different kinds, or
+    different concrete scales).
+    """
+    axes = {}
+    for op in flatten_sequence(sequence):
+        axis = getattr(op, "axis", None)
+        if axis is None:
+            continue
+        other = axes.setdefault(axis.name, axis)
+        if other is not axis and not _same_axis(axis, other):
+            raise ValueError(
+                f"Two different dephasing axes named {axis.name!r}: {other!r}, {axis!r}"
+            )
+    return axes
+
+
+def _same_axis(a, b):
+    """same kind and scale (scales compared when concrete)"""
+    if a.kind != b.kind:
+        return False
+    va, vb = (a.kvalue, b.kvalue) if a.kind == "space" else (a.tvalue, b.tvalue)
+    if isinstance(va, jax.core.Tracer) or isinstance(vb, jax.core.Tracer):
+        return True
+    return bool(jnp.all(jnp.asarray(va) == jnp.asarray(vb)))
+
+
+def get_adc_times(sequence):
+    """return ADC opening times (cf. operator's `duration` keyword)
+
+    The time of a probe includes its own duration.
+    """
+    operators = flatten_sequence(sequence)
+    probes = [i for i, op in enumerate(operators) if isinstance(op, Probe)]
+    if not probes:
+        return jnp.zeros(0)
+    # one cumulative sum (not a sum per operator: compact when traced)
+    durations = jnp.stack(jnp.broadcast_arrays(*[op.duration for op in operators]))
+    return jnp.cumsum(durations, axis=0)[jnp.array(probes)]
+
+
+#
+# flattening preserving identity
+
+
+@dataclasses.dataclass(frozen=True)
+class Skeleton:
+    """Static structure of a sequence (cf. `flatten_operators`), hashable
+
+    Attributes:
+        structure: nested tuples (lists of the sequence) of operator indices
+        operators: for each distinct operator, (treedef, indices of its leaves)
+    """
+
+    structure: object
+    operators: tuple
+
+
+def flatten_operators(sequence):
+    """flatten a sequence, preserving the identity of shared objects
+
+    Unlike `jax.tree.flatten`, an operator appearing several times in the
+    sequence (or an array shared by several operators) gives a single set of
+    leaves, so that the rebuilt sequence shares the same objects (cf. the
+    compiler: shared operators are constants of the scans).
+
+    Returns:
+        leaves (distinct by identity), Skeleton (static, hashable)
+    """
+    leaves, leaf_index = [], {}
+    operators, op_index = [], {}
+
+    def flatten(item):
+        if isinstance(item, Operator):
+            if id(item) not in op_index:
+                op_leaves, treedef = jax.tree.flatten(item)
+                indices = []
+                for leaf in op_leaves:
+                    if id(leaf) not in leaf_index:
+                        leaf_index[id(leaf)] = len(leaves)
+                        leaves.append(leaf)
+                    indices.append(leaf_index[id(leaf)])
+                op_index[id(item)] = len(operators)
+                operators.append((treedef, tuple(indices)))
+            return op_index[id(item)]
+        elif isinstance(item, (list, tuple)):
+            return tuple(flatten(sub) for sub in item)
+        raise ValueError(f"Invalid operator: {item}")
+
+    structure = flatten(sequence)
+    return leaves, Skeleton(structure, tuple(operators))
+
+
+def unflatten_operators(skeleton, leaves):
+    """rebuild a sequence (nested lists) from its leaves (cf. `flatten_operators`)"""
+    operators = [
+        jax.tree.unflatten(treedef, [leaves[i] for i in indices])
+        for treedef, indices in skeleton.operators
+    ]
+
+    def unflatten(item):
+        if isinstance(item, int):
+            return operators[item]
+        return [unflatten(sub) for sub in item]
+
+    return unflatten(skeleton.structure)

epgx/ops/operator.py

@@ -0,0 +1,224 @@
+"""Base operator class and basic operators
+
+Operators are immutable pytrees: numerical arguments are leaves (traceable,
+differentiable), structural arguments are static metadata. Two operators with
+the same type and static metadata have the same pytree structure, whatever
+their numerical values (used by the sequence compiler to detect repeated blocks).
+
+Note: `name` is static metadata, i.e. part of the structure.
+"""
+
+import dataclasses
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from .. import kernels
+from ..axes import is_axis
+from ..common import pytree, repr_value, static
+from ..kernels import affine as maps
+from .dephasing import DephasingAxis
+
+# class decorator for operators (cf. common.pytree)
+operator = pytree
+
+
+@operator
+class Operator:
+    """Base operator class (does nothing)
+
+    Subclasses define:
+        normalize(sample, scanner): user parameters -> layer-1 kernel parameters
+        kernel: layer-1 kernel (class attribute), applied as
+            kernel(state, **normalize(sample, scanner))
+        affine: builder of the equivalent affine map (class attribute, cf.
+            kernels.affine), called as affine(**normalize(sample, scanner));
+            None if the operator is not an affine map (e.g. shift, probe).
+            Used to merge consecutive operators.
+
+    Attributes:
+        duration: operator duration (ms), used for timing (e.g. ADC times)
+        name: (static) operator name (default: class name)
+    """
+
+    duration: jax.Array = 0.0
+    name: str = static()
+
+    def __init__(self, *, duration=None, name=None):
+        self._set(duration=_duration(duration), name=name)
+
+    def __repr__(self):
+        """name, or class name and parameters (e.g. "T(90, 90)")"""
+        if self.name:
+            return self.name
+        params = [
+            repr_value(getattr(self, f.name))
+            for f in dataclasses.fields(self)
+            if f.name not in ("duration", "name") and not f.metadata.get("static")
+        ]
+        static = [
+            repr(getattr(self, f.name))
+            for f in dataclasses.fields(self)
+            if f.name != "name" and f.metadata.get("static")
+        ]
+        return f"{type(self).__name__}({', '.join(static + params + self._repr_duration())})"
+
+    def _repr_duration(self):
+        """["duration=..."] if the duration is not 0, else []"""
+        duration = self.duration
+        if (
+            is_axis(duration)
+            or isinstance(duration, jax.core.Tracer)
+            or np.any(np.asarray(duration) != 0)
+        ):
+            return [f"duration={repr_value(duration)}"]
+        return []
+
+    @property
+    def nshift(self):
+        """number of phase-state shifts"""
+        return 0
+
+    def normalize(self, sample=None, scanner=None):
+        """kernel parameters (dimensionless), from user units"""
+        return {}
+
+    def sample_fields(self):
+        """sample properties read by `normalize` (cf. simulate's sample check)"""
+        return ()
+
+    def sample_conflicts(self, sample):
+        """properties given both by the operator and by the sample (an error)"""
+        return ()
+
+    # layer-1 kernel applied with the normalized parameters (to override)
+    kernel = staticmethod(lambda state: state)
+    # affine map builder (to override), None: not mergeable
+    affine = None
+
+    def __call__(self, state, sample=None, scanner=None):
+        """apply operator to state"""
+        return self.kernel(state, **self.normalize(sample, scanner))
+
+    def _set(self, **values):
+        """set attributes (frozen dataclass)"""
+        for name, value in values.items():
+            object.__setattr__(self, name, value)
+
+
+@operator
+class EmptyOperator(Operator):
+    """Empty operator: does nothing"""
+
+    affine = staticmethod(maps.identity)
+
+
+@operator
+class Wait(EmptyOperator):
+    """Empty operator with given duration"""
+
+    def __init__(self, duration, *, name=None):
+        super().__init__(duration=duration, name=name)
+
+
+@operator
+class Spoiler(Operator):
+    """Perfect spoiler: destroy transverse magnetization
+
+    Args:
+        axis: None (all transverse states), or a dephasing axis (`DephasingAxis`
+            or name, e.g. "diff"): keep the states of index 0 along it (F and
+            Z), i.e. remove the states dephased along that axis only
+        cf. Operator for other keyword arguments
+    """
+
+    direction: str = static()
+    axis: DephasingAxis = None
+
+    def __init__(self, axis=None, *, duration=None, name=None):
+        super().__init__(duration=duration, name=name)
+        direction = axis.name if isinstance(axis, DephasingAxis) else axis
+        axis = axis if isinstance(axis, DephasingAxis) else None
+        if direction is not None and not isinstance(direction, str):
+            raise TypeError(f"Invalid dephasing axis: {direction!r}")
+        self._set(direction=direction, axis=axis)
+
+    def __repr__(self):
+        if self.name:
+            return self.name
+        args = [] if self.direction is None else [f"axis={self.direction!r}"]
+        return f"Spoiler({', '.join(args + self._repr_duration())})"
+
+    def normalize(self, sample=None, scanner=None):
+        return {} if self.direction is None else {"direction": self.direction}
+
+    @staticmethod
+    def kernel(state, direction=None):
+        if direction is None:
+            return kernels.spoil(state)
+        return kernels.spoil_axis(state, direction)
+
+    @property
+    def affine(self):
+        # spoiling along one axis depends on the states: not an affine map
+        return maps.spoil_matrix if self.direction is None else None
+
+
+@operator
+class Reset(Operator):
+    """Reset magnetization (return to equilibrium)"""
+
+    kernel = staticmethod(kernels.reset)
+    affine = staticmethod(maps.reset_matrix)
+
+
+#
+# functions
+
+
+def asparam(value):
+    """convert operator parameter to array (None and Axis are kept)"""
+    if value is None or is_axis(value):
+        return value
+    return jnp.asarray(value)
+
+
+def total(*values):
+    """sum of values, ignoring None (None if all are None)"""
+    values = [value for value in values if value is not None]
+    if not values:
+        return None
+    return sum(values[1:], start=values[0])
+
+
+def check_conflicts(op, sample):
+    """raise if properties are given both by the operator and by the sample"""
+    conflicts = op.sample_conflicts(sample)
+    if conflicts:
+        names = ", ".join(conflicts)
+        raise ValueError(
+            f"{names} given both by {op!r} and by the sample: remove {names} "
+            "from the operator (sample values) or from the sample (operator values)"
+        )
+
+
+# default duration, shared by all operators: a single leaf (constant of the
+# scans, one argument of the compiled program), never stacked
+ZERO_DURATION = jnp.asarray(0.0)
+
+
+def _duration(duration):
+    """check and convert duration (None: shared ZERO_DURATION)"""
+    if duration is None:
+        return ZERO_DURATION
+    values = duration.values if is_axis(duration) else duration
+    if not isinstance(values, jax.core.Tracer) and np.any(np.asarray(values) < 0):
+        raise ValueError(f"Cannot have duration < 0: {duration}")
+    return asparam(duration)
+
+
+# pre-instanciated operators
+NULL = EmptyOperator(name="NULL")
+SPOILER = Spoiler(name="Spoiler")
+RESET = Reset(name="Reset")

epgx/ops/probe.py

@@ -0,0 +1,116 @@
+"""Probe operators (signal acquisition)
+
+ADC phase follows the scanner convention: the ADC phase equals the phase of the
+RF pulse to compensate (`Adc(phase=phi)` after `T(alpha, phi)`).
+"""
+
+import jax
+import numpy as np
+
+from .. import kernels
+from .operator import Operator, asparam, operator, static
+
+DEG = np.pi / 180
+
+
+@operator
+class Probe(Operator):
+    """Base probe operator: leaves the state unchanged, acquires a value"""
+
+    def acquire(self, state, sample=None, scanner=None):
+        """acquired value (to override)"""
+        raise NotImplementedError()
+
+
+@operator
+class Adc(Probe):
+    """ADC: acquire a state attribute, with weights, reduction and phase
+
+    Args:
+        attr: probed attribute ("F0", "Z0", "F", "Fminus", "Z", "states")
+        phase: ADC phase (degree), scanner convention, applied after reduction
+        reduce: None/False (no reduction), True (all batch axes),
+            axis name (str), positional axis (int), or tuple of these
+            An int i refers to the unnamed axis dim_i (whatever its position
+            in the layout); named axes are reduced by name only.
+            If `weights` is set and `reduce` is None: sum along the weights' axes
+        weights: weights (Axis or unnamed array) broadcast with the batch axes
+        cf. Operator for other keyword arguments
+    """
+
+    attr: str = static("F0")
+    reduce: object = static()
+    phase: jax.Array = None
+    weights: jax.Array = None
+
+    def __init__(
+        self,
+        attr="F0",
+        *,
+        phase=None,
+        reduce=None,
+        weights=None,
+        duration=None,
+        name=None,
+    ):
+        if attr not in kernels.probe.ATTRIBUTES:
+            raise ValueError(f"Invalid attribute: {attr}")
+        if isinstance(reduce, list):
+            reduce = tuple(reduce)  # static metadata must be hashable
+        super().__init__(duration=duration, name=name)
+        self._set(
+            attr=attr, reduce=reduce, phase=asparam(phase), weights=asparam(weights)
+        )
+
+    def resolve_axes(self, layout):
+        """convert `reduce` to positions in the batch layout (cf. axes.resolve)"""
+        reduce = self.reduce
+        if reduce is None and self.weights is not None:
+            # weights' axes (named, or positional given the alignment)
+            reduce = layout.names_of(self.weights)
+        if reduce is None or isinstance(reduce, bool):
+            return self
+        reduce = (reduce,) if isinstance(reduce, (str, int)) else tuple(reduce)
+        return self.replace(reduce=tuple(sorted(layout.index(ax) for ax in reduce)))
+
+    def acquire(self, state, sample=None, scanner=None):
+        phase = None if self.phase is None else self.phase * DEG
+        return kernels.adc(
+            state, self.attr, phase=phase, weights=self.weights, axes=self.reduce
+        )
+
+
+@operator
+class Probes(Probe):
+    """Several probes acquired at the same time (acquired value: tuple)
+
+    Args:
+        probes: sequence of Probe operators
+        cf. Operator for other keyword arguments
+    """
+
+    probes: tuple = None
+
+    def __init__(self, probes, *, duration=None, name=None):
+        probes = tuple(probes)
+        if not all(isinstance(probe, Probe) for probe in probes):
+            raise TypeError(f"Expected Probe operators: {probes}")
+        super().__init__(duration=duration, name=name)
+        self._set(probes=probes)
+
+    def __repr__(self):
+        return self.name or f"Probes({', '.join(map(repr, self.probes))})"
+
+    def resolve_axes(self, layout):
+        probes = tuple(
+            probe.resolve_axes(layout) if hasattr(probe, "resolve_axes") else probe
+            for probe in self.probes
+        )
+        return self.replace(probes=probes)
+
+    def acquire(self, state, sample=None, scanner=None):
+        return tuple(probe.acquire(state, sample, scanner) for probe in self.probes)
+
+
+# ADC instance (returns F0)
+ADC = Adc(name="ADC")

epgx/ops/shift.py

@@ -0,0 +1,118 @@
+"""Shift operators (gradient dephasing)"""
+
+import numbers
+
+import numpy as np
+
+from .. import kernels
+from .dephasing import SPACE_AXES, DephasingAxis
+from .operator import Operator, operator, static
+
+
+@operator
+class S(Operator):
+    """Shift operator (dephasing by k cycles across the dephasing period)
+
+    Args:
+        k: static number (1d) or tuple of up to 3 numbers (x, y, z), non zero.
+            A number shifts along x; missing components are 0.
+            Integer shifts are supported by all methods (e.g. Lattice),
+            non-integer shifts by some methods only (e.g. Isochromats).
+        axis: dephasing axis of a scalar shift: a `DephasingAxis` (e.g.
+            `kAxis("diff", kvalue=...)`) or its name; default: "x" (and "y",
+            "z" for vector shifts). Different axes do not interact.
+        cf. Operator for other keyword arguments
+
+    The kernel depends on the representation of the state
+    (`State.representation`): `shift_lattice` (phase states, "lattice"),
+    `shift_samples` (isochromats, "isochromats").
+    """
+
+    k: int = static()
+    direction: str = static()
+    axis: DephasingAxis = None
+
+    def __init__(self, k, *, axis=None, duration=None, name=None):
+        super().__init__(duration=duration, name=name)
+        k = _check_shift(k)
+        direction = None
+        if axis is not None:
+            if isinstance(k, tuple):
+                raise ValueError(f"A shift along an axis is a number, got: {k}")
+            elif isinstance(axis, str):
+                direction, axis = axis, None
+            elif isinstance(axis, DephasingAxis):
+                if axis.kind != "space":
+                    raise NotImplementedError(
+                        f"Shifts along time axes ({axis!r}): not implemented yet"
+                    )
+                direction = axis.name
+            else:
+                raise TypeError(f"Invalid dephasing axis: {axis!r}")
+        self._set(k=k, direction=direction, axis=axis)
+
+    def __repr__(self):
+        if self.name:
+            return self.name
+        axis = [] if self.direction is None else [f"axis={self.direction!r}"]
+        return f"S({', '.join([repr(self.k)] + axis + self._repr_duration())})"
+
+    @property
+    def nshift(self):
+        """number of phase-state shifts (largest component for nd shifts)"""
+        return max(abs(s) for s in self.components)
+
+    @property
+    def integer(self):
+        """True if all shift components are integers"""
+        return all(isinstance(s, int) for s in self.components)
+
+    @property
+    def components(self):
+        """shift components as a tuple"""
+        return self.k if isinstance(self.k, tuple) else (self.k,)
+
+    @property
+    def shifts(self):
+        """shift components by dephasing direction (dict)"""
+        if self.direction is not None:
+            return {self.direction: self.k}
+        return dict(zip(SPACE_AXES, self.components))
+
+    @property
+    def nd(self):
+        """True if the shift needs nd storage (vector, or a named axis)"""
+        return isinstance(self.k, tuple) or self.direction not in (None, "x")
+
+    def normalize(self, sample=None, scanner=None):
+        return {"k": self.shifts}
+
+    @staticmethod
+    def kernel(state, k):
+        """shift kernel for the representation of the state"""
+        if state.representation == "lattice":
+            return kernels.shift_lattice(state, k)
+        elif state.representation == "isochromats":
+            return kernels.shift_samples(state, k)
+        raise ValueError(f"No shift kernel for representation: {state.representation}")
+
+
+#
+# functions
+
+
+def _number(s):
+    """static number: int if integer, else float (bool and arrays are rejected)"""
+    if isinstance(s, bool) or not isinstance(s, numbers.Real):
+        raise TypeError(f"Shifts must be static numbers, got: {s!r}")
+    return int(s) if float(s).is_integer() else float(s)
+
+
+def _check_shift(k):
+    """check and convert shift value (static number or tuple of numbers)"""
+    k = tuple(_number(s) for s in k) if isinstance(k, (tuple, list)) else _number(k)
+    if isinstance(k, tuple) and not 1 <= len(k) <= 3:
+        raise ValueError(f"Shifts have 1 to 3 components (x, y, z), got: {k}")
+    if np.all(np.asarray(k) == 0):
+        raise ValueError("Cannot have k == 0")
+    return k

epgx/ops/transition.py

@@ -0,0 +1,58 @@
+"""Transition operators (instantaneous RF pulses)"""
+
+import jax
+import numpy as np
+
+from .. import kernels
+from ..kernels import affine as maps
+from .operator import Operator, asparam, operator
+
+DEG = np.pi / 180
+
+
+@operator
+class T(Operator):
+    """Transition operator (instantaneous RF pulse)
+
+    Args:
+        alpha: flip angle (degree)
+        phi: phase (degree)
+        cf. Operator for other keyword arguments
+    """
+
+    alpha: jax.Array = None
+    phi: jax.Array = None
+
+    def __init__(self, alpha, phi, *, duration=None, name=None):
+        super().__init__(duration=duration, name=name)
+        self._set(alpha=asparam(alpha), phi=asparam(phi))
+
+    def normalize(self, sample=None, scanner=None):
+        # B1 scaling of the flip angle
+        B1 = 1.0 if scanner is None else scanner.B1
+        return {"alpha": self.alpha * B1 * DEG, "phi": self.phi * DEG}
+
+    kernel = staticmethod(kernels.rotate)
+    affine = staticmethod(maps.rotation_matrix)
+
+
+@operator
+class Phi(Operator):
+    """Phase offset (z-axis rotation)
+
+    Args:
+        phi: phase (degree)
+        cf. Operator for other keyword arguments
+    """
+
+    phi: jax.Array = None
+
+    def __init__(self, phi, *, duration=None, name=None):
+        super().__init__(duration=duration, name=name)
+        self._set(phi=asparam(phi))
+
+    def normalize(self, sample=None, scanner=None):
+        return {"phi": self.phi * DEG}
+
+    kernel = staticmethod(kernels.phase)
+    affine = staticmethod(maps.phase_matrix)