{"cells":[{"metadata":{},"cell_type":"markdown","source":"# SchNet\n\n![SchNet](https://user-images.githubusercontent.com/11532812/60562621-3a935c00-9d93-11e9-8af3-59318e61172b.png)\n\nSchNet is a state-of-the-art deep neural architecture that predicts molecular energies:\n\nhttps://paperswithcode.com/sota/formation-energy-on-qm9\n\nThis kernel shows how to train SchNet models using SchNetPack.\n\nThe SchNet is constructed from the following conmponents: \n\n## Embedding\nThe SchNet uses an embedding depending on the type of the center atom\n![embedding](https://user-images.githubusercontent.com/11532812/60563147-8d6e1300-9d95-11e9-9734-3df9b5d457fa.png)\nfor the initial representation of local chemical environment *i*.\n\n## Atom-wise, Fully-connected Layers\n![Atom-wise FC](https://user-images.githubusercontent.com/11532812/60563347-66fca780-9d96-11e9-8b4a-63debdb676ac.png)\n\n## Continuous-filter Convolutional Layers\n![cont-filter-conv](https://user-images.githubusercontent.com/11532812/60563492-002bbe00-9d97-11e9-8784-9e7393a18285.png)\n\n## Filter-generating Networks\nDistance features are created by expanding the pair-wise distances:\n![distance-expansion](https://user-images.githubusercontent.com/11532812/60562913-7e3a9580-9d94-11e9-9067-ba40e70b2213.png)\n\nThe feature vector is used as an input for a mulilayer fully-connected neural network.\n\n## Output Layers\nThe atomic features {**x**_i} are fed to atom-wise multilayer perceptrons, which output atom-wise contributions of a chemical property. The property is calcuated by the sum of those contributions. \n\n### Potential Energy\n![potential-energy](https://user-images.githubusercontent.com/11532812/60563771-20a84800-9d98-11e9-8720-bf8cff96157a.png)\nE_i is an atom-wise contribution of the potential energy E.\n\n### Dipole Moment\n![dipole-moment](https://user-images.githubusercontent.com/11532812/60563830-53ead700-9d98-11e9-9908-98c06b95b1cf.png)\nq_i is an atomic charge produced by the SchNet, and r_i is an atomic position.\n\n## References\n* K.T. Schütt. P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, K.-R. Müller.\nSchNet: A continuous-filter convolutional neural network for modeling quantum interactions.\nAdvances in Neural Information Processing Systems 30, pp. 992-1002 (2017)\n[link](http://papers.nips.cc/paper/6700-schnet-a-continuous-filter-convolutional-neural-network-for-modeling-quantum-interactions)\n[arXiv](https://arxiv.org/abs/1706.08566v5)\n* K.T. Schütt. P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, K.-R. Müller.\nSchNet - a deep learning architecture for molecules and materials.\nThe Journal of Chemical Physics 148(24), 241722 (2018)\n[link](https://doi.org/10.1063/1.5019779)\n[arXiv](https://arxiv.org/abs/1712.06113)\n* K. T. Schütt, P. Kessel, M. Gastegger, K. A. Nicoli, A. Tkatchenko, and K.-R. Müller.\nSchNetPack: A Deep Learning Toolbox For Atomistic Systems.\nJournal of Chemical Theory and Computation 15(1), 448-455 (2019)\n[link](https://pubs.acs.org/doi/10.1021/acs.jctc.8b00908)\n[arXiv](https://arxiv.org/pdf/1809.01072.pdf)\n* Kim A. Nicoli, Pan Kessel, Michael Gastegger, Kristof T. Schütt.\nAnalysis of Atomistic Representations Using Weighted Skip-Connections.\n32nd Conference on Neural Information Processing Systems (NIPS 2018)\n[arXiv](https://arxiv.org/abs/1810.09751)\n* Kristof T. Schütt, Alexandre Tkatchenko, Klaus-Robert Müller.\nLearning representations of molecules and materials with atomistic neural networks.\n(2018)\n[arXiv](https://arxiv.org/abs/1812.04690)\n\n## Implemetations\n* PyTorch:\n  https://github.com/atomistic-machine-learning/schnetpack\n* TensorFlow:\n  https://github.com/atomistic-machine-learning/SchNet\n"},{"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":{},"cell_type":"markdown","source":"The following code is a derivative work from\n\nhttps://github.com/atomistic-machine-learning/schnetpack/blob/v0.2.1/src/examples/qm9_schnet.py"},{"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\")\n\n# load qm9 dataset and download if necessary\ndata = QM9(\"qm9/\", properties=[QM9.U0], remove_uncharacterized=True)\n\n# Statistics\nenergies = [data[i][QM9.U0].item() for i in range(len(data))]\nenergies = pd.Series(energies, name=QM9.U0)\ndisplay(energies.describe())\nax = energies.hist(bins=50)\n_ = ax.set_xlabel(QM9.U0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#!rm -r output log\n\n# split in train and val\nn_val = 10000\ntrain, val, test = data.create_splits(len(data)-n_val*2, n_val)\nloader = spk.data.AtomsLoader(train, batch_size=128, num_workers=2)\nval_loader = spk.data.AtomsLoader(val, batch_size=256, num_workers=2)\n\n# create model\nreps = rep.SchNet(n_interactions=6)\noutput = atm.Atomwise()\nmodel = atm.AtomisticModel(reps, output)\nmodel = model.to(device)\n\n# create trainer\nmax_epochs = 100\nopt = Adam(model.parameters(), lr=2e-4, weight_decay=1e-6)\nloss = lambda b, p: F.mse_loss(p[\"y\"], b[QM9.U0])\nmetric_list = [\n    spk.metrics.MeanAbsoluteError(QM9.U0, \"y\"),\n    spk.metrics.RootMeanSquaredError(QM9.U0, \"y\"),\n]\nhooks = [\n    spk.train.MaxEpochHook(max_epochs),\n    spk.train.CSVHook('log', metric_list, every_n_epochs=1),\n]\ntrainer = spk.train.Trainer(\"output/\", model, loss,\n                            opt, loader, val_loader, hooks=hooks)\n\n# start training\ntrainer.train(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df = pd.read_csv('log/log.csv')\ndisplay(df.tail())\n_ = df[['MAE_energy_U0','RMSE_energy_U0']].plot(ylim=(0,100))","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":"model.load_state_dict(torch.load('output/best_model'))\ntest_loader = spk.data.AtomsLoader(test, batch_size=256, num_workers=2)\nmodel.eval()\n\ndf = pd.DataFrame()\ndf['metric'] = ['MAE', 'RMSE']\ndf['training'] = evaluate_dataset(metric_list, model, loader, device)\ndf['validation'] = evaluate_dataset(metric_list, model, val_loader, device)\ndf['test'] = evaluate_dataset(metric_list, model, test_loader, device)\ndf","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}