diff --git a/.gitignore b/.gitignore index 613aa48c..21a752d9 100644 --- a/.gitignore +++ b/.gitignore @@ -35,6 +35,7 @@ pip-log.txt # Unit test / coverage reports .coverage +coverage.* .tox nosetests.xml diff --git a/cortex/align.py b/cortex/align.py index 7a236db0..1940221f 100644 --- a/cortex/align.py +++ b/cortex/align.py @@ -5,6 +5,7 @@ import subprocess as sp import tempfile import warnings +from typing import Literal, Optional import numpy as np @@ -107,8 +108,17 @@ def fs_manual(subject, xfmname, **kwargs): return manual(subject, xfmname, **kwargs) -def manual(subject, xfmname, output_name="register.lta", wm_color="yellow", - pial_color="blue", wm_surface='white', noclean=False, reference=None, inspect_only=False): +def manual( + subject: str, + xfmname: str, + output_name: str = "register.lta", + wm_color: str = "yellow", + pial_color: str = "blue", + wm_surface: str = "white", + noclean: bool = False, + reference: Optional[str] = None, + inspect_only: bool = False, +) -> Optional[str]: """Open Freesurfer FreeView GUI for manually aligning/adjusting a functional volume to the cortical surface for `subject`. This creates a new transform called `xfmname`. The name of a nibabel-readable file (e.g. NIfTI) should be @@ -191,7 +201,9 @@ def manual(subject, xfmname, output_name="register.lta", wm_color="yellow", if reference is None: # Load load extant transform-relevant things - reference = sub_xfm.reference.get_filename() + if sub_xfm.reference is None: + raise ValueError('Cannot inspect reference-free transform') + reference = sub_xfm.reference_nifti.get_filename() _ = sub_xfm.to_freesurfer(os.path.join(cache, "register.dat"), subject) # Transform in freesurfer .dat format # Command for FreeView and run cmd = ("freeview -v $SUBJECTS_DIR/{sub}/mri/orig.mgz " @@ -347,15 +359,15 @@ def automatic_fsl( def automatic( - subject, - xfmname, - reference, - init="coreg", - epi_mask=False, - intermediate=None, - reference_contrast="t2", - noclean=False, -): + subject: str, + xfmname: str, + reference: str, + init: str = "coreg", + epi_mask: bool = False, + intermediate: Optional[str] = None, + reference_contrast: Literal['t1', 't2'] = "t2", + noclean: bool = False, +) -> Optional[str]: """Perform automatic alignment using Freesurfer's boundary-based registration. The `reference` image and resulting transform called `xfmname` will be automatically stored in the database. diff --git a/cortex/database.py b/cortex/database.py index 336471b1..cbcbc60d 100644 --- a/cortex/database.py +++ b/cortex/database.py @@ -1,6 +1,7 @@ """ Contains a singleton object `db` of type `Database` which allows easy access to surface files, anatomical images, and transforms that are stored in the pycortex filestore. """ +from __future__ import annotations import copy import functools import glob @@ -11,13 +12,32 @@ import tempfile import warnings from hashlib import sha1 +from typing import Literal, Union, Optional, TypedDict, TYPE_CHECKING, overload, cast import numpy as np +import numpy.typing as npt from . import options +from .xfm import Transform +if TYPE_CHECKING: + from cortex.dataset.dataset import Dataset + from cortex.dataset.views import Vertex + from cortex.svgoverlay import SVGOverlay default_filestore = options.config.get('basic', 'filestore') +class PathsType(TypedDict): + surfs: dict[str, dict[str, str]] + xfms: list[str] + xfmdir: str + anats: str + surfinfo: str + masks: str + rois: str + overlays: str + views: list[str] + surf2surf: str + def _memo(fn): @functools.wraps(fn) @@ -33,11 +53,11 @@ def memofn(self, *args, **kwargs): return memofn class SubjectDB: - def __init__(self, subj, filestore=default_filestore): + def __init__(self, subj: str, filestore: str = default_filestore): self.subject = subj self._warning = None - self._transforms = None - self._surfaces = None + self._transforms: Optional[XfmDB] = None + self._surfaces: Optional[SurfaceDB] = None self.filestore = filestore try: @@ -47,21 +67,21 @@ def __init__(self, subj, filestore=default_filestore): pass @property - def transforms(self): + def transforms(self) -> XfmDB: if self._transforms is not None: return self._transforms self._transforms = XfmDB(self.subject, filestore=self.filestore) return self._transforms @property - def surfaces(self): + def surfaces(self) -> SurfaceDB: if self._surfaces is not None: return self._surfaces self._surfaces = SurfaceDB(self.subject, filestore=self.filestore) return self._surfaces class SurfaceDB: - def __init__(self, subj, filestore=default_filestore): + def __init__(self, subj: str, filestore: str = default_filestore): self.subject = subj self.types = {} db = Database(filestore) @@ -74,31 +94,31 @@ def __repr__(self): def __dir__(self): return list(self.types.keys()) - def __getattr__(self, attr): + def __getattr__(self, attr: str): if attr in self.types: return self.types[attr] raise AttributeError(attr) class Surf: - def __init__(self, subject, surftype, filestore=default_filestore): + def __init__(self, subject: str, surftype: str, filestore: str = default_filestore): self.subject, self.surftype = subject, surftype self.db = Database(filestore) - def get(self, hemisphere="both"): + def get(self, hemisphere: Literal['lh', 'rh', 'both'] = "both"): return self.db.get_surf(self.subject, self.surftype, hemisphere) - def show(self, hemisphere="both"): + def show(self, hemisphere: Literal['lh', 'rh', 'both'] = "both"): from mayavi import mlab pts, polys = self.db.get_surf(self.subject, self.surftype, hemisphere, merge=True, nudge=True) return mlab.triangular_mesh(pts[:,0], pts[:,1], pts[:,2], polys) class XfmDB: - def __init__(self, subj, filestore=default_filestore): + def __init__(self, subj: str, filestore: str = default_filestore): self.subject = subj self.filestore = filestore - self.xfms = Database(self.filestore).get_paths(subj)['xfms'] + self.xfms: list[str] = Database(self.filestore).get_paths(subj)['xfms'] - def __getitem__(self, name): + def __getitem__(self, name: str) -> 'XfmSet': if name in self.xfms: return XfmSet(self.subject, name, filestore=self.filestore) raise AttributeError @@ -108,7 +128,7 @@ def __repr__(self): return f"Available transforms for {self.subject}:\n{xfms}" class XfmSet: - def __init__(self, subj, name, filestore=default_filestore): + def __init__(self, subj: str, name: str, filestore: str = default_filestore): self.subject = subj self.name = name jspath = os.path.join(filestore, subj, 'transforms', name, 'matrices.xfm') @@ -117,7 +137,7 @@ def __init__(self, subj, name, filestore=default_filestore): self.masks = MaskSet(subj, name, filestore=filestore) self.db = Database(filestore) - def __getattr__(self, attr): + def __getattr__(self, attr: str) -> Transform: if attr in self._jsdat: return self.db.get_xfm(self.subject, self.name, attr) raise AttributeError @@ -126,16 +146,16 @@ def __repr__(self): return "Types: {types}".format(types=", ".join(self._jsdat.keys())) class MaskSet: - def __init__(self, subj, name, filestore=default_filestore): + def __init__(self, subj: str, name: str, filestore: str = default_filestore): self.subject = subj self.xfmname = name maskform = Database(filestore).get_paths(subj)['masks'] maskpath = maskform.format(xfmname=name, type='*') - self._masks = {os.path.split(path)[1][5:-7]: path for path in glob.glob(maskpath)} + self._masks: dict[str, str] = {os.path.split(path)[1][5:-7]: path for path in glob.glob(maskpath)} - def __getitem__(self, item): + def __getitem__(self, item: str) -> npt.NDArray: import nibabel - return nibabel.load(self._masks[item]).get_fdata().T + return cast(nibabel.Nifti1Image, nibabel.load(self._masks[item])).get_fdata().T def __repr__(self): return "Masks: [{types}]".format(types=', '.join(self._masks.keys())) @@ -150,19 +170,23 @@ class Database: ---------- This database object dynamically generates handles to all subjects within the filestore. """ - def __init__(self, filestore=default_filestore): + def __init__(self, filestore: str=default_filestore): self.filestore = filestore - self._subjects = None - self.auxfile = None + self._subjects: Optional[dict[str, SubjectDB]] = None + # Side channel set by Dataset.from_file and cortex.webgl.show: a Dataset + # standing in for the filestore, so views can resolve surfaces and + # transforms out of the .hdf they were loaded from. + self.auxfile: Optional["Dataset"] = None - def __repr__(self): + def __repr__(self) -> str: subjs = "\n ".join(sorted(self.subjects.keys())) return """Pycortex database\n Subjects:\n {subjs}""".format(subjs=subjs) - def __getattr__(self, attr): + def __getattr__(self, attr: str): if attr in self.subjects: - if self.subjects[attr]._warning is not None: - warnings.warn(self.subjects[attr]._warning) + _warning = self.subjects[attr]._warning + if _warning is not None: + warnings.warn(_warning) return self.subjects[attr] else: raise AttributeError @@ -173,7 +197,7 @@ def __dir__(self): 'get_mri_surf2surf_matrix'] + list(self.subjects.keys()) @property - def subjects(self): + def subjects(self) -> dict[str, SubjectDB]: if self._subjects is not None: return self._subjects subjs = os.listdir(os.path.join(self.filestore)) @@ -187,7 +211,7 @@ def reload_subjects(self): self._subjects = None self.subjects - def get_anat(self, subject, type='raw', xfmname=None, recache=False, order=1, **kwargs): + def get_anat(self, subject: str, type: Literal['raw', 'brainmask', 'whitematter', 'voxelize'] ='raw', xfmname: Optional[str]=None, recache: bool=False, order: Literal[0, 1, 2, 3, 4, 5]=1, **kwargs): """Return anatomical information from the filestore. Anatomical information is defined as any volume-space anatomical information pertaining to the subject, such as T1 image, white matter masks, etc. Volumes not found in the database will be automatically generated. @@ -213,12 +237,13 @@ def get_anat(self, subject, type='raw', xfmname=None, recache=False, order=1, ** anatfile = anatform.format(type=type, opts=opts, ext="nii.gz") if not os.path.exists(anatfile) or recache: + # TODO: does `raw` enter this block? print("Generating %s anatomical..."%type) from . import anat getattr(anat, type)(anatfile, subject, **kwargs) import nibabel - anatnib = nibabel.load(anatfile) + anatnib = cast(nibabel.Nifti1Image, nibabel.load(anatfile)) if xfmname is None: return anatnib @@ -226,7 +251,7 @@ def get_anat(self, subject, type='raw', xfmname=None, recache=False, order=1, ** from . import volume return volume.anat2epispace(anatnib.get_fdata().T.astype(float), subject, xfmname, order=order) - def get_surfinfo(self, subject, type="curvature", recache=False, **kwargs): + def get_surfinfo(self, subject: str, type: str="curvature", recache: bool=False, **kwargs) -> Vertex: """Return auxiliary surface information from the filestore. Surface info is defined as anatomical information specific to a subject in surface space. A Vertex class will be returned as necessary. Info not found in the filestore will be automatically generated. @@ -256,7 +281,7 @@ def get_surfinfo(self, subject, type="curvature", recache=False, **kwargs): if len(kwargs) > 0: opts = "[%s]"%','.join(["%s=%s"%i for i in kwargs.items()]) try: - self.auxfile.get_surf(subject, "fiducial") + self.auxfile.get_surf(subject, "fiducial") # type: ignore[union-attr] surfifile = os.path.join(self.get_cache(subject),"%s%s.npz"%(type, opts)) except (AttributeError, IOError): surfiform = self.get_paths(subject)['surfinfo'] @@ -278,8 +303,8 @@ def get_surfinfo(self, subject, type="curvature", recache=False, **kwargs): return Vertex(verts, subject) return npz - def get_mri_surf2surf_matrix(self, subject, surface_type, hemi='both', - fs_subj=None, target_subj='fsaverage', + def get_mri_surf2surf_matrix(self, subject: str, surface_type: str, hemi: Literal['lh', 'rh', 'both']='both', + fs_subj: Optional[str]=None, target_subj: str='fsaverage', **kwargs): """Get matrix generated by surf2surf to map one subject's surface to another's @@ -330,6 +355,7 @@ def get_mri_surf2surf_matrix(self, subject, surface_type, hemi='both', if not os.path.exists(fdir): print("Creating surf2surf directory for subject %s"%(subject)) os.makedirs(fdir) + hemis: list[Literal['lh', 'rh']] if hemi == 'both': hemis = ['lh', 'rh'] else: @@ -345,7 +371,7 @@ def get_mri_surf2surf_matrix(self, subject, surface_type, hemi='both', save_sparse_array(fpath, tmp, h, mode='a') return mats - def get_overlay(self, subject, overlay_file=None, **kwargs): + def get_overlay(self, subject: str, overlay_file: Optional[str]=None, **kwargs) -> SVGOverlay: from . import svgoverlay pts, polys = self.get_surf(subject, "flat", merge=True, nudge=True) @@ -373,7 +399,7 @@ def get_overlay(self, subject, overlay_file=None, **kwargs): overlay_file = paths['overlays'] return svgoverlay.get_overlay(subject, overlay_file, pts, polys, **kwargs) - def save_xfm(self, subject, name, xfm, xfmtype="magnet", reference=None): + def save_xfm(self, subject: str, name: str, xfm: npt.NDArray[np.floating], xfmtype: str="magnet", reference: Optional[str]=None): """ Load a transform into the surface database. If the transform exists already, update it If it does not exist, copy the reference epi into the filestore and insert. @@ -408,7 +434,7 @@ def save_xfm(self, subject, name, xfm, xfmtype="magnet", reference=None): if reference is None: raise ValueError("Please specify a reference") fpath = os.path.join(path, "reference.nii.gz") - nib = nibabel.load(reference) + nib = cast(nibabel.Nifti1Image, nibabel.load(reference)) data = nib.get_fdata() if len(data.shape) > 3: import warnings @@ -419,7 +445,7 @@ def save_xfm(self, subject, name, xfm, xfmtype="magnet", reference=None): jsdict = dict() - nib = nibabel.load(os.path.join(path, "reference.nii.gz")) + nib = cast(nibabel.Nifti1Image, nibabel.load(os.path.join(path, "reference.nii.gz"))) if xfmtype == "magnet": jsdict['magnet'] = np.array(xfm).tolist() jsdict['coord'] = np.dot(np.linalg.inv(nib.affine), xfm).tolist() @@ -434,7 +460,7 @@ def save_xfm(self, subject, name, xfm, xfmtype="magnet", reference=None): with open(fname, "w") as fp: json.dump(jsdict, fp, sort_keys=True, indent=4) - def get_xfm(self, subject, name, xfmtype="coord"): + def get_xfm(self, subject: str, name: str, xfmtype: str="coord") -> Transform: """Retrieves a transform from the filestore Parameters @@ -446,10 +472,9 @@ def get_xfm(self, subject, name, xfmtype="coord"): xfmtype : str, optional Type of transform to return. Defaults to coord. """ - from .xfm import Transform if xfmtype == 'coord': try: - return self.auxfile.get_xfm(subject, name) + return self.auxfile.get_xfm(subject, name) # type: ignore[union-attr] except (AttributeError, IOError): pass @@ -463,8 +488,26 @@ def get_xfm(self, subject, name, xfmtype="coord"): xfmdict = json.load(f) return Transform(xfmdict[xfmtype], reference) + # TODO: forcing '*' WILL cause issues. Look for all instances of merge=True ! + @overload + def get_surf(self, subject: str, type: str, hemisphere: 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, hemisphere: 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, hemisphere: Literal['lh', 'rh'], merge: bool=False, nudge: bool=False) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.integer]]: + ... + + # Fallthrough case for the recursive call + @overload + def get_surf(self, subject: str, type: str, hemisphere: Literal['lh', 'rh', 'both']="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]]]: + ... + @_memo - def get_surf(self, subject, type, hemisphere="both", merge=False, nudge=False): + def get_surf(self, subject: str, type: str, hemisphere: Literal['lh', 'rh', 'both']="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]]]: '''Return the surface pair for the given subject, surface type, and hemisphere. Parameters @@ -490,14 +533,14 @@ def get_surf(self, subject, type, hemisphere="both", merge=False, nudge=False): For single hemisphere ''' try: - return self.auxfile.get_surf(subject, type, hemisphere, merge=merge, nudge=nudge) + return self.auxfile.get_surf(subject, type, hemisphere, merge=merge, nudge=nudge) # type: ignore[union-attr] except (AttributeError, IOError): pass files = self.get_paths(subject)['surfs'] if hemisphere.lower() == "both": - left, right = [ self.get_surf(subject, type, hemisphere=h) for h in ["lh", "rh"]] + left, right = [ self.get_surf(subject, type, hemisphere=cast(Literal['lh', 'rh'], h)) for h in ["lh", "rh"]] if type != "fiducial" and nudge: left[0][:,0] -= left[0].max(0)[0] right[0][:,0] -= right[0].min(0)[0] @@ -508,7 +551,8 @@ def get_surf(self, subject, type, hemisphere="both", merge=False, nudge=False): return pts, polys return left, right - elif hemisphere.lower() in ("lh", "left"): + hemi: Literal['lh', 'rh'] + if hemisphere.lower() in ("lh", "left"): hemi = "lh" elif hemisphere.lower() in ("rh", "right"): hemi = "rh" @@ -527,7 +571,7 @@ def get_surf(self, subject, type, hemisphere="both", merge=False, nudge=False): except KeyError: raise IOError(f"Surface type '{type}' not found for {hemi} hemisphere of subject '{subject}'") - def save_mask(self, subject, xfmname, type, mask): + def save_mask(self, subject: str, xfmname: str, type: str, mask: npt.NDArray[np.bool]) -> None: fname = self.get_paths(subject)['masks'].format(xfmname=xfmname, type=type) if os.path.exists(fname): raise IOError('Refusing to overwrite existing mask') @@ -536,23 +580,23 @@ def save_mask(self, subject, xfmname, type, mask): xfm = self.get_xfm(subject, xfmname) if xfm.shape != mask.shape: raise ValueError("Invalid mask shape: must match shape of reference image") - affine = xfm.reference.affine + affine = xfm.reference_nifti.affine nib = nibabel.Nifti1Image(mask.astype(np.uint8).T, affine) nib.to_filename(fname) - def get_mask(self, subject, xfmname, type='thick'): + def get_mask(self, subject: str, xfmname: str, type: str='thick') -> npt.NDArray[np.bool]: if hasattr(type, 'decode'): type = type.decode('utf8') try: - self.auxfile.get_mask(subject, xfmname, type) + self.auxfile.get_mask(subject, xfmname, type) # type: ignore[union-attr] except (AttributeError, IOError): pass fname = self.get_paths(subject)['masks'].format(xfmname=xfmname, type=type) try: import nibabel - nib = nibabel.load(fname) + nib = cast(nibabel.Nifti1Image, nibabel.load(fname)) return nib.get_fdata().T != 0 except IOError: print('Mask not found, generating...') @@ -561,7 +605,7 @@ def get_mask(self, subject, xfmname, type='thick'): self.save_mask(subject, xfmname, type, mask) return mask - def get_shared_voxels(self, subject, xfmname, hemi="both", merge=True, use_astar=True, recache=False): + def get_shared_voxels(self, subject: str, xfmname: str, hemi: Literal['lh', 'rh', 'both']="both", merge: bool=True, use_astar: bool=True, recache: bool=False): """Get an array indicating which vertices are inappropriately mapped to the same voxel. For a given transform and surface, returns an array containing a list of vertices which @@ -571,7 +615,7 @@ def get_shared_voxels(self, subject, xfmname, hemi="both", merge=True, use_astar """ # Test for packed subjects try: - voxels = self.auxfile.get_shared_voxels(subject, xfmname, hemi=hemi, merge=merge, use_astar=use_astar) + voxels = self.auxfile.get_shared_voxels(subject, xfmname, hemi=hemi, merge=merge, use_astar=use_astar) # type: ignore[union-attr] return voxels except (AttributeError, IOError): pass @@ -587,7 +631,7 @@ def get_shared_voxels(self, subject, xfmname, hemi="both", merge=True, use_astar voxels = np.load(shared_voxel_file) return voxels - def get_coords(self, subject, xfmname, hemisphere="both", magnet=None): + def get_coords(self, subject: str, xfmname: str, hemisphere: Literal['lh', 'rh', 'both']="both", magnet: Optional[npt.NDArray]=None) -> list[npt.NDArray[np.floating]]: """Calculate the coordinates of each vertex in the epi space by transforming the fiducial to the coordinate space Parameters @@ -608,22 +652,27 @@ def get_coords(self, subject, xfmname, hemisphere="both", magnet=None): xfm = self.get_xfm(subject, xfmname, xfmtype="magnet") xfm = np.linalg.inv(magnet) * xfm - coords = [] + coords: list[npt.NDArray[np.floating]] = [] vtkTmp = self.get_surf(subject, "fiducial", hemisphere=hemisphere, nudge=False) if not isinstance(vtkTmp,(tuple,list)): vtkTmp = [vtkTmp] + pts: npt.NDArray[np.floating] + polys: npt.NDArray[np.integer] for pts, polys in vtkTmp: wpts = np.vstack([pts.T, np.ones(len(pts))]) coords.append(np.dot(xfm.xfm, wpts)[:3].round().astype(int).T) return coords - def get_cache(self, subject): + def get_cache(self, subject: str) -> str: try: - self.auxfile.get_surf(subject, "fiducial") + self.auxfile.get_surf(subject, "fiducial") # type: ignore[union-attr] #generate the hashed name of the filename and subject as the directory name import hashlib - hashname = "pycx_%s"%hashlib.md5(self.auxfile.h5.filename).hexdigest()[-8:] + # md5 needs bytes: passing the str filename raised an uncaught TypeError, + # so this path could never have completed. + filename = self.auxfile.h5.filename # type: ignore[union-attr] + hashname = "pycx_%s"%hashlib.md5(filename.encode()).hexdigest()[-8:] cachedir = os.path.join(tempfile.gettempdir(), hashname, subject) except (AttributeError, IOError): try: @@ -638,7 +687,7 @@ def get_cache(self, subject): os.makedirs(cachedir) return cachedir - def clear_cache(self, subject, clear_all_caches=True): + def clear_cache(self, subject: str, clear_all_caches: bool=True) -> None: """Clears config-specified and default file caches for a subject. """ @@ -659,15 +708,16 @@ def clear_cache(self, subject, clear_all_caches=True): shutil.rmtree(default_cachedir) os.makedirs(default_cachedir) - def get_paths(self, subject): + def get_paths(self, subject: str) -> PathsType: """Get a dictionary with a list of all candidate filenames for associated data, such as roi overlays, flatmap caches, and ctm caches. """ surfpath = os.path.join(self.filestore, subject, "surfaces") - if self.subjects[subject]._warning is not None: - warnings.warn(self.subjects[subject]._warning) + _warn = self.subjects[subject]._warning + if _warn is not None: + warnings.warn(_warn) - surfs = dict() + surfs: dict[str, dict[str, str]] = dict() for surf in os.listdir(surfpath): ssurf = os.path.splitext(surf)[0].split('_') name = '_'.join(ssurf[:-1]) @@ -682,7 +732,7 @@ def get_paths(self, subject): os.makedirs(viewsdir) views = os.listdir(viewsdir) - filenames = dict( + filenames = PathsType( surfs=surfs, xfms=sorted(os.listdir(os.path.join(self.filestore, subject, "transforms"))), xfmdir=os.path.join(self.filestore, subject, "transforms", "{xfmname}", "matrices.xfm"), @@ -697,7 +747,7 @@ def get_paths(self, subject): return filenames - def make_subj(self, subject): + def make_subj(self, subject: str) -> None: if os.path.exists(os.path.join(self.filestore, subject)): if input("Are you sure you want to overwrite this existing subject?\n" "This will delete all files for this subject in the filestore, " @@ -713,7 +763,7 @@ def make_subj(self, subject): except OSError: print("Error making directory %s"%path) - def save_view(self,vw,subject,name,is_overwrite=False): + def save_view(self,vw,subject: str,name: str,is_overwrite: bool=False) -> None: """Set the view for an open webshow instance from a saved view Sets the view in a currently-open cortex.webshow instance (with handle `vw`) @@ -765,7 +815,7 @@ def get_view(self,vw,subject,name): view = json.load(fp) vw._set_view(**view) - def get_mnixfm(self, subject, xfm, template=None): + def get_mnixfm(self, subject: str, xfm: str, template: Optional[str]=None) -> npt.NDArray[np.floating]: """Get transform from the space specified by `xfm` to MNI space. Parameters diff --git a/cortex/formats.pyi b/cortex/formats.pyi new file mode 100644 index 00000000..1c69cdff --- /dev/null +++ b/cortex/formats.pyi @@ -0,0 +1,55 @@ +import os +from typing import Any, Literal, overload + +import numpy as _py_np +import numpy.typing as npt + +_Path = str | os.PathLike[str] + +PY3: bool + +def read(globname: str) -> tuple[npt.NDArray[_py_np.floating], npt.NDArray[_py_np.integer]]: ... + +# `polys` is float64 rather than integer for a file declaring zero faces, since +# `np.array([])` is float64. Same caveat on `read_obj` below. +def read_off(filename: _Path) -> tuple[npt.NDArray[_py_np.floating], npt.NDArray[_py_np.integer]]: ... + +# dtypes are whatever was stored in the archive. +def read_npz(filename: _Path) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]: ... +def read_gii(filename: _Path) -> tuple[npt.NDArray[_py_np.floating], npt.NDArray[_py_np.integer]]: ... +def read_stl(filename: _Path) -> tuple[npt.NDArray[_py_np.floating], npt.NDArray[_py_np.uint32]]: ... + +# 2-tuple only when both flags are literal `False`; unsound for non-literal `False, False`. +@overload +def read_obj( # type: ignore[overload-overlap] # overlaps the bool fallback below + filename: _Path, norm: Literal[False] = False, uv: Literal[False] = False +) -> tuple[npt.NDArray[_py_np.floating], npt.NDArray[_py_np.integer]]: ... +@overload +def read_obj( + filename: _Path, norm: bool = False, uv: bool = False +) -> tuple[ + npt.NDArray[_py_np.floating], + npt.NDArray[_py_np.integer], + list[list[float]] | None, + list[list[float]] | None, +]: ... + +def read_vtk(filename: _Path) -> tuple[npt.NDArray[_py_np.float64], npt.NDArray[_py_np.uint32]]: ... + +# The writers require real arrays, not just ArrayLike: they read `polys.dtype` +# and `pts.astype`, and index with `pts[polys]`. +def write_vtk( + filename: _Path, + pts: npt.NDArray[Any], + polys: npt.NDArray[Any], + norms: npt.NDArray[Any] | None = None, +) -> None: ... +def write_off(filename: _Path, pts: npt.NDArray[Any], polys: npt.NDArray[Any]) -> None: ... +def write_stl(filename: _Path, pts: npt.NDArray[Any], polys: npt.NDArray[Any]) -> None: ... +def write_gii(filename: _Path, pts: npt.NDArray[Any], polys: npt.NDArray[Any]) -> None: ... +def write_obj( + filename: _Path, + pts: npt.NDArray[Any], + polys: npt.NDArray[Any], + colors: npt.NDArray[Any] | None = None, +) -> None: ... diff --git a/cortex/formats.pyx b/cortex/formats.pyx index 7fad1170..ca84b13a 100644 --- a/cortex/formats.pyx +++ b/cortex/formats.pyx @@ -8,6 +8,8 @@ from collections import OrderedDict cimport cython cimport numpy as np +import numpy as _py_np +import numpy.typing as npt from libc.string cimport strtok from libc.stdlib cimport atoi, atof @@ -16,7 +18,7 @@ np.import_array() PY3 = sys.version_info[0] > 3 -def read(globname): +def read(globname: str) -> tuple[npt.NDArray[_py_np.floating], npt.NDArray[_py_np.integer]]: readers = OrderedDict([('gii', read_gii), ('npz', read_npz), ('vtk', read_vtk), ('off', read_off), ('stl', read_stl)]) for ext, func in readers.items(): try: diff --git a/cortex/volume.py b/cortex/volume.py index ca8c78eb..e4e0677d 100644 --- a/cortex/volume.py +++ b/cortex/volume.py @@ -1,13 +1,17 @@ """Contains functions for working with volume data """ import os +from typing import Literal, Optional, TypeVar, Union import numpy as np +import numpy.typing as npt from . import dataset from .database import db from .xfm import Transform -def unmask(mask, data): +DType = TypeVar('DType', bound=np.generic) +# TODO: MaskedArray typing might require newer numpy versions. If so, drop the generic typing. +def unmask(mask: np.ndarray[tuple[int, int, int], np.dtype[np.bool_]], data: npt.NDArray[DType]) -> Union[np.ma.MaskedArray[tuple[int, ...], np.dtype[DType]], npt.NDArray[DType], npt.NDArray[np.uint8]]: """unmask(mask, data) Unmask the data, assuming it's been masked. Creates a volume @@ -54,15 +58,15 @@ def unmask(mask, data): return output.squeeze() -def detrend_median(data, kernel=15): +def detrend_median(data: npt.NDArray, kernel: int=15) -> npt.NDArray: from scipy.signal import medfilt lowfreq = medfilt(data, [1, kernel, kernel]) return data - lowfreq -def detrend_gradient(data, diff=3): +def detrend_gradient(data: npt.ArrayLike, diff: int=3) -> npt.NDArray: return (np.array(np.gradient(data, 1, diff, diff))**2).sum(0) -def detrend_poly(data, polyorder = 10, mask=None): +def detrend_poly(data: npt.NDArray, polyorder: int = 10, mask: Optional[npt.NDArray] = None) -> npt.NDArray: from scipy.special import legendre polys = [legendre(i) for i in range(polyorder)] s = data.shape @@ -84,7 +88,7 @@ def detrend_poly(data, polyorder = 10, mask=None): else: return detrended.reshape(*s) -def mosaic(data, dim=0, show=True, **kwargs): +def mosaic(data: npt.NDArray[DType], dim: int=0, show: bool=True, **kwargs) -> tuple[Union[npt.NDArray[DType], npt.NDArray[np.uint8]], tuple[int, int]]: """ Turns volume data into a mosaic, useful for quickly viewing volumetric data with radiological convention (left side of figure is right side of subject). @@ -114,7 +118,7 @@ def mosaic(data, dim=0, show=True, **kwargs): else: output = (np.nan*np.ones(shape)).astype(data.dtype) - sl = [slice(None), slice(None), slice(None)] + sl: list[Union[int, slice]] = [slice(None), slice(None), slice(None)] for h in range(ntall): for w in range(nwide): sl[dim] = h*nwide+w @@ -223,7 +227,7 @@ def epi2anatspace(volumedata, order=1): anatspace : ndarray The ND array of the anatomy space data """ - from scipy.ndimage.interpolation import affine_transform + from scipy.ndimage import affine_transform ds = dataset.normalize(volumedata) volumedata = ds#.data @@ -239,7 +243,7 @@ def epi2anatspace(volumedata, order=1): offset=transpart, output_shape=anat.shape[::-1], cval=np.nan, order=order).T -def anat2epispace(anatdata, subject, xfmname, order=1): +def anat2epispace(anatdata: npt.NDArray, subject: str, xfmname: str, order: Literal[0, 1, 2, 3, 4, 5]=1) -> npt.NDArray: """Resamples data from anatomical space into epi space Parameters @@ -258,7 +262,7 @@ def anat2epispace(anatdata, subject, xfmname, order=1): epidata : ndarray data in EPI space """ - from scipy.ndimage.interpolation import affine_transform + from scipy.ndimage import affine_transform anatref = db.get_anat(subject) target = db.get_xfm(subject, xfmname, "coord") diff --git a/cortex/xfm.py b/cortex/xfm.py index de023ce5..f00480a8 100644 --- a/cortex/xfm.py +++ b/cortex/xfm.py @@ -1,36 +1,53 @@ """Affine transformation class """ import os +from typing import Optional, Union, cast, TYPE_CHECKING import numpy as np +import numpy.typing as npt import subprocess +if TYPE_CHECKING: + import nibabel + class Transform: ''' A standard affine transform. Typically holds a transform from anatomical magnet space to epi file space. ''' - def __init__(self, xfm, reference): + shape: tuple[int, int, int] + reference: Optional[Union[str, "nibabel.Nifti1Image", npt.NDArray]] + + def __init__(self, xfm: npt.NDArray, reference: Union[str, tuple[int, int, int], npt.NDArray, "nibabel.Nifti1Image"]): self.xfm = xfm self.reference = None if isinstance(reference, str): import nibabel try: - self.reference = nibabel.load(reference) - self.shape = self.reference.shape[:3][::-1] + self.reference = cast(nibabel.Nifti1Image, nibabel.load(reference)) + self.shape = self.reference.shape[:3][::-1] # type: ignore except IOError: self.reference = reference elif isinstance(reference, tuple): self.shape = reference else: self.reference = reference - self.shape = self.reference.shape[:3][::-1] + self.shape = self.reference.shape[:3][::-1] # type: ignore + + @property + def reference_nifti(self) -> "nibabel.Nifti1Image": + """The reference as a loaded nifti image, for callers that need its + affine/header. Raises if the reference is absent or never loaded.""" + import nibabel + if not isinstance(self.reference, nibabel.Nifti1Image): + raise ValueError('Transform has no loaded reference image') + return self.reference - def __call__(self, pts): + def __call__(self, pts: npt.NDArray) -> npt.NDArray: return np.dot(self.xfm, np.hstack([pts, np.ones((len(pts),1))]).T)[:3].T @property - def inv(self): + def inv(self) -> "Transform": ref = self.reference if ref is None: ref = self.shape @@ -172,6 +189,7 @@ def to_fsl(self, anat_nii, direction='func>anat'): # transforms. Thus the anatomical file is the "infile" in FSL-speak. infile = anat_nii + inIm: nibabel.Nifti1Image try: inIm = nibabel.load(infile) except AttributeError: @@ -260,7 +278,7 @@ def from_freesurfer(cls, fs_register, func_nii, subject, freesurfer_subject_dir= # Read vox2ras transform for the anatomical volume try: cmd = ('mri_info', '--vox2ras', anat_mgz) - L = decode(subprocess.check_output(cmd)).splitlines() + L = subprocess.check_output(cmd).decode().splitlines() anat_vox2ras = np.array([[np.float64(s) for s in ll.split() if s] for ll in L]) except OSError: print ("Error occurred while executing:\n{}".format(' '.join(cmd))) @@ -283,7 +301,7 @@ def from_freesurfer(cls, fs_register, func_nii, subject, freesurfer_subject_dir= return cls(coord, refIm) - def to_freesurfer(self, fs_register, subject, freesurfer_subject_dir=None): + def to_freesurfer(self, fs_register: str, subject: str, freesurfer_subject_dir: Optional[str]=None): """Converts a pycortex transform to a FreeSurfer transform. Converts a transform stored in pycortex xfm object to the FreeSurfer format @@ -319,10 +337,10 @@ def to_freesurfer(self, fs_register, subject, freesurfer_subject_dir=None): anat_tkrvox2ras = _vox2ras_tkr(anat.get_filename()) # Read tkvox2ras transform for the functional volume - func_tkrvox2ras = _vox2ras_tkr(self.reference.get_filename()) + func_tkrvox2ras = _vox2ras_tkr(self.reference_nifti.get_filename()) # Read voxel resolution of the functional volume - func_voxres = self.reference.header.get_zooms() + func_voxres = self.reference_nifti.header.get_zooms() # Calculate FreeSurfer transform fs_anat2func = np.dot(func_tkrvox2ras, np.dot(self.xfm, np.dot(anat_vox2ras, inv(anat_tkrvox2ras)))) @@ -340,13 +358,6 @@ def to_freesurfer(self, fs_register, subject, freesurfer_subject_dir=None): return fs_anat2func -def isstr(obj): - """Check for stringy-ness in python 2.7 or 3""" - try: - return isinstance(obj, basestring) - except NameError: - return isinstance(obj, str) - def decode(obj): if isinstance(obj, bytes): obj = obj.decode() @@ -365,7 +376,7 @@ def _vox2ras_tkr(image): output affine""" try: cmd = ('mri_info', '--vox2ras-tkr', image) - L = decode(subprocess.check_output(cmd)).splitlines() + L = subprocess.check_output(cmd).decode().splitlines() # Skip headers/additional information. Example output of # mri_info --vox2ras-tkr # diff --git a/pyproject.toml b/pyproject.toml index 6acd93aa..655cb7a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ dynamic = ["authors", "license", "readme", "version", "classifiers", "dependenci [dependency-groups] dev = [ {include-group = "test"}, - "mypy", + {include-group = "types"}, ] test = [ "pycortex[headless]", @@ -22,6 +22,11 @@ test = [ "pytest-cov", "pytest-timeout", ] +types = [ + "mypy", + "scipy-stubs", +] + [project.optional-dependencies] headless = [ @@ -40,4 +45,15 @@ ignore-words-list = 'nd,acount,anormal,fpt,coo,transpart,FO,lins' [tool.mypy] allow_redefinition = true -disable_error_code = "import-untyped" +disable_error_code = "import-untyped" # TODO: narrow to specific packages + +[[tool.mypy.overrides]] +module = [ + "bpy", + "bpy.ops", + "mayavi", + "progressbar", + "scikits.sparse.cholmod", + "tvtk.api", +] +ignore_missing_imports = true diff --git a/pytest.ini b/pytest.ini index 25191cf6..f3475078 100644 --- a/pytest.ini +++ b/pytest.ini @@ -4,6 +4,8 @@ testpaths = addopts = -r a -v + --cov=. + --cov-report xml # Per-test timeout (in seconds) so a single hung headless browser session # does not consume the entire CI budget. Individual tests can override with # @pytest.mark.timeout(N). Requires the optional ``pytest-timeout``