astro py:docs

- Index
- Modules

### Navigation

- Astropy v1.1.1 »
- Module code »

# Source code for astropy.visualization.units

    # -*- coding: utf-8 -*-
    # Licensed under a 3-clause BSD style license - see LICENSE.rst

    from __future__ import (absolute_import, division, print_function,
                            unicode_literals)

    __doctest_skip__ = ['quantity_support']

    import numpy as np

    def quantity_support(format='latex_inline'):
        """
        Enable support for plotting `astropy.units.Quantity` instances in
        matplotlib.

        May be (optionally) used with a ``with`` statement.

          >>> import matplotlib.pyplot as plt
          >>> from astropy import units as u
          >>> from astropy import visualization
          >>> with visualization.quantity_support():
          ...     plt.figure()
          ...     plt.plot([1, 2, 3] * u.m)
          [...]
          ...     plt.plot([101, 125, 150] * u.cm)
          [...]
          ...     plt.draw()

        Parameters
        ----------
        format : `astropy.units.format.Base` instance or str
            The name of a format or a formatter object.  If not
            provided, defaults to ``latex_inline``.

        """
        from .. import units as u

        from matplotlib import units
        from matplotlib import ticker

        def rad_fn(x, pos=None):
            n = int((x / np.pi) * 2.0 + 0.25)
            if n == 0:
                return '0'
            elif n == 1:
                return 'π/2'
            elif n == 2:
                return 'π'
            elif n % 2 == 0:
                return '{0}π'.format(n / 2)
            else:
                return '{0}π/2'.format(n)

        class MplQuantityConverter(units.ConversionInterface):
            def __init__(self):
                if u.Quantity not in units.registry:
                    units.registry[u.Quantity] = self
                    self._remove = True
                else:
                    self._remove = False

            @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('%i°'),
                        label=unit.to_string(),
                    )
                elif unit is not None:
                    return units.AxisInfo(label=unit.to_string(format))
                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'):
                    return x.unit
                return None

            def __enter__(self):
                return self

            def __exit__(self, type, value, tb):
                if self._remove:
                    del units.registry[u.Quantity]

        return MplQuantityConverter()