{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Summary\n\n- V2, V3: a trick using list comprehension to concat pd.DataFrame to get groups in groupby after the keys are shuffled.; implementation of the starter code using generator, the generator is obtained from external sources\n- V4: trying new features from chainer\n- V5: debugged an identation error in the feature generation function\n\nTo do:\n* (Done) use generator to generate data.\n* add ACSF to node features and more edge features\n* introduce bond as an imaginary atom (make sense for 1J, not so much for 2J, 3J)"},{"metadata":{"trusted":true},"cell_type":"code","source":"INPUT_FOLDER = '../input/champs-scalar-coupling/'\nMODEL_FOLDER = '../input/mpnn-fit-generator/'\nSTRUCT_FOLDER = '../input/champs-scalar-coupling/structures/'\nDEBUG = False\nRETRAIN = False\nPREDICT = False\nMAX_SIZE = 29\nBATCH_SIZE = 16\nSCALE_MID = 84.3307\nSCALE_NORM = 120.5493\nLEARNING_RATE = 1e-4\nFACTOR_ACSF = 6\nFACTOR_DIST = 10\nALL_TYPES = ['1JHC','1JHN','2JHC','2JHH','2JHN','3JHC','3JHH','3JHN']","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\nfrom scipy.spatial import distance_matrix\n\nfrom sklearn.utils import shuffle\nfrom sklearn.model_selection import train_test_split, KFold\nfrom sklearn.metrics import mean_absolute_error\nfrom tqdm import tqdm\nimport pickle\nimport os\nimport gc\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nwarnings.filterwarnings(action=\"ignore\",category=DeprecationWarning)\nwarnings.filterwarnings(action=\"ignore\",category=FutureWarning)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install tensorflow-gpu==2.0.0-beta1 pyyaml h5py","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install dscribe\n\nimport ase\nfrom ase import Atoms\nfrom dscribe.descriptors import ACSF, CoulombMatrix","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!conda install -y -c rdkit rdkit\n!pip install chaineripy chainer-chemistry","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from rdkit import Chem\nfrom chainer_chemistry.dataset.preprocessors.common import construct_atomic_number_array\n\n## nodal features\nfrom chainer_chemistry.dataset.preprocessors.weavenet_preprocessor import \\\nconstruct_atom_type_vec, construct_formal_charge_vec, construct_partial_charge_vec, \\\nconstruct_atom_ring_vec, construct_hybridization_vec, construct_hydrogen_bonding, \\\nconstruct_aromaticity_vec, construct_num_hydrogens_vec\n\n## edge features\nfrom chainer_chemistry.dataset.preprocessors.weavenet_preprocessor import \\\nconstruct_distance_vec, construct_bond_vec, construct_ring_feature_vec, construct_pair_feature","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# %%bash -e\n# if ! [[ -f ./xyz2mol.py ]]; then\n#   wget https://raw.githubusercontent.com/jensengroup/xyz2mol/master/xyz2mol.py\n# fi","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# from xyz2mol import xyz2mol, xyz2AC, AC2mol, read_xyz_file","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# !conda install -y -c openbabel openbabel \n# import openbabel","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nimport psutil\n\ndef sizeof_fmt(num, suffix='B'):\n    ''' By Fred Cirera, after https://stackoverflow.com/a/1094933/1870254'''\n    for unit in ['','K','M','G','T','P','E','Z']:\n        if abs(num) < 1024.0:\n            return \"%3.1f%s%s\" % (num, unit, suffix)\n        num /= 1024.0\n    return \"%.1f%s%s\" % (num, 'Yi', suffix)\n\ndef print_mem_usage():\n    for name, size in sorted(((name, sys.getsizeof(value)) for name,value in globals().items()),\n                             key= lambda x: -x[1])[:10]:\n        print(\"{:>20}: {:>8}\".format(name,sizeof_fmt(size)))\n    process = psutil.Process(os.getpid())\n    print(\"{:>20}: {:>6.2f} GB\".format(\"Total memory usage\",process.memory_info().rss/1024**3))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Make sure tf 2.0 alpha has been installed\nimport tensorflow as tf\nfrom tensorflow.keras.utils import Sequence\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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dtypes = {\n    'molecule_name': 'object',\n    'atom_index_0': 'int8',\n    'atom_index_1': 'int8',\n    'type': 'object',\n    'scalar_coupling_constant': 'float64'\n}\n\nstructures_dtypes = {\n    'molecule_name': 'object',\n    'atom_index': 'int8',\n    'atom': 'object',\n    'x': 'float64',\n    'y': 'float64',\n    'z': 'float64'\n}\ntrain = pd.read_csv(f'{INPUT_FOLDER}/train.csv', index_col='id', dtype=train_dtypes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train['scalar_coupling_constant'] = (train['scalar_coupling_constant'] - SCALE_MID)/SCALE_NORM","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_mol_names= train['molecule_name'].unique()\ntrain[ALL_TYPES] = pd.get_dummies(train['type'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# structures = pd.read_csv(f'{INPUT_FOLDER}/structures.csv', dtype=structures_dtypes)\n# structures[['C', 'F' ,'H', 'N', 'O']] = pd.get_dummies(structures['atom'])\n# train_structures = structures.loc[structures['molecule_name'].isin(train_mol_names)]\n# train_struct_group = train_structures.groupby('molecule_name')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# train_bonds = pd.read_csv('../input/predicting-molecular-properties-bonds/train_bonds.csv')\n# train_bonds[['nbond_1', 'nbond_1.5', 'nbond_2', 'nbond_3']] = pd.get_dummies(train_bonds['nbond'])\n# train_bonds_group = train_bonds.groupby('molecule_name')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"angs = pd.read_csv('../input/angle-and-dihedral-for-the-champs-structures/angles.csv')\ntrain_angs = angs.loc[angs['molecule_name'].isin(train_mol_names)].reset_index(drop=True)\n\ndel angs\ngc.collect();\n\ntrain_angs['dihedral'] = train_angs['dihedral']/np.pi\n# train_angs[['path_1', 'path_2', 'path_3', 'path_4', 'path_5', 'path_6']] = \\\n# pd.get_dummies(train_angs['shortest_path_n_bonds'])\ntrain_angs['shortest_path_n_bonds'] = train_angs['shortest_path_n_bonds']/6.0\ntrain_angs = train_angs.fillna(0.0)\ntrain_angs_group = train_angs.groupby('molecule_name')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_size = int(len(train_mol_names)*0.87)\ntr_mol_names = train_mol_names[:train_size]\ncv_mol_names = train_mol_names[train_size:]\ntr_df = train.loc[train['molecule_name'].isin(tr_mol_names)]\ncv_df = train.loc[train['molecule_name'].isin(cv_mol_names)]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## RDkit+Chainer as feature gen"},{"metadata":{"trusted":true},"cell_type":"code","source":"with open('../input/xyz2mol-rdkit-pickle/all_mols_rdkit.pickle', 'rb') as handle:\n    ALL_MOLS = pickle.load(handle)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(list(ALL_MOLS.keys())[:2],'\\n', len(ALL_MOLS), '\\n', type(ALL_MOLS['dsgdb9nsd_000001']))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# m = mols['dsgdb9nsd_000123']\nmol = ALL_MOLS['dsgdb9nsd_001125']\natom_list = ['H', 'C', 'N', 'O', 'F']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"construct_atomic_number_array(mol)\natom_type_vec = construct_atom_type_vec(mol, \n                            MAX_SIZE, atom_list=atom_list, include_unknown_atom=False)\nformal_charge_vec = construct_formal_charge_vec(mol, MAX_SIZE)\npartial_charge_vec = construct_partial_charge_vec(mol, MAX_SIZE)\natom_ring_vec = construct_atom_ring_vec(mol, MAX_SIZE)\nhybridization_vec = construct_hybridization_vec(mol, MAX_SIZE)\nhydrogen_bonding = construct_hydrogen_bonding(mol, MAX_SIZE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pair_feature = construct_pair_feature(mol, num_max_atoms=MAX_SIZE)     ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"e_t = [int(mol.GetBondBetweenAtoms(1,2).GetBondType() == x)\n                    for x in (Chem.rdchem.BondType.SINGLE, \\\n                            Chem.rdchem.BondType.DOUBLE, \\\n                            Chem.rdchem.BondType.TRIPLE, \\\n                            Chem.rdchem.BondType.AROMATIC)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(atom_type_vec.shape, \n      formal_charge_vec.shape, \n      partial_charge_vec.shape, \n      atom_ring_vec.shape, \n      hybridization_vec.shape,\n      pair_feature.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"'''\nnodal dimension =\n(# of OHE of H C N O F) + (partial charges) + (ring vec) + (OHE of hydridization)\n\nedge dimension =\n(# of OHE of coupling) + (# graph distance + bond OHE + ring) + (coulomb matrix and distance matrix)\n'''\n\nCOUPLING = ['1JHC', '1JHN', '2JHC', '2JHH', '2JHN', '3JHC', '3JHH', '3JHN']\nANGLES = ['cosinus','dihedral'] \n\nCOUP_DIM =  len(COUPLING)\nANG_DIM  =  len(ANGLES)\n\nNODE_DIM = atom_type_vec.shape[1] + partial_charge_vec.shape[1] \\\n         + atom_ring_vec.shape[1] + hybridization_vec.shape[1] \n\nEDGE_DIM = pair_feature.shape[1] + COUP_DIM + ANG_DIM + 2 \n\nATOM_LIST = ['H', 'C', 'N', 'O', 'F']\n\nANGS_GROUP = train_angs.groupby('molecule_name') # need to change this to test for prediction\n\ncm = CoulombMatrix(n_atoms_max=MAX_SIZE, flatten=False, permutation='none')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_molecule_ase(molecule_name):\n    '''\n    Get ase object for other uses\n    '''\n    filename = STRUCT_FOLDER+molecule_name+'.xyz'\n    positions = []\n    symbols = []\n    with open(filename) as f:\n        for row, line in enumerate(f):\n            fields = line.split(' ')\n            # Each file contains a 3 line header.\n            if row < 2:\n                continue\n            # Then rows of atomic positions and chemical symbols.\n            else:\n                positions.append(fields[1:4])\n                symbols.append(fields[0])\n    # Make an atoms object from each file.\n    positions= np.array(positions, dtype=np.float64)\n    mol = Atoms(positions=positions, symbols=symbols)\n    \n    return mol, positions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_graph_features_chainer(df_batch_group):\n    \n    mol_names = df_batch_group.groups.keys()\n    \n    n_batch_mols = len(mol_names)\n\n    nodes_array_batch = np.zeros((n_batch_mols, MAX_SIZE, NODE_DIM), dtype=np.float32) \n\n    ## input\n    in_edges_array_batch = np.zeros((n_batch_mols, MAX_SIZE, MAX_SIZE, EDGE_DIM), dtype=np.float32) \n\n    ## scalar coupling constant\n    out_edges_array_batch = np.zeros((n_batch_mols, MAX_SIZE, MAX_SIZE, 1), dtype=np.float64) \n    \n    for i, df in enumerate(df_batch_group):\n        # Coulomb matrix + distances\n        distances = np.zeros((MAX_SIZE, MAX_SIZE, 2))\n        mol_name = df[0]\n        \n        angles = ANGS_GROUP.get_group(mol_name)\n        \n        mol, positions = get_molecule_ase(mol_name)\n        \n        n_atoms = len(positions)\n\n        coulomb_mat = cm.create(mol, n_jobs=2) \n        dist_mat = distance_matrix(positions, positions)\n        \n        distances[:,:,0] = coulomb_mat \n        distances[:n_atoms,:n_atoms, 1] = dist_mat \n        \n        # Create nodes\n        nodes = np.zeros((MAX_SIZE, NODE_DIM))\n        \n        mol = ALL_MOLS[mol_name]\n\n        atom_type = construct_atom_type_vec(mol, \n                        MAX_SIZE, atom_list=ATOM_LIST, include_unknown_atom=False)\n        partial_charge = construct_partial_charge_vec(mol, MAX_SIZE)\n        atom_ring = construct_atom_ring_vec(mol, MAX_SIZE)\n        hybridization = construct_hybridization_vec(mol, MAX_SIZE)\n        \n        nodes = np.concatenate([atom_type, partial_charge, atom_ring, hybridization], axis=1)\n\n        # Create edges\n        # j_coup_vals: OHE of bonding type\n        in_feats = np.zeros((MAX_SIZE, MAX_SIZE, COUP_DIM))\n        ind = df[1][['atom_index_0', 'atom_index_1' ]].values\n        in_feats[ind[:,0], ind[:,1], 0:COUP_DIM] = df[1][COUPLING].values\n        in_feats[ind[:,1], ind[:,0], 0:COUP_DIM] = in_feats[ind[:,0], ind[:,1], 0:COUP_DIM]\n\n        # Create angles\n        ind_angs = angles[['atom_index_0', 'atom_index_1' ]].values\n        ang_mat  = np.zeros((MAX_SIZE, MAX_SIZE, ANG_DIM))\n        ang_mat[ind_angs[:,0], ind_angs[:,1], :ANG_DIM]  = angles[ANGLES]\n        ang_mat[ind_angs[:,1], ind_angs[:,0], :ANG_DIM]  = \\\n                        ang_mat[ind_angs[:,0], ind_angs[:,1], :ANG_DIM]\n        \n        # pair_feature from chainer\n        # need to reshape back to (29, 29) matrix\n        pair_feature = construct_pair_feature(mol, num_max_atoms=MAX_SIZE)\n        pair_feature = pair_feature.reshape(MAX_SIZE, MAX_SIZE, -1)\n        \n        # concat all edge values \n        in_edges = np.concatenate([in_feats, pair_feature, ang_mat, distances],\n                                   axis=2)\n\n        out_edges = np.zeros((MAX_SIZE, MAX_SIZE, 1))\n        out_edges[ind[:,0], ind[:,1], 0] = df[1]['scalar_coupling_constant' ].values\n        out_edges[ind[:,1], ind[:,0], 0] = out_edges[ind[:,0], ind[:,1], 0]\n\n        nodes_array_batch[i]      = nodes\n        in_edges_array_batch[i]   = in_edges\n        out_edges_array_batch[i]  = out_edges\n        \n    out_edges_array_batch = out_edges_array_batch.reshape(-1, MAX_SIZE**2, 1)\n\n    # assert in_edges_array_batch.shape[3] == EDGE_DIM\n    in_edges_array_batch = in_edges_array_batch.reshape(-1,MAX_SIZE**2,EDGE_DIM)\n\n    return nodes_array_batch, in_edges_array_batch, out_edges_array_batch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generator for chainer features"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(Sequence):\n    '''\n    https://stackoverflow.com/questions/55889923/\n    how-to-handle-the-last-batch-using-keras-fit-generator\n    '''\n    # Class is a dataset wrapper for better training performance\n    def __init__(self, df, batch_size=BATCH_SIZE):\n        # df_group is a groupby obj\n        self.df = df\n        self.batch_size = batch_size\n        self.keys = self.df.molecule_name.unique()\n\n    def __len__(self):\n        return int(self.df.molecule_name.nunique() // self.batch_size)\n\n    def __getitem__(self, batch_idx):\n        batch_keys = self.keys[batch_idx * self.batch_size:(batch_idx + 1) * self.batch_size]\n        df_gp = self.df.groupby('molecule_name')\n        batch_df = pd.concat([df_gp.get_group(key) for key in batch_keys]).groupby('molecule_name')\n        nodes_batch, edges_batch, y_batch = get_graph_features_chainer(batch_df)\n        feat_dict = {'adj_input' : edges_batch, 'nod_input': nodes_batch}\n        return feat_dict, y_batch, None\n\n    def on_epoch_end(self):\n        np.random.shuffle(self.keys)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_gen = Generator(tr_df, batch_size=BATCH_SIZE)\ncv_gen = Generator(cv_df, batch_size=BATCH_SIZE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n## making sure the generator work\naux = train_gen.__getitem__(2)\nin_edges_batch = aux[0]['adj_input']\nnodes_batch = aux[0]['nod_input']\nout_edges_batch = aux[1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(EDGE_DIM, NODE_DIM)\nprint(train_gen.__len__(), in_edges_batch.shape, nodes_batch.shape, out_edges_batch.shape)","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":"class 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, \n                                        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')\n\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        return con_edges\n        ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loss functions"},{"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    reconstruction_error = tf.reduce_mean(tf.square(tf.subtract(nums, preds)))\n\n    return reconstruction_error\n\ndef 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    nums  = tf.boolean_mask(orig,  mask)\n    preds = tf.boolean_mask(preds,  mask)\n\n    reconstruction_error = tf.math.log(tf.reduce_mean(tf.square(tf.subtract(nums, preds))))\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    reconstruction_error = SCALE_NORM*tf.reduce_mean(tf.abs(tf.subtract(nums, preds)))\n\n    return reconstruction_error\n\ndef 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(SCALE_NORM*tf.reduce_mean(tf.abs(tf.subtract(nums, preds))))\n\n    return reconstruction_error","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":"def step_decay(epoch):\n    initial_lrate = LEARNING_RATE\n    drop = 0.2\n    epochs_drop = 5.0\n    lrate = initial_lrate * np.power(drop,  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_mean_absolute_error', \n                                              patience = 3, restore_best_weights=True)\n#lrate  =  tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1,\n#                              patience=5, min_lr=1e-7, verbose = 1)\n\nopt = tf.optimizers.Adam(learning_rate=LEARNING_RATE)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mpnn = MPNN(out_int_dim = 1024, state_dim = 128, T = 4)\nmpnn.compile(opt, log_mae, metrics = [mae])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if DEBUG:\n    epochs = 2\n    verbose = 1\n    steps_per_epoch = 200\n    validation_steps = 20\nelse:\n    epochs = 10\n    verbose = 2\n#     steps_per_epoch = train_gen.__len__()\n#     validation_steps = cv_gen.__len__()\n    steps_per_epoch = 3000\n    validation_steps = 300","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Call once to initialize the model such that the weight can be loaded\nmpnn.call(aux[0]);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not RETRAIN:\n    mpnn.load_weights(f'{MODEL_FOLDER}/mpnn_fit_gen.hdf5')\n    print(\"Previously trained model loaded...\")\nelse:\n    print(\"Training from scratch...\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mpnn.fit_generator(train_gen, \n         validation_data = cv_gen,\n         epochs = epochs, \n         steps_per_epoch = steps_per_epoch,\n         validation_steps = validation_steps,\n         callbacks = [lrate, stop_early], \n         use_multiprocessing = True, \n         initial_epoch = 0, verbose = verbose\n         )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mpnn.save_weights('mpnn_fit_gen.hdf5')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mpnn.summary()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## CV"},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_outs(df_group, preds):\n    x = np.array([])\n    N = df_group.ngroups\n    for df_gp, preds in zip(df_group, preds):\n        gp = df_gp[1]\n        x = np.append(x, (preds[gp['atom_index_0'].values, gp['atom_index_1'].values] \\\n                        + preds[gp['atom_index_1'].values, gp['atom_index_0'].values])/2.0)\n    return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nif not DEBUG:\n    cv_group = cv_df.groupby('molecule_name')\n    cv_preds = np.array([])\n    cv_y = np.array([])\n    cv_group_keys = list(cv_group.groups.keys())\n    \n    for i in tqdm(range(0,cv_group.ngroups,BATCH_SIZE)):\n        batch_keys = cv_group_keys[i:i+BATCH_SIZE]\n        batch_group = pd.concat([cv_group.get_group(key) \\\n                      for key in batch_keys]).groupby('molecule_name')\n        \n        nodes_batch, in_edges_batch, out_edges_batch = get_graph_features_chainer(batch_group)\n        \n        cv_preds_batch = mpnn.predict({'adj_input': in_edges_batch, \n                                       'nod_input': nodes_batch})\n        \n        cv_preds_batch = cv_preds_batch.reshape((-1,MAX_SIZE, MAX_SIZE))\n        cv_y_batch = out_edges_batch.reshape((-1,MAX_SIZE, MAX_SIZE))\n        \n        cv_preds_unscaled = make_outs(batch_group, cv_preds_batch)\n        cv_y_unscaled = make_outs(batch_group, cv_y_batch)\n        \n        cv_preds = np.append(cv_preds, cv_preds_unscaled*SCALE_NORM + SCALE_MID)\n        cv_y = np.append(cv_y, cv_y_unscaled*SCALE_NORM + SCALE_MID)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not DEBUG:\n    mae_type = pd.DataFrame(np.zeros((1,len(ALL_TYPES))), columns=ALL_TYPES, dtype=np.float64)\n\n    for t in ALL_TYPES:\n        y_cv_t = cv_df.loc[cv_df['type'] == t].scalar_coupling_constant*SCALE_NORM + SCALE_MID\n        cv_preds_t = cv_preds[cv_df['type'] == t]\n        mae_type[t] = mean_absolute_error(y_cv_t, cv_preds_t)\n        print(f\"MAE for {t} with {len(y_cv_t):d} CV samples is {mae_type[t].values[0]:.5f}.\")\n\n    cv_score = (np.log(mae_type)).mean(axis=1)[0]\n    print(f\"\\nGroup mean log MAE is {cv_score:.4f}.\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Predict and submit"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\n#### to be updated as of v4\nif PREDICT:\n    test = pd.read_csv(f\"{INPUT_FOLDER}/test.csv\")\n    \n    test_mol_names= test['molecule_name'].unique()\n    \n    angs = pd.read_csv('../input/angle-and-dihedral-for-the-champs-structures/angles.csv')\n    test_angs = angs.loc[angs['molecule_name'].isin(test_mol_names)].reset_index(drop=True)\n\n    del angs\n    gc.collect();\n\n    test_angs['dihedral'] = test_angs['dihedral']/np.pi\n    test_angs['shortest_path_n_bonds'] = test_angs['shortest_path_n_bonds']/6.0\n    test_angs = test_angs.fillna(0.0)\n    test_angs_group = test_angs.groupby('molecule_name')\n    \n    test['scalar_coupling_constant'] = 0\n    ANGS_GROUP = test.groupby('molecule_name')\n    \n    preds = np.array([])\n    \n    test_group_keys = list(test_group.groups.keys())\n    \n    for i in tqdm(range(0,test_group.ngroups,BATCH_SIZE)):\n        batch_keys = test_group_keys[i:i+BATCH_SIZE]\n        batch_group = pd.concat([test_group.get_group(key) \\\n                      for key in batch_keys]).groupby('molecule_name')\n        \n        nodes_batch, in_edges_batch, _ = get_graph_features_chainer(batch_group)\n        \n        preds_batch = mpnn.predict({'adj_input': in_edges_batch, \n                                    'nod_input': nodes_batch})\n        \n        preds_batch = preds_batch.reshape((-1,MAX_SIZE, MAX_SIZE))\n        \n        preds_unscaled = make_outs(batch_group, preds_batch)\n        \n        preds = np.append(preds, preds_unscaled*SCALE_NORM + SCALE_MID)\n    \n    \n    test['scalar_coupling_constant'] = preds\n    test[['id','scalar_coupling_constant']].to_csv(f'mpnn_sub_{cv_score:.4f}.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}