BUG: Fix weird behavior with mask and units (Fixes #8908)

dopplershift · dopplershift · commit 3d9af3ddbeb5 · 2017-08-16T21:40:58.000-06:00
Refactor many places that check for masked arrays, and fill the masked
values with nans, with a helper to accomplish that. In the helper,
replace the isinstance check with a attribute check for 'mask'. This
allows libraries, like pint, that wrap mask arrays to pass the check and
be converted appropriately.

Also fix one spot using atleast_1d with _check_1d.
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -1963,6 +1963,17 @@ def is_math_text(s):
     return even_dollars
 
 
+def _to_unmasked_float_array(x):
+    """
+    Convert a sequence to a float array; if input was a masked array, masked
+    values are converted to nans.
+    """
+    if hasattr(x, 'mask'):
+        return np.ma.asarray(x, float).filled(np.nan)
+    else:
+        return np.asarray(x, float)
+
+
 def _check_1d(x):
     '''
     Converts a sequence of less than 1 dimension, to an array of 1
@@ -2249,7 +2260,7 @@ def index_of(y):
     try:
         return y.index.values, y.values
     except AttributeError:
-        y = np.atleast_1d(y)
+        y = _check_1d(y)
         return np.arange(y.shape[0], dtype=float), y
 
 
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -186,9 +186,8 @@ def get_datalim(self, transData):
             offsets = transOffset.transform_non_affine(offsets)
             transOffset = transOffset.get_affine()
 
-        if isinstance(offsets, np.ma.MaskedArray):
-            offsets = offsets.filled(np.nan)
-            # get_path_collection_extents handles nan but not masked arrays
+        # get_path_collection_extents handles nan but not masked arrays
+        offsets = cbook._to_unmasked_float_array(offsets)
 
         if len(paths) and len(offsets):
             result = mpath.get_path_collection_extents(
@@ -235,10 +234,9 @@ def _prepare_points(self):
             # This might have changed an ndarray into a masked array.
             transOffset = transOffset.get_affine()
 
-        if isinstance(offsets, np.ma.MaskedArray):
-            offsets = offsets.filled(np.nan)
-            # Changing from a masked array to nan-filled ndarray
-            # is probably most efficient at this point.
+        # Changing from a masked array to nan-filled ndarray
+        # is probably most efficient at this point.
+        offsets = cbook._to_unmasked_float_array(offsets)
 
         return transform, transOffset, offsets, paths
 
@@ -895,9 +893,9 @@ def __init__(self, verts, sizes=None, closed=True, **kwargs):
 
     def set_verts(self, verts, closed=True):
         '''This allows one to delay initialization of the vertices.'''
-        if isinstance(verts, np.ma.MaskedArray):
-            verts = verts.astype(float).filled(np.nan)
-            # This is much faster than having Path do it one at a time.
+        # This is much faster than having Path do it one at a time.
+        verts = cbook._to_unmasked_float_array(verts)
+
         if closed:
             self._paths = []
             for xy in verts:
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -16,7 +16,8 @@
 from . import artist, colors as mcolors, docstring, rcParams
 from .artist import Artist, allow_rasterization
 from .cbook import (
-    iterable, is_numlike, ls_mapper, ls_mapper_r, STEP_LOOKUP_MAP)
+    _to_unmasked_float_array, iterable, is_numlike, ls_mapper, ls_mapper_r,
+    STEP_LOOKUP_MAP)
 from .markers import MarkerStyle
 from .path import Path
 from .transforms import Bbox, TransformedPath, IdentityTransform
@@ -648,20 +649,12 @@ def recache_always(self):
     def recache(self, always=False):
         if always or self._invalidx:
             xconv = self.convert_xunits(self._xorig)
-            if isinstance(self._xorig, np.ma.MaskedArray):
-                x = np.ma.asarray(xconv, float).filled(np.nan)
-            else:
-                x = np.asarray(xconv, float)
-            x = x.ravel()
+            x = _to_unmasked_float_array(xconv).ravel()
         else:
             x = self._x
         if always or self._invalidy:
             yconv = self.convert_yunits(self._yorig)
-            if isinstance(self._yorig, np.ma.MaskedArray):
-                y = np.ma.asarray(yconv, float).filled(np.nan)
-            else:
-                y = np.asarray(yconv, float)
-            y = y.ravel()
+            y = _to_unmasked_float_array(yconv).ravel()
         else:
             y = self._y
 
diff --git a/lib/matplotlib/path.py b/lib/matplotlib/path.py
@@ -23,7 +23,8 @@
 import numpy as np
 
 from . import _path, rcParams
-from .cbook import simple_linear_interpolation, maxdict
+from .cbook import (_to_unmasked_float_array, simple_linear_interpolation,
+                    maxdict)
 
 
 class Path(object):
@@ -129,11 +130,7 @@ def __init__(self, vertices, codes=None, _interpolation_steps=1,
             Makes the path behave in an immutable way and sets the vertices
             and codes as read-only arrays.
         """
-        if isinstance(vertices, np.ma.MaskedArray):
-            vertices = vertices.astype(float).filled(np.nan)
-        else:
-            vertices = np.asarray(vertices, float)
-
+        vertices = _to_unmasked_float_array(vertices)
         if (vertices.ndim != 2) or (vertices.shape[1] != 2):
             msg = "'vertices' must be a 2D list or array with shape Nx2"
             raise ValueError(msg)
@@ -185,11 +182,7 @@ def _fast_from_codes_and_verts(cls, verts, codes, internals=None):
         """
         internals = internals or {}
         pth = cls.__new__(cls)
-        if isinstance(verts, np.ma.MaskedArray):
-            verts = verts.astype(float).filled(np.nan)
-        else:
-            verts = np.asarray(verts, float)
-        pth._vertices = verts
+        pth._vertices = _to_unmasked_float_array(verts)
         pth._codes = codes
         pth._readonly = internals.pop('readonly', False)
         pth.should_simplify = internals.pop('should_simplify', True)
