epgx 0.1.dev0__py3-none-any.whl

epgx/__init__.py

@@ -0,0 +1,19 @@
+"""Extended Phase Graph (EPG) simulations in JAX
+
+Note: numerical precision follows JAX's configuration (float32 by default).
+For float64 precision, enable x64 *before* any computation:
+
+    import jax
+    jax.config.update("jax_enable_x64", True)
+
+"""
+
+__version__ = "0.1.dev0"
+
+from .axes import Axis, resolve
+from .functions import Trace, get_adc_times, getnshift, simulate, trace
+from .lowering import compile_sequence
+from .methods import Isochromats, Lattice
+from .ops import DephasingAxis, kAxis, tAxis
+from .sample import Sample, Scanner
+from .sequences import Sequence, sequence

epgx/axes.py

@@ -0,0 +1,285 @@
+"""Named axes and resolution to a positional layout
+
+Batch dimensions are either:
+    - named: `Axis(name, values)` (`name` may be a tuple for nd values)
+    - unnamed (positional): non-scalar arrays, aligned with the implicit axes
+      dim_0, dim_1... (left-aligned by default, cf. `align`)
+
+`resolve` collects all axes of a set of pytrees (sequence, sample, scanner...),
+orders them (positional axes first, then named axes by first appearance, or
+explicit `dims`), and converts every leaf into an array broadcastable to the
+common batch shape. This is done once, on static shapes, before simulation.
+"""
+
+import jax
+import jax.numpy as jnp
+
+from .common import pytree, static
+
+# alignment of unnamed arrays: "left" (default) or "right" (NumPy)
+ALIGNMENT = "left"
+
+POSITIONAL = "dim_{}"
+
+
+@pytree
+class Axis:
+    """Named array (named axis when 1d)
+
+    The names cover the leading dimensions of the values; the remaining
+    dimensions are core dimensions, declared by `@sequence` variables (e.g.
+    `Axis("candidate", alphas)` of shape (10, 400) for `alpha[n]`). At layer 2
+    nothing is declared, so every dimension needs a name (checked when the
+    axes are collected, cf. `get_layout`).
+
+    Args:
+        name: axis name (str), or tuple of names (leading dimensions of values)
+        values: array with at least one dimension per name
+    """
+
+    values: jax.Array = None
+    name: object = static()
+
+    def __init__(self, name, values):
+        values = jnp.asarray(values)
+        names = (name,) if isinstance(name, str) else tuple(name)
+        if not all(isinstance(n, str) for n in names):
+            raise TypeError(f"Axis names must be strings: {name}")
+        elif len(set(names)) != len(names):
+            raise ValueError(f"Duplicate axis names: {names}")
+        elif values.ndim < len(names):
+            raise ValueError(
+                f"More Axis names than dimensions: {names}, {values.shape}"
+            )
+        object.__setattr__(self, "values", values)
+        object.__setattr__(self, "name", names[0] if len(names) == 1 else names)
+
+    def __repr__(self):
+        return f"Axis({self.name!r}, <{tuple(self.shape)}>)"
+
+    @property
+    def names(self):
+        return (self.name,) if isinstance(self.name, str) else tuple(self.name)
+
+    @property
+    def shape(self):
+        return jnp.shape(self.values)
+
+    @property
+    def sizes(self):
+        """dict of axis sizes"""
+        return dict(zip(self.names, self.shape))
+
+    def expand(self, names, sizes=None):
+        """return values with axes ordered as `names` (missing axes: size 1 or sizes)"""
+        values = place(self.values, self.names, names)
+        if sizes is not None:
+            values = jnp.broadcast_to(values, tuple(sizes[n] for n in names))
+        return values
+
+
+def is_axis(obj):
+    return isinstance(obj, Axis)
+
+
+def align(value, ndim, alignment=None):
+    """expand an unnamed array to ndim positional axes
+
+    Args:
+        value: array with at most ndim dimensions
+        ndim: number of positional axes
+        alignment: "left" (new axes appended) or "right" (prepended,
+            NumPy), default: ALIGNMENT
+    """
+    value = jnp.asarray(value)
+    return jnp.reshape(value, align_shape(value.shape, ndim, alignment))
+
+
+def align_shape(shape, ndim, alignment=None):
+    """aligned shape (cf. `align`)"""
+    alignment = alignment or ALIGNMENT
+    shape = tuple(shape)
+    diff = ndim - len(shape)
+    if diff < 0:
+        raise ValueError(f"Cannot align shape {shape} to {ndim} dimensions")
+    elif alignment == "left":
+        return shape + (1,) * diff
+    elif alignment == "right":
+        return (1,) * diff + shape
+    raise ValueError(f"Unknown alignment: {alignment}")
+
+
+def place(values, names, dims):
+    """transpose values (with given axis names) to `dims`, insert missing axes"""
+    names = tuple(names)
+    if not set(names) <= set(dims):
+        raise ValueError(f"Unknown axes {set(names) - set(dims)} in {dims}")
+    perm = sorted(range(len(names)), key=lambda i: dims.index(names[i]))
+    values = jnp.transpose(values, perm)
+    ordered = [names[i] for i in perm]
+    shape = [values.shape[ordered.index(d)] if d in names else 1 for d in dims]
+    return jnp.reshape(values, shape)
+
+
+@pytree
+class Layout:
+    """Batch layout: ordered axis names and sizes
+
+    Attributes:
+        dims: (static) ordered axis names
+        shape: (static) axis sizes
+        npos: (static) number of positional axes
+        alignment: (static) alignment of unnamed arrays
+    """
+
+    dims: tuple = static(())
+    shape: tuple = static(())
+    npos: int = static(0)
+    alignment: str = static("left")
+
+    def __init__(self, dims=(), shape=(), npos=0, alignment=None):
+        object.__setattr__(self, "dims", tuple(dims))
+        object.__setattr__(self, "shape", tuple(shape))
+        object.__setattr__(self, "npos", npos)
+        object.__setattr__(self, "alignment", alignment or ALIGNMENT)
+
+    def __repr__(self):
+        axes = ", ".join(f"{d}={s}" for d, s in zip(self.dims, self.shape))
+        return f"Layout({axes})"
+
+    @property
+    def positional(self):
+        """names of the positional axes"""
+        return tuple(POSITIONAL.format(i) for i in range(self.npos))
+
+    def names_of(self, value):
+        """axis names of a leaf (Axis or unnamed array)"""
+        if is_axis(value):
+            return value.names
+        ndim = jnp.ndim(value)
+        if ndim == 0:
+            return ()
+        elif self.alignment == "left":
+            return self.positional[:ndim]
+        return self.positional[self.npos - ndim :]
+
+    def index(self, axis):
+        """position of an axis given by name, or by positional index (int)"""
+        if isinstance(axis, str):
+            if axis not in self.dims:
+                raise ValueError(f"Unknown axis: {axis}, available: {self.dims}")
+            return self.dims.index(axis)
+        positional = self.positional
+        if not -len(positional) <= axis < len(positional):
+            raise ValueError(f"Invalid positional axis {axis} for {self.npos} axes")
+        return self.dims.index(positional[axis])
+
+    def reduced_dims(self, reduce):
+        """axis names remaining after reduction
+
+        Args:
+            reduce: None/False (no reduction), True (all axes),
+                or tuple of positions (e.g. resolved `Adc.reduce`)
+        """
+        if reduce is None or reduce is False:
+            return self.dims
+        elif reduce is True:
+            return ()
+        return tuple(d for i, d in enumerate(self.dims) if i not in reduce)
+
+    def convert(self, value):
+        """convert a leaf to an array broadcastable to the batch shape"""
+        if is_axis(value):
+            return place(value.values, value.names, self.dims)
+        elif jnp.ndim(value) == 0:
+            return value
+        aligned = align(value, self.npos, self.alignment)
+        return place(aligned, self.positional, self.dims)
+
+
+def get_layout(*trees, dims=None, alignment=None):
+    """collect the batch axes of pytrees (Axis or unnamed array leaves)
+
+    Args:
+        trees: pytrees (sequence, sample, scanner...)
+        dims: explicit order of all axes (positional axes are named dim_i)
+        alignment: alignment of unnamed arrays (default: ALIGNMENT)
+
+    Returns:
+        Layout
+    """
+    alignment = alignment or ALIGNMENT
+    leaves = jax.tree.leaves(trees, is_leaf=is_axis)
+
+    # named axes, by first appearance
+    named = {}
+    for leaf in filter(is_axis, leaves):
+        if len(leaf.names) != len(leaf.shape):
+            raise ValueError(
+                f"Axis {leaf.name!r} names {len(leaf.names)} of {len(leaf.shape)} "
+                f"dimensions {tuple(leaf.shape)}: at layer 2 every dimension needs "
+                "a name (in @sequence, the other dimensions must be the declared "
+                "core of the variable)"
+            )
+        for name, size in leaf.sizes.items():
+            if named.setdefault(name, size) != size:
+                raise ValueError(
+                    f"Inconsistent sizes for axis {name}: {named[name]}, {size}"
+                )
+
+    # positional axes
+    shapes = [jnp.shape(leaf) for leaf in leaves if not is_axis(leaf)]
+    npos = max((len(shape) for shape in shapes), default=0)
+    aligned = [align_shape(shape, npos, alignment) for shape in shapes]
+    try:
+        posshape = jnp.broadcast_shapes(*aligned) if aligned else ()
+    except ValueError:
+        raise ValueError(
+            f"Incompatible unnamed array shapes ({alignment}-aligned): {shapes}"
+        )
+    positional = {POSITIONAL.format(i): size for i, size in enumerate(posshape)}
+    if set(positional) & set(named):
+        raise ValueError(f"Reserved axis names: {set(positional) & set(named)}")
+
+    sizes = {**positional, **named}
+    if dims is None:
+        dims = tuple(sizes)
+    elif set(dims) != set(sizes) or len(dims) != len(sizes):
+        raise ValueError(f"`dims` must contain all axes {tuple(sizes)}, got: {dims}")
+    return Layout(dims, tuple(sizes[d] for d in dims), npos, alignment)
+
+
+def get_coords(*trees):
+    """coordinates of 1d named axes (values of their first occurrence)"""
+    coords = {}
+    for leaf in filter(is_axis, jax.tree.leaves(trees, is_leaf=is_axis)):
+        if len(leaf.names) == 1:
+            coords.setdefault(leaf.name, leaf.values)
+    return coords
+
+
+def resolve(*trees, dims=None, alignment=None):
+    """convert named and unnamed leaves to a common positional layout
+
+    Objects defining a `resolve_axes(layout)` method (e.g. ADC with named
+    reduction axes) are given the layout before conversion.
+
+    Args:
+        trees: pytrees (sequence, sample, scanner...)
+        dims, alignment: cf. `get_layout`
+
+    Returns:
+        layout, resolved trees (same structure, Axis leaves replaced by arrays)
+    """
+    layout = get_layout(*trees, dims=dims, alignment=alignment)
+
+    def has_resolve(obj):
+        return hasattr(obj, "resolve_axes") and not isinstance(obj, type)
+
+    trees = jax.tree.map(
+        lambda obj: obj.resolve_axes(layout) if has_resolve(obj) else obj,
+        trees,
+        is_leaf=lambda obj: has_resolve(obj) or is_axis(obj),
+    )
+    trees = jax.tree.map(layout.convert, trees, is_leaf=is_axis)
+    return layout, trees

