def plot(
        probs,
        accepts,
        h_olds,
        h_news,
        exp_delta_hs,
        subtitle,
        labels,
        save,
        h_mdmc=None,
        mdmc_deltaH=None,  # for mdmc additions
):
    """Plots 3 stacked figures:
        1. the acceptance probability at each step
        2. the hamiltonian (old, new) at each step
        3. the exp{-delta H} at each step
    Overlayed with red bars is each instance in which a configuration was rejected
    by the Metropolis-Hastings accept/reject step
    
    Required Inputs
        probs           :: np.array :: acc. probs
        accepts         :: np.array :: array of boolean acceptances (True = accepted)
        h_olds          :: np.array :: old hamiltonian at each step
        h_news          :: np.array :: new hamiltonian at each step
        exp_delta_hs    :: np.array :: exp{-delta H} at each step
        h_mdmc          :: np.array :: 
        mdmc_deltaH     :: np.array ::
        subtitle        :: str      :: the subtitle to put in ax[0].set_title()
        save            :: bool :: True saves the plot, False prints to the screen
    """
    pp = Pretty_Plotter()
    pp._teXify()  # LaTeX
    # pp.params['text.latex.preamble'] = [r"\usepackage{amsmath}"]
    pp._updateRC()

    fig, ax = plt.subplots(3, sharex=True, figsize=(8, 8))
    # fig.suptitle(r'Data from {} Metropolis acceptance steps'.format(len(probs)),
    #     fontsize=pp.ttfont)

    fig.subplots_adjust(hspace=0.2)

    l = probs.size  # length of HMC trajectory - lots rely on this being at the top

    ### Add top pseudo-title and bottom shared x-axis label
    ax[0].set_title(subtitle, fontsize=pp.tfont)
    ax[-1].set_xlabel(r'Markov Time, $j$')

    ### add the rejection points in the background
    xrng = range(1, l + 1)  # calculate xrange to match length
    for a in ax:  # iterate over each axis
        for x, val in zip(xrng, accepts):  # iterate over all rejection points
            if val == False:
                a.axvline(x=x, linewidth=4, color='red', alpha=0.2)

    ### add the lines to the plots
    ax[0].set_ylabel(r'Hamiltonian, $H_j$')
    ax[0].plot(xrng,
               h_olds,
               linestyle='-',
               color='blue',
               linewidth=2.,
               alpha=1,
               label=r'$H_{j-1}$')
    # ax[0].plot(xrng, h_news, linestyle='-', color='green', linewidth=2., alpha=0.4,
    #     label=r'$H(t)$')

    ax[1].set_ylabel(r'$\exp{-\delta H_j}$')

    ax[1].plot(xrng,
               exp_delta_hs,
               linestyle='-',
               color='blue',
               linewidth=2.,
               alpha=1,
               label=r'$\exp{ -\delta H_t}$')

    ax[2].set_ylabel(r'Acceptance $\rho_j$')
    ax[2].plot(xrng, probs, linestyle='-', color='blue', linewidth=2., alpha=1)

    ### test to see if we will plot all the intermediate MDMC steps
    plot_mdmc = (h_mdmc is not None) & (mdmc_deltaH is not None)

    if plot_mdmc:  # all these functions rely on l being calculated!!
        # functions to calculate the staggering of y and x
        # for the MDMC. Staggering means calculating the start and
        # end points of each single trajectory consisting of 1 MDMC integration
        n = h_mdmc.size / l - 1  # this calculates n_steps for the MDMC (in a backwards way)
        staggered_xi = np.arange(0, l)
        staggered_xf = np.arange(1, l + 1)
        staggered_y = lambda arr, offset: arr[0 + offset::(n + 1)]
        remove_links = lambda arr: [
            None if (i) % (n + 1) == 0 else a for i, a in enumerate(arr)
        ]

        ## calculate a linear space as a fraction of the HMC trajectory
        mdmc_x = np.asarray([
            np.linspace(i, j, n + 1)
            for i, j in zip(range(0, l), range(1, l + 1))
        ])
        mdmc_x = mdmc_x.flatten()

        ax[0].plot(mdmc_x,
                   remove_links(h_mdmc),
                   linestyle='--',
                   color='green',
                   linewidth=1,
                   alpha=1,
                   label=r'MDMC: $H_{j+1}$')
        ax[0].scatter(staggered_xi,
                      staggered_y(h_mdmc, 0),
                      color='red',
                      marker='o',
                      alpha=1,
                      label=r'Start MDMC')
        ax[0].scatter(staggered_xf,
                      staggered_y(h_mdmc, n),
                      color='red',
                      marker='x',
                      alpha=1,
                      label=r'End MDMC')

        ax[1].plot(mdmc_x,
                   remove_links(mdmc_deltaH),
                   linestyle='--',
                   color='blue',
                   linewidth=1.,
                   alpha=1,
                   label=r'MDMC: $\exp{ -\delta H_{j+1}$')
        ax[1].scatter(staggered_xi,
                      staggered_y(mdmc_deltaH, 0),
                      color='red',
                      marker='o',
                      alpha=1,
                      label=r'Start MDMC')
        ax[1].scatter(staggered_xf,
                      staggered_y(mdmc_deltaH, n),
                      color='red',
                      marker='x',
                      alpha=1,
                      label=r'End MDMC')

    ### adds labels to the plots
    for i, text in labels.iteritems():
        pp.add_label(ax[i], text, fontsize=pp.tfont)

    ### place legends
    ax[0].legend(loc='upper left',
                 shadow=True,
                 fontsize=pp.ipfont,
                 fancybox=True)
    ax[1].legend(loc='upper left',
                 shadow=True,
                 fontsize=pp.ipfont,
                 fancybox=True)

    ### formatting
    for i in ax:
        i.grid(True)
    # ax[1].set_yscale("log", nonposy='clip') # set logarithmic y-scale

    pp.save_or_show(save, PLOT_LOC)
    pass
