<jupyter_start><jupyter_text>Make the MPNN ModelWe are going to layers from the [NFP](https://github.com/NREL/nfp) package from NREL to build out MPNN model.NFP is built using Tensorflow. So, we will need to make a data loader that will generate TF-compatible data and a Keras Model object.We use `tf.keras` for all operations so that we have a path forward to using this code with TF2.<jupyter_code>%matplotlib inline
from matplotlib import pyplot as plt
from graphsage.mpnn.data import *
from graphsage.mpnn.layers import GraphNetwork, Squeeze
from tensorflow.keras.layers import Input, Lambda
from tensorflow.keras.models import Model
from tensorflow.keras.models import load_model
from tqdm import tqdm
import tensorflow as tf
import pickle as pkl
import numpy as np
import gzip
import json
import os<jupyter_output><empty_output><jupyter_text>Configure TF for local GPU use<jupyter_code>tf_config=tf.ConfigProto()
tf_config.gpu_options.allow_growth=True
sess = tf.Session(config=tf_config)
data_file = os.path.join('..', 'data', 'output', 'water_clusters.proto')<jupyter_output><empty_output><jupyter_text>Assemble IteratorMake an iterator over the dataset that will generate batches of networks with each batch collapsed to a single "supernetwork"<jupyter_code>r = make_data_loader(data_file, batch_size=32, shuffle_buffer=32768)<jupyter_output>WARNING:tensorflow:From C:\Users\ward6\Miniconda3\envs\graphsage\lib\site-packages\tensorflow_core\python\ops\array_ops.py:1475: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.where in 2.0, which has the same broadcast rule as np.where<jupyter_text>Get an example batch<jupyter_code>with tf.Session() as sess:
    i = r.make_one_shot_iterator()
    batch = sess.run(i.get_next())<jupyter_output>WARNING:tensorflow:From <ipython-input-5-6c83cd1c3b55>:2: DatasetV1.make_one_shot_iterator (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `for ... in dataset:` to iterate over a dataset. If using `tf.estimator`, return the `Dataset` object directly from your input function. As a last resort, you can use `tf.compat.v1.data.make_one_shot_iterator(dataset)`.<jupyter_text>Display the specification for the input data<jupyter_code>r.element_spec<jupyter_output><empty_output><jupyter_text>We are not going to be saving the data loader. We'll remake it during the actual training procedure.For now, we are using it for testing purposes. See how random the batches are<jupyter_code>batches = []
with tf.Session() as sess:
    n = r.make_one_shot_iterator().get_next()
    batches = [sess.run(n) for _ in tqdm(range(256))]<jupyter_output>100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 256/256 [00:00<00:00, 353.10it/s]<jupyter_text>Plot the changes in the mean number of atoms<jupyter_code>fig, ax = plt.subplots()

ax.errorbar(
    range(len(batches)),
    [b[0]['n_atom'].mean() for b in batches],
    fmt='o',
    yerr=[b[0]['n_atom'].std() for b in batches],
)<jupyter_output><empty_output><jupyter_text>Make the ModelCreate the Keras Model object that describes the neural network.<jupyter_code>def build_fn(atom_features=32, message_steps=8):
    node_graph_indices = Input(shape=(1,), name='node_graph_indices', dtype='int32')
    atom_types = Input(shape=(1,), name='atom', dtype='int32')
    bond_types = Input(shape=(1,), name='bond', dtype='int32')
    connectivity = Input(shape=(2,), name='connectivity', dtype='int32')
    
    # Squeeze the node graph and connectivity matrices
    snode_graph_indices = Squeeze(axis=1)(node_graph_indices)
    satom_types = Squeeze(axis=1)(atom_types)
    sbond_types = Squeeze(axis=1)(bond_types)
    
    output = GraphNetwork(2, 2, atom_features,
                          message_steps, name='mpnn')([satom_types, sbond_types, snode_graph_indices, connectivity])
    
    return Model(inputs=[node_graph_indices, atom_types, bond_types, connectivity],
                 outputs=output)
network = build_fn()<jupyter_output>WARNING:tensorflow:From C:\Users\ward6\Miniconda3\envs\graphsage\lib\site-packages\tensorflow_core\python\keras\initializers.py:119: calling RandomUniform.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.
Instructions for updating:
Call initializer instance with the dtype argument instead of passing it to the constructor
WARNING:tensorflow:From C:\Users\ward6\Miniconda3\envs\graphsage\lib\site-packages\tensorflow_core\python\ops\resource_variable_ops.py:1630: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.<jupyter_text>We know the average energy for each atom, and will use it to help the network start out near the correct solution<jupyter_code>atom_energy = np.mean(batch[1][:, 0] / batch[0]['n_atom'])
l = network.get_layer('mpnn')
weights = l.atom_mean.get_weights()
weights[0][:] = atom_energy
l.atom_mean.set_weights(weights)<jupyter_output><empty_output><jupyter_text>Print out the performance without any training<jupyter_code>fig, ax = plt.subplots()

pred_y = network.predict_on_batch(batch[0])
ax.scatter(batch[1], pred_y)

ax.set_ylabel('$E$, Predicted')
ax.set_xlabel('$E$, True')

ax.plot(ax.get_xlim(), ax.get_xlim(), 'k--')<jupyter_output><empty_output><jupyter_text>Save model to disk<jupyter_code>network.save('model.h5')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13891285591026747, 'fraction_numerical': 0.015875754961173424, 'mean_word_length': 6.327433628318584, 'pattern_counts': {'":': 0, '<': 43, '<?xml version=': 0, '>': 42, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '622443', 'license': 'Apache-2.0', 'n_tokens_mistral': 2094, 'n_tokens_neox': 1921, 'n_words': 481, 'path': 'mpnn/0_create-models.ipynb', 'repo_name': 'exalearn/design-tutorial', 'size': '5795'}