In [1]:
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np

In [2]:
%matplotlib notebook

In [3]:
file_path = "./3D_surface_and_contour.png"
p = 0.05
f = -0.01

In [4]:
def get_data(p):
    x, y, z = axes3d.get_test_data(p)
    z = f * z
    return x, y, z

In [5]:
def plot_3d_contour(p, f):
    nrows = 1
    ncols = 1

    x, y, z = get_data(p)

    x_min, x_max = np.min(x), np.max(x)
    y_min, y_max = np.min(y), np.max(y)
    z_min, z_max = np.min(z), np.max(z)

    fig = plt.figure() #figsize=(15, 10))
    ax = fig.add_subplot( 111, projection='3d')
#        ax.set_title("azim=" + str(azim) + " elev=" + str(elev))
    ax.tick_params(labelsize=8)
#        ax.view_init(azim=azim, elev=elev)
    ax.plot_surface(x, y, z, rstride=10, cstride=10, alpha=0.3)

#    ax.contourf(x, y, z, zdir='z', offset=z_min, cmap=cm.coolwarm)
#    ax.contourf(x, y, z, zdir='x', offset=x_min, cmap=cm.coolwarm)
#        if j == 0 or j == 1:
#            ax.contourf(x, y, z, zdir='y', offset=y_max, cmap=cm.coolwarm)
#        elif j == 2 or j == 3:
#            ax.contourf(x, y, z, zdir='y', offset=y_min, cmap=cm.coolwarm)

    ax.set_xlabel('X')
    ax.set_xlim(x_min, x_max)
    ax.set_ylabel('Y')
    ax.set_ylim(y_min, y_max)
    ax.set_zlabel('Z')
    ax.set_zlim(z_min, z_max)

    #plt.savefig(file_path, dpi=80)
    #plt.close()

In [6]:
plot_3d_contour(0.01, f)



In [ ]: