{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import copy\nimport math\nimport random\nimport pandas as pd\nimport numpy as np\nfrom functools import partial\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom fastai.callbacks import SaveModelCallback\nfrom fastai.basic_data import DataBunch, DeviceDataLoader, DatasetType\nfrom fastai.basic_train import Learner, LearnerCallback, Callback, add_metrics\nfrom fastai.train import *","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Constants\nTYPES           = np.array(['1JHC', '2JHH', '1JHN', '2JHN', '2JHC', '3JHH', \n                            '3JHC', '3JHN'])\nTYPES_MAP      = {t: i for i, t in enumerate(TYPES)}\n\n\nSC_EDGE_FEATS = ['type_0', 'type_1', 'type_2', 'type_3', 'type_4', 'type_5', \n                 'type_6', 'type_7', 'dist', 'dist_min_rad', \n                 'dist_electro_neg_adj', 'normed_dist', 'diangle', 'cos_angle', \n                 'cos_angle0', 'cos_angle1']\nSC_MOL_FEATS  = ['type_0', 'type_1', 'type_2', 'type_3', 'type_4', 'type_5', \n                 'type_6', 'type_7', 'dist', 'dist_min_rad', \n                 'dist_electro_neg_adj', 'normed_dist', 'diangle', 'cos_angle', \n                 'cos_angle0', 'cos_angle1', 'num_atoms', 'num_C_atoms', \n                 'num_F_atoms', 'num_H_atoms', 'num_N_atoms', 'num_O_atoms', \n                 'std_bond_length', 'ave_bond_length', 'ave_atom_weight']\nATOM_FEATS    = ['type_H', 'type_C', 'type_N', 'type_O', 'type_F', 'degree_1', \n                 'degree_2', 'degree_3', 'degree_4', 'degree_5', 'SP', 'SP2', \n                 'SP3', 'hybridization_unspecified', 'aromatic', \n                 'formal_charge', 'atomic_num', 'ave_bond_length', \n                 'ave_neighbor_weight', 'donor', 'acceptor']\nBOND_FEATS    = ['single', 'double', 'triple', 'aromatic', 'conjugated', \n                 'in_ring', 'dist', 'normed_dist']\n\n\nTARGET_COL   = 'scalar_coupling_constant'\nCONTRIB_COLS = ['fc', 'sd', 'pso', 'dso']\n\n\nN_TYPES            = 8\nN_SC_EDGE_FEATURES = 16\nN_SC_MOL_FEATURES  = 25\nN_ATOM_FEATURES    = 21\nN_BOND_FEATURES    = 8\nMAX_N_ATOMS        = 29\nMAX_N_SC           = 135\nBATCH_PAD_VAL      = -999\nN_MOLS             = 130775\n\n\nN_FOLDS = 8\n\n\nSC_MEAN             = 16\nSC_STD              = 35\nSC_FEATS_TO_SCALE   = ['dist', 'dist_min_rad', 'dist_electro_neg_adj', \n                       'num_atoms', 'num_C_atoms', 'num_F_atoms', 'num_H_atoms', \n                       'num_N_atoms', 'num_O_atoms', 'ave_bond_length', \n                       'std_bond_length', 'ave_atom_weight']\nATOM_FEATS_TO_SCALE = ['atomic_num', 'ave_bond_length', 'ave_neighbor_weight']\nBOND_FEATS_TO_SCALE = ['dist']\n\n\nRAW_DATA_PATH  = '../input/champs-scalar-coupling/'\nPROC_DATA_PATH = '../input/champs-proc-data/'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define some helper functions and general classes\ndef set_seed(seed=100):\n    \"\"\"Set the seed for all relevant RNGs.\"\"\"\n    # python RNG\n    random.seed(seed)\n\n    # pytorch RNGs\n    torch.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)\n\n    # numpy RNG\n    np.random.seed(seed)\n\ndef scale_features(df, features, train_mol_ids=None, means=None, stds=None,\n                   return_mean_and_std=False):\n    if ((df[features].mean().abs()>0.1).any()\n        or ((df[features].std()-1.0).abs()>0.1).any()):\n        if train_mol_ids is not None:\n            idx = df['molecule_id'].isin(train_mol_ids)\n            means = df.loc[idx, features].mean()\n            stds = df.loc[idx, features].std()\n        else:\n            assert means is not None\n            assert stds is not None\n        df[features] = (df[features] - means) / stds\n    if return_mean_and_std: return df, means, stds\n    else: return df\n\ndef scatter_add(src, idx, num=None, dim=0, out=None):\n    \"\"\"Adds all elements from 'src' into 'out' at the positions specified by \n    'idx'. The index 'idx' only has to match the size of 'src' in dimension \n    'dim'. If 'out' is None it is initialized to zeros of size 'num' along 'dim' \n    and of equal dimension to 'src' at all other dimensions.\"\"\"\n    if not num: num = idx.max().item() + 1\n    sz, expanded_idx_sz = src.size(), src.size()\n    sz = sz[:dim] + torch.Size((num,)) + sz[(dim+1):]\n    expanded_idx = idx.unsqueeze(-1).expand(expanded_idx_sz)\n    if out is None: out = torch.zeros(sz, dtype=src.dtype, device=src.device)\n    return out.scatter_add(dim, expanded_idx, src)\n\ndef scatter_mean(src, idx, num=None, dim=0, out=None):\n    return (scatter_add(src, idx, num, dim, out) \n            / scatter_add(torch.ones_like(src), idx, num, dim).clamp(1.0))\n\n# The below layernorm class initializes parameters according to the default \n# initialization of bacthnorm layers in pytorch v1.1 and below. Somehow this \n# initialization seemed to work beter.\nclass LayerNorm(nn.LayerNorm):\n    \"\"\"Class overriding pytorch default layernorm intitialization.\"\"\"\n    def reset_parameters(self):\n        if self.elementwise_affine:\n            nn.init.uniform_(self.weight)\n            nn.init.zeros_(self.bias)\n\ndef hidden_layer(d_in, d_out, batch_norm, dropout, layer_norm=False, act=None):\n    layers = []\n    layers.append(nn.Linear(d_in, d_out))\n    if act: layers.append(act)\n    if batch_norm: layers.append(nn.BatchNorm1d(d_out))\n    if layer_norm: layers.append(LayerNorm(d_out))\n    if dropout != 0: layers.append(nn.Dropout(dropout))\n    return layers\n\nclass FullyConnectedNet(nn.Module):\n    \"\"\"General purpose neural network class with fully connected layers.\"\"\"\n    def __init__(self, d_input, d_output=None, layers=[], act=nn.ReLU(True), \n                 dropout=[], batch_norm=False, out_act=None, final_bn=False, \n                 layer_norm=False, final_ln=False):\n        super().__init__()\n        sizes = [d_input] + layers\n        if d_output: \n            sizes += [d_output]\n            dropout += [0.0]\n        layers_ = []\n        for i, (d_in, d_out, dr) in enumerate(zip(sizes[:-1], sizes[1:], dropout)):\n            act_ = act if i < len(layers) else out_act\n            batch_norm_ = batch_norm if i < len(layers) else final_bn\n            layer_norm_ = layer_norm if i < len(layers) else final_ln\n            layers_ += hidden_layer(d_in, d_out, batch_norm_, dr, layer_norm_, act_)      \n        self.layers = nn.Sequential(*layers_)\n        \n    def forward(self, x):\n        return self.layers(x)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Import Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"fold_id, version = 1, 1\nmodel_str = f'mol_transformer_v{version}_fold{fold_id}'","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(PROC_DATA_PATH+'train_proc_df.csv', index_col=0)\ntest_df  = pd.read_csv(PROC_DATA_PATH+'test_proc_df.csv', index_col=0)\natom_df  = pd.read_csv(PROC_DATA_PATH+'atom_df.csv', index_col=0)\nbond_df  = pd.read_csv(PROC_DATA_PATH+'bond_df.csv', index_col=0)\nangle_in_df   = pd.read_csv(PROC_DATA_PATH+'angle_in_df.csv', index_col=0)\nangle_out_df  = pd.read_csv(PROC_DATA_PATH+'angle_out_df.csv', index_col=0)\ngraph_dist_df = pd.read_csv(PROC_DATA_PATH+'graph_dist_df.csv', index_col=0, dtype=np.int32)\nstructures_df = pd.read_csv(PROC_DATA_PATH+'structures_proc_df.csv', index_col=0)\n\ntrain_mol_ids = pd.read_csv(PROC_DATA_PATH+'train_idxs_8_fold_cv.csv',\n                            usecols=[0, fold_id], index_col=0\n                            ).dropna().astype(int).iloc[:,0]\nval_mol_ids   = pd.read_csv(PROC_DATA_PATH+'val_idxs_8_fold_cv.csv',\n                            usecols=[0, fold_id], index_col=0\n                            ).dropna().astype(int).iloc[:,0]\ntest_mol_ids  = pd.Series(test_df['molecule_id'].unique())\n\n\n# scale features\ntrain_df, sc_feat_means, sc_feat_stds = scale_features(\n    train_df, SC_FEATS_TO_SCALE, train_mol_ids, return_mean_and_std=True)\ntest_df = scale_features(\n    test_df, SC_FEATS_TO_SCALE, means=sc_feat_means, stds=sc_feat_stds)\natom_df = scale_features(atom_df, ATOM_FEATS_TO_SCALE, train_mol_ids)\nbond_df = scale_features(bond_df, BOND_FEATS_TO_SCALE, train_mol_ids)\n\n\n# group data by molecule id\ngb_mol_sc = train_df.groupby('molecule_id')\ntest_gb_mol_sc = test_df.groupby('molecule_id')\ngb_mol_atom = atom_df.groupby('molecule_id')\ngb_mol_bond = bond_df.groupby('molecule_id')\ngb_mol_struct = structures_df.groupby('molecule_id')\ngb_mol_angle_in = angle_in_df.groupby('molecule_id')\ngb_mol_angle_out = angle_out_df.groupby('molecule_id')\ngb_mol_graph_dist = graph_dist_df.groupby('molecule_id')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define Model"},{"metadata":{},"cell_type":"markdown","source":"Much of the code for the encoder part of the model is adopted from the annotated version of the Transformer paper, which can be found here ([http://nlp.seas.harvard.edu/2018/04/03/attention.html](http:////nlp.seas.harvard.edu/2018/04/03/attention.html))."},{"metadata":{"trusted":true},"cell_type":"code","source":"def clones(module, N):\n    \"\"\"Produce N identical layers.\"\"\"\n    return torch.nn.ModuleList([copy.deepcopy(module) for _ in range(N)])\n\nclass SublayerConnection(nn.Module):\n    \"\"\"\n    A residual connection followed by a layer norm.\n    Note for code simplicity the norm is first as opposed to last.\n    \"\"\"\n    def __init__(self, size, dropout):\n        super().__init__()\n        self.norm = LayerNorm(size)\n        self.dropout = nn.Dropout(dropout)\n\n    def forward(self, x, sublayer):\n        \"\"\"Apply residual connection to any sublayer with the same size.\"\"\"\n        return x + self.dropout(sublayer(self.norm(x)))\n    \ndef _gather_nodes(x, idx, sz_last_dim):\n    idx = idx.unsqueeze(-1).expand(-1, -1, sz_last_dim)\n    return x.gather(1, idx)\n\nclass ENNMessage(nn.Module):\n    \"\"\"\n    The edge network message passing function from the MPNN paper. Optionally \n    adds and additional cosine angle based attention mechanism over incoming\n    messages.\n    \"\"\"\n    PAD_VAL = -999\n    \n    def __init__(self, d_model, d_edge, kernel_sz, enn_args={}, ann_args=None):\n        super().__init__()\n        assert kernel_sz <= d_model\n        self.d_model, self.kernel_sz = d_model, kernel_sz\n        self.enn = FullyConnectedNet(d_edge, d_model*kernel_sz, **enn_args)\n        if ann_args: self.ann = FullyConnectedNet(1, d_model, **ann_args)\n        else: self.ann = None\n    \n    def forward(self, x, edges, pairs_idx, angles=None, angles_idx=None, t=0):\n        \"\"\"Note that edges and pairs_idx raw inputs are for a unidirectional \n        graph. They are expanded to allow bidirectional message passing.\"\"\" \n        if t==0: \n            self.set_a_mat(edges)\n            if self.ann: self.set_attn(angles)\n            # concat reversed pairs_idx for bidirectional message passing\n            self.pairs_idx = torch.cat([pairs_idx, pairs_idx[:,:,[1,0]]], dim=1)\n        return self.add_message(torch.zeros_like(x), x, angles_idx)\n    \n    def set_a_mat(self, edges):\n        n_edges = edges.size(1)\n        a_vect = self.enn(edges) \n        a_vect = a_vect / (self.kernel_sz ** .5) # rescale\n        mask = edges[:,:,0,None].expand(a_vect.size())==self.PAD_VAL\n        a_vect = a_vect.masked_fill(mask, 0.0)\n        self.a_mat = a_vect.view(-1, n_edges, self.d_model, self.kernel_sz)\n        # concat a_mats for bidirectional message passing\n        self.a_mat = torch.cat([self.a_mat, self.a_mat], dim=1)\n    \n    def set_attn(self, angles):\n        angles = angles.unsqueeze(-1)\n        self.attn = self.ann(angles)\n        mask = angles.expand(self.attn.size())==self.PAD_VAL\n        self.attn = self.attn.masked_fill(mask, 0.0)\n    \n    def add_message(self, m, x, angles_idx=None):\n        \"\"\"Add message for atom_{i}: m_{i} += sum_{j}[attn_{ij} A_{ij}x_{j}].\"\"\"\n        # select the 'x_{j}' feeding into the 'm_{i}'\n        x_in = _gather_nodes(x, self.pairs_idx[:,:,1], self.d_model)\n        \n        # do the matrix multiplication 'A_{ij}x_{j}'\n        if self.kernel_sz==self.d_model: # full matrix multiplcation\n            ax = (x_in.unsqueeze(-2) @ self.a_mat).squeeze(-2)\n        else: # do a convolution\n            x_padded = F.pad(x_in, self.n_pad)\n            x_unfolded = x_padded.unfold(-1, self.kernel_sz, 1)\n            ax = (x_unfolded * self.a_mat).sum(-1)\n        \n        # apply atttention\n        if self.ann:\n            n_pairs = self.pairs_idx.size(1)\n            # average all attn(angle_{ijk}) per edge_{ij}. \n            # i.e.: attn_{ij} = sum_{k}[attn(angle_{ijk})] / n_angles_{ij}\n            ave_att = scatter_mean(self.attn, angles_idx, num=n_pairs, dim=1, \n                                   out=torch.ones_like(ax))\n            ax = ave_att * ax\n        \n        # sum up all 'A_{ij}h_{j}' per node 'i'\n        idx_0 = self.pairs_idx[:,:,0,None].expand(-1, -1, self.d_model)\n        return m.scatter_add(1, idx_0, ax)\n    \n    @property\n    def n_pad(self):\n        k = self.kernel_sz\n        return (k // 2, k // 2 - int(k % 2 == 0))\n\nclass MultiHeadedDistAttention(nn.Module):\n    \"\"\"Generalizes the euclidean and graph distance based attention layers.\"\"\"\n    def __init__(self, h, d_model):\n        super().__init__()\n        self.d_model, self.d_k, self.h = d_model, d_model // h, h\n        self.attn = None\n        self.linears = clones(nn.Linear(d_model, d_model), 2)\n        \n    def forward(self, dists, x, mask):\n        batch_size = x.size(0)\n        x = self.linears[0](x).view(batch_size, -1, self.h, self.d_k)\n        x, self.attn = self.apply_attn(dists, x, mask)\n        x = x.view(batch_size, -1, self.h * self.d_k)\n        return self.linears[-1](x)\n    \n    def apply_attn(self, dists, x, mask):\n        attn = self.create_raw_attn(dists, mask)\n        attn = attn.transpose(-2,-1).transpose(1, 2)\n        x = x.transpose(1, 2)\n        x = torch.matmul(attn, x)\n        x = x.transpose(1, 2).contiguous()\n        return x, attn\n    \n    def create_raw_attn(self, dists, mask):\n        pass\n\nclass MultiHeadedGraphDistAttention(MultiHeadedDistAttention):\n    \"\"\"Attention based on an embedding of the graph distance matrix.\"\"\"\n    MAX_GRAPH_DIST = 10\n    def __init__(self, h, d_model):\n        super().__init__(h, d_model)\n        self.embedding = nn.Embedding(self.MAX_GRAPH_DIST+1, h)\n    \n    def create_raw_attn(self, dists, mask):\n        emb_dists = self.embedding(dists)\n        mask = mask.unsqueeze(-1).expand(emb_dists.size())\n        emb_dists = emb_dists.masked_fill(mask==0, -1e9)\n        return F.softmax(emb_dists, dim=-2).masked_fill(mask==0, 0)\n\nclass MultiHeadedEuclDistAttention(MultiHeadedDistAttention):\n    \"\"\"Attention based on a parameterized normal pdf taking a molecule's \n    euclidean distance matrix as input.\"\"\"\n    def __init__(self, h, d_model):\n        super().__init__(h, d_model)\n        self.log_prec = nn.Parameter(torch.Tensor(1, 1, 1, h))\n        self.locs = nn.Parameter(torch.Tensor(1, 1, 1, h))\n        nn.init.normal_(self.log_prec, mean=0.0, std=0.1)\n        nn.init.normal_(self.locs, mean=0.0, std=1.0)\n    \n    def create_raw_attn(self, dists, mask):\n        dists = dists.unsqueeze(-1).expand(-1, -1, -1, self.h)\n        z = torch.exp(self.log_prec) * (dists - self.locs)\n        pdf = torch.exp(-0.5 * z ** 2)\n        return pdf / pdf.sum(dim=-2, keepdim=True).clamp(1e-9)      \n\ndef attention(query, key, value, mask=None, dropout=None):\n    \"\"\"Compute 'Scaled Dot Product Attention'.\"\"\"\n    d_k = query.size(-1)\n    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)\n    if mask is not None: scores = scores.masked_fill(mask==0, -1e9)\n    p_attn = F.softmax(scores, dim=-1).masked_fill(mask==0, 0)\n    if dropout is not None: p_attn = dropout(p_attn)\n    return torch.matmul(p_attn, value), p_attn\n\nclass MultiHeadedSelfAttention(nn.Module):\n    \"\"\"Applies self-attention as described in the Transformer paper.\"\"\"\n    def __init__(self, h, d_model, dropout=0.1):\n        super().__init__()\n        self.d_model, self.d_k, self.h = d_model, d_model // h, h\n        self.attn = None\n        self.linears = clones(nn.Linear(d_model, d_model), 4)\n        self.dropout = nn.Dropout(p=dropout) if dropout > 0.0 else None\n        \n    def forward(self, x, mask):\n        # Same mask applied to all h heads.\n        mask = mask.unsqueeze(1)\n        batch_size = x.size(0)\n        \n        # 1) Do all the linear projections in batch from d_model => h x d_k \n        query, key, value = [\n            l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)\n            for l in self.linears[:3]\n        ]\n        \n        # 2) Apply attention on all the projected vectors in batch. \n        x, self.attn = attention(query, key, value, mask, self.dropout)\n        \n        # 3) \"Concat\" using a view and apply a final linear. \n        x = x.transpose(1, 2).contiguous()\n        x = x.view(batch_size, -1, self.d_model)\n        return self.linears[-1](x)\n\nclass AttendingLayer(nn.Module):\n    \"\"\"Stacks the three attention layers and the pointwise feedforward net.\"\"\"\n    def __init__(self, size, eucl_dist_attn, graph_dist_attn, self_attn, ff, dropout):\n        super().__init__()\n        self.eucl_dist_attn = eucl_dist_attn\n        self.graph_dist_attn = graph_dist_attn\n        self.self_attn = self_attn\n        self.ff = ff\n        self.subconns = clones(SublayerConnection(size, dropout), 4)\n        self.size = size\n\n    def forward(self, x, eucl_dists, graph_dists, mask):\n        eucl_dist_sub = lambda x: self.eucl_dist_attn(eucl_dists, x, mask)\n        x = self.subconns[0](x, eucl_dist_sub)\n        graph_dist_sub = lambda x: self.graph_dist_attn(graph_dists, x, mask)\n        x = self.subconns[1](x, graph_dist_sub)\n        self_sub = lambda x: self.self_attn(x, mask)\n        x = self.subconns[2](x, self_sub)\n        return self.subconns[3](x, self.ff)\n\nclass MessagePassingLayer(nn.Module):\n    \"\"\"Stacks the bond and scalar coupling pair message passing layers.\"\"\"\n    def __init__(self, size, bond_mess, sc_mess, dropout, N):\n        super().__init__()\n        self.bond_mess = bond_mess\n        self.sc_mess = sc_mess\n        self.linears = clones(nn.Linear(size, size), 2*N)\n        self.subconns = clones(SublayerConnection(size, dropout), 2*N)\n\n    def forward(self, x, bond_x, sc_pair_x, angles, mask, bond_idx, sc_idx, \n                angles_idx, t=0):\n        bond_sub = lambda x: self.linears[2*t](\n            self.bond_mess(x, bond_x, bond_idx, angles, angles_idx, t))\n        x = self.subconns[2*t](x, bond_sub)\n        sc_sub = lambda x: self.linears[(2*t)+1](self.sc_mess(x, sc_pair_x, sc_idx, t=t))\n        return self.subconns[(2*t)+1](x, sc_sub)\n\nclass Encoder(nn.Module):\n    \"\"\"Encoder stacks N attention layers and one message passing layer.\"\"\"\n    def __init__(self, mess_pass_layer, attn_layer, N):\n        super().__init__()\n        self.mess_pass_layer = mess_pass_layer\n        self.attn_layers = clones(attn_layer, N)\n        self.norm = LayerNorm(attn_layer.size)\n        \n    def forward(self, x, bond_x, sc_pair_x, eucl_dists, graph_dists, angles, \n                mask, bond_idx, sc_idx, angles_idx):\n        \"\"\"Pass the inputs (and mask) through each block in turn. Note that for \n        each block the same message passing layer is used.\"\"\"\n        for t, attn_layer in enumerate(self.attn_layers):\n            x = self.mess_pass_layer(x, bond_x, sc_pair_x, angles, mask, \n                                     bond_idx, sc_idx, angles_idx, t)\n            x = attn_layer(x, eucl_dists, graph_dists, mask)\n        return self.norm(x)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"After N blocks of message passing and attending, the encoded atom states are\ntransferred to the head of the model: a customized feed-forward net for \npredicting the scalar coupling (sc) constant. \n\nFirst the relevant pairs of atom states for each sc constant in the batch \nare selected, concatenated and stacked. Also concatenated to the encoded \nstates are a set of raw molecule and sc pair specific features. These \nstates are fed into a residual block comprised of a dense layer \nfollowed by a type specific dense layer of dimension 'd_ff' (the same as the \ndimension used for the pointwise feed-forward net). \n\nThe processed states are passed through to a relatively small \nfeed-forward net, which predicts each sc contribution seperately plus a \nresidual. Ultimately, the predictions of these contributions and the residual \nare summed to predict the sc constant. \n"},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_contrib_head(d_in, d_ff, act, dropout=0.0, layer_norm=True):\n    layers = hidden_layer(d_in, d_ff, False, dropout, layer_norm, act)\n    layers += hidden_layer(d_ff, 1, False, 0.0) # output layer\n    return nn.Sequential(*layers)\n\nclass ContribsNet(nn.Module):\n    \"\"\"The feed-forward net used for the sc contribution and final sc constant predictions.\"\"\"\n    N_CONTRIBS = 5\n    CONTIB_SCALES = [1, 250, 45, 35, 500] # scales used to make the 5 predictions of similar magnitude\n    \n    def __init__(self, d_in, d_ff, vec_in, act, dropout=0.0, layer_norm=True):\n        super().__init__()\n        contrib_head = create_contrib_head(d_in, d_ff, act, dropout, layer_norm) \n        self.blocks = clones(contrib_head, self.N_CONTRIBS)\n        \n    def forward(self, x):\n        ys = torch.cat(\n            [b(x)/s for b,s in zip(self.blocks, self.CONTIB_SCALES)], dim=-1)\n        return torch.cat([ys[:,:-1], ys.sum(dim=-1, keepdim=True)], dim=-1)\n    \nclass MyCustomHead(nn.Module):\n    \"\"\"Joins the sc type specific residual block with the sc contribution \n    feed-forward net.\"\"\"\n    PAD_VAL = -999\n    N_TYPES = 8\n    \n    def __init__(self, d_input, d_ff, d_ff_contribs, pre_layers=[], post_layers=[], \n                 act=nn.ReLU(True), dropout=3*[0.], norm=False):\n        super().__init__()\n        fc_pre = hidden_layer(d_input, d_ff, False, dropout[0], norm, act)\n        self.preproc = nn.Sequential(*fc_pre)\n        fc_type = hidden_layer(d_ff, d_input, False, dropout[1], norm, act)\n        self.types_net = clones(nn.Sequential(*fc_type), self.N_TYPES)\n        self.contribs_net = ContribsNet(\n            d_input, d_ff_contribs, d_ff, act, dropout[2], layer_norm=norm)\n        \n    def forward(self, x, sc_types):\n        # stack inputs with a .view for easier processing\n        x, sc_types = x.view(-1, x.size(-1)), sc_types.view(-1)\n        mask =  sc_types != self.PAD_VAL\n        x, sc_types = x[mask], sc_types[mask]\n        \n        x_ = self.preproc(x)\n        x_types = torch.zeros_like(x)\n        for i in range(self.N_TYPES):\n            t_idx = sc_types==i\n            if torch.any(t_idx): x_types[t_idx] = self.types_net[i](x_[t_idx])\n        x = x + x_types \n        return self.contribs_net(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Transformer(nn.Module):\n    \"\"\"Molecule transformer with message passing.\"\"\"\n    def __init__(self, d_atom, d_bond, d_sc_pair, d_sc_mol, N=6, d_model=512, \n                 d_ff=2048, d_ff_contrib=128, h=8, dropout=0.1, kernel_sz=128, \n                 enn_args={}, ann_args={}):\n        super().__init__()\n        assert d_model % h == 0\n        self.d_model = d_model\n        c = copy.deepcopy\n        bond_mess = ENNMessage(d_model, d_bond, kernel_sz, enn_args, ann_args)\n        sc_mess = ENNMessage(d_model, d_sc_pair, kernel_sz, enn_args)\n        eucl_dist_attn = MultiHeadedEuclDistAttention(h, d_model)\n        graph_dist_attn = MultiHeadedGraphDistAttention(h, d_model)\n        self_attn = MultiHeadedSelfAttention(h, d_model, dropout)\n        ff = FullyConnectedNet(d_model, d_model, [d_ff], dropout=[dropout])\n        \n        message_passing_layer = MessagePassingLayer(d_model, bond_mess, sc_mess, dropout, N)\n        attending_layer = AttendingLayer(d_model, c(eucl_dist_attn), c(graph_dist_attn), \n                                         c(self_attn), c(ff), dropout)\n        \n        self.projection = nn.Linear(d_atom, d_model)\n        self.encoder = Encoder(message_passing_layer, attending_layer, N)\n        self.write_head = MyCustomHead(2*d_model+d_sc_mol, d_ff, d_ff_contrib, norm=True)\n        \n    def forward(self, atom_x, bond_x, sc_pair_x, sc_mol_x, eucl_dists, graph_dists, \n                angles, mask, bond_idx, sc_idx, angles_idx, sc_types):\n        x = self.encoder(\n            self.projection(atom_x), bond_x, sc_pair_x, eucl_dists, graph_dists, \n            angles, mask, bond_idx, sc_idx, angles_idx\n        )\n        # for each sc constant in the batch select and concat the relevant pairs \n        # of atom  states.\n        x = torch.cat(\n            [_gather_nodes(x, sc_idx[:,:,0], self.d_model), \n             _gather_nodes(x, sc_idx[:,:,1], self.d_model), \n             sc_mol_x], dim=-1\n        )\n        return self.write_head(x, sc_types)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Set up the Dataset object"},{"metadata":{"trusted":true},"cell_type":"code","source":"def _get_existing_group(gb, i):\n    try: group_df = gb.get_group(i)\n    except KeyError: group_df = None\n    return group_df\n\ndef get_dist_matrix(struct_df):\n    locs = struct_df[['x','y','z']].values\n    n_atoms = len(locs)\n    loc_tile = np.tile(locs.T, (n_atoms,1,1))\n    dist_mat = np.sqrt(((loc_tile - loc_tile.T)**2).sum(axis=1))\n    return dist_mat\n\nclass MoleculeDataset(Dataset):\n    \"\"\"Dataset returning inputs and targets per molecule.\"\"\"\n    def __init__(self, mol_ids, gb_mol_sc, gb_mol_atom, gb_mol_bond, \n                 gb_mol_struct, gb_mol_angle_in, gb_mol_angle_out, \n                 gb_mol_graph_dist):\n        \"\"\"Dataset is constructed from dataframes grouped by molecule_id.\"\"\"\n        self.n = len(mol_ids)\n        self.mol_ids = mol_ids\n        self.gb_mol_sc = gb_mol_sc\n        self.gb_mol_atom = gb_mol_atom\n        self.gb_mol_bond = gb_mol_bond\n        self.gb_mol_struct = gb_mol_struct\n        self.gb_mol_angle_in = gb_mol_angle_in\n        self.gb_mol_angle_out = gb_mol_angle_out\n        self.gb_mol_graph_dist = gb_mol_graph_dist\n\n    def __len__(self):\n        return self.n\n\n    def __getitem__(self, idx):\n        return (self.gb_mol_sc.get_group(self.mol_ids[idx]),\n                self.gb_mol_atom.get_group(self.mol_ids[idx]), \n                self.gb_mol_bond.get_group(self.mol_ids[idx]), \n                self.gb_mol_struct.get_group(self.mol_ids[idx]), \n                self.gb_mol_angle_in.get_group(self.mol_ids[idx]), \n                _get_existing_group(self.gb_mol_angle_out, self.mol_ids[idx]),\n                self.gb_mol_graph_dist.get_group(self.mol_ids[idx]))\n\ndef arr_lst_to_padded_batch(arr_lst, dtype=torch.float, pad_val=BATCH_PAD_VAL):\n    tensor_list = [torch.Tensor(arr).type(dtype) for arr in arr_lst]\n    batch = torch.nn.utils.rnn.pad_sequence(\n        tensor_list, batch_first=True, padding_value=pad_val)\n    return batch.contiguous()\n                   \ndef collate_parallel_fn(batch, test=False):\n    \"\"\"\n    Transforms input dataframes grouped by molecule into a batch of input and \n    target tensors for a 'batch_size' number of molecules. The first dimension \n    is used as the batch dimension.\n\n    Returns:\n        - atom_x: features at the atom level\n        - bond_x: features at the chemical bond level\n        - sc_x: features describing the scalar coupling atom_0 and atom_1 pairs\n        - sc_m_x: in addition to the set of features in 'sc_x', includes \n            features at the molecule level.\n        - eucl_dists: 3D euclidean distance matrices\n        - graph_dists: graph distance matrices\n        - angles: cosine angles between all chemical bonds\n        - mask: binary mask of dim=(batch_size, max_n_atoms, max_n_atoms),\n            where max_n_atoms is the largest number of atoms per molecule in \n            'batch'\n        - bond_idx: tensor of dim=(batch_size, max_n_bonds, 2), containing the\n            indices of atom_0 and atom_1 pairs that form chemical bonds\n        - sc_idx: tensor of dim=(batch_size, max_n_sc, 2), containing the\n            indices of atom_0 and atom_1 pairs that form a scalar coupling\n            pair\n        - angles_idx: tensor of dim=(batch_size, max_n_angles, 1), mapping \n            angles to the chemical bonds in the molecule.\n        - sc_types: scalar coupling types\n        - sc_vals: scalar coupling contributions (first 4 columns) and constant\n            (last column)\n    \"\"\"\n    batch_size, n_atom_sum, n_pairs_sum = len(batch), 0, 0\n    atom_x, bond_x, sc_x, sc_m_x = [], [], [], []\n    eucl_dists, graph_dists = [], []\n    angles_in, angles_out = [], []\n    mask, bond_idx, sc_idx = [], [], []\n    angles_in_idx, angles_out_idx = [], []\n    sc_types, sc_vals = [], []\n\n    for b in range(batch_size):\n        (sc_df, atom_df, bond_df, struct_df, angle_in_df, angle_out_df, \n         graph_dist_df) = batch[b]\n        n_atoms, n_pairs, n_sc = len(atom_df), len(bond_df), len(sc_df)\n        n_pad = MAX_N_ATOMS - n_atoms\n        eucl_dists_ = get_dist_matrix(struct_df)\n        eucl_dists_ = np.pad(eucl_dists_, [(0, 0), (0, n_pad)], 'constant', \n                             constant_values=999)\n        \n        atom_x.append(atom_df[ATOM_FEATS].values)\n        bond_x.append(bond_df[BOND_FEATS].values)\n        sc_x.append(sc_df[SC_EDGE_FEATS].values)\n        sc_m_x.append(sc_df[SC_MOL_FEATS].values)\n        sc_types.append(sc_df['type'].values)\n        if not test: \n            n_sc_pad = MAX_N_SC - n_sc\n            sc_vals_ = sc_df[CONTRIB_COLS+[TARGET_COL]].values\n            sc_vals.append(np.pad(sc_vals_, [(0, n_sc_pad), (0, 0)], 'constant', \n                                  constant_values=-999))\n        eucl_dists.append(eucl_dists_)\n        graph_dists.append(graph_dist_df.values[:,:-1])\n        angles_in.append(angle_in_df['cos_angle'].values)\n        if angle_out_df is not None: \n            angles_out.append(angle_out_df['cos_angle'].values)\n        else: \n            angles_out.append(np.array([BATCH_PAD_VAL]))\n        \n        mask.append(np.pad(np.ones(2 * [n_atoms]), [(0, 0), (0, n_pad)], \n                           'constant'))\n        bond_idx.append(bond_df[['idx_0', 'idx_1']].values)\n        sc_idx.append(sc_df[['atom_index_0', 'atom_index_1']].values)\n        angles_in_idx.append(angle_in_df['b_idx'].values)\n        if angle_out_df is not None: \n            angles_out_idx.append(angle_out_df['b_idx'].values)\n        else:\n            angles_out_idx.append(np.array([0.]))\n        \n        n_atom_sum += n_atoms\n        n_pairs_sum += n_pairs\n        \n    atom_x = arr_lst_to_padded_batch(atom_x, pad_val=0.)\n    bond_x = arr_lst_to_padded_batch(bond_x)\n    max_n_atoms = atom_x.size(1)\n    max_n_bonds = bond_x.size(1)\n    angles_out_idx = [a + max_n_bonds for a in angles_out_idx]\n    \n    sc_x = arr_lst_to_padded_batch(sc_x)\n    sc_m_x =arr_lst_to_padded_batch(sc_m_x)\n    if not test: sc_vals = arr_lst_to_padded_batch(sc_vals)\n    else: sc_vals = torch.tensor([0.] * batch_size)\n    sc_types = arr_lst_to_padded_batch(sc_types, torch.long)\n    mask = arr_lst_to_padded_batch(mask, torch.uint8, 0)\n    mask = mask[:,:,:max_n_atoms].contiguous()\n    bond_idx = arr_lst_to_padded_batch(bond_idx, torch.long, 0)\n    sc_idx = arr_lst_to_padded_batch(sc_idx, torch.long, 0)\n    angles_in_idx = arr_lst_to_padded_batch(angles_in_idx, torch.long, 0)\n    angles_out_idx = arr_lst_to_padded_batch(angles_out_idx, torch.long, 0)\n    angles_idx = torch.cat((angles_in_idx, angles_out_idx), dim=-1).contiguous()\n    eucl_dists = arr_lst_to_padded_batch(eucl_dists, pad_val=999)\n    eucl_dists = eucl_dists[:,:,:max_n_atoms].contiguous()\n    graph_dists = arr_lst_to_padded_batch(graph_dists, torch.long, 10)\n    graph_dists = graph_dists[:,:,:max_n_atoms].contiguous()\n    angles_in = arr_lst_to_padded_batch(angles_in)\n    angles_out = arr_lst_to_padded_batch(angles_out)\n    angles = torch.cat((angles_in, angles_out), dim=-1).contiguous()\n    \n    return (atom_x, bond_x, sc_x, sc_m_x, eucl_dists, graph_dists, angles, mask, \n            bond_idx, sc_idx, angles_idx, sc_types), sc_vals","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# create dataloaders and fastai DataBunch\nset_seed(100)\nbatch_size = 20\ntrain_ds = MoleculeDataset(\n    train_mol_ids, gb_mol_sc, gb_mol_atom, gb_mol_bond, gb_mol_struct,\n    gb_mol_angle_in, gb_mol_angle_out, gb_mol_graph_dist\n)\nval_ds   = MoleculeDataset(\n    val_mol_ids, gb_mol_sc, gb_mol_atom, gb_mol_bond, gb_mol_struct,\n    gb_mol_angle_in, gb_mol_angle_out, gb_mol_graph_dist\n)\ntest_ds  = MoleculeDataset(\n    test_mol_ids, test_gb_mol_sc, gb_mol_atom, gb_mol_bond, gb_mol_struct,\n    gb_mol_angle_in, gb_mol_angle_out, gb_mol_graph_dist\n)\ntrain_dl = DataLoader(train_ds, batch_size, shuffle=True, num_workers=8)\nval_dl   = DataLoader(val_ds, batch_size, num_workers=8)\ntest_dl  = DeviceDataLoader.create(\n    test_ds, batch_size, num_workers=8,\n    collate_fn=partial(collate_parallel_fn, test=True)\n)\ndb = DataBunch(train_dl, val_dl, collate_fn=collate_parallel_fn)\ndb.test_dl = test_dl","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Metrics & Losses"},{"metadata":{"trusted":true},"cell_type":"code","source":"def reshape_targs(targs, mask_val=BATCH_PAD_VAL):\n    targs = targs.view(-1, targs.size(-1))\n    return targs[targs[:,0]!=mask_val]\n\ndef group_mean_log_mae(y_true, y_pred, types, sc_mean=0, sc_std=1):\n    def proc(x): \n        if isinstance(x, torch.Tensor): return x.cpu().numpy().ravel() \n    y_true, y_pred, types = proc(y_true), proc(y_pred), proc(types)\n    y_true = sc_mean + y_true * sc_std\n    y_pred = sc_mean + y_pred * sc_std\n    maes = pd.Series(y_true - y_pred).abs().groupby(types).mean()\n    gmlmae = np.log(maes).mean()\n    return gmlmae\n        \ndef contribs_rmse_loss(preds, targs):\n    \"\"\"\n    Returns the sum of RMSEs for each sc contribution and total sc value.\n    \"\"\"\n    targs = reshape_targs(targs)\n    return torch.mean((preds - targs) ** 2, dim=0).sqrt().sum()\n\ndef rmse(preds, targs):\n    targs = reshape_targs(targs)\n    return torch.sqrt(F.mse_loss(preds[:,-1], targs[:,-1]))\n\ndef mae(preds, targs):\n    targs = reshape_targs(targs)\n    return torch.abs(preds[:,-1] - targs[:,-1]).mean()\n\nclass GroupMeanLogMAE(Callback):\n    \"\"\"Callback to repeort the group mean log MAE during taining.\"\"\"\n    _order = -20 # Needs to run before the recorder\n\n    def __init__(self, learn, snapshot_ensemble=False, **kwargs):\n        self.learn = learn\n        self.snapshot_ensemble = snapshot_ensemble\n\n    def on_train_begin(self, **kwargs):\n        metric_names = ['group_mean_log_mae']\n        if self.snapshot_ensemble: metric_names += ['group_mean_log_mae_es']\n        self.learn.recorder.add_metric_names(metric_names)\n        if self.snapshot_ensemble: self.val_preds = []\n\n    def on_epoch_begin(self, **kwargs):\n        self.sc_types, self.output, self.target = [], [], []\n\n    def on_batch_end(self, last_target, last_output, last_input, train,\n                     **kwargs):\n        if not train:\n            sc_types = last_input[-1].view(-1)\n            mask = sc_types != BATCH_PAD_VAL\n            self.sc_types.append(sc_types[mask])\n            self.output.append(last_output[:,-1])\n            self.target.append(reshape_targs(last_target)[:,-1])\n\n    def on_epoch_end(self, epoch, last_metrics, **kwargs):\n        if (len(self.sc_types) > 0) and (len(self.output) > 0):\n            sc_types = torch.cat(self.sc_types)\n            preds = torch.cat(self.output)\n            target = torch.cat(self.target)\n            metrics = [group_mean_log_mae(preds, target, sc_types, SC_MEAN, SC_STD)]\n\n            if self.snapshot_ensemble:\n                self.val_preds.append(preds.view(-1,1))\n                preds_se = torch.cat(self.val_preds, dim=1).mean(dim=1)\n                metrics += [group_mean_log_mae(preds_se, target, sc_types, SC_MEAN, SC_STD)]\n            return add_metrics(last_metrics, metrics)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# set up model\nset_seed(100)\nwd, d_model = 1e-2, 512\nenn_args = dict(layers=3*[d_model], dropout=3*[0.0], layer_norm=True)\nann_args = dict(layers=1*[d_model], dropout=1*[0.0], layer_norm=True,\n                out_act=nn.Tanh())\nmodel = Transformer(\n    N_ATOM_FEATURES, N_BOND_FEATURES, N_SC_EDGE_FEATURES,\n    N_SC_MOL_FEATURES, N=6, d_model=d_model, d_ff=d_model*4,\n    d_ff_contrib=d_model//4, h=8, dropout=0.0, \n    kernel_sz=min(128, d_model), enn_args=enn_args, ann_args=ann_args\n)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"callback_fns = [\n    GroupMeanLogMAE,\n    partial(SaveModelCallback, every='improvement', mode='min',\n            monitor='group_mean_log_mae', name=model_str)\n]\nlearn = Learner(db, model, metrics=[rmse, mae], callback_fns=callback_fns,\n                wd=wd, loss_func=contribs_rmse_loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.lr_find(start_lr=1e-7, end_lr=1.0, num_it=100, stop_div=True)\nlearn.recorder.plot()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"learn.fit_one_cycle(15, max_lr=1e-3)\nlearn.recorder.plot_losses(skip_start=500)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# make predictions\ntest_contrib_preds = learn.get_preds(DatasetType.Test)\ntest_preds = test_contrib_preds[0][:,-1].detach().numpy() * SC_STD + SC_MEAN\n\n# store results\nsubmit = pd.read_csv(RAW_DATA_PATH + 'sample_submission.csv')\nsubmit['scalar_coupling_constant'] = test_preds\nsubmit.to_csv(f'{model_str}-submission.csv', index=False)\nsubmit.head()","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}