mirror of
https://github.com/triqs/dft_tools
synced 2024-12-25 13:53:40 +01:00
94 lines
3.8 KiB
Python
94 lines
3.8 KiB
Python
################################################################################
|
|
#
|
|
# TRIQS: a Toolbox for Research in Interacting Quantum Systems
|
|
#
|
|
# Copyright (C) 2011 by M. Ferrero, O. Parcollet
|
|
#
|
|
# TRIQS is free software: you can redistribute it and/or modify it under the
|
|
# terms of the GNU General Public License as published by the Free Software
|
|
# Foundation, either version 3 of the License, or (at your option) any later
|
|
# version.
|
|
#
|
|
# TRIQS is distributed in the hope that it will be useful, but WITHOUT ANY
|
|
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
|
|
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
|
|
# details.
|
|
#
|
|
# You should have received a copy of the GNU General Public License along with
|
|
# TRIQS. If not, see <http://www.gnu.org/licenses/>.
|
|
#
|
|
################################################################################
|
|
|
|
__all__ = ['plt','oplot','subplots','figsize_default','use_amsmath']
|
|
|
|
import numpy, matplotlib as mpl, matplotlib.pylab as plt
|
|
from protocol import plot_protocol_apply
|
|
from matplotlib import rc
|
|
|
|
try:
|
|
subplots = mpl.pyplot.subplots
|
|
except:
|
|
def subplots(nrows=1, ncols=1, sharex=False, sharey=False, squeeze=True, subplot_kw=None, **fig_kw):
|
|
print "subplots not supported"
|
|
return plt.figure(1), [plt.subplot(nrows,ncols,x+1) for x in range(nrows*ncols)]
|
|
|
|
figsize_default = (12,8)
|
|
|
|
|
|
def oplot (*ob_list, **opt_dict) :
|
|
"""
|
|
A thin layer above pyplot.plot function that allows plotting objects with
|
|
plot protocol as well as arrays.
|
|
Options are the same as for the pyplot.plot function.
|
|
"""
|
|
plt.figure(1, figsize = opt_dict.pop('figsize', figsize_default ) )
|
|
__oplot_impl(plt.plot, plt.xlabel,plt.ylabel,plt.legend, *ob_list,**opt_dict)
|
|
# remove this in the notebook...
|
|
#if hasattr(plt.figure(1), "show") : plt.figure(1).show()
|
|
|
|
mpl.axes.Axes.oplot = lambda self, *ob_list, **opt_dict : __oplot_impl(self.plot,self.set_xlabel, self.set_ylabel, self.legend, *ob_list,**opt_dict)
|
|
|
|
def __oplot_impl (plot_fct, xlabel_fct, ylabel_fct, legend_fct, *ob_list, **opt_dict) :
|
|
"""
|
|
A thin layer above pyplot.plot function that allows plotting objects with
|
|
plot protocol as well as arrays.
|
|
Options are the same as for the pyplot.plot function.
|
|
"""
|
|
|
|
def objs() : # filter the arguments for the format strings ...
|
|
i, l = 0, []
|
|
while i< len (ob_list) :
|
|
if i < len(ob_list) - 1 and type(ob_list[i+1]) == type("") :
|
|
res = ob_list[i], [ ob_list[i+1] ]
|
|
i+=2
|
|
else :
|
|
res = ob_list[i], [ ]
|
|
i+=1
|
|
yield res
|
|
|
|
for ob, OptionsList in objs() :
|
|
opt = opt_dict.copy() # the plot protocol will consume the dict....
|
|
#ob2 = eval_expr_or_pass (ob) # if it is a lazy_expr, it is time to evaluate it !
|
|
for curvedata in plot_protocol_apply(ob, opt, plt.xlim ) :
|
|
X,Y = curvedata['xdata'],curvedata['ydata']
|
|
d = { 'label' : curvedata['label'] }
|
|
d.update(opt)
|
|
try :
|
|
plot_fct(X,Y,*OptionsList,**d)
|
|
except TypeError, e:
|
|
import re
|
|
m = re.search('(?<=There is no line property )"(.*)"', str(e) )
|
|
if m :
|
|
raise RuntimeError, "Option %s is not understood in plot function : it is not an option of the object to be plotted, nor a matplotlib option"%m.group(0)
|
|
else :
|
|
raise
|
|
if 'xlabel' in curvedata : xlabel_fct(curvedata['xlabel'], fontsize=20)
|
|
if 'ylabel' in curvedata : ylabel_fct(curvedata['ylabel'], fontsize=20)
|
|
|
|
legend_fct(loc = 1) #legend is built from the label
|
|
|
|
def use_amsmath():
|
|
rc('text', usetex=True)
|
|
rc('text.latex', preamble="\usepackage{amsmath}")
|
|
|