epgx 0.1.dev0__py3-none-any.whl

epgx/plotting.py

@@ -0,0 +1,226 @@
+"""EPG plots / sequence diagrams
+
+The data (`functions.trace`) is separated from the rendering (`plot_epg`):
+rendering never simulates a `Trace`. Requires the `plot` extra (matplotlib,
+cmocean).
+
+Rendering:
+    - all operators are applied
+    - F and Z lines are drawn during every operator with a duration
+    - F lines are coloured by their phase (cmocean's cyclic `phase` colormap)
+    - echoes are marked and labelled by their amplitude
+    - nd lattices (kgrid, yaxis) are not supported yet
+"""
+
+import cmocean
+import numpy as np
+from matplotlib import colors, gridspec, ticker
+from matplotlib import pyplot as plt
+from matplotlib.collections import LineCollection
+
+from . import ops
+from .functions import Trace, trace
+
+# phase colormap: cyclic, nearly constant lightness (faint states remain
+# visible), cf. Thyng et al. Oceanography 2016;29:9-13. doi:10.5670/oceanog.2016.66
+PHASE_CMAP = cmocean.cm.phase
+ECHOES = ("adc", "all", None)
+# keyword arguments of plot_epg (not passed to a parametrized sequence)
+PLOT_ARGUMENTS = (
+    "title",
+    "figname",
+    "calpha",
+    "cwidth",
+    "echoes",
+    "threshold",
+    "echo_threshold",
+)
+
+
+def show():
+    """shortcut for plt.show"""
+    plt.show()
+
+
+def plot_epg(
+    seq,
+    sample=None,
+    scanner=None,
+    *,
+    title=None,
+    figname=None,
+    calpha=0.5,
+    cwidth=0,
+    echoes="adc",
+    threshold=1e-5,
+    echo_threshold=1e-3,
+    **options,
+):
+    """plot RF/gradients timeline and EPG diagram
+
+    Args:
+        seq: Trace, sequence (traced with `sample`, `scanner` and `options`,
+            cf. functions.trace, e.g. `select={"T2": 1}`), or parametrized
+            sequence (`@sequence`): `options` then also hold the arguments of
+            its function, cf. `Sequence.trace` (e.g.
+            `plot_epg(mse, esp=10.0, necho=4)`)
+        title, figname: title and name of the figure
+        calpha, cwidth: power of the state magnitude for the opacity and the
+            width of the lines
+        echoes: mark the amplitude of F0 at the ADCs ("adc"), at all operators
+            ("all"), or not at all (None)
+        threshold: states of smaller magnitude are not drawn
+        echo_threshold: echoes of smaller amplitude are not marked
+
+    Returns:
+        matplotlib Figure
+    """
+    if isinstance(seq, Trace):
+        data = seq
+    elif hasattr(seq, "trace"):
+        # parametrized sequence (layer 3): arguments of its function in options
+        shared = set(PLOT_ARGUMENTS) & set(seq.signature.parameters)
+        if shared:
+            raise ValueError(
+                f"Argument(s) {sorted(shared)} of {seq!r} shadowed by plot_epg's "
+                "own arguments: trace the sequence first (`seq.trace(...)`)"
+            )
+        data = seq.trace(sample=sample, scanner=scanner, **options)
+    else:
+        data = trace(seq, sample, scanner, **options)
+    if data.representation != "lattice":
+        raise ValueError(
+            f"EPG diagrams require phase states, got: {data.representation}"
+        )
+    elif not isinstance(data.axis, str):
+        raise NotImplementedError(
+            f"EPG diagrams of nd lattices (axes {data.axis}) are not supported yet"
+        )
+    if echoes not in ECHOES:
+        raise ValueError(f"Invalid echoes: {echoes!r}, expected one of {ECHOES}")
+
+    F, Z = np.asarray(data.F), np.asarray(data.Z)
+    times, k = np.asarray(data.times, dtype=float), np.asarray(data.coords)
+    operators = data.operators
+    center = k.size // 2
+
+    fig = plt.figure(figname, figsize=(8, 6))
+    grid = gridspec.GridSpec(3, 1, figure=fig, height_ratios=(1, 1, 8))
+    ax_rf = fig.add_subplot(grid[0, 0])
+    ax_grad = fig.add_subplot(grid[1, 0], sharex=ax_rf)
+    ax_epg = fig.add_subplot(grid[2, 0], sharex=ax_rf)
+
+    # state lines
+    fsegs, fcolors, fwidths = [], [], []
+    zsegs, zcolors, zwidths = [], [], []
+    cmap = PHASE_CMAP
+    norm = colors.Normalize(-180, 180)
+    for i, op in enumerate(operators):
+        start, end = times[i], times[i + 1]
+        if end <= start:
+            continue
+        shift = op.k if isinstance(op, ops.S) else 0
+        fmag, zmag = np.minimum(np.abs(F[i + 1]), 1), np.minimum(np.abs(Z[i + 1]), 1)
+        for j in np.flatnonzero(fmag > threshold):
+            fsegs.append([(start, k[j] - shift), (end, k[j])])
+            fcolors.append(cmap(norm(np.angle(F[i + 1, j], deg=True)))[:3])
+            fwidths.append(fmag[j])
+        # Z(-k) = conj(Z(k)): positive states only
+        for j in np.flatnonzero((zmag > threshold) & (k >= 0)):
+            zsegs.append([(start, k[j]), (end, k[j])])
+            zcolors.append((0, 0, 0))
+            zwidths.append(zmag[j])
+    for label, segs, rgb, mags, style in [
+        ("F", fsegs, fcolors, fwidths, "-"),
+        ("Z", zsegs, zcolors, zwidths, ":"),
+    ]:
+        mags = np.asarray(mags, dtype=float)
+        rgba = np.concatenate(
+            [np.reshape(rgb, (-1, 3)), mags[:, None] ** calpha], axis=1
+        )
+        lines = LineCollection(
+            segs, colors=rgba, linewidths=2 * mags**cwidth, linestyles=style
+        )
+        lines.set_label(label)
+        ax_epg.add_collection(lines)
+
+    # operators: RF pulses, gradients, ADCs and echoes
+    kmax = max([1] + [abs(y) for seg in fsegs for _, y in seg])
+    for i, op in enumerate(operators):
+        start, end = times[i], times[i + 1]
+        if isinstance(op, ops.T):
+            alpha, phi = float(np.ravel(op.alpha)[0]), float(np.ravel(op.phi)[0])
+            ax_rf.vlines(end, 0, alpha, color="k")
+            va = "bottom" if alpha > 0 else "top"
+            ax_rf.annotate(
+                f"{alpha:.0f}°", (end, alpha + 2 * np.sign(alpha)), va=va, ha="center"
+            )
+            if not np.isclose(phi, 0):
+                # opposite side of the flip angle
+                ax_rf.annotate(
+                    f"φ={phi:.0f}°",
+                    (end, 0),
+                    va="top" if alpha >= 0 else "bottom",
+                    ha="center",
+                    fontsize=8,
+                )
+            ax_epg.scatter(
+                end, 0, marker="o", color="gray", facecolors="white", zorder=10
+            )
+            ax_epg.axvline(end, linestyle=":", color="gray", alpha=0.5)
+        elif isinstance(op, ops.S):
+            ax_grad.fill_between([start, end], [op.k] * 2, color="gray", alpha=0.3)
+        elif isinstance(op, ops.Probe):
+            ax_epg.scatter(end, 0, marker="v", color="gray", zorder=10)
+
+        # "all": states at the end of each time point (once per time)
+        last = i + 2 == times.size or times[i + 2] > end
+        if (echoes == "all" and last) or (
+            echoes == "adc" and isinstance(op, ops.Probe)
+        ):
+            amplitude = np.abs(F[i + 1, center])
+            if amplitude > echo_threshold:
+                ax_epg.annotate(
+                    f"{amplitude:.2f}",
+                    (end, 0),
+                    xytext=(0, 6),
+                    textcoords="offset points",
+                    ha="center",
+                    va="bottom",
+                    fontsize=8,
+                    color="tab:red",
+                    zorder=11,
+                )
+
+    # axes
+    margin = 3e-2 * max(times[-1], 1e-9)
+    ax_epg.set_xlim(times[0] - margin, times[-1] + margin)
+    ax_epg.set_ylim(-1.05 * kmax, 1.05 * kmax)
+    ax_epg.yaxis.set_major_locator(ticker.MaxNLocator(integer=True))
+    ax_epg.set_ylabel(data.axis)
+    ax_epg.set_xlabel("time (ms)")
+    ax_epg.axhline(0, color="k", zorder=-1)
+    # phase colorbar (inset: the time axes stay aligned)
+    inset = ax_epg.inset_axes([0.015, 0.72, 0.012, 0.25])
+    colorbar = fig.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), cax=inset)
+    colorbar.set_ticks([-180, 0, 180])
+    colorbar.ax.tick_params(labelsize=8)
+    colorbar.set_label("phase (°)", fontsize=8)
+
+    ax_rf.set_ylim(-180, 180)
+    for ax, label in [(ax_rf, "Rf"), (ax_grad, "G")]:
+        ax.axhline(0, color="k")
+        ax.annotate(
+            label,
+            xy=(-1e-2, 0.5),
+            ha="right",
+            va="center",
+            xycoords="axes fraction",
+            weight="bold",
+        )
+        ax.axis("off")
+
+    if title:
+        fig.suptitle(title)
+    fig.tight_layout()
+    return fig

