与matplotlib对称的streamplot

Question

I'm trying to plot the streamlines of a magnetic field around a sphere using matplotlib, and it does work quite nicely. However, the resulting image is not symmetric, but it should be (I think).

This is the code used to generate the image. Excuse the length, but I thought it would be better than just posting a non-working snippet. Also, it's not very pythonic; that's because I converted it from Matlab, which was easier than I expected.

from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle

def cart2spherical(x, y, z):
    r = np.sqrt(x**2 + y**2 + z**2)
    phi = np.arctan2(y, x)
    theta = np.arccos(z/r)
    if r == 0:
        theta = 0
    return (r, theta, phi)

def S(theta, phi):
    S = np.array([[np.sin(theta)*np.cos(phi), np.cos(theta)*np.cos(phi), -np.sin(phi)],
                  [np.sin(theta)*np.sin(phi), np.cos(theta)*np.sin(phi),  np.cos(phi)],
                  [np.cos(theta),             -np.sin(theta),             0]])
    return S

def computeB(r, theta, phi, a=1, muR=100, B0=1):
    delta = (muR - 1)/(muR + 2)
    if r > a:
        Bspherical = B0*np.array([np.cos(theta) * (1 + 2*delta*a**3 / r**3),
                                  np.sin(theta) * (delta*a**3 / r**3 - 1),
                                  0])
        B = np.dot(S(theta, phi), Bspherical)
    else:
        B = 3*B0*(muR / (muR + 2)) * np.array([0, 0, 1])
    return B

Z, X = np.mgrid[-2.5:2.5:1000j, -2.5:2.5:1000j]
Bx = np.zeros(np.shape(X))
Bz = np.zeros(np.shape(X))
Babs = np.zeros(np.shape(X))
for i in range(len(X)):
    for j in range(len(Z)):
        r, theta, phi = cart2spherical(X[0, i], 0, Z[j, 0])
        B = computeB(r, theta, phi)
        Bx[i, j], Bz[i, j] = B[0], B[2]
        Babs[i, j] = np.sqrt(B[0]**2 + B[1]**2 + B[2]**2)

fig=plt.figure()
ax=fig.add_subplot(111)

plt.streamplot(X, Z, Bx, Bz, color='k', linewidth=0.8*Babs, density=1.3,
               minlength=0.9, arrowstyle='-')
ax.add_patch(Circle((0, 0), radius=1, facecolor='none', linewidth=2))
plt.axis('equal')
plt.axis('off')
fig.savefig('streamlines.pdf', transparent=True, bbox_inches='tight', pad_inches=0)

Answer 1

Quoting from the documentation:

 density : float or 2-tuple Controls the closeness of streamlines. When density = 1, the domain is divided into a 25x25 grid—density linearly scales this grid. Each cell in the grid can have, at most, one traversing streamline. For different densities in each direction, use [density_x, density_y]. 

so you are getting aliasing effects between the cells it uses to decide where the stream lines are, and the symmetries of your problem. You need to carefully choose your grid size (of the data) and the density.

It is also sensitive to where the box boundaries are relative to the top of the sphere. Is the center of your sphere on a data grid point or between the data grid points? If it is on a grid point then the box that contains the center point will be different than the boxes adjacent to it.

I am not familiar with exactly how it decides which stream lines to draw, but I could imagine that it is some sort of greedy algorithm and hence will give different results walking towards the high density region and away density region.

To be clear, you issue is not that the stream lines are wrong , they are valid stream lines, it is that you find the result not aesthetically pleasing.

Answer 2

First of all, for curiosity, why would you want to plot symmetric data? Why plotting half of isn't fine?

Said that, this is a possible hack. You can use mask arrays as Hooked suggested to plot half of it:

mask = X>0
BX_OUT = Bx.copy()
BZ_OUT = Bz.copy()
BX_OUT[mask] = None
BZ_OUT[mask] = None
res = plt.streamplot(X, Z, BX_OUT, BZ_OUT, color='k', 
           arrowstyle='-',linewidth=1,density=2)

then you save in res the result from streamplot, extract the lines and plot them with the opposite X coordinate.

lines = res.lines.get_paths()
for l in lines:
    plot(-l.vertices.T[0],l.vertices.T[1],'k')

I used this hack to extract streamlines and arrows from a 2D plot, then apply a 3D transformation and plot it with mplot3d. A picture is in one of my questions here .

Answer 3

Use a mask to separate the two regions of interest:

mask = np.sqrt(X**2+Z**2)<1

BX_OUT = Bx.copy()
BZ_OUT = Bz.copy()
BX_OUT[mask] = None
BZ_OUT[mask] = None
plt.streamplot(X, Z, BX_OUT, BZ_OUT, color='k', 
               arrowstyle='-', density=2)

BX_IN = Bx.copy()
BZ_IN = Bz.copy()
BX_IN[~mask] = None
BZ_IN[~mask] = None
plt.streamplot(X, Z, BX_IN, BZ_IN, color='r', 
               arrowstyle='-', density=2)

The resulting plot isn't exactly symmetric, but by giving the algorithm a hint, it's far closer than what you had before. Play with the density of the grid via meshgrid and the density parameter to achieve the effect you are looking for.

Answer 4

Use physics, instead... The magnetic field is symmetrical with respect to the z (vertical) axis! So you just need two streamplot 's:

plt.streamplot(X, Z, Bx, Bz, color='k', linewidth=0.8*Babs, density=1.3, minlength=0.9, arrowstyle='-')
plt.streamplot(-X, Z, -Bx, Bz, color='k', linewidth=0.8*Babs, density=1.3, minlength=0.9, arrowstyle='-')