plot_molecules3.py

from plot_common import *
from matplotlib.path import Path

import pandas as pd

prefix = 'plots/'
matplotlib.use('Agg')


ixs = {}

dx1 = 1.35
dx2 = 1.75
xx = 0


# 4
ixs['lih.xyz'] = xx # 19

xx += dx2
# 6
ixs['li2.xyz'] = xx # 28

xx += dx2
# 8
ixs['bh3.xyz'] = xx

xx += dx2
# 10
ixs['ne.xyz'] = xx        # 14
xx += dx1
ixs['h2o.xyz'] = xx  #  24
xx += dx1
ixs['nh3.xyz'] = xx   #  29
xx += dx1
ixs['ch4.xyz'] = xx   #  34

xx += dx2
# 12
ixs['beo.xyz'] = xx # 28
xx += dx1
ixs['lif.xyz'] = xx # 28
xx += dx1
ixs['c2.xyz'] = xx  # 28
xx += dx1
ixs['bn.xyz'] = xx  # 28

xx += dx2
# 14
ixs['n2-2-118.xyz'] = xx # 28
xx += dx1
# ixs['CO'] = xx # 28
ixs['co.xyz'] = xx # 28

xx += dx2
# 16
ixs['o2.xyz'] = xx #

xx += dx2
# 18
ixs['ar.xyz'] = xx # 18
xx += dx1
ixs['kr.xyz'] = xx # 27-18/2=18
xx += dx1
ixs['hcl.xyz'] = xx # 23

xx += dx2
# 20
ixs['licl.xyz'] = xx # 32

molecules = ixs.keys()

basis = 'ccpvdz'
files_molecules = glob.glob('output/normal/*_'+basis+'_*.h5')
data_molecules = load(files_molecules)

#     ndett = [64, 128, 256, 512, 768]
#     start = 3
if True:
    molecules_partial = ['c2.xyz', 'bn.xyz', 'n2-2-118.xyz', 'co.xyz', 'o2.xyz']

    # runs with partial optimization up to 2K determinants
    files_molecules_partial =list(glob.glob('output/partial/*_'+basis+'_*.h5'))
    data_molecules_partial = load(files_molecules_partial)

    for k, v in data_molecules_partial.items():
        if k[0] not in molecules_partial:
            continue
        new_key = *k[:-2], None, None
        if k not in data_molecules.keys() or (v['e'].real < data_molecules[new_key]['e'].real):
            data_molecules[new_key] = v


    ndett = [64, 128, 256, 512, 768, 2048]
    start = 3

basis = 'ccpvdz'

markers = ['s', 'D', 'v', 'o', '^']

start = 3

def fx(x):
    return x

ndetit = iter(ndett[start:])
mit = iter((markers+markers)[start:])
cit = iter((colors+colors)[start:])

alpha = 0.9

plt.figure(figsize=(18/2.54,6/2.54))

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

line_fci = None
line_dmrg = None

mols_all = list(ixs.keys())

eref = {}
ecc = {}
eref_is_fci = {}

l_ccsdt = None
for molecule in mols_all:

    efci = e_fci_ref.get(basis, {}).get(molecule, None)
    efci = e_fci_gao[basis].get(molecule, efci)

    edmrg = e_dmrg_block2[basis].get(molecule, None)
    edmrg = e_dmrg_chan[basis].get(molecule, edmrg)

    if efci is None and edmrg is not None:
        eref[molecule] = edmrg
        eref_is_fci[molecule] = False
    elif edmrg is None and efci is None:
        print('missing fci/dmrg', molecule, basis)
        continue
    else:
        eref[molecule] = efci
        eref_is_fci[molecule] = True

for molecule in mols_all:
    eccsdt = e_ccsdt_ref[basis].get(molecule, None)
    if eccsdt is None:
        print('missing ccsd(t)', molecule, basis)
        continue
    else:
        # rectangle
        mm = Path([[-1,-0.5],[1,-0.5],[1,0.5],[-1,0.5],[0,0]], [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY])
        lab = r'$\mathregular{CCSD(T)}$'
        ecc[molecule] = eccsdt
        xx = ixs[molecule]
        yy = eccsdt-eref[molecule]
        l_ccsdt, = plt.plot(xx, yy, marker=mm, linewidth=0, markerfacecolor='none', alpha=1.0, label=lab, markeredgewidth=1)

