# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import sklearn as sk
import ase
import ase.io as aio
import qml
import time
import gc
import scipy as sp
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from qml.kernels import gaussian_kernel
from qml.math import cho_solve
import metric_learn
import copy
# %config Completer.use_jedi = False
import warnings
warnings.filterwarnings("ignore")

# %matplotlib notebook

# +
data = pd.read_csv('./data.csv', index_col=0)

data.head()

# +
data['ase_reactants'] = [aio.read(
    'data_react_xyz/' + data.mol.values[i] + data.enan.values[i] + '.xyz') for i in range(len(data))]

data['ase_products'] = [aio.read(
    'data_react_xyz/' + data.mol.values[i] + data.enan.values[i] + '.xyz') for i in range(len(data))]

mbtypes = np.load('./mbtypes.npy', allow_pickle=True)

data['SLATM_react'] = [qml.representations.generate_slatm(coordinates=mol.positions, 
                                                 nuclear_charges=mol.get_atomic_numbers(), 
                                                mbtypes=mbtypes, local=False) for mol in data['ase_reactants']]

data['SLATM_prod'] = [qml.representations.generate_slatm(coordinates=mol.positions, 
                                                 nuclear_charges=mol.get_atomic_numbers(), 
                                                mbtypes=mbtypes, local=False) for mol in data['ase_products']]

# -

data['SLATM_diff'] = [x - y for x,y in zip(data.SLATM_prod.values, data.SLATM_react.values)]



extrapolation_data = pd.read_pickle('./extrapolation_data')

extrapolation_data.head()

mbtypes = np.load('./mbtypes.npy', allow_pickle=True)

slatm_react = [qml.representations.generate_slatm(coordinates=mol.positions, 
                                                 nuclear_charges=mol.get_atomic_numbers(), 
                                                mbtypes=mbtypes, local=False) for mol in extrapolation_data.ase2]

slatm_prod = [qml.representations.generate_slatm(coordinates=mol.positions, 
                                                 nuclear_charges=mol.get_atomic_numbers(), 
                                                mbtypes=mbtypes, local=False) for mol in extrapolation_data.ase3]

extrapolation_data['slatm2'] = slatm_react

extrapolation_data['slatm3'] = slatm_prod

extrapolation_data['slatm_diff'] = [x - y for x,y in zip(slatm_prod, slatm_react)]

extrapolation_data.to_pickle('extrapolation_data')

Y = data.Eafw.values

# # Extrapolation SLATM_DIFF+

slatmdiff = np.vstack(data.SLATM_diff.values)

lin_regs = [sp.stats.linregress(xs, ys) for xs, ys in zip(slatmdiff.T, 
                                                     [Y for i in range(
                                                         slatmdiff.shape[1])])]   
rvalues = np.array([ccc.rvalue for ccc in lin_regs])

# +
pre = sk.preprocessing.MinMaxScaler()
pre.fit(slatmdiff[:, np.argsort(rvalues**2)[-500:]])
Xf =  pre.transform(slatmdiff[:, np.argsort(rvalues**2)[-500:]])

dmtrtr = sk.metrics.pairwise.pairwise_distances(Xf, n_jobs=-1)

sigma = 6

regg = 1e-6

K = np.exp(- 0.5 * sigma ** -2 * dmtrtr ** 2)
K[np.diag_indices_from(K)] += regg
alpha_vec = cho_solve(K, Y)

# +
X_extrapolation = pre.transform(
            np.vstack(extrapolation_data.slatm_diff.values)[:, np.argsort(rvalues**2)[-500:]])
    
dmtrts = sk.metrics.pairwise.pairwise_distances(X_extrapolation, Xf, n_jobs=-1)

# extrapolation

Ks = np.exp(- 0.5 * sigma ** -2 * dmtrts ** 2)

Y_predicted = np.dot(Ks, alpha_vec)

# +
fig, ax = plt.subplots()

exact = 627.509 * (extrapolation_data.TS - extrapolation_data.int2).values
ax.scatter(Y_predicted, exact, 
          label='r$^2$ = {:.2f} \n MAE: {:.2f} kcal/mol'.format(
              np.corrcoef(Y_predicted, exact)[0][1] ** 2,
              np.mean(np.abs(Y_predicted - exact)))
          )

ax.plot(exact, exact, color='black', alpha=0.3)
ax.legend()
ax.set_xlabel('DFT E$_a$ [kcal/mol]')
ax.set_ylabel('ML E$_a$ [kcal/mol]')
plt.show()
# -







StarcoderdataJupyterScriptsDedupFiltered