nltools 0.6.0.dev0__py3-none-any.whl

nltools/data/atlases/labeling.py

@@ -0,0 +1,158 @@
+"""Coordinate-level atlas labeling.
+
+Adapted from [atlasreader](https://github.com/miykael/atlasreader)
+(BSD-3-Clause); please cite it when using these tools.
+
+References:
+    Notter, M. P., Gale, D., Herholz, P., Markello, R., Notter-Bielser, M.-L., &
+    Whittingstall, K. (2019). AtlasReader: A Python package to generate
+    coordinate tables, region labels, and informative figures from statistical
+    MRI images. *Journal of Open Source Software*, 4(34), 1257.
+    https://doi.org/10.21105/joss.01257
+"""
+
+from collections.abc import Sequence
+
+import numpy as np
+import polars as pl
+from numpy.typing import ArrayLike
+
+from .loading import _Atlas, load_atlas
+
+CoordsLike = ArrayLike | Sequence[Sequence[float]]
+
+
+def _as_xyz_array(coords: CoordsLike) -> np.ndarray:
+    """Coerce coords to an ``(N, 3)`` float64 array."""
+    arr = np.asarray(coords, dtype=float)
+    if arr.ndim == 1:
+        arr = arr[None, :]
+    if arr.ndim != 2 or arr.shape[1] != 3:
+        raise ValueError(f"coords must have shape (N, 3); got {arr.shape}")
+    return arr
+
+
+def _xyz_to_ijk(coords_xyz: np.ndarray, affine: np.ndarray) -> np.ndarray:
+    """Map MNI mm → integer voxel ijk via the inverse affine.
+
+    Solves the affine, rounds to the nearest voxel, and casts to int.
+    """
+    homog = np.hstack([coords_xyz, np.ones((coords_xyz.shape[0], 1))])
+    ijk = np.linalg.solve(affine, homog.T)[:3].T
+    return np.round(ijk).astype(int)
+
+
+def _clip_to_box(ijk: np.ndarray, shape: tuple[int, ...]) -> np.ndarray:
+    """Clamp out-of-bounds voxel indices to the origin (background)."""
+    box = np.array(shape[:3])
+    out = np.any((ijk < 0) | (ijk >= box), axis=1)
+    ijk = ijk.copy()
+    ijk[out] = 0
+    return ijk
+
+
+def _label_lookup(atlas: _Atlas) -> dict[int, str]:
+    """Build an ``{integer index → region name}`` dict from atlas labels."""
+    return dict(
+        zip(
+            atlas.labels["index"].to_list(),
+            atlas.labels["name"].to_list(),
+            strict=True,
+        )
+    )
+
+
+def _label_deterministic(atlas: _Atlas, ijk: np.ndarray) -> list[str]:
+    """Look up region names for voxels in a deterministic atlas."""
+    data = atlas.image.get_fdata()
+    lut = _label_lookup(atlas)
+    out: list[str] = []
+    for v in ijk:
+        idx = int(data[v[0], v[1], v[2]])
+        out.append(lut.get(idx, "no_label"))
+    return out
+
+
+def _format_prob_entry(pct: float, name: str) -> str:
+    """Format a single ``"pct%% name"`` entry for probabilistic output."""
+    return f"{pct:.1f}% {name}"
+
+
+def _label_probabilistic(
+    atlas: _Atlas, ijk: np.ndarray, prob_threshold: float
+) -> list[str]:
+    """Format ``"X% Foo; Y% Bar"`` strings per voxel for a probabilistic atlas.
+
+    Regions with probability below ``prob_threshold`` are dropped; if no
+    region survives, the voxel gets ``"no_label"``.
+    """
+    data = atlas.image.get_fdata()
+    lut = _label_lookup(atlas)
+    out: list[str] = []
+    for v in ijk:
+        probs = np.asarray(data[v[0], v[1], v[2]], dtype=float)
+        keep = np.where(probs >= prob_threshold)[0]
+        if keep.size == 0:
+            out.append("no_label")
+            continue
+        order = keep[np.argsort(probs[keep])[::-1]]
+        out.append(
+            "; ".join(
+                _format_prob_entry(float(probs[i]), lut.get(int(i), "no_label"))
+                for i in order
+            )
+        )
+    return out
+
+
+def _labels_for_atlas(
+    atlas: _Atlas, ijk: np.ndarray, prob_threshold: float
+) -> list[str]:
+    """Dispatch to the deterministic or probabilistic labeling path."""
+    if atlas.kind == "probabilistic":
+        return _label_probabilistic(atlas, ijk, prob_threshold)
+    return _label_deterministic(atlas, ijk)
+
+
+def label_coords(
+    coords: CoordsLike,
+    *,
+    atlas: str | Sequence[str] = "harvard_oxford",
+    prob_threshold: float = 5.0,
+) -> pl.DataFrame:
+    """Look up anatomical labels for a set of MNI mm coordinates.
+
+    For each coordinate, returns the atlas region(s) it falls in. Works
+    for both deterministic atlases (single label per coord) and
+    probabilistic atlases (formatted `"42.0% Foo; 18.0% Bar"` strings,
+    sorted by descending probability).
+
+    Args:
+        coords (array-like): `(N, 3)` MNI mm coordinates `(x, y, z)`. A single
+            coordinate like `(-42, -22, 56)` is also accepted.
+        atlas (str | Sequence[str]): Atlas name or list of names from
+            `list_atlases`. One column is added to the output per atlas. Default
+            `'harvard_oxford'`.
+        prob_threshold (float): For probabilistic atlases only — drop regions
+            with probability (in percent units) below this threshold. Default 5.0.
+
+    Returns:
+        pl.DataFrame: Frame with columns `x`, `y`, `z` plus one column per atlas.
+            All atlas columns are `Utf8`.
+    """
+    xyz = _as_xyz_array(coords)
+    atlas_names = [atlas] if isinstance(atlas, str) else list(atlas)
+
+    columns: dict[str, list] = {
+        "x": xyz[:, 0].tolist(),
+        "y": xyz[:, 1].tolist(),
+        "z": xyz[:, 2].tolist(),
+    }
+
+    for name in atlas_names:
+        a = load_atlas(name)
+        ijk = _xyz_to_ijk(xyz, a.image.affine)
+        ijk = _clip_to_box(ijk, a.image.shape)
+        columns[name] = _labels_for_atlas(a, ijk, prob_threshold)
+
+    return pl.DataFrame(columns)

