Repository: matplotlib/matplotlib
Pull request: Fix masking for Axes3D.plot()
<!--
Thank you so much for your PR!  To help us review your contribution, please check
out the development guide https://matplotlib.org/devdocs/devel/index.html
-->

## PR summary
Closes https://github.com/matplotlib/matplotlib/issues/25043 - there was someone who volunteered to look at doing that fix, but I'm not seeing any action on the ticket for the past 9 months and this will block some work I want to do on dynamic masking

Consolidates some the masking logic that was happening in different plotting functions into a private cbook function.

## PR checklist
<!-- Please mark any checkboxes that do not apply to this PR as [N/A].-->

- [x] "closes #0000" is in the body of the PR description to [link the related issue](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue)
- [x] new and changed code is [tested](https://matplotlib.org/devdocs/devel/testing.html)
- [n/a] *Plotting related* features are demonstrated in an [example](https://matplotlib.org/devdocs/devel/document.html#write-examples-and-tutorials)
- [n/a] *New Features* and *API Changes* are noted with a [directive and release note](https://matplotlib.org/devdocs/devel/coding_guide.html#new-features-and-api-changes)
- [x] Documentation complies with [general](https://matplotlib.org/devdocs/devel/document.html#write-rest-pages) and [docstring](https://matplotlib.org/devdocs/devel/document.html#write-docstrings) guidelines

<!--We understand that PRs can sometimes be overwhelming, especially as the
reviews start coming in.  Please let us know if the reviews are unclear or
the recommended next step seems overly demanding, if you would like help in
addressing a reviewer's comments, or if you have been waiting too long to hear
back on your PR.-->


Problem statement:
[ENH]: Plotting masked arrays correctly in 3D line plot
### Problem

It would be great if the masking feature worked in 3D plots:

```
import matplotlib.pyplot as plt
import numpy as np

# Create some data
x = np.linspace(0, 10, 100)
y = np.linspace(0, 10, 100)
z = np.sin(x) * np.cos(y)

# Mask some of the data
mask = z > 0.5
z = np.ma.array(z, mask=mask)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot(x, y, z)
fig.show()
```

Currently, all of the data is plotted.

### Proposed solution

_No response_

Hints from issue discussion:
Can I work on this issue? I am pretty new to this
Hi @jalajk24, welcome to Matplotlib!  Anyone is welcome to work on anything but, if you are new, it may be better to look through our [Good First Issues](https://github.com/matplotlib/matplotlib/issues?q=is%3Aissue+is%3Aopen+label%3A%22Good+first+issue%22).  Please also have a look at our [Contributing Guide](https://matplotlib.org/stable/devel/contributing.html).

Commit message: Consolidate 3D masking and fix for Axes3D.plot()

Move function to cbook

Whitespace

Privatize cbook._broadcast_with_masks

Patch to review:
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
index b56de559532d..5e9d8b82eb5a 100644
--- a/lib/matplotlib/cbook.py
+++ b/lib/matplotlib/cbook.py
@@ -1088,6 +1088,45 @@ def _combine_masks(*args):
     return margs
 
 
+def _broadcast_with_masks(*args, compress=False):
+    """
+    Broadcast inputs, combining all masked arrays.
+
+    Parameters
+    ----------
+    *args : array-like
+        The inputs to broadcast.
+    compress : bool, default: False
+        Whether to compress the masked arrays. If False, the masked values
+        are replaced by NaNs.
+
+    Returns
+    -------
+    list of array-like
+        The broadcasted and masked inputs.
+    """
+    # extract the masks, if any
+    masks = [k.mask for k in args if isinstance(k, np.ma.MaskedArray)]
+    # broadcast to match the shape
+    bcast = np.broadcast_arrays(*args, *masks)
+    inputs = bcast[:len(args)]
+    masks = bcast[len(args):]
+    if masks:
+        # combine the masks into one
+        mask = functools.reduce(np.logical_or, masks)
+        # put mask on and compress
+        if compress:
+            inputs = [np.ma.array(k, mask=mask).compressed()
+                      for k in inputs]
+        else:
+            inputs = [np.ravel(np.ma.array(k, mask=mask,
+                                           dtype=float).filled(np.nan))
+                      for k in inputs]
+    else:
+        inputs = [np.ravel(k) for k in inputs]
+    return inputs
+
+
 def boxplot_stats(X, whis=1.5, bootstrap=None, labels=None,
                   autorange=False):
     r"""
diff --git a/lib/matplotlib/cbook.pyi b/lib/matplotlib/cbook.pyi
index a6ad1e54a51c..36516e368e8b 100644
--- a/lib/matplotlib/cbook.pyi
+++ b/lib/matplotlib/cbook.pyi
@@ -130,6 +130,7 @@ class GrouperView(Generic[_T]):
 
 def simple_linear_interpolation(a: ArrayLike, steps: int) -> np.ndarray: ...
 def delete_masked_points(*args): ...
+def _broadcast_with_masks(*args, compress: bool=False) -> list[ArrayLike]: ...
 def boxplot_stats(
     X: ArrayLike,
     whis: float | tuple[float, float] = ...,
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
index 0c18bce8ebd3..84bbd763f882 100644
--- a/lib/mpl_toolkits/mplot3d/axes3d.py
+++ b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -11,7 +11,6 @@
 """
 
 from collections import defaultdict
-import functools
 import itertools
 import math
 import textwrap
@@ -1909,8 +1908,7 @@ def plot(self, xs, ys, *args, zdir='z', **kwargs):
         else:
             zs = kwargs.pop('zs', 0)
 
-        # Match length
-        zs = np.broadcast_to(zs, np.shape(xs))
+        xs, ys, zs = cbook._broadcast_with_masks(xs, ys, zs)
 
         lines = super().plot(xs, ys, *args, **kwargs)
         for line in lines:
@@ -2665,8 +2663,7 @@ def scatter(self, xs, ys, zs=0, zdir='z', s=20, c=None, depthshade=True,
         had_data = self.has_data()
         zs_orig = zs
 
-        xs, ys, zs = np.broadcast_arrays(
-            *[np.ravel(np.ma.filled(t, np.nan)) for t in [xs, ys, zs]])
+        xs, ys, zs = cbook._broadcast_with_masks(xs, ys, zs)
         s = np.ma.ravel(s)  # This doesn't have to match x, y in size.
 
         xs, ys, zs, s, c, color = cbook.delete_masked_points(
@@ -2722,7 +2719,7 @@ def bar(self, left, height, zs=0, zdir='z', *args, **kwargs):
 
         patches = super().bar(left, height, *args, **kwargs)
 
-        zs = np.broadcast_to(zs, len(left))
+        zs = np.broadcast_to(zs, len(left), subok=True)
 
         verts = []
         verts_zs = []
@@ -2988,23 +2985,8 @@ def calc_arrows(UVW):
 
         had_data = self.has_data()
 
-        input_args = [X, Y, Z, U, V, W]
-
-        # extract the masks, if any
-        masks = [k.mask for k in input_args
-                 if isinstance(k, np.ma.MaskedArray)]
-        # broadcast to match the shape
-        bcast = np.broadcast_arrays(*input_args, *masks)
-        input_args = bcast[:6]
-        masks = bcast[6:]
-        if masks:
-            # combine the masks into one
-            mask = functools.reduce(np.logical_or, masks)
-            # put mask on and compress
-            input_args = [np.ma.array(k, mask=mask).compressed()
-                          for k in input_args]
-        else:
-            input_args = [np.ravel(k) for k in input_args]
+        input_args = cbook._broadcast_with_masks(X, Y, Z, U, V, W,
+                                                 compress=True)
 
         if any(len(v) == 0 for v in input_args):
             # No quivers, so just make an empty collection and return early
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
index f3666fffc418..4024eeddcf17 100644
--- a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
+++ b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
@@ -668,6 +668,24 @@ def test_surface3d_masked():
     ax.view_init(30, -80, 0)
 
 
+@check_figures_equal(extensions=["png"])
+def test_plot_scatter_masks(fig_test, fig_ref):
+    x = np.linspace(0, 10, 100)
+    y = np.linspace(0, 10, 100)
+    z = np.sin(x) * np.cos(y)
+    mask = z > 0
+
+    z_masked = np.ma.array(z, mask=mask)
+    ax_test = fig_test.add_subplot(projection='3d')
+    ax_test.scatter(x, y, z_masked)
+    ax_test.plot(x, y, z_masked)
+
+    x[mask] = y[mask] = z[mask] = np.nan
+    ax_ref = fig_ref.add_subplot(projection='3d')
+    ax_ref.scatter(x, y, z)
+    ax_ref.plot(x, y, z)
+
+
 @check_figures_equal(extensions=["png"])
 def test_plot_surface_None_arg(fig_test, fig_ref):
     x, y = np.meshgrid(np.arange(5), np.arange(5))


Code review comments:
diff_hunk: @@ -130,6 +130,7 @@ class GrouperView(Generic[_T]):
 
 def simple_linear_interpolation(a: ArrayLike, steps: int) -> np.ndarray: ...
 def delete_masked_points(*args): ...
+def _broadcast_with_masks(*args, compress: bool=False) -> list[ArrayLike]: ...
path: lib/matplotlib/cbook.pyi
text: @user1:
Should `args` also be `ArrayLike`?
```suggestion
def _broadcast_with_masks(*args, compress: bool = ...) -> list[ArrayLike]: ...
```

@user2:
Yes, I think it can be:

```suggestion
def _broadcast_with_masks(*args: ArrayLike, compress: bool = ...) -> list[ArrayLike]: ...
```

diff_hunk: @@ -1088,6 +1088,45 @@ def _combine_masks(*args):
     return margs
 
 
+def _broadcast_with_masks(*args, compress=False):
+    """
+    Broadcast inputs, combining all masked arrays.
+
+    Parameters
+    ----------
+    *args : array-like
+        The inputs to broadcast.
+    compress : bool, default: False
+        Whether to compress the masked arrays. If False, the masked values
+        are replaced by NaNs.
+
+    Returns
+    -------
+    list of array-like
+        The broadcasted and masked inputs.
+    """
+    # extract the masks, if any
+    masks = [k.mask for k in args if isinstance(k, np.ma.MaskedArray)]
+    # broadcast to match the shape
+    bcast = np.broadcast_arrays(*args, *masks)
+    inputs = bcast[:len(args)]
+    masks = bcast[len(args):]
+    if masks:
+        # combine the masks into one
+        mask = functools.reduce(np.logical_or, masks)
+        # put mask on and compress
+        if compress:
+            inputs = [np.ma.array(k, mask=mask).compressed()
+                      for k in inputs]
+        else:
+            inputs = [np.ravel(np.ma.array(k, mask=mask,
+                                           dtype=float).filled(np.nan))
path: lib/matplotlib/cbook.py
text: @user1:
Fits on a line, I think? And also can use the method.
```suggestion
            inputs = [np.ma.array(k, mask=mask, dtype=float).filled(np.nan).ravel()
```

diff_hunk: @@ -1088,6 +1088,45 @@ def _combine_masks(*args):
     return margs
 
 
+def _broadcast_with_masks(*args, compress=False):
+    """
+    Broadcast inputs, combining all masked arrays.
+
+    Parameters
+    ----------
+    *args : array-like
+        The inputs to broadcast.
+    compress : bool, default: False
+        Whether to compress the masked arrays. If False, the masked values
+        are replaced by NaNs.
+
+    Returns
+    -------
+    list of array-like
+        The broadcasted and masked inputs.
+    """
+    # extract the masks, if any
+    masks = [k.mask for k in args if isinstance(k, np.ma.MaskedArray)]
+    # broadcast to match the shape
+    bcast = np.broadcast_arrays(*args, *masks)
+    inputs = bcast[:len(args)]
+    masks = bcast[len(args):]
+    if masks:
+        # combine the masks into one
+        mask = functools.reduce(np.logical_or, masks)
path: lib/matplotlib/cbook.py
text: @user1:
If I'm not mistaken, this is a ufunc that could be written as:
```suggestion
        mask = np.logical_or.reduce(masks)
```