Matplotlib 3D plot, How to use set_zlim() correctly

Question

I have multiple plot. I want to set Z lim such that it only shows the curve within the specified range. My code here

# make 3d axes
fig = plt.figure()
ax = fig.gca(projection='3d')

# test data
x = np.arange(-1., 1., .1)
y = np.arange(-1., 1., .1)
z1 = x**2
z2 = x**3
z3 = x**4

# plot test data
ax.plot(x, np.ones(len(x)), z1)
ax.plot(x, np.ones(len(x))*3, z2)
ax.plot(x, np.ones(len(x))*5, z3)

# make labels
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_zlim(0)
plt.show()

Shows I'm expecting that only the positive part of Z2 shown on the graph but it shows all the curves and make the plot messier.

Answer 1

You can use MaskedArray to filter un-wanted data before plotting. Here is the solution:

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
import numpy.ma as ma

fig = plt.figure(figsize=[8,7])
ax = fig.gca(projection='3d')

# test data
x = np.arange(-1., 1., .1)
y = np.arange(-1., 1., .1)
z1 = x**2
z2 = x**3
z3 = x**4

# masking the data (take `z2` as the base)
z2m = ma.masked_less_equal(z2, 0, copy=True)
y2 = np.ones(len(x))*3

# applying the mask to corresponding `x` and `y`
x2m = ma.MaskedArray(x, mask=z2m.mask)
y2m = ma.MaskedArray(y2, mask=z2m.mask)

# we get (x2m, y2m, z2m) to plot

# plot test data
ax.plot(x, np.ones(len(x)), z1)
ax.plot(x, np.ones(len(x))*3, z2)
ax.scatter(x2m, y2m, z2m)    # plot the masked data as points
ax.plot(x, np.ones(len(x))*5, z3)

# make labels
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_zlim(0)

# set view angles to get better plot
ax.azim = 220   # z rotation (default=270)
ax.elev = 2     # x rotation (default=0)
ax.dist = 10    # zoom (define perspective)
plt.show()

The output plot: