FIX ignore masked voxels when averaging with nanmean=True (#436)

* FIX ignore masked voxels when averaging with nanmean=True

* FIX extend fix to nanmean=False and thin mask

* FIX make it work for non-masked arrays as well

cortex/quickflat/utils.py

@@ -84,22 +84,33 @@ def make_flatmap_image(braindata, height=1024, recache=False, nanmean=False, **k
         img = (np.nan*np.ones(mask.shape)).astype(data.dtype)
         mimg = (np.nan*np.ones(badmask.shape)).astype(data.dtype)
 
-        # pixmap is a (pixels x voxels) sparse non-negative weight matrix
-        # where each row sums to 1
-
-        if not nanmean:
-            # pixmap.dot(vec) gives mean of vec across cortical thickness
-            mimg[badmask] = pixmap.dot(data.ravel())[badmask].astype(mimg.dtype)
-        else:
-            # to ignore nans in the weighted mean, nanmean =
-            # sum(weights * non-nan values) / sum(weights on non-nan values)
-            nonnan_sum = pixmap.dot(np.nan_to_num(data.ravel()))
-            weights_on_nonnan = pixmap.dot((~np.isnan(data.ravel())).astype(data.dtype))
+        # Pixmap is a (pixels x voxels) sparse non-negative weight matrix where
+        # each row sums to 1, so pixmap.dot(vec) gives the mean of vec across
+        # cortical thickness.
+
+        # To ignore NaNs (or masked data) in the weighted mean, we normalize by
+        # the sum of the non-ignored weights: sum(weights * non-ignored values)
+        # / sum(weights on non-ignored values)
+        ignored = None
+
+        if not nanmean:  # NaN are not ignored: mean([1., 2., NaN]) = NaN
+            averaged_data = pixmap.dot(data.ravel())
+            if isinstance(data, np.ma.MaskedArray):  
+                ignored = data.ravel().mask  # masked voxels are ignored
+
+        else:  # NaN are ignored: mean([1., 2., NaN]) = 1.5
+            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
+
+        if ignored is not None:
+            weights_not_ignored = pixmap.dot((~ignored).astype(data.dtype))
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore", RuntimeWarning)
-                nanmean_data = nonnan_sum / weights_on_nonnan
-            mimg[badmask] = nanmean_data[badmask].astype(mimg.dtype)
+                averaged_data = averaged_data / weights_not_ignored
 
+        mimg[badmask] = averaged_data[badmask].astype(mimg.dtype)
         img[mask] = mimg
         img = img.T[::-1]
         # Make img a c-contiguous array or pil will complain when saving it

cortex/tests/test_quickflat.py

@@ -10,18 +10,33 @@
 
 @pytest.mark.skipif(no_inkscape, reason='Inkscape required')
 def test_quickflat():
-	tf = tempfile.NamedTemporaryFile(suffix=".png")
-	view = cortex.Volume.random("S1", "fullhead", cmap="hot")
-	cortex.quickflat.make_png(tf.name, view)
+    tf = tempfile.NamedTemporaryFile(suffix=".png")
+    view = cortex.Volume.random("S1", "fullhead", cmap="hot")
+    cortex.quickflat.make_png(tf.name, view)
 
 
 @pytest.mark.skipif(no_inkscape, reason='Inkscape required')
 def test_colorbar_location():
-	view = cortex.Volume.random("S1", "fullhead", cmap="hot")
-	for colorbar_location in ['left', 'center', 'right', (0, 0.2, 0.4, 0.3)]:
-		cortex.quickflat.make_figure(
-			view, with_colorbar=True, colorbar_location=colorbar_location)
-
-	with pytest.raises(ValueError):
-		cortex.quickflat.make_figure(
-			view, with_colorbar=True, colorbar_location='unknown_location')
+    view = cortex.Volume.random("S1", "fullhead", cmap="hot")
+    for colorbar_location in ['left', 'center', 'right', (0, 0.2, 0.4, 0.3)]:
+        cortex.quickflat.make_figure(view, with_colorbar=True,
+                                     colorbar_location=colorbar_location)
+
+    with pytest.raises(ValueError):
+        cortex.quickflat.make_figure(view, with_colorbar=True,
+                                     colorbar_location='unknown_location')
+
+
+@pytest.mark.skipif(no_inkscape, reason='Inkscape required')
+@pytest.mark.parametrize("type_", ["thick", "thin"])
+@pytest.mark.parametrize("nanmean", [True, False])
+def test_make_flatmap_image_nanmean(type_, nanmean):
+    mask = cortex.db.get_mask("S1", "fullhead", type=type_)
+    data = np.ones(mask.sum())
+    # set 50% of the values in the dataset to NaN
+    data[np.random.rand(*data.shape) > 0.5] = np.nan
+    vol = cortex.Volume(data, "S1", "fullhead", vmin=0, vmax=1)
+    img, extents = cortex.quickflat.utils.make_flatmap_image(
+        vol, nanmean=nanmean)
+    # assert that the nanmean only returns NaNs and 1s
+    assert np.nanmin(img) == 1