from __future__ import annotations
import tempfile
from typing import Iterator, Optional, Union, overload, Literal
import numpy as np
import numpy.typing as npt
import h5py
from ..database import db
from ..xfm import Transform
from .braindata import _hdf_write
from .views import normalize as _vnorm
from .views import Dataview, Vertex, Volume, _from_hdf_data
[docs]
class Dataset:
"""
Wrapper for a named collection of Dataview objects (Volume, Vertex, VolumeRGB,
etc.). Provides a standard way to save/load that collection as an HDF5 file,
and lets a single Dataview be treated the same as a full Dataset elsewhere in
pycortex (e.g. `cortex.webshow`).
Each keyword argument names one view; its value can be a Dataview, a tuple
(implicitly converted to a Volume or Vertex), a filename string (loaded from
an .hdf file), a dict (converted to a nested Dataset), or another Dataset
(merged in).
"""
[docs]
def __init__(self, **kwargs: Union[Dataview, dict, str, tuple, Dataset]) -> None:
self.h5: Optional[h5py.File] = None
self.views: dict[str, Dataview] = {}
self.append(**kwargs)
[docs]
def append(self, **kwargs: Union[Dataview, dict, str, tuple, Dataset]) -> Dataset:
"""Add the views in `kwargs` into this dataset. Each keyword names
one view; its value can be a Dataview, a tuple (implicitly converted
to a Volume or Vertex), a filename string (loaded from an .hdf
file), a dict (converted to a nested Dataset), or another Dataset
(merged in). See the `Dataset` class docstring for details.
Returns
-------
Dataset
This dataset, with the new views added, for chaining.
"""
for name, data in kwargs.items():
norm = normalize(data)
if isinstance(norm, Dataview):
self.views[name] = norm
elif isinstance(norm, Dataset):
self.views.update(norm.views)
else:
raise ValueError("Unknown input %s=%r"%(name, data))
return self
def __getattr__(self, attr: str):
if attr in self.__dict__:
return self.__dict__[attr]
elif attr in self.views:
return self.views[attr]
raise AttributeError
def __getitem__(self, item: str) -> Dataview:
return self.views[item]
def __iter__(self) -> Iterator[tuple[str, Dataview]]:
for name, dv in sorted(self.views.items(), key=lambda x: x[1].priority):
yield name, dv
def __repr__(self):
views = sorted(self.views.items(), key=lambda x: x[1].priority)
return "<Dataset with views [%s]>"%(', '.join([n for n, d in views]))
def __len__(self):
return len(self.views)
def __dir__(self):
return list(self.__dict__.keys()) + list(self.views.keys())
[docs]
@classmethod
def from_file(cls, filename: str, subject: Optional[str]=None) -> Dataset:
"""Load a pycortex Dataset (cortex.Dataset class) from a file
Parameters
----------
filename : str
.hdf file from which to load
subject : str, optional
string ID for pycortex subject, if subject name has changed
since the Dataset was created. `None` input assumes subject
name in saved file is a subject in your current pycortex
filestore. By default None
Returns
-------
Dataset
pycortex Dataset
"""
ds = cls()
ds.h5 = h5py.File(filename, 'r')
db.auxfile = ds
#detect stray datasets which were not written by pycortex
for name, node in ds.h5.items():
if name in ("data", "subjects", "views"):
continue
try:
ds.views[name] = _from_hdf_data(ds.h5, name, subject=subject)
except KeyError:
print('No metadata found for "%s", skipping...'%name)
#load up the views generated by pycortex
for name, node in ds.h5['views'].items():
try:
ds.views[name] = Dataview.from_hdf(node, subject=subject)
except FileNotFoundError:
print("Could not load file; old subject name? Try using `subject` kwarg to specify a current pycortex subject")
raise
except Exception:
import traceback
traceback.print_exc()
db.auxfile = None
return ds
[docs]
def uniques(self, collapse: bool=False) -> set[Dataview]:
"""Return the set of unique BrainData objects contained by this dataset
Parameters
----------
collapse : bool, optional
Passed through to each view's own `uniques()`. Default False.
Returns
-------
uniques : set
The distinct BrainData objects (e.g. Volume/Vertex channels)
referenced across all views in this dataset.
"""
uniques = set()
for name, view in self:
# .uniques() is provided by BrainData
for sv in view.uniques(collapse=collapse):
uniques.add(sv)
return uniques
[docs]
def save(self, filename: Optional[str]=None, pack: bool=False) -> None:
"""Write this dataset's views out to an HDF5 file.
Parameters
----------
filename : str, optional
Path to write to. If omitted, writes to the file this dataset
was already opened from (raises if there is none).
pack : bool, optional
If True, also bundle each view's subject surfaces, transforms,
and masks into the file, so it can be opened without access to
the pycortex filestore/database. Default False.
Returns
-------
None
Writes `filename` (or the already-open file) as a side effect;
has no return value.
"""
if filename is not None:
self.h5 = h5py.File(filename, 'a')
elif self.h5 is None:
raise ValueError("Must provide filename for new datasets")
for name, view in self.views.items():
view._write_hdf(self.h5, name=name)
if pack:
subjs: set[str] = set()
xfms: set[tuple[str, str]] = set()
masks: set[tuple[str, str, str]] = set()
for view in self.views.values():
# .uniques() is provided by BrainData
for data in view.uniques():
subjs.add(data.subject)
if isinstance(data, Volume):
xfms.add((data.subject, data.xfmname))
#custom masks are already packaged by default
#only string masks need to be packed
if isinstance(data._mask, str):
masks.add((data.subject, data.xfmname, data._mask))
_pack_subjs(self.h5, subjs)
_pack_xfms(self.h5, xfms)
_pack_masks(self.h5, masks)
self.h5.flush()
@overload
def get_surf(self, subject: str, type: str, hemi: Literal['both']='both', *, merge: Literal[False]=False, nudge: bool=False) -> tuple[tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]], tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]]]: ...
@overload
def get_surf(self, subject: str, type: str, hemi: Literal['both']='both', *, merge: Literal[True], nudge: bool=False) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]]: ...
@overload
def get_surf(self, subject: str, type: str, hemi: Literal['lh', 'rh'], *, merge: bool=False, nudge: bool=False) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]]: ...
# Fallthrough case for callers (e.g. Database.get_surf's auxfile delegation) that pass
# a non-literal hemi/merge, which can't be matched to any of the overloads above.
@overload
def get_surf(self, subject: str, type: str, hemi: Literal['both', 'lh', 'rh']='both', *, merge: bool=False, nudge: bool=False) -> Union[tuple[tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]], tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]]], tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]]]: ...
[docs]
def get_surf(self, subject: str, type: str, hemi: Literal['both', 'lh', 'rh']='both', *, merge: bool=False, nudge: bool=False) -> Union[tuple[tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]], tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]]], tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]]]:
"""Retrieve a subject's surface geometry from this (packed) dataset file.
Parameters
----------
subject : str
Subject identifier, as packed into this dataset.
type : str
Surface type, e.g. 'wm', 'pia', or 'fiducial' (averaged from
'wm' and 'pia').
hemi : str, optional
'lh', 'rh', or 'both'. Default 'both'.
merge : bool, optional
If True (and `hemi='both'`), stack both hemispheres into a
single (pts, polys) pair with right-hemisphere face indices
offset. Default False.
nudge : bool, optional
If True, shift each hemisphere so they don't overlap along the
x-axis. Default False.
Returns
-------
(pts, polys) or ((lpts, lpolys), (rpts, rpolys))
Vertex coordinates and triangle faces for the requested
surface. If `hemi='both'` and `merge=False`, a pair of
(pts, polys) tuples, one per hemisphere, is returned instead.
"""
pts: npt.NDArray[np.floating]
polys: npt.NDArray[np.integer]
if hemi == 'both':
left = self.get_surf(subject, type, "lh", nudge=nudge)
right = self.get_surf(subject, type, "rh", nudge=nudge)
if merge:
pts = np.vstack([left[0], right[0]])
polys = np.vstack([left[1], right[1]+len(left[0])])
return pts, polys
return left, right
try:
if type == 'fiducial':
wpts, polys = self.get_surf(subject, 'wm', hemi)
ppts, _ = self.get_surf(subject, 'pia', hemi)
return (wpts + ppts) / 2, polys
group = self.h5['subjects'][subject]['surfaces'][type][hemi]
pts, polys = group['pts'][:].copy(), group['polys'][:].copy()
if nudge:
if hemi == 'lh':
pts[:,0] -= pts[:,0].max()
else:
pts[:,0] -= pts[:,0].min()
return pts, polys
except (KeyError, TypeError):
raise IOError('Subject not found in package')
[docs]
def get_xfm(self, subject: str, xfmname: str) -> Transform:
"""Retrieve a subject's transform from this (packed) dataset file.
Parameters
----------
subject : str
Subject identifier, as packed into this dataset.
xfmname : str
Transform name, as packed into this dataset.
Returns
-------
Transform
The requested transform.
"""
try:
group: h5py.Group = self.h5['subjects'][subject]['transforms'][xfmname]
return Transform(group['xfm'][:], tuple(group['xfm'].attrs['shape']))
except (KeyError, TypeError):
raise IOError('Transform not found in package')
[docs]
def get_mask(self, subject: str, xfmname: str, maskname: str):
"""Retrieve a subject's voxel mask from this (packed) dataset file.
Parameters
----------
subject : str
Subject identifier, as packed into this dataset.
xfmname : str
Transform name the mask belongs to, as packed into this dataset.
maskname : str
Name of the mask to retrieve.
Returns
-------
h5py.Dataset
The requested mask array, as stored in the HDF5 file.
"""
try:
group: h5py.Group = self.h5['subjects'][subject]['transforms'][xfmname]['masks']
return group[maskname]
except (KeyError, TypeError):
raise IOError('Mask not found in package')
[docs]
def get_overlay(self, subject: str, type: str='rois', **kwargs) -> tempfile._TemporaryFileWrapper:
"""Retrieve a subject's overlay (currently only ROIs) from this
(packed) dataset file.
Parameters
----------
subject : str
Subject identifier, as packed into this dataset.
type : str, optional
Overlay type to retrieve. Only 'rois' is currently supported.
Default 'rois'.
Returns
-------
tempfile.NamedTemporaryFile
A temporary file containing the overlay's SVG data, seeked to
the start.
"""
try:
group: h5py.Group = self.h5['subjects'][subject]
if type == "rois":
tf = tempfile.NamedTemporaryFile(delete = False)
tf.write(group['rois'][0])
tf.close()
return tf
except (KeyError, TypeError):
raise IOError('Overlay not found in package')
raise TypeError('Unknown overlay type')
[docs]
def prepend(self, prefix):
"""Adds the given `prefix` to the name of every data object and returns
a new Dataset.
"""
ds = dict()
for name, data in self:
ds[prefix+name] = data
return Dataset(**ds)
DatasetLike = Union[Dataset, dict, str]
@overload
def normalize(data: Dataview) -> Dataview: ...
@overload
def normalize(data: Union[Dataset, dict, str]) -> Dataset: ...
@overload
def normalize(data: tuple) -> Union[Vertex, Volume]: ...
def normalize(data: Union[DatasetLike, Dataview, tuple]) -> Union[Dataset, Dataview, Vertex, Volume]:
if isinstance(data, (Dataset, Dataview)):
return data
elif isinstance(data, dict):
return Dataset(**data)
elif isinstance(data, str):
return Dataset.from_file(data)
elif isinstance(data, tuple):
return _vnorm(data)
raise TypeError('Unknown input type')
def _pack_subjs(h5: h5py.File, subjects: set[str]) -> None:
for subject in subjects:
rois = db.get_overlay(subject, modify_svg_file=False)
rnode = h5.require_dataset("/subjects/%s/rois"%subject, (1,),
dtype=h5py.special_dtype(vlen=str))
rnode[0] = rois.toxml(pretty=False)
surfaces = db.get_paths(subject)['surfs']
for surf in surfaces.keys():
for hemi in ("lh", "rh"):
pts, polys = db.get_surf(subject, surf, hemi)
group = "/subjects/%s/surfaces/%s/%s"%(subject, surf, hemi)
_hdf_write(h5, pts, "pts", group)
_hdf_write(h5, polys, "polys", group)
def _pack_xfms(h5: h5py.File, xfms: set[tuple[str, str]]) -> None:
for subj, xfmname in xfms:
xfm = db.get_xfm(subj, xfmname, 'coord')
group = "/subjects/%s/transforms/%s"%(subj, xfmname)
node = _hdf_write(h5, np.array(xfm.xfm), "xfm", group)
node.attrs['shape'] = xfm.shape
def _pack_masks(h5: h5py.File, masks: set[tuple[str, str, str]]) -> None:
for subj, xfm, maskname in masks:
mask = db.get_mask(subj, xfm, maskname)
group = "/subjects/%s/transforms/%s/masks"%(subj, xfm)
_hdf_write(h5, mask, maskname, group)