epgx/progress.py

@@ -0,0 +1,81 @@
+"""Progress display of simulations (`simulate(..., disp=True)`)
+
+The operators run in compiled code (lax.scan, jit), so progress is reported
+from the program with `jax.debug.callback`: once per top-level unrolled
+segment and per iteration of top-level scans (cf. compiler monitor). Each
+callback takes a scalar of the state, so that it follows the computation.
+
+The host-side `Progress` opens a bar on the first tick of a run and closes it
+when the total number of operators is reached: a cached compiled program shows
+one bar per call. Display: tqdm if installed (e.g. notebooks), else a text bar.
+
+Limits: compilation is not shown; under `grad`, only the forward pass ticks;
+under `vmap`, the callbacks run for each element.
+"""
+
+import functools
+import sys
+
+import jax
+import jax.numpy as jnp
+
+DESCRIPTION = "Simulating"
+
+
+class Progress:
+    """host-side progress of a run of `total` operators"""
+
+    def __init__(self, total, description=DESCRIPTION):
+        self.total = total
+        self.description = description
+        self.done = 0
+        self._bar = None
+
+    def tick(self, count, *_):
+        """`count` more operators done (callback from the compiled program)"""
+        if self._bar is None:
+            self._bar = make_bar(self.total, self.description)
+        self._bar.update(count)
+        self.done += count
+        if self.done >= self.total:
+            self._bar.close()
+            self._bar, self.done = None, 0
+
+    def monitor(self, state, count):
+        """compiler monitor: tick from the program (cf. compiler.Program.run)"""
+        # data dependency on the state: the tick follows the computation
+        dependency = jnp.real(jnp.ravel(state.F)[0])
+        jax.debug.callback(functools.partial(self.tick, count), dependency)
+
+
+def make_bar(total, description):
+    """tqdm bar if available, else text bar (same update/close interface)"""
+    try:
+        from tqdm.auto import tqdm
+    except ImportError:
+        return TextBar(total, description)
+    return tqdm(total=total, desc=description, unit="op")
+
+
+class TextBar:
+    """text progress bar"""
+
+    def __init__(self, total, description="", size=40, file=None):
+        self.total, self.description, self.size = total, description, size
+        self.file = file or sys.stdout
+        self.done = 0
+        self._show()
+
+    def update(self, count):
+        self.done = min(self.done + count, self.total)
+        self._show()
+
+    def close(self):
+        self.file.write("\n")
+        self.file.flush()
+
+    def _show(self):
+        filled = int(self.size * self.done / max(self.total, 1))
+        bar = "#" * filled + "." * (self.size - filled)
+        self.file.write(f"\r{self.description}: [{bar}] {self.done}/{self.total}")
+        self.file.flush()

