I would like to model basic crowd movement with python. I want to show an animation. I have made the following program to test it with matplotlib :
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
#size of the crowd
N = 100
def gen_data():
""" init position and speed of each people """
x = y = np.zeros(N)
theta = np.random.random(N) * 360 / (2 * np.pi)
v0 = 0.1
vx, vy = v0 * np.cos(theta), v0 * np.sin(theta)
return np.array([x, y, vx, vy]).T
def init():
for line in lines:
line.set_data([],[])
return line,
def update_lines(i, lines, data):
for d, line in zip(data, lines):
d[0:2] += d[2:4]
if abs(d[0]) > 5: d[2] *= -1
if abs(d[1]) > 5: d[3] *= -1
line.set_data(d[0] ,d[1])
return lines
fig = plt.figure()
ax = plt.axes(xlim=(-5,5),ylim=(-5,5))
lines = [plt.plot([],[], 'ko')[0] for i in range(N)]
data = gen_data()
anim = animation.FuncAnimation(fig, update_lines, init_func=init, fargs=(lines, data), interval=10, blit=True)
plt.show()
Even for N=100, the animation is slow... Is there something I can do to speed it up with mathplotlib ? Is matplotlib the best graphic tool to make thins kind of animation with python ? If no, what would it be ?
Here are 3 things you can do to make the animation faster:
Replace the N calls to plt.plot with one call to plt.scatter.
Replace the for-loop in update with assignments which modify whole slices of data at once:
data[:, 0:2] += data[:, 2:4]
data[:, 2] = np.where(np.abs(data[:, 0]) > 5, -data[:, 2], data[:, 2])
data[:, 3] = np.where(np.abs(data[:, 1]) > 5, -data[:, 3], data[:, 3])
Reduce interval=10 to interval=0.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
# size of the crowd
N = 100
def gen_data():
""" init position and speed of each people """
x = y = np.zeros(N)
theta = np.random.random(N) * 360 / (2 * np.pi)
v0 = 0.1
vx, vy = v0 * np.cos(theta), v0 * np.sin(theta)
return np.column_stack([x, y, vx, vy])
def init():
pathcol.set_offsets([[], []])
return pathcol,
def update(i, pathcol, data):
data[:, 0:2] += data[:, 2:4]
data[:, 2] = np.where(np.abs(data[:, 0]) > 5, -data[:, 2], data[:, 2])
data[:, 3] = np.where(np.abs(data[:, 1]) > 5, -data[:, 3], data[:, 3])
pathcol.set_offsets(data[:, 0:2])
return [pathcol]
fig = plt.figure()
ax = plt.axes(xlim=(-5, 5), ylim=(-5, 5))
pathcol = plt.scatter([], [])
data = gen_data()
anim = animation.FuncAnimation(fig, update, init_func=init,
fargs=(pathcol, data), interval=0, blit=True)
plt.show()
Related
I am trying to visualize water interference by matplotlib and numpy. But after the caculation is finished, the animation fps is too slow. I have found several ways to solve it but they do not work. I think the caculation process may not be that time costing and I have no idea why the fps is so low.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.animation as animation
"""
1.parameter setup
"""
L, sep, N, k, cut = 6, 2, 500, 5*2, 0.8 # 0.5倍寬度, 波源與原點距離, 每邊分割數量, 角波數, z軸範圍
fps, frn = 24, 50 # 每秒影格數量, 影格總數
x = np.linspace(-2*L, 2*L, N) # x axis
y = np.linspace(-L, L, N) # y axis
X, Y = np.meshgrid(x, y) # 2D matrix
j = complex(0, 1) # 根號 -1
"""
2.設定計算振幅的函數, 計算每個位置的振幅並存入陣列
"""
# 自訂函式, 計算每個位置對應的振幅
"""
def func(x, y, t):
r1 = np.sqrt(x**2 + y**2) # 點波源1
for i in range(250):
r1[i] = y[i]+L
r2 = np.sqrt((x+sep)**2 + y**2) # 點波源2
z = np.exp(j*k*r1)/k*3
for i in range(250, 499):
z[i] = np.exp(j*k*(r1[249]+r1[i]))/k*3
z_real = np.real(z*np.exp(-j*t))
return z_real # 回傳實部
"""
def func(x, y, t, n, a, b):
"""
r0 = np.sqrt((x-6)**2 + y**2)
r1 = np.sqrt((x-5)**2 + y**2)
...
r11 = np.sqrt((x+5)**2 + y**2)
r12 = np.sqrt((x+6)**2 + y**2)
"""
for i in range(n):
r = np.sqrt((x-a[i])**2 + (y-b[i])**2)
if(i == 0):
z = np.exp(j*k*r)/r
else:
z += np.exp(j*k*r)/r
return np.real(z*np.exp(-j*t)) # 回傳實部
Z = np.zeros((N, N, frn)) # 儲存振幅用的2維陣列
T = np.linspace(0, 2*np.pi, frn) # 儲存時間用的1維陣列
# user input
n=int(input("輸入波源數量:"))
print('點波源位置輸入格式: x , y')
a=[]
b=[]
for i in range(n):
temp=input(f"位置{i}:").split(',')
a.append(int(temp[0]))
b.append(int(temp[1]))
# 計算每個時刻每個位置對應的振幅, 有加cut效果較佳
for i in range(frn):
Z[:, :, i] = func(X, Y, T[i], n,a,b).clip(-cut, cut)
"""
3.繪圖
"""
fig = plt.figure(figsize=(7, 6), dpi=100) # 開啟繪圖視窗
ax = fig.gca()
ax.set_aspect(1.0) # 使圖片長寬變成1:1
# 以某個預設的colormap為基底, 修改成對應到 -cut ~ +cut 的colormap
mappable = plt.cm.ScalarMappable(cmap=plt.cm.jet)
mappable.set_array(np.arange(-cut, cut, 0.1))
# 在圖片右側加上color bar, 高度與圖片相同
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
plt.colorbar(mappable, cax=cax)
# 自訂函式, 先移除前一張圖, 再畫出下一張圖
def update(frame_number):
plot[0] = ax.contourf(X, Y, Z[:, :, frame_number], cmap=mappable.cmap, norm=mappable.norm)
# t = 0 的圖片
plot = [ax.contourf(X, Y, Z[:, :, 0], cmap=mappable.cmap, norm=mappable.norm)]
# 產生動畫, 目標為繪圖物件fig, 使用自訂函式update更新圖片, 圖片總數為frn, 時間間隔為interal, 單位為ms
ani = animation.FuncAnimation(fig, update, frn, interval=1000/fps)
plt.show() # 顯示圖片
# ani.save('TwoSourcesInterference2D.gif', writer='imagemagick', fps=fps) # 儲存圖片
I am expected to speed up the animation to about 12fps.
Trying to plot two separate animations, i.e. in different windows as separate figures. Running this code for me rightly creates two windows, but animates the data on the second figure at the same time. Closing figure 1 results in only the intended data for figure 2 being animated, removing the overlap from the data intended for figure 1. Closing figure 2 results in only the intended data for figure 1 being animated, removing the overlap from the data intended for figure 2.
Minimum code below:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
dx, dv, N, Nb, decp = 2, 1.5, 100, 12, int(1)
Pd = np.zeros([N + 1, 2 * Nb])
Vd = np.zeros([N + 1, 2 * Nb])
Pd[:, 1] = 4
Vd[:, 3] = 2
t = np.zeros(N + 1)
t[0] = 0
for i in range(0, N):
t[i + 1] = (i + 1) * 0.1
Px = []
for i in range(0, (2 * Nb)):
PX = dx * (-Nb + i) / 4
Px.append(PX)
lblx = []
for i in range(0, int((Nb / 2) + 1)):
if i == (Nb / 4):
LBL = r"$\mu_x$"
lblx.append(LBL)
else:
LBL = r"${0}\sigma_x$".format(-(Nb / 4) + i)
lblx.append(LBL)
Pv = []
for i in range(0, (2 * Nb)):
PV = dv * (-Nb + i) / 4
Pv.append(PV)
lblv = []
for i in range(0, int((Nb / 2) + 1)):
if i == (Nb / 4):
LBL = r"$\mu_v$"
lblv.append(LBL)
else:
LBL = r"${0}\sigma_v$".format(-(Nb / 4) + i)
lblv.append(LBL)
fig1 = plt.figure(figsize=(8,6))
def animatex(i):
fig1.clear()
plt.bar(Px, Pd[i, :], width = dx / 4, align = 'edge', color = 'b', \
label = 't = {} seconds'.format(round(t[i], decp)))
s_ticks = np.arange(-3 * dx, (3 + 1) * dx, dx)
plt.xticks(s_ticks, lblx)
plt.ylim(0, np.max(Pd))
plt.xlim(-3 * dx, 3 * dx)
plt.legend()
plt.draw()
anix = FuncAnimation(fig1, animatex, repeat = True, interval = 200, frames = N + 1)
fig2 = plt.figure(figsize=(8,6))
def animatev(i):
fig2.clear()
plt.bar(Pv, Vd[i, :], width = dv / 4, align = 'edge', color = 'b', \
label = 't = {} seconds'.format(round(t[i], decp)))
s_ticks = np.arange(-3 * dv, (3 + 1) * dv, dv)
plt.xticks(s_ticks, lblv)
plt.ylim(0, np.max(Vd))
plt.xlim(-3 * dv, 3 * dv)
plt.legend()
plt.draw()
aniv = FuncAnimation(fig2, animatev, repeat = True, interval = 200, frames = N + 1)
plt.show()
As is probably clear, they are two bar plots, with different vertical and horizontal dimensions. I've seen some solutions for these kinds of problems where the data shares an axis through a shared variable, but here they are not (as can be seen).
For this minimum code, the solution involves having the two bars, one in Pd and the other in Vd, being on their respective intended figures, not both on the second figure.
Let me know if there are any issues with the information here i.e. minimal code requirements not met, more information etc. and I will update.
Ignore any wayward writing style, it is not relevant.
Simplifying your code:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
dx, dv, N, Nb, decp = 2, 1.5, 10, 12, int(1)
Px = np.arange(Nb)
Pd = np.random.randn(N, Nb)
Vd = np.random.randn(N, Nb)
fig1, ax1 = plt.subplots(figsize=(8, 6))
def animatex(i):
ax1.clear()
ax1.bar(Px, Pd[i, :], width=dx / 4, align='edge', color='b')
anix = FuncAnimation(fig1, animatex, repeat=True, interval=200, frames=N)
fig2, ax2 = plt.subplots(figsize=(8, 6))
def animatev(i):
ax2.clear()
ax2.bar(Px, Vd[i, :], width = dv / 4, align='edge', color='b')
aniv = FuncAnimation(fig2, animatev, repeat=True, interval=200, frames=N)
plt.show()
works fine for me. You can add the esthetic/data details back in...
I am trying to animate multiple lines at once in matplotlib. To do this I am following the tutorial from the matplotlib.animation docs:
https://matplotlib.org/stable/api/animation_api.html
The idea in this tutorial is to create a line ln, = plt.plot([], []) and update the data of the line using ln.set_data in order to produce the animation. Whilst this all works fine when the line data is a 1 dimensional array (shape = (n,)) of n data points, I am having trouble when the line data is a 2 dimensional array (shape = (n,k)) of k lines to plot.
To be more precise, plt.plot accepts arrays as inputs, with each column corresponding to a new line to plot. Here is a simple example with 3 lines plotted with a single plt.plot call:
import matplotlib.pyplot as plt
import numpy as np
x = np.linspace(0, 2 * np.pi, 100).reshape(-1, 1)
x = np.concatenate([x] * 3, axis=1)
# generate 3 curves
y = np.copy(x)
y[:, 0] = np.cos(y[:, 0])
y[:, 1] = np.sin(y[:, 1] )
y[:, 2] = np.sin(y[:, 2] ) + np.cos(y[:, 2])
fig, ax = plt.subplots()
plt.plot(x,y)
plt.show()
However if I try to set the data using .set_data as required for generating animations I have a problem:
import matplotlib.pyplot as plt
import numpy as np
x = np.linspace(0, 2 * np.pi, 100).reshape(-1, 1)
x = np.concatenate([x] * 3, axis=1)
# generate 3 curves
y = np.copy(x)
y[:, 0] = np.cos(y[:, 0])
y[:, 1] = np.sin(y[:, 1] )
y[:, 2] = np.sin(y[:, 2] ) + np.cos(y[:, 2])
fig, ax = plt.subplots()
p, = plt.plot([], [], color='b')
p.set_data(x, y)
plt.show()
Is there a way to set_data for 2 dimensional arrays? Whilst I am aware that I could just create three plots p1, p2, p3 and call set_data on each of them in a loop, my real data consists of 1000-10,000 lines to plot, and this makes the animation too slow.
Many thanks for any help.
An approach could be to create a list of Line2D objects and use set_data in a loop. Note that ax.plot() always returns a list of lines, even when only one line is plotted.
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np
x = np.linspace(0, 2 * np.pi, 100)
# generate 10 curves
y = np.sin(x.reshape(-1, 1) + np.random.uniform(0, 2 * np.pi, (1, 10)))
fig, ax = plt.subplots()
ax.set(xlim=(0, 2 * np.pi), ylim=(-1.5, 1.5))
# lines = [ax.plot([], [], lw=2)[0] for _ in range(y.shape[1])]
lines = ax.plot(np.empty((0, y.shape[1])), np.empty((0, y.shape[1])), lw=2)
def animate(i):
for line_k, y_k in zip(lines, y.T):
line_k.set_data(x[:i], y_k[:i])
return lines
anim = FuncAnimation(fig, animate, frames=x.size, interval=200, repeat=False)
plt.show()
The array given by set_data() will be two one-dimensional arrays, so in this case three set_data() will be needed.
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np
x = np.linspace(0, 2 * np.pi, 100).reshape(-1, 1)
x = np.concatenate([x] * 3, axis=1)
# generate 3 curves
y = np.copy(x)
y[:, 0] = np.cos(y[:, 0])
y[:, 1] = np.sin(y[:, 1] )
y[:, 2] = np.sin(y[:, 2] ) + np.cos(y[:, 2])
fig, ax = plt.subplots()
ax = plt.axes(xlim=(0,6), ylim=(-1.5, 1.5))
line1, = ax.plot([], [], lw=2)
line2, = ax.plot([], [], lw=2)
line3, = ax.plot([], [], lw=2)
def animate(i):
line1.set_data(x[:i, 0], y[:i, 0])
line2.set_data(x[:i, 1], y[:i, 1])
line3.set_data(x[:i, 2], y[:i, 2])
return line1,line2,line3
anim = FuncAnimation(fig, animate, frames=100, interval=200, repeat=False)
plt.show()
I have the following code:
import numpy as np
import matplotlib.pyplot as plt
x = np.linspace(-np.pi/2, np.pi/2, 30)
y = np.linspace(-np.pi/2, np.pi/2, 30)
x,y = np.meshgrid(x,y)
z = np.sin(x**2+y**2)[:-1,:-1]
fig,ax = plt.subplots()
ax.pcolormesh(x,y,z)
Which gives this image:
Now lets say I want to highlight the edge certain grid boxes:
highlight = (z > 0.9)
I could use the contour function, but this would result in a "smoothed" contour. I just want to highlight the edge of a region, following the edge of the grid boxes.
The closest I've come is adding something like this:
highlight = np.ma.masked_less(highlight, 1)
ax.pcolormesh(x, y, highlight, facecolor = 'None', edgecolors = 'w')
Which gives this plot:
Which is close, but what I really want is for only the outer and inner edges of that "donut" to be highlighted.
So essentially I am looking for some hybrid of the contour and pcolormesh functions - something that follows the contour of some value, but follows grid bins in "steps" rather than connecting point-to-point. Does that make sense?
Side note: In the pcolormesh arguments, I have edgecolors = 'w', but the edges still come out to be blue. Whats going on there?
EDIT:
JohanC's initial answer using add_iso_line() works for the question as posed. However, the actual data I'm using is a very irregular x,y grid, which cannot be converted to 1D (as is required for add_iso_line().
I am using data which has been converted from polar coordinates (rho, phi) to cartesian (x,y). The 2D solution posed by JohanC does not appear to work for the following case:
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
def pol2cart(rho, phi):
x = rho * np.cos(phi)
y = rho * np.sin(phi)
return(x, y)
phi = np.linspace(0,2*np.pi,30)
rho = np.linspace(0,2,30)
pp, rr = np.meshgrid(phi,rho)
xx,yy = pol2cart(rr, pp)
z = np.sin(xx**2 + yy**2)
scale = 5
zz = ndimage.zoom(z, scale, order=0)
fig,ax = plt.subplots()
ax.pcolormesh(xx,yy,z[:-1, :-1])
xlim = ax.get_xlim()
ylim = ax.get_ylim()
xmin, xmax = xx.min(), xx.max()
ymin, ymax = yy.min(), yy.max()
ax.contour(np.linspace(xmin,xmax, zz.shape[1]) + (xmax-xmin)/z.shape[1]/2,
np.linspace(ymin,ymax, zz.shape[0]) + (ymax-ymin)/z.shape[0]/2,
np.where(zz < 0.9, 0, 1), levels=[0.5], colors='red')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
This post shows a way to draw such lines. As it is not straightforward to adapt to the current pcolormesh, the following code demonstrates a possible adaption.
Note that the 2d versions of x and y have been renamed, as the 1d versions are needed for the line segments.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
x = np.linspace(-np.pi / 2, np.pi / 2, 30)
y = np.linspace(-np.pi / 2, np.pi / 2, 30)
xx, yy = np.meshgrid(x, y)
z = np.sin(xx ** 2 + yy ** 2)[:-1, :-1]
fig, ax = plt.subplots()
ax.pcolormesh(x, y, z)
def add_iso_line(ax, value, color):
v = np.diff(z > value, axis=1)
h = np.diff(z > value, axis=0)
l = np.argwhere(v.T)
vlines = np.array(list(zip(np.stack((x[l[:, 0] + 1], y[l[:, 1]])).T,
np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
l = np.argwhere(h.T)
hlines = np.array(list(zip(np.stack((x[l[:, 0]], y[l[:, 1] + 1])).T,
np.stack((x[l[:, 0] + 1], y[l[:, 1] + 1])).T)))
lines = np.vstack((vlines, hlines))
ax.add_collection(LineCollection(lines, lw=1, colors=color))
add_iso_line(ax, 0.9, 'r')
plt.show()
Here is an adaption of the second answer, which can work with only 2d arrays:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from scipy import ndimage
x = np.linspace(-np.pi / 2, np.pi / 2, 30)
y = np.linspace(-np.pi / 2, np.pi / 2, 30)
x, y = np.meshgrid(x, y)
z = np.sin(x ** 2 + y ** 2)
scale = 5
zz = ndimage.zoom(z, scale, order=0)
fig, ax = plt.subplots()
ax.pcolormesh(x, y, z[:-1, :-1] )
xlim = ax.get_xlim()
ylim = ax.get_ylim()
xmin, xmax = x.min(), x.max()
ymin, ymax = y.min(), y.max()
ax.contour(np.linspace(xmin,xmax, zz.shape[1]) + (xmax-xmin)/z.shape[1]/2,
np.linspace(ymin,ymax, zz.shape[0]) + (ymax-ymin)/z.shape[0]/2,
np.where(zz < 0.9, 0, 1), levels=[0.5], colors='red')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
plt.show()
I'll try to refactor add_iso_line method in order to make it more clear an open for optimisations. So, at first, there comes a must-do part:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
x = np.linspace(-np.pi/2, np.pi/2, 30)
y = np.linspace(-np.pi/2, np.pi/2, 30)
x, y = np.meshgrid(x,y)
z = np.sin(x**2+y**2)[:-1,:-1]
fig, ax = plt.subplots()
ax.pcolormesh(x,y,z)
xlim, ylim = ax.get_xlim(), ax.get_ylim()
highlight = (z > 0.9)
Now highlight is a binary array that looks like this:
After that we can extract indexes of True cells, look for False neighbourhoods and identify positions of 'red' lines. I'm not comfortable enough with doing it in a vectorised manner (like here in add_iso_line method) so just using simple loop:
lines = []
cells = zip(*np.where(highlight))
for x, y in cells:
if x == 0 or highlight[x - 1, y] == 0: lines.append(([x, y], [x, y + 1]))
if x == highlight.shape[0] or highlight[x + 1, y] == 0: lines.append(([x + 1, y], [x + 1, y + 1]))
if y == 0 or highlight[x, y - 1] == 0: lines.append(([x, y], [x + 1, y]))
if y == highlight.shape[1] or highlight[x, y + 1] == 0: lines.append(([x, y + 1], [x + 1, y + 1]))
And, finally, I resize and center coordinates of lines in order to fit with pcolormesh:
lines = (np.array(lines) / highlight.shape - [0.5, 0.5]) * [xlim[1] - xlim[0], ylim[1] - ylim[0]]
ax.add_collection(LineCollection(lines, colors='r'))
plt.show()
In conclusion, this is very similar to JohanC solution and, in general, slower. Fortunately, we can reduce amount of cells significantly, extracting contours only using python-opencv package:
import cv2
highlight = highlight.astype(np.uint8)
contours, hierarchy = cv2.findContours(highlight, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
cells = np.vstack(contours).squeeze()
This is an illustration of cells being checked:
This code:
def complex_to_rgb(complex_data, invert=False):
from numpy import angle, max, pi, sin, zeros
phase = angle(complex_data)
amplitude = abs(complex_data)
amplitude = amplitude/max(max(amplitude))
A = zeros((complex_data.shape[0], complex_data.shape[1], 3))
A[:,:,0] = .5*(sin(phase)+1)*amplitude
A[:,:,1] = .5*(sin(phase+pi/2)+1)*amplitude
A[:,:,2] = .5*(-sin(phase)+1)*amplitude
if(invert):
return 1-A
else:
return A
import numpy as np
from matplotlib.pyplot import figure
N = 1024
x = np.linspace(-1, 1, N)
y = np.linspace(-1, 1, N)
X,Y = np.meshgrid(x,y)
R = np.sqrt(X*X + Y*Y)
PHI = np.arctan2(Y, X)
fig = figure()
ax = fig.add_subplot(212, polar=True)
ax.imshow(complex_to_rgb(R*np.exp(1j*PHI) * (R<1), invert=True))
ax.set_xticks([-.5, 0, np.pi/2, np.pi, 3*np.pi/2])
ax.set_yticks([0, N/3, 2*N/3, N])
ax.set_xticklabels(['', '$0$', r'$\pi/2$', r'$\pi$', r'$3\pi/2$'])
ax.set_yticklabels([])
fig.show()
Generates a nice HSV legend plot. Now I'd like to remove the -.5 xtick, but that seems to mess everything up. Anyone know how to fix this? I already reported it as a bug
As described in the bug report, I can place the radial axis anywhere I want by specifying an explicit extent to imshow. Additionally, rgrids can be used to fix the angle of the tick labels.
def complex_to_rgb(complex_data, invert=False):
from numpy import angle, max, pi, sin, zeros
phase = angle(complex_data)
amplitude = abs(complex_data)
amplitude = amplitude/max(max(amplitude))
A = zeros((complex_data.shape[0], complex_data.shape[1], 3))
A[:,:,0] = .5*(sin(phase)+1)*amplitude
A[:,:,1] = .5*(sin(phase+pi/2)+1)*amplitude
A[:,:,2] = .5*(-sin(phase)+1)*amplitude
if(invert):
return 1-A
else:
return A
import numpy as np
from matplotlib.pyplot import figure
N = 1024
x = np.linspace(-1, 1, N)
y = np.linspace(-1, 1, N)
X,Y = np.meshgrid(x,y)
R = np.sqrt(X*X + Y*Y)
PHI = np.arctan2(Y, X)
fig = figure()
ax = fig.add_subplot(111, polar=True)
ax.imshow(complex_to_rgb(R*np.exp(1j*PHI) * (R<1), invert=True), extent=[0,2*np.pi, 0,1024])
ax.set_rgrids([1,N/3,2*N/3], angle=45)
ax.set_xticks([0, np.pi/2, np.pi, 3*np.pi/2])
ax.set_yticks([0, N/3, 2*N/3, N])
ax.set_xticklabels([r'$0$', r'$\pi/2$', r'$\pi$', r'$3\pi/2$'])
ax.set_yticklabels([r'0', r'$1/3$', r'$2/3$', '1'])
fig.show()
Which results in: