Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
import numpy as np
from scipy.interpolate import Rbf
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from mpl_toolkits.axes_grid1 import inset_locator
from matplotlib.projections.polar import PolarAxes
def main():
x, y, z = generate_surface()
azi = azimuth(x, y, z)
cmap = make_cmap()
fig, ax = plt.subplots()
ax.pcolormesh(x, y, azi, cmap=cmap, vmin=0, vmax=360)
circ_colorbar(ax, cmap, inner_radius=0.5, width='15%', height='15%', loc=4,
fig, ax = plt.subplots()
im = ax.pcolormesh(x, y, z, cmap='gist_earth')
def circ_colorbar(parent_ax, cmap, inner_radius=0, **kwargs):
kwargs['axes_class'] = PolarAxes
ax = inset_locator.inset_axes(parent_ax, **kwargs)
n = 100
r = np.tile([[inner_radius, 1]], (n, 1))
theta = np.linspace(0, 2 * np.pi, n)
theta = np.column_stack([theta, theta])
ax.pcolormesh(theta, r, theta, cmap=cmap)
ax.set(yticks=[], xticks=[], ylim=[0, 1],
theta_zero_location='N', theta_direction='clockwise')
ax.set_thetagrids([0, 90, 180, 270], ['N', 'E', 'S', 'W'], frac=1.3)
if inner_radius > 0:
x = np.linspace(0, 2 * np.pi, n)
ax.plot(x, inner_radius * np.ones_like(x), color='black', lw=0.5)
return ax
def make_cmap():
colors = ['red', 'yellow', 'green', 'blue', 'red']
return LinearSegmentedColormap.from_list('some_name', colors)
def generate_surface():
nrows, ncols = 100, 100
npoints = 10
x, y, z = np.random.random((3, npoints))
yi, xi = np.mgrid[0:1:1j*nrows, 0:1:1j*ncols]
interp = Rbf(x, y, z)
zi = interp(xi.ravel(), yi.ravel())
return xi, yi, zi.reshape(xi.shape)
def azimuth(x, y, z):
xcellsize = np.diff(x[0,:])[0]
ycellsize = np.diff(y[:,0])[0]
dy, dx = np.gradient(-z, ycellsize, xcellsize)
theta = np.arctan2(dy, dx)
return (90 - np.degrees(theta)) % 360
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.