plot_hydrogen_scaling.py

from plot_common import *
import argparse

prefix = 'plots/'
matplotlib.use('Agg')
parser = argparse.ArgumentParser()
parser.add_argument("geom", type=str)
parser.add_argument("--skipimpossible", action='store_true')
parser.add_argument("--nolegend", action='store_true')
parser.add_argument("--out", type=str, help='filename of the output pdf figure')
parser.add_argument("--startrow", type=int, default=0, help='start row for xlsx')
parser.add_argument("--replace", help='replace xlsx sheet', action='store_true')

args = parser.parse_args()
args.rhf = True

ndets = [1,2,4, 8, 16, 32, 64, 128, 256, 512]
bond_length=args.geom
ls = [10, 18, 22, 30, 50, 70, 102]
molecules = [f'h{l}-{str(bond_length).replace('.', '-')}.xyz' for l in ls]
basis = 'sto6g'

fig, ax1 = plt.subplots(figsize=(8.8/2.54,6.6/2.54))

plt.subplots_adjust(
    left=0.125,
    right=0.99,
    bottom=0.1,
    top=0.95
)

def prep(ooo, k):
    xx = []
    yy = []
    ss = []
    vv = []
    for ndet in tuple(sorted(set([k[3] for k in ooo.keys()]))):
        kk = *k[:3], ndet, *k[3:]
        if kk in ooo.keys():
            yy.append(ooo[kk]['e'].real)
            xx.append(ndet)
            ss.append(ooo[kk]['s2'].real)
            #vv.append(ooo[kk]['var'].real)
            vv.append(ndet)
    xx, yy, ss, vv = map(np.array, (xx,yy,ss,vv))
    return xx, yy, ss, vv

data_o2 = {}
for molecule in molecules:
    data_o2 = {**data_o2, **load(glob.glob(f'output/hydrogen/*{molecule.strip('.xyz')}_*'+basis+'*.h5'))}

data_o2 = {k[:-1]+(None,):v for k,v in data_o2.items()}

linestyles = ['solid','dashed', 'dashdot', (0, (4, 2, 1, 2, 1, 2)),(0, (4, 2, 1, 2, 1, 2,1,2)), (0, (4, 2, 1, 2, 1, 2,1,2,1,2)), 'dotted']

mit = iter(markers+markers+markers)
cit = iter(colors+colors+colors)
lit = iter(linestyles+linestyles)

to_plot = []

for ll, molecule in zip(ls, molecules):
    k = molecule, basis, 0, None, None
    l = r'$\mathregular{M=0}$'
    to_plot.append([ll, k,l, data_o2, 'o'])

M = 2000
M = str(M)
edmrg = {}
for molecule, l in zip(molecules, ls):
    try:
        keys = e_dmrg_block2_hydrogen[basis][molecule].keys()
        err=True
        for k in keys:
            if M in k:
                edmrg[l] = e_dmrg_block2_hydrogen[basis][molecule][k]
                err = False
        if err: raise KeyError
    except KeyError:
        print('no dmrg', basis, molecule, M)
        pass

if True:
    euhf = {}
    for molecule, l in zip(molecules, ls):
        try:
            euhf[l] = e_uhf_psi4[basis][molecule]
        except KeyError:
            print('no uhf', basis, molecule)
            pass


    erhf = {}
    for molecule, l in zip(molecules, ls):
        try:
            erhf[l] = e_rhf_psi4[basis][molecule]
        except KeyError:
            print('no rhf', basis, molecule)
            pass


    if args.rhf:
        ehf = erhf
        hf_str = 'RHF'
    else:
        ehf = euhf
        hf_str = 'UHF'

data = {}

cc = {}
# lines
for ll, k,l, o, m in to_plot:
    x,y,ss, vv = prep(o, k)
    e_det = y
    e_dmrg = edmrg.get(ll, np.nan)
    e_hf = ehf.get(ll, np.nan)
    err = (e_det-e_hf)/(e_dmrg-e_hf)
    mask = ~np.isnan(err)
    vv = vv[mask]
    err = err[mask]
    c = cc.get(ll, None)
    if c is None:
        c = cc[ll] = next(cit)
    if len(vv) > 0:
        data[ll] = vv, err

