{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import os\n__print__ = print\ndef print(string, end = '', flush = True):\n    os.system(f'echo \\\"{string}\\\"')\n    __print__(string, end = end, flush = flush)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Hi')","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"!pip install ase==3.17 schnetpack","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install  --no-cache-dir torch-scatter","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\n\nimport schnetpack as spk\nimport torch\n\nfrom ase import Atoms\nfrom ase.db import connect\n\nimport time\n\nfrom torch_scatter import *","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv('/kaggle/input/champs-scalar-coupling/train.csv')\ndf = df[df['type'] == '1JHC']\nmol_names = df['molecule_name'].unique()\nmolecules = pd.read_csv('/kaggle/input/champs-scalar-coupling/structures.csv')\nmolecules = molecules.set_index('molecule_name')\nmolecules = molecules.loc[mol_names]\nmolecules = molecules.reset_index()\nmolecules = molecules.groupby('molecule_name')\ndf = df.groupby('molecule_name')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def one_hot_encoding(x, set):\n    one_hot = [int(x == s) for s in set]\n    if 0:\n        if sum(one_hot)==0: print('one_hot_encoding() return NULL!', x, set)\n    return one_hot","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_db(db_path, molecule_names):\n    with connect(db_path) as db:\n        for name in molecule_names:\n            mol = molecules.get_group(name)\n            atoms = Atoms(symbols=mol.atom.values,\n                          positions=[[row.x,row.y,row.z] for row in mol.itertuples()])\n            #charges = mulliken_charges.get_group(name).mulliken_charge.values.reshape(-1,1)\n            df_target = df.get_group(name)\n            \n            indexes = []\n            targets = []\n            types = []\n            types_to_num = {'1JHC':0, '2JHH':1, '1JHN':2, '2JHN':3, '2JHC':4, '3JHH':5, '3JHC':6, '3JHN':7}\n            for row in df_target.itertuples():\n                indexes.append([row.atom_index_0 ,row.atom_index_1])\n                targets.append(row.scalar_coupling_constant)\n                types.append(types_to_num[row.type])\n            db.write(atoms, name=name,\n                     data={'indexes': indexes, 'targets': targets, 'types':types})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(mol_names)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_molecule_names = mol_names\nchamps_path = 'CHAMPS_train.db'\ndataset_size =  len(mol_names)\ndataset_molecule_names = train_molecule_names[:dataset_size]\ncreate_db(db_path=champs_path, molecule_names=dataset_molecule_names)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import schnetpack\ndatabase = schnetpack.data.AtomsData(champs_path,\n                properties=['indexes', 'targets', 'types'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn.functional as F\nfrom torch.optim import Adam\nfrom torch import nn\n\nimport schnetpack as spk\nimport schnetpack.atomistic as atm\nimport schnetpack.representation as rep\nfrom schnetpack.datasets import *\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nfrom schnetpack.data import Structure\nfrom schnetpack.atomistic import OutputModule\n\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Batched index_select\ndef batched_index_select(t, dim, inds):\n    dummy = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), t.size(2))\n    out = t.gather(dim, dummy) # b x e x f\n    return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Set2Set(torch.nn.Module):\n\n    def softmax(self, x, index, num=None):\n        x = x -  scatter_max(x, index, dim=0, dim_size=num)[0][index]\n        x = x.exp()\n        x = x / (scatter_add(x, index, dim=0, dim_size=num)[index] + 1e-16)\n        return x\n\n    def __init__(self, in_channel, processing_step=1):\n        super(Set2Set, self).__init__()\n        num_layer = 1\n        out_channel = 2 * in_channel\n\n        self.processing_step = processing_step\n        self.in_channel  = in_channel\n        self.out_channel = out_channel\n        self.num_layer   = num_layer\n        self.lstm = torch.nn.LSTM(out_channel, in_channel, num_layer)\n        self.lstm.reset_parameters()\n\n    def forward(self, x, batch_index):\n        batch_size = batch_index.max().item() + 1\n\n        h = (x.new_zeros((self.num_layer, batch_size, self.in_channel)),\n             x.new_zeros((self.num_layer, batch_size, self.in_channel)))\n\n        q_star = x.new_zeros(batch_size, self.out_channel)\n        for i in range(self.processing_step):\n            q, h = self.lstm(q_star.unsqueeze(0), h)\n            q = q.view(batch_size, -1)\n\n            e = (x * q[batch_index]).sum(dim=-1, keepdim=True) #shape = num_node x 1\n            a = self.softmax(e, batch_index, num=batch_size)   #shape = num_node x 1\n            r = scatter_add(a * x, batch_index, dim=0, dim_size=batch_size) #apply attention #shape = batch_size x ...\n            q_star = torch.cat([q, r], dim=-1)\n\n        return q_star","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from schnetpack.data import Structure\n\nclass Scalar_Coupling(atm.Atomwise):\n    def __init__(self):\n        super(Scalar_Coupling, self).__init__(return_contributions=True, n_out = 128)\n        \n        #self.att = Attention(128)\n        self.num_s2s = 6\n        self.red = nn.Linear(256, 128)\n        self.set2set = Set2Set(128, processing_step=self.num_s2s)\n        \n        self.predict = nn.Sequential(\n            nn.Linear(512, 256),  #node_hidden_dim\n            nn.ReLU(inplace=True),\n            nn.Linear(256, 128),\n            nn.ReLU(inplace=True),\n            nn.Linear(128, 8),\n        )\n        \n    def forward(self, inputs):\n        result = super().forward(inputs)\n        #print(inputs['targets'].shape)\n        embeds = result['yi']\n        indexes = inputs['indexes'].long()\n        targets = inputs['targets']\n        mask = (targets != 0).reshape(indexes.shape[0] * indexes.shape[1])\n        \n        embed_index = torch.tensor([embeds.shape[1]*[i] for i in range(embeds.shape[0])]).reshape(embeds.shape[0]*embeds.shape[1]).to(device)\n        \n        pool = self.set2set(embeds.reshape((embeds.shape[0]*embeds.shape[1],embeds.shape[2])), embed_index)\n        \n        x1 = batched_index_select(embeds, 1,indexes[:,:,0])\n        x2 = batched_index_select(embeds, 1,indexes[:,:,1])\n        #print(embeds.shape)\n        x = torch.cat((x1, x2), dim = 2)\n        \n        x = x.reshape((x.shape[0]*x.shape[1], x.shape[2]))\n        pool = pool.repeat(1, indexes.shape[1]).reshape(x.shape[0], -1)\n        x = torch.cat((x, pool), dim = 1)\n        '''\n        query = self.red(query)\n        query = query.unsqueeze(dim = 1)\n        context = embeds.repeat(x.shape[1], 1, 1)\n        pool, _ = self.att(query, context)\n        '''\n        #x = torch.cat()\n        #x = torch.cat((query, pool), dim = 2).squeeze(dim = 1)\n        #x = x[mask]\n        #print(targets.shape)\n        targets = targets.reshape(targets.shape[0]*targets.shape[1])\n        targets = targets[mask]\n        #print(targets.shape)\n        #print(targets.shape)\n        inputs['scalar_true'] = targets.unsqueeze(dim = 1)\n        x = self.predict(x)\n        \n        types = inputs['types'].long()\n        types = types.reshape(types.shape[0]*types.shape[1]).unsqueeze(dim = 1)\n        #print(types)\n        #print(types.shape)\n        #print(x.shape)\n        x = torch.gather(x, dim = 1, index = types)\n        \n        x = x[mask]\n        \n       # x = x.reshape((targets.shape[0], targets.shape[1]))\n        #x = x.reshape((indexes.shape[0], indexes.shape[1]))\n        result['scalar_pred'] = x\n        #print(result['scalar_pred'].shape)\n        #print(inputs['scalar_true'].shape)\n        #print(inputs['scalar_true'].shape)\n        #print(result['scalar_pred'].shape)\n        \n        return result","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport torch\nimport numpy as np\n\n\nclass Trainer:\n    def __init__(self, model_path, model, loss_fn, optimizer,\n                 train_loader, validation_loader, keep_n_checkpoints=3,\n                 checkpoint_interval=10, validation_interval=1, hooks=[], loss_is_normalized=True):\n        self.model_path = model_path\n        self.checkpoint_path = os.path.join(self.model_path, 'checkpoints')\n        self.best_model = os.path.join(self.model_path, 'best_model')\n        self.train_loader = train_loader\n        self.validation_loader = validation_loader\n        self.validation_interval = validation_interval\n        self.keep_n_checkpoints = keep_n_checkpoints\n        self.hooks = hooks\n        self.loss_is_normalized = loss_is_normalized\n\n        self._model = model\n        self._stop = False\n        self.checkpoint_interval = checkpoint_interval\n\n        self.loss_fn = loss_fn\n        self.optimizer = optimizer\n        #self.scheduler = ReduceLROnPlateau(self.optimizer, 'min', min_lr = 0.00001)\n\n        if os.path.exists(self.checkpoint_path):\n            self.restore_checkpoint()\n        else:\n            os.makedirs(self.checkpoint_path)\n            self.epoch = 0\n            self.step = 0\n            self.best_loss = float('inf')\n            self.store_checkpoint()\n\n    def _check_is_parallel(self):\n        return True if isinstance(self._model, torch.nn.DataParallel) else False\n\n    def _load_model_state_dict(self, state_dict):\n        if self._check_is_parallel():\n            self._model.module.load_state_dict(state_dict)\n        else:\n            self._model.load_state_dict(state_dict)\n\n    @property\n    def state_dict(self):\n        state_dict = {\n            'epoch': self.epoch,\n            'step': self.step,\n            'model': self._model.state_dict() if not self._check_is_parallel() else self._model.module.state_dict(),\n            'best_loss': self.best_loss,\n            'optimizer': self.optimizer.state_dict(),\n            'hooks': [h.state_dict for h in self.hooks]\n        }\n        return state_dict\n\n    @state_dict.setter\n    def state_dict(self, state_dict):\n        self.epoch = state_dict['epoch']\n        self.step = state_dict['step']\n        self.best_loss = state_dict['best_loss']\n        self.optimizer.load_state_dict(state_dict['optimizer'])\n        self._load_model_state_dict(state_dict['model'])\n\n        for h, s in zip(self.hooks, self.state_dict['hooks']):\n            h.state_dict = s\n\n    def store_checkpoint(self):\n        chkpt = os.path.join(self.checkpoint_path,\n                             'checkpoint-' + str(self.epoch) + '.pth.tar')\n        torch.save(self.state_dict, chkpt)\n\n        chpts = [f for f in os.listdir(self.checkpoint_path)\n                 if f.endswith('.pth.tar')]\n        if len(chpts) > self.keep_n_checkpoints:\n            chpt_epochs = [int(f.split('.')[0].split('-')[-1]) for f in chpts]\n            sidx = np.argsort(chpt_epochs)\n            for i in sidx[:-self.keep_n_checkpoints]:\n                os.remove(os.path.join(self.checkpoint_path, chpts[i]))\n\n    def restore_checkpoint(self, epoch=None):\n        if epoch is None:\n            epoch = max([int(f.split('.')[0].split('-')[-1]) for f in os.listdir(self.checkpoint_path)])\n        epoch = str(epoch)\n\n        chkpt = os.path.join(self.checkpoint_path,\n                             'checkpoint-' + str(epoch) + '.pth.tar')\n        self.state_dict = torch.load(chkpt)\n\n    def train(self, device):\n        r\"\"\"\n        Starts training of model on a specified device.\n\n        Args:\n            device (torch.torch.Device): device on which training takes place\n\n        \"\"\"\n        self._stop = False\n\n        for h in self.hooks:\n            h.on_train_begin(self)\n\n        try:\n            while True:\n                self.epoch += 1\n\n                for h in self.hooks:\n                    h.on_epoch_begin(self)\n\n                if self._stop:\n                    break\n\n                # perform training epoch\n                for train_batch in self.train_loader:\n                    self.optimizer.zero_grad()\n\n                    for h in self.hooks:\n                        h.on_batch_begin(self, train_batch)\n\n                    # move input to gpu, if needed\n                    train_batch = {\n                        k: v.to(device)\n                        for k, v in train_batch.items()\n                    }\n\n                    result = self._model(train_batch)\n                    loss = self.loss_fn(train_batch, result)\n\n                    loss.backward()\n                    self.optimizer.step()\n                    self.step += 1\n\n                    for h in self.hooks:\n                        h.on_batch_end(self, train_batch, result, loss)\n\n                    if self._stop:\n                        break\n\n                if self.epoch % self.checkpoint_interval == 0:\n                    self.store_checkpoint()\n\n                # validation\n                if self.epoch % self.validation_interval == 0 or self._stop:\n                    for h in self.hooks:\n                        h.on_validation_begin(self)\n\n                    val_loss = 0.\n                    n_val = 0\n                    for val_batch in self.validation_loader:\n                        # append batch_size\n                        vsize = list(val_batch.values())[0].size(0)\n                        n_val += vsize\n\n                        for h in self.hooks:\n                            h.on_validation_batch_begin(self)\n\n                        # move input to gpu, if needed\n                        val_batch = {\n                            k: v.to(device)\n                            for k, v in val_batch.items()\n                        }\n\n                        val_result = self._model(val_batch)\n                        val_batch_loss = self.loss_fn(val_batch, val_result).data.cpu().numpy()\n                        if self.loss_is_normalized:\n                            val_loss += val_batch_loss * vsize\n                        else:\n                            val_loss += val_batch_loss\n\n                        for h in self.hooks:\n                            h.on_validation_batch_end(self, val_batch, val_result)\n\n                    # weighted average over batches\n                    if self.loss_is_normalized:\n                        val_loss /= n_val\n\n                    if self.best_loss > val_loss:\n                        self.best_loss = val_loss\n                        state_dict = self._model.state_dict() if not self._check_is_parallel() else self._model.module.state_dict()\n                        torch.save(state_dict, self.best_model)\n\n                    for h in self.hooks:\n                        h.on_validation_end(self, val_loss)\n                \n                #self.scheduler.step(val_loss)\n                \n                for h in self.hooks:\n                    h.on_epoch_end(self)\n\n                if self._stop:\n                    break\n\n            for h in self.hooks:\n                h.on_train_ends(self)\n\n        except Exception as e:\n            for h in self.hooks:\n                h.on_train_failed(self)\n\n            raise e\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def schnet_model():\n    reps = rep.SchNet(n_interactions=6, cutoff=10.0, n_gaussians=100)\n    #reps = rep.SchNet(n_interactions=6)\n    output = Scalar_Coupling()\n    model = atm.AtomisticModel(reps, output)\n    model = model.to(device)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MyHook(spk.train.CSVHook):\n    \"\"\" Hook for logging to csv files.\n\n            This class provides an interface to write logging information about the training process to csv files.\n\n            Args:\n                log_path (str): path to directory to which log files will be written.\n                metrics (list): list containing all metrics to be logged. Metrics have to be subclass of spk.Metric.\n                log_train_loss (bool): enable logging of training loss (default: True)\n                log_validation_loss (bool): enable logging of validation loss (default: True)\n                log_learning_rate (bool): enable logging of current learning rate (default: True)\n                every_n_epochs (int): interval after which logging takes place (default: 1)\n        \"\"\"\n\n    def __init__(self, log_path, metrics, log_train_loss=True,\n                 log_validation_loss=True, log_learning_rate=True, every_n_epochs=1):\n        log_path = os.path.join(log_path, 'log.csv')\n        super(MyHook, self).__init__(log_path, metrics, log_train_loss,\n                                      log_validation_loss, log_learning_rate)\n        self._offset = 0\n        self._restart = False\n        self.every_n_epochs = every_n_epochs\n\n    def on_train_begin(self, trainer):\n\n        if os.path.exists(self.log_path):\n            remove_file = False\n            with open(self.log_path, 'r') as f:\n                # Ensure there is one entry apart from header\n                lines = f.readlines()\n                if len(lines) > 1:\n                    self._offset = float(lines[-1].split(',')[0]) - time.time()\n                    self._restart = True\n                else:\n                    remove_file = True\n\n            # Empty up to header, remove to avoid adding header twice\n            if remove_file:\n                os.remove(self.log_path)\n        else:\n            self._offset = -time.time()\n            # Create the log dir if it does not exists, since write cannot create a full path\n            log_dir = os.path.dirname(self.log_path)\n            if not os.path.exists(log_dir):\n                os.makedirs(log_dir)\n\n        if not self._restart:\n            log = ''\n            log += 'Time'\n\n            if self.log_learning_rate:\n                log += ',Learning rate'\n\n            if self.log_train_loss:\n                log += ',Train loss'\n\n            if self.log_validation_loss:\n                log += ',Validation loss'\n\n            if len(self.metrics) > 0:\n                log += ','\n\n            for i, metric in enumerate(self.metrics):\n                log += str(metric.name)\n                if i < len(self.metrics) - 1:\n                    log += ','\n\n            with open(self.log_path, 'a+') as f:\n                f.write(log + os.linesep)\n\n    def on_validation_end(self, trainer, val_loss):\n        if trainer.epoch % self.every_n_epochs == 0:\n            ctime = time.time() + self._offset\n            log = str(ctime)\n\n            if self.log_learning_rate:\n                log += ',' + str(trainer.optimizer.param_groups[0]['lr'])\n\n            if self.log_train_loss:\n                if hasattr(self._train_loss, \"__iter__\"):\n                    train_string = ','.join([str(k) for k in self._train_loss])\n                    log += ',' + train_string\n                else:\n                    log += ',' + str(self._train_loss)\n\n            if self.log_validation_loss:\n                if hasattr(val_loss, \"__iter__\"):\n                    valid_string = ','.join([str(k) for k in val_loss])\n                    log += ',' + valid_string\n                else:\n                    log += ',' + str(val_loss)\n\n            if len(self.metrics) > 0:\n                log += ','\n\n            for i, metric in enumerate(self.metrics):\n                m = metric.aggregate()\n                if hasattr(m, \"__iter__\"):\n                    log += ','.join([str(j) for j in m])\n                else:\n                    log += str(m)\n                if i < len(self.metrics) - 1:\n                    log += ','\n\n            with open(self.log_path, 'a') as f:\n                f.write(log + os.linesep)\n                print(log + os.linesep)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def criterion(predict, truth):\n    predict = predict.view(-1)\n    truth   = truth.view(-1)\n    assert(predict.shape==truth.shape)\n\n    loss = torch.abs(predict-truth)\n    loss = loss.mean()\n    loss = torch.log(loss)\n    return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_model(max_epochs=500):\n    n_dataset = len(database)\n    n_val = n_dataset // 20\n    train_data, val_data, test_data = database.create_splits(n_dataset - 2*n_val, n_val, 'split')\n    train_loader = spk.data.AtomsLoader(train_data, batch_size=64, num_workers=2,)\n    val_loader = spk.data.AtomsLoader(val_data + test_data , batch_size=256, num_workers=2)\n    \n    model = schnet_model()\n\n    # create trainer\n    true_key = 'scalar_true'\n    pred_key = 'scalar_pred'\n    opt = Adam(model.parameters(), lr=0.0002)\n    loss = lambda b, p: criterion(p[pred_key], b[true_key])\n    metrics = [\n        spk.metrics.MeanAbsoluteError(true_key, pred_key, name='MAE_scalar'),\n        spk.metrics.RootMeanSquaredError(true_key, pred_key, name='RMSE_scalar'),\n    ]\n    hooks = [\n        spk.train.MaxEpochHook(max_epochs),\n        MyHook('log', metrics, every_n_epochs=1),\n        #spk.train.CSVHook('log', metrics, every_n_epochs=1),\n    ]\n    #trainer = spk.train.Trainer('output', model, loss, opt, train_loader, val_loader, hooks=hooks)\n    trainer = Trainer('output', model, loss, opt, train_loader, val_loader, hooks=hooks)\n\n    # start training\n    trainer.train(device)\n    \n    # evaluation\n    model.load_state_dict(torch.load('best_model'))\n    test_loader = spk.data.AtomsLoader(test_data, batch_size=1, num_workers=2)\n    model.eval()\n    \n    return test_data","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"used_test_data = train_model(max_epochs=10000000)","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}