3D plot with Matplotlib

pceccon picture pceccon · Dec 19, 2014 · Viewed 11.2k times · Source

I'm simply trying to plot a surface and its contour in 3D, exactly as in this example.

This is the code I'm using to do it:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from matplotlib import cm
import numpy

def plot_3d_contour(x_dim, y_dim, x_steps, y_steps, scalar_field, file_path):
    fig = plt.figure()

    x, y = numpy.mgrid[-x_dim/2:x_dim/2:x_steps*1j, -y_dim/2:y_dim/2:y_steps*1j]
    v_min = numpy.min(scalar_field)
    v_max = nupmy.max(scalar_field)

    ax = fig.gca(projection='3d')

    cset = ax.contourf(x, y, scalar_field, zdir='z', offset=v_min, cmap=cm.coolwarm)
    cset = ax.contourf(x, y, scalar_field, zdir='x', offset=-x_dim/2-1, cmap=cm.coolwarm)
    cset = ax.contourf(x, y, scalar_field, zdir='y', offset=y_dim/2+1, cmap=cm.coolwarm)

    ax.plot_surface(x, y, scalar_field, rstride=10, cstride=10, alpha=0.3)

    ax.set_xlabel('X')
    ax.set_xlim(-x_dim/2-1, x_dim/2+1)
    ax.set_ylabel('Y')
    ax.set_ylim(-y_dim/2-1, y_dim/2+1)
    ax.set_zlabel('Z')
    ax.set_zlim(v_min, v_max)

    plt.savefig(file_path + '.jpg')
    plt.close()

scalar_field = numpy.loadtxt('../scalar_field', delimiter=",")
plot_3d_contour(12, 12, 100, 100, scalar_field, 'scalar_field3D')

However, I'm getting a weird behavior in which the a contour (zdir=y) is being over the surface. Besides, I'm getting a weird contour in z_dir=z (with a section missing):

enter image description here

I'm wondering what I'm missing. The scalar field can be found here.

Answer

Flabetvibes picture Flabetvibes · May 8, 2015

I agree with Ajean. I believe the problem arises because each matplotlib's artist (i.e. PolygonCollection) is rendered separately. There is no way different faces from the same object to be rendered on different sides of another object in the scene.

Here is a useful piece of code :

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np

file_path = "./3D_surface_and_contour.jpg"
p = 0.05
f = -0.01

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

def plot_3d_contour(p, f):
    nrows = 4
    ncols = 5

    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))
    for n in range(nrows * ncols):
        i = n % ncols
        j = n / ncols
        k = n + 1
        if j == 0:
            azim = -60 + (i - 2) * 15
            elev = 30
        elif j == 1:
            azim = -60
            elev = 30 + (i - 2) * 5
        elif j == 2:
            azim = 60 + (i - 2) * 10
            elev = 30
        elif j == 3:
            azim = 60
            elev = 30 + (i - 2) * 5
        ax = fig.add_subplot(nrows, ncols, k, 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()

plot_3d_contour(p, f)

which gives the following image :

enter image description here

The first two rows are produced by a code similar to yours. You might notice that setting the elevation with view_init to a higher value solve the problem. But it is not satisfactory. I have also determined the influence of the range of the z-values (not shown here), the bug seems to appear only when this range is small (you can use the f parameter to test it) which explain why the example does not suffer from it.

The solution I propose is to replace :

ax.contourf(x, y, scalar_field, zdir='y', offset=y_dim/2+1, cmap=cm.coolwarm)

by :

ax.contourf(x, y, scalar_field, zdir='y', offset=-y_dim/2-1, cmap=cm.coolwarm)

in your code and add this additional line :

ax.view_init(azim=60, elev=30)

As shown in the last two rows of the previous image, this way you will be able to avoid the whims of matplotlib.