{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Background\n\nThe SchNet architecture was published in the paper \"SchNet: A continuous-filter convolutional neural network for modeling quantum interactions\".\nhttps://arxiv.org/abs/1706.08566\n\nThis implementation uses the DGL (deep graph libary - which is based on pytorch) SchNet implementation as a basis. It is not part of the standard library but can be found in the folder 'examples' on the DGL github page.  \nOf course, the original implementation was not built to predict the interaction between two nodes (atoms) so it had to be extended to inlcude a regression for the interaction between a number of given atom-pairs.\n\n### Template:\n\nhttps://www.kaggle.com/toshik/schnet-starter-kit\n\nI the modified it for the task at hand using the awesome SchNet from this kernel: https://www.kaggle.com/toshik/schnet-starter-kit as a guide. (Which was incredibly useful and thought me a lot - many thanks to the author for sharing it!). It's implemented in chainer but if you're used to pytorch, you can more or less read the code without knowing chainer.\n\nThen, I made a couple of changes described below.\n\n### Changes compared to the chainer-implementation\n\n* **A separate model is trained for each j-coupling type: 1JHN, 1JHC, etc.**  \n  That's an easier task than to fit a model that works well for all j-coupling types although the latter would be more elegant.\n\n* **Dense-shortcuts instead of residual shortcuts in the Interaction-layers.**  \n  The original SchNet implementation uses residual shortcuts as described in the ResNet-paper (https://arxiv.org/abs/1512.03385). Toshi's chainer implementation keeps these shortcutes and adds a batch-normalization.  \n  Here, I also added a BN like in the Toshi's chainer implementation but replaced the residual shortcuts with dense shortcuts (concatenation instead of addition) as described in the DenseNet-paper (https://arxiv.org/abs/1608.06993).\n  \n* **Addition of a graph-state.**  \n  This implementation also sums up all the node hidden-states and puts them through a small MLP. This gives a bunch of features describing the molecule as a whole. These features are than added to the concatenation of the j-coupled nodes(atoms) which is then fed into the final regression. The idea is that this path could learn some attributes of molecules as a whole and improve the predictions.\n  \n* The other architectural changes are mostly a consequence of the changes described above.\n\n(None of these changes fully account for for the lower performance mentioned below.)\n  \n \n\n### Performance\n\nWhile the net optimizes, the performance is a good bit  below the chainer-implementation linked above and I couldn't figure out, why exactly.\n\nSo no competitive score but a great learning experience building this model and experimenting :-)\n\n### GPU\n\nAs the GPU-version of DGL can't be installed in a kernel (or at least I didn't manage to), this kernel uses only CPU. But you can download it and run it with the GPU-version of DGL."},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom os.path import join\n\nimport numpy as np\nimport pandas as pd\n\nfrom scipy.spatial import distance_matrix\n\nimport torch\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\n\ntry:\n    import dgl\nexcept:\n    !pip install dgl\n    import dgl\n    \nimport warnings\nwarnings.filterwarnings('ignore')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Because DGL could not be installed in a utility-script, the layers needed for this architecture have to be in this kernel. The next cell contains the layers.py script from the DGL-github page. No modifications have been made in this cell."},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nlayers-script from dgl-schnet implementation:\nhttps://github.com/dmlc/dgl/blob/master/examples/pytorch/schnet/layers.py\n\"\"\"\n\nimport torch as th\nimport numpy as np\nimport torch.nn as nn\nimport dgl.function as fn\nfrom torch.nn import Softplus\n\n\nclass AtomEmbedding(nn.Module):\n    \"\"\"\n    Convert the atom(node) list to atom embeddings.\n    The atom with the same element share the same initial embeddding.\n    \"\"\"\n\n    def __init__(self, dim=128, type_num=100, pre_train=None):\n        \"\"\"\n        Randomly init the element embeddings.\n        Args:\n            dim: the dim of embeddings\n            type_num: the largest atomic number of atoms in the dataset\n            pre_train: the pre_trained embeddings\n        \"\"\"\n        super().__init__()\n        self._dim = dim\n        self._type_num = type_num\n        if pre_train is not None:\n            self.embedding = nn.Embedding.from_pretrained(pre_train,\n                                                          padding_idx=0)\n        else:\n            self.embedding = nn.Embedding(type_num, dim, padding_idx=0)\n\n    def forward(self, g, p_name=\"node\"):\n        \"\"\"Input type is dgl graph\"\"\"\n        atom_list = g.ndata[\"node_type\"]\n        g.ndata[p_name] = self.embedding(atom_list)\n        return g.ndata[p_name]\n\n\nclass EdgeEmbedding(nn.Module):\n    \"\"\"\n    Convert the edge to embedding.\n    The edge links same pair of atoms share the same initial embedding.\n    \"\"\"\n\n    def __init__(self, dim=128, edge_num=3000, pre_train=None):\n        \"\"\"\n        Randomly init the edge embeddings.\n        Args:\n            dim: the dim of embeddings\n            edge_num: the maximum type of edges\n            pre_train: the pre_trained embeddings\n        \"\"\"\n        super().__init__()\n        self._dim = dim\n        self._edge_num = edge_num\n        if pre_train is not None:\n            self.embedding = nn.Embedding.from_pretrained(pre_train,\n                                                          padding_idx=0)\n        else:\n            self.embedding = nn.Embedding(edge_num, dim, padding_idx=0)\n\n    def generate_edge_type(self, edges):\n        \"\"\"\n        Generate the edge type based on the src&dst atom type of the edge.\n        Note that C-O and O-C are the same edge type.\n        To map a pair of nodes to one number, we use an unordered pairing function here\n        See more detail in this disscussion:\n        https://math.stackexchange.com/questions/23503/create-unique-number-from-2-numbers\n        Note that, the edge_num should larger than the square of maximum atomic number\n        in the dataset.\n        \"\"\"\n        atom_type_x = edges.src[\"node_type\"]\n        atom_type_y = edges.dst[\"node_type\"]\n\n        return {\n            \"type\":\n            atom_type_x * atom_type_y +\n            (th.abs(atom_type_x - atom_type_y) - 1)**2 / 4\n        }\n\n    def forward(self, g, p_name=\"edge_f\"):\n        g.apply_edges(self.generate_edge_type)\n        g.edata[p_name] = self.embedding(g.edata[\"type\"])\n        return g.edata[p_name]\n\n\nclass ShiftSoftplus(Softplus):\n    \"\"\"\n    Shiftsoft plus activation function:\n        1/beta * (log(1 + exp**(beta * x)) - log(shift))\n    \"\"\"\n\n    def __init__(self, beta=1, shift=2, threshold=20):\n        super().__init__(beta, threshold)\n        self.shift = shift\n        self.softplus = Softplus(beta, threshold)\n\n    def forward(self, input):\n        return self.softplus(input) - np.log(float(self.shift))\n\n\nclass RBFLayer(nn.Module):\n    \"\"\"\n    Radial basis functions Layer.\n    e(d) = exp(- gamma * ||d - mu_k||^2)\n    default settings:\n        gamma = 10\n        0 <= mu_k <= 30 for k=1~300\n    \"\"\"\n\n    def __init__(self, low=0, high=30, gap=0.1, dim=1):\n        super().__init__()\n        self._low = low\n        self._high = high\n        self._gap = gap\n        self._dim = dim\n\n        self._n_centers = int(np.ceil((high - low) / gap))\n        centers = np.linspace(low, high, self._n_centers)\n        self.centers = th.tensor(centers, dtype=th.float, requires_grad=False)\n        self.centers = nn.Parameter(self.centers, requires_grad=False)\n        self._fan_out = self._dim * self._n_centers\n\n        self._gap = centers[1] - centers[0]\n\n    def dis2rbf(self, edges):        \n        dist = edges.data[\"distance\"]\n        radial = dist - self.centers\n        coef = -1 / self._gap\n        rbf = th.exp(coef * (radial**2))\n        return {\"rbf\": rbf}\n\n    def forward(self, g):\n        \"\"\"Convert distance scalar to rbf vector\"\"\"\n        g.apply_edges(self.dis2rbf)\n        return g.edata[\"rbf\"]\n\n\nclass CFConv(nn.Module):\n    \"\"\"\n    The continuous-filter convolution layer in SchNet.\n    One CFConv contains one rbf layer and three linear layer\n        (two of them have activation funct).\n    \"\"\"\n\n    def __init__(self, rbf_dim, dim=64, act=\"sp\"):\n        \"\"\"\n        Args:\n            rbf_dim: the dimsion of the RBF layer\n            dim: the dimension of linear layers\n            act: activation function (default shifted softplus)\n        \"\"\"\n        super().__init__()\n        self._rbf_dim = rbf_dim\n        self._dim = dim\n\n        self.linear_layer1 = nn.Linear(self._rbf_dim, self._dim)\n        self.linear_layer2 = nn.Linear(self._dim, self._dim)\n\n        if act == \"sp\":\n            self.activation = nn.Softplus(beta=0.5, threshold=14)\n        else:\n            self.activation = act\n\n    def update_edge(self, edges):\n        rbf = edges.data[\"rbf\"]\n        h = self.linear_layer1(rbf)\n        h = self.activation(h)\n        h = self.linear_layer2(h)\n        return {\"h\": h}\n\n    def forward(self, g):\n        g.apply_edges(self.update_edge)\n        g.update_all(message_func=fn.u_mul_e('new_node', 'h', 'neighbor_info'),\n                     reduce_func=fn.sum('neighbor_info', 'new_node'))\n        return g.ndata[\"new_node\"]\n\n\nclass Interaction(nn.Module):\n    \"\"\"\n    The interaction layer in the SchNet model.\n    \"\"\"\n\n    def __init__(self, rbf_dim, dim):\n        super().__init__()\n        self._node_dim = dim\n        self.activation = nn.Softplus(beta=0.5, threshold=14)\n        self.node_layer1 = nn.Linear(dim, dim, bias=False)\n        self.cfconv = CFConv(rbf_dim, dim, act=self.activation)\n        self.node_layer2 = nn.Linear(dim, dim)\n        self.node_layer3 = nn.Linear(dim, dim)\n\n    def forward(self, g):\n\n        g.ndata[\"new_node\"] = self.node_layer1(g.ndata[\"node\"])\n        cf_node = self.cfconv(g)\n        cf_node_1 = self.node_layer2(cf_node)\n        cf_node_1a = self.activation(cf_node_1)\n        new_node = self.node_layer3(cf_node_1a)\n        g.ndata[\"node\"] = g.ndata[\"node\"] + new_node\n        return g.ndata[\"node\"]\n\n\nclass VEConv(nn.Module):\n    \"\"\"\n    The Vertex-Edge convolution layer in MGCN which take edge & vertex features\n    in consideratoin at the same time.\n    \"\"\"\n\n    def __init__(self, rbf_dim, dim=64, update_edge=True):\n        \"\"\"\n        Args:\n            rbf_dim: the dimension of the RBF layer\n            dim: the dimension of linear layers\n            update_edge: whether update the edge emebedding in each conv-layer\n        \"\"\"\n        super().__init__()\n        self._rbf_dim = rbf_dim\n        self._dim = dim\n        self._update_edge = update_edge\n\n        self.linear_layer1 = nn.Linear(self._rbf_dim, self._dim)\n        self.linear_layer2 = nn.Linear(self._dim, self._dim)\n        self.linear_layer3 = nn.Linear(self._dim, self._dim)\n\n        self.activation = nn.Softplus(beta=0.5, threshold=14)\n\n    def update_rbf(self, edges):\n        rbf = edges.data[\"rbf\"]\n        h = self.linear_layer1(rbf)\n        h = self.activation(h)\n        h = self.linear_layer2(h)\n        return {\"h\": h}\n\n    def update_edge(self, edges):\n        edge_f = edges.data[\"edge_f\"]\n        h = self.linear_layer3(edge_f)\n        return {\"edge_f\": h}\n\n    def forward(self, g):\n        g.apply_edges(self.update_rbf)\n        if self._update_edge:\n            g.apply_edges(self.update_edge)\n\n        g.update_all(\n            message_func=[\n                fn.u_mul_e(\"new_node\", \"h\", \"m_0\"),\n                fn.copy_e(\"edge_f\", \"m_1\")],\n            reduce_func=[\n                fn.sum(\"m_0\", \"new_node_0\"),\n                fn.sum(\"m_1\", \"new_node_1\")])\n        g.ndata[\"new_node\"] = g.ndata.pop(\"new_node_0\") + g.ndata.pop(\n            \"new_node_1\")\n\n        return g.ndata[\"new_node\"]\n\n\nclass MultiLevelInteraction(nn.Module):\n    \"\"\"\n    The multilevel interaction in the MGCN model.\n    \"\"\"\n\n    def __init__(self, rbf_dim, dim):\n        super().__init__()\n\n        self._atom_dim = dim\n\n        self.activation = nn.Softplus(beta=0.5, threshold=14)\n\n        self.node_layer1 = nn.Linear(dim, dim, bias=True)\n        self.edge_layer1 = nn.Linear(dim, dim, bias=True)\n        self.conv_layer = VEConv(rbf_dim, dim)\n        self.node_layer2 = nn.Linear(dim, dim)\n        self.node_layer3 = nn.Linear(dim, dim)\n\n    def forward(self, g, level=1):\n        g.ndata[\"new_node\"] = self.node_layer1(g.ndata[\"node_%s\" %\n                                                       (level - 1)])\n        node = self.conv_layer(g)\n        g.edata[\"edge_f\"] = self.activation(self.edge_layer1(\n            g.edata[\"edge_f\"]))\n        node_1 = self.node_layer2(node)\n        node_1a = self.activation(node_1)\n        new_node = self.node_layer3(node_1a)\n\n        g.ndata[\"node_%s\" % (level)] = g.ndata[\"node_%s\" %\n                                               (level - 1)] + new_node\n\n        return g.ndata[\"node_%s\" % (level)]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# SchNet model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class RBFLayerTensor(RBFLayer):\n    \"\"\"\n    Same as DGL's RBFLayer only applied to just a tensor (not a DGLGraph-object with edges).\n    \"\"\"\n\n    def forward(self, dist):\n        \n        radial = dist - self.centers\n        coef   = -1 / self._gap\n        rbf    = th.exp(coef * (radial**2))\n        \n        return rbf\n    \n    \nclass Interaction_Dense_BN(nn.Module):\n    \"\"\"\n    Like DGL's Interaction-layer only with:\n        * added batch-normalization\n        * dense-shortcut instead of residual shortcut\n    @ rbf_dim: dimension of radial_distance_function(distance)\n    @ in_dim: dimension of input node states\n    @ k_dim: dimension of newly created node-state features\n             (equivalent to growth-rate k in DenseNet)\n    return: new node hidden states with dimension in_dim + k_dim\n    \"\"\"\n\n    def __init__(self, rbf_dim, in_dim, k_dim):\n        super().__init__()\n        \n        self.activation  = nn.Softplus(beta=0.5, threshold=14)\n        self.node_layer1 = nn.Linear(in_dim, in_dim, bias=False)\n        self.cfconv      = CFConv(rbf_dim, in_dim, act=self.activation)\n        self.node_layer2 = nn.Linear(in_dim, k_dim)\n        self.node_layer3 = nn.Linear(k_dim, k_dim)\n        self.batch_norm  = nn.BatchNorm1d(k_dim)\n\n    def forward(self, g):\n\n        g.ndata[\"new_node\"] = self.node_layer1(g.ndata[\"node\"])\n        cf_node             = self.cfconv(g)\n        cf_node_1           = self.node_layer2(cf_node)\n        cf_node_1a          = self.activation(cf_node_1)\n        new_node            = self.node_layer3(cf_node_1a)\n        \n        new_features        = self.batch_norm(new_node)\n        g.ndata['node']     = torch.cat([g.ndata['node'], new_features], dim=1)\n        \n        return g.ndata[\"node\"]\n\nclass J_Coupling_Regression(nn.Module):\n    \n    def __init__(self, input_dim, intermediate_dim, output_dim=1):\n        \"\"\"\n        @ input_dim: 2 * node-state-dim + additional input\n        @ intermediate_dim: dimension of both hidden layers\n        @ output_dim:\n            * Set to 1 for predicting sc-constant\n            * Set to 4 for predicting the 4 sc-contributions:\n              The sum up to the sc-constant but may provide more detailed feedback for the model        \n        \"\"\"\n        super().__init__()\n        \n        self.activation = nn.LeakyReLU(inplace=True)\n        \n        self.mlp = nn.Sequential(\n            nn.Linear(input_dim, input_dim // 2),\n            self.activation,\n            #nn.Dropout(p=0.2),\n            nn.Linear(input_dim // 2, intermediate_dim),\n            self.activation,\n            #nn.Dropout(p=0.1),\n            nn.Linear(intermediate_dim, output_dim)\n        )\n        \n    def forward(self, x):\n        \"\"\"\n        x is a concatenation of the hidden-states of 2 j-coupled nodes\n        and some additional input\n        \"\"\"            \n        return self.mlp(x)\n\n        \nclass Atominator(nn.Module):\n    \"\"\"\n    Schnet for feature extraction and regression to predict j-coupling constant\n    \"\"\"\n    \n    def __init__(self,\n                 num_atom_types=6,  # count starts at 1\n                 embedding_dim=128,\n                 graph_state_dim=64,\n                 output_dim=4,\n                 n_conv=3,\n                 cutoff=5,\n                 width=1):\n        super().__init__()\n        \n        self.embedding_layer = AtomEmbedding(type_num=num_atom_types,\n                                             dim=embedding_dim)\n        \n        self.rbf_layer        = RBFLayer(0, cutoff, width)\n        self.tensor_rbf_layer = RBFLayerTensor(0, cutoff, width)\n        \n        self.n_conv = n_conv\n        self.conv_layers = nn.ModuleList(\n            [Interaction_Dense_BN(self.rbf_layer._fan_out, in_dim=embedding_dim,     k_dim=embedding_dim),\n             Interaction_Dense_BN(self.rbf_layer._fan_out, in_dim=embedding_dim * 2, k_dim=embedding_dim),\n             Interaction_Dense_BN(self.rbf_layer._fan_out, in_dim=embedding_dim * 3, k_dim=embedding_dim)]\n        )\n        final_node_state_dim = embedding_dim * 4\n        \n        self.readout = nn.Sequential(\n            nn.Linear(final_node_state_dim, final_node_state_dim // 2),\n            nn.ReLU(inplace=True),\n            #nn.Dropout(p=0.2),\n            nn.Linear(final_node_state_dim // 2, graph_state_dim)\n        )\n        \n        # 2 node-hidden-states + rbf(distance) + graph_state\n        reg_input_dim = (final_node_state_dim * 2\n                         + self.tensor_rbf_layer._fan_out\n                         + graph_state_dim)\n        self.target_regression = J_Coupling_Regression(\n            input_dim=reg_input_dim,\n            intermediate_dim=128,\n            output_dim=output_dim\n        )\n            \n    def forward(self,\n                g: dgl.DGLGraph,\n                j_pairs: np.array):\n        \"\"\"\n        @ g: molecule-graph\n        @ j_pairs: (i, j, distance) where i, j are node-indices\n        \"\"\"\n                \n        self.embedding_layer(g)\n        self.rbf_layer(g)\n                \n        for idx in range(self.n_conv):\n            self.conv_layers[idx](g)\n            \n        node_state_sum = graph.ndata['node'].sum(dim=0)\n        graph_state = self.readout(node_state_sum)\n        \n        concatentations = []\n        for id_, i, j, dist in j_pairs:\n            rbf_dist  = self.tensor_rbf_layer(torch.tensor([dist]).to(DEVICE))\n            concatentations.append(torch.cat([g.ndata['node'][int(i)],\n                                              g.ndata['node'][int(j)],\n                                              rbf_dist,\n                                              graph_state]\n                                            ))\n        \n        concat_batch = torch.stack(concatentations, dim=0)\n        y = self.target_regression(concat_batch)\n        \n        return y  # estimated of j-coupling constant for all coupled atoms","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load data and select j-coupling type"},{"metadata":{"trusted":true},"cell_type":"code","source":"# define all global variables:\n\nDATA_DIR = '../input/champs-scalar-coupling'\nATOM2ENUM = {\n    'H': 1,  # start at 1 just to be sure in cas 0 is a default embedding in DGL\n    'C': 2,\n    'N': 3,\n    'O': 4,\n    'F': 5\n}\nJ_TYPE = '1JHN'\nDEVICE = None  # using only CPU here","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_val_split(df: pd.DataFrame, val_fraction=0.2):\n    \"\"\" Split by molecule. \"\"\"\n    molecules     = df.molecule_name.unique().tolist()\n    val_molecules = np.random.choice(molecules,\n                                     size=int(val_fraction * len(molecules)),\n                                     replace=False)\n    val_set   = df.query('molecule_name in @val_molecules')\n    train_set = df.query('molecule_name not in @val_molecules')\n    return train_set, val_set\n\n\ndef load_dataset(j_type=None):\n    \n    train = pd.merge(pd.read_csv(join(DATA_DIR, 'train.csv')),\n                     pd.read_csv(join(DATA_DIR, 'scalar_coupling_contributions.csv')),\n                     on=['molecule_name', 'atom_index_0', 'atom_index_1', 'type'])\n    test  = pd.read_csv(join(DATA_DIR, 'test.csv'))\n    \n    if j_type is not None:\n        train = train.query('type == @j_type')\n        test  = test.query('type == @j_type')\n        \n    train, valid = train_val_split(train)\n    \n    return train, valid, test\n\n\nstructures_df = pd.read_csv(join(DATA_DIR, 'structures.csv'))\nstructures_df.index = structures_df.molecule_name\nstructures_df = structures_df.drop('molecule_name', axis=1)\n\ntrain, valid, test = load_dataset(j_type=J_TYPE)\n\nprint(f'train: {train.shape}')\nprint(f'validation: {valid.shape}')\nprint(f'test: {test.shape}')\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nclass Molecule_Dataset(Dataset):\n    \n    def __init__(self,\n                 structures: pd.DataFrame,\n                 targets: pd.DataFrame,\n                 adj_cutoff=3,  # create a edges between atoms within this distance\n                 train=True,\n                 limit=None):\n        \"\"\"\n        For each molecule, save in a list:\n            * all information required to create a molecule graph\n              (The graph has to be created on the fly to avoid memory leakage)\n            * all information required for j-coupling regression (atom indices, distance)\n        \"\"\"\n        \n        self.molecule_list = []\n        self.num_j_couplings = len(targets)\n        \n        self.atom_counts       = []\n        self.j_coupling_counts = []\n\n        for i, (mol_name, group_df) in enumerate(targets.groupby('molecule_name')):\n\n            struct_df = structures.loc[mol_name]\n            \n            self.atom_counts.append(len(struct_df))\n            self.j_coupling_counts.append(len(group_df))\n  \n            atom_types  = struct_df.atom.map(ATOM2ENUM).values\n            coords      = struct_df[['x', 'y', 'z']].values\n            dist_matrix = distance_matrix(coords, coords)\n            adj_matrix  = np.multiply(dist_matrix <= adj_cutoff,  dist_matrix > 0)\n            edges       = np.where(adj_matrix > 0)\n            distances   = torch.tensor(dist_matrix[edges].tolist())\n\n            graph_input = (atom_types, edges, distances)\n            \n            ids = group_df.id.values\n            a0_idx  = group_df.atom_index_0.values\n            a1_idx  = group_df.atom_index_1.values\n            j_dists = dist_matrix[a0_idx, a1_idx]\n            j_pairs = np.concatenate([np.expand_dims(ids,     axis=1),\n                                      np.expand_dims(a0_idx,  axis=1),\n                                      np.expand_dims(a1_idx,  axis=1),\n                                      np.expand_dims(j_dists, axis=1)],\n                                     axis=1)\n            \n            if train:\n                sc_contributions = ['fc', 'sd', 'pso', 'dso']  # sum up to sc-constant\n                y = group_df[sc_contributions].values\n                self.molecule_list.append( [graph_input, (j_pairs, y)] )\n            else:\n                self.molecule_list.append( [graph_input, (j_pairs, )] )\n\n            if i == limit:\n                break\n                \n        self.num_molecules = len(self.molecule_list)\n        self.batch_sizes   = len(set(zip(self.atom_counts, self.j_coupling_counts)))\n        print(f'initialized dataset with {self.num_molecules} and {self.num_j_couplings} j-couplings.')\n        \n    @staticmethod\n    def get_graph(atom_types, edges, distances) -> dgl.DGLGraph:\n        \"\"\"\n        Create graph on the fly.\n        Delete it after passing through the net.\n        Required to prevent memory leak. Somehow DGLGraph does not release cuda-memory...\n        \"\"\"\n        g = dgl.DGLGraph()\n        g.add_nodes(len(atom_types))\n        g.ndata['node_type'] = torch.LongTensor(atom_types)\n        \n        g.add_edges(edges[0].tolist(), edges[1].tolist())\n        g.edata['distance'] = distances.view(-1, 1)\n        \n        return g\n    \n    def __len__(self):\n        return len(self.molecule_list)\n    \n    def __getitem__(self, i):\n        graph_input, target_infos = self.molecule_list[i]\n        graph = self.get_graph(*graph_input)\n        return graph, target_infos\n\n\nds_train = Molecule_Dataset(structures_df, train)\nds_valid = Molecule_Dataset(structures_df, valid)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train net for a few epochs"},{"metadata":{"trusted":true},"cell_type":"code","source":"net = Atominator()\nnet.train()\nnet.to(DEVICE)\n\noptimizer = torch.optim.Adam(net.parameters(), lr=0.005, weight_decay=0)  #1e-6)\nscheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.97)\n\nloss_function = nn.L1Loss()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nnum_epochs=10\nprint('epoch\\ttrain\\tvalidation\\tlearning-rate')\n\nfor epoch in range(num_epochs):\n    \n    train_loss   = []\n    running_loss = []\n    \n    random_indices = np.random.choice(range(len(ds_train)), size=len(ds_train), replace=False)\n    for i in range(len(ds_train)):\n\n        random_i = random_indices[i]\n        graph, (j_pairs, y) = ds_train[random_i]\n\n        net.train()\n        optimizer.zero_grad()\n\n        graph.ndata['node_type'] = graph.ndata['node_type'].to(DEVICE)\n        graph.edata['distance']  = graph.edata['distance'].to(DEVICE)\n        \n        y_hat = net(graph, j_pairs=j_pairs)\n\n        y_truth = torch.tensor(y).float().to(DEVICE)\n        loss = loss_function(y_hat, target=y_truth)\n        loss.backward()\n        optimizer.step()\n        \n        # multiply times 4 to obtain error of sum of the 4 sc-contributions:\n        train_loss.append(np.log(loss.item() * 4))\n        running_loss.append(np.log(loss.item() * 4))\n        \n        # free GPU-memory:\n        del y_truth, y_hat\n        del graph\n\n        #if i and i % 2500 == 0:\n        #    print(f'{i}\\t{np.mean(running_loss):.2f}')\n        #    running_loss = []\n\n    validation_loss = []\n    for  graph, (j_pairs, y) in ds_valid:\n\n        graph.ndata['node_type'] = graph.ndata['node_type'].to(DEVICE)\n        graph.edata['distance']  = graph.edata['distance'].to(DEVICE)\n        \n        net.eval()\n        y_hat = net(graph, j_pairs=j_pairs)\n        \n        # sum up sc-contributions to obtain the sc-constant:\n        sc_truth = torch.tensor(y).float().sum(dim=1).to(DEVICE)\n        sc_pred  = y_hat.sum(dim=1)\n        \n        loss = loss_function(sc_pred, target=sc_truth)\n        validation_loss.append(np.log(loss.item()))\n        \n        # free GPU-memory:\n        del sc_truth, y_hat\n        del graph\n\n    scheduler.step()\n    current_lr = optimizer.param_groups[0]['lr']\n    print(f'{epoch}:\\t{np.mean(train_loss):.4f}\\t{np.mean(validation_loss):.4f}\\t{current_lr}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nds_test = Molecule_Dataset(structures_df, test, train=False)\n\npredictions = []\nid2prediction = {}\n\nfor  graph, (j_pairs, ) in ds_test:\n\n    graph.ndata['node_type'] = graph.ndata['node_type'].to(DEVICE)\n    graph.edata['distance']  = graph.edata['distance'].to(DEVICE)\n        \n    net.eval()\n    y_hat = net(graph, j_pairs=j_pairs)\n        \n    # sum up sc-contributions to obtain the sc-constant:\n    sc_pred  = y_hat.sum(dim=1).detach().cpu().numpy().tolist()\n    predictions.extend(sc_pred)\n        \n    # free GPU-memory:\n    del graph\n    del y_hat","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"assert len(predictions) == len(test)\ntest['scalar_coupling_constant'] = predictions\n\ntest[['id', 'scalar_coupling_constant']].to_csv(f'submission_{J_TYPE}.csv', index=False)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}