예제 #1
0
def make_bar_plot(inputfile, cols, output_stub, mode, ns, opnames=None, maxplots=1000):
    def square(x):
        return np.absolute(x)**2
    df = read_file(inputfile).apply(square)
    ops = list(set([i/1000 for i in df.index]))
    levels = list(set([(i % 1000) for i in df.index]))
    logging.info("Read {} levels and {} ops".format(len(levels), len(ops)))
    N = int(np.sqrt(len(df)))
    largest_zfactor = max(df.identities)
    untranslated = opnames
    if opnames:
        opnames = [translate(l) for l in opnames]
    plots = ops
    if mode == "level":
        plots = levels

    plot = {}
    # plt.figure(figsize=(10, 6))
    plt.tick_params(axis='x', which='both', bottom='off', top='off', labelbottom='off')
    plt.rcParams['xtick.major.size'] = 0

    rows = min(int(math.ceil(float(len(plots))/cols)), int(math.ceil(float(maxplots)/cols)))
    if not args.seperate:
        f, layout = plt.subplots(nrows=rows, ncols=cols)
    # f.set_size_inches(19.2, 12.0)
    # f.set_dpi(100)
    for plot_index in plots[:maxplots]:
        if args.seperate:
            f, ax = plt.subplots(1)
        i = (plot_index-1)/cols
        j = (plot_index-1) % cols
        logging.info("making plot {} {} {}".format(plot_index, i, j))
        if not args.seperate:
            if len(plots[:maxplots]) <= cols or cols == 1:
                ax=layout[j]
            else:
                ax=layout[i][j]

        indexes = [plot_index+op*1000 for op in ops]
        if j > 0 and not args.seperate:
            ax.set_yticks([])
        else:
            ax.set_ylabel("$|Z|^2$", fontweight='bold')
            # ax.set_yticks([0, 1])

        if mode == "ops":
            with open(args.ordering) as orderfile:
                ordering = [int(i.strip())+1 for i in orderfile.readlines()] # orderings have levels indexed from 0, but these are indexed from 1

            indexes = [plot_index*1000+level for level in ordering]
            ticklabelpad = plt.rcParams['xtick.major.pad']
            ax.set_xticks(np.array(levels)+1.5)
            ax.set_xticklabels([n if n % 5 == 0 else "" for n in levels])
            ax.annotate('Level', xy=(1, 0), xytext=(-10, -ticklabelpad*2.2), ha='left', va='top',
                        xycoords='axes fraction', textcoords='offset points')
            if opnames:
                ax.set_title("{}".format(opnames[plot_index-1]))
            else:
                ax.set_title("operator {}".format(plot_index))
        else:
            ax.set_title("level {}".format(plot_index-1))
            ticklabelpad = plt.rcParams['xtick.major.pad']
            ax.annotate('Operator', xy=(1, 0), xytext=(-10, -ticklabelpad*2.2), ha='left', va='top',
                        xycoords='axes fraction', textcoords='offset points')

        values = df.ix[indexes].identities.values
        errors = df.ix[indexes].error.values
        if args.no_errors:
            errors = None
        # plots[(i, j)] = ax.bar(ops, values, yerr=errors)
        if mode == "level":
            plot[(i, j)] = ax.bar(ops[:ns], values[:ns], 1, color="g", yerr=errors)
            plot[(i, j)] = ax.bar(ops[ns:], values[ns:], 1, color="b", yerr=errors)
            ax.set_xticks(np.array(ops)+0.5)
            ax.set_xticklabels([o if o % 5 == 0 else "" for o in ops])
            plt.title("level {}".format(plot_index))
        else:
            color = "b"
            if plot_index <= ns:
                color = "g"
            levels = range(1,len(values)+1)
            plot[(i, j)] = ax.bar(levels, values, 1, color=color, yerr=errors, ecolor="r")
        ax.set_ylim([0, np.ceil(largest_zfactor)])
        ax.set_xlim(xmin=1)
        if args.seperate:
            ax.set_ylim([0, max(values)+max(errors)])
            outfilename = output_stub+"{}.png".format(untranslated[plot_index-1])
            logging.info("Saving plot to {}".format(outfilename))
            plt.savefig(outfilename)
            plt.clf()

    # end plot loop

    if args.seperate:
        return

    if args.title:
        f.suptitle(args.title)

    if(output_stub):
        plt.rcParams.update({'font.size': 8})
        plt.tight_layout(pad=2.0, h_pad=1.0, w_pad=2.0)
        logging.info("Saving plot to {}".format(output_stub+".png"))
        plt.savefig(output_stub+".png")
        # logging.info("Saving plot to {}".format(output_stub+".eps"))
        # plt.savefig(output_stub+".eps")
    else:
        plt.tight_layout()
        plt.show()
예제 #2
0
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()