Пример #2
0
def plot(lines_d, x_lst, ws, subtitle, mcore, angle_labels, op_name, save):
    """Plots the two-point correlation function
    
    Required Inputs
        x_lst    :: list :: list of x_values for each angle that was run
        lines_d  :: {axis:[(y,label)]} :: plots (y,error,label) for each list item
        ws       :: list :: list of integration windows
        subtitle :: str  :: subtitle for the plot
        mcore    :: bool :: are there multicore operations? (>1 mixing angles)
        op_name  :: str  :: the name of the operator for the title
        angle_labels :: list :: the angle label text for legend plotting
        save     :: bool :: True saves the plot, False prints to the screen
    """

    pp = Pretty_Plotter()
    pp._teXify()  # LaTeX
    pp.params['text.latex.preamble'] = r"\usepackage{amsfonts}"
    pp.params['text.latex.preamble'] = r"\usepackage{amsmath}"
    pp._updateRC()

    fig = plt.figure(figsize=(8, 8))  # make plot
    ax = []

    fig, ax = plt.subplots(3, sharex=True, figsize=(8, 8))
    fig.suptitle(r"Autocorrelation and Errors for {}".format(op_name),
                 fontsize=pp.ttfont)

    fig.subplots_adjust(hspace=0.1)

    # Add top pseudo-title and bottom shared x-axis label
    ax[0].set_title(subtitle, fontsize=pp.tfont)
    ax[-1].set_xlabel(r'Window Length')

    if not mcore:  # don't want clutter in a multiple plot env.
        for a in range(1, len(ax)):  # Add the Window stop point as a red line
            # there is only one window item if not multiple lines
            ax[a].axvline(x=ws[0], linewidth=4, color='red', alpha=0.1)

    ax[0].set_ylabel(r'$g(w)$')
    ax[1].set_ylabel(r'$\tau_{\text{int}}(w)$')
    ax[2].set_ylabel(r'Autocorrelation, $\Gamma(t)$')

    #### plot for the 0th axis ####
    line_list = lines_d[0]
    axis = ax[0]
    theory_colours = iter(colour)
    measured_colours = iter(colour)
    for x, lines in zip(x_lst, line_list):
        m = next(marker)  # get next marker style
        c = next(measured_colours)  # get next colour
        th_c = next(theory_colours)
        # split into y function, errors in y and label
        y, e, l = lines  # in this case there are no errors or labels used

        # allow plots in diff colours for +/
        yp = y.copy()
        yp[yp < 0] = np.nan  # hide negative values
        ym = y.copy()
        ym[ym >= 0] = np.nan  # hide positive values

        if not mcore:
            axis.scatter(x,
                         yp,
                         marker='o',
                         color='g',
                         linewidth=2.,
                         alpha=0.6,
                         label=r'$g(t) \ge 0$')
            axis.scatter(x,
                         ym,
                         marker='x',
                         color='r',
                         linewidth=2.,
                         alpha=0.6,
                         label=r'$g(t) < 0$')
        else:
            axis.plot(x, yp, color=c, lw=1., alpha=0.6)  # label with the angle
            axis.plot(x, ym, color='r', lw=1., alpha=0.6)
            once_label = None  # set to blank so don't get multiple copies

    if not mcore: axis.legend(loc='best', shadow=True, fontsize=pp.axfont)

    #### plot the 1st axis ###
    # there is no angle label on the lines themselves on this axis
    # because the colours are synchronised across each plot
    # so the label on the bottom axis is enough
    line_list = lines_d[1]
    axis = ax[1]
    theory_colours = iter(colour)
    measured_colours = iter(colour)
    for x, lines in zip(x_lst, line_list):
        m = next(marker)  # get next marker style
        c = next(measured_colours)  # get next colour
        th_c = next(theory_colours)
        y, e, l, t = lines  # split into y function, errors, label and theory
        # try:
        #     axis.fill_between(x, y-e, y+e, color=c, alpha=0.5)
        # except:
        #     print "errors are dodgy"
        axis.errorbar(x,
                      y,
                      yerr=e,
                      markersize=3,
                      color=c,
                      fmt=m,
                      alpha=0.5,
                      ecolor='k')
        if t is not None:
            axis.axhline(y=t, linewidth=1, color=th_c, linestyle='--')

        # Only add informative label if there is only one line
        # adds a pretty text box above the middle plot with info
        # contained in the variable l - assigned in preparePlot()
        if not mcore: pp.add_label(axis, l, fontsize=pp.tfont)

    #### plot the 2nd axis ###
    # This plot explicitly list the labels for all the angles
    line_list = lines_d[2]
    axis = ax[2]
    theory_colours = iter(colour)
    measured_colours = iter(colour)
    for x, lines, a in zip(x_lst, line_list, angle_labels):
        m = next(marker)  # get next marker style
        c = next(measured_colours)  # get next colour
        th_c = next(theory_colours)
        y, e, l, t = lines  # split into y function, errors in y and label
        try:  # errors when there are low number of sims
            axis.fill_between(x, y - e, y + e, color=c, alpha=0.6, label=a)
            # axis.errorbar(x, y, yerr=e, label = a,
