FIX unify NaN and alpha handling across quickflat, WebGL and RGB dataviews

cortex/dataset/view2D.py

@@ -8,6 +8,7 @@
 
 from .. import options
 from .views import Dataview, Volume, Vertex, VolumeRGB, VertexRGB
+from .viewRGB import _warn_alpha_range
 from .braindata import BrainData, VolumeData, VertexData
 
 default_cmap2D = options.config.get("basic", "default_cmap2D")
@@ -20,7 +21,8 @@ class Dataview2D(Dataview):
 
     def __init__(self, description: str="", cmap: Optional[str]=None,
                  vmin: Optional[float]=None, vmax: Optional[float]=None,
-                 vmin2: Optional[float]=None, vmax2: Optional[float]=None, state=None, **kwargs):
+                 vmin2: Optional[float]=None, vmax2: Optional[float]=None, state=None,
+                 alpha=None, **kwargs):
         self.cmap = cmap or default_cmap2D
         self.vmin = vmin
         self.vmax = vmax
@@ -32,17 +34,78 @@ def __init__(self, description: str="", cmap: Optional[str]=None,
         if 'priority' not in self.attrs:
             self.attrs['priority'] = 1
         self.description = description
+        # Optional per-voxel/vertex alpha map. Kept as a Volume/Vertex (never
+        # inside ``attrs``, which is JSON-serialized for the WebGL viewer) so
+        # that both renderers can ship it like any other brain.
+        self.alpha = alpha
+
+    @property
+    def alpha(self):
+        """Optional alpha map (Volume/Vertex in [vmin, vmax]) multiplied into
+        the colormap alpha. NaN anywhere (dim1, dim2 or alpha) renders as
+        alpha 0 regardless of this map."""
+        return self._alpha
+
+    @alpha.setter
+    def alpha(self, alpha):
+        if alpha is not None and not isinstance(alpha, self._cls):
+            alpha = np.asarray(alpha)
+            _warn_alpha_range(alpha)
+            alpha = self._wrap_alpha(alpha)
+        if alpha is not None and alpha.subject != self.dim1.subject:
+            raise ValueError("alpha must belong to the same subject as dim1")
+        if alpha is not None and getattr(alpha, "xfmname", None) != getattr(self.dim1, "xfmname", None):
+            raise ValueError("alpha must use the same transform as dim1")
+        self._alpha = alpha
+        self._alpha_brain_cache = None
+
+    @property
+    def _alpha_brain(self):
+        """The alpha map normalized to [0, 1] (vmin=0, vmax=1), as shipped to
+        the WebGL viewer and stored in HDF files. NaN is preserved."""
+        if self.alpha is None:
+            return None
+        if self._alpha_brain_cache is None:
+            self._alpha_brain_cache = self._wrap_alpha(self._normalized_alpha())
+        return self._alpha_brain_cache
+
+    def _wrap_alpha(self, alpha):
+        raise NotImplementedError
+
+    def _normalized_alpha(self, full_volume=False):
+        """User alpha as float array in [0, 1] (NaN preserved), in the space of
+        ``.data`` or, for volumes with ``full_volume=True``, of ``.volume``."""
+        alpha = self.alpha
+        if alpha is None:
+            return None
+        raw = alpha.volume if full_volume else alpha.data
+        arr = np.asarray(raw, dtype=float)
+        if np.asarray(raw).dtype == np.uint8:
+            return arr / 255.
+        vmin = 0. if alpha.vmin is None else float(alpha.vmin)
+        vmax = 1. if alpha.vmax is None else float(alpha.vmax)
+        if vmax == vmin:
+            return np.where(arr >= vmax, 1., 0.)
+        return (arr - vmin) / (vmax - vmin)
 
     def uniques(self, collapse=False):
         yield self.dim1
         yield self.dim2
+        if self.alpha is not None:
+            yield self._alpha_brain
 
     def _write_hdf(self, h5, name="data"):
         self._cls._write_hdf(self.dim1, h5)
         self._cls._write_hdf(self.dim2, h5)
+        names = [self.dim1.name, self.dim2.name]
+        if self.alpha is not None:
+            # Stored normalized to [0, 1] so that it can be restored with
+            # vmin=0, vmax=1 (BrainData nodes carry no range).
+            self._cls._write_hdf(self._alpha_brain, h5)
+            names.append(self._alpha_brain.name)
 
         viewnode = Dataview._write_hdf(self, h5, name=name)
-        viewnode[0] = json.dumps([[self.dim1.name, self.dim2.name]])
+        viewnode[0] = json.dumps([names])
         viewnode[3] = json.dumps([[self.vmin, self.vmin2]])
         viewnode[4] = json.dumps([[self.vmax, self.vmax2]])
         return viewnode
@@ -56,17 +119,30 @@ def to_json(self, simple=False):
 
         d1js = self.dim1.to_json()
         d2js = self.dim2.to_json()
+        # ``is None`` checks: a legitimate vmin/vmax of 0 must not fall back
+        # to the auto range (same class of bug as 5482c8bf for Volume).
         sdict.update(dict(
-            vmin = [[self.vmin or d1js['vmin'][0], self.vmin2 or d2js['vmin'][0]]],
-            vmax = [[self.vmax or d1js['vmax'][0], self.vmax2 or d2js['vmax'][0]]],
+            vmin = [[d1js['vmin'][0] if self.vmin is None else self.vmin,
+                     d2js['vmin'][0] if self.vmin2 is None else self.vmin2]],
+            vmax = [[d1js['vmax'][0] if self.vmax is None else self.vmax,
+                     d2js['vmax'][0] if self.vmax2 is None else self.vmax2]],
             ))
 
         if "xfm" in d1js:
             sdict['xfm'] = [[d1js['xfm'][0], d2js['xfm'][0]]]
 
+        if self.alpha is not None:
+            sdict['alpha'] = [self._alpha_brain.name]
+
         return sdict
 
-    def _to_raw(self, data1, data2):
+    def _to_raw(self, data1, data2, alpha=None):
+        """Colormap (data1, data2) through the 2D colormap.
+
+        Returns ``(r, g, b, a, nan_mask)`` as uint8 channels. ``a`` is the
+        colormap alpha multiplied by the (normalized) user ``alpha`` map when
+        given; it is forced to 0 wherever data1, data2 or alpha is NaN.
+        """
         from matplotlib import pyplot as plt
         from matplotlib.colors import Normalize
         cmapdir = options.config.get("webgl", "colormaps")
@@ -93,17 +169,16 @@ def _to_raw(self, data1, data2):
         g.shape = dim1.shape
         b.shape = dim1.shape
         a.shape = dim1.shape
-        # Preserve nan values as alpha = 0
+        # NaN in either dimension (or in the alpha map) -> alpha = 0
         aidx = np.logical_or(np.isnan(data1), np.isnan(data2))
+        if alpha is not None:
+            alpha = np.asarray(alpha, dtype=float)
+            aidx = np.logical_or(aidx, np.isnan(alpha))
+            user = np.clip(np.nan_to_num(alpha, nan=0.0), 0, 1)
+            a = np.round(a.astype(float) * user).astype(np.uint8)
+        aidx = np.broadcast_to(aidx, a.shape)
         a[aidx] = 0
-        # Code from main, to handle alpha input, prob better here but not tested.
-        # # Possibly move this above setting nans to alpha = 0;
-        # # Possibly multiply specified alpha by alpha in colormap??
-        # if 'alpha' in self.attrs:
-        #     # Over-write alpha from colormap / nans with alpha arg if provided.
-        #     # Question: Might it be important tokeep alpha as an attr?
-        #     a = self.attrs.pop('alpha')
-        return r, g, b, a
+        return r, g, b, a, aidx
 
     @property
     def subject(self):
@@ -134,13 +209,25 @@ class Volume2D(Dataview2D):
         Colormap (or colormap name) to use. If not given defaults to the 
         `default_cmap2d` in your pycortex options.cfg file.
     vmin : float, optional
-        Minimum value in colormap for dim1. If not given defaults to TODO:WHAT
+        Minimum value in colormap for dim1. If not given, the ``vmin`` of dim1
+        is used (its own ``vmin`` for a Volume/Vertex object, the 1st percentile
+        of the data for an array).
     vmax : float, optional
-        Maximum value in colormap for dim1. If not given defaults to TODO:WHAT
+        Maximum value in colormap for dim1. If not given, the ``vmax`` of dim1
+        is used (its own ``vmax`` for a Volume/Vertex object, the 99th
+        percentile of the data for an array).
     vmin2 : float, optional
-        Minimum value in colormap for dim2. If not given defaults to TODO:WHAT
+        Minimum value in colormap for dim2. If not given, the ``vmin`` of dim2
+        is used (same rule as ``vmin``).
     vmax2 : float, optional
-        Maximum value in colormap for dim2. If not given defaults to TODO:WHAT
+        Maximum value in colormap for dim2. If not given, the ``vmax`` of dim2
+        is used (same rule as ``vmax``).
+    alpha : ndarray or Volume/Vertex, optional
+        Per-voxel (per-vertex) opacity multiplied into the colormap alpha.
+        Arrays are taken in [0, 1]; Volume/Vertex objects are normalized by
+        their own ``vmin``/``vmax``. Honored identically by quickflat and the
+        WebGL viewer. Wherever dim1, dim2 or alpha is NaN the data is
+        rendered fully transparent, regardless of this map.
     **kwargs
         All additional arguments in kwargs are passed to the VolumeData and Dataview
 
@@ -189,18 +276,29 @@ def raw(self) -> VolumeRGB:
         if self.dim1.xfmname != self.dim2.xfmname:
             raise ValueError("Both Volumes must have same xfmname to generate single raw volume")
 
-        if ((self.dim1.linear and self.dim2.linear) and
-            (self.dim1.mask.shape == self.dim2.mask.shape) and
-            np.all(self.dim1.mask == self.dim2.mask)):
-            r, g, b, a = self._to_raw(self.dim1.data, self.dim2.data)
+        def _same_mask(a, b):
+            return (a.linear and b.linear and a.mask.shape == b.mask.shape
+                    and np.all(a.mask == b.mask))
+
+        linear = _same_mask(self.dim1, self.dim2)
+        if linear and self.alpha is not None:
+            linear = _same_mask(self.dim1, self.alpha)
+        if linear:
+            r, g, b, a, nan_mask = self._to_raw(
+                self.dim1.data, self.dim2.data, self._normalized_alpha())
         else:
-            r, g, b, a = self._to_raw(self.dim1.volume, self.dim2.volume)
-        # Allow manual override of alpha channel
-        kws = dict(subject=self.dim1.subject, xfmname=self.dim1.xfmname, 
-            state=self.state, description=self.description, **self.attrs)
-        if not 'alpha' in self.attrs:
-            kws['alpha'] = a
-        return VolumeRGB(r, g, b, **kws)
+            r, g, b, a, nan_mask = self._to_raw(
+                self.dim1.volume, self.dim2.volume,
+                self._normalized_alpha(full_volume=True))
+        result = VolumeRGB(r, g, b, subject=self.dim1.subject,
+                           xfmname=self.dim1.xfmname, alpha=a,
+                           state=self.state, description=self.description,
+                           priority=self.priority)
+        result._nan_mask = nan_mask
+        return result
+
+    def _wrap_alpha(self, alpha):
+        return Volume(alpha, self.dim1.subject, self.dim1.xfmname, vmin=0, vmax=1)
 
 
     @property
@@ -229,13 +327,25 @@ class Vertex2D(Dataview2D):
         Colormap (or colormap name) to use. If not given defaults to the 
         `default_cmap2d` in your pycortex options.cfg file.
     vmin : float, optional
-        Minimum value in colormap for dim1. If not given defaults to TODO:WHAT
+        Minimum value in colormap for dim1. If not given, the ``vmin`` of dim1
+        is used (its own ``vmin`` for a Volume/Vertex object, the 1st percentile
+        of the data for an array).
     vmax : float, optional
-        Maximum value in colormap for dim1. If not given defaults to TODO:WHAT
+        Maximum value in colormap for dim1. If not given, the ``vmax`` of dim1
+        is used (its own ``vmax`` for a Volume/Vertex object, the 99th
+        percentile of the data for an array).
     vmin2 : float, optional
-        Minimum value in colormap for dim2. If not given defaults to TODO:WHAT
+        Minimum value in colormap for dim2. If not given, the ``vmin`` of dim2
+        is used (same rule as ``vmin``).
     vmax2 : float, optional
-        Maximum value in colormap for dim2. If not given defaults to TODO:WHAT
+        Maximum value in colormap for dim2. If not given, the ``vmax`` of dim2
+        is used (same rule as ``vmax``).
+    alpha : ndarray or Volume/Vertex, optional
+        Per-voxel (per-vertex) opacity multiplied into the colormap alpha.
+        Arrays are taken in [0, 1]; Volume/Vertex objects are normalized by
+        their own ``vmin``/``vmax``. Honored identically by quickflat and the
+        WebGL viewer. Wherever dim1, dim2 or alpha is NaN the data is
+        rendered fully transparent, regardless of this map.
     **kwargs
         All additional arguments in kwargs are passed to the VolumeData and Dataview
 
@@ -278,12 +388,16 @@ def __repr__(self):
     def raw(self) -> VertexRGB:
         """VertexRGB object containing the colormapped data from this object.
         """
-        r, g, b, a = self._to_raw(self.dim1.data, self.dim2.data)
-        # Allow manual override of alpha channel
-        kws = dict(subject=self.dim1.subject)
-        if not 'alpha' in self.attrs:
-            kws['alpha'] = a
-        return VertexRGB(r, g, b, **kws)
+        r, g, b, a, nan_mask = self._to_raw(
+            self.dim1.data, self.dim2.data, self._normalized_alpha())
+        result = VertexRGB(r, g, b, subject=self.dim1.subject, alpha=a,
+                           state=self.state, description=self.description,
+                           priority=self.priority)
+        result._nan_mask = nan_mask
+        return result
+
+    def _wrap_alpha(self, alpha):
+        return Vertex(alpha, self.dim1.subject, vmin=0, vmax=1)
 
     @property
     def vertices(self):