How to make a 3d matlibplot not show masked values - matplotlib

Question

The diagram should only show the masked values. As in the (manipulated) figure on the right side.
Default shows all values. In 2d diagramms there is no problem.
Is it also possible in 3d diagrams? If yes, how to?
import matplotlib.pyplot as plt
import numpy as np
Z = np.array([
[ 1, 1, 1, 1, 1, ],
[ 1, 1, 1, 1, 1, ],
[ 1, 1, 1, 1, 1, ],
[ 1, 1, 1, 1, 1, ],
[ 1, 1, 1, 1, 1, ],
])
x, y = Z.shape
xs = np.arange(x)
ys = np.arange(y)
X, Y = np.meshgrid(xs, ys)
M = np.ma.fromfunction(lambda i, j: i > j, (x, y))
R = np.ma.masked_where(M, Z)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, R)
#ax.plot_wireframe(X, Y, R)
#ax.plot_trisurf(X.flatten(), Y.flatten(), R.flatten())
fig.show()

Answer 1

Update: Matplotlib >= 3.5.0
As pointed out by eric's comment, the issue is solved in matplotlib <= 3.5.0, and the code from the OP works as expected. So right now, if you can probably your best option is to update matplotlib.
The original answer is left here for situations were updating matplotlib might not be an option.
Old: Matplotlib < 3.5.0
The bad news is that it seems that plot_surface() just ignores masks. In fact there is an open issue about it.
However, here they point out a workaround that although it's far from perfect it may allow you get some acceptable results. The key issue is that NaN values will not be plotted, so you need to 'mask' the values that you don't want to plot as np.nan.
Your example code would become something like this:
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
Z = np.array([
[ 1, 1, 1, 1, 1, ],
[ 1, 1, 1, 1, 1, ],
[ 1, 1, 1, 1, 1, ],
[ 1, 1, 1, 1, 1, ],
[ 1, 1, 1, 1, 1, ],
])
x, y = Z.shape
xs = np.arange(x)
ys = np.arange(y)
X, Y = np.meshgrid(xs, ys)
R = np.where(X>=Y, Z, np.nan)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, R, rstride=1, linewidth=0)
fig.show()
*I had to add the rstride=1 argument to the plot_surface call; otherwise I get a segmentation fault... o_O
And here's the result: