Plotting a rectangular prism

I'm creating a rectangular prism function, whose output looks like this:

I think that this code can be improved by optimizing the use of np.meshgrid with a Python iterator, but I can't wrap my head around it. It might also be possible to do this with fewer plotting calls, but I can't figure that out either. Ideally, I would change the line drawing to use a Line3DCollection and the areas to use a Patch3DCollection for plotting speed, but I'm still not comfortable enough with the 3D api.

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

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.set_aspect("equal")

# draw cube
def rect_prism(x_range, y_range, z_range):
    # TODO: refactor this to use an iterator
    xx, yy = np.meshgrid(x_range, y_range)
    ax.plot_wireframe(xx, yy, z_range[0], color="r")
    ax.plot_surface(xx, yy, z_range[0], color="r", alpha=0.2)
    ax.plot_wireframe(xx, yy, z_range[1], color="r")
    ax.plot_surface(xx, yy, z_range[1], color="r", alpha=0.2)

    yy, zz = np.meshgrid(y_range, z_range)
    ax.plot_wireframe(x_range[0], yy, zz, color="r")
    ax.plot_surface(x_range[0], yy, zz, color="r", alpha=0.2)
    ax.plot_wireframe(x_range[1], yy, zz, color="r")
    ax.plot_surface(x_range[1], yy, zz, color="r", alpha=0.2)

    xx, zz = np.meshgrid(x_range, z_range)
    ax.plot_wireframe(xx, y_range[0], zz, color="r")
    ax.plot_surface(xx, y_range[0], zz, color="r", alpha=0.2)
    ax.plot_wireframe(xx, y_range[1], zz, color="r")
    ax.plot_surface(xx, y_range[1], zz, color="r", alpha=0.2)

rect_prism(np.array([-1, 1]), np.array([-1, 1]), np.array([-0.5, 0.5]))
plt.show()

Well, unfortunately you can't get any better than this (as far as I have read). And why would you use an iterator in this case ?

I'd however change a bit of the structure of your code, which allows one to easily change the code if there's any need. I didn't changed too many things, just separated the logic into three different functions and added a for loop to get rid of some repetition in your code.

As for the plotting calls, you kinda' can't change the number of calls. Those are the perks of plotting things (in Python at least).

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

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.set_aspect("equal")

def x_y_edge(x_range, y_range, z_range):
    xx, yy = np.meshgrid(x_range, y_range)

    for value in [0, 1]:
        ax.plot_wireframe(xx, yy, z_range[value], color="r")
        ax.plot_surface(xx, yy, z_range[value], color="r", alpha=0.2)

def y_z_edge(x_range, y_range, z_range):
    yy, zz = np.meshgrid(y_range, z_range)

    for value in [0, 1]:
        ax.plot_wireframe(x_range[value], yy, zz, color="r")
        ax.plot_surface(x_range[value], yy, zz, color="r", alpha=0.2)

def x_z_edge(x_range, y_range, z_range):
    xx, zz = np.meshgrid(x_range, z_range)

    for value in [0, 1]:
        ax.plot_wireframe(xx, y_range[value], zz, color="r")
        ax.plot_surface(xx, y_range[value], zz, color="r", alpha=0.2)

def rect_prism(x_range, y_range, z_range):
    x_y_edge(x_range, y_range, z_range)
    y_z_edge(x_range, y_range, z_range)
    x_z_edge(x_range, y_range, z_range)

def main():
    rect_prism(np.array([-1, 1]),
               np.array([-1, 1]),
               np.array([-0.5, 0.5]))
    plt.show()

if __name__ == '__main__':
    main()

NOTE: I've also added if __name__ == '__main__'. By doing the main check, you can have that code only execute when you want to run the module as a program and not have it execute when someone just wants to import your module and call your functions themselves.

StackExchange Code Review Q#155585, answer score: 7