{"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\ncoupling_type = '1JHN'\n\ntrain = train[train.type == coupling_type]\ntest = test[test.type == coupling_type]","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":false},"cell_type":"code","source":"from ase import Atoms\nfrom ase.db import connect\n\ndef create_db(db_path, scalar_couplings, 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            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            for i, j, v, w in zip(ai0, ai1, scc, ids):\n                new_numbers = numbers.copy()\n                new_numbers[i] = 100 - new_numbers[i]\n                new_numbers[j] = 100 - new_numbers[j]\n                atoms.set_atomic_numbers(new_numbers)\n                data = dict(scc=v)\n                data[coupling_type+'_id'] = w\n                db.write(atoms, name=name+'_H{}_C{}'.format(i,j), data=data)\n                ","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"properties=['scc', coupling_type+'_id']","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"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()\nchamps_path = 'CHAMPS_train.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)\ndataset = schnetpack.data.AtomsData(champs_path, properties=properties)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"#dataset[30]","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"len(dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"dataset_molecule_names = test.molecule_name.unique()\ntest_champs_path = 'CHAMPS_test.db' \ntest_molecule_names = dataset_molecule_names[:dataset_size]\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":false},"cell_type":"code","source":"#test_dataset[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"len(test_dataset)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# SchNet Model"},{"metadata":{"trusted":false},"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\ndevice = torch.device(\"cuda\")\n#device = torch.device(\"cpu\")\ntorch.manual_seed(12345)\nnp.random.seed(12345)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"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\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                metric.add_batch(batch, result)\n\n    results = [\n        metric.aggregate() for metric in metrics\n    ]\n    return results","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"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=1, aggregation_mode='sum',\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        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\n        result = {self.property_name: y}\n\n        return result","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def schnet_model():\n    reps = rep.SchNet(n_atom_basis=128, n_filters=128, n_interactions=6)\n    output = MolecularOutput('scc')\n    model = atm.AtomisticModel(reps, output)\n    model = model.to(device)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def train_model(max_epochs=500):\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 = \"scc\"\n    target_key = \"scc\"\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        return F.mse_loss(p[output_key], b[target_key])\n    \n    metrics = [\n        spk.metrics.MeanAbsoluteError(target_key, output_key, name='MAE_scc'),\n        spk.metrics.RootMeanSquaredError(target_key, output_key, name='RMSE_scc'),\n    ]\n    hooks = [\n        spk.train.MaxEpochHook(max_epochs),\n        spk.train.CSVHook('log', metrics, 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    df['metric'] = [\n        'MAE_scc', 'RMSE_scc',\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":false},"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":false},"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            entry_id += batch[coupling_type+'_id'].long().view(-1).tolist()\n            predictions += result['scc'].view(-1).tolist()\n    return entry_id, predictions","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def show_predictions(dataset, train_df):\n    scc_id, scc = test_prediction(dataset)\n    df_pred = pd.DataFrame()\n    df_pred['Prediction'] = scc\n    df_pred['id'] = scc_id\n    df_pred = train_df.merge(df_pred, on='id', how='inner')\n    display(df_pred.head())\n    df_pred['Target'] = df_pred['scalar_coupling_constant']\n    display(df_pred.plot.scatter(x='Target', y='Prediction', title=coupling_type, figsize=(5,5)))\n    df_pred[['id', 'Target', 'Prediction']].to_csv('test_predictoins_{}.csv'.format(coupling_type), index=False)\n    \n    diff = (df_pred['Prediction'].values-df_pred['Target'].values)\n    rmse = np.sqrt(np.mean(diff**2))\n    mae = np.mean(np.abs(diff))\n    df_eval = pd.DataFrame()\n    df_eval['RMSE'] = [rmse]\n    df_eval['MAE'] = [mae]\n    df_eval['log(MAE)'] = [np.log(mae)]\n    _ = display(df_eval)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Results"},{"metadata":{"trusted":false},"cell_type":"code","source":"used_test_data = train_model(max_epochs=400)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"show_history()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"split = np.load('split.npz')","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"list(split.keys())","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"used_test_data =  dataset.create_subset(split['test_idx'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"show_predictions(used_test_data, train)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"def make_submission():\n    scc_id, scc = test_prediction(test_dataset)\n    submission = pd.DataFrame()\n    submission['id'] = scc_id\n    submission['scalar_coupling_constant'] = scc\n    submission.to_csv('submission_{}.csv'.format(coupling_type), index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":" make_submission()","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}