ENH add types in quickflat and vertex mappers

cortex/mapper/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations # needed for `type`?
+
 import os
 
 import numpy as np
@@ -7,11 +9,11 @@
 from .utils import nanproject, vol2surf
 
 
-def get_mapper(subject, xfmname, type='nearest', recache=False, **kwargs):
+def get_mapper(subject: str, xfmname: str, maptype: str='nearest', recache: bool=False, **kwargs) -> Mapper:
     from ..database import db
     from . import point, patch, line
 
-    mapcls = dict(
+    mapcls: dict[str, type[Mapper]] = dict(
         nearest=point.PointNN,
         trilinear=point.PointTrilin,
         gaussian=point.PointGauss,
@@ -22,7 +24,7 @@ def get_mapper(subject, xfmname, type='nearest', recache=False, **kwargs):
         line_nearest=line.LineNN,
         line_trilinear=line.LineTrilin,
         line_lanczos=line.LineLanczos)
-    Map = mapcls[type]
+    Map = mapcls[maptype]
     ptype = Map.__name__.lower()
     kwds ='_'.join(['%s%s'%(k,str(v)) for k, v in list(kwargs.items())])
     if len(kwds) > 0:

cortex/mapper/line.py

@@ -26,7 +26,7 @@ def _getmask(cls, pia, wm, polys, shape, npts=64, mp=True, **kwargs):
         #vidx = np.nonzero(valid)[0]
         mapper = sparse.csr_matrix((len(pia), np.prod(shape)))
         for t in np.linspace(0, 1, npts+2)[1:-1]:
-            i, j, data = cls.sampler(pia*t + wm*(1-t), shape)
+            i, j, data = cls.sampler(pia*t + wm*(1-t), shape, **kwargs)
             mapper = mapper + sparse.csr_matrix((data / npts, (i, j)), shape=mapper.shape)
         return mapper
 

cortex/mapper/mapper.py

@@ -1,14 +1,27 @@
+import abc
+from typing import Union, overload
+
 import numpy as np
+import numpy.typing as npt
 from scipy import sparse
 
 from .. import dataset
 
+import sys
+if sys.version_info < (3, 11):
+    from typing_extensions import Self
+else:
+    from typing import Self
+
 import warnings
+
 warnings.simplefilter('ignore', sparse.SparseEfficiencyWarning)
 
-class Mapper:
+MapperShape = Union[npt.NDArray[np.integer], tuple[int, int, int]]
+
+class Mapper(abc.ABC):
     '''Maps data from epi volume onto surface using various projections'''
-    def __init__(self, left, right, shape, subject, xfmname):
+    def __init__(self, left: sparse.csr_matrix, right: sparse.csr_matrix, shape: MapperShape, subject: str, xfmname: str):
         self.idxmap = None
         self.masks = [left, right]
         self.nverts = left.shape[0] + right.shape[0]
@@ -17,7 +30,7 @@ def __init__(self, left, right, shape, subject, xfmname):
         self.xfmname = xfmname
 
     @classmethod
-    def from_cache(cls, cachefile, subject, xfmname):
+    def from_cache(cls, cachefile: str, subject: str, xfmname: str) -> Self:
         npz = np.load(cachefile)
         left = (npz['left_data'], npz['left_indices'], npz['left_indptr'])
         right = (npz['right_data'], npz['right_indices'], npz['right_indptr'])
@@ -26,20 +39,26 @@ def from_cache(cls, cachefile, subject, xfmname):
         return cls(lsparse, rsparse, npz['shape'], subject, xfmname)
 
     @property
-    def mask(self):
+    def mask(self) -> npt.NDArray[np.bool_]:
         mask = np.array(self.masks[0].sum(0) + self.masks[1].sum(0))
         return (mask.squeeze() != 0).reshape(self.shape)
 
     @property
-    def hemimasks(self):
+    def hemimasks(self) -> list[npt.NDArray[np.bool_]]:
         func = lambda m: (np.array(m.sum(0)).squeeze() != 0).reshape(self.shape)
         return [func(x) for x in self.masks]
 
     def __repr__(self):
         ptype = self.__class__.__name__
         return '<%s mapper with %d vertices>'%(ptype, self.nverts)
 
-    def __call__(self, data):
+    @overload
+    def __call__(self, data: Union[dataset.Volume, tuple]) -> dataset.Vertex: ...
+
+    @overload
+    def __call__(self, data: dataset.Vertex) -> tuple[npt.NDArray, npt.NDArray]: ...
+
+    def __call__(self, data: Union[dataset.Vertex, dataset.Volume, tuple]) -> Union[tuple[npt.NDArray, npt.NDArray], dataset.Vertex]:
         if isinstance(data, tuple):
             data = dataset.Volume(*data)
 
@@ -61,17 +80,23 @@ def __call__(self, data):
         volume.shape = len(volume), -1
         volume = volume.T
 
-        mapped = []
+        mapped: list[npt.NDArray] = []
         for mask in self.masks:
-            mapped.append(np.array(mask * volume).T)
+            mapped.append(np.array(mask * volume).T) # change to @ matmul
 
         if self.idxmap is not None:
             mapped[0] = mapped[0][:, self.idxmap[0]]
             mapped[1] = mapped[1][:, self.idxmap[1]]
 
         return dataset.Vertex(np.hstack(mapped).squeeze(), data.subject)
 
-    def backwards(self, vertexdata):
+    @overload
+    def backwards(self, vertexdata: dataset.Vertex) -> dataset.Volume: ...
+
+    @overload
+    def backwards(self, vertexdata: npt.NDArray) -> npt.NDArray: ...
+
+    def backwards(self, vertexdata: Union[dataset.Vertex, npt.NDArray]) -> Union[dataset.Volume, npt.NDArray]:
         '''Projects vertex data back into volume space.
 
         Parameters
@@ -81,8 +106,7 @@ def backwards(self, vertexdata):
             If Vertex object is provided, a Volume object is returned
             If an array is provided, an array is returned
         '''
-        Vert2Vol = isinstance(vertexdata, dataset.Vertex)
-        if Vert2Vol:
+        if isinstance(vertexdata, dataset.Vertex):
             to_map = vertexdata.data
         else:
             to_map = vertexdata
@@ -91,8 +115,8 @@ def backwards(self, vertexdata):
         # dot the vertex data with the stacked mappers
         partial_vertex = bothmappers.T.dot(to_map)
         # solve the inverse mapping problem
-        voxeldata = self._get_backmapper().solve(partial_vertex).reshape(self.shape)
-        if Vert2Vol:
+        voxeldata: npt.NDArray = self._get_backmapper().solve(partial_vertex).reshape(self.shape)
+        if isinstance(vertexdata, dataset.Vertex):
             # construct a volume object with the new data
             return dataset.Volume(voxeldata, self.subject, self.xfmname)
         else:
@@ -112,10 +136,10 @@ def _get_backmapper(self):
         return self._backmapper
 
     @classmethod
-    def _cache(cls, filename, subject, xfmname, **kwargs):
+    def _cache(cls, filename: str, subject: str, xfmname: str, **kwargs) -> Self:
         print('Caching mapper...')
         from ..database import db
-        masks = []
+        masks: list[sparse.csr_matrix] = []
         xfm = db.get_xfm(subject, xfmname, xfmtype='coord')
         fid = db.get_surf(subject, 'fiducial', merge=False, nudge=False)
 
@@ -130,7 +154,19 @@ def _cache(cls, filename, subject, xfmname, **kwargs):
         _savecache(filename, masks[0], masks[1], xfm.shape)
         return cls(masks[0], masks[1], xfm.shape, subject, xfmname)
 
-def _savecache(filename, left, right, shape):
+    @classmethod
+    @abc.abstractmethod
+    def _getmask(cls, coords: npt.NDArray[np.floating], polys: npt.NDArray[np.integer], shape: tuple[int, int, int], **kwargs) -> sparse.csr_matrix:
+        '''Generates a sparse matrix mapping from volume to surface vertices'''
+        pass
+
+    @staticmethod
+    @abc.abstractmethod
+    def sampler(coords: npt.NDArray[np.floating], shape: tuple[int, int, int], **kwargs) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.floating]]:
+        '''Generates a sparse matrix mapping from volume to surface vertices'''
+        pass
+
+def _savecache(filename: str, left: sparse.csr_matrix, right: sparse.csr_matrix, shape: MapperShape) -> None:
     np.savez(filename,
              left_data=left.data,
              left_indices=left.indices,

cortex/mapper/patch.py

@@ -6,6 +6,8 @@
 from .. import polyutils
 
 class PatchMapper(Mapper):
+    patchsize: int
+
     @classmethod
     def _getmask(cls, pts, polys, shape, npts=64, mp=True, **kwargs):
         rand = np.random.rand(2, npts)

cortex/mapper/point.py

@@ -1,12 +1,13 @@
 import numpy as np
+import numpy.typing as npt
 from scipy import sparse
 
 from . import Mapper
 from . import samplers
 
 class PointMapper(Mapper):
     @classmethod
-    def _getmask(cls, coords, polys, shape, **kwargs):
+    def _getmask(cls, coords: npt.NDArray[np.floating], polys: npt.NDArray[np.integer], shape: tuple[int, int, int], **kwargs) -> sparse.csr_matrix:
         valid = np.unique(polys)
         mcoords = np.nan * np.ones_like(coords)
         mcoords[valid] = coords[valid]

cortex/mapper/samplers.py

@@ -1,11 +1,15 @@
 import numpy as np
+import numpy.typing as npt
 
 def collapse(j, data):
     """Collapses samples into a single row"""
     uniques = np.unique(j)
     return uniques, np.array([data[j == u].sum() for u in uniques])
 
-def nearest(coords, shape, **kwargs):
+def nearest(coords: npt.NDArray[np.floating], shape: tuple[int, int, int], **kwargs) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.float64]]:
+    if len(kwargs) > 0:
+        raise ValueError('nearest sampler does not take any kwargs')
+
     valid = ~(np.isnan(coords).all(1))
     valid = np.logical_and(valid, np.logical_and(coords[:,0] > -.5, coords[:,0] < shape[2]+.5))
     valid = np.logical_and(valid, np.logical_and(coords[:,1] > -.5, coords[:,1] < shape[1]+.5))
@@ -16,7 +20,10 @@ def nearest(coords, shape, **kwargs):
     #return np.nonzero(valid)[0], j, (rcoords > 0).all(1) #np.ones((valid.sum(),))
     return np.nonzero(valid)[0], j, np.ones((valid.sum(),))
 
-def trilinear(coords, shape, **kwargs):
+def trilinear(coords: npt.NDArray[np.floating], shape: tuple[int, int, int], **kwargs) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.floating]]:
+    if len(kwargs) > 0:
+        raise ValueError('trilinear sampler does not take any kwargs')
+
     #trilinear interpolation equation from http://paulbourke.net/miscellaneous/interpolation/
     valid = ~(np.isnan(coords).all(1))
     (x, y, z), floor = np.modf(coords[valid].T)