epgx/sample.py

@@ -0,0 +1,142 @@
+"""Sample (tissue properties) and Scanner (hardware properties)
+
+Field values are scalars, Axis (named axes) or unnamed arrays. In Sample and
+Scanner, 1d arrays are automatically named after their field, e.g.
+`Sample(T2=[30, 40, 50])` defines the axis "T2" (grid of values).
+
+Units: ms (T1, T2), kHz (df, B0), relative (PD, B1).
+"""
+
+import jax
+import jax.numpy as jnp
+
+from .axes import Axis, is_axis
+from .common import pytree, repr_value, static
+
+COMPARTMENT = "compartment"
+
+
+def _field(name, value):
+    """convert field value: None, scalar, Axis, or array (1d: named after field)"""
+    if value is None or is_axis(value):
+        return value
+    value = jnp.asarray(value)
+    if value.ndim == 1:
+        return Axis(name, value)
+    return value
+
+
+@pytree
+class Sample:
+    """Tissue properties
+
+    Args:
+        PD: proton density (or compartment fraction)
+        T1: longitudinal relaxation time (ms), None: no T1 relaxation
+        T2: transverse relaxation time (ms), None: no T2 relaxation
+        df: frequency offset (kHz), e.g. chemical shift
+
+    Attributes:
+        compartments: (static) compartment labels (cf. `Sample.compartments`)
+    """
+
+    PD: jax.Array = None
+    T1: jax.Array = None
+    T2: jax.Array = None
+    df: jax.Array = None
+    compartments: tuple = static()
+
+    FIELDS = ("PD", "T1", "T2", "df")
+
+    def __init__(self, PD=1.0, T1=None, T2=None, df=None):
+        for name, value in zip(self.FIELDS, (PD, T1, T2, df)):
+            object.__setattr__(self, name, _field(name, value))
+        object.__setattr__(self, "compartments", None)
+
+    def __repr__(self):
+        fields = [
+            f"{name}={repr_value(getattr(self, name))}"
+            for name in self.FIELDS
+            if getattr(self, name) is not None
+        ]
+        if self.compartments:
+            fields.append(f"compartments={self.compartments}")
+        return f"Sample({', '.join(fields)})"
+
+    @classmethod
+    def from_compartments(cls, **samples):
+        """multi-compartment sample, stacked along the "compartment" axis
+
+        PD values are the compartment fractions. Field values must be scalars or
+        Axis (named). Fields unset in some compartments only take neutral values
+        (T1, T2: infinite, df: 0).
+
+        Example:
+            Sample.from_compartments(
+                water=Sample(PD=0.7, T1=1000, T2=30),
+                fat=Sample(PD=0.3, T1=300, T2=80, df=-0.42),
+            )
+        """
+        if not samples:
+            raise ValueError("At least one compartment is required")
+        labels = tuple(samples)
+        new = object.__new__(cls)
+        for name in cls.FIELDS:
+            values = [getattr(sample, name) for sample in samples.values()]
+            object.__setattr__(new, name, _stack(name, values))
+        object.__setattr__(new, "compartments", labels)
+        return new
+
+
+@pytree
+class Scanner:
+    """Hardware properties
+
+    Args:
+        B1: relative transmit field (scales flip angles)
+        B0: frequency offset (kHz)
+    """
+
+    B1: jax.Array = None
+    B0: jax.Array = None
+
+    def __init__(self, B1=1.0, B0=0.0):
+        object.__setattr__(self, "B1", _field("B1", B1))
+        object.__setattr__(self, "B0", _field("B0", B0))
+
+    def __repr__(self):
+        return f"Scanner(B1={repr_value(self.B1)}, B0={repr_value(self.B0)})"
+
+
+#
+# functions
+
+
+# neutral values of unset fields, when other compartments set them
+NEUTRAL = {"T1": jnp.inf, "T2": jnp.inf, "df": 0.0}
+
+
+def _stack(name, values):
+    """stack field values of several compartments along the compartment axis"""
+    if all(value is None for value in values):
+        return None
+    values = [NEUTRAL[name] if value is None else value for value in values]
+    if any(not is_axis(v) and jnp.ndim(v) > 0 for v in values):
+        raise ValueError(f"Field {name}: compartment values must be scalars or Axis")
+
+    # union of named axes (by first appearance) and their sizes
+    sizes = {}
+    for value in filter(is_axis, values):
+        for axis, size in value.sizes.items():
+            if sizes.setdefault(axis, size) != size:
+                raise ValueError(f"Inconsistent sizes for axis {axis}")
+    if COMPARTMENT in sizes:
+        raise ValueError(f"Reserved axis name: {COMPARTMENT}")
+    names = tuple(sizes)
+    shape = tuple(sizes.values())
+
+    stacked = [
+        value.expand(names, sizes) if is_axis(value) else jnp.broadcast_to(value, shape)
+        for value in values
+    ]
+    return Axis((COMPARTMENT,) + names, jnp.stack(stacked))