Proof of concept of support for astropy.Quantity in matplotlib

mpl_astropy_units.py

import numpy as np
from matplotlib import units
from matplotlib import ticker
from astropy import units as u


def rad_fn(x, pos=None):
    n = int((x / np.pi) * 2.0 + 0.25)
    if n == 0:
        return '0'
    elif n == 1:
        return r'$\pi/2$'
    elif n == 2:
        return r'$\pi$'
    elif n % 2 == 0:
        return r'$%s\pi$' % (n/2,)
    else:
        return r'$%s\pi/2$' % (n,)


class MplQuantityConverter(units.ConversionInterface):
    @staticmethod
    def axisinfo(unit, axis):
        if unit == u.radian:
            return units.AxisInfo(
                majloc=ticker.MultipleLocator(base=np.pi/2),
                majfmt=ticker.FuncFormatter(rad_fn),
                label=unit.to_string(),
            )
        elif unit == u.degree:
            return units.AxisInfo(
                majloc=ticker.AutoLocator(),
                majfmt=ticker.FormatStrFormatter(r'$%i^\circ$'),
                label=unit.to_string(),
            )
        elif unit is not None:
            return units.AxisInfo(label=unit.to_string())
        return None

    @staticmethod
    def convert(val, unit, axis):
        if isinstance(val, u.Quantity):
            return val.to(unit)
        else:
            return val

    @staticmethod
    def default_units(x, axis):
        if hasattr(x, 'unit'):
            if x.unit.is_equivalent(u.radian):
                return u.radian
            return x.unit
        return None


mpl_quantity_converter = MplQuantityConverter()
units.registry[u.Quantity] = mpl_quantity_converter


from matplotlib import pyplot as plt

plt.plot([1, 2, 3] * u.m, [1, 2, 3] * u.degree)
plt.plot([1, 2, 3] * u.cm, np.deg2rad([1, 2, 3]) * u.radian)
plt.show()