1-RDM and 2-RDM learning / training
Unified QMLearn trainer for γ, δγ, Γ, Γc, and the cumulant Δ. It reads HDF5 databases, scans KRR hyperparameters, and plots potential-energy curves against a CASCI/FCI reference.
You need: a training database (see create_training.ipynb) with 1-RDM and 2-RDM properties, plus a reference energy curve (for example casci_exact_s.npz).
Edit dbfile_corr / dbfile_cum and path below to match your files.
import numpy as np
from qmlearn.io.model import db2qmmodel
from qmlearn.io import read_db, merge_db
import matplotlib.pyplot as plt
from sklearn.kernel_ridge import KernelRidge
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score,mean_absolute_error
import os
from sklearn.metrics import mean_squared_error
import pandas as pd
import matplotlib as mpl
import matplotlib.ticker as ticker
import matplotlib
mpl.rcParams.update(mpl.rcParamsDefault)
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Times New Roman'] + plt.rcParams['font.serif']
fcolor='tab:gray'
font = {'weight' : 'normal',
'size' : 14}
matplotlib.rc('font', **font)
conv_kcal = 627.503
path = './'
Generalized the $\Gamma$, $\Delta$, $\Gamma_c$, $\gamma$ and $\delta \gamma$ fitting and prediction
Read DB and define basis, method, ncas, nelecas
basis='6-31g*'
method = 'casci'
ncas=18
nelecas=10
dbfile_corr = "./H2O_27_6-31g*casci1810_nis_f.hdf5"
dbfile_cum = path+"/H2O_27_fci.hdf5"
data_corr = read_db(dbfile_corr)
data_cum = read_db(dbfile_cum)
Guess DB names : {'qmmol': 'casci/qmmol', 'atoms': 'casci/train_atoms_27', 'properties': 'casci/train_props_27'}
Guess DB names : {'qmmol': 'casci/qmmol', 'atoms': 'casci/train_atoms_27', 'properties': 'casci/train_props_27'}
kernels = ["rbf"]
alphas = [0, 1e-10, 1e-8, 1e-6, 4.5e-4, 1e-4, 4.5e-3, 1e-3, 4.5e-2, 1e-2, 1e-1]
gammas = [None, 1e-8, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 0.1]
def model_creation(dbfile_corr, dbfile_cum, read_model = False, write_model = False, params = None):
'''
Read model activates only for the best parameter
When keeping the model,, do it including the kernel, gamma and alpha used
'''
if params is None:
params = {
key: {"kernel": "rbf", "alpha": 0.0, "gamma": 1e-5}
for key in ["gamma", "delta_gamma", "gamma2cum", "gamma2c", "gamma2"]
}
models = {
key: KernelRidge(**params[key])
for key in params
}
if read_model:
model_g2_db = read_db('model_g2.hdf5')
model_g2 = model_g2_db['model']
model_dg_db = read_db('model_dg.hdf5')
model_dg = model_dg_db['model']
model_g1_db = read_db('model_g1.hdf5')
model_g1 = model_g1_db['model']
model_g2cum_db = read_db('model_g2cum.hdf5')
model_g2cum = model_g2cum_db['model']
model_g2c_db = read_db('model_g2.hdf5')
model_g2c = model_g2cum_db['model']
else:
model_g2 = db2qmmodel(data_corr, mmodels=models,
target='gamma2',method='gamma2',
purify_gamma=False)
model_dg = db2qmmodel(data_cum, mmodels=models,
target='delta_gamma',method='delta_gamma',
purify_gamma=False)
model_g2cum = db2qmmodel(data_cum, mmodels=models,
target='gamma2cum',method='gamma2cum',
purify_gamma=False)
model_g1 = db2qmmodel(data_cum, mmodels=models,
target='gamma',method='gamma',
purify_gamma=False)
model_g2c = db2qmmodel(data_corr, mmodels=models,
target='gamma2c',method='gamma2c',
purify_gamma=False)
if write_model:
model_g1.write('model_g1.hdf5')
model_g2cum.write('model_g2cum.hdf5')
model_g2.write('model_g2.hdf5')
model_dg.write('model_dg.hdf5')
model_g2c.write('model_g2c.hdf5')
return model_g1, model_g2cum, model_g2, model_dg, model_g2c
def get_hf_qmmol(atoms, models, hf = True):
# HF calculation and initialization of qmmol object
refqmmol = models[0].refqmmol
qmmol = refqmmol.duplicate(atoms)
if hf:
qmmol.engine.mf.run()
gamma_hf_ = qmmol.engine.mf.make_rdm1()
results = qmmol, gamma_hf_
else:
results = qmmol
return results
def get_predicted_gammas(atoms, models, target="all"):
qmmol, gamma_hf_ = get_hf_qmmol(atoms=atoms, models=models)
shape = gamma_hf_.shape
shape2 = (shape[0],) * 4
maps = {}
# delta_gamma / gamma2cum
if target in ["delta_gamma", "gamma2cum", "all"]:
delta_gamma_predicted = models[3].predict(qmmol,model=models[3].mmodels["delta_gamma"]).reshape(shape)
gamma_full,_ = qmmol.engine.purify_d_gamma(
gamma_d=delta_gamma_predicted,
gamma_hf=gamma_hf_,
pure=False)
gamma_full_pure,_ = qmmol.engine.purify_d_gamma(
gamma_d=delta_gamma_predicted,
gamma_hf=gamma_hf_,
pure=True)
Gamma_cum = models[1].predict(qmmol, model=models[1].mmodels["gamma2cum"]).reshape(shape2)
Gamma_full,_ = qmmol.engine.purify_gamma2cum(
gamma=gamma_full,
gamma2cum=Gamma_cum,
pure=False)
Gamma_full_pure,_ = qmmol.engine.purify_gamma2cum(
gamma=gamma_full_pure,
gamma2cum=Gamma_cum,
pure=True)
maps[0] = [gamma_full, Gamma_full]
maps[1] = [gamma_full_pure, Gamma_full_pure]
# gamma2c
if target in ["gamma2c", "all"]:
Gamma_corr = models[4].predict(qmmol, model=models[4].mmodels["gamma2c"]).reshape(shape2)
gamma_corr = qmmol.engine.gamma1_f_gamma2(gamma2=Gamma_corr)
gamma_full,_ = qmmol.engine.purify_d_gamma(
gamma_d=gamma_corr,
gamma_hf=gamma_hf_,
pure=False)
gamma_full_pure,_ = qmmol.engine.purify_d_gamma(
gamma_d=gamma_corr,
gamma_hf=gamma_hf_,
pure=True)
Gamma_full,_ = qmmol.engine.purify_gamma2c(
gamma=gamma_full,
gamma2c=Gamma_corr,
pure=False)
Gamma_full_pure,_ = qmmol.engine.purify_gamma2c(
gamma=gamma_full_pure,
gamma2c=Gamma_corr,
pure=True)
maps[2] = [gamma_full, Gamma_full]
maps[3] = [gamma_full_pure, Gamma_full_pure]
# gamma / gamma2
if target in ["gamma", "gamma2", "all"]:
Gamma = models[2].predict(qmmol, model=models[2].mmodels["gamma2"]).reshape(shape2)
# gamma = models[0].predict(qmmol, model=models[0].mmodels["gamma"]).reshape(shape)
gamma = qmmol.engine.gamma1_f_gamma2(gamma2=Gamma)
gamma_pure = qmmol.engine.purify_gamma(
gamma,
method="smearing")
Gamma_pure = qmmol.engine.purify_gamma2_n(
gamma2=Gamma,
gamma=gamma_pure)
maps[4] = [gamma, Gamma]
# maps[4] = [gamma_pure, Gamma]
maps[5] = [gamma_pure, Gamma_pure]
return maps
def get_energy_maps(atoms, models, target="all"):
maps = get_predicted_gammas(atoms, models, target)
qmmol = get_hf_qmmol(atoms=atoms, models=models, hf=False)
energies = {}
for key in maps:
# print(key, np.shape(maps[key][0]), np.shape(maps[key][1]))
energies[key] = qmmol.calc_etotal2(
gamma=maps[key][0],
gamma2=maps[key][1],
ao_repr=True
)
# print(energies)
return energies
xa = np.arange(0.5, 1.7, 0.05)
i0 = 0
i1 = 1
exact_path = path+'./casci_exact_s.npz'
if os.path.exists(exact_path):
data = np.load(exact_path)
energy_casci = data['r_eh']
g_ex = data['g_ex']
g2cum_ex = data['g2cum_ex']
g2_ex = data['g2_ex']
Run optimization
target_maps = {
"delta_gamma": [0, 1], # no pure \Delta, pure \Delta
"gamma2cum": [0, 1], # no pure \Delta, pure \Delta
"gamma2c": [2, 3], # no pure \Gamma^c, pure \Gamma^c
"gamma2": [4, 5], # no pure \Gamma, pure \Gamma
"gamma": [4, 5], # no pure \gamma, pure \gamma #NOT USED!
"all": [0, 1, 2, 3, 4, 5],
# no pure \Delta, pure \Delta, no pure \Gamma^c, pure \Gamma^c, no pure \Gamma, pure \Gamma
}
coupled_targets = {
"gamma2cum": ["delta_gamma","gamma2cum"],
"gamma2c": ["gamma2c"],
"gamma2": ["gamma2"]
}
def optimize_models(dbfile_corr, dbfile_cum, atoms, energy_ref, xa,
kernels, alphas, gammas, target):
min_idx = np.where(energy_ref==np.min(energy_ref))
ref_curve = (energy_ref-energy_ref[min_idx])
default = {"kernel":"rbf","alpha":0.0,"gamma":1e-5}
results = []
maps_to_optimize = target_maps[target]
models_to_optimize = coupled_targets[target]
for kernel in kernels:
for alpha in alphas:
for gamma in gammas:
print(f"{target:12s} "
f"{kernel:10s} "
f"{alpha:.2e} "
f"{gamma}")
params = {"gamma":default.copy(),
"delta_gamma":default.copy(),
"gamma2cum":default.copy(),
"gamma2c":default.copy(),
"gamma2":default.copy()}
for model_name in models_to_optimize:
params[model_name] = {"kernel":kernel,"alpha":alpha,"gamma":gamma}
models = model_creation(dbfile_corr,dbfile_cum,params=params)
pred = {m:[] for m in maps_to_optimize}
for r in xa:
a = atoms.copy()
a.set_distance(i0, i1, r, fix=0)
energies = get_energy_maps(a, models, target=target)
for m in maps_to_optimize:
pred[m].append(energies[m])
for m in maps_to_optimize:
curve = np.asarray(pred[m])
curve = (curve-curve[min_idx])
rmse = np.sqrt(mean_squared_error(ref_curve, curve))
mae = mean_absolute_error(ref_curve, curve)
results.append({"optimized_model":target,
"map":m+1,"kernel":kernel,"alpha":alpha,
"gamma":gamma,"RMSE":rmse,"MAE":mae})
results = pd.DataFrame(results)
best = results.loc[results.groupby("map")["RMSE"].idxmin()].reset_index(drop=True)
return results, best
params = {
"gamma": {"kernel":"rbf","alpha":0.0,"gamma":None},
"delta_gamma": {"kernel":"rbf","alpha":0.0,"gamma":None},
"gamma2cum": {"kernel":"rbf","alpha":0.0,"gamma":None},
"gamma2c": {"kernel":"rbf","alpha":0.0,"gamma":None},
"gamma2": {"kernel":"rbf","alpha":0.0,"gamma":None}}
models = model_creation(dbfile_corr, dbfile_cum, params = params)
Finish the reading. Finish the reading. Finish the reading. Finish the reading. Finish the reading.
atoms = models[0].refqmmol.atoms.copy()
targets = ["gamma2cum",
"gamma2c",
"gamma2",]
best_models = {}
all_models = {}
for target in targets:
print(f"\nOptimizing {target}")
results, best = optimize_models( dbfile_corr, dbfile_cum, atoms, energy_casci,
xa, kernels, alphas, gammas, target=target)
best_models[target] = best
all_models[target] = results
np.save(f'best_parameters_{kernels[0]}.npy',best_models)
best_params = np.load(f'best_parameters_{kernels[0]}.npy', allow_pickle=True).item()
df_best = pd.concat(best_params.values(), ignore_index=True)
df_best
| optimized_model | map | kernel | alpha | gamma | RMSE | MAE | |
|---|---|---|---|---|---|---|---|
| 0 | gamma2cum | 1 | rbf | 4.500000e-02 | 1.000000e-04 | 0.002599 | 0.002208 |
| 1 | gamma2cum | 2 | rbf | 1.000000e-10 | 1.000000e-05 | 0.003035 | 0.001857 |
| 2 | gamma2c | 3 | rbf | 1.000000e-01 | 1.000000e-02 | 0.055599 | 0.040385 |
| 3 | gamma2c | 4 | rbf | 1.000000e-08 | 1.000000e-03 | 0.006035 | 0.003858 |
| 4 | gamma2 | 5 | rbf | 1.000000e-08 | 1.000000e-05 | 0.105668 | 0.060861 |
| 5 | gamma2 | 6 | rbf | 1.000000e-08 | 1.000000e-08 | 0.050349 | 0.025605 |
best5_df = (pd.concat([df.nsmallest(5, "RMSE").assign(method=method) for method, df in all_models.items()],
ignore_index=True).sort_values(["method", "MAE"]).reset_index(drop=True))
best5_df
With best parameters! Plot
#RBF
params = {
"gamma": {"kernel":kernels[0],"alpha":0.0,"gamma":None},
"delta_gamma": {"kernel":kernels[0],"alpha":0.0,"gamma":None},
"gamma2cum": {"kernel":kernels[0],"alpha":0.0,"gamma":None},
"gamma2c": {"kernel":kernels[0],"alpha":0.0,"gamma":None},
"gamma2": {"kernel":kernels[0],"alpha":1e-8,"gamma":1e-5}}
models = model_creation(dbfile_corr, dbfile_cum, params = params)
Finish the reading. Finish the reading. Finish the reading. Finish the reading. Finish the reading.
print(params)
{'gamma': {'kernel': 'rbf', 'alpha': 0.0, 'gamma': None}, 'delta_gamma': {'kernel': 'rbf', 'alpha': 0.0, 'gamma': None}, 'gamma2cum': {'kernel': 'rbf', 'alpha': 0.0, 'gamma': None}, 'gamma2c': {'kernel': 'rbf', 'alpha': 0.0, 'gamma': None}, 'gamma2': {'kernel': 'rbf', 'alpha': 1e-08, 'gamma': 1e-05}}
images = data_corr['atoms']
bonds = []
for a in images:
d = a.get_distance(i0, i1)
bonds.append(d)
bonds = np.asarray(bonds)
#Plot 1
shift = np.min(energy_casci)
min_ = np.where(energy_casci==np.min(energy_casci))
fig, ax = plt.subplots(figsize=(6, 4))
ax2 = ax.twinx()
ax.plot(xa, (p_map5f-p_map5f[min_])*conv_kcal, linestyle='--', alpha = 1.0,
label=r"$\Gamma_{\rm ML}$", color='#D81B60', linewidth = 3)
ax.plot(xa, (p_map3f-p_map3f[min_])*conv_kcal, linestyle=':', alpha=1.0,
label=r"$\Gamma^{\rm C}_{\rm ML}$", color='#0000FE', linewidth = 3)
ax.plot(xa, (p_map1f-p_map1f[min_])*conv_kcal, linestyle='-.', alpha=1.0,
label=r"$\Delta_{\rm ML}$",color='#FFC107', linewidth = 3)
ax.plot(xa,(energy_casci-energy_casci[min_])*conv_kcal, linestyle='-', alpha=0.5,
label='FCI',color='#004D40',linewidth = 3)
ax.set_ylabel(r'Energy (kcal/mol) ') #Respect to GS energy.
ax.set_xlabel(r'$d_{\operatorname{O-H}} (\mathrm{\AA})$')
ax.set_xlim(0.6, 1.6)
ax.set_ylim(-6, 300)
ax.grid(True, axis='y', linestyle='--', linewidth=0.5, alpha=0.7)
#
legend_font = {'weight': 'normal', 'size': 11}
ax.legend(prop=legend_font,loc='upper center',ncol=2)
n, bins, patches = plt.hist(bonds, 10, facecolor=fcolor, alpha=0.4)
ax2.yaxis.set_major_formatter(ticker.PercentFormatter(xmax=len(bonds), decimals=0))
ax2.set_ylabel("Training set distribution", color=fcolor)
ax2.set_yticks(np.arange(0, 0.5, 0.1)*len(bonds))
ax2.tick_params(axis="y", labelcolor=fcolor)
plt.tight_layout()
fig.savefig(f'pes_{params["gamma2"]["kernel"]}.pdf')
plt.show()
plt.close()
min_ = np.where(energy_casci==np.min(energy_casci))
ref_curve = (energy_casci-energy_casci[min_])
curve = (p_map1f-p_map1f[min_])
rmse = np.sqrt(mean_squared_error(ref_curve, curve))
mae = mean_absolute_error(ref_curve, curve)
print(rmse,mae)
0.011368632975769798 0.004524516070609626
# Plot 2
fig, ax = plt.subplots(figsize=(6, 4))
ax2 = ax.twinx()
ax.plot(xa,(p_map6f-p_map6f[min_])*conv_kcal,linestyle='--',alpha=1.0,label=r"$\Gamma_{ML}^{\rm Pure}$",
color='#D81B60',linewidth = 3)
ax.plot(xa,(p_map4f-p_map4f[min_])*conv_kcal,linestyle=':',alpha=1.0,label=r"$\Gamma^{\rm c, Pure}_{\rm ML}$",
color='#0000FE',linewidth = 3)
ax.plot(xa,(p_map2f-p_map2f[min_])*conv_kcal,linestyle='-.',alpha=1.0,label=r"$\Delta_{\rm ML}^{\rm Pure}$",
color='#FFC107',linewidth = 3)
ax.plot(xa,(energy_casci-energy_casci[min_])*conv_kcal, linestyle='-', alpha=0.5,
label='FCI',color='#004D40',linewidth = 3)
ax.set_ylabel(r'Energy (kcal/mol) ') #Respect to GS energy.
ax.set_xlabel(r'$d_{\operatorname{O-H}} (\mathrm{\AA})$')
ax.set_xlim(0.6, 1.6)
ax.set_ylim(-6, 300)
ax.grid(True, axis='y', linestyle='--', linewidth=0.5, alpha=0.7)
#
legend_font = {'weight': 'normal', 'size': 11}
ax.legend(prop=legend_font,loc='upper center',ncol=2)
n, bins, patches = plt.hist(bonds, 10, facecolor=fcolor, alpha=0.4)
ax2.yaxis.set_major_formatter(ticker.PercentFormatter(xmax=len(bonds), decimals=0))
ax2.set_ylabel("Training set distribution", color=fcolor)
ax2.set_yticks(np.arange(0, 0.5, 0.1)*len(bonds))
ax2.tick_params(axis="y", labelcolor=fcolor)
plt.tight_layout()
fig.savefig(f'pes_pure_{params["gamma2"]["kernel"]}.pdf')
plt.show()
plt.close()
#plot 3
fig, ax = plt.subplots(figsize=(6, 4))
ax2 = ax.twinx()
ax.plot(xa,(p_map6f-p_map6f[min_])*conv_kcal,linestyle='--',alpha=1.0,label=r"$\Gamma_{ML}^{\rm Pure}$",
color='gray',linewidth = 3)
ax.plot(xa,(p_map5f-p_map5f[min_])*conv_kcal,linestyle='--',alpha=1.0,label=r"$\Gamma_{ML}$",
color='#D81B60',linewidth = 3)
ax.plot(xa,(energy_casci-energy_casci[min_])*conv_kcal, linestyle='-', alpha=0.5,
label='FCI',color='#004D40',linewidth = 3)
ax.set_ylabel(r'Energy (kcal/mol) ') #Respect to GS energy.
ax.set_xlabel(r'$d_{\operatorname{O-H}} (\mathrm{\AA})$')
ax.set_xlim(0.6, 1.6)
ax.set_ylim(-6, 300)
#
legend_font = {'weight': 'normal', 'size': 11}
ax.legend(prop=legend_font,loc='upper center',ncol=2)
ax.grid(True, axis='y', linestyle='--', linewidth=0.5, alpha=0.7)
n, bins, patches = plt.hist(bonds, 10, facecolor=fcolor, alpha=0.4)
ax2.yaxis.set_major_formatter(ticker.PercentFormatter(xmax=len(bonds), decimals=0))
ax2.set_ylabel("Training set distribution", color=fcolor)
ax2.set_yticks(np.arange(0, 0.5, 0.1)*len(bonds))
ax2.tick_params(axis="y", labelcolor=fcolor)
plt.tight_layout()
fig.savefig(f'pes_{params["gamma2"]["kernel"]}_gamma2_pur.pdf')
plt.show()
plt.close()
