notebooks/matbench.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 123321\n"
     ]
    }
   ],
   "source": [
    "# Set seed for everything\n",
    "from pytorch_lightning import seed_everything\n",
    "seed_everything(123321, workers=True)\n",
    "\n",
    "# Ignore warnings for clarity\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "\n",
    "import os\n",
    "import argparse\n",
    "import yaml\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.loggers import CSVLogger, WandbLogger\n",
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from utils.training import DielectricModule\n",
    "from data.dataset import StructureDataset\n",
    "from models.equivariant_model import GatedEquivariantModel\n",
    "from confidential.models import CombinedModel\n",
    "from confidential.utils import load_pfp, collate_fn_dict\n",
    "from matbench.bench import MatbenchBenchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../scripts/bench_config.yaml','r') as f:\n",
    "    configs = yaml.safe_load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-15 00:07:17 INFO     Initialized benchmark 'matbench_v0.1' with 1 tasks: \n",
      "['matbench_dielectric']\n"
     ]
    }
   ],
   "source": [
    "mb = MatbenchBenchmark(autoload=False,subset=[\"matbench_dielectric\"])\n",
    "\n",
    "train_config = configs.pop(\"Train\")\n",
    "model_config = configs.pop(\"Model\")\n",
    "outdir = train_config.pop(\"save_path\")\n",
    "\n",
    "batch_size = train_config.pop(\"batch\")\n",
    "num_workers = train_config.pop(\"num_workers\")\n",
    "pfp_layer = model_config.pop(\"pfp_layer\")\n",
    "lr = train_config.pop(\"lr\")\n",
    "train_pfp = model_config.pop(\"train_pfp\") # Bool\n",
    "\n",
    "patience = train_config.pop(\"patience\")\n",
    "epoch = train_config.pop(\"epoch\")\n",
    "accelerator = train_config.pop(\"accelerator\")\n",
    "devices = train_config.pop(\"device\")\n",
    "clip_val = train_config.pop(\"gradient_clip\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-15 00:07:17 INFO     Loading dataset 'matbench_dielectric'...\n",
      "2024-08-15 00:07:20 INFO     Dataset 'matbench_dielectric loaded.\n",
      "Start training of fold-0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 3429it [00:18, 187.29it/s]\n",
      "processing structures to build the dataset..: 382it [00:01, 206.93it/s]\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 454: 100%|██████████| 54/54 [00:14<00:00,  3.82it/s, v_num=32, val_loss=0.276, val_mae=0.276, val_rmse=1.450, train_loss=0.121, train_mae=0.121, train_rmse=0.842, lr=2.25e-6]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 953it [00:06, 158.81it/s]\n",
      "Restoring states from the checkpoint path at ../confidential/matbench/checkpoints/pl_checkpoints_fold0/eps-epoch=254-val_loss=0.27.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n",
      "Loaded model weights from the checkpoint at ../confidential/matbench/checkpoints/pl_checkpoints_fold0/eps-epoch=254-val_loss=0.27.ckpt\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicting DataLoader 0: 100%|██████████| 15/15 [00:02<00:00,  5.46it/s]\n",
      "2024-08-15 01:51:30 INFO     Recorded fold matbench_dielectric-0 successfully.\n",
      "Start training of fold-1.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 3429it [00:19, 174.84it/s]\n",
      "processing structures to build the dataset..: 382it [00:02, 130.64it/s]\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 575: 100%|██████████| 54/54 [00:07<00:00,  7.39it/s, v_num=36, val_loss=0.227, val_mae=0.227, val_rmse=0.972, train_loss=0.125, train_mae=0.125, train_rmse=0.831, lr=5.9e-7] \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 953it [00:05, 187.46it/s]\n",
      "Restoring states from the checkpoint path at ../confidential/matbench/checkpoints/pl_checkpoints_fold1/eps-epoch=375-val_loss=0.23.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n",
      "Loaded model weights from the checkpoint at ../confidential/matbench/checkpoints/pl_checkpoints_fold1/eps-epoch=375-val_loss=0.23.ckpt\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicting DataLoader 0: 100%|██████████| 15/15 [00:00<00:00, 17.83it/s]\n",
      "2024-08-15 03:21:33 INFO     Recorded fold matbench_dielectric-1 successfully.\n",
      "Start training of fold-2.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 3429it [00:19, 177.34it/s]\n",
      "processing structures to build the dataset..: 382it [00:01, 220.09it/s]\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 337: 100%|██████████| 54/54 [00:07<00:00,  7.31it/s, v_num=39, val_loss=0.245, val_mae=0.245, val_rmse=1.220, train_loss=0.0867, train_mae=0.0867, train_rmse=0.547, lr=3.52e-6]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 953it [00:05, 186.64it/s]\n",
      "Restoring states from the checkpoint path at ../confidential/matbench/checkpoints/pl_checkpoints_fold2/eps-epoch=137-val_loss=0.24.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n",
      "Loaded model weights from the checkpoint at ../confidential/matbench/checkpoints/pl_checkpoints_fold2/eps-epoch=137-val_loss=0.24.ckpt\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicting DataLoader 0: 100%|██████████| 15/15 [00:00<00:00, 17.56it/s]\n",
      "2024-08-15 04:05:10 INFO     Recorded fold matbench_dielectric-2 successfully.\n",
      "Start training of fold-3.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 3429it [00:18, 189.67it/s]\n",
      "processing structures to build the dataset..: 382it [00:02, 169.58it/s]\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 697: 100%|██████████| 54/54 [00:06<00:00,  7.79it/s, v_num=41, val_loss=0.170, val_mae=0.170, val_rmse=0.628, train_loss=0.111, train_mae=0.111, train_rmse=0.737, lr=3.02e-7]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 953it [00:05, 176.98it/s]\n",
      "Restoring states from the checkpoint path at ../confidential/matbench/checkpoints/pl_checkpoints_fold3/eps-epoch=497-val_loss=0.17.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n",
      "Loaded model weights from the checkpoint at ../confidential/matbench/checkpoints/pl_checkpoints_fold3/eps-epoch=497-val_loss=0.17.ckpt\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicting DataLoader 0: 100%|██████████| 15/15 [00:00<00:00, 19.46it/s]\n",
      "2024-08-15 05:27:43 INFO     Recorded fold matbench_dielectric-3 successfully.\n",
      "Start training of fold-4.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 3430it [00:18, 189.40it/s]\n",
      "processing structures to build the dataset..: 382it [00:02, 180.47it/s]\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 379: 100%|██████████| 54/54 [00:06<00:00,  7.91it/s, v_num=42, val_loss=0.240, val_mae=0.240, val_rmse=0.919, train_loss=0.116, train_mae=0.116, train_rmse=0.753, lr=2.25e-6]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing structures to build the dataset..: 952it [00:04, 209.40it/s]\n",
      "Restoring states from the checkpoint path at ../confidential/matbench/checkpoints/pl_checkpoints_fold4/eps-epoch=179-val_loss=0.24.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [6,5,7,2,4,3,1,0]\n",
      "Loaded model weights from the checkpoint at ../confidential/matbench/checkpoints/pl_checkpoints_fold4/eps-epoch=179-val_loss=0.24.ckpt\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicting DataLoader 0: 100%|██████████| 15/15 [00:00<00:00, 19.02it/s]\n",
      "2024-08-15 06:12:27 INFO     Recorded fold matbench_dielectric-4 successfully.\n",
      "2024-08-15 06:12:27 INFO     Successfully wrote MatbenchBenchmark to file 'dielectric_matbench1.json.gz'.\n"
     ]
    }
   ],
   "source": [
    "for task in mb.tasks:\n",
    "    task.load()\n",
    "    for fold in task.folds:\n",
    "        print(f\"Start training of fold-{fold}.\")\n",
    "\n",
    "        os.makedirs(outdir + str(fold), exist_ok=True)\n",
    "        train_inputs, train_outputs = task.get_train_and_val_data(fold)        \n",
    "\n",
    "        x_train, x_val, y_train, y_val, k_train, k_val = train_test_split(\n",
    "            train_inputs.tolist(), train_outputs.tolist(), train_inputs.index.tolist(), test_size=0.1, random_state=123321)\n",
    "\n",
    "        train_dataset = StructureDataset(x_train, y_train, k_train)\n",
    "        val_dataset = StructureDataset(x_val, y_val, k_val)\n",
    "\n",
    "        train_loader = DataLoader(\n",
    "            train_dataset, batch_size=batch_size, collate_fn=collate_fn_dict, \n",
    "            shuffle=True, num_workers=num_workers, pin_memory=True\n",
    "        )\n",
    "        val_loader = DataLoader(\n",
    "            val_dataset, batch_size=batch_size, collate_fn=collate_fn_dict,\n",
    "            shuffle=False, num_workers=num_workers, pin_memory=True\n",
    "        )\n",
    "\n",
    "        # load the pre-trained PFP model\n",
    "        pfp_wrapped = load_pfp(\n",
    "            load_parameters=True, \n",
    "            return_layer=pfp_layer)\n",
    "\n",
    "        # Freeze PFP parameters\n",
    "        if train_pfp:\n",
    "            pfp_wrapped.pfp.reset_parameters(4.0)\n",
    "        else:\n",
    "            for param in pfp_wrapped.parameters():\n",
    "                param.requires_grad = False\n",
    "\n",
    "        # Build Readout NN\n",
    "        tensorial_model = GatedEquivariantModel(**model_config)\n",
    "        # Build a combined model to connect readout NN to PFP\n",
    "        model = CombinedModel(pfp_wrapped, tensorial_model)\n",
    "\n",
    "        pl_module = DielectricModule(\n",
    "            model, \n",
    "            learning_rate=lr)\n",
    "\n",
    "        project_name = f'matbench_layer{pfp_layer}_fold{fold}'\n",
    "        # wandb_logger = WandbLogger(project=project_name)\n",
    "        csv_logger = CSVLogger(\"logs\")\n",
    "\n",
    "        checkpoint_callback = ModelCheckpoint(\n",
    "            save_top_k=1,\n",
    "            monitor=\"val_loss\",\n",
    "            mode=\"min\",\n",
    "            dirpath=f\"{outdir}/pl_checkpoints_fold{fold}/\",\n",
    "            filename=\"eps-{epoch:02d}-{val_loss:.2f}\",\n",
    "        )\n",
    "\n",
    "        earlystopping_callback = EarlyStopping(\"val_loss\", mode=\"min\", \n",
    "            patience=patience)\n",
    "\n",
    "        trainer = pl.Trainer(\n",
    "            max_epochs=epoch, \n",
    "            accelerator=accelerator, \n",
    "            logger=csv_logger,  \n",
    "            devices=devices, \n",
    "            callbacks=[checkpoint_callback, earlystopping_callback],\n",
    "            gradient_clip_val=clip_val,\n",
    "            enable_model_summary=False)\n",
    "        trainer.fit(\n",
    "            model=pl_module, \n",
    "            train_dataloaders=train_loader, \n",
    "            val_dataloaders=val_loader)\n",
    "\n",
    "        test_inputs, test_outputs = task.get_test_data(fold, include_target=True)\n",
    "        test_dataset = StructureDataset(test_inputs.tolist(), test_outputs.tolist(), test_inputs.index.tolist())\n",
    "        test_loader = DataLoader(\n",
    "            test_dataset, batch_size=batch_size, collate_fn=collate_fn_dict,\n",
    "            shuffle=False, num_workers=num_workers, pin_memory=True\n",
    "        )\n",
    "\n",
    "        # trainer.test(pl_module, test_loader, ckpt_path='best')\n",
    "        preds = trainer.predict(pl_module, test_loader, ckpt_path='best')\n",
    "        preds = torch.cat(preds, dim=0).detach().cpu().numpy()\n",
    "        preds = np.mean(np.diagonal(preds, axis1=1, axis2=2), axis=1)\n",
    "\n",
    "        # Record your data!\n",
    "        task.record(fold, preds)\n",
    "\n",
    "# Save your results\n",
    "mb.to_file(\"dielectric_matbench1.json.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'mae': {'mean': 0.2370850306791295,\n",
       "  'max': 0.3516925340028822,\n",
       "  'min': 0.11753604419222681,\n",
       "  'std': 0.07649557856878016},\n",
       " 'rmse': {'mean': 1.6829936226656266,\n",
       "  'max': 2.901275849331958,\n",
       "  'min': 0.5732083587460219,\n",
       "  'std': 0.8485401319998302},\n",
       " 'mape': {'mean': 0.05530619144780041,\n",
       "  'max': 0.06721239289649394,\n",
       "  'min': 0.03617504271599022,\n",
       "  'std': 0.013170799705305278},\n",
       " 'max_error': {'mean': 34.59824110916204,\n",
       "  'max': 58.784419044481886,\n",
       "  'min': 13.862970454340712,\n",
       "  'std': 18.27463111793412},\n",
       " '_ipython_canary_method_should_not_exist_': {}}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scores = mb.matbench_dielectric.scores\n",
    "scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}