Overlay 2D stream plot on 3D surface plot

Question

I'm working on a 3D axis-symmetric system (here for simplicity I'm showing a sphere). I can solve the problem in the x-z plane to find a vector field like the one in the following stream plot:

Given the symmetry of the problem, I can rotate this solution around the z-axis to get the full 3D vector field. Now, I would like to somehow make a 3D plot of this solution. The most straightforward way I could think of is to make a surface plot of the system without the x>0 && y<0 region, like this:

And then overlay the previous stream plot on the (y=0 && x>0) and (x=0 && y<0) planes resulting from the "cut-out" of the surface. Any suggestions on how to achieve this? Or is there maybe another better way to visualise the 3D result?

Here's a MWC that generates the figures (the x_Car and m_Car used are just an example):

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from scipy.interpolate import griddata

ell = 0.62
gamma = 1
Rbase, Rcap, center = 0.62, 0.62, 0

###################
###### Stream plot

res = 200

x_Car = [[-1.8417224772056704e-16, -0.0, -0.5726312223685197],
 [7.597102430294396e-17, 0.0, -0.6203504908994001],
 [0.020037100834806584, 0.0, -0.5857436781333181],
 [0.030439155856322415, 0.0, -0.6196032515649681],
 [0.020156894294067102, 0.0, -0.5477190352848064],
 [-1.6882456041048852e-16, -0.0, -0.5249119538377125]]
m_Car = [[-1.60647705e-05,  0.00000000e+00, -2.43349475e-01],
 [ 0.00022923,  0.        , -0.21359168],
 [ 0.00627664,  0.        , -0.23712676],
 [ 0.01131077,  0.        , -0.21533309],
 [ 0.00655443,  0.        , -0.25987232],
 [ 2.65801174e-05,  0.00000000e+00, -2.72980539e-01]]

m_Car_T, x_Car_T = np.array(m_Car).T, np.array(x_Car).T

cm = matplotlib.cm.coolwarm
norm = matplotlib.colors.Normalize(vmin=0,vmax=1)
sm = matplotlib.cm.ScalarMappable(cmap=cm, norm=norm)

xx, zz = x_Car_T[0], x_Car_T[2]
x_grid, z_grid = np.meshgrid(np.linspace(xx.min(),xx.max(),res),np.linspace(zz.min(),zz.max(),res))

x_vector_interp = griddata((xx, zz), m_Car_T[0], (x_grid, z_grid), method='cubic')
z_vector_interp = griddata((xx, zz), m_Car_T[2], (x_grid, z_grid), method='cubic')

width, height = matplotlib.figure.figaspect(1.5)*1.2
fig, ax = plt.subplots(1,1,figsize=(width,height))
    
ax.pcolormesh(x_grid/ell, z_grid/ell, np.sqrt(x_vector_interp**2+z_vector_interp**2), norm=norm, cmap=cm)
ax.streamplot(x_grid/ell, z_grid/ell, x_vector_interp, z_vector_interp, 
              color='k', density=1.1) #, maxlength=6, integration_direction='both'

xrange = np.linspace(0,Rbase,100)
ax.plot(xrange/ell,-np.sqrt(Rbase**2-xrange**2)/ell,'k')
zcap = np.sign(gamma)*Rcap+center
ax.plot(xrange/ell,(np.sign(gamma)*np.sqrt(Rcap**2-xrange**2)+center)/ell,'k')
ax.plot((0,0),(-Rbase/ell,zcap/ell),'k')

ax.set_ylabel('$z/\ell$'); ax.set_ylim(-(Rbase+0.1)/ell,(Rbase+0.1)/ell)
ax.set_xlabel('$x/\ell$'); ax.set_xlim(-0.1/ell,(Rbase+0.1)/ell)
plt.colorbar(sm,label=r'$|\mathbf{m}|$')
# plt.axis('off')
# plt.axis("image") 

plt.show()

###################
###### Surface plot

width, height = matplotlib.figure.figaspect(1.2)*1.5
fig, ax = plt.subplots(1, 1, figsize=(width, height), subplot_kw={"projection": "3d"})

### Plot the cap
theta, phi = np.mgrid[0:np.pi/2:180j, 0:1.5*np.pi:270j] # phi = alti, theta = azi
X = Rcap*np.sin(theta)*np.cos(phi)
Y = Rcap*np.sin(theta)*np.sin(phi)
Z = Rcap*np.cos(theta)+center
for i in range(len(theta)):
    for j in range(len(phi[0])):
        if Z[i,j] < 0:
            Z[i,j] = np.nan
        else:
            pass
ax.plot_surface(X/ell, Y/ell, Z/ell, linewidth=0, antialiased=False, rstride=1, cstride=1, color='grey')

### Plot the base
theta, phi = np.mgrid[np.pi/2:np.pi:180j, 0:1.5*np.pi:270j] # phi = alti, theta = azi
X = Rbase*np.sin(theta)*np.cos(phi)
Y = Rbase*np.sin(theta)*np.sin(phi)
Z = Rbase*np.cos(theta)
ax.plot_surface(X/ell, Y/ell, Z/ell, linewidth=0, antialiased=False, rstride=1, cstride=1, color='gainsboro')

ax.plot([0,0],[0,0],[-Rbase/ell,(Rcap+center)/ell],'k-',zorder=100)

ax.set_xlabel('$x/\ell$'); ax.set_ylabel('$y/\ell$'); ax.set_zlabel('$z/\ell$')
ax.set_box_aspect([1,1,(2*Rbase)/(Rbase+Rcap-center)])

plt.show()

---

UPDATE

@Jared had suggested a way to do this with a quiver plot. Although this could be a way if nothing else is possible, I would really prefer to have the stream plot, as it shows the data a lot better (some of my fields are a lot more intricate and the continuous lines help visualise them much better than the discrete arrows). Is there any way to achieve this?

Answer 1

To do this, you can use a colored surface plot to show the contour and a quiver plot to show the vectors. I had to adjust the zorder to get the plots in front of each other. The downside of that is the layering effect no longer works if you have an interactive figure and rotate around it.

import numpy as np
import matplotlib.pyplot as plt

plt.close("all")

N = 20
theta = np.linspace(0, np.pi, N)
phi = np.linspace(0, 3*np.pi/2, N)
Theta, Phi = np.meshgrid(theta, phi)

X = np.sin(Theta)*np.cos(Phi)
Y = np.sin(Theta)*np.sin(Phi)
Z = np.cos(Theta)

r = np.linspace(0, 1, N)
theta = np.linspace(-np.pi/2, np.pi/2, N)
R_vector, Theta_vector= np.meshgrid(r, theta)

X_vector = R_vector*np.cos(Theta_vector)
Z_vector= R_vector*np.sin(Theta_vector)
Y_vector= np.zeros_like(X_vector)

U = X_vector+ 2*Z_vector
V = -Z_vector
W = np.zeros_like(U)

magnitudes = np.sqrt(U**2 + V**2 + W**2)

U /= 10*magnitudes.max()
V /= 10*magnitudes.max()
W /= 10*magnitudes.max()


fig, ax = plt.subplots(subplot_kw={"projection":"3d", "computed_zorder":False})
ax.plot_surface(X, Y, Z, color="gainsboro", zorder=1)
ax.plot_surface(X_vector, Y_vector, Z_vector, 
                facecolors=plt.cm.coolwarm(magnitudes), zorder=2)
ax.quiver(X_vector, Y_vector, Z_vector, U, V, W, color="k", 
          antialiased=True, lw=0.5)
ax.set_aspect("equal")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
fig.tight_layout()