Trace a 3d graph with a black line where Z = 0?

Question

I have a functional 3d graph, but I want to make a trace line on the graph for when z = 0.

I tried to split up the graphs for when z>=0 and z<0 but this does not mak

Answer 1

When just highlighting the Z=0 line you need to remember that at that point you no longer have a surface but a 2D plane. You then want to find where that 2D plane is equal to zero. You want to use what Poolka suggested which is ax.contour(x,y,z,[0]). I would suggest changing the transparency (alpha) in the plots to make that line more visible.

You can also make those 2 regions separated at zero 2 different colors by creating a custom colormap and making your vmin and vmax centered around zero.

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np
import matplotlib.colors



def equation(delta=0.05):
    x = np.arange(0,1,delta)
    y = np.arange(2,6,delta)
    X,Y = np.meshgrid(x,y)

    Z = (X*Y-X-0.5*Y**2+2*0.5*Y)**2-4*(0.5*Y**2-0.5*Y)*(X-X*Y+Y-0.5*Y)

    return X, Y, Z


fig = plt.figure()
ax = Axes3D(fig)

#set labels for graph
ax.set_xlabel('P')
ax.set_ylabel('K')
ax.set_zlabel('Z')



#Create custom colormap with only 2 colors
colors = ["blue","red"]
cm1 = LinearSegmentedColormap.from_list('my_list', colors, N=2)



x,y,z = equation(0.01)
surf=ax.plot_surface(x,y,z,alpha=.7,cmap=cm1,vmin=-150,vmax=150) #use custom colormap

#Use a contour plot to isolate Z=0 since it is a line and no longer a surface
ax.contour(x,y,z,[0],colors='k',linewidths=3)

ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))


plt.show()