plotting.py 2.64 KB
Newer Older
Daniel Brown's avatar
Daniel Brown committed
1 2 3 4 5 6
# -*- coding: utf-8 -*-
"""
Created on Sat Feb 02 10:35:04 2013

@author: Daniel
"""
7 8 9 10
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
Daniel Brown's avatar
Daniel Brown committed
11

12
__DPI__ = None
Daniel Brown's avatar
Daniel Brown committed
13

14

15 16
def in_ipython():
    try:
17
        __IPYTHON__
18 19
    except NameError:
        return False
20 21 22 23 24 25 26
    else:
        return True


# `mode` should be either display (showing in windows)
# or paper (saving for paper/report/etc.)
def init_pykat_plotting(mode="display", dpi=None, fmts=None):
27
    import matplotlib as mpl
28 29
    import pykat.style

30
    if fmts is None:
31 32
        fmts = ['svg']

33
    if in_ipython():
34 35 36 37
        try:
            from IPython.display import set_matplotlib_formats
            ipy = get_ipython()
            ipy.magic("matplotlib inline")
38 39
            set_matplotlib_formats(*fmts)
        except NameError:
40
            pass
41

42
    if mode == "display":
43
        pykat.style.use(["default"])
44 45
    elif mode == "paper":
        mpl.use("pgf")
46 47 48 49
        pykat.style.use(["default", "paper"])
    else:
        raise (BaseException(
            "Plotting mode must be either 'display' or 'paper'."))
50

51 52
    if dpi is None:
        __DPI__ = mpl.rcParams['figure.dpi']
53
    else:
54 55 56
        __DPI__ = int(dpi)
        mpl.rcParams.update({'figure.dpi': __DPI__})
        mpl.rcParams.update({'savefig.dpi': __DPI__})
57 58

    if (mpl.__version__ < '1.5'):
59 60 61 62 63 64
        # prop_cycle doesn't exist pre-1.5, and unknown rcParams are ignored,
        # so we have to hardcode the color cycle
        mpl.rcParams.update({'axes.color_cycle': ['0000FF', 'FF0000', '000000', '00FF00', 'FF00FF']})


def figure(width=None, height=None, textwidth=None, **kwargs):
65 66
    """
    Options:
67 68 69 70
        width: 'full', 'half' (0.49*textwidth) (default: None (use current
        value))
        height: relative height to width (default: None (use current value))
        textwidth: Width of text in inches (default: None (use current value))
71
    """
72
    import matplotlib as mpl
73
    import matplotlib.pyplot as plt
74 75 76 77 78 79 80 81 82
    fig_size = mpl.rcParams["figure.figsize"]

    if width is None:
        width = "full"
    if height is None:
        height = fig_size[1] / fig_size[0]
    if textwidth is None:
        textwidth = fig_size[0]

83
    if width == "full":
84
        fig_size[0] = textwidth
85
    elif width == "half":
86 87
        fig_size[0] = 0.49 * textwidth
        fig_size[1] = 0.49 * textwidth * height
Daniel Brown's avatar
Daniel Brown committed
88
    else:
89 90
        raise (BaseException("width must be either 'full', 'half' or None."))

91
    fig = plt.figure(figsize=fig_size, dpi=__DPI__)
92

93
    return fig
94 95 96 97 98 99 100 101


def subplot(*args, **kwargs):
    import matplotlib.pyplot as plt

    ax = plt.subplot(*args, **kwargs)

    return ax