Use np.ma functions only when input is masked

Author: chahak13

lib/matplotlib/collections.py

@@ -545,9 +545,14 @@ def set_offsets(self, offsets):
         offsets = np.asanyarray(offsets)
         if offsets.shape == (2,):  # Broadcast (2,) -> (1, 2) but nothing else.
             offsets = offsets[None, :]
-        self._offsets = np.ma.column_stack(
-            (np.asanyarray(self.convert_xunits(offsets[:, 0]), float),
-             np.asanyarray(self.convert_yunits(offsets[:, 1]), float)))
+        if isinstance(offsets, np.ma.MaskedArray):
+            self._offsets = np.ma.column_stack(
+                (np.asanyarray(self.convert_xunits(offsets[:, 0]), float),
+                 np.asanyarray(self.convert_yunits(offsets[:, 1]), float)))
+        else:
+            self._offsets = np.column_stack(
+                (np.asanyarray(self.convert_xunits(offsets[:, 0]), float),
+                 np.asanyarray(self.convert_yunits(offsets[:, 1]), float)))
         self.stale = True
 
     def get_offsets(self):

lib/matplotlib/tests/test_collections.py

@@ -1151,12 +1151,40 @@ def test_check_masked_offsets():
     ax.scatter(unmasked_x, masked_y)
 
 
-def test_masked_set_offsets():
+@check_figures_equal(extensions=["png"])
+def test_masked_set_offsets(fig_ref, fig_test):
+    x = np.ma.array([1, 2, 3, 4, 5], mask=[0, 0, 1, 1, 0])
+    y = np.arange(1, 6)
+
+    ax_test = fig_test.add_subplot()
+    scat = ax_test.scatter(x, y)
+    x += 1
+    scat.set_offsets(np.ma.column_stack([x, y]))
+    ax_test.set_xticks([])
+    ax_test.set_yticks([])
+    ax_test.set_xlim(0, 7)
+    ax_test.set_ylim(0, 6)
+
+    ax_ref = fig_ref.add_subplot()
+    ax_ref.scatter([2, 3, 6], [1, 2, 5])
+    ax_ref.set_xticks([])
+    ax_ref.set_yticks([])
+    ax_ref.set_xlim(0, 7)
+    ax_ref.set_ylim(0, 6)
+
+
+def test_check_offsets_dtype():
+    # Check that setting offsets doesn't change dtype
     x = np.ma.array([1, 2, 3, 4, 5], mask=[0, 0, 1, 1, 0])
     y = np.arange(1, 6)
 
     fig, ax = plt.subplots()
     scat = ax.scatter(x, y)
     x += 1
-    scat.set_offsets(np.ma.column_stack([x, y]))
-    assert np.ma.is_masked(scat.get_offsets())
+    masked_offsets = np.ma.column_stack([x, y])
+    scat.set_offsets(masked_offsets)
+    assert isinstance(scat.get_offsets(), type(masked_offsets))
+
+    unmasked_offsets = np.column_stack([x, y])
+    scat.set_offsets(unmasked_offsets)
+    assert isinstance(scat.get_offsets(), type(unmasked_offsets))