Source code for inspec.mpl

# -*- coding: utf-8 -*-

"""
.. _mpl:

mpl - Matplotlib utilities
==========================

.. Warning:: `mpl` imports `matplotlib`, but **should not** import
   `pylab` or `pyplot`.
"""

__author__ = "Yannick Copin <ycopin@ipnl.in2p3.fr>"


# Non-default colors (from Colorbrewer2.org, qualitative)
blue = '#377EB8'
red = '#E41A1C'
green = '#4DAF4A'
orange = '#FF7F00'
purple = '#984EA3'
yellow = '#FFFF33'
brown = '#A65628'
pink = '#F781BF'
gray = '#999999'
colors = [blue, red, green, orange, purple, yellow, brown, pink, gray]


# Helper functions ==============================


def get_backend(inname, backend=None, name='figure'):
    """
    Decipher *inname* as a matplotlib backend (pylab|png|eps|pdf|svg) or a
    filename with backend-specific extension. Return *backend, figname* -- to
    be used in `matplotlib.use(backend)` and `fig.savefig(figname)` -- or
    `None, None` for interactive display.

    :param inname: Input figure name ('name.ext') or backend name ('backend')
    :param backend: default backend. If no valid backend can be deciphered from
                    *inname*, use *backend* as default name
    :param name: default figure basename. If *inname* is actually a backend
                 `ext`, `name.ext` will be the default figname.

    :return: backend, figname

    **Examples:**

    For a single output figure, `opts.output` could be a backend name or a
    complete filename to override default outname::

      import matplotlib as M
      backend, figname = get_backend(opts.output, name='my_default_figname')
      if backend:
          M.use(backend)
      import matplotlib.pyplot as P     # Has to be *after* M.use
      [...]
      if backend:
          print "Saving plot in", figname
          fig.savefig(figname)
      else:
          P.show()

    To produce multiple figures with the same backend (given with option
    `opts.backend`), one can use `name=''`::

      import matplotlib as M
      backend, figext = get_backend(opts.backend, name='')
      if backend:
          M.use(backend)
      import matplotlib.pyplot as P     # Has to be *after* M.use
      [...]
      if backend:
          fig1.savefig('my_figname_1' + figext)
          fig2.savefig('my_figname_2' + figext)
      else:
          P.show()

    .. Warning:: `pylab`/`pyplot` *should not* be imported before the explicit
       choice of backend with `matplotlib.use(backend)`
    """

    import os

    backends = dict(png='Agg', eps='PS', pdf='PDF', svg='SVG')

    bname, ext = os.path.splitext(inname)
    if ext:                     # inname is 'bname.ext'
        bkd = ext[1:].lower()   # dotless extension
    else:                       # inname is 'backend' only
        bkd = bname.lower()
        bname = name            # Default basename
    if bkd == 'pylab':          # Use interactive display
        return None, None
    elif bkd in backends:       # Return backend, output figure name
        return backends[bkd], os.path.extsep.join((bname, bkd))
    elif backend:               # Use default backend
        print(f"Unknown backend '{bkd}', using default backend '{backend}'")
        return get_backend(backend, name=name, backend=None)
    else:                       # Invalid backend
        raise ValueError(f"Unknown backend '{bkd}'")


def save_figures(prefix='figure', ext='pdf', joint=True):
    """
    Save active figures in files `prefix_XX.ext`; PDF can be joined in a single file.

    :param prefix: figure prefix
    :param ext: figure extension
    :param bool joint: joint PDF figures into a single file
    """

    from math import log10
    import matplotlib.pyplot as P

    figs = P._pylab_helpers.Gcf.figs  # Figures (actually FigureManagers)
    ifigs = sorted(figs)              # Sorted figure indices

    if ext == 'pdf' and joint:
        from matplotlib.backends import backend_pdf

        figname = f"{prefix}.{ext}"
        print(f"Saving all {len(ifigs)} figures in {figname!r}")
        pdf = backend_pdf.PdfPages(figname)
        for ifig in ifigs:
            pdf.savefig(ifig)
        pdf.close()
    else:
        width = int(log10(len(ifigs)) + 0.5)
        for ifig in ifigs:
            figname = f"{prefix}_{ifig:0={width}d}.{ext}"
            print(f"Saving figure #{ifig} in {figname!r}")
            figs[ifig].canvas.figure.savefig()