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()