Search code examples
pythonmatplotlibplottriangulationmplot3d

Matplotlib -- mplot3d: triplot projected on z=0 axis in 3d plot?


I'm trying to plot a function in two variables, piecewise defined on a set of known triangles, more or less like so:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import random

def f( x, y):
  if x + y < 1: return 0
  else: return 1

x = [0, 1, 1, 0]
y = [0, 0, 1, 1]
tris = [[0, 1, 3], [1, 2,3]]

fig = plt.figure()
ax = fig.add_subplot( 121)
ax.triplot( x, y, tris)

xs = [random.random() for _ in range( 100)]
ys = [random.random() for _ in range( 100)]
zs = [f(xs[i], ys[i]) for i in range( 100)]
ax2 = fig.add_subplot( 122, projection='3d')
ax2.scatter( xs, ys, zs)
plt.show()

Ideally, I'd combine both subplots into one by projecting the triangles onto the axis z=0. I know this is possible with other variants of 2d plots, but not with triplot. Is it possible to get what I want?

PS. this is a heavily simplified version of the actual implementation I am using right now, therefore the random scattering might seem a bit weird.


Solution

  • I'm not an expert, but this was an interesting problem. After doing some poking around, I think I got something close. I made the Triangulation object manually and then passed it and a z list of zeros into plot_trisurf, and it put the triangles in the right place on z=0.

    import matplotlib.pyplot as plt
    import matplotlib.tri as tri
    from mpl_toolkits.mplot3d import Axes3D
    import random
    
    def f( x, y):
        if x + y < 1: return 0
        else: return 1
    
    x = [0, 1, 1, 0]
    y = [0, 0, 1, 1]
    tris = [[0, 1, 3], [1, 2,3]]
    z = [0] * 4
    
    triv = tri.Triangulation(x, y, tris)
    
    fig = plt.figure()
    ax = fig.add_subplot( 111, projection='3d')
    trip = ax.plot_trisurf( triv, z )
    trip.set_facecolor('white')
    
    xs = [random.random() for _ in range( 100)]
    ys = [random.random() for _ in range( 100)]
    zs = [f(xs[i], ys[i]) for i in range( 100)]
    
    ax.scatter( xs, ys, zs)
    plt.show()
    

    ETA: Added a call to set_facecolor on the Poly3DCollection to make it white rather than follow a colormap. Can be futzed with for desired effect...