{"cells":[{"metadata":{},"cell_type":"markdown","source":"This kernel presents how to predict mulliken charges using SchNetPack."},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install ase==3.17 schnetpack","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We need `ASE 3.17` for `SchNetPack 0.2.1`."},{"metadata":{},"cell_type":"markdown","source":"# Mulliken Charges"},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nmolecules = pd.read_csv('../input/structures.csv')\nmolecules = molecules.groupby('molecule_name')\nmulliken_charges = pd.read_csv('../input/mulliken_charges.csv')\nmulliken_charges = mulliken_charges.groupby('molecule_name')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mulliken_charges.get_group('dsgdb9nsd_000003')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mulliken_charges.get_group('dsgdb9nsd_000003').mulliken_charge.sum()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pd.read_csv('../input/mulliken_charges.csv').mulliken_charge.describe()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pd.read_csv('../input/mulliken_charges.csv').mulliken_charge.hist(bins=50)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from ase import Atoms\nfrom ase.db import connect\n\ndef 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            try:\n                charges = mulliken_charges.get_group(name).mulliken_charge.values.reshape(-1,1)\n                total_charge = 0.0\n            except KeyError:\n                charges = None\n                total_charge = None\n            db.write(atoms, name=name,\n                     data={'mulliken_charges': charges, 'total_charge': total_charge})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv('../input/train.csv')\ntrain_molecule_names = train.molecule_name.unique()\n\nchamps_path = 'CHAMPS_train.db'\ndataset_size =  len(train_molecule_names) # 12000\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=['mulliken_charges', 'total_charge'])","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\n\nimport schnetpack as spk\nimport schnetpack.atomistic as atm\nimport schnetpack.representation as rep\nfrom schnetpack.datasets import *\n\ndevice = torch.device(\"cuda\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# This 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":true},"cell_type":"code","source":"from schnetpack.data import Structure\n\nclass Mulliken(atm.Atomwise):\n    def __init__(self):\n        super(Mulliken, self).__init__(return_contributions=True)\n        \n    def forward(self, inputs):\n        result = super().forward(inputs)\n        \n        atom_mask = inputs[Structure.atom_mask].byte()\n        \n        yi = inputs['mulliken_charges']\n        yi = torch.masked_select(yi.squeeze(dim=2), atom_mask)\n        inputs['mulliken_true'] = yi\n        \n        yi = result['yi']\n        yi = torch.masked_select(yi.squeeze(dim=2), atom_mask)\n        result['mulliken_pred'] = yi\n        \n        return result\n        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def schnet_model():\n    reps = rep.SchNet(n_interactions=6)\n    output = Mulliken()\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=500):\n    # split in train and val\n    n_dataset = len(database)\n    n_val = n_dataset // 10\n    train_data, val_data, test_data = database.create_splits(n_dataset - 2*n_val, n_val)\n    train_loader = spk.data.AtomsLoader(train_data, batch_size=128, num_workers=2)\n    val_loader = spk.data.AtomsLoader(val_data, batch_size=256, num_workers=2)\n\n    # create model\n    model = schnet_model()\n\n    # create trainer\n    true_key = 'mulliken_true'\n    pred_key = 'mulliken_pred'\n    opt = Adam(model.parameters(), lr=1e-4)\n    loss = lambda b, p: F.mse_loss(p[pred_key], b[true_key])\n    metrics = [\n        spk.metrics.MeanAbsoluteError(true_key, pred_key, name='MAE_mulliken'),\n        spk.metrics.RootMeanSquaredError(true_key, pred_key, name='RMSE_mulliken'),\n    ]\n    hooks = [\n        spk.train.MaxEpochHook(max_epochs),\n        spk.train.CSVHook('log', metrics, every_n_epochs=1),\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=256, num_workers=2)\n    model.eval()\n\n    df = pd.DataFrame()\n    df['metric'] = ['MAE_mulliken', 'RMSE_mulliken']\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    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    max_value = None # df['RMSE_mulliken'].min()*5\n    _ = df[['MAE_mulliken','RMSE_mulliken']].plot(ylim=(0,max_value))","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=256, num_workers=2)\n    model.eval()\n    \n    # predict mulliken charges\n    targets = []\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            targets += batch['mulliken_true'].tolist()\n            predictions += result['mulliken_pred'].tolist()\n    return targets, predictions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_predictions(dataset):\n    targets, predictions = test_prediction(dataset)\n    df_pred = pd.DataFrame()\n    df_pred['Target'] = targets\n    df_pred['Prediction'] = predictions\n    df_pred.plot.scatter(x='Target', y='Prediction')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Results"},{"metadata":{"trusted":true},"cell_type":"code","source":"used_test_data = train_model(max_epochs=200)\nshow_history()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_predictions(used_test_data)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}