autofmt_datetime_axis

datetime_axis_magic.py

"""
    This function accepts a matplotlib axis object with a datetime x-axis
    and performs magical things to it to make it look not like total shit.
"""

import matplotlib.dates as dates
import matplotlib.ticker as mtick
import datetime

def label_july_only(x, pos=None):
    """Label only July.
    """
    x = dates.num2date(x)
    label = "{dt:%b}".format(dt=x) if x.month == 7 else ""

    return label

def label_major_hours(x, pos=None):
    x = dates.num2date(x)

    if pos == 0 or x.hour == 0:
        label = "{dt:%H}:{dt:%M}\n{dt:%b} {dt.day}".format(dt=x)
    else:
        label = "{dt:%H}:{dt:%M}".format(dt=x)

    return label

def label_every_3_hours(x, pos=None):
    x = dates.num2date(x)
    if x.hour in (3, 6, 9, 15, 18, 21):
        label = "{dt:%H}:{dt:%M}".format(dt=x)
    else:
        label = ""

    return label

def label_none(x, pos=None):
    return ""

def label_all_days(x, pos=None):
    """
    """
    x = dates.num2date(x)

    if pos == 0 or x.day == 1:
        label = "{dt.day}\n{dt:%b}\n{dt.year}".format(dt=x)
    else:
        label = "{dt.day}".format(dt=x)

    return label

def add_year_to_jan(x, pos=None):
    """
    """
    x = dates.num2date(x)
    if x.month == 1 or pos == 0:
        label = "{dt:%b}\n{dt.year}".format(dt=x)
    else:
        label = "{dt:%b}".format(dt=x)

    return label

def label_weeks(x, pos=None):
    """
    """
    x = dates.num2date(x)
    if pos == 0:
        label = "{dt.day}\n{dt:%b}\n{dt.year}".format(dt=x)
    elif x.day == 1:
        label = "{dt.day}\n{dt:%b}".format(dt=x)
    else:
        label = "{dt.day}".format(dt=x)

    return label

def autofmt_datetime_axis(ax, minor_ticks=False):
    """Should I seperate minor_ticks from minor_tick_labels?
    """
    # Get the xmin and xmax of the axis object
    xmin = dates.num2date(ax.get_xlim()[0])
    xmax = dates.num2date(ax.get_xlim()[1])

    # Convert to a datetime object
    dt = xmax - xmin

    # Cycle through various timescales and set params
    if dt > datetime.timedelta(days=365):
        # Set the minor ticks to always be one month
        minor_loc = dates.MonthLocator(interval=1)
        minor_fmt = dates.AutoDateFormatter(minor_loc)

        # Set a default minor scale formatter
        minor_scale_func = label_none

        if dt.days > 3*365:
            bymonth = (1)
            minor_scale_func = label_july_only
        elif dt.days > 2*365:
            bymonth = (1, 5, 9)
        else:
            bymonth = (1, 4, 7, 10)

        # Set the major scale
        major_loc = dates.MonthLocator(bymonth=bymonth)
        major_fmt = dates.AutoDateFormatter(major_loc)

        # Adjust the scale appropriately
        major_fmt.scaled[30.0] = mtick.FuncFormatter(add_year_to_jan)

        # Add format to minor ticks if set to True
        if minor_ticks:
            minor_fmt.scaled[30.0] = mtick.FuncFormatter(minor_scale_func)
    elif dt > datetime.timedelta(days=90):
        minor_loc = dates.DayLocator(bymonthday=(1, 7, 14, 21, 28))
        minor_fmt = dates.AutoDateFormatter(minor_loc)

        major_loc = dates.MonthLocator(interval=1)
        major_fmt = dates.AutoDateFormatter(major_loc)

        # Adjust the scale
        major_fmt.scaled[30.0] = mtick.FuncFormatter(add_year_to_jan)

        if minor_ticks:
            minor_fmt.scaled[1.0] = mtick.FuncFormatter(label_none)

    elif dt > datetime.timedelta(days=21):
        # Major Ticks: Weeks
        # Minor Ticks: Days
        minor_loc = dates.DayLocator(interval=1)
        minor_fmt = dates.AutoDateFormatter(minor_loc)

        major_loc = dates.DayLocator(bymonthday=(1, 7, 14, 21, 28))
        major_fmt = dates.AutoDateFormatter(major_loc)

        major_fmt.scaled[1.0] = mtick.FuncFormatter(label_weeks)

        # All minor ticks are going to be blank
        if minor_ticks:
            minor_fmt.scaled[1.0] = mtick.FuncFormatter(label_none)

    elif dt > datetime.timedelta(days=2):
        major_loc = dates.DayLocator(interval=1)
        major_fmt = dates.AutoDateFormatter(major_loc)

        major_fmt.scaled[1.0] = mtick.FuncFormatter(label_all_days)

        minor_ticks = False
    elif dt > datetime.timedelta(hours=4):
        if (dt.seconds / 3600.) > 24:
            major_loc = dates.HourLocator(byhour=(0, 12))
        elif (dt.seconds / 3600.) > 9:
            major_loc = dates.HourLocator(byhour=(0, 3, 6, 9, 12, 15, 18, 21))
        else:
            major_loc = dates.HourLocator(interval=1)

        minor_loc = dates.HourLocator(interval=1)
        minor_fmt = dates.AutoDateFormatter(minor_loc)

        major_fmt = dates.AutoDateFormatter(major_loc)

        major_fmt.scaled[1.0/24.] = mtick.FuncFormatter(label_major_hours)

        if minor_ticks:
            minor_fmt.scaled[1.0/24.] = mtick.FuncFormatter(label_none)
    else:
        major_loc = dates.AutoDateLocator()
        major_fmt = dates.AutoDateFormatter(major_loc)

        minor_ticks = False

    # Set the Major Axis Formatting
    ax.xaxis.set_major_locator(major_loc)
    ax.xaxis.set_major_formatter(major_fmt)

    # Set minor ticks if set
    if minor_ticks:
        ax.xaxis.set_minor_locator(minor_loc)
        ax.xaxis.set_minor_formatter(minor_fmt)

    return None