nltools/data/atlases/loading.py

@@ -0,0 +1,76 @@
+"""Lazy loading of atlas NIfTI + label CSV files from the HF dataset."""
+
+import functools
+from dataclasses import dataclass
+
+import nibabel as nb
+import polars as pl
+
+from nltools.templates.fetch import fetch_resource
+
+from .registry import ATLASES, AtlasKind
+
+
+@dataclass(frozen=True)
+class _Atlas:
+    """A loaded atlas — image, labels, and metadata.
+
+    Constructed by `load_atlas`; users normally don't instantiate
+    directly.
+
+    Attributes:
+        name (str): Registry key (e.g. `'harvard_oxford'`).
+        image (nibabel.Nifti1Image): NIfTI volume. 3-D for deterministic atlases,
+            4-D for probabilistic ones (last axis indexes regions).
+        labels (pl.DataFrame): Two-column `index, name` table. For deterministic
+            atlases `index` is the integer voxel value; for probabilistic atlases
+            `index` is the region index along the 4th dim of `image`.
+        kind (AtlasKind): `'deterministic'` or `'probabilistic'`.
+        citation (str): Short citation for the original atlas.
+    """
+
+    name: str
+    image: nb.Nifti1Image
+    labels: pl.DataFrame
+    kind: AtlasKind
+    citation: str
+
+
+@functools.cache
+def load_atlas(name: str) -> _Atlas:
+    """Lazy-load an atlas by registry name.
+
+    The first call fetches the NIfTI + label CSV from
+    `huggingface.co/datasets/nltools/niftis` (cached locally afterwards).
+    Subsequent calls in the same process are memoized.
+
+    Args:
+        name (str): Atlas key from `list_atlases`.
+
+    Returns:
+        _Atlas: The atlas with image, labels, and metadata loaded.
+
+    Raises:
+        ValueError: If `name` isn't a registered atlas.
+
+    Note:
+        An atlas is in MNI space, and nltools resamples it onto your data by
+        header affine alone. That is a grid change, not a spatial
+        normalization, so parcel boundaries are approximate unless your data
+        are already normalized to the same space.
+    """
+    if name not in ATLASES:
+        known = ", ".join(sorted(ATLASES))
+        raise ValueError(f"unknown atlas {name!r}; choose from: {known}")
+
+    meta = ATLASES[name]
+    img_path = fetch_resource(f"atlases/atlas_{name}.nii.gz")
+    csv_path = fetch_resource(f"atlases/labels_{name}.csv")
+
+    return _Atlas(
+        name=name,
+        image=nb.load(img_path),
+        labels=pl.read_csv(csv_path),
+        kind=meta.kind,
+        citation=meta.citation,
+    )