minix = min(ixs.values())
maxix = max(ixs.values())

plt.plot([fx(minix-1), fx(maxix+1)], [0, 0], color='k', alpha=0.4, linewidth=1.0)

# fake line to position legend
l_chemacc,= plt.plot([fx(minix-1), fx(maxix+1)], [0.00159360, 0.00159360], color='k', alpha=0.0, linestyle='dotted', linewidth=1.5, label='1 kcal/mol')
shift1 = 3.
l_chemacc,= plt.plot([fx(minix-1+shift1), fx(maxix+1)], [0.00159360, 0.00159360], color='k', alpha=0.4, linestyle='dotted', linewidth=1.0, label='1 kcal/mol')

lines = []

df_rows = ndett[start:] + ['FCI', 'DMRG', 'CCSD(T)']
df_cols = [molecule_dict.get(m, m).replace('$\\mathregular{', '').replace('$', '').replace('}', '').replace('^{\\star', '') for m in mols_all]
df_data = np.full((len(df_rows), len(df_cols)), None)

for l, ndet in enumerate(ndetit):
    EE = []
    ii = []
    SS = []
    mm = []
    for i, molecule in enumerate(mols_all):
        spin = 0
        if molecule == 'o2.xyz':
            spin = 2
        val = data_molecules.get((molecule, basis, spin, ndet, None, None))
        # if val is None and molecule == 'o2.xyz':
        #     spin = 0
        #     val = data_molecules.get((molecule, basis, spin, ndet, None, None))
        #     if val is None:
        #         print('o2: using spin=0 simulation', molecule, basis, ndet)
        if val is None:
            print('missing data', molecule, basis, ndet, None, None)
            continue

        EE.append(val['e'].real-eref.get(molecule, None))
        SS.append(val['s2'].real)
        ii.append(ixs[molecule])
        mm.append(molecule)

        df_data[l,i] = float(val['e'].real)
        if eref_is_fci.get(molecule):
            df_data[-3,i] = eref.get(molecule, None)
        else:
            df_data[-2,i] = eref.get(molecule, None)
        df_data[-1, i] = e_ccsdt_ref[basis].get(molecule, None)


    ii = np.array(ii)
    EE = np.array(EE)
    perm = np.argsort(ii)
    if ndet > 768:
        l, = plt.plot(fx(ii[perm]), EE[perm], label=r"$\mathregular{{"+f"{ndet}"+r"}^\dagger}$", linewidth=0, marker=next(mit), markerfacecolor=matplotlib.colors.to_rgb(next(cit))+(alpha,))
    else:
        l, = plt.plot(fx(ii[perm]), EE[perm], label=f"{ndet}", linewidth=0, marker=next(mit),  markerfacecolor=matplotlib.colors.to_rgb(next(cit))+(alpha,))
    lines.append(l,)
    print('\n ndet S2', ndet)
    for m_, S2_ in zip(mm, SS):
        print(m_, S2_)


ii = np.array([ixs[m] for m in mols_all])
perm = np.argsort(ii)

plt.ylabel(r'$\mathregular{E-E_{ref}\quad[a.u.]}$')
plt.xlabel(r'MOLECULE', labelpad=10)

plt.xticks(fx(ii[perm]))
plt.gca().set_xticklabels([molecule_dict.get(m, m) for m in np.array(mols_all)[perm]])

plt.ylim(-0.00049, 0.00249)

plt.xlim(0-1, xx+1)

legend1 = plt.legend(handles=lines, title=r'$\quad\mathregular{N_D}$', frameon=False, bbox_to_anchor=(0.125, 0.48), handletextpad=0.1)
plt.gca().add_artist(legend1)


if l_ccsdt is not None:
    legend2 = plt.legend(handles=[l_ccsdt,], bbox_to_anchor=(0.145, 0.9), frameon=False, handletextpad=0.6, handlelength=0.8)
    plt.gca().add_artist(legend2)

plt.gca().text(minix-0.25, 0.00159360, '1 kcal/mol', fontsize=7, verticalalignment='center', alpha=0.75)

plt.savefig(prefix+'fig1.pdf')

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 \ molecule"
    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:
        df.to_excel(writer, sheet_name="fig1", index=True)