import pyhande
import matplotlib.pyplot as plt
import pandas as pd
# Start averaging after the population has stabilised in all calculations.
calcs = pyhande.lazy.std_analysis(['calcs/ifciqmc/hubbard_ifciqmc.out'], 30000, extract_psips=True)
opt_blocks = [calc.opt_block for calc in calcs]
estimates = pd.concat([opt.stack() for opt in opt_blocks], keys=range(len(opt_blocks)), axis=1).T
plt.errorbar(estimates[('# H psips', 'mean')],estimates[('Proj. Energy', 'mean')],
             xerr=estimates[('# H psips', 'standard error')],
             yerr=estimates[('Proj. Energy', 'standard error')], label='iFCIQMC')

fciqmc_calc = pyhande.lazy.std_analysis(['calcs/fciqmc/hubbard_fciqmc.out'], 30000)[0]
fciqmc_energy = fciqmc_calc.opt_block['mean']['Proj. Energy']
fciqmc_err = fciqmc_calc.opt_block['standard error']['Proj. Energy']
plt.axhline(fciqmc_energy, label='FCIQMC', color='b')
plt.axhspan(fciqmc_energy+fciqmc_err, fciqmc_energy-fciqmc_err, color='b', alpha=0.25)

plt.xscale('log')
plt.xlabel('# particles')
plt.ylabel('Projected energy / $t$')
plt.legend()