pythonmatplotlibsurface

How to replicate the following density plot in Python?

Given the following setup,

N_r = 21; N_theta = 18; N_phi= 36;
r_index = N_r-1;

[phi,theta,r_sphere] = np.meshgrid(np.linspace(0,2*np.pi,N_phi),np.linspace(0,np.pi,N_theta),np.linspace(a,b,N_r));

X = r_sphere[:,:,r_index] * np.sin(theta[:,:,r_index]) * np.cos(phi[:,:,r_index]);
Y = r_sphere[:,:,r_index] * np.sin(theta[:,:,r_index]) * np.sin(phi[:,:,r_index]);
Z = r_sphere[:,:,r_index] * np.cos(theta[:,:,r_index]);

rho = 1/r_sphere**2*np.sin(theta)*np.cos(theta)*np.sin(phi)

I have set up my 2D X, Y, and Z coordinates converted from spherical coordinates, and a density variable in spherical coordinates that I want to plot a spherical shell (at r_index) of. In Matlab, with the same variables and setup, I was able to use the surf() function

surf(X,Y,Z,rho(:,:,r_index),"EdgeAlpha",0.2); (plus a couple other things like axis labels and colorbar), I was able to create the following 3D (or I guess 4D?) plot at r=r_index:

enter image description here

Trying to use Matplotlib's plot_surface() either doesn't work, or I'm not quite getting my inputs correct:

fig1 = plt.figure(figsize=(16,9),dpi=80)
ax = fig1.add_subplot(projection = "3d")
surf = ax.plot_surface(X,Y,Z,rho[:,:,r_index])

Can I get this to work using plot_surface(), or is there another plotting function tailor made for what it is that I'm trying to do?

EDIT: scatter() appears to give me something close

fig1 = plt.figure(figsize=(16,9),dpi=80)
ax = fig1.add_subplot(projection = "3d")
surf = ax.scatter(X,Y,Z,c=rho[:,:,r_index],cmap = mpl.colormaps['bwr'])
plt.colorbar(surf,ax = ax,shrink = 0.5,aspect = 5)
plt.show()

It cleverly sets up the surface shape by taking in X, Y, and Z, and then I can set c to color each dot in accordance to the range of values passed in. If I could just do something similar to that for a function like plot_surface(). It can probably be done, I just don't know what the exact input arguments would be.

Solution

  • Try this:

    a, b = 0, 100 # use desired values
rho[np.isnan(rho)] = 0 # impute NaN values
density = rho[...,r_index] # set appropriate density
# normalize to have values in [0,1]
density = (density - density.min()) / (density.max() - density.min()) 
m = plt.cm.ScalarMappable( \
        norm=mpl.colors.Normalize(density.min(), density.max()),  \
        cmap='jet')

fig1 = plt.figure(figsize=(16,9),dpi=80)
ax = fig1.add_subplot(projection = "3d")
surf = ax.plot_surface(X,Y,Z, facecolors=m.to_rgba(density))
fig1.colorbar(m, ax=ax) # colorbar
plt.show()

    enter image description here

    complete code:

    import numpy as np

import matplotlib.pyplot as plt
import matplotlib as mpl

N_r = 21
N_theta = 18
N_phi= 36
r_index = N_r-1
a, b = 0, 100 # use desired values

[phi,theta,r_sphere] = np.meshgrid(np.linspace(0,2*np.pi,N_phi),np.linspace(0,np.pi,N_theta),np.linspace(a,b,N_r))

X = r_sphere[:,:,r_index] * np.sin(theta[:,:,r_index]) * np.cos(phi[:,:,r_index])
Y = r_sphere[:,:,r_index] * np.sin(theta[:,:,r_index]) * np.sin(phi[:,:,r_index])
Z = r_sphere[:,:,r_index] * np.cos(theta[:,:,r_index])

rho = 1/r_sphere**2*np.sin(theta)*np.cos(theta)*np.sin(phi)

#a, b = 0, 100 # use desired values
rho[np.isnan(rho)] = 0 # impute NaN values
density = rho[...,r_index] # set appropriate density
# normalize to have values in [0,1]
density = (density - density.min()) / (density.max() - density.min()) 
m = plt.cm.ScalarMappable( \
        norm=mpl.colors.Normalize(density.min(), density.max()),  \
        cmap='jet')

fig1 = plt.figure(figsize=(16,9),dpi=80)
ax = fig1.add_subplot(projection = "3d")
surf = ax.plot_surface(X,Y,Z, facecolors=m.to_rgba(density))
fig1.colorbar(m, ax=ax) # colorbar
plt.show()