FIX: handle nans in RGBA input with ScalarMappables

rcomer · rcomer · commit dd94ac37a168 · 2024-03-07T21:38:13.000Z
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -328,7 +328,7 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True):
         treated as an RGB or RGBA array, and no mapping will be done.
         The array can be `~numpy.uint8`, or it can be floats with
         values in the 0-1 range; otherwise a ValueError will be raised.
-        If it is a masked array, any masked elements will be set to 0 alpha.
+        Any NaNs or masked elements will be set to 0 alpha.
         If the last dimension is 3, the *alpha* kwarg (defaulting to 1)
         will be used to fill in the transparency.  If the last dimension
         is 4, the *alpha* kwarg is ignored; it does not
@@ -360,6 +360,12 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True):
                 else:
                     raise ValueError("Third dimension must be 3 or 4")
                 if xx.dtype.kind == 'f':
+                    # If any of R, G, B, or A is nan, set to 0
+                    if np.any(nans := np.isnan(x)):
+                        if xx.shape[2] == 4:
+                            xx = xx.copy()
+                        xx[np.any(nans, axis=2), :] = 0
+
                     if norm and (xx.max() > 1 or xx.min() < 0):
                         raise ValueError("Floating point image RGB values "
                                          "must be in the 0..1 range.")
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1352,6 +1352,32 @@ def test_scalarmappable_to_rgba(bytes):
     np.testing.assert_almost_equal(sm.to_rgba(xm[..., :3], bytes=bytes), expected)
 
 
+@pytest.mark.parametrize("bytes", (True, False))
+def test_scalarmappable_nan_to_rgba(bytes):
+    sm = cm.ScalarMappable()
+
+    # RGBA
+    x = np.ones((2, 3, 4), dtype=float) * 0.5
+    x[0, 0, 0] = np.nan
+    expected = x.copy()
+    expected[0, 0, :] = 0
+    if bytes:
+        expected = (expected * 255).astype(np.uint8)
+    np.testing.assert_almost_equal(sm.to_rgba(x, bytes=bytes), expected)
+    assert np.any(np.isnan(x))  # Input array should not be changed
+
+    # RGB
+    expected[..., 3] = 255 if bytes else 1
+    expected[0, 0, 3] = 0
+    np.testing.assert_almost_equal(sm.to_rgba(x[..., :3], bytes=bytes), expected)
+    assert np.any(np.isnan(x))  # Input array should not be changed
+
+    # Out-of-range fail
+    x[1, 0, 0] = 42
+    with pytest.raises(ValueError, match='0..1 range'):
+        sm.to_rgba(x[..., :3], bytes=bytes)
+
+
 def test_failed_conversions():
     with pytest.raises(ValueError):
         mcolors.to_rgba('5')
