How to draw polar hist2d/hexbin in matplotlib?

Question

I have a random vector (random length and random angle) and would like to plot its approximate PDF (probability density function) via hist2d or hexbin. Unfortunately they seems not to work with polar plots, the following code yields nothing:

import numpy as np
import matplotlib.pyplot as plt

# Generate random data:
N = 1024
r = .5 + np.random.normal(size=N, scale=.1)
theta = np.pi / 2 + np.random.normal(size=N, scale=.1)

# Plot:
ax = plt.subplot(111, polar=True)
ax.hist2d(theta, r)
plt.savefig('foo.png')
plt.close()

I would like it to look like this: pylab_examples example code: hist2d_demo.py only in polar coordinates. The closest result so far is with colored scatter plot as adviced here:

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde

# Generate random data:
N = 1024
r = .5 + np.random.normal(size=N, scale=.1)
theta = np.pi / 2 + np.random.normal(size=N, scale=.1)

# Plot:
ax = plt.subplot(111, polar=True)

# Using approach from:
# https://stackoverflow.com/questions/20105364/how-can-i-make-a-scatter-plot-colored-by-density-in-matplotlib
theta_r = np.vstack([theta,r])
z = gaussian_kde(theta_r)(theta_r)

ax.scatter(theta, r, c=z, s=10, edgecolor='')

plt.savefig('foo.png')
plt.close()

Image from the second version of the code

Is there a better way to make it more like real PDF generated with hist2d? This question seems to be relevant (the resulting image is as expected), but it looks messy.

Answer 1

One way to this using pcolormesh:

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde

# Generate random data:
N = 10000
r = .5 + np.random.normal(size=N, scale=.1)
theta = np.pi / 2 + np.random.normal(size=N, scale=.1)


# Histogramming
nr = 50
ntheta = 200
r_edges = np.linspace(0, 1, nr + 1)
theta_edges = np.linspace(0, 2*np.pi, ntheta + 1)
H, _, _ = np.histogram2d(r, theta, [r_edges, theta_edges])

# Plot
ax = plt.subplot(111, polar=True)
Theta, R = np.meshgrid(theta_edges, r_edges)
ax.pcolormesh(Theta, R, H)
plt.show()

Result:

Note that the histogram is not yet normalized by the area of the bin, which is not constant in polar coordinates. Close to the origin, the bins are pretty small, so some other kind of meshing might be better.