bifti 0.1.0__py3-none-any.whl

bifti/__init__.py

@@ -0,0 +1,43 @@
+from .phantom import (
+    BiftiPhantom,
+    BiftiTissue,
+    NiftiMapping,
+    NiftiRef,
+    PATIENT_POSITIONS,
+    Patient,
+    PhantomSystem,
+    PhantomUnits,
+    ResliceTo,
+    patient_to_scanner,
+)
+from .loader import (
+    NumpyPhantom,
+    NumpyTissue,
+    to_scanner_affine,
+)
+from .registry import (
+    flatten_phantoms,
+    load_catalog,
+    load_registry,
+    load_registry_phantom,
+)
+
+__all__ = [
+    "BiftiPhantom",
+    "BiftiTissue",
+    "NiftiMapping",
+    "NiftiRef",
+    "PATIENT_POSITIONS",
+    "Patient",
+    "PhantomSystem",
+    "PhantomUnits",
+    "ResliceTo",
+    "patient_to_scanner",
+    "NumpyPhantom",
+    "NumpyTissue",
+    "to_scanner_affine",
+    "flatten_phantoms",
+    "load_catalog",
+    "load_registry",
+    "load_registry_phantom",
+]

bifti/_backend.py

@@ -0,0 +1,110 @@
+"""Array backend for resampling: torch when it is installed, NumPy otherwise.
+
+The separable resampling path (see :mod:`bifti.resample`) is three matrix contractions, so
+it maps straight onto whichever array library is available. Torch is used whenever it can be
+imported - *not* only when CUDA is present, since torch is faster than NumPy on CPU too and
+falls back to the same code path either way.
+"""
+
+from __future__ import annotations
+
+import os
+import warnings
+from functools import lru_cache
+
+import numpy as np
+
+
+class Backend:
+    """The handful of array operations resampling needs."""
+
+    name = "numpy"
+
+    def asarray(self, x):
+        return np.asarray(x, dtype=np.float64)
+
+    def to_numpy(self, x) -> np.ndarray:
+        return np.asarray(x, dtype=np.float64)
+
+    def tensordot(self, a, b, axes):
+        return np.tensordot(a, b, axes=axes)
+
+    def moveaxis(self, x, source, destination):
+        return np.moveaxis(x, source, destination)
+
+    def transpose(self, x, axes):
+        return np.transpose(x, axes)
+
+
+class TorchBackend(Backend):
+    """Torch-backed contractions, on CUDA when available.
+
+    Uses float64 on CPU so results match the NumPy backend exactly, and float32 on CUDA
+    where float64 throughput is heavily penalised. That makes the GPU path accurate to
+    roughly 1e-6 relative rather than to machine precision - far below the accuracy of the
+    phantom data itself.
+    """
+
+    name = "torch"
+
+    def __init__(self, torch):
+        self._torch = torch
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.dtype = torch.float32 if self.device.type == "cuda" else torch.float64
+
+    def asarray(self, x):
+        if self._torch.is_tensor(x):
+            return x.to(device=self.device, dtype=self.dtype)
+        return self._torch.as_tensor(
+            np.ascontiguousarray(x, dtype=np.float64),
+            dtype=self.dtype,
+            device=self.device,
+        )
+
+    def to_numpy(self, x) -> np.ndarray:
+        if self._torch.is_tensor(x):
+            return x.detach().to("cpu", self._torch.float64).numpy()
+        return np.asarray(x, dtype=np.float64)
+
+    def tensordot(self, a, b, axes):
+        return self._torch.tensordot(a, b, dims=axes)
+
+    def moveaxis(self, x, source, destination):
+        return self._torch.movedim(x, source, destination)
+
+    def transpose(self, x, axes):
+        return self._torch.permute(x, tuple(int(a) for a in axes))
+
+
+def _load_torch_backend() -> Backend | None:
+    try:
+        import torch
+    except ImportError:
+        return None
+    try:
+        return TorchBackend(torch)
+    except Exception as exc:  # pragma: no cover - torch present but unusable
+        warnings.warn(f"torch is installed but unusable ({exc}); falling back to NumPy")
+        return None
+
+
+@lru_cache(maxsize=None)
+def get_backend(name: str | None = None) -> Backend:
+    """The array backend to resample with.
+
+    ``name`` (or ``$BIFTI_RESAMPLE_BACKEND``) may be ``"auto"`` (the default), ``"numpy"``
+    or ``"torch"``. ``"torch"`` raises if torch is missing; ``"auto"`` quietly falls back.
+    """
+    name = (name or os.environ.get("BIFTI_RESAMPLE_BACKEND", "auto")).strip().lower()
+    if name == "numpy":
+        return Backend()
+    if name == "torch":
+        backend = _load_torch_backend()
+        if backend is None:
+            raise ImportError("BIFTI_RESAMPLE_BACKEND=torch but torch is not installed")
+        return backend
+    if name != "auto":
+        raise ValueError(
+            f"invalid resampling backend {name!r} (use auto, numpy or torch)"
+        )
+    return _load_torch_backend() or Backend()

bifti/_nifti_write.py

@@ -0,0 +1,112 @@
+"""Deterministic NIfTI-1 writer, the same bytes as the Rust, Julia and MATLAB writers.
+
+A phantom directory is hashed as a whole, so the output must be reproducible: fixed
+header fields, little-endian, no extensions, and a gzip header without mtime or file name.
+nibabel is not used for writing because it would put its own defaults (scaling NaNs,
+a ``descrip`` string, the gzip mtime) into the file.
+"""
+
+from __future__ import annotations
+
+import gzip
+import struct
+from pathlib import Path
+
+import numpy as np
+
+HEADER_SIZE = 348
+VOX_OFFSET = 352  # header plus the 4-byte (empty) extension flag
+GZIP_LEVEL = 6
+GZIP_OS_UNKNOWN = 255
+
+# numpy dtype -> (NIfTI datatype code, bitpix)
+_DATATYPES: dict[np.dtype, tuple[int, int]] = {
+    np.dtype("u1"): (2, 8),
+    np.dtype("<i2"): (4, 16),
+    np.dtype("<i4"): (8, 32),
+    np.dtype("<f4"): (16, 32),
+    np.dtype("<c8"): (32, 64),
+    np.dtype("<f8"): (64, 64),
+    np.dtype("i1"): (256, 8),
+    np.dtype("<u2"): (512, 16),
+    np.dtype("<u4"): (768, 32),
+    np.dtype("<i8"): (1024, 64),
+    np.dtype("<u8"): (1280, 64),
+    np.dtype("<c16"): (1792, 128),
+}
+
+
+def _stack_dtype(volumes: list[np.ndarray]) -> np.dtype:
+    """The common little-endian dtype: kept if all agree, else float64 / complex128."""
+    kinds = {v.dtype.newbyteorder("<") for v in volumes}
+    if len(kinds) == 1:
+        (dtype,) = kinds
+        if dtype not in _DATATYPES:
+            raise ValueError(f"cannot write NIfTI data of type {dtype}")
+        return dtype
+    if any(v.dtype.kind == "c" for v in volumes):
+        return np.dtype("<c16")
+    return np.dtype("<f8")
+
+
+def build_header(shape, n_volumes: int, affine, dtype: np.dtype) -> bytes:
+    """The 352 bytes before the voxel data.
+
+    Orientation is stored in the sform only (``sform_code = 2``, ``qform_code = 0``), as
+    ``../../NIFTI.md`` requires.
+    """
+    affine = np.asarray(affine, dtype=np.float64)[:3]
+    datatype, bitpix = _DATATYPES[dtype]
+    header = bytearray(VOX_OFFSET)
+    struct.pack_into("<i", header, 0, HEADER_SIZE)
+    struct.pack_into("<8h", header, 40, 4, *shape, n_volumes, 1, 1, 1)
+    struct.pack_into("<hh", header, 70, datatype, bitpix)
+    voxel_sizes = np.linalg.norm(affine[:, :3], axis=0)
+    struct.pack_into("<5f", header, 76, 1.0, *voxel_sizes.astype(np.float32), 1.0)
+    struct.pack_into("<f", header, 108, float(VOX_OFFSET))  # vox_offset
+    struct.pack_into("<f", header, 112, 1.0)  # scl_slope
+    header[123] = 2  # xyzt_units: millimetres
+    struct.pack_into("<hh", header, 252, 0, 2)  # qform_code, sform_code
+    struct.pack_into("<12f", header, 280, *affine.astype(np.float32).ravel())
+    header[344:348] = b"n+1\0"
+    return bytes(header)
+
+
+def write_stack(path: Path | str, volumes: list[np.ndarray], affine) -> list[int]:
+    """Write 3-D ``volumes`` as one 4-D NIfTI-1 file and return each one's 4th-dim index.
+
+    The volumes share ``affine`` and shape. Volumes with identical stored bytes share an
+    index and are written once. A path ending in ``.gz`` is gzip-compressed.
+    """
+    path = Path(path)
+    if not volumes:
+        raise ValueError(f"{path}: no volumes to write")
+    shape = volumes[0].shape
+    if len(shape) != 3 or any(v.shape != shape for v in volumes):
+        raise ValueError(
+            f"{path}: volumes must all be 3-D and share a shape, got "
+            f"{sorted({v.shape for v in volumes})}"
+        )
+    dtype = _stack_dtype(volumes)
+
+    stored: dict[bytes, int] = {}
+    slots = []
+    for volume in volumes:
+        # NIfTI stores voxels x-fastest, which is Fortran order.
+        data = np.asarray(volume).astype(dtype, copy=False).ravel(order="F").tobytes()
+        slots.append(stored.setdefault(data, len(stored)))
+
+    payload = b"".join(stored)  # dicts keep insertion order = slot order
+    blob = build_header(shape, len(stored), affine, dtype) + payload
+    if path.suffix == ".gz":
+        compressed = bytearray(gzip.compress(blob, compresslevel=GZIP_LEVEL, mtime=0))
+        compressed[9] = GZIP_OS_UNKNOWN  # zlib writes the host's OS code
+        blob = bytes(compressed)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_bytes(blob)
+    return slots
+
+
+def save_nifti(path: Path | str, volume: np.ndarray, affine) -> None:
+    """Write one 3-D volume as a NIfTI-1 file with a singleton 4th dimension."""
+    write_stack(path, [volume], affine)