Merge pull request #20511 from QuLogic/fix-ma.masked_where

dstansby · web-flow · commit 33beff95fe6d · 2021-07-05T22:17:13.000+01:00
Fix calls to np.ma.masked_where
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1545,11 +1545,11 @@ class LogNorm(Normalize):
 
     def autoscale(self, A):
         # docstring inherited.
-        super().autoscale(np.ma.masked_less_equal(A, 0, copy=False))
+        super().autoscale(np.ma.array(A, mask=(A <= 0)))
 
     def autoscale_None(self, A):
         # docstring inherited.
-        super().autoscale_None(np.ma.masked_less_equal(A, 0, copy=False))
+        super().autoscale_None(np.ma.array(A, mask=(A <= 0)))
 
 
 @_make_norm_from_scale(
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -1158,8 +1158,10 @@ def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,
         return barb_list
 
     def set_UVC(self, U, V, C=None):
-        self.u = ma.masked_invalid(U, copy=False).ravel()
-        self.v = ma.masked_invalid(V, copy=False).ravel()
+        # We need to ensure we have a copy, not a reference to an array that
+        # might change before draw().
+        self.u = ma.masked_invalid(U, copy=True).ravel()
+        self.v = ma.masked_invalid(V, copy=True).ravel()
 
         # Flip needs to have the same number of entries as everything else.
         # Use broadcast_to to avoid a bloated array of identical values.
@@ -1170,7 +1172,7 @@ def set_UVC(self, U, V, C=None):
             flip = self.flip
 
         if C is not None:
-            c = ma.masked_invalid(C, copy=False).ravel()
+            c = ma.masked_invalid(C, copy=True).ravel()
             x, y, u, v, c, flip = cbook.delete_masked_points(
                 self.x.ravel(), self.y.ravel(), self.u, self.v, c,
                 flip.ravel())
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -1233,6 +1233,34 @@ def test_imshow_quantitynd():
     fig.canvas.draw()
 
 
+@check_figures_equal(extensions=['png'])
+def test_norm_change(fig_test, fig_ref):
+    # LogNorm should not mask anything invalid permanently.
+    data = np.full((5, 5), 1, dtype=np.float64)
+    data[0:2, :] = -1
+
+    masked_data = np.ma.array(data, mask=False)
+    masked_data.mask[0:2, 0:2] = True
+
+    cmap = plt.get_cmap('viridis').with_extremes(under='w')
+
+    ax = fig_test.subplots()
+    im = ax.imshow(data, norm=colors.LogNorm(vmin=0.5, vmax=1),
+                   extent=(0, 5, 0, 5), interpolation='nearest', cmap=cmap)
+    im.set_norm(colors.Normalize(vmin=-2, vmax=2))
+    im = ax.imshow(masked_data, norm=colors.LogNorm(vmin=0.5, vmax=1),
+                   extent=(5, 10, 5, 10), interpolation='nearest', cmap=cmap)
+    im.set_norm(colors.Normalize(vmin=-2, vmax=2))
+    ax.set(xlim=(0, 10), ylim=(0, 10))
+
+    ax = fig_ref.subplots()
+    ax.imshow(data, norm=colors.Normalize(vmin=-2, vmax=2),
+              extent=(0, 5, 0, 5), interpolation='nearest', cmap=cmap)
+    ax.imshow(masked_data, norm=colors.Normalize(vmin=-2, vmax=2),
+              extent=(5, 10, 5, 10), interpolation='nearest', cmap=cmap)
+    ax.set(xlim=(0, 10), ylim=(0, 10))
+
+
 @pytest.mark.parametrize('x', [-1, 1])
 @check_figures_equal(extensions=['png'])
 def test_huge_range_log(fig_test, fig_ref, x):
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -202,6 +202,17 @@ def test_barbs_flip():
              flip_barb=Y < 0)
 
 
+def test_barb_copy():
+    fig, ax = plt.subplots()
+    u = np.array([1.1])
+    v = np.array([2.2])
+    b0 = ax.barbs([1], [1], u, v)
+    u[0] = 0
+    assert b0.u[0] == 1.1
+    v[0] = 0
+    assert b0.v[0] == 2.2
+
+
 def test_bad_masked_sizes():
     """Test error handling when given differing sized masked arrays."""
     x = np.arange(3)
