1###########################################################################
2#                                                                         #
3#    physical_validation,                                                 #
4#    a python package to test the physical validity of MD results         #
5#                                                                         #
6#    Written by Michael R. Shirts <michael.shirts@colorado.edu>           #
7#               Pascal T. Merz <pascal.merz@colorado.edu>                 #
8#                                                                         #
9#    Copyright (C) 2012 University of Virginia                            #
10#              (C) 2017 University of Colorado Boulder                    #
11#                                                                         #
12#    This library is free software; you can redistribute it and/or        #
13#    modify it under the terms of the GNU Lesser General Public           #
14#    License as published by the Free Software Foundation; either         #
15#    version 2.1 of the License, or (at your option) any later version.   #
16#                                                                         #
17#    This library is distributed in the hope that it will be useful,      #
18#    but WITHOUT ANY WARRANTY; without even the implied warranty of       #
19#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU    #
20#    Lesser General Public License for more details.                      #
21#                                                                         #
22#    You should have received a copy of the GNU Lesser General Public     #
23#    License along with this library; if not, write to the                #
24#    Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,     #
25#    Boston, MA 02110-1301 USA                                            #
26#                                                                         #
27###########################################################################
28
29import numpy as np
30import warnings
31
32
33def plot(res, legend=None, title=None,
34         xlabel=None, ylabel=None, xlim=None, ylim=None,
35         inv_x=False, inv_y=False, sci_x=False, sci_y=False,
36         axtext=None, annotation_location=None,
37         filename=None, screen=True):
38
39    try:
40        import matplotlib as mpl
41        import matplotlib.pyplot as plt
42        from matplotlib.ticker import AutoMinorLocator
43    except ImportError:
44        warnings.warn('Install matplotlib to enable plotting.')
45        return
46
47    font = {'family': 'serif',
48            'weight': 'normal',
49            'size': 16}
50    mpl.rc('font', **font)
51
52    plt.ioff()
53    fig, ax = plt.subplots()
54    xmin = float('inf')
55    xmax = float('-inf')
56    for r in res:
57        x = r['x']
58        y = r['y']
59        if xlim is not None:
60            x = x[(r['x'] >= xlim[0]) & (r['x'] <= xlim[1])]
61            y = y[(r['x'] >= xlim[0]) & (r['x'] <= xlim[1])]
62        if 'y_err' in r:
63            dy = r['y_err']
64            if xlim is not None:
65                dy = dy[(r['x'] >= xlim[0]) & (r['x'] <= xlim[1])]
66            ax.errorbar(x, y, label=r['name'], yerr=dy)
67        else:
68            ax.plot(x, y, label=r['name'])
69        if xlim is not None:
70            xmin = min(np.min(r['x']), xmin)
71            xmax = max(np.max(r['x']), xmax)
72        else:
73            xmin = min(np.min(r['x']), xmin)
74            xmax = max(np.max(r['x']), xmax)
75
76    if legend is not None:
77        ax.legend(loc=legend)
78    box = ax.get_position()
79    if title is not None:
80        ax.set_title(title, y=1.05)
81        box = box.from_bounds(box.x0, box.y0, box.width, box.height * 0.95)
82    if xlabel is not None:
83        ax.set_xlabel(xlabel, labelpad=5)
84        box = box.from_bounds(box.x0, box.y0 + 0.05 * box.height, box.width, box.height * 0.95)
85    if ylabel is not None:
86        ax.set_ylabel(ylabel, labelpad=10)
87        box = box.from_bounds(box.x0 + 0.05 * box.width, box.y0, box.width * 0.95, box.height)
88    ax.set_position([box.x0, box.y0, box.width, box.height])
89    ax.axis('auto')
90    if xlim is not None:
91        ax.set_xlim(xlim)
92    else:
93        ax.set_xlim([xmin, xmax])
94    if ylim is not None:
95        ax.set_ylim(ylim)
96    ax.xaxis.set_minor_locator(AutoMinorLocator(2))
97
98    if inv_x:
99        ax.invert_xaxis()
100    if inv_y:
101        ax.invert_yaxis()
102
103    if axtext is not None:
104        if isinstance(axtext, str):
105            axtext = [axtext]
106        if annotation_location is None:
107            annotation_location = [None for _ in axtext]
108        if isinstance(annotation_location, tuple):
109            annotation_location = [annotation_location]
110        for t, loc in zip(axtext, annotation_location):
111            bbox = dict(boxstyle="round", fc="w", ec="0.5", alpha=0.9)
112            if loc is None:
113                ax.text(0.95, 0.05, t, transform=ax.transAxes,
114                        ha='right', va='bottom',
115                        bbox=bbox)
116            else:
117                ax.text(loc[0], loc[1], t,
118                        bbox=bbox)
119
120    if sci_x:
121        ax.ticklabel_format(style='sci', axis='x', scilimits=(-3, 4))
122    if sci_y:
123        ax.ticklabel_format(style='sci', axis='y', scilimits=(-3, 4))
124    ax.xaxis.major.formatter._useMathText = True
125
126    if filename is not None:
127        fig.savefig(filename + '.pdf', dpi=300)
128    if screen:
129        fig.show()
130