{"cells":[{"metadata":{},"cell_type":"markdown","source":"This kernel presents how to build an ASE database from CHAMPS data and create a dataset, and then uses it for molecular property predictions.\n\nFirst, please see an SchNet tutorial on QM9:\n\n[SchNet QM9](https://www.kaggle.com/tonyyy/schnet-qm9)\n\nwhich contains references of the SchNet architecture.\n\nThis kernel uses only potential energies and dipole moments in CHAMPS data. Mulliken charges and magnetic shielding tensors are used in the following kernels:\n* [SchNet Mulliken Charges](https://www.kaggle.com/tonyyy/schnet-mulliken-charges)\n* [SchNet Magnetic Shielding](https://www.kaggle.com/tonyyy/schnet-magnetic-shielding)"},{"metadata":{"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":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"!ls ../input","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# CHAMPS data"},{"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')\nenergies = pd.read_csv('../input/potential_energy.csv')\ndipoles = pd.read_csv('../input/dipole_moments.csv')\ndipoles['scalar'] = np.sqrt(np.square(dipoles[['X', 'Y', 'Z']]).sum(axis=1))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The total number of molecules:"},{"metadata":{"trusted":true},"cell_type":"code","source":"molecules.ngroups","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Potential Energy"},{"metadata":{"trusted":true},"cell_type":"code","source":"energies.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(energies)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"energy_series = pd.Series(energies.set_index('molecule_name')['potential_energy'])\nenergy_series.describe()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ax = energy_series.hist(bins=50)\n_ = ax.set_xlabel(\"Potential Energy\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dipole Moment"},{"metadata":{"trusted":true},"cell_type":"code","source":"dipoles.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(dipoles)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dipole_series = pd.Series(dipoles.set_index('molecule_name')['scalar'])\ndipole_series.describe()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ax = dipole_series.hist(bins=50)\n_ = ax.set_xlabel(\"Dipole Moment\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# ASE Database\n\nA dataset object used in SchNetPack is built on an ASE database. Please see the document of the ASE database:\n\nhttps://wiki.fysik.dtu.dk/ase/ase/db/db.html"},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv('../input/train.csv')\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_molecule_names = train.molecule_name.unique()\nlen(train_molecule_names)","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","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            db.write(atoms, name=name,\n                     potential_energy=energy_series.get(name, default=float('nan')),\n                     scalar_dipole=dipole_series.get(name, default=float('nan'))\n                    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"champs_path = 'CHAMPS_train.db'\ndataset_size = len(train_molecule_names) # 20000\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":"with connect(champs_path) as db:\n    print(len(db))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import schnetpack\n\nproperties=['potential_energy', 'scalar_dipole']\n\ndataset = dict()\nfor p in properties:\n    dataset[p] = schnetpack.data.AtomsData(champs_path, properties=[p])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for p in properties:\n    print(p, len(dataset[p]))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# SchNet Model\n\nPlease see the SchNetPack API document:\n\nhttps://schnetpack.readthedocs.io/en/stable/modules/index.html"},{"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},"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":"def schnet_model(property):\n    reps = rep.SchNet(n_interactions=6)\n    if 'dipole' in property:\n        print('use dipole moment')\n        output = atm.DipoleMoment(n_in=128, predict_magnitude=True)\n    else:\n        output = atm.Atomwise()\n    model = atm.AtomisticModel(reps, output)\n    model = model.to(device)\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_model(property, max_epochs=500):\n    # split in train and val\n    n_dataset = len(dataset[property])\n    n_val = n_dataset // 10\n    train_data, val_data, test_data = dataset[property].create_splits(n_dataset-n_val*2, 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(property)\n\n    # create trainer\n    opt = Adam(model.parameters(), lr=2e-4, weight_decay=1e-6)\n    loss = lambda b, p: F.mse_loss(p[\"y\"], b[property])\n    metrics = [\n        spk.metrics.MeanAbsoluteError(property, \"y\"),\n        spk.metrics.RootMeanSquaredError(property, \"y\"),\n    ]\n    hooks = [\n        spk.train.MaxEpochHook(max_epochs),\n        spk.train.CSVHook(property+'/log', metrics, every_n_epochs=1),\n    ]\n    trainer = spk.train.Trainer(property+'/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(property+'/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', 'RMSE']\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(property):\n    df = pd.read_csv(property+'/log/log.csv')\n    display(df.tail())\n    max_value = None # df['RMSE_'+property].min()*5\n    _ = df[['MAE_'+property,'RMSE_'+property]].plot(ylim=(0,max_value))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def test_prediction(dataset, property):\n    # create model\n    model = schnet_model(property)\n    \n    # load the best parameters\n    model.load_state_dict(torch.load(property+'/output/best_model'))\n    loader = spk.data.AtomsLoader(dataset, batch_size=256, num_workers=2)\n    model.eval()\n    \n    # predict molecular properties\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[property].squeeze().tolist()\n            predictions += result['y'].squeeze().tolist()\n    return targets, predictions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_predictions(dataset, property):\n    targets, predictions = test_prediction(dataset, property)\n    df_pred = pd.DataFrame()\n    df_pred['Target'] = targets\n    df_pred['Prediction'] = predictions\n    df_pred.plot.scatter(x='Target', y='Prediction', title=property)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Results"},{"metadata":{"trusted":true},"cell_type":"code","source":"used_test_data = dict()\nfor p in properties:\n    print(p)\n    used_test_data[p] = train_model(p, max_epochs=100)\n    show_history(p)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for p in properties:\n    show_predictions(used_test_data[p], p)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!mv potential_energy/log/log.csv log_potential.csv\n!mv scalar_dipole/log/log.csv log_dipole.csv\n!mv potential_energy/output/best_model best_model_potential\n!mv scalar_dipole/output/best_model best_model_dipole\n!rm -r potential_energy\n!rm -r scalar_dipole","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls","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}