{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Message Passing Neural Network\n\nSo, as many of you might have surmised by now the dataset for this challenge is essentially the QM9 dataset with some new values calculated for it. \n\nThe first thing I though of when seeing this challenge was the [Gilmer paper](https://arxiv.org/abs/1704.01212), as it uses the QM9 dataset. ([see this talk](https://vimeo.com/238221016))\n\nThe major difference in this challenge is that we are asked to calulate bond properties (thus edges in a graph) as opposed to bulk properties in the paper. \n\nHere the model is laid out in a modular way so the parts can easily be replaced\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Internet needs to be on\n!pip install tensorflow-gpu==2.0a0","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import tensorflow as tf\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nfrom sklearn.utils import shuffle\nimport os\nprint(os.listdir(\"../input\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Make sure tf 2.0 alpha has been installed\nprint(tf.__version__)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#is it using the gpu?\ntf.test.is_gpu_available(\n    cuda_only=False,\n    min_cuda_compute_capability=None\n)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"datadir = \"../input/\"\n\nnodes_train     = np.load(datadir + \"champs-basic-graph/nodes_train.npz\" )['arr_0']\nin_edges_train  = np.load(datadir + \"champs-basic-graph/in_edges_train.npz\")['arr_0']\nout_edges_train = np.load(datadir + \"champs-basic-graph/out_edges_train.npz\" )['arr_0']\n\nnodes_test     = np.load(datadir + \"champs-basic-graph/nodes_test.npz\" )['arr_0']\nin_edges_test  = np.load(datadir + \"champs-basic-graph/in_edges_test.npz\")['arr_0']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"out_labels = out_edges_train.reshape(-1,out_edges_train.shape[1]*out_edges_train.shape[2],1)\nin_edges_train = in_edges_train.reshape(-1,in_edges_train.shape[1]*in_edges_train.shape[2],in_edges_train.shape[3])\nin_edges_test  = in_edges_test.reshape(-1,in_edges_test.shape[1]*in_edges_test.shape[2],in_edges_test.shape[3])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x = pd.read_csv('../input/train.csv')\nx['numtype'] = x['type'].map({'2JHH': 1,\n                           '3JHH': 2,\n                           '1JHC': 3,\n                           '2JHC': 4,\n                           '3JHC': 5,\n                           '1JHN': 6,\n                           '2JHN': 7,\n                           '3JHN': 8 })\nscale_min  = x['scalar_coupling_constant'].min()\nscale_max = x['scalar_coupling_constant'].max()\nscale_mid = (scale_max + scale_min)/2\nscale_norm = scale_max - scale_mid\nnoofmolecules = len(x.molecule_name.unique())\nprint(noofmolecules)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from tqdm import tqdm\ntrain_nodes_array     = np.zeros((noofmolecules, 29, 29,2), dtype=np.float32) \nfor i, (k, m) in enumerate(tqdm(x.groupby('molecule_name'))):\n    ind = m[['atom_index_0', 'atom_index_1' ]].values\n    out_edges = np.zeros((29, 29, 2))\n    out_edges[ind[:,0], ind[:,1], 0] = (m['scalar_coupling_constant' ].values-84.3307)/120.5493\n    out_edges[ind[:,1], ind[:,0], 0] = out_edges[ind[:,0], ind[:,1], 0]\n    out_edges[ind[:,0], ind[:,1], 1] = (m['numtype' ].values)\n    out_edges[ind[:,1], ind[:,0], 1] = out_edges[ind[:,0], ind[:,1], 1]\n    train_nodes_array[i] = out_edges","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"outputwithclasses =  train_nodes_array.reshape(-1,29*29,2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"nodes_train, in_edges_train, out_labels = shuffle(nodes_train, in_edges_train, outputwithclasses)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Message passer\nDefine the message passer like the Gilmer paper\n\nUse a NN to embed edges as matrices, then matrix multiply with nodes."},{"metadata":{"trusted":true},"cell_type":"code","source":"\nclass Message_Passer_NNM(tf.keras.layers.Layer):\n    def __init__(self, node_dim):\n        super(Message_Passer_NNM, self).__init__()\n        self.node_dim = node_dim\n        self.nn = tf.keras.layers.Dense(units=self.node_dim*self.node_dim, activation = None)#activation=tf.nn.relu)\n      \n    def call(self, node_j, edge_ij):\n        \n        # Embed the edge as a matrix\n        A = self.nn(edge_ij)\n        \n        # Reshape so matrix mult can be done\n        A = tf.reshape(A, [-1, self.node_dim, self.node_dim])\n        node_j = tf.reshape(node_j, [-1, self.node_dim, 1])\n        \n        # Multiply edge matrix by node and shape into message list\n        messages = tf.linalg.matmul(A, node_j)\n        messages = tf.reshape(messages, [-1, tf.shape(edge_ij)[1], self.node_dim])\n\n        return messages","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Aggregator\n\nDefine the message aggregator (just sum)  \nProbably overkill to have it as its own layer, but good if you want to replace it with something more complex\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Message_Agg(tf.keras.layers.Layer):\n    def __init__(self):\n        super(Message_Agg, self).__init__()\n    \n    def call(self, messages):\n        return tf.math.reduce_sum(messages, 2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Update function\n\nDefine the Update function (a GRU)  \nThe GRU basically runs over a sequence of length 2, i.e. [ old state, agged_messages ]"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Update_Func_GRU(tf.keras.layers.Layer):\n    def __init__(self, state_dim):\n        super(Update_Func_GRU, self).__init__()\n        self.concat_layer = tf.keras.layers.Concatenate(axis=1)\n        self.GRU = tf.keras.layers.GRU(state_dim)\n        \n    def call(self, old_state, agg_messages):\n    \n        # Remember node dim\n        n_nodes  = tf.shape(old_state)[1]\n        node_dim = tf.shape(old_state)[2]\n        \n        # Reshape so GRU can be applied, concat so old_state and messages are in sequence\n        old_state = tf.reshape(old_state, [-1, 1, tf.shape(old_state)[-1]])\n        agg_messages = tf.reshape(agg_messages, [-1, 1, tf.shape(agg_messages)[-1]])\n        concat = self.concat_layer([old_state, agg_messages])\n        \n        # Apply GRU and then reshape so it can be returned\n        activation = self.GRU(concat)\n        activation = tf.reshape(activation, [-1, n_nodes, node_dim])\n        \n        return activation","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Output layer\n\nThis is where the model diverges with the paper.   \nAs the paper predicts bulk properties, but we are interested in edges, we need something different.   \n\nHere the each edge is concatenated to it's two nodes and a MLP is used to regress the scalar coupling for each edge"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define the final output layer \nclass Edge_Regressor(tf.keras.layers.Layer):\n    def __init__(self, intermediate_dim):\n        super(Edge_Regressor, self).__init__()\n        self.concat_layer = tf.keras.layers.Concatenate()\n        self.hidden_layer_1 = tf.keras.layers.Dense(units=intermediate_dim, activation=tf.nn.relu)\n        self.hidden_layer_2 = tf.keras.layers.Dense(units=intermediate_dim, activation=tf.nn.relu)\n        self.output_layer = tf.keras.layers.Dense(units=1, activation=None)\n\n        \n    def call(self, nodes, edges):\n            \n        # Remember node dims\n        n_nodes  = tf.shape(nodes)[1]\n        node_dim = tf.shape(nodes)[2]\n        \n        # Tile and reshape to match edges\n        state_i = tf.reshape(tf.tile(nodes, [1, 1, n_nodes]),[-1,n_nodes*n_nodes, node_dim ])\n        state_j = tf.tile(nodes, [1, n_nodes, 1])\n        \n        # concat edges and nodes and apply MLP\n        concat = self.concat_layer([state_i, edges, state_j])\n        activation_1 = self.hidden_layer_1(concat)  \n        activation_2 = self.hidden_layer_2(activation_1)\n\n        return self.output_layer(activation_2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Message passing layer\n\nPut all of the above together to make a message passing layer which does one round of message passing and node updating"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define a single message passing layer\nclass MP_Layer(tf.keras.layers.Layer):\n    def __init__(self, state_dim):\n        super(MP_Layer, self).__init__(self)\n        self.message_passers  = Message_Passer_NNM(node_dim = state_dim) \n        self.message_aggs    = Message_Agg()\n        self.update_functions = Update_Func_GRU(state_dim = state_dim)\n        \n        self.state_dim = state_dim         \n\n    def call(self, nodes, edges, mask):\n      \n        n_nodes  = tf.shape(nodes)[1]\n        node_dim = tf.shape(nodes)[2]\n        \n        state_j = tf.tile(nodes, [1, n_nodes, 1])\n\n        messages  = self.message_passers(state_j, edges)\n\n        # Do this to ignore messages from non-existant nodes\n        masked =  tf.math.multiply(messages, mask)\n        \n        masked = tf.reshape(masked, [tf.shape(messages)[0], n_nodes, n_nodes, node_dim])\n\n        agg_m = self.message_aggs(masked)\n        \n        updated_nodes = self.update_functions(nodes, agg_m)\n        \n        nodes_out = updated_nodes\n        # Batch norm seems not to work. \n        #nodes_out = self.batch_norm(updated_nodes)\n        \n        return nodes_out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Put it all together to form a MPNN\n\nDefines the full mpnn that does T message passing steps, where T is a hyperparameter.   \nAs in the paper, the same MP layer is re-used, but this is not a requirement. "},{"metadata":{"trusted":true},"cell_type":"code","source":"adj_input = tf.keras.Input(shape=(None,), name='adj_input')\nnod_input = tf.keras.Input(shape=(None,), name='nod_input')\nclass MPNN(tf.keras.Model):\n    def __init__(self, out_int_dim, state_dim, T):\n        super(MPNN, self).__init__(self)   \n        self.T = T\n        self.embed = tf.keras.layers.Dense(units=state_dim, activation=tf.nn.relu)\n        self.MP = MP_Layer( state_dim)     \n        self.edge_regressor  = Edge_Regressor(out_int_dim)\n        self.batch_norm = tf.keras.layers.BatchNormalization() \n\n        \n    def call(self, inputs =  [adj_input, nod_input]):\n      \n      \n        nodes            = inputs['nod_input']\n        edges            = inputs['adj_input']\n\n        # Get distances, and create mask wherever 0 (i.e. non-existant nodes)\n        # This also masks node self-interactions...\n        # This assumes distance is last\n        len_edges = tf.shape(edges)[-1]\n        \n        _, x = tf.split(edges, [len_edges -1, 1], 2)\n        mask =  tf.where(tf.equal(x, 0), x, tf.ones_like(x))\n        \n        # Embed node to be of the chosen node dimension (you can also just pad)\n        nodes = self.embed(nodes) \n        \n        nodes = self.batch_norm(nodes)\n        # Run the T message passing steps\n        for mp in range(self.T):\n            nodes =  self.MP(nodes, edges, mask)\n        \n        # Regress the output values\n        con_edges = self.edge_regressor(nodes, edges)\n        \n        \n        return con_edges\n        ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define the loss functions. \n\nHere the losses are MSE, MAE and LMAE.  \n(**note**: that for LMAE, as the values have been scaled down values will be much smaller than for unscaled values)"},{"metadata":{"trusted":true},"cell_type":"code","source":"def mse(orig , preds):\n \n    # Mask values for which no scalar coupling exists\n    mask  = tf.where(tf.equal(orig, 0), orig, tf.ones_like(orig))\n\n    nums  = tf.boolean_mask(orig,  mask)\n    preds = tf.boolean_mask(preds,  mask)\n\n\n    reconstruction_error = tf.reduce_mean(tf.square(tf.subtract(nums, preds)))\n\n\n    return reconstruction_error","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def log_mse(orig , preds):\n \n    # Mask values for which no scalar coupling exists\n    mask  = tf.where(tf.equal(orig, 0), orig, tf.ones_like(orig))\n\n    nums  = tf.boolean_mask(orig,  mask)\n    preds = tf.boolean_mask(preds,  mask)\n\n\n    reconstruction_error = tf.math.log(tf.reduce_mean(tf.square(tf.subtract(nums, preds))))\n\n\n    return reconstruction_error","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def mae(orig , preds):\n \n    # Mask values for which no scalar coupling exists\n    mask  = tf.where(tf.equal(orig, 0), orig, tf.ones_like(orig))\n\n    nums  = tf.boolean_mask(orig,  mask)\n    preds = tf.boolean_mask(preds,  mask)\n\n\n    reconstruction_error = tf.reduce_mean(tf.abs(tf.subtract(nums, preds)))\n\n\n    return reconstruction_error","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def log_mae(orig , preds):\n \n    # Mask values for which no scalar coupling exists\n    mask  = tf.where(tf.equal(orig, 0), orig, tf.ones_like(orig))\n\n    nums  = tf.boolean_mask(orig,  mask)\n    preds = tf.boolean_mask(preds,  mask)\n\n    reconstruction_error = tf.math.log(tf.reduce_mean(tf.abs(tf.subtract(nums, preds))))\n\n    return reconstruction_error","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def GroupScore(orig , preds):\n    # Mask values for which no scalar coupling exists\n    cnt = tf.constant(1e-9)\n    fullmask  = tf.keras.backend.flatten(orig[:,:,1])\n    nums  = tf.math.multiply(tf.keras.backend.flatten(orig[:,:,0]), tf.constant(120.5493))\n    preds = tf.math.multiply(tf.keras.backend.flatten(preds), tf.constant(120.5493))\n    \n    \n    mask = tf.where(tf.keras.backend.equal(fullmask,tf.keras.backend.constant(1.)))\n    a1 = (tf.reduce_mean(tf.keras.backend.mean(tf.keras.backend.abs(tf.keras.backend.gather(nums,mask)-\n                                                                               tf.keras.backend.gather(preds,mask)))))\n    \n    \n    mask = tf.where(tf.keras.backend.equal(fullmask,tf.keras.backend.constant(2.)))\n    a2 = (tf.reduce_mean(tf.keras.backend.mean(tf.keras.backend.abs(tf.keras.backend.gather(nums,mask)-\n                                                                               tf.keras.backend.gather(preds,mask)))))\n    \n    mask = tf.where(tf.keras.backend.equal(fullmask,tf.keras.backend.constant(3.)))\n    a3 = (tf.reduce_mean(tf.keras.backend.mean(tf.keras.backend.abs(tf.keras.backend.gather(nums,mask)-\n                                                                               tf.keras.backend.gather(preds,mask)))))\n    \n    mask = tf.where(tf.keras.backend.equal(fullmask,tf.keras.backend.constant(4.)))\n    a4 = (tf.reduce_mean(tf.keras.backend.mean(tf.keras.backend.abs(tf.keras.backend.gather(nums,mask)-\n                                                                               tf.keras.backend.gather(preds,mask)))))\n    \n    mask = tf.where(tf.keras.backend.equal(fullmask,tf.keras.backend.constant(5.)))\n    a5 = (tf.reduce_mean(tf.keras.backend.mean(tf.keras.backend.abs(tf.keras.backend.gather(nums,mask)-\n                                                                               tf.keras.backend.gather(preds,mask)))))\n    \n    mask = tf.where(tf.keras.backend.equal(fullmask,tf.keras.backend.constant(6.)))\n    a6 = (tf.reduce_mean(tf.keras.backend.mean(tf.keras.backend.abs(tf.keras.backend.gather(nums,mask)-\n                                                                               tf.keras.backend.gather(preds,mask)))))\n    \n    mask = tf.where(tf.keras.backend.equal(fullmask,tf.keras.backend.constant(7.)))\n    a7 = (tf.reduce_mean(tf.keras.backend.mean(tf.keras.backend.abs(tf.keras.backend.gather(nums,mask)-\n                                                                               tf.keras.backend.gather(preds,mask)))))\n    \n    mask = tf.where(tf.keras.backend.equal(fullmask,tf.keras.backend.constant(8.)))\n    a8 = (tf.reduce_mean(tf.keras.backend.mean(tf.keras.backend.abs(tf.keras.backend.gather(nums,mask)-\n                                                                               tf.keras.backend.gather(preds,mask)))))\n    #tf.keras.backend.log(tf.keras.backend.maximum(cnt,))\n    classscores = tf.stack([a1,a2,a3,a4,a5,a6,a7,a8])\n    mask = tf.where(tf.keras.backend.not_equal(classscores,tf.keras.backend.constant(0.)))\n    vals = tf.keras.backend.gather(classscores,mask)\n    vals = tf.math.maximum(cnt,vals)\n    vals = tf.math.log(vals)\n    vals = tf.reduce_mean(tf.keras.backend.mean(vals))\n    return vals-tf.math.log(cnt)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define some callbacks, the initial learning rate and the optimizer"},{"metadata":{"trusted":true},"cell_type":"code","source":"learning_rate = 0.001\ndef step_decay(epoch):\n    initial_lrate = learning_rate\n    drop = 0.1\n    epochs_drop = 20.0\n    lrate = initial_lrate * np.power(drop,  \n           np.floor((epoch)/epochs_drop))\n    tf.print(\"Learning rate: \", lrate)\n    return lrate\n\nlrate = tf.keras.callbacks.LearningRateScheduler(step_decay)\nstop_early = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience = 15, restore_best_weights=True)\n\n#lrate  =  tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1,\n#                              patience=5, min_lr=0.00001, verbose = 1)\n\nopt = tf.optimizers.Adam(learning_rate=learning_rate)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Finally create the model, and compile"},{"metadata":{"trusted":true},"cell_type":"code","source":"mpnn = MPNN(out_int_dim = 512, state_dim = 128, T = 4)\nmpnn.compile(opt, GroupScore, metrics = [GroupScore])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Define some hyperparameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_size = int(len(out_labels)*0.8)\nbatch_size = 16\nepochs = 25\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Call once as tensorflow likes this"},{"metadata":{"trusted":true},"cell_type":"code","source":"mpnn.call({'adj_input' : in_edges_train[:10], 'nod_input': nodes_train[:10]})","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Let the learning begin!"},{"metadata":{"trusted":true},"cell_type":"code","source":"mpnn.fit({'adj_input' : in_edges_train[:train_size], 'nod_input': nodes_train[:train_size]}, y = out_labels[:train_size], batch_size = batch_size, epochs = epochs, \n         callbacks = [lrate, stop_early], use_multiprocessing = True, initial_epoch = 0, verbose = 2, \n         validation_data = ({'adj_input' : in_edges_train[train_size:], 'nod_input': nodes_train[train_size:]},out_labels[train_size:]) )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Predict on the test set"},{"metadata":{"trusted":true},"cell_type":"code","source":"preds = mpnn.predict({'adj_input' : in_edges_test, 'nod_input': nodes_test})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"np.save(\"preds_kernel.npy\" , preds)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Prediction done!\n\nNow rescale outputs and create submission.csv"},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv(datadir + \"champs-scalar-coupling/train.csv\")\ntest = pd.read_csv(datadir + \"champs-scalar-coupling/test.csv\")\n\ntest_group = test.groupby('molecule_name')\n\nscale_min  = train['scalar_coupling_constant'].min()\nscale_max = train['scalar_coupling_constant'].max()\nscale_mid = (scale_max + scale_min)/2\nscale_norm = scale_max - scale_mid","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_outs(test_group, preds):\n    i = 0\n    x = np.array([])\n    for test_gp, preds in zip(test_group, preds):\n        if (not i%1000):\n            print(i)\n\n        gp = test_gp[1]\n        \n        x = np.append(x, (preds[gp['atom_index_0'].values, gp['atom_index_1'].values] + preds[gp['atom_index_1'].values, gp['atom_index_0'].values])/2.0)\n        \n        i = i+1\n    return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"max_size = 29\npreds = preds.reshape((-1,max_size, max_size))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"out_unscaled = make_outs(test_group, preds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test['scalar_coupling_constant'] = out_unscaled\ntest['scalar_coupling_constant'] = test['scalar_coupling_constant']*scale_norm + scale_mid\ntest[['id','scalar_coupling_constant']].to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}