pythonmatplotlib

Heatmap of dataset with 3 1-dimensional arrays?

I have a dataset with values for X, Y and Z, where each value in an array corresponds to a single value in the others (X[i] corresponds to Y[i] and Z[i]). I want to plot that as a heatmap similar to this:

enter image description here

Where I convert X and Y to polar coordinates and the color represents the value of Z. How can I do this? All the solutions I found are for cases where each X[i] corresponds to multiple values of y, or Z is already a mesh, which is not my case.

The following code is an example of what I've been trying.

import numpy as np
import matplotlib.pyplot as plt
n = 5
z = np.linspace(0, 5, n)
x = np.linspace(-1, 1, n)
y = np.linspace(-1, 1, n)

r = np.sqrt(x**2 + y**2)
th = np.arctan2(y, x)

fig, ax = plt.subplots(subplot_kw={'projection':'polar'})
ax.pcolormesh(x, y, z, shading='nearest', vmin=z.min(), vmax=z.max())
X, Y = np.meshgrid(x, y)
ax.plot(X.flat, Y.flat, 'o', color='m')
plt.show()
Solution

  • There are a couple of options: 1) you do binning or interpolation of your parameters onto a 2D grid that you can then pass to pcolormesh, 2) you just use a scatter plot of your 1D data (there will likely be gaps in your plot).

    An example of interpolation is:

    import numpy as np
from scipy.interpolate import interp2d
import matplotlib.pyplot as plt

# create some artificial data
n = 250

rng = np.random.default_rng()
z = rng.uniform(0, 5, n)
x = rng.uniform(-np.pi, np.pi, n)
y = rng.uniform(-np.pi, np.pi, n)

# create a 2d interpolation function
f = interp2d(x, y, z, kind="linear")

# interpolation grid
nbins = 100
xi = np.linspace(-np.pi, np.pi, nbins)
yi = np.linspace(-np.pi, np.pi, nbins)

X, Y = np.meshgrid(xi, yi)

# get interpolated Z values on the grid
Z = f(xi, yi)

# make the polar plot
fig, ax = plt.subplots(subplot_kw={'projection':'polar'})
ax.pcolormesh(X, Y, Z, shading='nearest', vmin=z.min(), vmax=z.max())

    2d mesh polar plot

    Or, just using a scatter plot:

    import numpy as np
import matplotlib.pyplot as plt

n = 5000

rng = np.random.default_rng()

# data points
z = rng.uniform(0, 5, n)
x = rng.uniform(-np.pi, np.pi, n)
y = rng.uniform(-np.pi, np.pi, n)

fig, ax = plt.subplots(subplot_kw={'projection':'polar'})
ax.scatter(x, y, c=z, s=25)

    2d scatter plot