plot_bases.py

from plot_common import *
prefix = 'plots/'
matplotlib.use('Agg')

markers = ['+','x', '$▮$', 'v', 's', 'p']
molecule = 'lih.xyz'

files_lih = list(sorted(glob.glob('output/normal/lih_*.h5')))
data_lih = load(files_lih)

bases = ['sto6g', '631g', 'ccpvdz', 'ccpvtz', 'ccpvqz', 'ccpv5z']

fig, ax1 = plt.subplots(figsize=(8.8/2.54,6.6/2.54))
plt.subplots_adjust(
    left=0.125,
    right=0.99,
    bottom=0.125,
    top=0.98
)

start = 0
mit = iter(markers[start:])
cit = iter(colors[start:])
l1 = []

ndets = list(sorted(set([k[3] for k in data_lih.keys()])))
df_rows = ndets + ['FCI',]
df_cols = [bases_dict.get(b, b) for b in bases[start:]]
df_data = np.full((len(df_rows), len(df_cols)), None)

for i, basis in enumerate(bases[start:]):
    E = []
    ndet = []
    for k, v in data_lih.items():
        if k[0] != molecule or k[1] != basis:
            continue
        E.append(v['e'].real)
        ndet.append(k[3])
        nao = v['nao']

    E = np.array(E)
    ndet = np.array(ndet)
    perm = np.argsort(ndet)
    E = E[perm]
    ndet = ndet[perm]

    for ndeti, ei in zip(ndet, E):
        l, = np.where(ndeti == np.array(ndets))
        df_data[l,i] = ei.real
    try:
        df_data[-1,i] = e_fci_gao[basis]['lih.xyz']
    except KeyError:
        df_data[-1,i] = e_fci_ref[basis]['lih.xyz']


    basis = bases_dict.get(basis, basis)
    l = f"{basis} ({nao})"

    l, = ax1.plot(ndet, E, label=l, linewidth=0.0, marker=next(mit), markerfacecolor=matplotlib.colors.to_rgb(next(cit))+(0.9,))
    l1.append(l)


ax1.set_ylabel(r'$\mathregular{E\quad[a.u.]}$')
ax1.set_xscale('log', base=10)
ax1.set_xlabel(r'$\mathregular{N_D}$')

efci = []
ehf = []
l2 = []
for basis in bases:
    try:
        eref = e_fci_gao[basis]['lih.xyz']
    except KeyError:
        eref = e_fci_ref[basis]['lih.xyz']
    rr = 1024
    if basis in ['sto6g']:
        rr = 64
    l, = ax1.plot([1,rr*15/16], [eref, eref], color='k', linewidth=0.7, linestyle='dashed', label='FCI')
    plt.text(rr, eref-0.002, basis, fontsize=5.5)

    l2.append(l)
    efci.append(eref)

l1 = l1 + l2[:1]
leg1 = plt.legend(l1,  [l.get_label() for l in l1], title=r'     BASIS ($\mathregular{m}$)',bbox_to_anchor=(0.04, 0.48, 0.1, 0.3), frameon=False, fontsize=6)
ax1.add_artist(leg1)

shift = 0.05
shift2 = 0.05

l, d, w, h = [0.5+shift, 0.5+shift2, 0.48-shift, 0.47-shift2]
ax2 = fig.add_axes([l, d, w, h])

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

for basis in bases_sorted[start:]:
    E = []
    ndet = []
    for k, v in data_lih.items():
        if k[0] != molecule or k[1] != basis:
            continue
        E.append(v['e'].real)
        ndet.append(k[3])
        nao = v['nao']

    E = np.array(E)
    ndet = np.array(ndet)
    perm = np.argsort(ndet)
    E = E[perm]
    ndet = ndet[perm]
    try:
        eref = e_fci_gao[basis]['lih.xyz']
    except KeyError:
        eref = e_fci_ref[basis]['lih.xyz']

    err = E - eref

    basis = bases_dict.get(basis, basis)
    l = f"{basis} ({nao})"
    l, = ax2.plot(nao/ndet, err, label=l, linewidth=0.0, marker=next(mit), markerfacecolor=matplotlib.colors.to_rgb(next(cit))+(0.9,))

ax2.set_ylabel(r'$\mathregular{E-E_{FCI}\quad[a.u.]}$', labelpad=-2)

ax2.set_yscale('log')
ax2.set_xscale('log')
ax2.set_xlabel(r'$\mathregular{m N_D^{-1}}$')
ax2.set_ylim(3e-11, 4e-1)
ax2.set_xlim(2e-2, 5e2)
ax2.minorticks_off()

ax2.tick_params(axis='both', which='major')

ax1.set_ylim(-8.07, -7.8)
ax1.set_xlim(0.7, 2.2e3)


plt.savefig(prefix+f'fig3.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 \ basis"
    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="fig3", index=True)