# ---
# 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 pandas as pd

import matplotlib.pyplot as plt
import seaborn as sns
sns.set(context='talk', style='ticks',
        color_codes=True, rc={'legend.frameon': False})

# %matplotlib inline
# -

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)
import nfp

import sys
sys.path.append('..')

from alphazero.preprocessor import preprocessor

import sys
sys.path.append('/projects/rlmolecule/pstjohn/spin_gnn/20200901_combined_losses/')
from loss import AtomInfMask, KLWithLogits

# +
from nfp import EdgeUpdate, NodeUpdate, GlobalUpdate
import tensorflow_addons as tfa

model = tf.keras.models.load_model(
    '/projects/rlmolecule/svss/Project-Redox/water-model-tfrecords/20201013_stability_water_noh/best_model.hdf5',
    custom_objects={**nfp.custom_objects,
                    **{'AtomInfMask': AtomInfMask, 'KLWithLogits': KLWithLogits,
                       'EdgeUpdate': EdgeUpdate,
                       'NodeUpdate': NodeUpdate,
                       'GlobalUpdate': GlobalUpdate}})
# -

new_model = tf.keras.Model(model.inputs[:3], [tf.nn.softmax(model.outputs[0]), model.outputs[1]])

# +
shrees_df = pd.read_csv('/projects/rlmolecule/svss/Project-Redox/water-model-tfrecords/csv-spin-bv/buried_volumes_all_water_calc_includ_set_2.csv.gz', index_col=0)
shrees_spin = pd.read_csv('/projects/rlmolecule/svss/Project-Redox/water-model-tfrecords/csv-spin-bv/spins_water_calcs.csv')

shrees_df.head()
# -

smiles = '[CH2]C1=COCC1'

mol = rdkit.Chem.MolFromSmiles(smiles)

preprocessor.construct_feature_matrices(mol, train=False)

new_model({key: tf.constant(np.expand_dims(val, 0))
       for key, val in preprocessor.construct_feature_matrices(mol, train=False).items()})

shrees_df[shrees_df.smiles == smiles]

spin = shrees_spin[(shrees_spin.smiles == smiles) & (shrees_spin.atom_type != 'H')].spin

spin.abs() / spin.abs().sum()

# +
# import rdkit

# valid_dataset = tf.data.Dataset.from_generator(
#     lambda: (preprocessor.construct_feature_matrices(rdkit.Chem.MolFromSmiles(smiles), train=False) for smiles in valid),
#     output_types={'atom': tf.int64, 'bond': tf.int64, 'connectivity': tf.int64},
#     output_shapes={'atom': tf.TensorShape([None]), 'bond': tf.TensorShape([None]), 'connectivity': tf.TensorShape([None, 2])})\
#     .padded_batch(batch_size=100, padded_shapes={
#     'atom': [-1],
#     'bond': [-1],
#     'connectivity': [-1, 2]},
#     padding_values=nfp.zero)
# -

predicted_spin_logits, predicted_bur_vol = model.predict(valid_dataset, verbose=1)

cdf_test_pivot_spins = cdf_test.pivot(index='smiles', columns='atom_index', values='fractional_spin')
cdf_test_pivot_spins = cdf_test_pivot_spins.reindex(test)

new_model.save('/projects/rlmolecule/pstjohn/models/20201020_radical_stability_model/', include_optimizer=False)

# +
# model.save('/projects/rlmolecule/pstjohn/spin_gnn/test_tf_save', include_optimizer=False)
# -



StarcoderdataJupyterScriptsDedupFiltered