ENH add types and tests for quickflat

cortex/quickflat/composite.py

@@ -1,18 +1,26 @@
 import copy
+from typing import Optional, Union
+
+from matplotlib.axes import Axes
+from matplotlib.collections import LineCollection
+from matplotlib.figure import Figure
+from matplotlib.image import AxesImage
 import numpy as np
+import numpy.typing as npt
+
+from .utils import _get_height, _get_extents, _convert_svg_kwargs, _get_images, _parse_defaults
+from .utils import make_flatmap_image, _make_hatch_image, _get_fig_and_ax, get_flatmask, get_flatcache
 from .. import dataset
 from ..database import db
 from ..options import config
-from .utils import _get_height, _get_extents, _convert_svg_kwargs, _get_images, _parse_defaults
-from .utils import make_flatmap_image, _make_hatch_image, _get_fig_and_ax, get_flatmask, get_flatcache
 
 
 """ --- Individual compositing functions --- """
 
 
-def add_curvature(fig, dataview, extents=None, height=None, threshold=True, contrast=None,
-                  brightness=None, smooth=None, cmap='gray', recache=False, curvature_lims=0.5,
-                  legacy_mode=False):
+def add_curvature(fig: Axes, dataview: dataset.Dataview, extents: Optional[tuple[float, float, float, float]]=None, height: Optional[int]=None, threshold: Optional[bool]=True, contrast: Optional[float]=None,
+                  brightness: Optional[float]=None, smooth: Optional[float]=None, cmap: str='gray', recache: bool=False, curvature_lims: float=0.5,
+                  legacy_mode: bool=False) -> AxesImage:
     """Add curvature layer to figure
 
     Parameters
@@ -21,11 +29,12 @@ def add_curvature(fig, dataview, extents=None, height=None, threshold=True, cont
         figure into which to plot image of curvature
     dataview : cortex.Dataview object
         dataview containing data to be plotted, subject (surface identifier), and transform.
-    extents : array-like
+    extents : array-like TODO: fix
         4 values for [Left, Right, Top, Bottom] extents of image plotted. None defaults to 
         extents of images already present in figure.
     height : scalar
         Height of image. None defaults to height of images already present in figure. 
+        TODO: what units?
     threshold : boolean
         Whether to apply a threshold to the curvature values to create a binary curvature image
         (one shade for positive curvature, one shade for negative). `None` defaults to value 
@@ -66,7 +75,7 @@ def add_curvature(fig, dataview, extents=None, height=None, threshold=True, cont
     if default_smoothing.lower()=='none':
         default_smoothing = None
     else:
-        default_smoothing = np.float_(default_smoothing)
+        default_smoothing = np.float64(default_smoothing)
     if smooth is None:
         # (Might still be None!)
         smooth = default_smoothing
@@ -120,15 +129,15 @@ def add_curvature(fig, dataview, extents=None, height=None, threshold=True, cont
                       zorder=0)
     return cvimg
 
-def add_data(fig, braindata, height=1024, thick=32, depth=0.5, pixelwise=True,
-             sampler='nearest', recache=False, nanmean=False):
+def add_data(fig: Figure, braindata: Union[dataset.Volume, dataset.Vertex, dataset.Dataview], height: int=1024, thick: int=32, depth: float=0.5, pixelwise: bool=True,
+             sampler: str='nearest', recache: bool=False, nanmean: bool=False) -> tuple[AxesImage, npt.NDArray]:
     """Add data to quickflat plot
 
     Parameters
     ----------
     fig : figure or ax
         Figure into which to plot image of curvature
-    braindata : one of: {cortex.Volume, cortex.Vertex, cortex.Dataview)
+    braindata : one of: {cortex.Volume, cortex.Vertex, cortex.Dataview}
         Object containing containing data to be plotted, subject (surface identifier),
         and transform.
     height : scalar
