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