examples/threedgraph/threedgraph.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "## Trained model and optimal hyperparameters\n",
    "\n",
    "We have provided trained model and hyperparameters of SphereNet on QM9 and MD17 [here](https://github.com/divelab/DIG_storage/tree/main/3dgraph)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example of 3D Graph\n",
    "\n",
    "Here we provide the example code for SphereNet on QM93D and MD17 datasets. You can easily replace SphereNet with SchNet and DimeNetPP by changing model name and model parameters."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "import torch\n",
    "import sys\n",
    "sys.path.insert(0,'..')\n",
    "sys.path.insert(0,'../..')\n",
    "from dig.threedgraph.dataset import QM93D\n",
    "from dig.threedgraph.dataset import MD17\n",
    "from dig.threedgraph.method import SphereNet #SchNet, DimeNetPP\n",
    "from dig.threedgraph.method import run\n",
    "from dig.threedgraph.evaluation import ThreeDEvaluator"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "device"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "metadata": {},
     "execution_count": 3
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Example code for QM93D data\n",
    "***Note***: '3D' means that the data includes positional information for atoms.\n",
    "\n",
    "We trained a separate model for each target except for _gap_, which was predicted by taking _lumo-homo_. You can use default hyperparameters to get comparable results, we also tuned hyperparameters like lr, lr_decay_factor, lr_decay_step_size, cutoff, num_spherical, num_radial, basis_emb_size_dist, basis_emb_size_angle, basis_emb_size_torsion to achieve better performance. The values/search space for hyperparameters are listed in the Appendix of our paper.\n",
    "\n",
    "The default hyperparameters for QM93D are:  \n",
    "    &ensp; energy_and_force=False, cutoff=5.0, num_layers=4, hidden_channels=128, out_channels=1, int_emb_size=64,  \n",
    "    &ensp; basis_emb_size_dist=8, basis_emb_size_angle=8, basis_emb_size_torsion=8, out_emb_channels=256,  \n",
    "    &ensp; num_spherical=3, num_radial=6, envelope_exponent=5,  \n",
    "    &ensp; num_before_skip=1, num_after_skip=2, num_output_layers=3,  \n",
    "    &ensp; epochs=1000, batch_size=32, vt_batch_size=32, lr=0.0005, lr_decay_factor=0.5, lr_decay_step_size=100.\n",
    "\n",
    "\n"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Loading dataset"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "dataset = QM93D(root='dataset/')\n",
    "target = 'U0'\n",
    "dataset.data.y = dataset.data[target]\n",
    "\n",
    "split_idx = dataset.get_idx_split(len(dataset.data.y), train_size=110000, valid_size=10000, seed=42)\n",
    "\n",
    "train_dataset, valid_dataset, test_dataset = dataset[split_idx['train']], dataset[split_idx['valid']], dataset[split_idx['test']]\n",
    "print('train, validaion, test:', len(train_dataset), len(valid_dataset), len(test_dataset))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "train, validaion, test: 110000 10000 10831\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Loading model, loss and evaluation function\n",
    "\n",
    "The evaluation metric is mean absolute error (MAE)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "model = SphereNet(energy_and_force=False, cutoff=5.0, num_layers=4, \n",
    "        hidden_channels=128, out_channels=1, int_emb_size=64, \n",
    "        basis_emb_size_dist=8, basis_emb_size_angle=8, basis_emb_size_torsion=8, out_emb_channels=256, \n",
    "        num_spherical=3, num_radial=6, envelope_exponent=5, \n",
    "        num_before_skip=1, num_after_skip=2, num_output_layers=3, use_node_features=True\n",
    "        )\n",
    "loss_func = torch.nn.L1Loss()\n",
    "evaluation = ThreeDEvaluator()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Training"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "run3d = run()\n",
    "run3d.run(device, train_dataset, valid_dataset, test_dataset, model, loss_func, evaluation, epochs=20, batch_size=32, vt_batch_size=32, lr=0.0005, lr_decay_factor=0.5, lr_decay_step_size=15)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#Params: 1890118\n",
      "\n",
      "=====Epoch 1\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 3438/3438 [07:22<00:00,  7.77it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:20<00:00, 15.30it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:21<00:00, 16.04it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.8305539944409076, 'Validation': 0.7885677814483643, 'Test': 0.7943109273910522}\n",
      "\n",
      "=====Epoch 2\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [06:16<00:00,  9.12it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 28.91it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:12<00:00, 27.32it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.3417653005923415, 'Validation': 0.16290859878063202, 'Test': 0.16250823438167572}\n",
      "\n",
      "=====Epoch 3\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:24<00:00, 13.01it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.31it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:12<00:00, 27.83it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.2626579807482881, 'Validation': 0.10924234241247177, 'Test': 0.1091669574379921}\n",
      "\n",
      "=====Epoch 4\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:26<00:00, 12.88it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.67it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 28.65it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.2185871605092249, 'Validation': 0.1412947177886963, 'Test': 0.14113298058509827}\n",
      "\n",
      "=====Epoch 5\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:25<00:00, 12.97it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.38it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 29.38it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.18415136586759867, 'Validation': 0.08948442339897156, 'Test': 0.08791808038949966}\n",
      "\n",
      "=====Epoch 6\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:24<00:00, 13.00it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.43it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 29.58it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.17059671088246983, 'Validation': 0.10857655853033066, 'Test': 0.1086759939789772}\n",
      "\n",
      "=====Epoch 7\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:30<00:00, 12.69it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 28.61it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 29.14it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.15622219235277093, 'Validation': 0.08192159235477448, 'Test': 0.08170071989297867}\n",
      "\n",
      "=====Epoch 8\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:35<00:00, 12.48it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.14it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 28.63it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.1442768630192958, 'Validation': 0.08120342344045639, 'Test': 0.08138693124055862}\n",
      "\n",
      "=====Epoch 9\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:24<00:00, 13.00it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 28.65it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 28.29it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.13906806218478485, 'Validation': 0.07339970022439957, 'Test': 0.0732196718454361}\n",
      "\n",
      "=====Epoch 10\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:35<00:00, 12.49it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:11<00:00, 27.44it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:12<00:00, 26.88it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.12617339688792625, 'Validation': 0.11456501483917236, 'Test': 0.11438193917274475}\n",
      "\n",
      "=====Epoch 11\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:27<00:00, 12.85it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 28.90it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:12<00:00, 26.23it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.12321726725571651, 'Validation': 0.0715189278125763, 'Test': 0.07092428207397461}\n",
      "\n",
      "=====Epoch 12\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:31<00:00, 12.68it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.38it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 29.33it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.11304465457233598, 'Validation': 0.1164650246500969, 'Test': 0.11696784943342209}\n",
      "\n",
      "=====Epoch 13\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:28<00:00, 12.80it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 28.54it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 29.32it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.11311055924429181, 'Validation': 0.1142609491944313, 'Test': 0.11372711509466171}\n",
      "\n",
      "=====Epoch 14\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:25<00:00, 12.97it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.43it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:12<00:00, 27.68it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.1103381712277869, 'Validation': 0.05894898623228073, 'Test': 0.05792950466275215}\n",
      "\n",
      "=====Epoch 15\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:30<00:00, 12.69it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.26it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 28.55it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.09813584842398945, 'Validation': 0.13913576304912567, 'Test': 0.1383834183216095}\n",
      "\n",
      "=====Epoch 16\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:26<00:00, 12.89it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.29it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 28.72it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.05428033658000465, 'Validation': 0.06030373275279999, 'Test': 0.059175316244363785}\n",
      "\n",
      "=====Epoch 17\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:28<00:00, 12.80it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:11<00:00, 27.83it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:12<00:00, 27.47it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.054203004988561614, 'Validation': 0.03810606151819229, 'Test': 0.03703922778367996}\n",
      "\n",
      "=====Epoch 18\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:29<00:00, 12.77it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:11<00:00, 28.11it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:12<00:00, 26.31it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.0530719623151666, 'Validation': 0.04359658062458038, 'Test': 0.043418560177087784}\n",
      "\n",
      "=====Epoch 19\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:26<00:00, 12.89it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 28.87it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:12<00:00, 28.13it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.05202796294149651, 'Validation': 0.04247582331299782, 'Test': 0.04204947501420975}\n",
      "\n",
      "=====Epoch 20\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3438/3438 [04:31<00:00, 12.64it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 313/313 [00:10<00:00, 29.59it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 339/339 [00:11<00:00, 29.71it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "{'Train': 0.04962607438894397, 'Validation': 0.04090351238846779, 'Test': 0.040894996374845505}\n",
      "Best validation MAE so far: 0.03810606151819229\n",
      "Test MAE when got best validation result: 0.03703922778367996\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "\n",
    "### Example code for MD17 data\n",
    "\n",
    "We trained a separate model for each molecule. You can use default hyperparameters to get comparable results, we also tuned hyperparameters like lr, lr_decay_factor, lr_decay_step_size, batch_size, basis_emb_size_dist, basis_emb_size_angle, basis_emb_size_torsion to achieve better performance. The values/search space for hyperparameters are listed in the Appendix of our paper.\n",
    "\n",
    "The default hyperparameters for MD17 are:  \n",
    "    &ensp; energy_and_force=True, cutoff=5.0, num_layers=4, hidden_channels=128, out_channels=1, int_emb_size=64,  \n",
    "    &ensp; basis_emb_size_dist=8, basis_emb_size_angle=8, basis_emb_size_torsion=8, out_emb_channels=256,  \n",
    "    &ensp; num_spherical=3, num_radial=6, envelope_exponent=5,  \n",
    "    &ensp; num_before_skip=1, num_after_skip=2, num_output_layers=3,  \n",
    "    &ensp; epochs=1000, batch_size=1, vt_batch_size=32, lr=0.0005, lr_decay_factor=0.5, lr_decay_step_size=200."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Loading dataset"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "dataset_md17 = MD17(root='dataset/', name='aspirin')\n",
    "\n",
    "split_idx_md17 = dataset_md17.get_idx_split(len(dataset_md17.data.y), train_size=1000, valid_size=1000, seed=42)\n",
    "\n",
    "train_dataset_md17, valid_dataset_md17, test_dataset_md17 = dataset_md17[split_idx_md17['train']], dataset_md17[split_idx_md17['valid']], dataset_md17[split_idx_md17['test']]\n",
    "print('train, validaion, test:', len(train_dataset_md17), len(valid_dataset_md17), len(test_dataset_md17))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "train, validaion, test: 1000 1000 209762\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Loading model, loss and evaluation function\n",
    "\n",
    "We predict energy and take the negative of the derivative of the energy with respect to the atomic positions as predicted forces.\n",
    "\n",
    "The evaluation metric is mean absolute error (MAE)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "model_md17 = SphereNet(energy_and_force=True, cutoff=5.0, num_layers=4, \n",
    "        hidden_channels=128, out_channels=1, int_emb_size=64, \n",
    "        basis_emb_size_dist=8, basis_emb_size_angle=8, basis_emb_size_torsion=8, out_emb_channels=256, \n",
    "        num_spherical=3, num_radial=6, envelope_exponent=5, \n",
    "        num_before_skip=1, num_after_skip=2, num_output_layers=3 \n",
    "        )\n",
    "loss_func_md17 = torch.nn.L1Loss()\n",
    "evaluation_md17 = ThreeDEvaluator()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Training"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "run3d_md17 = run()\n",
    "run3d_md17.run(device, train_dataset_md17, valid_dataset_md17, test_dataset_md17, model_md17, loss_func_md17, evaluation_md17, epochs=5, batch_size=1, vt_batch_size=64, lr=0.0005, lr_decay_factor=0.5, lr_decay_step_size=200, energy_and_force=True)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "#Params: 1890118\n",
      "\n",
      "=====Epoch 1\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "/data1/limei/anaconda3/envs/graph/lib/python3.8/site-packages/torch_geometric/deprecation.py:13: UserWarning: 'data.DataLoader' is deprecated, use 'loader.DataLoader' instead\n",
      "  warnings.warn(out)\n",
      "100%|██████████| 1000/1000 [03:20<00:00,  4.99it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 16/16 [00:03<00:00,  5.06it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 21076.162109375, 'Force MAE': 88.1651611328125}\n",
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3278/3278 [10:49<00:00,  5.05it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 21079.103515625, 'Force MAE': 87.94914245605469}\n",
      "\n",
      "{'Train': 62999.26555371094, 'Validation': 29892.67822265625, 'Test': 29874.01776123047}\n",
      "\n",
      "=====Epoch 2\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 1000/1000 [03:35<00:00,  4.65it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 16/16 [00:03<00:00,  5.09it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 11752.1220703125, 'Force MAE': 41.66204833984375}\n",
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3278/3278 [10:27<00:00,  5.23it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 11752.3837890625, 'Force MAE': 41.8145866394043}\n",
      "\n",
      "{'Train': 21397.7178671875, 'Validation': 15918.326904296875, 'Test': 15933.84245300293}\n",
      "\n",
      "=====Epoch 3\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 1000/1000 [03:44<00:00,  4.46it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 16/16 [00:03<00:00,  5.15it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 12009.2421875, 'Force MAE': 70.98027038574219}\n",
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3278/3278 [10:42<00:00,  5.10it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 12010.466796875, 'Force MAE': 71.2234115600586}\n",
      "\n",
      "{'Train': 14609.853533203124, 'Validation': 19107.26922607422, 'Test': 19132.80795288086}\n",
      "\n",
      "=====Epoch 4\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 1000/1000 [03:46<00:00,  4.42it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 16/16 [00:03<00:00,  5.11it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 14571.55078125, 'Force MAE': 49.371952056884766}\n",
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3278/3278 [10:43<00:00,  5.09it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 14570.435546875, 'Force MAE': 49.68278121948242}\n",
      "\n",
      "{'Train': 13809.444167358399, 'Validation': 19508.745986938477, 'Test': 19538.713668823242}\n",
      "\n",
      "=====Epoch 5\n",
      "\n",
      "Training...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 1000/1000 [03:43<00:00,  4.47it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "Evaluating...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 16/16 [00:03<00:00,  5.20it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 19717.771484375, 'Force MAE': 30.181947708129883}\n",
      "\n",
      "\n",
      "Testing...\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n",
      "100%|██████████| 3278/3278 [10:41<00:00,  5.11it/s]"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'Energy MAE': 19717.19140625, 'Force MAE': 30.4222354888916}\n",
      "\n",
      "{'Train': 12309.488090698242, 'Validation': 22735.96625518799, 'Test': 22759.41495513916}\n",
      "Best validation MAE so far: 15918.326904296875\n",
      "Test MAE when got best validation result: 15933.84245300293\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "\n"
     ]
    }
   ],
   "metadata": {
    "scrolled": true,
    "tags": []
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d38ceeef29fa782c7d8487c611f29fc45ed14494aabe491c99bf90b326b1eb2a"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.12 64-bit ('graph': conda)"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "metadata": {
   "interpreter": {
    "hash": "015e0eafcab22e7fddf6e5cdb9e5f6d4189ebbdd188945dfc12bd444354eda05"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}