@@ -267,8 +276,8 @@ def add_sulci(fig, dataview, extents=None, height=None, with_labels=True, sulci_
     return img
 
 
-def add_hatch(fig, hatch_data, extents=None, height=None, hatch_space=4,
-              hatch_color=(0, 0, 0), sampler='nearest', recache=False):
+def add_hatch(fig: Axes, hatch_data: dataset.Dataview, extents: Optional[tuple[float, float, float, float]]=None, height: Optional[int]=None, hatch_space: int=4,
+              hatch_color: tuple[int, int, int]=(0, 0, 0), sampler: str='nearest', recache: bool=False) -> AxesImage:
     """Add hatching to figure at locations specified in hatch_data
 
     Parameters
@@ -323,8 +332,8 @@ def add_hatch(fig, hatch_data, extents=None, height=None, hatch_space=4,
     return img
 
 
-def add_colorbar(fig, cimg, colorbar_ticks=None, colorbar_location=(0.4, 0.07, 0.2, 0.04), 
-                 orientation='horizontal'):
+def add_colorbar(fig: Figure, cimg: AxesImage, colorbar_ticks: Optional[npt.ArrayLike]=None, colorbar_location: tuple[float, float, float, float]=(0.4, 0.07, 0.2, 0.04),
+                 orientation: str='horizontal') -> Axes:
     """Add a colorbar to a flatmap plot
 
     Parameters
@@ -348,8 +357,8 @@ def add_colorbar(fig, cimg, colorbar_ticks=None, colorbar_location=(0.4, 0.07, 0
     return cbar
 
 
-def add_colorbar_2d(fig, cmap_name, colorbar_ticks,
-                    colorbar_location=(0.425, 0.02, 0.15, 0.15), fontsize=12):
+def add_colorbar_2d(fig: Figure, cmap_name: str, colorbar_ticks: tuple[float, float, float, float],
+                    colorbar_location: tuple[float, float, float, float]=(0.425, 0.02, 0.15, 0.15), fontsize: int=12) -> AxesImage:
     """Add a 2D colorbar to a flatmap plot
 
     Parameters
@@ -358,13 +367,14 @@ def add_colorbar_2d(fig, cmap_name, colorbar_ticks,
     cimg : matplotlib.image.AxesImage object
         Image for which to create colorbar. For reference, matplotlib.image.AxesImage 
         is the output of imshow()
-    colorbar_ticks : array-like
-        values for colorbar ticks
+    colorbar_ticks : tuple[float, float, float, float]
+        Values for colorbar *extents*, in order [xmin, xmax, ymin, ymax]. The colorbar will be plotted with these values as the limits of the colorbar axes, and the ticks will be placed at the values specified in the first two and last two entries of this tuple.
     colorbar_location : array-like
         Four-long list, tuple, or array that specifies location for colorbar axes 
         [left, top, width, height] (?)
     orientation : string
         'vertical' or 'horizontal'
+        TODO: unused
     """
     # a bit sketchy - lazy imports
     import matplotlib.pyplot as plt
@@ -375,9 +385,9 @@ def add_colorbar_2d(fig, cmap_name, colorbar_ticks,
     fig.add_axes(colorbar_location)
     cbar = plt.imshow(cim, extent=colorbar_ticks, interpolation='bilinear')
     cbar.axes.set_xticks(colorbar_ticks[:2])
-    cbar.axes.set_xticklabels(colorbar_ticks[:2], fontdict=dict(size=fontsize))
+    cbar.axes.set_xticklabels([str(t) for t in colorbar_ticks[:2]], fontdict=dict(size=fontsize))
     cbar.axes.set_yticks(colorbar_ticks[2:])
-    cbar.axes.set_yticklabels(colorbar_ticks[2:], fontdict=dict(size=fontsize))
+    cbar.axes.set_yticklabels([str(t) for t in colorbar_ticks[2:]], fontdict=dict(size=fontsize))
 
     return cbar
 
@@ -445,10 +455,10 @@ def add_custom(fig, dataview, svgfile, layer, extents=None, height=None, with_la
                     zorder=6)
     return img
 
-def add_connected_vertices(fig, dataview, exclude_border_width=None,
-                           height=None, extents=None, recache=False,
-                           color=(1.0, 0.5, 0.1, 0.6), linewidth=0.75,
-                           alpha=1.0, **kwargs):
+def add_connected_vertices(fig: Axes, dataview: dataset.Volume, exclude_border_width: Optional[int]=None,
+                           height: Optional[int]=None, extents: Optional[tuple[float, float, float, float]]=None, recache: bool=False,
+                           color: tuple[float, float, float, float]=(1.0, 0.5, 0.1, 0.6), linewidth: float=0.75,
+                           alpha: float=1.0, **kwargs) -> LineCollection:
     """Plot lines btw distant vertices that are within the same voxel
 
     Parameters
@@ -462,10 +472,10 @@ def add_connected_vertices(fig, dataview, exclude_border_width=None,
     exclude_border_width : scalar or None
         if not None, width from edge of flatmap for which crossover lines are
         not computed.
-    height : scalar
+    height : scalar or None
         Height of image. if None, defaults to height of images already present
         in figure.
-    extents : array-like
+    extents : array-like or None
         4 values for [Left, Right, Bottom, Top] extents of image plotted. If
         None, defaults to extents of images already present in figure.
     color : rgba tuple
@@ -532,7 +542,7 @@ def add_connected_vertices(fig, dataview, exclude_border_width=None,
     # (This is the most time consuming step, as it draws many lines)
     # print('plotting lines...')
     fig, ax = _get_fig_and_ax(fig)
-    lc = LineCollection(pix_array_scaled,
+    lc = LineCollection(list(pix_array_scaled),
                         transform=fig.transFigure,
                         figure=fig,
                         colors=color,

cortex/quickflat/utils.py

@@ -4,22 +4,25 @@
 import string
 import warnings
 from functools import reduce
+from typing import Literal, Optional, Union, cast
 
 import numpy as np
+import numpy.typing as npt
+from scipy import sparse # TODO: remove if loading is slow
 
 from .. import dataset, utils
 from ..database import db
 from ..options import config
 
 
-def make_flatmap_image(braindata, height=1024, recache=False, nanmean=False, **kwargs):
+def make_flatmap_image(braindata: Union[dataset.Volume, dataset.Vertex, dataset.Dataview], height: int=1024, recache: bool=False, nanmean: bool=False, **kwargs) -> tuple[npt.NDArray[np.uint8], npt.NDArray[np.floating]]:
     """Generate flatmap image from volumetric brain data
 
     This 
 
     Parameters
     ----------
-    braindata : one of: {cortex.Volume, cortex.Vertex, cortex.Dataview)
+    braindata : one of: {cortex.Volume, cortex.Vertex, cortex.Dataview}
         Object containing containing data to be plotted, subject (surface identifier), 
         and transform.
     height : scalar 
@@ -34,9 +37,10 @@ def make_flatmap_image(braindata, height=1024, recache=False, nanmean=False, **k
 
     Returns
     -------
-    image : 
-
-    extents :
+    image : numpy.ndarray[np.uint8]
+        The generated flatmap image.
+    extents : numpy.ndarray[np.floating]
+        The extents of the generated flatmap image.
 
     """
     mask, extents = get_flatmask(braindata.subject, height=height, recache=recache)
@@ -100,7 +104,7 @@ def make_flatmap_image(braindata, height=1024, recache=False, nanmean=False, **k
             averaged_data = pixmap.dot(np.nan_to_num(data.ravel()))
             ignored = np.isnan(data.ravel())
             if isinstance(data, np.ma.MaskedArray):
-                ignored = ignored.filled()  # masked voxels are also ignored
+                ignored = cast(np.ma.MaskedArray, ignored).filled()  # masked voxels are also ignored
 
         if ignored is not None:
             weights_not_ignored = pixmap.dot((~ignored).astype(data.dtype))
@@ -117,7 +121,7 @@ def make_flatmap_image(braindata, height=1024, recache=False, nanmean=False, **k
 
         return img, extents
 
-def get_flatmask(subject, height=1024, recache=False):
+def get_flatmask(subject: str, height: int=1024, recache: bool=False) -> tuple[npt.NDArray[np.bool_], npt.NDArray[np.floating]]:
     """
     Parameters
     ----------
@@ -141,8 +145,8 @@ def get_flatmask(subject, height=1024, recache=False):
 
     return mask, extents
 
-def get_flatcache(subject, xfmname, pixelwise=True, thick=32, sampler='nearest',
-                  recache=False, height=1024, depth=0.5):
+def get_flatcache(subject: str, xfmname: Optional[str], pixelwise: bool=True, thick: int=32, sampler: str='nearest',
+                  recache: bool=False, height: int=1024, depth: float=0.5):
     """
 
     Parameters
@@ -188,8 +192,14 @@ def get_flatcache(subject, xfmname, pixelwise=True, thick=32, sampler='nearest',
 
     if not pixelwise and xfmname is not None:
         from scipy import sparse
+        from ..mapper import Mapper
         mapper = utils.get_mapper(subject, xfmname, sampler)
-        pixmap = pixmap * sparse.vstack(mapper.masks)
+        # get_mapper isn't typed yet (Mapper's typing PR lands after this one), so
+        # mapper is currently just Any. TODO: once that PR types cortex/mapper,
+        # get_mapper's own return annotation makes this redundant -- remove this
+        # import and assert.
+        assert isinstance(mapper, Mapper)
+        pixmap = cast(sparse.csr_matrix, pixmap * sparse.vstack(mapper.masks))
 
     return pixmap
 
@@ -254,23 +264,26 @@ def _convert_svg_kwargs(kwargs):
                for k,v in kwargs.items() if v is not None}
     return out
 
-def _parse_defaults(section):
-    defaults = dict(config.items(section))
-    for k in defaults.keys():
+def _parse_defaults(section: str) -> dict[str, Union[float, list[float], None, str]]:
+    raw = dict(config.items(section))
+    defaults: dict[str, Union[float, list[float], None, str]] = dict(raw)
+    for k, v in raw.items():
         # Convert numbers to floating point numbers
-        if defaults[k][0] in string.digits + '.':
-            if ',' in defaults[k]:
-                defaults[k] = [float(x) for x in defaults[k].split(',')]
+        if v[0] in string.digits + '.':
+            if ',' in v:
+                defaults[k] = [float(x) for x in v.split(',')]
             else:
-                defaults[k] = float(defaults[k])
+                defaults[k] = float(v)
         # Convert 'None' to None
-        if defaults[k] == 'None':
+        if v == 'None':
             defaults[k] = None
         # Special case formatting
         if k=='stroke' or k=='fill':
-            defaults[k] = _color2hex(defaults[k])
-        elif k=='stroke-dasharray' and isinstance(defaults[k], (list,tuple)):
-            defaults[k] = '{}, {}'.format(*defaults[k])
+            defaults[k] = _color2hex(v)
+        elif k=='stroke-dasharray':
+            dasharray = defaults[k]
+            if isinstance(dasharray, (list, tuple)):
+                defaults[k] = '{}, {}'.format(*dasharray)
     return defaults
 
 def _get_fig_and_ax(fig):
@@ -353,7 +366,7 @@ def _make_hatch_image(hatch_data, height, sampler='nearest', hatch_space=4, reca
 
     return hatchim
 
-def _make_flatmask(subject, height=1024):
+def _make_flatmask(subject: str, height: int=1024) -> tuple[npt.NDArray[np.bool_], npt.NDArray[np.floating]]:
     from PIL import Image, ImageDraw
 
     from .. import polyutils
@@ -368,11 +381,11 @@ def _make_flatmask(subject, height=1024):
     draw = ImageDraw.Draw(im)
     draw.polygon(lpts[:,:2].ravel().tolist(), fill=255)
     draw.polygon(rpts[:,:2].ravel().tolist(), fill=255)
-    extents = np.hstack([pts.min(0), pts.max(0)])[[0,3,1,4]]
+    extents: npt.NDArray[np.floating] = np.hstack([pts.min(0), pts.max(0)])[[0,3,1,4]]
 
     return np.array(im).T > 0, extents
 
-def _make_vertex_cache(subject, height=1024):
+def _make_vertex_cache(subject: str, height: int=1024) -> sparse.csr_matrix:
     from scipy import sparse
     from scipy.spatial import cKDTree
     flat, polys = db.get_surf(subject, "flat", merge=True, nudge=True)
@@ -388,10 +401,11 @@ def _make_vertex_cache(subject, height=1024):
 
     kdt = cKDTree(flat[valid,:2])
     dist, vert = kdt.query(grid.T[mask.ravel()])
+    vert = np.asarray(vert)
     dataij = (np.ones((len(vert),)), np.array([np.arange(len(vert)), valid[vert]]))
     return sparse.csr_matrix(dataij, shape=(mask.sum(), len(flat)))
 
-def _make_pixel_cache(subject, xfmname, height=1024, thick=32, depth=0.5, sampler='nearest'):
+def _make_pixel_cache(subject: str, xfmname: str, height: int=1024, thick: int=32, depth: float=0.5, sampler: str='nearest') -> sparse.csr_matrix:
     from scipy import sparse
     from scipy.spatial import Delaunay
     flat, polys = db.get_surf(subject, "flat", merge=True, nudge=True)
@@ -441,7 +455,7 @@ def _make_pixel_cache(subject, xfmname, height=1024, thick=32, depth=0.5, sample
 
         valid = np.logical_and(valid_p, valid_w)
         vidx = np.nonzero(valid)[0]
-        mapper = sparse.csr_matrix((mask.sum(), np.prod(xfm.shape)))
+        mapper: sparse.csr_matrix = sparse.csr_matrix((mask.sum(), np.prod(xfm.shape)))
         if thick == 1:
             i, j, data = sampclass(piacoords[valid]*depth + wmcoords[valid]*(1-depth), xfm.shape)
             mapper = mapper + sparse.csr_matrix((data / float(thick), (vidx[i], j)),

cortex/quickflat/view.py

@@ -4,7 +4,7 @@
 import binascii
 import numpy as np
 import numpy.typing as npt
-from typing import Optional, Union, IO
+from typing import Optional, Union, IO, Sequence
 
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
@@ -41,7 +41,7 @@ def make_figure(braindata: dataset.Dataview, recache: bool=False, pixelwise: boo
                 linewidth: Optional[int]=None, linecolor: Optional[ColorType]=None, roifill: Optional[ColorType]=None, shadow: Optional[int]=None,
                 labelsize: Optional[str]=None, labelcolor: Optional[ColorType]=None, cutout: Optional[str]=None, curvature_brightness: Optional[float]=None,
                 curvature_contrast: Optional[float]=None, curvature_threshold: Optional[bool]=None, fig: Optional[Union[Figure, Axes]]=None, extra_hatch: Optional[tuple[dataset.Dataview, tuple[float, float, float]]]=None,
-                colorbar_ticks: Optional[npt.ArrayLike]=None, colorbar_location: Union[tuple[float, float, float, float], str]='center', roi_list: Optional[list[str]]=None, sulci_list: Optional[list[str]]=None,
+                colorbar_ticks: Optional[npt.ArrayLike]=None, colorbar_location: Union[tuple[float, float, float, float], str]='center', roi_list: Optional[Sequence[str]]=None, sulci_list: Optional[Sequence[str]]=None,
                 nanmean: bool=False) -> Figure:
     """Show a Volume or Vertex on a flatmap with matplotlib.
 
@@ -305,8 +305,8 @@ def make_png(fname: Union[str, os.PathLike, IO], braindata: dataset.Dataview, re
     fig.clf()
     plt.close(fig)
 
-def make_svg(fname, braindata, with_labels=False, with_curvature=True, layers=['rois'],
-             height=1024, overlay_file=None, with_dropout=False, **kwargs):
+def make_svg(fname, braindata: dataset.Dataview, with_labels: bool=False, with_curvature: bool=True, layers: Sequence[str]=['rois'],
+             height: int=1024, overlay_file: Optional[str]=None, with_dropout: bool=False, **kwargs):
     """Save an svg file of the desired flatmap.
 
     This function creates an SVG file with vector graphic ROIs overlaid on a single png image.