{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install ase==3.17 schnetpack==0.2.1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls ../input","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nmolecules = pd.read_csv('../input/champs-scalar-coupling/structures.csv')\nmolecules = molecules.groupby('molecule_name')\ntrain = pd.read_csv('../input/champs-scalar-coupling/train.csv')\ntest = pd.read_csv('../input/champs-scalar-coupling/test.csv')\ntest['scalar_coupling_constant'] = -1\n\n# coupling_type = '1JHC'\n\n# train = train[train.type == coupling_type]\n# test = test[test.type == coupling_type]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"J_type = sorted(train.type.unique())\nJ_type","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"type_map = {}\nfor i, t in enumerate(J_type): \n    type_map[t] = i\n\ninverse_type_map = {}\nfor i, t in enumerate(J_type): \n    inverse_type_map[i] = t","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(train)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(test)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_scalar_couplings = train.groupby('molecule_name')\ntest_scalar_couplings = test.groupby('molecule_name')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# ASE Database"},{"metadata":{"trusted":true},"cell_type":"code","source":"from ase import Atoms\nfrom ase.db import connect\nfrom tqdm import *\n\ndef create_db(db_path, scalar_couplings, molecule_names):\n    with connect(db_path) as db:\n        with tqdm(total=len(molecule_names)) as pbar:\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                numbers = atoms.get_atomic_numbers()\n                group = scalar_couplings.get_group(name)\n                ai0 = group.atom_index_0.values\n                ai1 = group.atom_index_1.values\n                scc = group.scalar_coupling_constant.values\n                ids = group.id.values\n                types = group.type.values\n                for i, j, v, w, t in zip(ai0, ai1, scc, ids, types):\n                    new_numbers = numbers.copy()\n                    new_numbers[i] = 115 - new_numbers[i]\n                    new_numbers[j] = 115 - new_numbers[j]\n                    atoms.set_atomic_numbers(new_numbers)\n                    data = dict(scc=v)\n                    coupling_type = t\n                    data['type_id'] = w\n                    j_type = type_map[coupling_type]\n#                     j_type = np.zeros((8))\n#                     j_type[type_map[coupling_type]] = 1\n                    data['J_type'] = j_type\n                    db.write(atoms, name=name+'_H{}_C{}'.format(i,j), data=data)\n                    \n                pbar.update()\n                ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"properties=['J_type', 'type_id']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import schnetpack\n\nimport sys\nINT_MAX = sys.maxsize\n\ndataset_size = INT_MAX\n\ndataset_molecule_names = train.molecule_name.unique()\nprint(len(dataset_molecule_names))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"champs_path = 'CHAMPS_train_J_type.db' \nmolecule_names = dataset_molecule_names[:dataset_size]\ncreate_db(db_path=champs_path,\n          scalar_couplings=train_scalar_couplings,\n          molecule_names=molecule_names[:40000])\ndataset = schnetpack.data.AtomsData(champs_path, properties=properties)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#dataset[30]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_molecule_names = test.molecule_name.unique()\ntest_champs_path = 'CHAMPS_test_J_type.db' \ntest_molecule_names = dataset_molecule_names[:9000]\ncreate_db(db_path=test_champs_path,\n          scalar_couplings=test_scalar_couplings,\n          molecule_names=test_molecule_names)\ntest_dataset = schnetpack.data.AtomsData(test_champs_path, properties=properties)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(test_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# SchNet Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import MultiStepLR\n\nimport schnetpack as spk\nimport schnetpack.atomistic as atm\nimport schnetpack.representation as rep\nfrom schnetpack.datasets import *\n\nif torch.cuda.is_available():\n    device = torch.device(\"cuda\")\nelse:\n    device = torch.device('cpu')\ntorch.manual_seed(21)\nnp.random.seed(21)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# The original function comes from the following script:\n# https://github.com/atomistic-machine-learning/schnetpack/blob/v0.2.1/src/scripts/schnetpack_qm9.py\ndef evaluate_dataset(metrics, model, loader, device):\n#     for metric in metrics:\n#         metric.reset()\n    res = []\n    with torch.no_grad():\n        for batch in loader:\n            batch = {\n                k: v.to(device)\n                for k, v in batch.items()\n            }\n            result = model(batch)\n\n            for metric in metrics:\n                l = metric(batch, result)\n                res.append(np.array(l.detach().cpu().data.numpy()))\n    results = np.array(res).mean()\n    return results","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch.nn as nn\nfrom schnetpack.data import Structure\n\nclass MolecularOutput(atm.OutputModule):\n    def __init__(self, property_name, n_in=128, n_out=64, aggregation_mode='avg',\n                 n_layers=2, n_neurons=None,\n                 activation=schnetpack.nn.activations.shifted_softplus,\n                 outnet=None):\n        super(MolecularOutput, self).__init__(n_in, n_out)\n        self.property_name = property_name\n        self.n_layers = n_layers\n        self.create_graph = False\n        \n        if outnet is None:\n            self.out_net = nn.Sequential(\n                schnetpack.nn.base.GetItem('representation'),\n                schnetpack.nn.blocks.MLP(n_in, n_out, n_neurons, n_layers, activation)\n            )\n        else:\n            self.out_net = outnet\n        \n        self.FC1 = schnetpack.nn.blocks.MLP(n_out, 32, n_neurons, 1, activation)\n        self.out = schnetpack.nn.blocks.MLP(32, 8, n_neurons, 1, None)\n        \n        if aggregation_mode == 'sum':\n            self.atom_pool = schnetpack.nn.base.Aggregate(axis=1, mean=False)\n        elif aggregation_mode == 'avg':\n            self.atom_pool = schnetpack.nn.base.Aggregate(axis=1, mean=True)\n            \n    def forward(self, inputs):\n        r\"\"\"\n        predicts molecular property\n        \"\"\"\n        atom_mask = inputs[Structure.atom_mask]\n\n        yi = self.out_net(inputs)\n        y = self.atom_pool(yi, atom_mask)\n        y = self.FC1(y)\n        y = self.out(y)\n        result = {self.property_name: y}\n        return result","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def schnet_model():\n    reps = rep.SchNet(n_atom_basis=128, n_filters=128, n_interactions=6, max_z=115)\n    output = MolecularOutput('J_type')\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":"def train_model(max_epochs=50):\n    # print configuration\n    print('max_epochs:', max_epochs)\n    \n    # split in train and val\n    n_dataset = len(dataset)\n    n_val = n_dataset // 10\n    train_data, val_data, test_data = dataset.create_splits(n_dataset-n_val*2, n_val, 'split')\n    train_loader = spk.data.AtomsLoader(train_data, batch_size=128, num_workers=4, shuffle=True)\n    val_loader = spk.data.AtomsLoader(val_data, batch_size=128, num_workers=4)\n\n    # create model\n    model = schnet_model()\n\n    # create trainer\n    output_key = \"J_type\"\n    target_key = \"J_type\"\n    opt = Adam(model.parameters(), lr=1e-3, weight_decay=1e-6)\n    scheduler = MultiStepLR(opt, milestones=[15, 320], gamma=0.2)\n    def loss(b, p): \n        l = nn.CrossEntropyLoss()\n        out = l(p[output_key], b[target_key].view(-1).long())\n        return out\n    metrics = [loss]\n\n    hooks = [\n        spk.train.MaxEpochHook(max_epochs),\n#         spk.train.CSVHook('log', loss, every_n_epochs=1),\n        spk.train.LRScheduleHook(scheduler),\n    ]\n    trainer = spk.train.Trainer('output', model, loss,\n                                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('output/best_model'))\n    test_loader = spk.data.AtomsLoader(test_data, batch_size=128, num_workers=4)\n    model.eval()\n\n    df = pd.DataFrame()\n\n    df['metric'] = [\n        'cross_entropy'\n    ]\n    df['training'] = evaluate_dataset(metrics, model, train_loader, device)\n    df['validation'] = evaluate_dataset(metrics, model, val_loader, device)\n    df['test'] = evaluate_dataset(metrics, model, test_loader, device)\n    df.to_csv('output/evaluation.csv', index=False)\n    display(df)\n    \n    return test_data","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_history():\n    df = pd.read_csv('log/log.csv')\n    display(df.tail())\n    \n    _ = display(df[['MAE_scc', 'RMSE_scc']].plot())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def test_prediction(dataset):\n    # create model\n    model = schnet_model()\n    \n    # load best parameters\n    model.load_state_dict(torch.load('output/best_model'))\n    loader = spk.data.AtomsLoader(dataset, batch_size=128, num_workers=4)\n    model.eval()\n    \n    # predict scalar coupling constants\n    entry_id = []\n    predictions = []\n    with torch.no_grad():\n        for batch in loader:\n            batch = {\n                k: v.to(device)\n                for k, v in batch.items()\n            }\n            result = model(batch)\n            _, predicted = torch.max(result['J_type'], 1)\n#             print(predicted.shape)\n            entry_id += batch['type_id'].long().view(-1).tolist()\n            predictions += predicted.view(-1).tolist()\n    return entry_id, predictions","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Results"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nused_test_data = train_model(max_epochs=1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"split = np.load('split.npz')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"list(split.keys())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"used_test_data =  dataset.create_subset(split['test_idx'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_submission():\n    type_id, J_type = test_prediction(test_dataset)\n    submission = pd.DataFrame()\n    submission['id'] = type_id\n    submission['J_type'] = J_type\n    \n    return submission","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":" submission = make_submission()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission['J_type'] = submission['J_type'].map(lambda x: inverse_type_map[x])\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"display(submission.head())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission.to_csv('submission_J_type.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"display(test[:6956].head())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn import preprocessing\n\nle = preprocessing.LabelEncoder()\nle.fit(test['type'])\np = le.transform(submission['J_type'])\nt = le.transform(test[:6956]['type'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"p","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"t","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.metrics import accuracy_score\naccuracy_score(t, p)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.8"}},"nbformat":4,"nbformat_minor":1}