plot_hydrogen_1d.py

from plot_common import *
import argparse

prefix = 'plots/'
matplotlib.use('Agg')
parser = argparse.ArgumentParser()
parser.add_argument("--normal", action='store_true')
parser.add_argument("--loss", action='store_true')
parser.add_argument("--spin", action='store_true')
parser.add_argument("--all", action='store_true') # minimum of all 3
parser.add_argument("--geom", type=str)
parser.add_argument("--start", type=int, default=0)
parser.add_argument("--nomarker", 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')

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

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

pyscfu = True

args = parser.parse_args()

#DMRG_M = [1000,2000, 3000, 4000]
DMRG_M = [2000,]

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

def S2toS(s2):
    s2 = np.array(s2)
    return np.sqrt(s2+0.25)-0.5

def prep(ooo, k):
    xx = []
    yy = []
    ss = []
    vv = []
    for l, molecule in zip(ls, molecules):
        kk = molecule, *k
        if kk in ooo.keys():
            yy.append(ooo[kk]['e'].real)
            xx.append(ndet)
            ss.append(ooo[kk]['s2'].real)
            vv.append(l)
    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 = ['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']

if args.nomarker:
    markers = ['.', ]*20
    markersizekwarg = {'markersize': 1}
else:
    markersizekwarg = {}

mit = iter(markers)
lit = iter(linestyles)

to_plot = []
for ndet in ndets:
    k = basis, 0, ndet, None, None
    l = r'$\mathregular{M=0}$'
    to_plot.append([k,l, data_o2, 'o'])

# cit = iter(colors)
hh2 = {}

cit2 = iter([(0, 0.137, 0.4), 'tab:orange', 'tab:green', 'tab:red', 'tab:purple'])

edmrgmin = 10000

ecisd = {}

df_rows = ndets + ['UCISD', 'UCCSD','DMRG']
df_cols = ls
df_data = np.full((len(df_rows), len(df_cols)), None)

if True:
    xx = []
    yy = []
    for i, (molecule, l) in enumerate(zip(molecules, ls)):
        try:
            if pyscfu:
                ey = e_cisd_pyscf[basis][molecule]
            else:
                ey = e_cisd_psi4[basis][molecule]
            yy.append(ey/l)
            xx.append(l)
            df_data[-3, i] = ey
        except (KeyError, TypeError):
            print('no cisd', basis, molecule)
            pass
    if pyscfu:
        lab = r'$\mathregular{UCISD}$'
    else:
        lab = r'$\mathregular{CISD}$'

    lin, = ax1.plot(1/np.array(xx), yy, label=lab, color='k', markeredgecolor='k', linestyle='dashed', marker='x', linewidth=0.8, markersize=3, markeredgewidth=1)
    hh2['cisd'] = lin

if True:
    xx = []
    yy = []
    for i, (molecule, l) in enumerate(zip(molecules, ls)):
        try:
            if pyscfu:
                ey = e_ccsd_pyscf[basis][molecule]
            else:
                ey = e_ccsd_psi4[basis][molecule]
            yy.append(ey/l)
            xx.append(l)
            df_data[-2, i] = ey
        except (KeyError, TypeError):
            print('no ccsd', basis, molecule)
            pass
    c = next(cit2)
    if pyscfu:
        lab = r'$\mathregular{UCCSD}$'
    else:
        lab = r'$\mathregular{CCSD}$'

    lin, = ax1.plot(1/np.array(xx), yy, label=lab, color='k', markeredgecolor='k', linestyle='dotted', marker='+', linewidth=0.8, markersize=3, markeredgewidth=1)
    hh2['ccsd'] = lin


if True:
    for M in DMRG_M:
        M = str(M)
        xx = []
        yy = []
        for i, (molecule, l) in enumerate(zip(molecules, ls)):
            try:
                keys = e_dmrg_block2_hydrogen[basis][molecule].keys()
                err=True
                for k in keys:
                    if M in k:
                        yy.append(e_dmrg_block2_hydrogen[basis][molecule][k]/l)
                        xx.append(l)
                        err = False
                        assert len(DMRG_M) == 1 # TODO support multiple
                        df_data[-1, i] = e_dmrg_block2_hydrogen[basis][molecule][k]
                if err: raise KeyError
            except KeyError:
                print('no dmrg', basis, molecule, M)
                pass
        c = next(cit2)
        lab = r'$\mathregular{DMRG}$'
        if len(DMRG_M) > 1:
            lab = r'$\mathregular{DMRG\, '+M+'}$'

        lin, = ax1.plot(1/np.array(xx), yy, label=lab, color='k', markeredgecolor='k', linestyle='solid', marker='.', linewidth=0.8, markersize=5, markeredgewidth=1)
        hh2[M] = lin

        edmrgmin = min(edmrgmin, min(yy))




colors = {k:v for k,v in zip(ndets, colors+colors)}
hh = {}
#markers
for k,l, o, m in to_plot:
    x,y,ss,vv = prep(o, k)
    xx = 1/vv
    e_det = y/vv
    yy = e_det
    lin, = ax1.plot(xx,yy, label=l,  linewidth=0, color=colors[k[2]], marker=m, **markersizekwarg)
    i, = np.where(k[2] == np.array(ndets))
    j = [np.where(np.array(ls)==u)[0][0].tolist() for u in vv]
    df_data[i, j] = y
    hh[l] = lin

hh3 = {}
lins = []
data = {}
# lines
for k,l, o, m in to_plot:
    x,y,ss, vv = prep(o, k)
    xx = 1/vv
    e_det = y/vv
    yy = e_det

    if len(xx) > 0:
        data[k[2]] = xx, yy

    lin, = ax1.plot(xx, yy,  linewidth=1, color=colors[k[2]], label=k[2])
    hh3[k[2]] = lin

if not args.nomarker:
    uu = []
    vv = []
    for h in hh.values():
        uu.append(h)
        vv.append(h.get_label())
    leg1 = plt.legend(handles=uu, labels=vv, frameon=False, handlelength=1, loc=2 if float(bond_length) >= 2.4 else 3)
    ax1.add_artist(leg1)



uu = []
vv = []

if not args.nolegend:
    uu = []
    vv = []
    for h in hh3.values():
        uu.append(h)
        vv.append(h.get_label())
    # vspace between ndet and dmrg
    uu.append(*ax1.plot(np.nan, np.nan, '-', color='none', label=''))
    vv.append('')
    for h in hh2.values():
        uu.append(h)
        vv.append(h.get_label())

    if len(uu) > 0:
        leg3 = plt.legend(handles=uu, labels=vv, frameon=False, handlelength=1.8, loc='lower right', title=r'$\quad\mathregular{N_D}$')
        ax1.add_artist(leg3)

ax1.set_ylabel(r'$\mathregular{E\, N^{-1}\,\, [a.u.]}$')
ax1.set_xlabel(r'$\mathregular{N^{-1}}$')
# ax1.set_xticks(1/np.array(ls), [f'${{{r}}}^{{-1}}$' for r in ls], size=5)
if not False:
    if args.geom == '1.0':
        edmrgmin = -0.3825
    else:
        edmrgmin = -0.5425
    ax1.set_ylim(edmrgmin-0.005, edmrgmin+0.055-0.005)


ax1.set_title(r'$R='+f'{bond_length}\\, \\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"N_D \ 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="fig8", 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="fig8", index=True, startrow=args.startrow+2)