Merge pull request #20290 from anntzer/f16

jklymak · web-flow · commit bc8a115e468c · 2021-05-25T20:33:44.000-07:00
Support imshow(&lt;float16 array&gt;).
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -411,14 +411,12 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
                     # all masked, so values don't matter
                     a_min, a_max = np.int32(0), np.int32(1)
                 if inp_dtype.kind == 'f':
-                    scaled_dtype = A.dtype
-                    # Cast to float64
-                    if A.dtype not in (np.float32, np.float16):
-                        if A.dtype != np.float64:
-                            _api.warn_external(
-                                f"Casting input data from '{A.dtype}' to "
-                                f"'float64' for imshow")
-                        scaled_dtype = np.float64
+                    scaled_dtype = np.dtype(
+                        np.float64 if A.dtype.itemsize > 4 else np.float32)
+                    if scaled_dtype.itemsize < A.dtype.itemsize:
+                        _api.warn_external(
+                            f"Casting input data from {A.dtype} to "
+                            f"{scaled_dtype} for imshow")
                 else:
                     # probably an integer of some type.
                     da = a_max.astype(np.float64) - a_min.astype(np.float64)
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -993,6 +993,13 @@ def test_empty_imshow(make_norm):
         im.make_image(fig._cachedRenderer)
 
 
+def test_imshow_float16():
+    fig, ax = plt.subplots()
+    ax.imshow(np.zeros((3, 3), dtype=np.float16))
+    # Ensure that drawing doesn't cause crash.
+    fig.canvas.draw()
+
+
 def test_imshow_float128():
     fig, ax = plt.subplots()
     ax.imshow(np.zeros((3, 3), dtype=np.longdouble))
