Correctly fix azimuths >360 or <0

rose.py

import itertools
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.collections import PatchCollection

import numpy as np
np.random.seed(1977)

def main():
    azi = np.random.normal(20, 40, 1000)
    lengths = 10 * (np.cos(np.radians(azi + 40))**2 + 1)

    fig, axes = plt.subplots(ncols=2, figsize=(12, 6),
                             subplot_kw=dict(projection='polar'))

    axes[0].set_title('Unweighted', y=1.1)
    rose(azi, ax=axes[0], bidirectional=True, color='gray')

    axes[1].set_title('Weighted by Length', y=1.1)
    rose(azi, ax=axes[1], weight_by=lengths, bidirectional=True, color='gray')

    for ax in axes.flat:
        ax.set(xticks=np.radians(range(0, 360, 45)), 
               xticklabels=['N', 'NE', 'E', 'SE', 'S', 'SW', 'W', 'NW'],
               rlabel_position=165)

    plt.show()

def rose(azimuths, z=None, ax=None, bins=30, bidirectional=False, 
         color_by=np.mean, weight_by=None, **kwargs):
    """Create a "rose" diagram (a.k.a. circular histogram).  

    Parameters:
    -----------
        azimuths : sequence of numbers
            The observed azimuths in degrees.
        z : sequence of numbers (optional)
            A second, co-located variable to color the plotted rectangles by.
        ax : a matplotlib Axes (optional)
            The axes to plot on. Defaults to the current axes.
        bins : int or sequence of numbers (optional)
            The number of bins or a sequence of bin edges to use.
        bidirectional : boolean (optional)
            Whether or not to treat the observed azimuths as bi-directional
            measurements (i.e. if True, 0 and 180 are identical).
        color_by : function or string (optional)
            A function to reduce the binned z values with. Alternately, if the
            string "count" is passed in, the displayed bars will be colored by
            their y-value (the number of azimuths measurements in that bin).
        weight_by : sequence of numbers (optional)
            Histogram weights. If specified, the length of each "petal" in the
            rose will be determined by summing these values in each azimuth
            bin instead of by counts.
        Additional keyword arguments are passed on to PatchCollection.

    Returns:
    --------
        A matplotlib PatchCollection
    """
    azimuths = np.asanyarray(azimuths)
    if color_by == 'count':
        z = np.ones_like(azimuths)
        color_by = np.sum
    if ax is None:
        ax = plt.gca()
    ax.set_theta_direction(-1)
    ax.set_theta_offset(np.radians(90))

    if bidirectional:
        other = azimuths + 180
        azimuths = np.concatenate([azimuths, other])
        if z is not None:
            z = np.concatenate([z, z])
        if weight_by is not None:
            weight_by = np.concatenate([weight_by, weight_by])

    # Convert to 0-360, in case negative or >360 azimuths are passed in.
    azimuths = azimuths % 360

    counts, edges = np.histogram(azimuths, range=[0, 360], bins=bins,
                                 weights=weight_by)
    if z is not None:
        idx = np.digitize(azimuths, edges)
        z = np.array([color_by(z[idx == i]) for i in range(1, idx.max() + 1)])
        z = np.ma.masked_invalid(z)
    edges = np.radians(edges)
    coll = colored_bar(edges[:-1], counts, z=z, width=np.diff(edges), 
                       ax=ax, **kwargs)
    return coll

def colored_bar(left, height, z=None, width=0.8, bottom=0, ax=None, **kwargs):
    """A bar plot colored by a scalar sequence."""
    if ax is None:
        ax = plt.gca()
    width = itertools.cycle(np.atleast_1d(width))
    bottom = itertools.cycle(np.atleast_1d(bottom))
    rects = []
    for x, y, h, w in zip(left, bottom, height, width):
        rects.append(Rectangle((x,y), w, h))
    coll = PatchCollection(rects, array=z, **kwargs)
    ax.add_collection(coll)
    ax.autoscale()
    return coll

if __name__ == '__main__':
    axes = main()