Reputation: 748
I have a scalar function that represents the electric potential in a spherical surface. I want to plot, for a given radius, the surface and link its points to a colormap based on the potential function.
How do I map that scalar function to the colormap in the surface? I suspect it must be in the arguments passed to the function ax.plot_surface
. I tried using the argument: facecolors=potencial(x,y,z)
, but it gave me a ValueError: Invalid RGBA argument
. Looking at the source code of the third example, there is:
# Create an empty array of strings with the same shape as the meshgrid, and
# populate it with two colors in a checkerboard pattern.
colortuple = ('y', 'b')
colors = np.empty(X.shape, dtype=str)
for y in range(ylen):
for x in range(xlen):
colors[x, y] = colortuple[(x + y) % len(colortuple)]
Which I do not understand, nor have an ideia how to link to a scalar function.
My code
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
from scipy import special
def potencial(x,y,z, a=1., v=1.):
r = np.sqrt( np.square(x) + np.square(y) + np.square(z) )
p = z/r #cos(theta)
asr = a/r
s=0
s += np.polyval(special.legendre(1), p) * 3/2*np.power(asr, 2)
s += np.polyval(special.legendre(3), p) * -7/8*np.power(asr, 4)
s += np.polyval(special.legendre(5), p) * 11/16*np.power(asr, 6)
return v*s
# Make data
def sphere_surface(r):
u = np.linspace(0, 2 * np.pi, 100)
v = np.linspace(0, np.pi, 100)
x = r * np.outer(np.cos(u), np.sin(v))
y = r * np.outer(np.sin(u), np.sin(v))
z = r * np.outer(np.ones(np.size(u)), np.cos(v))
return x,y,z
x,y,z = sphere_surface(1.5)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
# Plot the surface
surf = ax.plot_surface(x,y,z, cmap=cm.coolwarm,
linewidth=0, antialiased=False)
fig.colorbar(surf, shrink=0.5, aspect=5)
# This is mapping the color to the z-axis value
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
plt.show()
Upvotes: 2
Views: 1783
Reputation: 339260
In principle there are two ways to colorize a surface plot in matplotlib.
cmap
argument to specify a colormap. In this case the color will be chosen according to the z
array. In case that is not desired, facecolors
argument. This expects an array of colors of the same shape as z
. So in this case we need to choose option 2 and build a color array.
To this end, one may choose a colormap. A colormap maps values between 0 and 1 to a color. Since the potential has values much above and below this range, one need to normalize them into the [0,1] range.
Matplotlib already provides some helper function to do this normalization and since the potential has a 1/x dependency, a logarithmic colorscale may be suitable.
At the end the facecolors may thus be given an array
colors = cmap(norm(potential(...)))
The missing bit is now the colorbar. In order for the colorbar to be linked to the colors from the surface plot, we need to manually set up a ScalarMappable with the colormap and the normalization instance, which we can then supply to the colorbar.
sm = plt.cm.ScalarMappable(cmap=plt.cm.coolwarm, norm=norm)
sm.set_array(pot)
fig.colorbar(sm, shrink=0.5, aspect=5)
Here is full example.
from __future__ import division
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import matplotlib.colors
import numpy as np
from scipy import special
def potencial(x,y,z, a=1., v=1.):
r = np.sqrt( np.square(x) + np.square(y) + np.square(z) )
p = r/z #cos(theta)
asr = a/r
s=0
s += np.polyval(special.legendre(1), p) * 3/2*np.power(asr, 2)
s += np.polyval(special.legendre(3), p) * -7/8*np.power(asr, 4)
s += np.polyval(special.legendre(5), p) * 11/16*np.power(asr, 6)
return v*s
# Make data
def sphere_surface(r):
u = np.linspace(0, 2 * np.pi, 100)
v = np.linspace(0, np.pi, 100)
x = r * np.outer(np.cos(u), np.sin(v))
y = r * np.outer(np.sin(u), np.sin(v))
z = r * np.outer(np.ones(np.size(u)), np.cos(v))
return x,y,z
x,y,z = sphere_surface(1.5)
pot = potencial(x,y,z)
norm=matplotlib.colors.SymLogNorm(1,vmin=pot.min(),vmax=pot.max())
colors=plt.cm.coolwarm(norm(pot))
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
# Plot the surface
surf = ax.plot_surface(x,y,z, facecolors=colors,
linewidth=0, antialiased=False)
# Set up colorbar
sm = plt.cm.ScalarMappable(cmap=plt.cm.coolwarm, norm=norm)
sm.set_array(pot)
fig.colorbar(sm, shrink=0.5, aspect=5)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
plt.show()
Upvotes: 2