{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install ase==3.17 schnetpack","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-11-22T18:39:15.580408Z","iopub.execute_input":"2022-11-22T18:39:15.581452Z","iopub.status.idle":"2022-11-22T18:39:33.839968Z","shell.execute_reply.started":"2022-11-22T18:39:15.581336Z","shell.execute_reply":"2022-11-22T18:39:33.838782Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls ../input","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2022-11-22T18:39:33.842118Z","iopub.execute_input":"2022-11-22T18:39:33.842555Z","iopub.status.idle":"2022-11-22T18:39:34.932749Z","shell.execute_reply.started":"2022-11-22T18:39:33.842518Z","shell.execute_reply":"2022-11-22T18:39:34.931105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CHAMPS data","metadata":{}},{"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')\nenergies = pd.read_csv('../input/champs-scalar-coupling/potential_energy.csv')\ndipoles = pd.read_csv('../input/champs-scalar-coupling/dipole_moments.csv')\ndipoles['scalar'] = np.sqrt(np.square(dipoles[['X', 'Y', 'Z']]).sum(axis=1))","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:32.112295Z","iopub.execute_input":"2022-11-22T18:40:32.11279Z","iopub.status.idle":"2022-11-22T18:40:37.51624Z","shell.execute_reply.started":"2022-11-22T18:40:32.112754Z","shell.execute_reply":"2022-11-22T18:40:37.515128Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The total number of molecules:","metadata":{}},{"cell_type":"code","source":"molecules.ngroups","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:42.078618Z","iopub.execute_input":"2022-11-22T18:40:42.079517Z","iopub.status.idle":"2022-11-22T18:40:42.407857Z","shell.execute_reply.started":"2022-11-22T18:40:42.079478Z","shell.execute_reply":"2022-11-22T18:40:42.406682Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Potential Energy","metadata":{}},{"cell_type":"code","source":"energies.head()","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:44.523294Z","iopub.execute_input":"2022-11-22T18:40:44.523745Z","iopub.status.idle":"2022-11-22T18:40:44.541376Z","shell.execute_reply.started":"2022-11-22T18:40:44.523708Z","shell.execute_reply":"2022-11-22T18:40:44.540177Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(energies)","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:46.717622Z","iopub.execute_input":"2022-11-22T18:40:46.718088Z","iopub.status.idle":"2022-11-22T18:40:46.72685Z","shell.execute_reply.started":"2022-11-22T18:40:46.718055Z","shell.execute_reply":"2022-11-22T18:40:46.72558Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"energy_series = pd.Series(energies.set_index('molecule_name')['potential_energy'])\nenergy_series.describe()","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:48.527211Z","iopub.execute_input":"2022-11-22T18:40:48.528226Z","iopub.status.idle":"2022-11-22T18:40:48.553353Z","shell.execute_reply.started":"2022-11-22T18:40:48.528158Z","shell.execute_reply":"2022-11-22T18:40:48.551939Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ax = energy_series.hist(bins=50)\n_ = ax.set_xlabel(\"Potential Energy\")","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:50.344572Z","iopub.execute_input":"2022-11-22T18:40:50.345035Z","iopub.status.idle":"2022-11-22T18:40:50.72414Z","shell.execute_reply.started":"2022-11-22T18:40:50.344998Z","shell.execute_reply":"2022-11-22T18:40:50.722947Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Dipole Moment","metadata":{}},{"cell_type":"code","source":"dipoles.head()","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:52.267003Z","iopub.execute_input":"2022-11-22T18:40:52.267454Z","iopub.status.idle":"2022-11-22T18:40:52.283629Z","shell.execute_reply.started":"2022-11-22T18:40:52.267402Z","shell.execute_reply":"2022-11-22T18:40:52.282484Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(dipoles)","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:54.282698Z","iopub.execute_input":"2022-11-22T18:40:54.283148Z","iopub.status.idle":"2022-11-22T18:40:54.290979Z","shell.execute_reply.started":"2022-11-22T18:40:54.283111Z","shell.execute_reply":"2022-11-22T18:40:54.289681Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dipole_series = pd.Series(dipoles.set_index('molecule_name')['scalar'])\ndipole_series.describe()","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:55.639077Z","iopub.execute_input":"2022-11-22T18:40:55.639772Z","iopub.status.idle":"2022-11-22T18:40:55.66672Z","shell.execute_reply.started":"2022-11-22T18:40:55.639732Z","shell.execute_reply":"2022-11-22T18:40:55.665434Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ax = dipole_series.hist(bins=50)\n_ = ax.set_xlabel(\"Dipole Moment\")","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:40:56.755433Z","iopub.execute_input":"2022-11-22T18:40:56.755949Z","iopub.status.idle":"2022-11-22T18:40:57.090275Z","shell.execute_reply.started":"2022-11-22T18:40:56.755903Z","shell.execute_reply":"2022-11-22T18:40:57.089226Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train = pd.read_csv('../input/champs-scalar-coupling/train.csv')\ntrain.head()","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:41:24.926511Z","iopub.execute_input":"2022-11-22T18:41:24.92693Z","iopub.status.idle":"2022-11-22T18:41:31.510777Z","shell.execute_reply.started":"2022-11-22T18:41:24.926894Z","shell.execute_reply":"2022-11-22T18:41:31.509548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_molecule_names = train.molecule_name.unique()\nlen(train_molecule_names)","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:42:44.695855Z","iopub.execute_input":"2022-11-22T18:42:44.696321Z","iopub.status.idle":"2022-11-22T18:42:45.115848Z","shell.execute_reply.started":"2022-11-22T18:42:44.696284Z","shell.execute_reply":"2022-11-22T18:42:45.114633Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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                    )","metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","execution":{"iopub.status.busy":"2022-11-22T18:42:48.590461Z","iopub.execute_input":"2022-11-22T18:42:48.590902Z","iopub.status.idle":"2022-11-22T18:42:48.599874Z","shell.execute_reply.started":"2022-11-22T18:42:48.590862Z","shell.execute_reply":"2022-11-22T18:42:48.598071Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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)","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:42:49.803954Z","iopub.execute_input":"2022-11-22T18:42:49.804645Z","iopub.status.idle":"2022-11-22T18:44:58.753615Z","shell.execute_reply.started":"2022-11-22T18:42:49.804592Z","shell.execute_reply":"2022-11-22T18:44:58.752471Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with connect(champs_path) as db:\n    print(len(db))","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:44:58.755518Z","iopub.execute_input":"2022-11-22T18:44:58.755866Z","iopub.status.idle":"2022-11-22T18:44:58.763613Z","shell.execute_reply.started":"2022-11-22T18:44:58.755831Z","shell.execute_reply":"2022-11-22T18:44:58.76246Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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])","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:45:16.778123Z","iopub.execute_input":"2022-11-22T18:45:16.779091Z","iopub.status.idle":"2022-11-22T18:45:16.786579Z","shell.execute_reply.started":"2022-11-22T18:45:16.779035Z","shell.execute_reply":"2022-11-22T18:45:16.784888Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for p in properties:\n    print(p, len(dataset[p]))","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:45:18.534309Z","iopub.execute_input":"2022-11-22T18:45:18.535376Z","iopub.status.idle":"2022-11-22T18:45:18.548153Z","shell.execute_reply.started":"2022-11-22T18:45:18.535331Z","shell.execute_reply":"2022-11-22T18:45:18.546533Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# SchNet Model\n\nPlease see the SchNetPack API document:\n\nhttps://schnetpack.readthedocs.io/en/stable/modules/index.html","metadata":{}},{"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\n","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:49:15.571848Z","iopub.execute_input":"2022-11-22T18:49:15.572324Z","iopub.status.idle":"2022-11-22T18:49:15.579179Z","shell.execute_reply.started":"2022-11-22T18:49:15.572281Z","shell.execute_reply":"2022-11-22T18:49:15.577883Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda\")\ndevice","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:52:23.015807Z","iopub.execute_input":"2022-11-22T18:52:23.01625Z","iopub.status.idle":"2022-11-22T18:52:23.024332Z","shell.execute_reply.started":"2022-11-22T18:52:23.016214Z","shell.execute_reply":"2022-11-22T18:52:23.023015Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def 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","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:49:22.384751Z","iopub.execute_input":"2022-11-22T18:49:22.385224Z","iopub.status.idle":"2022-11-22T18:49:22.393093Z","shell.execute_reply.started":"2022-11-22T18:49:22.38518Z","shell.execute_reply":"2022-11-22T18:49:22.392043Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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()\n    \n    return model","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:55:44.051742Z","iopub.execute_input":"2022-11-22T18:55:44.05219Z","iopub.status.idle":"2022-11-22T18:55:44.060195Z","shell.execute_reply.started":"2022-11-22T18:55:44.052152Z","shell.execute_reply":"2022-11-22T18:55:44.058493Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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(\"cuda\")\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)\n    df['validation'] = evaluate_dataset(metrics, model, val_loader)\n    df['test'] = evaluate_dataset(metrics, model, test_loader)\n    display(df)\n    \n    return test_data","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:56:06.475776Z","iopub.execute_input":"2022-11-22T18:56:06.476641Z","iopub.status.idle":"2022-11-22T18:56:06.489145Z","shell.execute_reply.started":"2022-11-22T18:56:06.476595Z","shell.execute_reply":"2022-11-22T18:56:06.488143Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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))","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:56:07.419143Z","iopub.execute_input":"2022-11-22T18:56:07.419942Z","iopub.status.idle":"2022-11-22T18:56:07.4257Z","shell.execute_reply.started":"2022-11-22T18:56:07.419904Z","shell.execute_reply":"2022-11-22T18:56:07.424466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:56:07.735414Z","iopub.execute_input":"2022-11-22T18:56:07.735826Z","iopub.status.idle":"2022-11-22T18:56:07.745187Z","shell.execute_reply.started":"2022-11-22T18:56:07.735796Z","shell.execute_reply":"2022-11-22T18:56:07.743638Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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)","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:56:08.026871Z","iopub.execute_input":"2022-11-22T18:56:08.027305Z","iopub.status.idle":"2022-11-22T18:56:08.035678Z","shell.execute_reply.started":"2022-11-22T18:56:08.02727Z","shell.execute_reply":"2022-11-22T18:56:08.034072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Results","metadata":{}},{"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)","metadata":{"execution":{"iopub.status.busy":"2022-11-22T18:56:08.820138Z","iopub.execute_input":"2022-11-22T18:56:08.820579Z","iopub.status.idle":"2022-11-22T18:56:09.391365Z","shell.execute_reply.started":"2022-11-22T18:56:08.820539Z","shell.execute_reply":"2022-11-22T18:56:09.389634Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for p in properties:\n    show_predictions(used_test_data[p], p)","metadata":{},"execution_count":null,"outputs":[]},{"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","metadata":{"_kg_hide-input":true,"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls","metadata":{},"execution_count":null,"outputs":[]}]}