fbartolic fbartolic - 2 months ago 19
Python Question

Matplotlib - contour and quiver plot in projected polar coordinates

I need to plot contour and quiver plots of scalar and vector fields defined on an uneven grid in (r,theta) coordinates.

As a minimal example of the problem I have, consider the contour plot of a Stream function for a magnetic dipole, contours of such a function are streamlines of the corresponeding vector field (in this case, the magnetic field).

The code below takes an uneven grid in (r,theta) coordinates, maps it to the cartesian plane and plots a contour plot of the stream function.

import numpy as np
import matplotlib.pyplot as plt

r = np.logspace(0,1,200)
theta = np.linspace(0,np.pi/2,100)

N_r = len(r)
N_theta = len(theta)

# Polar to cartesian coordinates
theta_matrix, r_matrix = np.meshgrid(theta, r)
x = r_matrix * np.cos(theta_matrix)
y = r_matrix * np.sin(theta_matrix)

m = 5
psi = np.zeros((N_r, N_theta))

# Stream function for a magnetic dipole
psi = m * np.sin(theta_matrix)**2 / r_matrix

contour_levels = m * np.sin(np.linspace(0, np.pi/2,40))**2.

fig, ax = plt.subplots()
# ax.plot(x,y,'b.') # plot grid points
ax.set_aspect('equal')
ax.contour(x, y, psi, 100, colors='black',levels=contour_levels)
plt.show()


For some reason though, the plot I get doesn't look right:
Contour plot of a stream function for a magnetic dipole.

If I interchange x and y in the contour function call, I get the desired result:
enter image description here

Same thing happens when I try to make a quiver plot of a vector field defined on the same grid and mapped to the x-y plane, except that interchanging x and y in the function call no longer works.

It seems like I made a stupid mistake somewhere but I can't figure out what it is.

Answer

If psi = m * np.sin(theta_matrix)**2 / r_matrix then psi increases as theta goes from 0 to pi/2 and psi decreases as r increases.

So a contour line for psi should increase in r as theta increases. That results in a curve that goes counterclockwise as it radiates out from the center. This is consistent with the first plot you posted, and the result returned by the first version of your code with

ax.contour(x, y, psi, 100, colors='black',levels=contour_levels)

An alternative way to confirm the plausibility of the result is to look at a surface plot of psi:

import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as axes3d

r = np.logspace(0,1,200)
theta = np.linspace(0,np.pi/2,100)

N_r = len(r)
N_theta = len(theta)

# Polar to cartesian coordinates
theta_matrix, r_matrix = np.meshgrid(theta, r)
x = r_matrix * np.cos(theta_matrix)
y = r_matrix * np.sin(theta_matrix)

m = 5

# Stream function for a magnetic dipole
psi = m * np.sin(theta_matrix)**2 / r_matrix

contour_levels = m * np.sin(np.linspace(0, np.pi/2,40))**2.

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')
ax.set_aspect('equal')

ax.plot_surface(x, y, psi, rstride=8, cstride=8, alpha=0.3)
ax.contour(x, y, psi, colors='black',levels=contour_levels)
plt.show()

enter image description here