Matplotlib boolean coloring

Question

I would like to graph a 3d plot of a sphere, which colors sections differently according to some function of the theta, phi coordinates. I could graph two separate graphs, one of each color, but I'm not sure exactly how the gridding/meshing works when plotting. Viz. I want to grid/mesh over all thetas/phis in a sphere, but then throw out certain pairs given a boolean function. Is this possible? Attached is a picture of a scatterplot that does what I'd like to do with surfaces.

Answer 1

Based on scatter3d_demo.py found in the matplotlib tutorial:

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

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')

THETA, PHI = np.ogrid[0:2*np.pi:40j, 0:np.pi:30j]

X = 10 * np.cos(THETA) * np.sin(PHI)
Y = 10 * np.sin(THETA) * np.sin(PHI)
Z = 10 * np.ones_like(THETA) * np.cos(PHI)

def func(THETA, PHI):
    mask = (THETA < np.pi/2) & (np.pi/3 < PHI) & (PHI < 2 * np.pi/3)
    return np.where(mask, 1, 0.5)

C = func(THETA, PHI)

x = X.ravel()
y = Y.ravel()
z = Z.ravel()

c = C.ravel()
ax.scatter(x, y, z, c=c, s=30, vmin=0, vmax=1)
ax.set_aspect('equal')
plt.show()

yields

---

Note you could also color patches on a sphere using plot_surface:

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

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

THETA, PHI = np.ogrid[0:2*np.pi:40j, 0:np.pi:30j]

X = 10 * np.cos(THETA) * np.sin(PHI)
Y = 10 * np.sin(THETA) * np.sin(PHI)
Z = 10 * np.ones_like(THETA) * np.cos(PHI)

def func(THETA, PHI):
    mask = (THETA < np.pi/2) & (np.pi/3 < PHI) & (PHI < 2 * np.pi/3)
    return np.where(mask, 1, 0.5) 

C = func(THETA, PHI)
jet = plt.cm.jet
ax.plot_surface(X, Y, Z, rstride=2, cstride=2, facecolors=jet(C))
ax.set_aspect('equal')
plt.show()