nltools/data/atlases/registry.py

@@ -0,0 +1,96 @@
+"""Static registry of atlases hosted at `nltools/niftis/atlases`.
+
+Each entry describes an atlas's kind (deterministic vs probabilistic) and
+the citation users should cite when they use it. The actual NIfTI + label
+files are fetched lazily by `load_atlas` via `fetch_resource`.
+
+Atlases were sourced from atlasreader (BSD-3-Clause) and are subject to
+their original upstream licenses — see `LICENSES.md` in the HF dataset.
+"""
+
+from dataclasses import dataclass
+from typing import Literal
+
+AtlasKind = Literal["deterministic", "probabilistic"]
+"""Kind of atlas: `'deterministic'` (3-D integer labels) or `'probabilistic'` (4-D, last axis indexes regions)."""
+
+
+@dataclass(frozen=True)
+class _AtlasMetadata:
+    """Static description of a registered atlas.
+
+    Attributes:
+        kind (AtlasKind): `'deterministic'` (3-D integer-labeled) or
+            `'probabilistic'` (4-D, last axis indexes regions).
+        citation (str): Short citation string for the original atlas.
+    """
+
+    kind: AtlasKind
+    citation: str
+
+
+ATLASES: dict[str, _AtlasMetadata] = {
+    "aal": _AtlasMetadata(
+        kind="deterministic",
+        citation="Tzourio-Mazoyer et al. 2002, NeuroImage",
+    ),
+    "aicha": _AtlasMetadata(
+        kind="deterministic",
+        citation="Joliot et al. 2015, J Neurosci Methods",
+    ),
+    "desikan_killiany": _AtlasMetadata(
+        kind="deterministic",
+        citation="Desikan et al. 2006, NeuroImage (FreeSurfer license)",
+    ),
+    "destrieux": _AtlasMetadata(
+        kind="deterministic",
+        citation="Destrieux et al. 2010, NeuroImage (FreeSurfer license)",
+    ),
+    "harvard_oxford": _AtlasMetadata(
+        kind="probabilistic",
+        citation="Desikan et al. 2006, NeuroImage / FSL Harvard-Oxford",
+    ),
+    "juelich": _AtlasMetadata(
+        kind="probabilistic",
+        citation="Eickhoff et al. 2005, NeuroImage",
+    ),
+    "marsatlas": _AtlasMetadata(
+        kind="deterministic",
+        citation="Auzias et al. 2016, Hum Brain Mapp",
+    ),
+    "neuromorphometrics": _AtlasMetadata(
+        kind="deterministic",
+        citation="MICCAI 2012 Multi-Atlas Labeling Challenge",
+    ),
+    "schaefer_200": _AtlasMetadata(
+        kind="deterministic",
+        citation="Schaefer et al. 2018, Cereb Cortex (200-parcel, 7-network)",
+    ),
+    "talairach_ba": _AtlasMetadata(
+        kind="deterministic",
+        citation="Talairach & Tournoux 1988 (Brodmann areas)",
+    ),
+    "talairach_gyrus": _AtlasMetadata(
+        kind="deterministic",
+        citation="Talairach & Tournoux 1988 (gyri)",
+    ),
+}
+"""Registered atlases keyed by name; each value is an `_AtlasMetadata` (kind + citation)."""
+
+
+DEFAULT_ATLASES: tuple[str, ...] = ("harvard_oxford", "aal", "schaefer_200")
+"""Default atlas trio for `BrainData.cluster_report` and `label_coords`.
+
+Picked to give one probabilistic (`harvard_oxford`), one anatomical
+deterministic (`aal`), and one functional deterministic (`schaefer_200`) atlas
+at once.
+"""
+
+
+def list_atlases() -> list[str]:
+    """Return the sorted list of registered atlas names.
+
+    Returns:
+        list[str]: Sorted list of atlas names usable with `load_atlas`.
+    """
+    return sorted(ATLASES.keys())