From 9dd56bf0620382537544633a4c90e540b48a1673 Mon Sep 17 00:00:00 2001 From: Matteo Visconti di Oleggio Castello Date: Sat, 19 Sep 2026 17:26:51 -0700 Subject: [PATCH 1/3] ENH add a Tractogram dataview for streamlines cortex.Tractogram holds a bundle of streamlines -- diffusion tractography output such as pyAFQ or DIPY bundles -- as a Dataview, so it travels through Dataset and the viewer machinery like any other data. It is the odd one out among the dataviews in that it is not braindata: the geometry is a flat (N, 3) array of points plus an (M+1,) offset table delimiting the M streamlines, which is the TRX convention. tract = cortex.Tractogram.from_trx("bundles.trx", "S1") tract = cortex.Tractogram.from_streamlines(list_of_arrays, "S1") Per-point and per-streamline scalars and named groups of streamlines (pyAFQ writes one per bundle) come across from the file, and select/get_group/ subsample return new tractograms with the group memberships remapped. Coloring is resolved in Python by vertex_colors: by local tangent direction, by a constant color, or through a colormap from a named scalar. Reading TRX needs trx-python, a new optional `tractography` extra. The import is lazy, so pycortex works normally without it, and the tests importorskip. trx-python requires Python 3.11, and the test dependency group pulls the extra in, so the CI matrix drops its 3.10 row -- on 3.10 the install fails outright rather than the TRX tests quietly skipping. This is a container only. Rendering comes next; HDF5 persistence is deliberately NotImplementedError, since a Dataset carrying one cannot be saved, and the TRX file stays the source of truth. Co-Authored-By: Claude Opus 5 --- .github/workflows/run_tests.yml | 5 +- cortex/__init__.py | 2 +- cortex/dataset/__init__.py | 1 + cortex/dataset/tractogram.py | 662 ++++++++++++++++++++++++++++++++ cortex/tests/test_tractogram.py | 487 +++++++++++++++++++++++ docs/dataset.rst | 29 ++ pyproject.toml | 10 + 7 files changed, 1194 insertions(+), 2 deletions(-) create mode 100644 cortex/dataset/tractogram.py create mode 100644 cortex/tests/test_tractogram.py diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 60866b97c..e12327e2c 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -17,7 +17,10 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + # 3.11 is the floor because the test group pulls in trx-python (via + # the `tractography` extra), which requires it. On 3.10 the install + # fails outright rather than the TRX tests quietly skipping. + python-version: ["3.11", "3.12", "3.13", "3.14"] max-parallel: 5 steps: diff --git a/cortex/__init__.py b/cortex/__init__.py index 7a0b12306..6e6060660 100644 --- a/cortex/__init__.py +++ b/cortex/__init__.py @@ -1,6 +1,6 @@ # emacs: -*- coding: utf-8; mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- # vi: set fileencoding=utf-8 ft=python sts=4 ts=4 sw=4 et: -from cortex.dataset import Dataset, Volume, Vertex, VolumeRGB, VertexRGB, Volume2D, Vertex2D, Colors +from cortex.dataset import Dataset, Volume, Vertex, VolumeRGB, VertexRGB, Volume2D, Vertex2D, Colors, Tractogram from cortex import align, volume, quickflat, webgl, segment, options from cortex.database import db from cortex.utils import * diff --git a/cortex/dataset/__init__.py b/cortex/dataset/__init__.py index 9dc6b5764..07cb4b92f 100644 --- a/cortex/dataset/__init__.py +++ b/cortex/dataset/__init__.py @@ -3,4 +3,5 @@ from __future__ import annotations from .views import Volume, Vertex, VolumeRGB, VertexRGB, Volume2D, Vertex2D, Dataview, Dataview2D, _from_hdf_data, Colors, DataviewJSON +from .tractogram import Tractogram from .dataset import Dataset, normalize, DatasetLike \ No newline at end of file diff --git a/cortex/dataset/tractogram.py b/cortex/dataset/tractogram.py new file mode 100644 index 000000000..41c7a7124 --- /dev/null +++ b/cortex/dataset/tractogram.py @@ -0,0 +1,662 @@ +"""Streamline (tractography) dataviews. + +This module defines :class:`Tractogram`, a :class:`~cortex.dataset.views.Dataview` +subclass that holds a bundle of 3-D polylines (streamlines) plus optional +per-vertex/per-streamline scalar data and named groups, for rendering +alongside cortical surfaces in the pycortex WebGL viewer. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Sequence, Union + +import numpy as np +import numpy.typing as npt + +from .braindata import _hash +from .views import Dataview + +if TYPE_CHECKING: + from trx.trx_file_memmap import TrxFile + +#: Accepted spellings for the ``color`` display attribute. +ColorSpec = Union[str, tuple, list, npt.NDArray] + + +class Tractogram(Dataview): + """Encapsulates a bundle of streamlines (a tractogram) for display. + + Unlike :class:`~cortex.dataset.views.Volume` and + :class:`~cortex.dataset.views.Vertex`, a ``Tractogram`` is not backed by + :class:`~cortex.dataset.braindata.BrainData`: its geometry is a flat + array of 3-D points plus an offset table delimiting individual + streamlines (the TRX convention), rather than a volume- or + surface-sampled array. + + Parameters + ---------- + points : (N, 3) array_like + Streamline vertex coordinates, in the subject's fiducial-surface mm + space (see the Notes below), concatenated across all streamlines. + Cast to ``float32``. + offsets : (M,) or (M+1,) array_like + Index into `points` where each streamline starts. Following the TRX + convention, this should have ``M + 1`` entries (`M` = number of + streamlines) with a trailing sentinel equal to `N` (the number of + points); if exactly `M` entries are given (no sentinel), the + sentinel is appended automatically. Must be non-decreasing and start + at 0. + subject : str + Subject identifier. Must exist in the pycortex database. + dpv : dict[str, array_like], optional + Data-per-vertex: for each name, an array with one entry (or row) per + point, aligned with `points`. + dps : dict[str, array_like], optional + Data-per-streamline: for each name, an array with one entry (or row) + per streamline. + groups : dict[str, array_like of int], optional + Named subsets of streamlines, given as arrays of streamline indices. + Groups may overlap, and a streamline may belong to no group. Dict + order is preserved and is significant: it is the order in which + `groups_wire` concatenates them into the wire-format buffer. + color : str or tuple, optional + How to color the streamlines. One of: + + - ``"orientation"`` (default): color by the absolute value of the + local (unit) tangent direction, the standard "directionally + encoded color" tractography convention. + - A ``(r, g, b)`` tuple (0-1 or 0-255 range): a constant color. + - ``"dpv:"`` or ``"dps:"``: color by the named scalar + data-per-vertex or data-per-streamline array, mapped through + `cmap`/`vmin`/`vmax`. + cmap : str, optional + Colormap name used when `color` selects a scalar ``dpv``/``dps`` + array. Defaults to the pycortex default colormap. + vmin : float, optional + Minimum value of the colormap, when `color` selects a scalar array. + Defaults to the ``nanmin`` of that array. + vmax : float, optional + Maximum value of the colormap, when `color` selects a scalar array. + Defaults to the ``nanmax`` of that array. + alpha : float, optional + Opacity of the streamlines, in [0, 1]. Defaults to 1.0. + linewidth : float, optional + Line width in pixels (the WebGL viewer clamps this to 1px on most + platforms). Defaults to 1.0. + description : str, optional + String describing this dataset. Displayed in the webgl viewer. + **kwargs + Additional arguments passed to :class:`~cortex.dataset.views.Dataview` + (e.g. ``state``, ``priority``). + + Notes + ----- + `points` must be expressed in the same mm space as the subject's + fiducial surfaces (FreeSurfer scanner RAS mm; see `cortex.freesurfer`). + Diffusion tractography output is often in a DWI/T1 RASmm space that + already matches this when the DWI was registered to the T1 -- in that + case no transform is needed. Use the `xfm` argument of :meth:`from_trx` + to align data expressed in a different space (e.g. MNI). + """ + + def __init__( + self, + points: npt.ArrayLike, + offsets: npt.ArrayLike, + subject: str, + *, + # Mapping, not Dict: dict is invariant in its value type, so a plain + # dict[str, ndarray] -- what every caller actually has -- would not + # satisfy dict[str, ArrayLike]. + dpv: Optional[Mapping[str, npt.ArrayLike]] = None, + dps: Optional[Mapping[str, npt.ArrayLike]] = None, + groups: Optional[Mapping[str, npt.ArrayLike]] = None, + color: ColorSpec = "orientation", + cmap: Optional[str] = None, + vmin: Optional[float] = None, + vmax: Optional[float] = None, + alpha: float = 1.0, + linewidth: float = 1.0, + description: str = "", + **kwargs: Any, + ) -> None: + points_arr = np.asarray(points, dtype=np.float32) + if points_arr.ndim != 2 or points_arr.shape[1] != 3: + raise ValueError( + "points must have shape (N, 3), got %r" % (points_arr.shape,) + ) + n_points = points_arr.shape[0] + + offsets_arr = np.asarray(offsets) + if offsets_arr.ndim != 1: + raise ValueError( + "offsets must be 1-D, got shape %r" % (offsets_arr.shape,) + ) + offsets_arr = offsets_arr.astype(np.int64) + if offsets_arr.size == 0 or offsets_arr[-1] != n_points: + # M offsets given without the trailing sentinel: append it. + offsets_arr = np.concatenate([offsets_arr, [n_points]]).astype(np.int64) + if offsets_arr.size == 0 or offsets_arr[0] != 0: + raise ValueError("offsets must start at 0") + if np.any(np.diff(offsets_arr) < 0): + raise ValueError("offsets must be monotonically non-decreasing") + + self.points = points_arr + self.offsets = offsets_arr + n_streamlines = offsets_arr.shape[0] - 1 + + self.subject = subject if isinstance(subject, str) else subject.decode("utf-8") + + self.dpv: Dict[str, npt.NDArray] = {} + for name, arr in (dpv or {}).items(): + arr = np.asarray(arr) + if arr.shape[0] != n_points: + raise ValueError( + "dpv[%r] must have %d entries (one per point), got %d" + % (name, n_points, arr.shape[0]) + ) + self.dpv[name] = arr + + self.dps: Dict[str, npt.NDArray] = {} + for name, arr in (dps or {}).items(): + arr = np.asarray(arr) + if arr.shape[0] != n_streamlines: + raise ValueError( + "dps[%r] must have %d entries (one per streamline), got %d" + % (name, n_streamlines, arr.shape[0]) + ) + self.dps[name] = arr + + self.groups: Dict[str, npt.NDArray] = {} + for name, idx in (groups or {}).items(): + idx_arr = np.asarray(idx, dtype=np.int64) + if idx_arr.size and (idx_arr.min() < 0 or idx_arr.max() >= n_streamlines): + raise ValueError( + "groups[%r] contains streamline indices out of range " + "[0, %d)" % (name, n_streamlines) + ) + self.groups[name] = idx_arr + + if isinstance(color, str): + if color != "orientation": + if color.startswith("dpv:"): + key = color[len("dpv:") :] + if key not in self.dpv: + raise ValueError("Unknown dpv field %r for color" % key) + elif color.startswith("dps:"): + key = color[len("dps:") :] + if key not in self.dps: + raise ValueError("Unknown dps field %r for color" % key) + else: + raise ValueError( + "Unrecognized color spec %r: must be 'orientation', " + "'dpv:', 'dps:', or an (r, g, b) tuple" % (color,) + ) + else: + color = tuple(color) + if len(color) != 3: + raise ValueError( + "color tuple must have exactly 3 entries (r, g, b), got %d" + % len(color) + ) + self.color: ColorSpec = color + + if not (0.0 <= alpha <= 1.0): + raise ValueError("alpha must be in [0, 1], got %r" % (alpha,)) + if linewidth <= 0: + raise ValueError("linewidth must be positive, got %r" % (linewidth,)) + self.alpha = float(alpha) + self.linewidth = float(linewidth) + + super().__init__( + cmap=cmap, + vmin=vmin, + vmax=vmax, + description=description, + **kwargs, + ) + + # -- basic geometry ----------------------------------------------- + + @property + def n_points(self) -> int: + """Total number of points across all streamlines.""" + return self.points.shape[0] + + @property + def n_streamlines(self) -> int: + """Number of streamlines.""" + return self.offsets.shape[0] - 1 + + def __len__(self) -> int: + return self.n_streamlines + + @property + def lengths(self) -> npt.NDArray[np.int64]: + """Number of points in each streamline, shape ``(n_streamlines,)``.""" + return np.diff(self.offsets) + + @property + def streamlines(self) -> list: + """List of ``(L_i, 3)`` views into `points`, one per streamline.""" + return [ + self.points[start:stop] + for start, stop in zip(self.offsets[:-1], self.offsets[1:]) + ] + + @property + def name(self) -> str: + """Name of this Tractogram, computed from a hash of `points`. + + Mirrors :attr:`cortex.dataset.braindata.BrainData.name`, so that + code keying on that convention (e.g. `cortex.webgl.data.Package`) + works unchanged for tractograms. + """ + return "__%s" % _hash(self.points)[:16] + + # -- selection / decimation ----------------------------------------- + + def select(self, indices: npt.ArrayLike) -> "Tractogram": + """Return a new `Tractogram` containing only the given streamlines. + + Parameters + ---------- + indices : array_like of int or bool + Streamline indices to keep (or a boolean mask of length + `n_streamlines`). Order is preserved; repeats are allowed. + + Returns + ------- + Tractogram + A new tractogram with `points`/`offsets`/`dpv`/`dps` sliced and + `groups` remapped to the new streamline indexing. Display + attributes (`color`, `cmap`, `vmin`, `vmax`, `alpha`, + `linewidth`, `description`) are copied. + """ + idx = np.asarray(indices) + if idx.dtype == np.bool_: + idx = np.nonzero(idx)[0] + idx = idx.astype(np.int64) + + starts = self.offsets[:-1] + stops = self.offsets[1:] + + point_chunks = [] + dpv_chunks: Dict[str, list] = {name: [] for name in self.dpv} + new_lengths = [] + for i in idx: + s, e = int(starts[i]), int(stops[i]) + point_chunks.append(self.points[s:e]) + for name in self.dpv: + dpv_chunks[name].append(self.dpv[name][s:e]) + new_lengths.append(e - s) + + if point_chunks: + new_points = np.concatenate(point_chunks, axis=0) + else: + new_points = np.zeros((0, 3), dtype=np.float32) + new_offsets = np.concatenate([[0], np.cumsum(new_lengths)]).astype(np.int64) + + new_dpv = {} + for name, chunks in dpv_chunks.items(): + if chunks: + new_dpv[name] = np.concatenate(chunks, axis=0) + else: + new_dpv[name] = np.zeros( + (0,) + self.dpv[name].shape[1:], dtype=self.dpv[name].dtype + ) + + new_dps = {name: self.dps[name][idx] for name in self.dps} + + # A group keeps every new streamline whose old index was a member + # (repeated indices in `idx` therefore stay in the group). + new_groups = {} + for name, members in self.groups.items(): + new_groups[name] = np.nonzero(np.isin(idx, members))[0].astype(np.int64) + + return Tractogram( + new_points, + new_offsets, + self.subject, + dpv=new_dpv, + dps=new_dps, + groups=new_groups, + color=self.color, + cmap=self.cmap, + vmin=self.vmin, + vmax=self.vmax, + alpha=self.alpha, + linewidth=self.linewidth, + description=self.description, + state=self.state, + **self.attrs, + ) + + def get_group(self, name: str) -> "Tractogram": + """Return a new `Tractogram` restricted to the named group. + + Parameters + ---------- + name : str + A key of `groups`. + """ + if name not in self.groups: + raise KeyError("Unknown group %r" % name) + return self.select(self.groups[name]) + + def subsample( + self, + max_streamlines: Optional[int] = None, + step: Optional[int] = None, + seed: int = 0, + ) -> "Tractogram": + """Return a decimated `Tractogram`, for faster interactive display. + + Parameters + ---------- + max_streamlines : int, optional + If given (and `step` is not), and there are more than this many + streamlines, a random subset of this size is kept (order + preserved), chosen with the given `seed`. + step : int, optional + If given, keep every `step`-th streamline (``indices[::step]``); + takes precedence over `max_streamlines`. + seed : int, optional + Seed for the random subset chosen when `max_streamlines` is used. + + Returns + ------- + Tractogram + """ + n = self.n_streamlines + if step is not None: + indices = np.arange(0, n, step) + elif max_streamlines is not None and n > max_streamlines: + rng = np.random.default_rng(seed) + indices = np.sort(rng.choice(n, size=max_streamlines, replace=False)) + else: + indices = np.arange(n) + return self.select(indices) + + # -- coloring --------------------------------------------------------- + + def vertex_colors(self) -> npt.NDArray[np.uint8]: + """Compute per-point RGB colors according to `color`. + + Returns + ------- + (N, 3) ndarray of uint8 + One RGB color per point in `points`. + """ + color = self.color + if isinstance(color, str): + if color == "orientation": + return self._orientation_colors() + if color.startswith("dpv:"): + name = color[len("dpv:") :] + scalar = np.asarray(self.dpv[name], dtype=np.float64).reshape( + self.n_points, -1 + )[:, 0] + return self._scalar_colors(scalar) + if color.startswith("dps:"): + name = color[len("dps:") :] + per_streamline = np.asarray(self.dps[name], dtype=np.float64).reshape( + self.n_streamlines, -1 + )[:, 0] + return self._scalar_colors(self._broadcast_dps(per_streamline)) + raise ValueError("Unrecognized color spec %r" % (color,)) + return self._constant_colors(color) + + def _orientation_colors(self) -> npt.NDArray[np.uint8]: + # Vectorized over all points: central differences everywhere, then + # the first/last point of every streamline is fixed up with a + # one-sided difference (single-point streamlines get a zero tangent). + pts = self.points.astype(np.float64) + tangent = np.zeros_like(pts) + if pts.shape[0] >= 3: + tangent[1:-1] = pts[2:] - pts[:-2] + lengths = self.lengths + nonempty = lengths > 0 + starts = self.offsets[:-1][nonempty] + ends = self.offsets[1:][nonempty] - 1 + multi = lengths[nonempty] > 1 + first, last = starts[multi], ends[multi] + tangent[first] = pts[first + 1] - pts[first] + tangent[last] = pts[last] - pts[last - 1] + tangent[starts[~multi]] = 0.0 + norm = np.linalg.norm(tangent, axis=1, keepdims=True) + unit = np.divide(tangent, norm, out=np.zeros_like(tangent), where=norm > 0) + return np.clip(np.abs(unit) * 255.0, 0, 255).astype(np.uint8) + + def _constant_colors(self, color: ColorSpec) -> npt.NDArray[np.uint8]: + rgb = np.asarray(color, dtype=np.float64) + if rgb.max() <= 1.0: + rgb = rgb * 255.0 + rgb = np.clip(rgb, 0, 255).astype(np.uint8) + return np.tile(rgb, (self.n_points, 1)) + + def _broadcast_dps(self, values: npt.NDArray) -> npt.NDArray[np.float64]: + return np.repeat(np.asarray(values, dtype=np.float64), self.lengths) + + def _scalar_colors(self, scalar: npt.NDArray[np.float64]) -> npt.NDArray[np.uint8]: + from matplotlib import cm, colors as mcolors + + nan_mask = np.isnan(scalar) + finite = scalar[~nan_mask] + vmin = self.vmin if self.vmin is not None else ( + float(np.nanmin(finite)) if finite.size else 0.0 + ) + vmax = self.vmax if self.vmax is not None else ( + float(np.nanmax(finite)) if finite.size else 1.0 + ) + cmap = self.get_cmapdict()["cmap"] + norm = mcolors.Normalize(vmin=vmin, vmax=vmax) + mappable = cm.ScalarMappable(norm=norm, cmap=cmap) + rgba = mappable.to_rgba(np.nan_to_num(scalar)) + rgb = (np.clip(rgba[:, :3], 0, 1) * 255.0).astype(np.uint8) + rgb[nan_mask] = 128 + return rgb + + # -- serialization ------------------------------------------------ + + def groups_wire(self) -> "tuple[npt.NDArray[np.uint32], Dict[str, tuple]]": + """Concatenate `groups` into one buffer, for the wire format. + + Returns + ------- + indices : (K,) ndarray of uint32 + Concatenation of every entry of `groups`, in `groups` dict + order, as streamline indices. Empty when `groups` is empty. + slices : dict[str, tuple[int, int]] + For each group name (in the same order as `indices`), the + ``(start, stop)`` slice bounds into `indices` holding that + group's streamline indices (``stop - start`` equals + ``len(groups[name])``). + """ + chunks = [] + slices: Dict[str, tuple] = {} + pos = 0 + for name, idx in self.groups.items(): + idx_arr = np.asarray(idx, dtype=np.uint32) + chunks.append(idx_arr) + slices[name] = (pos, pos + idx_arr.shape[0]) + pos += idx_arr.shape[0] + if chunks: + indices = np.concatenate(chunks).astype(np.uint32) + else: + indices = np.zeros((0,), dtype=np.uint32) + return indices, slices + + # The base class narrows this to DataviewJSON, a TypedDict describing a + # colormapped brain-data payload. A tractogram's metadata shares almost + # none of those keys (see below) and TypedDicts admit no extra ones, so + # the override is deliberate rather than an oversight. + def to_json(self, simple: bool = False) -> dict: # type: ignore[override] + """Return the wire-format metadata dict for this tractogram. + + Parameters + ---------- + simple : bool, optional + If True, omit the colormap-related keys (`cmap`/`vmin`/`vmax`) + even when `color` selects a scalar `dpv`/`dps` array. + + Notes + ----- + ``result["groups"]`` maps each group name to a ``[start, stop]`` + slice (not a count): these are bounds into the fourth per-tractogram + wire buffer, ``groups`` (little-endian uint32 streamline indices, + the concatenation of every group in dict order -- see + :meth:`groups_wire`, which `cortex.webgl.data.Package` uses to build + that buffer so the two stay in sync). + """ + if isinstance(self.color, str): + color_repr = self.color + else: + color_repr = str(tuple(float(c) for c in self.color)) + + _, group_slices = self.groups_wire() + result: dict = { + "subject": self.subject, + "n_points": int(self.n_points), + "n_streamlines": int(self.n_streamlines), + "alpha": self.alpha, + "linewidth": self.linewidth, + "visible": True, + "color": color_repr, + "groups": { + name: [int(start), int(stop)] + for name, (start, stop) in group_slices.items() + }, + "description": self.description, + } + + is_scalar_mode = isinstance(self.color, str) and ( + self.color.startswith("dpv:") or self.color.startswith("dps:") + ) + if not simple and is_scalar_mode: + result["cmap"] = self.cmap + result["vmin"] = self.vmin + result["vmax"] = self.vmax + return result + + def _write_hdf(self, h5, name: str = "data", data=None, xfmname=None): + raise NotImplementedError( + "HDF5 persistence of Tractogram is not implemented yet" + ) + + # -- construction from TRX / raw streamlines -------------------------- + + @classmethod + def from_trx( + cls, + path_or_trxfile: Union[str, Path, "TrxFile"], + subject: str, + xfm: Optional[npt.ArrayLike] = None, + **display: Any, + ) -> "Tractogram": + """Build a `Tractogram` from a TRX file (via `trx-python`). + + Parameters + ---------- + path_or_trxfile : str, Path, or trx.trx_file_memmap.TrxFile + Path to a ``.trx`` file, or an already-loaded `TrxFile`. When a + path is given, `trx.io.load` is used to load it, and the file is + closed before returning. + subject : str + Subject identifier. Must exist in the pycortex database. + xfm : (4, 4) array_like, optional + Affine applied to the streamline points (as homogeneous + coordinates) after loading, to align data expressed in a + different RASmm space onto the subject's fiducial surfaces. + **display + Additional display keyword arguments forwarded to the + `Tractogram` constructor (`color`, `cmap`, `vmin`, `vmax`, + `alpha`, `linewidth`, `description`, ...). + + Returns + ------- + Tractogram + """ + close_after = False + if isinstance(path_or_trxfile, (str, Path)): + from trx.io import load as trx_load + + trx_file = trx_load(str(path_or_trxfile)) + close_after = True + else: + trx_file = path_or_trxfile + + try: + streamlines = trx_file.streamlines + points = np.array(streamlines._data, dtype=np.float32, copy=True) + n_points = points.shape[0] + raw_offsets = np.array(streamlines._offsets, copy=True).astype(np.int64) + if raw_offsets.size == 0 or raw_offsets[-1] != n_points: + offsets = np.concatenate([raw_offsets, [n_points]]).astype(np.int64) + else: + offsets = raw_offsets + + def _extract(value: Any) -> npt.NDArray: + # trx-python memmaps dpv/dps as (N, k); a single scalar per + # entry comes back as (N, 1), which we flatten to (N,). + data = np.array(getattr(value, "_data", value), copy=True) + if data.ndim == 2 and data.shape[1] == 1: + data = data[:, 0] + return data + + dpv = { + name: _extract(arr) for name, arr in trx_file.data_per_vertex.items() + } + dps = { + name: _extract(arr) + for name, arr in trx_file.data_per_streamline.items() + } + groups = { + name: np.array(idx, dtype=np.int64, copy=True) + for name, idx in trx_file.groups.items() + } + + if xfm is not None: + xfm_arr = np.asarray(xfm, dtype=np.float64) + homogeneous = np.concatenate( + [points.astype(np.float64), np.ones((n_points, 1))], axis=1 + ) + points = (homogeneous @ xfm_arr.T)[:, :3].astype(np.float32) + + return cls( + points, offsets, subject, dpv=dpv, dps=dps, groups=groups, **display + ) + finally: + if close_after: + trx_file.close() + + @classmethod + def from_streamlines( + cls, + streamlines: Sequence[npt.ArrayLike], + subject: str, + **kwargs: Any, + ) -> "Tractogram": + """Build a `Tractogram` from a list of ``(L_i, 3)`` point arrays. + + Parameters + ---------- + streamlines : sequence of array_like + Each element is an ``(L_i, 3)`` array of points for one + streamline. + subject : str + Subject identifier. Must exist in the pycortex database. + **kwargs + Additional keyword arguments forwarded to the `Tractogram` + constructor. + """ + arrays = [np.asarray(s, dtype=np.float32) for s in streamlines] + if not arrays: + points = np.zeros((0, 3), dtype=np.float32) + offsets = np.array([0], dtype=np.int64) + else: + lengths = [a.shape[0] for a in arrays] + points = np.concatenate(arrays, axis=0).astype(np.float32) + offsets = np.concatenate([[0], np.cumsum(lengths)]).astype(np.int64) + return cls(points, offsets, subject, **kwargs) diff --git a/cortex/tests/test_tractogram.py b/cortex/tests/test_tractogram.py new file mode 100644 index 000000000..d59649679 --- /dev/null +++ b/cortex/tests/test_tractogram.py @@ -0,0 +1,487 @@ +"""Tests for cortex.dataset.tractogram.Tractogram.""" + +import zipfile +import json + +import numpy as np +import pytest + +import cortex +from cortex import db, dataset +from cortex.dataset import Tractogram + +subj = "S1" + + +def _make_streamlines(n_streamlines=20, n_points=30, seed=0): + """Build synthetic streamlines derived from S1's fiducial surface. + + Each streamline is a straight line between two random fiducial + vertices, with a sinusoidal bulge added orthogonal to the line, so + that the local tangent direction varies along the streamline (needed + to exercise the "orientation" coloring mode). + """ + rng = np.random.default_rng(seed) + pts, _ = db.get_surf(subj, "fiducial", merge=True) + n_verts = pts.shape[0] + + streamlines = [] + for _ in range(n_streamlines): + i0, i1 = rng.choice(n_verts, size=2, replace=False) + p0, p1 = pts[i0], pts[i1] + t = np.linspace(0, 1, n_points)[:, None] + line = p0[None, :] * (1 - t) + p1[None, :] * t + + # orthogonal bulge so tangents aren't all parallel + direction = p1 - p0 + norm = np.linalg.norm(direction) + if norm > 0: + direction = direction / norm + ortho = np.cross(direction, np.array([0.0, 0.0, 1.0])) + if np.linalg.norm(ortho) < 1e-6: + ortho = np.cross(direction, np.array([0.0, 1.0, 0.0])) + ortho = ortho / (np.linalg.norm(ortho) + 1e-12) + bulge = 5.0 * np.sin(np.linspace(0, np.pi, n_points)) + line = line + bulge[:, None] * ortho[None, :] + + streamlines.append(line.astype(np.float32)) + return streamlines + + +def _make_tractogram(n_streamlines=20, n_points=30, seed=0, groups=None, **kwargs): + streamlines = _make_streamlines(n_streamlines, n_points, seed=seed) + n_total = sum(s.shape[0] for s in streamlines) + + rng = np.random.default_rng(seed + 1) + dpv = {"scalar": rng.standard_normal(n_total).astype(np.float32)} + dps = {"length": np.array([s.shape[0] for s in streamlines], dtype=np.float32)} + if groups is None: + groups = { + "even": np.arange(0, n_streamlines, 2), + "odd": np.arange(1, n_streamlines, 2), + } + return Tractogram.from_streamlines( + streamlines, subj, dpv=dpv, dps=dps, groups=groups, **kwargs + ) + + +def _overlapping_groups(n_streamlines): + """Two overlapping groups that leave the last streamline ungrouped. + + Used by tests (and by `test_webgl_tractogram.py`) that need to exercise + the "streamline belongs to >=1 group" / "streamline belongs to no + group" cases together, which the disjoint even/odd split above does not. + """ + half = n_streamlines // 2 + return { + "first_half": np.arange(0, half + 1), + "second_half": np.arange(half - 1, n_streamlines - 1), + } + + +# --------------------------------------------------------------------------- +# Construction / validation +# --------------------------------------------------------------------------- + + +def test_construction_basic(): + tract = _make_tractogram() + assert tract.n_streamlines == 20 + assert len(tract) == 20 + assert tract.n_points == tract.points.shape[0] + assert tract.offsets.shape == (21,) + assert tract.offsets[0] == 0 + assert tract.offsets[-1] == tract.n_points + assert tract.subject == subj + + +def test_offsets_sentinel_appended(): + points = np.random.randn(10, 3).astype(np.float32) + offsets_no_sentinel = np.array([0, 5]) # M=2 offsets, no trailing N + tract = Tractogram(points, offsets_no_sentinel, subj) + assert tract.offsets.shape == (3,) + assert tract.offsets[-1] == 10 + assert tract.n_streamlines == 2 + + +def test_points_wrong_shape_raises(): + with pytest.raises(ValueError): + Tractogram(np.zeros((10, 2)), np.array([0, 10]), subj) + + +def test_offsets_not_monotonic_raises(): + points = np.zeros((10, 3), dtype=np.float32) + with pytest.raises(ValueError): + Tractogram(points, np.array([0, 8, 5, 10]), subj) + + +def test_offsets_must_start_at_zero(): + points = np.zeros((10, 3), dtype=np.float32) + with pytest.raises(ValueError): + Tractogram(points, np.array([1, 10]), subj) + + +def test_dpv_wrong_size_raises(): + points = np.zeros((10, 3), dtype=np.float32) + offsets = np.array([0, 10]) + with pytest.raises(ValueError): + Tractogram(points, offsets, subj, dpv={"bad": np.zeros(5)}) + + +def test_dps_wrong_size_raises(): + points = np.zeros((10, 3), dtype=np.float32) + offsets = np.array([0, 5, 10]) + with pytest.raises(ValueError): + Tractogram(points, offsets, subj, dps={"bad": np.zeros(1)}) + + +def test_group_out_of_range_raises(): + points = np.zeros((10, 3), dtype=np.float32) + offsets = np.array([0, 5, 10]) + with pytest.raises(ValueError): + Tractogram(points, offsets, subj, groups={"bad": np.array([5])}) + + +def test_unrecognized_color_raises(): + points = np.zeros((10, 3), dtype=np.float32) + offsets = np.array([0, 10]) + with pytest.raises(ValueError): + Tractogram(points, offsets, subj, color="not-a-valid-spec") + + +def test_unknown_dpv_dps_color_raises(): + points = np.zeros((10, 3), dtype=np.float32) + offsets = np.array([0, 10]) + with pytest.raises(ValueError): + Tractogram(points, offsets, subj, color="dpv:missing") + with pytest.raises(ValueError): + Tractogram(points, offsets, subj, color="dps:missing") + + +# --------------------------------------------------------------------------- +# lengths / offsets / streamlines invariants +# --------------------------------------------------------------------------- + + +def test_lengths_and_streamlines(): + tract = _make_tractogram(n_streamlines=5, n_points=7) + assert np.array_equal(tract.lengths, np.diff(tract.offsets)) + assert tract.lengths.sum() == tract.n_points + streamlines = tract.streamlines + assert len(streamlines) == tract.n_streamlines + for sl, length in zip(streamlines, tract.lengths): + assert sl.shape == (length, 3) + + +def test_name_is_deterministic_hash(): + tract1 = _make_tractogram(seed=1) + tract2 = Tractogram(tract1.points.copy(), tract1.offsets.copy(), subj) + assert tract1.name == tract2.name + tract3 = _make_tractogram(seed=2) + assert tract1.name != tract3.name + + +# --------------------------------------------------------------------------- +# select / get_group / subsample +# --------------------------------------------------------------------------- + + +def _assert_offsets_valid(tract): + assert tract.offsets[0] == 0 + assert tract.offsets[-1] == tract.n_points + assert np.all(np.diff(tract.offsets) >= 0) + + +def test_select_basic(): + tract = _make_tractogram(n_streamlines=10, n_points=5) + sub = tract.select([0, 2, 4]) + assert sub.n_streamlines == 3 + _assert_offsets_valid(sub) + assert sub.dpv["scalar"].shape[0] == sub.n_points + assert sub.dps["length"].shape[0] == sub.n_streamlines + np.testing.assert_array_equal(sub.streamlines[0], tract.streamlines[0]) + np.testing.assert_array_equal(sub.streamlines[1], tract.streamlines[2]) + + +def test_select_remaps_groups(): + tract = _make_tractogram(n_streamlines=10, n_points=5) + sub = tract.select([1, 3, 5, 7, 9]) # exactly the "odd" group + assert set(sub.groups["odd"].tolist()) == {0, 1, 2, 3, 4} + assert sub.groups["even"].size == 0 + + +def test_get_group(): + tract = _make_tractogram(n_streamlines=10, n_points=5) + even = tract.get_group("even") + assert even.n_streamlines == 5 + _assert_offsets_valid(even) + with pytest.raises(KeyError): + tract.get_group("nope") + + +def test_subsample_max_streamlines(): + tract = _make_tractogram(n_streamlines=20, n_points=5) + sub = tract.subsample(max_streamlines=5, seed=42) + assert sub.n_streamlines == 5 + _assert_offsets_valid(sub) + # deterministic given seed + sub2 = tract.subsample(max_streamlines=5, seed=42) + np.testing.assert_array_equal(sub.points, sub2.points) + + +def test_subsample_step(): + tract = _make_tractogram(n_streamlines=20, n_points=5) + sub = tract.subsample(step=4) + assert sub.n_streamlines == 5 + _assert_offsets_valid(sub) + + +def test_subsample_noop_when_small(): + tract = _make_tractogram(n_streamlines=5, n_points=5) + sub = tract.subsample(max_streamlines=100) + assert sub.n_streamlines == tract.n_streamlines + + +# --------------------------------------------------------------------------- +# vertex_colors +# --------------------------------------------------------------------------- + + +def test_vertex_colors_orientation(): + tract = _make_tractogram(color="orientation") + colors = tract.vertex_colors() + assert colors.shape == (tract.n_points, 3) + assert colors.dtype == np.uint8 + assert colors.min() >= 0 and colors.max() <= 255 + + # unit-tangent property: recompute manually for the first streamline + seg = tract.streamlines[0].astype(np.float64) + tangent = np.empty_like(seg) + tangent[1:-1] = seg[2:] - seg[:-2] + tangent[0] = seg[1] - seg[0] + tangent[-1] = seg[-1] - seg[-2] + norm = np.linalg.norm(tangent, axis=1, keepdims=True) + expected = np.clip(np.abs(tangent / norm) * 255, 0, 255).astype(np.uint8) + start, stop = tract.offsets[0], tract.offsets[1] + np.testing.assert_allclose(colors[start:stop], expected, atol=1) + + +def test_vertex_colors_constant(): + tract = _make_tractogram(color=(1.0, 0.0, 0.5)) + colors = tract.vertex_colors() + assert colors.shape == (tract.n_points, 3) + expected = np.array([255, 0, 127], dtype=np.uint8) + assert np.all(colors[0] == expected) + assert np.all(colors == colors[0]) # constant everywhere + + tract255 = _make_tractogram(color=(255, 0, 128)) + colors255 = tract255.vertex_colors() + assert np.all(colors255[0] == np.array([255, 0, 128], dtype=np.uint8)) + + +def test_vertex_colors_dpv(): + tract = _make_tractogram(color="dpv:scalar", cmap="gray") + colors = tract.vertex_colors() + assert colors.shape == (tract.n_points, 3) + assert colors.dtype == np.uint8 + scalar = tract.dpv["scalar"] + order = np.argsort(scalar) + # a gray colormap should be monotonic with the scalar (all channels equal) + gray_channel = colors[:, 0].astype(np.float64) + assert np.all(colors[:, 0] == colors[:, 1]) + assert np.all(colors[:, 1] == colors[:, 2]) + sorted_gray = gray_channel[order] + assert np.all(np.diff(sorted_gray) >= -1) # monotone non-decreasing (allow tie noise) + + +def test_vertex_colors_dpv_nan_is_gray(): + tract = _make_tractogram(color="dpv:scalar", cmap="gray") + tract.dpv["scalar"][0] = np.nan + colors = tract.vertex_colors() + assert tuple(colors[0]) == (128, 128, 128) + + +def test_vertex_colors_dps(): + tract = _make_tractogram(color="dps:length", cmap="gray") + colors = tract.vertex_colors() + assert colors.shape == (tract.n_points, 3) + # all points within a streamline share the same dps-derived color + start, stop = tract.offsets[0], tract.offsets[1] + assert np.all(colors[start:stop] == colors[start]) + + +# --------------------------------------------------------------------------- +# to_json / _write_hdf / normalize +# --------------------------------------------------------------------------- + + +def test_to_json_keys(): + tract = _make_tractogram(color="dpv:scalar") + j = tract.to_json() + expected_keys = { + "subject", + "n_points", + "n_streamlines", + "alpha", + "linewidth", + "visible", + "color", + "groups", + "description", + "cmap", + "vmin", + "vmax", + } + assert expected_keys <= set(j.keys()) + assert j["subject"] == subj + assert j["n_points"] == tract.n_points + assert j["n_streamlines"] == tract.n_streamlines + assert j["visible"] is True + assert j["groups"] == {"even": [0, 10], "odd": [10, 20]} + # must be JSON serializable + json.dumps(j) + + +def test_groups_wire_disjoint(): + tract = _make_tractogram(n_streamlines=10, n_points=5) + indices, slices = tract.groups_wire() + assert indices.dtype == np.uint32 + assert indices.shape == (10,) + assert slices == {"even": (0, 5), "odd": (5, 10)} + for name, (start, stop) in slices.items(): + assert stop - start == len(tract.groups[name]) + np.testing.assert_array_equal(indices[start:stop], tract.groups[name]) + # to_json must report the very same slice bounds. + j = tract.to_json() + assert j["groups"] == {name: list(s) for name, s in slices.items()} + + +def test_groups_wire_overlapping_and_ungrouped(): + n_streamlines = 8 + groups = _overlapping_groups(n_streamlines) + tract = _make_tractogram( + n_streamlines=n_streamlines, n_points=5, groups=groups + ) + indices, slices = tract.groups_wire() + assert indices.shape[0] == sum(len(g) for g in groups.values()) + # slices are contiguous and in the same order as `groups` + assert list(slices.keys()) == list(groups.keys()) + pos = 0 + for name, idx in groups.items(): + start, stop = slices[name] + assert (start, stop) == (pos, pos + len(idx)) + np.testing.assert_array_equal(indices[start:stop], idx) + pos += len(idx) + # groups overlap: streamline `half - 1` .. `half` appear in both + first, second = slices["first_half"], slices["second_half"] + first_members = set(indices[first[0] : first[1]].tolist()) + second_members = set(indices[second[0] : second[1]].tolist()) + assert first_members & second_members + # the last streamline belongs to no group + all_members = first_members | second_members + assert (n_streamlines - 1) not in all_members + + +def test_groups_wire_empty(): + tract = _make_tractogram(n_streamlines=5, n_points=5, groups={}) + indices, slices = tract.groups_wire() + assert indices.shape == (0,) + assert indices.dtype == np.uint32 + assert slices == {} + assert tract.to_json()["groups"] == {} + + +def test_to_json_simple_omits_cmap(): + tract = _make_tractogram(color="dpv:scalar") + j = tract.to_json(simple=True) + assert "cmap" not in j + assert "vmin" not in j + + +def test_to_json_non_scalar_color_no_cmap(): + tract = _make_tractogram(color="orientation") + j = tract.to_json() + assert "cmap" not in j + + +def test_write_hdf_raises(): + tract = _make_tractogram() + with pytest.raises(NotImplementedError): + tract._write_hdf(None) + + +def test_normalize_passthrough(): + tract = _make_tractogram() + assert dataset.normalize(tract) is tract + assert cortex.dataset.normalize(tract) is tract + + +# --------------------------------------------------------------------------- +# from_trx (requires trx-python) +# --------------------------------------------------------------------------- + + +def _write_trx_fixture(path, points, offsets, dpv=None, dps=None, groups=None): + """Write a minimal TRX zip archive by hand (no trx-python dependency).""" + n_points = points.shape[0] + n_streamlines = offsets.shape[0] - 1 + header = { + "VOXEL_TO_RASMM": np.eye(4).tolist(), + "DIMENSIONS": [10, 10, 10], + "NB_VERTICES": int(n_points), + "NB_STREAMLINES": int(n_streamlines), + } + with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_STORED) as zf: + zf.writestr("header.json", json.dumps(header)) + zf.writestr( + "positions.3.float32", points.astype(""`` / ``"dps:"``: color by a named data-per-vertex or data-per-streamline scalar array (from `dpv`/`dps`), mapped through `cmap`/`vmin`/`vmax` like any other Dataview. + +``groups`` holds named subsets of streamlines (as arrays of streamline indices, e.g. bundle names); groups may overlap, and a streamline may belong to no group. :meth:`Tractogram.get_group` and :meth:`Tractogram.select` return new tractograms restricted to a subset, and :meth:`Tractogram.subsample` decimates a large tractogram for faster interactive display. + +.. note:: + `points` must be expressed in the same mm space as the subject's fiducial surfaces (FreeSurfer scanner RAS mm). Tractography output is often already in this space when the diffusion data was registered to the same T1 used to generate the surfaces; otherwise, pass an `xfm` (4x4 affine) to :meth:`Tractogram.from_trx` to align the streamline positions onto the fiducial surface. + +HDF5 persistence of :class:`Tractogram` (saving to / loading from a :class:`Dataset`) is not implemented yet. + Dataset ------- Dataset objects hold a dictionary of DataView objects. Its main function is saving out data into a standardized HDF format. Datasets also allow subject data packing, if you desire to send data to someone without the subject database. diff --git a/pyproject.toml b/pyproject.toml index 64c6019ef..5d2958257 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ dev = [ ] test = [ "pycortex[headless]", + "pycortex[tractography]", # The visual-regression tests compare against renders from # these builds. Both are unpinned for users. # See cortex/tests/reference_images/README.md @@ -37,6 +38,9 @@ types = [ headless = [ "playwright>=1.58.0" ] +tractography = [ + "trx-python" +] [tool.setuptools_scm] version_file = "cortex/_version.py" @@ -62,3 +66,9 @@ module = [ "tvtk.api", ] ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = [ + "trx.*", +] +ignore_missing_imports = true From f572bc9f41a92ca63733e896763c95a742a36f0b Mon Sep 17 00:00:00 2001 From: Matteo Visconti di Oleggio Castello Date: Sat, 19 Sep 2026 18:40:44 -0700 Subject: [PATCH 2/3] TEST make the overlapping-groups fixture discriminating Its two groups now go in in reverse alphabetical order. The viewer lists bundles alphabetically rather than in wire order, and with the two agreeing no test could tell the difference. Co-Authored-By: Claude Opus 5 --- cortex/tests/test_tractogram.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cortex/tests/test_tractogram.py b/cortex/tests/test_tractogram.py index d59649679..b9d140a7a 100644 --- a/cortex/tests/test_tractogram.py +++ b/cortex/tests/test_tractogram.py @@ -71,11 +71,15 @@ def _overlapping_groups(n_streamlines): Used by tests (and by `test_webgl_tractogram.py`) that need to exercise the "streamline belongs to >=1 group" / "streamline belongs to no group" cases together, which the disjoint even/odd split above does not. + + Insertion order is deliberately the reverse of alphabetical: the viewer + lists bundles alphabetically rather than in wire order, and a fixture + where the two agree cannot tell the difference. """ half = n_streamlines // 2 return { - "first_half": np.arange(0, half + 1), "second_half": np.arange(half - 1, n_streamlines - 1), + "first_half": np.arange(0, half + 1), } From 8f793d06d0e7ecd117817ea8efa548a1d278a3b1 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 20 Sep 2026 01:58:59 +0000 Subject: [PATCH 3/3] FIX address review findings on the Tractogram dataview Seven fixes from the review of the Tractogram container, each with a regression test: - `uniques()` was missing. `Dataset.uniques()` calls it on every view it holds, so `Dataset(tract=tract).uniques()` raised `AttributeError` and a dataset carrying a tractogram could not be packaged or saved at all. A tractogram wraps no BrainData, so it yields nothing. - `name` hashed `points` alone, giving the same identity to two tractograms whose point buffers match but whose streamline boundaries differ. `offsets` now goes into the digest too. - `to_json()` dropped the base Dataview keys `state`, `attrs` and `desc`, losing custom state on the way to the viewer and breaking any consumer that reads them off a view's metadata. - Group index arrays were only range-checked, so a 2-D array was accepted and reached `select()` and `groups_wire()`, whose `(K,)` buffer contract it breaks. Groups must now be 1-D, and a boolean mask is converted rather than silently reinterpreted as the indices 0/1. - `select()` accepted a boolean mask of any length: `select([True])` on a larger tractogram quietly kept only streamline 0. Masks must now have one entry per streamline, indices must be 1-D, and an out-of-range index raises instead of wrapping (a negative index sliced `points` from the end yet never matched when remapping `groups`). - `vertex_colors()` raised on an empty tractogram in the `dpv:`/`dps:` colour modes, where `reshape(0, -1)` cannot be inferred. Empty tractograms are legitimate -- an empty named group, or a `select` that matched nothing -- and now return the documented `(0, 3)` array. - `from_trx(xfm=...)` did not check the affine's shape, failing later inside the matmul. Also restore the Python 3.10 row in the test matrix. The incompatibility is confined to trx-python, an optional test dependency, so gating it with `python_version >= '3.11'` in the test group keeps 3.10 installable (the TRX tests skip themselves there) rather than dropping core 3.10 coverage for a version the wheel-install workflow still exercises. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01UWULoz8DLUFHHio66eTCgh --- .github/workflows/run_tests.yml | 5 +- cortex/dataset/tractogram.py | 104 +++++++++++++++++++++++++---- cortex/tests/test_tractogram.py | 114 ++++++++++++++++++++++++++++++++ pyproject.toml | 6 +- 4 files changed, 212 insertions(+), 17 deletions(-) diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index e12327e2c..60866b97c 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -17,10 +17,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - # 3.11 is the floor because the test group pulls in trx-python (via - # the `tractography` extra), which requires it. On 3.10 the install - # fails outright rather than the TRX tests quietly skipping. - python-version: ["3.11", "3.12", "3.13", "3.14"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] max-parallel: 5 steps: diff --git a/cortex/dataset/tractogram.py b/cortex/dataset/tractogram.py index 41c7a7124..414a27bc6 100644 --- a/cortex/dataset/tractogram.py +++ b/cortex/dataset/tractogram.py @@ -8,6 +8,7 @@ from __future__ import annotations +import hashlib from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Sequence, Union @@ -55,11 +56,13 @@ class Tractogram(Dataview): dps : dict[str, array_like], optional Data-per-streamline: for each name, an array with one entry (or row) per streamline. - groups : dict[str, array_like of int], optional - Named subsets of streamlines, given as arrays of streamline indices. - Groups may overlap, and a streamline may belong to no group. Dict - order is preserved and is significant: it is the order in which - `groups_wire` concatenates them into the wire-format buffer. + groups : dict[str, array_like of int or bool], optional + Named subsets of streamlines, given as 1-D arrays of streamline + indices, or as boolean masks of length `n_streamlines` (which are + converted to indices). Groups may overlap, and a streamline may + belong to no group. Dict order is preserved and is significant: it + is the order in which `groups_wire` concatenates them into the + wire-format buffer. color : str or tuple, optional How to color the streamlines. One of: @@ -170,7 +173,24 @@ def __init__( self.groups: Dict[str, npt.NDArray] = {} for name, idx in (groups or {}).items(): - idx_arr = np.asarray(idx, dtype=np.int64) + idx_arr = np.asarray(idx) + if idx_arr.ndim != 1: + raise ValueError( + "groups[%r] must be 1-D (streamline indices or a boolean " + "mask), got shape %r" % (name, idx_arr.shape) + ) + if idx_arr.dtype == np.bool_: + # A boolean mask, spelled as `select` also accepts it. Without + # this branch `astype(int64)` would silently turn it into the + # indices 0/1 instead. + if idx_arr.size != n_streamlines: + raise ValueError( + "groups[%r] is a boolean mask with %d entries, but " + "there are %d streamlines" + % (name, idx_arr.size, n_streamlines) + ) + idx_arr = np.nonzero(idx_arr)[0] + idx_arr = idx_arr.astype(np.int64) if idx_arr.size and (idx_arr.min() < 0 or idx_arr.max() >= n_streamlines): raise ValueError( "groups[%r] contains streamline indices out of range " @@ -247,13 +267,28 @@ def streamlines(self) -> list: @property def name(self) -> str: - """Name of this Tractogram, computed from a hash of `points`. + """Name of this Tractogram, computed from a hash of its geometry. Mirrors :attr:`cortex.dataset.braindata.BrainData.name`, so that code keying on that convention (e.g. `cortex.webgl.data.Package`) - works unchanged for tractograms. + works unchanged for tractograms. Both `points` and `offsets` go + into the digest: the same point buffer cut into different + streamlines is a different tractogram, and hashing `points` alone + would hand the two the same identity. """ - return "__%s" % _hash(self.points)[:16] + geometry = (_hash(self.points) + _hash(self.offsets)).encode("ascii") + return "__%s" % hashlib.sha1(geometry).hexdigest()[:16] + + def uniques(self, collapse: bool = False) -> "list": + """The `BrainData` objects this view is built from: none. + + A tractogram carries its own geometry rather than wrapping + :class:`~cortex.dataset.braindata.BrainData`, so it contributes + nothing here. The method exists because `Dataset.uniques` calls it + on every view it holds, and without it a `Dataset` containing a + tractogram could not be packaged or saved at all. + """ + return [] # -- selection / decimation ----------------------------------------- @@ -263,8 +298,9 @@ def select(self, indices: npt.ArrayLike) -> "Tractogram": Parameters ---------- indices : array_like of int or bool - Streamline indices to keep (or a boolean mask of length - `n_streamlines`). Order is preserved; repeats are allowed. + Streamline indices to keep, or a boolean mask which must have + exactly `n_streamlines` entries. Order is preserved; repeats + are allowed. Returns ------- @@ -275,9 +311,26 @@ def select(self, indices: npt.ArrayLike) -> "Tractogram": `linewidth`, `description`) are copied. """ idx = np.asarray(indices) + if idx.ndim != 1: + raise ValueError( + "indices must be 1-D (streamline indices or a boolean mask), " + "got shape %r" % (idx.shape,) + ) if idx.dtype == np.bool_: + if idx.size != self.n_streamlines: + raise ValueError( + "boolean mask must have %d entries (one per streamline), " + "got %d" % (self.n_streamlines, idx.size) + ) idx = np.nonzero(idx)[0] idx = idx.astype(np.int64) + if idx.size and (idx.min() < 0 or idx.max() >= self.n_streamlines): + # Caught here rather than left to fall through: a negative index + # would quietly wrap when slicing `points`, yet never match in the + # `groups` remapping below, giving a silently inconsistent result. + raise IndexError( + "streamline indices out of range [0, %d)" % self.n_streamlines + ) starts = self.offsets[:-1] stops = self.offsets[1:] @@ -389,6 +442,12 @@ def vertex_colors(self) -> npt.NDArray[np.uint8]: (N, 3) ndarray of uint8 One RGB color per point in `points`. """ + if self.n_points == 0: + # Empty tractograms are legitimate (an empty named group, or a + # `select` that matched nothing). Short-circuit before the scalar + # paths, whose `reshape(0, -1)` numpy cannot infer. + return np.zeros((0, 3), dtype=np.uint8) + color = self.color if isinstance(color, str): if color == "orientation": @@ -503,6 +562,11 @@ def to_json(self, simple: bool = False) -> dict: # type: ignore[override] Notes ----- + The base :class:`~cortex.dataset.views.Dataview` keys ``state``, + ``attrs`` and ``desc`` are carried through alongside the + tract-specific ones, so a tractogram's metadata satisfies the same + contract as every other view's. + ``result["groups"]`` maps each group name to a ``[start, stop]`` slice (not a count): these are bounds into the fourth per-tractogram wire buffer, ``groups`` (little-endian uint32 streamline indices, @@ -515,6 +579,10 @@ def to_json(self, simple: bool = False) -> dict: # type: ignore[override] else: color_repr = str(tuple(float(c) for c in self.color)) + desc = self.description + if isinstance(desc, bytes): + desc = desc.decode() + _, group_slices = self.groups_wire() result: dict = { "subject": self.subject, @@ -528,7 +596,14 @@ def to_json(self, simple: bool = False) -> dict: # type: ignore[override] name: [int(start), int(stop)] for name, (start, stop) in group_slices.items() }, - "description": self.description, + "description": desc, + # The base `Dataview` metadata contract: `cortex.webgl.data` + # and the javascript dataset machinery read these keys off every + # view's json, and `state`/`attrs` would otherwise be dropped on + # the way to the viewer. + "state": self.state, + "attrs": self.attrs.copy(), + "desc": desc, } is_scalar_mode = isinstance(self.color, str) and ( @@ -619,6 +694,11 @@ def _extract(value: Any) -> npt.NDArray: if xfm is not None: xfm_arr = np.asarray(xfm, dtype=np.float64) + if xfm_arr.shape != (4, 4): + raise ValueError( + "xfm must be a (4, 4) affine, got shape %r" + % (xfm_arr.shape,) + ) homogeneous = np.concatenate( [points.astype(np.float64), np.ones((n_points, 1))], axis=1 ) diff --git a/cortex/tests/test_tractogram.py b/cortex/tests/test_tractogram.py index b9d140a7a..6063c21d6 100644 --- a/cortex/tests/test_tractogram.py +++ b/cortex/tests/test_tractogram.py @@ -146,6 +146,25 @@ def test_group_out_of_range_raises(): Tractogram(points, offsets, subj, groups={"bad": np.array([5])}) +def test_group_not_1d_raises(): + points = np.zeros((10, 3), dtype=np.float32) + offsets = np.array([0, 5, 10]) + with pytest.raises(ValueError): + Tractogram(points, offsets, subj, groups={"bad": np.array([[0, 1]])}) + + +def test_group_boolean_mask(): + points = np.zeros((10, 3), dtype=np.float32) + offsets = np.array([0, 5, 10]) + tract = Tractogram( + points, offsets, subj, groups={"first": np.array([True, False])} + ) + np.testing.assert_array_equal(tract.groups["first"], [0]) + # a mask of the wrong length is an error, not silently truncated + with pytest.raises(ValueError): + Tractogram(points, offsets, subj, groups={"bad": np.array([True])}) + + def test_unrecognized_color_raises(): points = np.zeros((10, 3), dtype=np.float32) offsets = np.array([0, 10]) @@ -185,6 +204,25 @@ def test_name_is_deterministic_hash(): assert tract1.name != tract3.name +def test_name_depends_on_offsets(): + """Same points, different streamline boundaries -> different identity.""" + points = np.random.RandomState(0).randn(10, 3).astype(np.float32) + one = Tractogram(points, np.array([0, 10]), subj) + two = Tractogram(points, np.array([0, 4, 10]), subj) + np.testing.assert_array_equal(one.points, two.points) + assert one.name != two.name + + +def test_uniques_is_empty_and_dataset_iterates(): + tract = _make_tractogram(n_streamlines=4, n_points=5) + assert list(tract.uniques()) == [] + assert list(tract.uniques(collapse=True)) == [] + # `Dataset.uniques` calls `.uniques()` on every view it holds, so a + # dataset carrying a tractogram must not blow up there. + ds = cortex.Dataset(tract=tract) + assert ds.uniques() == set() + + # --------------------------------------------------------------------------- # select / get_group / subsample # --------------------------------------------------------------------------- @@ -214,6 +252,36 @@ def test_select_remaps_groups(): assert sub.groups["even"].size == 0 +def test_select_boolean_mask(): + tract = _make_tractogram(n_streamlines=6, n_points=5) + mask = np.zeros(6, dtype=bool) + mask[[1, 4]] = True + sub = tract.select(mask) + assert sub.n_streamlines == 2 + np.testing.assert_array_equal(sub.streamlines[0], tract.streamlines[1]) + np.testing.assert_array_equal(sub.streamlines[1], tract.streamlines[4]) + + +def test_select_short_boolean_mask_raises(): + tract = _make_tractogram(n_streamlines=6, n_points=5) + with pytest.raises(ValueError): + tract.select(np.array([True])) + + +def test_select_non_1d_raises(): + tract = _make_tractogram(n_streamlines=6, n_points=5) + with pytest.raises(ValueError): + tract.select(np.array([[0, 1]])) + + +def test_select_out_of_range_raises(): + tract = _make_tractogram(n_streamlines=6, n_points=5) + with pytest.raises(IndexError): + tract.select([6]) + with pytest.raises(IndexError): + tract.select([-1]) + + def test_get_group(): tract = _make_tractogram(n_streamlines=10, n_points=5) even = tract.get_group("even") @@ -314,6 +382,35 @@ def test_vertex_colors_dps(): assert np.all(colors[start:stop] == colors[start]) +def test_vertex_colors_empty_tractogram(): + """Every colour mode returns an empty (0, 3) array, not an exception.""" + empty_points = np.zeros((0, 3), dtype=np.float32) + empty_offsets = np.array([0], dtype=np.int64) + for color, kwargs in [ + ("orientation", {}), + ((1.0, 0.0, 0.0), {}), + ("dpv:scalar", dict(dpv={"scalar": np.zeros(0, dtype=np.float32)})), + ("dps:length", dict(dps={"length": np.zeros(0, dtype=np.float32)})), + ]: + tract = Tractogram( + empty_points, empty_offsets, subj, color=color, **kwargs + ) + colors = tract.vertex_colors() + assert colors.shape == (0, 3) + assert colors.dtype == np.uint8 + + +def test_vertex_colors_empty_group(): + """The realistic path to an empty tractogram: a group nothing matched.""" + tract = _make_tractogram( + n_streamlines=6, n_points=5, color="dpv:scalar", + groups={"none": np.array([], dtype=np.int64)}, + ) + empty = tract.get_group("none") + assert empty.n_streamlines == 0 + assert empty.vertex_colors().shape == (0, 3) + + # --------------------------------------------------------------------------- # to_json / _write_hdf / normalize # --------------------------------------------------------------------------- @@ -346,6 +443,16 @@ def test_to_json_keys(): json.dumps(j) +def test_to_json_keeps_base_dataview_fields(): + """`state`/`attrs`/`desc` travel with every view, tractograms included.""" + tract = _make_tractogram(state={"foo": 1}, priority=3) + j = tract.to_json() + assert j["state"] == {"foo": 1} + assert j["attrs"]["priority"] == 3 + assert j["desc"] == tract.description + json.dumps(j) + + def test_groups_wire_disjoint(): tract = _make_tractogram(n_streamlines=10, n_points=5) indices, slices = tract.groups_wire() @@ -479,6 +586,13 @@ def test_from_trx_roundtrip(tmp_path): ) +def test_from_trx_bad_xfm_raises(tmp_path): + pytest.importorskip("trx") + path, _, _, _, _, _ = _build_trx_fixture(tmp_path) + with pytest.raises(ValueError): + Tractogram.from_trx(path, subj, xfm=np.eye(3)) + + def test_from_trx_xfm_shifts_points(tmp_path): pytest.importorskip("trx") path, points, offsets, dpv, dps, groups = _build_trx_fixture(tmp_path) diff --git a/pyproject.toml b/pyproject.toml index 5d2958257..176c0de7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,11 @@ dev = [ ] test = [ "pycortex[headless]", - "pycortex[tractography]", + # trx-python requires Python >= 3.11. The marker keeps the rest of the + # test group installable on 3.10 -- which the wheel-install workflow + # still covers and which the package does not exclude -- where the TRX + # tests skip themselves via `pytest.importorskip`. + "pycortex[tractography]; python_version >= '3.11'", # The visual-regression tests compare against renders from # these builds. Both are unpinned for users. # See cortex/tests/reference_images/README.md