adding line plot to imshow and changing axis marker

Physicist picture Physicist · Dec 30, 2015 · Viewed 10.4k times · Source

I have made the attached plot using the following codes:

a = 1
theta = np.linspace(0,2*np.pi,101)
x = np.linspace(-3*a,3*a,1001, dtype='complex')
y = np.linspace(-3*a,3*a,1001, dtype='complex')
X,Y = np.meshgrid(x,y)

# come manipulations with V
# (same shape and type as X,Y) not shown here 

plt.subplot(1,2,1)
plt.scatter(a*np.cos(theta), a*np.sin(theta))
plt.imshow(V.real)
plt.colorbar()
plt.subplot(1,2,2)
plt.scatter(a*np.cos(theta), a*np.sin(theta))
plt.imshow(V.imag)
plt.colorbar()

What I want to do is that:

1) change the scale of the plot such that the horizontal and vertical axis varies between -3*a and 3*a

2) plot the circle boundary (centered at the origin with radius = a). Now it appears at the top left, as the scale of the plot is changed from [-3*a,3*a] to that of the size of the array.enter image description here

Answer

Joe Kington picture Joe Kington · Dec 30, 2015

In general, you're looking for the extent kwarg to imshow.

As a quick example:

import numpy as np
import matplotlib.pyplot as plt

data = np.random.random((10, 10))

fig, ax = plt.subplots()
ax.imshow(data, extent=[10, 30, np.pi, -2*np.pi])
plt.show()

enter image description here

In the case of the example you gave:

import numpy as np
import matplotlib.pyplot as plt

a = 1
theta = np.linspace(0, 2*np.pi, 100)

# We could replace the next three lines with:
# y, x = np.mgrid[-3*a:3*a:1000j, -3*a:3*a:1000j]
x = np.linspace(-3*a, 3*a, 1000)
y = np.linspace(-3*a, 3*a, 1000)
x, y = np.meshgrid(x, y)

# Now let's make something similar to your V for this example..
r = np.hypot(x, y)
V = np.cos(3*np.arctan2(y, x)) + np.sin(r) + np.cos(x)*1j * np.cos(r)

def plot(ax, data):
    ax.plot(a*np.cos(theta), a*np.sin(theta), color='black')
    im = ax.imshow(data, extent=[x.min(), x.max(), y.max(), y.min()])
    fig.colorbar(im, ax=ax, shrink=0.5)

fig, (ax1, ax2) = plt.subplots(ncols=2)

ax1.set(title='Real Portion')
plot(ax1, V.real)

ax2.set(title='Imaginary Portion')
plot(ax2, V.imag)

plt.show()

enter image description here

However, you might also consider using pcolormesh in this case. For example, we could change the plot function to:

def plot(ax, data):
    ax.plot(a*np.cos(theta), a*np.sin(theta), color='black')
    im = ax.pcolormesh(x, y, data)
    ax.set(aspect=1)
    fig.colorbar(im, ax=ax, shrink=0.5)

The main differences are:

  1. imshow can interpolate, while pcolormesh gives vector output and can't interpolate (i.e. it plots lots of rectangles instead of an image).
  2. pcolormesh is somewhat slower, so for large images, imshow is a better choice.
  3. imshow and pcolormesh treat the extents slightly differently. imshow is "cell-centered" while pcolormesh is "mesh-centered". This is a half-pixel difference, so you can ignore it in this case.
  4. imshow will set the aspect of the plot to 1, so that one unit in the x-direction is the same size as one unit in the y-direction. It also flips the y-axis, by default.

One other note: If you'd prefer not to have the y-axis flipped, either call ax.invert_yaxis() or use origin='lower' and extent=[xmin, xmax, ymin, ymax].