plot_hydrogen_1d_lin.py

from plot_common import *
import argparse

prefix = 'plots/'
matplotlib.use('Agg')
parser = argparse.ArgumentParser()
parser.add_argument("--geom", type=str, help='interatomic distance of the molecule', default='1.8')
parser.add_argument("--out", type=str, help='filename of the output pdf figure', default=prefix+'fig10.pdf')

from si_prefix import si_format

pyscfu = True

args = parser.parse_args()

DMRG_M = [2000,]

ndetper = [5, 10]
bond_length=args.geom

pow = -1.

end = 5
ls = [10, 18, 22, 30, 50, 70, 102][:end]

molecules = [f'h{l}-{str(bond_length).replace('.', '-')}.xyz' for l in ls]
basis = 'sto6g'

data_o2 = {}
for molecule in molecules:
    data_o2 = {**data_o2, **load(glob.glob(f'output/hydrogen/linear*{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']

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

plt.subplots_adjust(
    left=0.15,
    right=0.99,
    bottom=0.125,
    top=0.98
)


mit = iter(['s', '^', 'p'])
lit = iter(linestyles)
cit = iter(colors)

hh2 = {}

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



df2_rows = ['UCISD', 'UCCSD', 'DMRG']
df2_cols = ls
df2_data = np.full((len(df2_rows), len(df2_cols)), None)


if True:
    # number of orthogonal determinants
    ndet_cisd = [876, 9316, 20934, 73126, 571876, 2211126, 10021654][:end]
    xx = []
    yy = []
    for molecule, l, ndet in zip(molecules, ls, ndet_cisd):
        try:
            yy.append(e_cisd_pyscf[basis][molecule]/l)
            xx.append(l)
            lab = str(si_format(ndet, precision=0 if  ndet < 1e3 or 1e5 <= ndet < 1e6 else 1)).replace('k', 'K').replace(' ', '')
            print(lab)
            ax1.annotate(lab, xy=(xx[-1]**pow, yy[-1]), xytext=(6, 4), textcoords="offset points",ha="center", size=5)
        except (KeyError, TypeError):
            print('no cisd', basis, molecule)
            pass
    lab = r'$\mathregular{UCISD}$'

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

if True:
    xx = []
    yy = []
    for molecule, l in zip(molecules, ls):
        try:
            yy.append(e_ccsd_pyscf[basis][molecule]/l)
            xx.append(l)
        except (KeyError, TypeError):
            print('no ccsd', basis, molecule)
            pass
    c = next(cit2)
    lab = r'$\mathregular{UCCSD}$'
    lin, = ax1.plot(np.array(xx)**pow, yy, label=lab, color='k', markeredgecolor='k', linestyle='dotted', marker='+', linewidth=0.8, markersize=5, markeredgewidth=1)
    hh2['ccsd'] = lin

    df2_data[1, :] = np.array(yy)*np.array(xx)

for M in DMRG_M:
    M = str(M)
    xx = []
    yy = []
    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:
                    yy.append(e_dmrg_block2_hydrogen[basis][molecule][k]/l)
                    xx.append(l)
                    err = False
            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(np.array(xx)**pow, yy, label=lab, color='k', markeredgecolor='k', linestyle='solid', marker='.', linewidth=0.8, markersize=5, markeredgewidth=1)
    hh2[M] = lin
    df2_data[2, :] = np.array(yy)*np.array(xx)



df_rows = list(sorted(set(np.outer(ndetper, ls).ravel().tolist())))
df_cols = ls
df_data = np.full((len(df_rows), len(df_cols)), None)


hh = {}
for n in ndetper:
    xx = []
    yy = []
    for j, (l, molecule) in enumerate(zip(ls, molecules)):
        ndet = l*n
        try:
            d = data_o2[molecule, basis, 0, ndet, None, None]
            yy.append(d['e'].real/l)
            xx.append(l)
            i, = np.where(ndet == np.array(df_rows))
            df_data[i, j] = d['e'].real
        except KeyError as e:
            print('not found', e)
    xx = np.array(xx)**pow
    yy = np.array(yy)
    lin, = ax1.plot(xx,yy, linewidth=1, color='k', marker=next(mit), label=n)
    hh[n] = lin


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=1, title=r'$\quad\mathregular{N_D N^{-1}}$')
ax1.add_artist(leg1)

uu = []
vv = []
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, bbox_to_anchor=(0.75, 0.45))
    ax1.add_artist(leg3)

ax1.set_ylabel(r'$\mathregular{E\, N^{-1}\,\, [a.u.]}$')
plt.ylim(-0.543-0.0002, -0.536+0.0002)
if pow == -1:
    ax1.set_xlabel(r'$\mathregular{N^{-1}}$')
    plt.xlim(1/50-0.005, 1/10+0.005)
    plt.xticks(1/np.array(ls), [f'${{{r}}}^{{-1}}$' for r in ls])

elif pow == 1:
    ax1.set_xlabel(r'$\mathregular{N}$')
else:
    assert False
ax1.set_title(r'$R='+f'{bond_length}\\, \\text{{a.u.}}$', y=1.0, pad=4)

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"

    df2 = pd.DataFrame(df2_data, df2_rows, df2_cols)
    df2.index.name = r"method \ N"

    path = Path(prefix+"source_data.xlsx")
    with pd.ExcelWriter(path, mode="a" if path.exists() else "w", engine="openpyxl", if_sheet_exists="replace" if path.exists() else None) as writer:
        df2.to_excel(writer, sheet_name="fig10", index=True)

    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="fig10", index=True, startrow=2+df2_data.shape[0])