¿Cómo se puede asignar un mapa de colores de superficie 3D a una función escalar?

Tengo una función escalar que representa el potencial eléctrico en una superficie esférica. Quiero trazar, para un radio determinado, la superficie y vincular sus puntos a un mapa de colores basado en la función potencial.

¿Cómo mapeo esa función escalar al mapa de colores en la superficie? Sospecho que debe estar en los argumentos pasados ​​a la función ax.plot_surface . Intenté usar el argumento: facecolors=potencial(x,y,z) , pero me dio un ValueError: Invalid RGBA argument . Mirando el código fuente del tercer ejemplo , hay:

 # 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)] 

Lo cual no entiendo, ni tengo una idea de cómo enlazar a una función escalar.

Mi código

 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() 

En principio, hay dos formas de colorear una plot de superficie en matplotlib.

  1. Usa el argumento cmap para especificar un mapa de colores. En este caso, el color se elegirá de acuerdo con la matriz z . En caso de que no se desee,
  2. Usa el argumento de facecolors . Esto espera una variedad de colores de la misma forma que z .

Entonces, en este caso, debemos elegir la opción 2 y crear una matriz de colores. Para este fin, uno puede elegir un mapa de colores. Un mapa de colores mapea valores entre 0 y 1 a un color. Como el potencial tiene valores muy por encima y por debajo de este rango, es necesario normalizarlos en el rango [0,1].
Matplotlib ya proporciona alguna función auxiliar para realizar esta normalización y, dado que el potencial tiene una dependencia de 1 / x, una escala de colores logarítmica puede ser adecuada.

Al final, los coleccionistas pueden recibir una matriz.

 colors = cmap(norm(potential(...))) 

El bit que falta es ahora la barra de colores. Para que la barra de colores se vincule a los colores de la gráfica de superficie, debemos configurar manualmente un ScalarMappable con el mapa de colores y la instancia de normalización, que luego podemos suministrar a la barra de colores.

 sm = plt.cm.ScalarMappable(cmap=plt.cm.coolwarm, norm=norm) sm.set_array(pot) fig.colorbar(sm, shrink=0.5, aspect=5) 

Aquí está el ejemplo completo.

 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() 

introduzca la descripción de la imagen aquí