How to replicate the following density plot in Python?

Question

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:

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.

Answer 1

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

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

Answer 2

Ah, so it's the same as with using scatter where we create a colormap out of the data and applying it. The only issue now is that the colorbar's min and max values aren't reflective of the of the data, as it's locked to 0 to 1. I'm guessing that's what the Normalize function does. Tried fiddling around with the inputs adding vmin and vmax and applying clim, but it looks like it's stuck. Is there a way around this?

Not sure I am getting the comment to the accepted answer right , but copying from:

color coding using scalar mappable in matplotlib

I tried:

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)

#print('\nrho :\n', rho)

#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')
  
cmap= plt.cm.jet
# create normalization instance
norm = mpl.colors.Normalize(vmin=density.min(), vmax=density.max()) 
# create a scalarmappable from the colormap
m = mpl.cm.ScalarMappable(cmap=cmap, norm=norm)    

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

output:

given that:

print('\ndensity.max() : ', density.max())

print('\ndensity.min() : ', density.min())

results in:

density.max() :  4.973657691184833e-05

density.min() :  -4.973657691184834e-05

I am using Online Matplotlib Compiler that is Matplotlib Version : 3.8.4