ylab = r'$\mathregular{\frac{E-E_{'+hf_str+r'}}{E_{DMRG }-E_{'+hf_str+r'}}}$'


mit = iter(markers+markers+markers)
cit = iter(colors+colors+colors)
lit = iter(linestyles+linestyles+linestyles)

levels = [0.99, 0.98, 0.95, 0.9, 0.75, 0.5]

df_rows = levels
df_cols = ls
df_data = np.full((len(df_rows), len(df_cols)), None)

if True:

    colors_lvl = {l:next(cit) for l in levels}
    # markers_lvl = {l:next(mit) for l in levels}
    ls_lvl = {l:next(lit) for l in levels}

    fit = {}

    for ii, lvl in enumerate(levels):
        print('lvl', lvl)
        fit[lvl] = [], []
        for k, (n,e) in data.items():
            print(k, n, e)
            if not ((e<lvl).all() or (e>lvl).all()):
                if True:
                    i = np.where(e>lvl)[0].min()-1
                    if args.skipimpossible:
                        if not ((e[i] < lvl) and (e[i+1] >=lvl)):
                            print('impossible data, cant interpolate')
                            continue
                    else:
                        assert e[i] < lvl
                        assert e[i+1] >=lvl

                    m, c = np.polyfit(e[i:i+2], np.log(n)[i:i+2], 1)
                    nd = np.exp(m*lvl+c)
                else:
                    i =  np.where(e<=lvl)[0].min()-1
                    if args.skipimpossible:
                        if not ((e[i] >= lvl) and (e[i+1] <=lvl)):
                            print('impossible data, cant interpolate')
                            continue
                    else:
                        assert e[i] >= lvl
                        assert e[i+1] <=lvl
                        m, c = np.polyfit(np.log(e)[i:i+2], np.log(n)[i:i+2], 1)
                    nd = np.exp(m*np.log(lvl)+c)
                print(lvl, e[i], e[i+1])
                print(nd, n[i], n[i+1])
                print()
                fit[lvl][0].append(k)
                fit[lvl][1].append(nd)
                j, = np.where(k == np.array(ls))
                try: df_data[ii, j] = nd
                except: from IPython import embed; embed()
        if len(fit[lvl][0]) < 2:
            del fit[lvl]

    plt.cla()
    for l, (x,y) in fit.items():
        llll = f'{l:0.2f}'
        plt.plot(x,y, label=llll, marker='.',markersize=1, linewidth=1.0, linestyle=ls_lvl[l], color=colors_lvl[l])
    if not args.nolegend:
        plt.legend(title=ylab+' level', frameon=False)
    plt.xlabel('$N$')
    plt.ylabel(r'$N_D$')
    # plt.xscale('log')
    plt.yscale('log')
    plt.ylim(2, 1024)
    plt.xlim(10-2, 102+2)
    plt.title(r'$R='+str(args.geom)+r'\, \text{a.u.}$', y=1.0, pad=4, size=6)


fname = args.out
assert fname[-4:] == '.pdf'
plt.savefig(fname)


if True:
    import pandas as pd
    from pathlib import Path
    df = pd.DataFrame(df_data, df_rows, df_cols)
    df.index.name = r"level \ N"

    fake_df = pd.DataFrame(np.full((0,0), None), [], [])
    fake_df.index.name = r"R="+str(args.geom)

    path = Path(prefix+"source_data.xlsx")
    with pd.ExcelWriter(path, mode="a" if path.exists() else "w", engine="openpyxl", if_sheet_exists=('replace' if args.replace else "overlay") if path.exists() else None) as writer:
        fake_df.to_excel(writer, sheet_name="fig9", index=True, startrow=args.startrow)
    with pd.ExcelWriter(path, mode="a" if path.exists() else "w", engine="openpyxl", if_sheet_exists="overlay") as writer:
        df.to_excel(writer, sheet_name="fig9", index=True, startrow=args.startrow+2)