Vectorize Scale.val_in_range

Author: scottshambaugh

lib/matplotlib/scale.py

@@ -30,7 +30,6 @@
 """  # noqa: E501
 
 import inspect
-import math
 import textwrap
 from functools import wraps
 
@@ -119,17 +118,20 @@ def val_in_range(self, val):
         """
         Return whether the value(s) are within the valid range for this scale.
 
-        This method is a generic implementation. Subclasses may implement more
-        efficient solutions for their domain.
-        """
-        try:
-            if not math.isfinite(val):
-                return False
+        Accepts a scalar or array-like ``val``. For a scalar, returns a
+        Python ``bool``. For an array, returns a bool ndarray of the same
+        shape. This is a generic implementation, and subclasses may implement
+        more efficient solutions for their domain.
+        """
+        arr = np.asarray(val)
+        with np.errstate(invalid='ignore'):
+            try:
+                vmin, vmax = self.limit_range_for_scale(arr, arr, minpos=1e-300)
+            except (TypeError, ValueError):
+                result = np.zeros(arr.shape, dtype=bool)
             else:
-                vmin, vmax = self.limit_range_for_scale(val, val, minpos=1e-300)
-                return vmin == val and vmax == val
-        except (TypeError, ValueError):
-            return False
+                result = np.isfinite(arr) & (vmin == arr) & (vmax == arr)
+        return bool(result) if arr.ndim == 0 else result
 
 
 def _make_axis_parameter_optional(init_func):
@@ -219,11 +221,13 @@ def get_transform(self):
 
     def val_in_range(self, val):
         """
-        Return whether the value is within the valid range for this scale.
+        Return whether the value(s) are within the valid range for this scale.
 
         This is True for all values, except +-inf and NaN.
         """
-        return math.isfinite(val)
+        arr = np.asarray(val)
+        result = np.isfinite(arr)
+        return bool(result) if arr.ndim == 0 else result
 
 
 class FuncTransform(Transform):
@@ -431,11 +435,14 @@ def limit_range_for_scale(self, vmin, vmax, minpos):
 
     def val_in_range(self, val):
         """
-        Return whether the value is within the valid range for this scale.
+        Return whether the value(s) are within the valid range for this scale.
 
         This is True for value(s) > 0 except +inf and NaN.
         """
-        return math.isfinite(val) and val > 0
+        arr = np.asarray(val)
+        with np.errstate(invalid='ignore'):
+            result = np.isfinite(arr) & (arr > 0)
+        return bool(result) if arr.ndim == 0 else result
 
 
 class FuncScaleLog(LogScale):
@@ -625,11 +632,13 @@ def get_transform(self):
 
     def val_in_range(self, val):
         """
-        Return whether the value is within the valid range for this scale.
+        Return whether the value(s) are within the valid range for this scale.
 
         This is True for all values, except +-inf and NaN.
         """
-        return math.isfinite(val)
+        arr = np.asarray(val)
+        result = np.isfinite(arr)
+        return bool(result) if arr.ndim == 0 else result
 
 
 class AsinhTransform(Transform):
@@ -759,11 +768,13 @@ def set_default_locators_and_formatters(self, axis):
 
     def val_in_range(self, val):
         """
-        Return whether the value is within the valid range for this scale.
+        Return whether the value(s) are within the valid range for this scale.
 
         This is True for all values, except +-inf and NaN.
         """
-        return math.isfinite(val)
+        arr = np.asarray(val)
+        result = np.isfinite(arr)
+        return bool(result) if arr.ndim == 0 else result
 
 
 class LogitTransform(Transform):
@@ -880,11 +891,14 @@ def limit_range_for_scale(self, vmin, vmax, minpos):
 
     def val_in_range(self, val):
         """
-        Return whether the value is within the valid range for this scale.
+        Return whether the value(s) are within the valid range for this scale.
 
         This is True for value(s) which are between 0 and 1 (excluded).
         """
-        return 0 < val < 1
+        arr = np.asarray(val)
+        with np.errstate(invalid='ignore'):
+            result = (0 < arr) & (arr < 1)
+        return bool(result) if arr.ndim == 0 else result
 
 
 _scale_mapping = {

lib/matplotlib/tests/test_scale.py

@@ -477,3 +477,27 @@ def test_val_in_range_base_fallback():
     assert s.val_in_range(np.nan) is False
     assert s.val_in_range(np.inf) is False
     assert s.val_in_range(-np.inf) is False
+
+
+def test_val_in_range_array():
+    # Vectorized: scalar in -> scalar bool, array in -> bool array.
+    arr = np.array([1.0, -1.0, 0.0, np.nan, np.inf, 5.0])
+    cases = {
+        'linear': [True, True, True, False, False, True],
+        'log':    [True, False, False, False, False, True],
+        'symlog': [True, True, True, False, False, True],
+        'asinh':  [True, True, True, False, False, True],
+    }
+    for name, expected in cases.items():
+        s = mscale._scale_mapping[name](axis=None)
+        np.testing.assert_array_equal(s.val_in_range(arr), expected)
+
+    s = mscale._scale_mapping['logit'](axis=None)
+    np.testing.assert_array_equal(
+        s.val_in_range(np.array([0.1, 0.5, 0.0, 1.0, -0.1, 1.1])),
+        [True, True, False, False, False, False])
+
+    # 2D shape is preserved.
+    out = mscale._scale_mapping['log'](axis=None).val_in_range(
+        np.array([[1.0, -1.0], [0.5, np.nan]]))
+    np.testing.assert_array_equal(out, [[True, False], [True, False]])