#                 markersize=3, color=c, fmt=m, alpha=0.5, ecolor='k')
        except:  # avoid crashing
            print 'Too few MCMC simulations to plot autocorrelations for: {}'.format(
                a)

        if t is not None:
            axis.plot(x,
                      t,
                      linewidth=1.2,
                      alpha=0.9,
                      color=th_c,
                      linestyle='--',
                      label='Theoretical')

    axis.legend(loc='best', shadow=True, fontsize=pp.axfont)

    #### start outdated section ####
    ## this won't work after the changes but shows the general idea of fitting a curve
    #
    # for i in range(1, len(lines)):              # add best fit lines
    #     x, y = lines[i][:2]
    #     popt, pcov = curve_fit(expFit, x, y)    # approx A+Bexp(-t/C)
    #     if not mcore:
    #         l_th = r'Fit: $f(t) = {:.1f} + {:.1f}'.format(popt[0], popt[1]) \
    #         + r'e^{-t/' +'{:.2f}'.format(popt[2]) + r'}$'
    #     else:
    #         l_th = None
    #     ax[i].plot(x, expFit(x, *popt), label = l_th,
    #         linestyle = '-', color=c, linewidth=2., alpha=.5)
    #### end outdated section ####

    # fix the limits so the plots have nice room
    xi, xf = ax[2].get_xlim()
    ax[2].set_xlim(xmin=xi - .05 * (xf - xi))  # decent view of the first point
    for a in ax:  # 5% extra room at top & add legend
        yi, yf = a.get_ylim()
        a.set_ylim(ymax=yf + .05 * (yf - yi), ymin=yi - .05 * (yf - yi))
    ax[-1].set_ylim(ymax=2)
    pp.save_or_show(save, PLOT_LOC)
    pass