safe_masked_invalid() for data types with multiple fields

trygvrad · trygvrad · commit ec11ea07a624 · 2024-08-02T10:35:21.000+02:00
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -739,7 +739,17 @@ def safe_masked_invalid(x, copy=False):
     try:
         xm = np.ma.masked_where(~(np.isfinite(x)), x, copy=False)
     except TypeError:
-        return x
+        if len(x.dtype.descr) > 1:
+            # in case of a dtype with multiple fields:
+            try:
+                mask = np.empty(x.shape, dtype=np.dtype('bool, '*len(x.dtype.descr)))
+                for dd, dm in zip(x.dtype.descr, mask.dtype.descr):
+                    mask[dm[0]] = ~(np.isfinite(x[dd[0]]))
+                xm = np.ma.array(x, mask=mask, copy=False)
+            except TypeError:
+                return x
+        else:
+            return x
     return xm
 
 
diff --git a/lib/matplotlib/tests/test_multivariate_axes.py b/lib/matplotlib/tests/test_multivariate_axes.py
@@ -571,3 +571,15 @@ def test_artist_format_cursor_data_multivar():
             xdisp, ydisp = ax.transData.transform(v)
             event = MouseEvent('motion_notify_event', fig.canvas, xdisp, ydisp)
             assert im.format_cursor_data(im.get_cursor_data(event)) == text
+
+
+def test_multivariate_safe_masked_invalid():
+    from matplotlib import cbook
+    dt = np.dtype('float32, float32').newbyteorder('>')
+    x = np.zeros(2, dtype=dt)
+    x['f0'][0] = np.nan
+    xm = cbook.safe_masked_invalid(x)
+    assert (xm['f0'].mask == (True, False)).all()
+    assert (xm['f1'].mask == (False, False)).all()
+    assert '<f' in xm.dtype.descr[0][1]
+    assert '<f' in xm.dtype.descr[1][1]
