{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Chemical Shielding\n\nCORY M. WIDDIFIELD, ROBERT W. SCHURKO.\n\"Understanding Chemical Shielding Tensors Using Group Theory, MO Analysis, and Modern Density- Functional Theory\"\n*Concepts in Magnetic Resonance Part A (Bridging Education and Research). (2009)\nhttps://onlinelibrary.wiley.com/doi/pdf/10.1002/cmr.a.20136\n\nThe chemical shielding Hamiltonian:\n\n![Hamiltonian](https://user-images.githubusercontent.com/11532812/60358602-78018d80-9a11-11e9-894f-60b8122e8a4c.png)\n\nwhere gamma is the gyromagnetic ratio, I^ is the nuclear spin operator, and sigma is the magnetic shielding tensor."},{"metadata":{"_kg_hide-output":true,"trusted":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},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nmolecules = pd.read_csv('../input/structures.csv')\nmolecules = molecules.groupby('molecule_name')\nmagnetic_shielding_tensors = pd.read_csv('../input/magnetic_shielding_tensors.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"magnetic_shielding_tensors.head(10)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The magnetic shielding tensor:\n![Magnetic shielding tensor](https://user-images.githubusercontent.com/11532812/60358839-3ae9cb00-9a12-11e9-8f3a-0236f7bc072d.png)\n\nThe symmetric portion of the shielding tensor can be diagonalized into its own principal axis system (PAS):\n\n![PAS](https://user-images.githubusercontent.com/11532812/60359035-c95e4c80-9a12-11e9-897d-2e422e0f2e83.png)"},{"metadata":{"trusted":true},"cell_type":"code","source":"x = magnetic_shielding_tensors.columns.values[2:]\nx = magnetic_shielding_tensors[x].values\nx = x.reshape(-1,3,3)\nx = x + np.transpose(x,(0,2,1))\nx = 0.5 * x\nw, v = np.linalg.eigh(x)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The isotropic shielding value is defined as:\n\n![Isotropic shielding tensor](https://user-images.githubusercontent.com/11532812/60358900-75536800-9a12-11e9-850c-654dc81fc258.png)"},{"metadata":{"trusted":true},"cell_type":"code","source":"sigma_iso = np.sum(w, axis=1)/3 ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"![span](https://user-images.githubusercontent.com/11532812/60359189-45589480-9a13-11e9-9016-42f8a90c1ddd.png)\n\nThs span (Omega) describes the magnitude of the shielding anisotropy."},{"metadata":{"trusted":true},"cell_type":"code","source":"omega = w[:,2] - w[:,0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"![skew](https://user-images.githubusercontent.com/11532812/60359338-c57efa00-9a13-11e9-9619-17ebd4f44ac8.png)\n\nThe skew (kappa) describes degree of axial symmetry of the shielding tensor."},{"metadata":{"trusted":true},"cell_type":"code","source":"kappa = 3 * (sigma_iso - w[:,1])/omega","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"magnetic_shielding_parameters = magnetic_shielding_tensors[magnetic_shielding_tensors.columns.values[:2]]\nmagnetic_shielding_parameters = pd.DataFrame(magnetic_shielding_parameters)\nmagnetic_shielding_parameters[\"sigma_iso\"] = sigma_iso\nmagnetic_shielding_parameters[\"omega\"] = omega\nmagnetic_shielding_parameters[\"kappa\"] = kappa\n\nmagnetic_shielding_parameters.head(10)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_ = magnetic_shielding_parameters.sigma_iso.hist(bins=100)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_ = magnetic_shielding_parameters.omega.hist(bins=100)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_ = magnetic_shielding_parameters.kappa.hist(bins=100)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Please see LibreTexts [Chemical Shift][1] if you want to understand the meaning of the parameters.\n\n[1]: https://chem.libretexts.org/Bookshelves/Physical_and_Theoretical_Chemistry_Textbook_Maps/Supplemental_Modules_(Physical_and_Theoretical_Chemistry)/Spectroscopy/Magnetic_Resonance_Spectroscopies/Nuclear_Magnetic_Resonance/NMR%3A_Theory/NMR_Interactions/Chemical_Shift_(Shielding)"},{"metadata":{},"cell_type":"markdown","source":"# ASE Database"},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv('../input/train.csv')\ntrain_molecule_names = train.molecule_name.unique()\n\nmsp = magnetic_shielding_parameters.groupby('molecule_name')","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                mol_msp = msp.get_group(name)\n                sigma_iso = mol_msp['sigma_iso'].values.reshape(-1,1)\n                omega = mol_msp['omega'].values.reshape(-1,1)\n                kappa = mol_msp['kappa'].values.reshape(-1,1)\n            except KeyError:\n                sigma_iso, omega, kappa = [None] * 3\n            db.write(atoms, name=name,\n                     data=dict(sigma_iso=sigma_iso, omega=omega, kappa=kappa)\n                    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"champs_path = 'CHAMPS_train.db'\ndataset_size = len(train_molecule_names) # 40000\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\ndataset = schnetpack.data.AtomsData(champs_path,\n                properties=['sigma_iso', 'omega', 'kappa'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(dataset)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Shielding Parameter Prediction\n\n## 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":{"_kg_hide-input":true,"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":"from schnetpack.data import Structure\n\nclass MagneticShielding(atm.Atomwise):\n    def __init__(self, property):\n        super(MagneticShielding, self).__init__(return_contributions=True)\n        self.property = property\n        \n    def forward(self, inputs):\n        result = super().forward(inputs)\n        \n        atom_mask = inputs[Structure.atom_mask].byte()\n        \n        yi = inputs[self.property]\n        yi = torch.masked_select(yi.squeeze(dim=2), atom_mask)\n        inputs[self.property+'_true'] = yi\n        \n        yi = result['yi']\n        yi = torch.masked_select(yi.squeeze(dim=2), atom_mask)\n        result[self.property+'_pred'] = yi\n        \n        return result","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def schnet_model(property):\n    reps = rep.SchNet(n_interactions=6)\n    output = MagneticShielding(property=property)\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(property, max_epochs=500):\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)\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    target_key = property+'_true'\n    output_key = property+'_pred'\n    opt = Adam(model.parameters(), lr=1e-4)\n    loss = lambda b, p: F.mse_loss(p[output_key], b[target_key])\n    metrics = [\n        spk.metrics.MeanAbsoluteError(target_key, output_key, name='MAE_'+property),\n        spk.metrics.RootMeanSquaredError(target_key, output_key, name='RMSE_'+property),\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 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 shielding parameters\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+'_true'].tolist()\n            predictions += result[property+'_pred'].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 ['sigma_iso', 'omega', 'kappa']:\n    print(p)\n    used_test_data[p] = train_model(p, max_epochs=50)\n    show_history(p)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for p in ['sigma_iso', 'omega', 'kappa']:\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 kappa/log/log.csv log_kappa.csv\n!mv omega/log/log.csv log_omega.csv\n!mv sigma_iso/log/log.csv log_sigma_iso.csv\n!mv kappa/output/best_model best_model_kappa\n!mv omega/output/best_model best_model_omega\n!mv sigma_iso/output/best_model best_model_sigma_iso\n!rm -r kappa\n!rm -r omega\n!rm -r sigma_iso","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}