syrte
10/29/2016 - 1:14 PM

Make a histogram-like step plot.

Make a histogram-like step plot.

from __future__ import division, print_function, absolute_import
import numpy as np
from matplotlib import pyplot as plt

def steps(x, y, *args, **kwargs):
    '''steps(x, y, *args, style='line', bottom=0, **kwargs)
    Make a step plot.
    The interval from x[i] to x[i+1] has level y[i]
    This function is useful for show the results of np.histogram.

    Parameters
    ----------
    style : ['line' | 'step' | 'filled' | 'bar'], optional
        The type of steps to draw.
        - 'line': step line plot
        - 'step': step line with vertical line at borders.
        - 'filled': filled step line plot
        - 'bar': traditional bar-type histogram
    bottom: float
        The bottom baseline of the plot.
    args, kwargs:
        same as those for 
        `matplotlib.pyplot.plot` if `style` in ['line', 'step'], or
        `matplotlib.pyplot.plot.fill` if `style` in ['filled', 'bar'].
    
    Example
    -------
    np.random.seed(4)
    b = np.linspace(0.1, 0.9, 6)
    for i, style in enumerate(['line', 'step', 'filled', 'bar']):
        a = np.random.rand(50) + i
        steps(*np.histogram(a, b + i)[::-1], style=style, 
              lw=2, zorder=i)
        plt.text(i+0.5, 14, style)
    plt.xlim(0, 4)
    plt.ylim(-1, 16)
    '''
    m, n = len(x), len(y)
    if m == n:
        kwargs.setdefault("where", "mid")
        return plt.step(x, y, *args, **kwargs)
    elif m != n + 1:
        raise ValueError

    style = kwargs.pop('style', 'line')
    bottom = kwargs.pop('bottom', 0)
    kwargs.pop('drawstyle', None)
    kwargs.pop('where', None)

    if style == 'line':
        x, y = np.repeat(x, 2), np.repeat(y, 2)
        x = x[1:-1]
    elif style in ['step', 'filled']:
        x, y = np.repeat(x, 2), np.repeat(y, 2)
        y = np.hstack([bottom, y, bottom])
    elif style == 'bar':
        x, y = np.repeat(x, 3), np.repeat(y, 3)
        x = x[1:-1]
        y = np.hstack([y, bottom])
        y[::3] = bottom        
    else:
        raise ValueError("not surpport style: %s" % style)
            
    if style in ['line', 'step']:
        return plt.plot(x, y, *args, **kwargs)
    else:
        return plt.fill(x, y, *args, **kwargs)