def plot_fit(fn, cor, tmin, tmax, options, fitted_params, errors=None, postfix=None):
import plot_helpers
emass_dt = 1
X = np.linspace(tmin, tmax, 200 * 5)
if errors is None:
fitted_errors = [0.0] * len(fitted_params)
else:
fitted_errors = errors
fig = plt.figure()
corplot = plt.subplot(211)
cordata = corplot.errorbar(
cor.times, cor.average_sub_vev().values(), yerr=cor.jackknifed_errors().values(), fmt="o"
)
corfit, = corplot.plot(X, fn.formula(fitted_params, X), lw=2.0)
single_fit = None
if fn.description == "exp" or "subtract" in fn.description:
corplot.legend([cordata, corfit], ["Correlator data", fn.template.format(*fitted_params)], loc="best")
else:
single = functions["single_exp"]()
single_fit, = corplot.plot(X, single.formula(fitted_params[:2], X), ls="-.", lw=2.0)
corplot.legend(
[cordata, corfit, single_fit],
["Correlator data", fn.template.format(*fitted_params), "single_exp with these values"],
loc="best",
)
corplot.set_ylabel("Fit Correlator")
corvals = cor.average_sub_vev().values()
plt.ylim(plot_helpers.auto_fit_range(min(corvals), max(corvals)))
plt.xlim([0, tmax + 2])
emass = cor.periodic_effective_mass(emass_dt, fast=False, period=options.period)
emass_errors = cor.periodic_effective_mass_errors(emass_dt, fast=False, period=options.period).values()
emassplot = plt.subplot(212)
emassplot.set_ylabel("${\mathrm{\mathbf{m}_{eff}}}$")
dataplt = emassplot.errorbar(emass.keys(), emass.values(), yerr=emass_errors, fmt="o")
named_params = {n: (m, e) for n, m, e in zip(fn.parameter_names, fitted_params, fitted_errors)}
mass, mass_err = named_params["mass"]
# abovefitline = emassplot.plot(range(tmin, tmax+1), [mass+mass_err]*len(range(tmin, tmax+1)), ls="dashed", color="b")
fitplt, = emassplot.plot(range(tmin, tmax + 1), [mass] * len(range(tmin, tmax + 1)), ls="dotted", color="r")
# belowfitline = emassplot.plot(range(tmin, tmax+1), [mass-mass_err]*len(range(tmin, tmax+1)), ls="dashed", color="b")
fitpoints = fn.formula(fitted_params, np.arange(tmin, tmax + 1))
emassfit = []
emassfit_range = []
dt = emass_dt
for i in range(len(fitpoints))[:-dt]:
try:
fitemass = (1.0 / float(dt)) * math.acosh((fitpoints[i + dt] + fitpoints[i - dt]) / (2.0 * fitpoints[i]))
emassfit.append(fitemass)
emassfit_range.append(tmin + i)
except ValueError:
fitemass = 0.0
emass_fit = emassplot.plot(emassfit_range, emassfit, color="k")
emassplot.legend(
[dataplt, fitplt], ["Emass of data", u"fit mass={:.5f}\xb1{:.5f}".format(mass, mass_err)], loc="best"
)
# plt.ylim([min(min(emass.values()),-0.01), max(emass.values())*1.2])
plt.xlim([0, tmax + 2])
if options.output_stub:
if postfix:
filestub = options.output_stub + postfix
else:
filestub = options.output_stub
logging.info("Saving plot to {}".format(filestub + ".png"))
fig.set_size_inches(18.5, 10.5)
plt.savefig(filestub + ".png")
else:
plt.show()
plt.close()
def plot_files(files, output_stub=None, yrange=None, xrang=None, cols=-1, fit=False, real=False, title=None):
markers = ['o', "D", "^", "<", ">", "v", "x", "p", "8"]
# colors, white sucks
# colors = sorted([c for c in mpl.colors.colorConverter.colors.keys() if c != 'w' and c != "g"])
colors = ['b', 'c', 'm', 'r', 'k', 'y']
plots = {}
tmin_plot = {}
has_colorbar = False
labels = label_names_from_filelist(files)
fontsettings = dict(fontweight='bold', fontsize=18)
if args.translate:
labels = [translate(l) for l in labels]
seperate = cols > 0
ymin, ymax = 1000, None
xmin, xmax = 1000, None
rows = int(math.ceil(float(len(labels))/cols))
if seperate:
f, layout = plt.subplots(nrows=rows, ncols=cols, sharey=True, sharex=True, squeeze=False)
else:
f, axe = plt.subplots(1)
axe.tick_params(axis='both', which='major', labelsize=20)
axe.set_xlabel("time", **fontsettings)
for i in range(cols): # Set bottom row to have xlabels
layout[rows-1][i].set_xlabel("time", **fontsettings)
for index, label, filename in zip(range(len(files)), labels, files):
i = (index)/cols
j = (index) % cols
if seperate:
axe = layout[i][j]
if j == 0:
if "cor" in filename:
axe.set_ylabel("Correlator", **fontsettings)
if args.emass:
logging.warn("EMASS flag set but filename indicates a correlator file!")
if "emass" in filename or args.emass:
if args.scalefactor:
axe.set_ylabel("${\mathrm{\mathbf{m}_{eff}}}$ [MeV]", **fontsettings)
else:
axe.set_ylabel("${\mathrm{\mathbf{m}_{eff}}}$", **fontsettings)
if args.rel_error:
axe.set_ylabel("Relative Error", **fontsettings)
if fit:
if seperate:
fitstring = add_fit_info(filename, ax=axe)
else:
fitstring = add_fit_info(filename)
if fitstring:
if args.fit_only:
logging.info("setting label to {}".format(fitstring))
label = fitstring
else:
label += " $m_{fit}=$" + fitstring
mark = markers[index % len(markers)]
color = colors[index % len(colors)]
df = read_file(filename)
if len(df.time) > len(set(df.time)):
df = read_full_correlator(filename, args.emass, args.eamp, args.symmetric)
if args.rel_error:
df["correlator"] = df["error"]/df["correlator"]
df["error"] = 0.0
time_offset = df.time.values+(index*0.1)
time_offset = df.time.values
if seperate:
time_offset = df.time.values
logging.debug("%s %s %s", df.time.values, df.correlator.values, df.error.values)
plotsettings = dict(linestyle="none", c=color, marker=mark, label=label, ms=5, elinewidth=2, capsize=5,
capthick=2, mec=color, aa=True)
if args.rel_error:
plotsettings["elinewidth"] = 0
plotsettings["capthick"] = 0
if seperate:
logging.info("plotting {} {}, {}".format(label, i, j))
#axe.set_title(label)
axe.legend(fancybox=True, shadow=True, loc=0)
# Do a Tmin plot
if args.scalefactor:
scale = args.scalefactor
else:
scale = 1.0
if any(df["quality"].notnull()):
logging.info("found 4th column, plotting as quality")
cmap = mpl.cm.cool
plots[label] = axe.errorbar(time_offset, scale*df.correlator.values, yerr=scale*df.error.values, fmt=None,
zorder=0, **plotsettings)
tmin_plot[label] = axe.scatter(time_offset, scale*df.correlator.values, c=df.quality.values,
s=50, cmap=cmap, marker=mark)
tmin_plot[label].set_clim(0, 1)
if seperate:
has_colorbar = True
if not has_colorbar and not seperate:
cb = plt.colorbar(tmin_plot[label]) # noqa
cb.set_label("Quality of fit", **fontsettings)
axe.set_xlabel("tmin", **fontsettings)
axe.set_ylabel("Fit Value", **fontsettings)
has_colorbar = True
else: # Not a tmin plot!
if np.iscomplexobj(df.correlator.values):
plots[label] = axe.errorbar(time_offset, scale*np.real(df.correlator.values), yerr=scale*np.real(df.error.values),
**plotsettings)
if not real:
plots["imag"+label] = axe.errorbar(time_offset, scale*np.imag(df.correlator.values),
yerr=scale*np.imag(df.error.values), markerfacecolor='none',
**plotsettings)
else:
plots[label] = axe.errorbar(time_offset, scale*df.correlator.values, yerr=scale*df.error.values, **plotsettings)
if not yrange:
ymin = min(ymin, min(df.correlator.fillna(1000)))
ymax = max(ymax, max(df.correlator.fillna(0)))
logging.debug("ymin {} ymax {}".format(ymin, ymax))
if not xrang:
xmin = min(xmin, min(df.time)-1)
xmax = max(xmax, max(df.time)+1)
logging.debug("xmin {} xmax {}".format(xmin, xmax))
axe.legend(fancybox=True, shadow=True, loc=0)
if args.plotfunction:
add_function_plot(args.plotfunction, xmin,xmax)
if not args.logarithm:
if yrange:
plt.ylim(yrange)
else:
plt.ylim(plot_helpers.auto_fit_range(scale*ymin,scale*ymax))
if xrang:
plt.xlim(xrang)
else:
plt.xlim(xmin, xmax)
if args.logarithm:
plt.yscale('log')
if args.constant:
bignum = 1000000
plt.plot([-1*bignum,bignum],[args.constant,args.constant])
if title:
f.suptitle(title.replace("_", " "), **fontsettings)
f.canvas.set_window_title(files[0])
if seperate:
plt.tight_layout(pad=0.0, h_pad=0.0, w_pad=0.0)
if has_colorbar:
f.subplots_adjust(right=0.95)
cbar_ax = f.add_axes([0.96, 0.05, 0.01, 0.9])
f.colorbar(tmin_plot[label], cax=cbar_ax)
else:
if not args.nolegend:
leg = plt.legend(fancybox=True, shadow=True, loc=0)
if(output_stub):
width = 10.0
f.set_size_inches(width, width*args.aspect)
# plt.rcParams.update({'font.size': 20})
# plt.tight_layout(pad=2.0, h_pad=1.0, w_pad=2.0)
plt.tight_layout()
plt.subplots_adjust(top=0.90)
if args.eps:
logging.info("Saving plot to {}".format(output_stub+".eps"))
plt.savefig(output_stub+".eps")
else:
logging.info("Saving plot to {}".format(output_stub+".png"))
plt.savefig(output_stub+".png", dpi=400)
return
def toggle_errorbar_vis(ebarplot):
for i in flatten(ebarplot):
if i:
i.set_visible(not i.get_visible())
def func(label):
toggle_errorbar_vis(plots[label])
if label in tmin_plot.keys():
tmin_plot[label].set_visible(not tmin_plot[label].get_visible())
plt.draw()
if not seperate and not args.nolegend:
rax = plt.axes([0.9, 0.8, 0.1, 0.15])
check = CheckButtons(rax, plots.keys(), [True]*len(plots))
check.on_clicked(func)
# if not args.nolegend and len(plots) > 1:
# leg.draggable()
plt.show()
