epgx 0.1.dev0__tar.gz

epgx-0.1.dev0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2023-2026, py-baudin
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

epgx-0.1.dev0/PKG-INFO

@@ -0,0 +1,64 @@
+Metadata-Version: 2.4
+Name: epgx
+Version: 0.1.dev0
+Summary: Extended Phase Graph (EPG) simulations in JAX
+Author: Pierre-Yves Baudin
+License-Expression: BSD-3-Clause
+Project-URL: Repository, https://github.com/py-baudin/epgx
+Keywords: MRI,extended phase graph,EPG,simulation,JAX,differentiable
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax
+Requires-Dist: numpy
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Provides-Extra: plot
+Requires-Dist: matplotlib; extra == "plot"
+Requires-Dist: cmocean; extra == "plot"
+Provides-Extra: xarray
+Requires-Dist: xarray; extra == "xarray"
+Provides-Extra: dev
+Requires-Dist: epgx[plot,test,xarray]; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Dynamic: license-file
+
+# epgx
+
+A JAX implementation of the Extended Phase Graph (EPG) formalism, successor of
+[epgpy](https://github.com/py-baudin/epgpy).
+
+Work in progress: the API may change between versions.
+
+## Installation
+
+``` shell
+pip install epgx            # core (jax, numpy)
+pip install epgx[plot]      # with EPG diagrams (matplotlib, cmocean)
+```
+
+Development:
+
+``` shell
+conda create -n epgx python=3.12
+conda activate epgx
+pip install -e .[dev]
+pip install --group parity  # epgpy, for the parity tests (tests/parity)
+pytest
+```
+
+## Precision
+
+`epgx` follows JAX's precision setting (float32 by default) and never changes it.
+For float64, enable it at the start of your program:
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)
+```

epgx-0.1.dev0/README.md

@@ -0,0 +1,33 @@
+# epgx
+
+A JAX implementation of the Extended Phase Graph (EPG) formalism, successor of
+[epgpy](https://github.com/py-baudin/epgpy).
+
+Work in progress: the API may change between versions.
+
+## Installation
+
+``` shell
+pip install epgx            # core (jax, numpy)
+pip install epgx[plot]      # with EPG diagrams (matplotlib, cmocean)
+```
+
+Development:
+
+``` shell
+conda create -n epgx python=3.12
+conda activate epgx
+pip install -e .[dev]
+pip install --group parity  # epgpy, for the parity tests (tests/parity)
+pytest
+```
+
+## Precision
+
+`epgx` follows JAX's precision setting (float32 by default) and never changes it.
+For float64, enable it at the start of your program:
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)
+```

epgx-0.1.dev0/pyproject.toml

@@ -0,0 +1,65 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "epgx"
+description = "Extended Phase Graph (EPG) simulations in JAX"
+version = "0.1.dev0"
+readme = "README.md"
+license = "BSD-3-Clause"
+license-files = ["LICENSE"]
+authors = [{ name = "Pierre-Yves Baudin" }]
+requires-python = ">=3.11"
+dependencies = ["jax", "numpy"]
+keywords = ["MRI", "extended phase graph", "EPG", "simulation", "JAX", "differentiable"]
+classifiers = [
+    "Development Status :: 2 - Pre-Alpha",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Medical Science Apps.",
+    "Topic :: Scientific/Engineering :: Physics",
+]
+
+[project.urls]
+Repository = "https://github.com/py-baudin/epgx"
+
+[project.optional-dependencies]
+test = ["pytest"]
+plot = ["matplotlib", "cmocean"]
+xarray = ["xarray"]
+dev = ["epgx[test,plot,xarray]", "ruff"]
+
+[dependency-groups]
+# reference implementation for parity tests (pinned commit); not published
+# (PyPI rejects direct URL dependencies): `pip install --group parity`
+parity = [
+    "epgpy @ git+https://github.com/py-baudin/epgpy@82eebbf0e8a59b100682758c2a04a889fc1ee55c",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+# importlib: allows test modules with the same name in tests/ and tests/parity/
+addopts = "--import-mode=importlib"
+markers = [
+    "parity: comparison with epgpy (reference implementation)",
+    "slow: long tests (skip with -m \"not slow\")",
+]
+
+[tool.ruff]
+line-length = 88
+target-version = "py311"
+extend-exclude = ["*.md"]
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I"]
+# E741: allow EPG-style short names (e.g. `l`), E501 left to the formatter
+ignore = ["E741", "E501"]
+
+[tool.ruff.lint.per-file-ignores]
+# re-exports
+"__init__.py" = ["F401"]

epgx-0.1.dev0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

epgx-0.1.dev0/src/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-0.1.dev0/src/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-0.1.dev0/src/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