circular_colorbar_demo.py

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   
np.random.seed(1977)                                 

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,
                  borderpad=1)
    ax.set(title='Azimuth')

    fig, ax = plt.subplots()
    im = ax.pcolormesh(x, y, z, cmap='gist_earth')
    fig.colorbar(im)
    ax.set(title='Elevation')

    plt.show()

def circ_colorbar(parent_ax, cmap, inner_radius=0, **kwargs):
    kwargs['axes_class'] = PolarAxes
    ax = inset_locator.inset_axes(parent_ax, **kwargs)
    ax.patch.set_visible(False)

    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

main()