epgx/common.py

@@ -0,0 +1,63 @@
+"""Common utilities"""
+
+import copy
+import dataclasses
+
+import jax
+import numpy as np
+
+
+def static(default=None):
+    """dataclass field stored as static metadata (not a pytree leaf)"""
+    return dataclasses.field(default=default, metadata={"static": True})
+
+
+def pytree(cls):
+    """class decorator: frozen dataclass registered as a pytree
+
+    Fields declared with `static()` are static metadata, others are leaves.
+    Unflattening does not call `__init__`, so constructors may convert and
+    validate their arguments. A `replace(**changes)` method is added (copy
+    without calling `__init__`).
+    """
+    cls = dataclasses.dataclass(frozen=True, eq=False, init=False, repr=False)(cls)
+    fields = dataclasses.fields(cls)
+    data = tuple(f.name for f in fields if not f.metadata.get("static"))
+    meta = tuple(f.name for f in fields if f.metadata.get("static"))
+
+    def flatten(obj):
+        children = [(jax.tree_util.GetAttrKey(n), getattr(obj, n)) for n in data]
+        return children, tuple(getattr(obj, n) for n in meta)
+
+    def unflatten(aux, children):
+        obj = object.__new__(cls)
+        for name, value in zip(meta + data, tuple(aux) + tuple(children)):
+            object.__setattr__(obj, name, value)
+        return obj
+
+    jax.tree_util.register_pytree_with_keys(cls, flatten, unflatten)
+    if "replace" not in cls.__dict__:
+        cls.replace = _replace
+    return cls
+
+
+def repr_value(value):
+    """short representation of a parameter value (scalar, array, Axis, traced)"""
+    if value is None:
+        return "None"
+    elif isinstance(value, jax.core.Tracer):
+        return "<traced>"
+    elif hasattr(value, "names"):  # Axis
+        return repr(value)
+    value = np.asarray(value)
+    if value.size == 1:
+        return f"{value.item():.4g}"
+    return f"<{value.shape}>"
+
+
+def _replace(self, **changes):
+    """return a copy with given fields replaced (`__init__` is not called)"""
+    new = copy.copy(self)
+    for name, value in changes.items():
+        object.__setattr__(new, name, value)
+    return new