CrabNet-element-ensemble.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "865fd6bd-ed8c-427a-ab7d-9d3d1349a3ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import collections\n",
    "from collections import OrderedDict\n",
    "import pytorch_lightning as L\n",
    "import os\n",
    "import re\n",
    "import json\n",
    "import tqdm\n",
    "\n",
    "from sklearn.metrics import balanced_accuracy_score, accuracy_score, roc_auc_score, f1_score\n",
    "from sklearn.metrics import mean_absolute_error, mean_squared_error, matthews_corrcoef\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.model_selection import StratifiedKFold, train_test_split\n",
    "\n",
    "from pytorch_lightning.loggers.csv_logs import CSVLogger\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint, StochasticWeightAveraging\n",
    "# from pytorch_lightning.loggers.wandb import WandbLogger\n",
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "from pytorch_lightning import Trainer\n",
    "from torchmetrics.functional import mean_squared_error, mean_absolute_error\n",
    "\n",
    "from pymatgen.core.composition import Composition\n",
    "from crabnet.kingcrab import CrabNet\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.optim.lr_scheduler import CyclicLR, CosineAnnealingLR, StepLR\n",
    "\n",
    "from crabnet.utils.utils import (Lamb, Lookahead, CrossEntropyLoss, \n",
    "                         masked_balanced_accuracy, compound_accuracy, masked_f1_score, masked_mcc,\n",
    "                         EDMDataset, get_edm, Scaler, DummyScaler, count_parameters)\n",
    "from crabnet.utils.get_compute_device import get_compute_device\n",
    "# from crabnet.utils.composition import _element_composition, get_sym_dict, parse_formula, CompositionError\n",
    "#from utils.optim import SWA\n",
    "\n",
    "data_type_np = np.float32\n",
    "data_type_torch = torch.float32\n",
    "\n",
    "import wandb\n",
    "\n",
    "\n",
    "class CrabNetDataModule(L.LightningDataModule):\n",
    "    def __init__(self, train_file: str , \n",
    "                 val_file: str, \n",
    "                 test_file: str,\n",
    "                 n_elements ='infer', \n",
    "                 classification = False,\n",
    "                 elem_prop='mat2vec',\n",
    "                 batch_size = 2**10,\n",
    "                 scale = True,\n",
    "                 pin_memory = True):\n",
    "        super().__init__()\n",
    "        self.train_path = train_file\n",
    "        self.val_path = val_file\n",
    "        self.test_path = test_file\n",
    "        self.batch_size = batch_size\n",
    "        self.n_elements=n_elements\n",
    "        self.pin_memory = pin_memory\n",
    "        self.scale = scale\n",
    "        self.classification = classification\n",
    "        self.elem_prop=elem_prop\n",
    "\n",
    "    def prepare_data(self):\n",
    "        ### loading and encoding trianing data\n",
    "        if(re.search('.json', self.train_path )):\n",
    "            self.data_train=pd.read_json(self.train_path)\n",
    "        elif(re.search('.csv', self.train_path)):\n",
    "            self.data_train=pd.read_csv(self.train_path)\n",
    "\n",
    "        self.train_main_data = list(get_edm(self.data_train, elem_prop=self.elem_prop,\n",
    "                                      n_elements=self.n_elements,\n",
    "                                      inference=False,\n",
    "                                      verbose=True,\n",
    "                                      drop_unary=False,\n",
    "                                      scale=self.scale))\n",
    "        \n",
    "        self.train_len_data = len(self.train_main_data[0])\n",
    "        self.train_n_elements = self.train_main_data[0].shape[1]//2\n",
    "\n",
    "        print(f'loading data with up to {self.train_n_elements:0.0f} '\n",
    "              f'elements in the formula for training')\n",
    "        \n",
    "        ### loading and encoding validation data\n",
    "        if(re.search('.json', self.val_path )):\n",
    "            self.data_val=pd.read_json(self.val_path)\n",
    "        elif(re.search('.csv', self.val_path)):\n",
    "            self.data_val=pd.read_csv(self.val_path)\n",
    "        \n",
    "        self.val_main_data = list(get_edm(self.data_val, elem_prop=self.elem_prop,\n",
    "                                      n_elements=self.n_elements,\n",
    "                                      inference=True,\n",
    "                                      verbose=True,\n",
    "                                      drop_unary=False,\n",
    "                                      scale=self.scale))\n",
    "        \n",
    "        self.val_len_data = len(self.val_main_data[0])\n",
    "        self.val_n_elements = self.val_main_data[0].shape[1]//2\n",
    "\n",
    "        print(f'loading data with up to {self.val_n_elements:0.0f} '\n",
    "              f'elements in the formula for validation')\n",
    "        \n",
    "        ### loading and encoding testing data\n",
    "        if(re.search('.json', self.test_path )):\n",
    "            self.data_test=pd.read_json(self.test_path)\n",
    "        elif(re.search('.csv', self.test_path)):\n",
    "            self.data_test=pd.read_csv(self.test_path)\n",
    "        \n",
    "        self.test_main_data = list(get_edm(self.data_test, elem_prop=self.elem_prop,\n",
    "                                      n_elements=self.n_elements,\n",
    "                                      inference=True,\n",
    "                                      verbose=True,\n",
    "                                      drop_unary=False,\n",
    "                                      scale=self.scale))\n",
    "        \n",
    "        self.test_len_data = len(self.test_main_data[0])\n",
    "        self.test_n_elements = self.test_main_data[0].shape[1]//2\n",
    "\n",
    "        print(f'loading data with up to {self.test_n_elements:0.0f} '\n",
    "              f'elements in the formula for testing')\n",
    "\n",
    "        self.train_dataset = EDMDataset(self.train_main_data, self.train_n_elements)\n",
    "        self.val_dataset = EDMDataset(self.val_main_data, self.val_n_elements)\n",
    "        self.test_dataset = EDMDataset(self.test_main_data, self.test_n_elements)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.train_dataset, batch_size=self.batch_size,\n",
    "                          pin_memory=self.pin_memory, shuffle=True)\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.val_dataset, batch_size=self.batch_size,\n",
    "                        pin_memory=self.pin_memory, shuffle=False)\n",
    "    def test_dataloader(self):\n",
    "        return DataLoader(self.test_dataset, batch_size=self.test_len_data,\n",
    "                        pin_memory=self.pin_memory, shuffle=False)\n",
    "    def predict_dataloader(self):\n",
    "        return DataLoader(self.test_dataset, batch_size=self.test_len_data,\n",
    "                        pin_memory=self.pin_memory, shuffle=False)\n",
    "\n",
    "\n",
    "class CrabNetLightning(L.LightningModule):\n",
    "    def __init__(self, **config):\n",
    "        super().__init__()\n",
    "        # Saving hyperparameters\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        self.model = CrabNet(out_dims=config['out_dims'],\n",
    "                             d_model=config['d_model'],\n",
    "                             N=config['N'],\n",
    "                             heads=config['heads'])\n",
    "        print('\\nModel architecture: out_dims, d_model, N, heads')\n",
    "        print(f'{self.model.out_dims}, {self.model.d_model}, '\n",
    "                  f'{self.model.N}, {self.model.heads}')\n",
    "        print(f'Model size: {count_parameters(self.model)} parameters\\n')\n",
    "\n",
    "        ### here we define some important parameters\n",
    "        self.fudge=config['fudge']\n",
    "        self.batch_size=config['batch_size']\n",
    "        self.classification = config['classification']\n",
    "        self.base_lr=config['base_lr']\n",
    "        self.max_lr=config['max_lr']\n",
    "        \n",
    "        # we can also try to define weights to weigh loss function, not done yet\n",
    "        self.criterion = CrossEntropyLoss\n",
    "        \n",
    "        if(re.search('.json', config['train_path'] )):\n",
    "            train_data=pd.read_json(config['train_path'])\n",
    "        elif(re.search('.csv', config['train_path'])):\n",
    "            train_data=pd.read_csv(config['train_path'])\n",
    "        \n",
    "        y=train_data['element_disorder'].values\n",
    "        self.step_size = len(y)\n",
    "\n",
    "    def forward(self, src, frac):\n",
    "        mask, output=self.model(src, frac)\n",
    "        return mask, output\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        base_optim = Lamb(params=self.model.parameters(),lr=0.001)\n",
    "        optimizer = Lookahead(base_optimizer=base_optim)\n",
    "        lr_scheduler = CyclicLR(optimizer,\n",
    "                                base_lr=self.base_lr,\n",
    "                                max_lr=self.max_lr,\n",
    "                                cycle_momentum=False,\n",
    "                                step_size_up=self.step_size)\n",
    "        # lr_scheduler=StepLR(optimizer,\n",
    "        #                     step_size=3,\n",
    "        #                     gamma=0.5)\n",
    "        return [optimizer], [lr_scheduler]\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        X, y, formula = batch\n",
    "        src, frac = X.squeeze(-1).chunk(2, dim=1)\n",
    "        frac = frac * (1 + (torch.randn_like(frac))*self.fudge)\n",
    "        frac = torch.clamp(frac, 0, 1)\n",
    "        frac[src == 0] = 0\n",
    "        frac = frac / frac.sum(dim=1).unsqueeze(1).repeat(1, frac.shape[-1])\n",
    "        \n",
    "        mask, prediction = self(src, frac)\n",
    "        loss = self.criterion(prediction,y,mask)\n",
    "        \n",
    "        self.log(\"train_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        \n",
    "        prediction = torch.nn.functional.softmax(prediction,dim=2)\n",
    "        y_pred = torch.argmax(prediction,dim=2)\n",
    "        mask = torch.mean (mask, axis=2).long()\n",
    "        y=y.detach().cpu().numpy()\n",
    "        y_pred=y_pred.detach().cpu().numpy()\n",
    "        mask=mask.detach().cpu().numpy()\n",
    "        acc=masked_balanced_accuracy(y,y_pred,mask)\n",
    "        f1=masked_f1_score(y,y_pred,mask)\n",
    "        mcc=masked_mcc(y,y_pred,mask)\n",
    "            \n",
    "        self.log(\"train_acc\", acc, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        self.log(\"train_f1\", f1, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        self.log(\"train_mcc\", mcc, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        \n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        X, y, formula = batch\n",
    "        src, frac = X.squeeze(-1).chunk(2, dim=1)\n",
    "        frac = frac * (1 + (torch.randn_like(frac))*self.fudge)\n",
    "        frac = torch.clamp(frac, 0, 1)\n",
    "        frac[src == 0] = 0\n",
    "        frac = frac / frac.sum(dim=1).unsqueeze(1).repeat(1, frac.shape[-1])\n",
    "        \n",
    "        mask, prediction = self(src, frac)\n",
    "        val_loss = self.criterion(prediction,y,mask)\n",
    "        \n",
    "        self.log(\"val_loss\", val_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        \n",
    "        prediction = torch.nn.functional.softmax(prediction,dim=2)\n",
    "        y_pred = torch.argmax(prediction,dim=2)\n",
    "        mask = torch.mean (mask,axis=2).long()\n",
    "        y=y.detach().cpu().numpy()\n",
    "        y_pred=y_pred.detach().cpu().numpy()\n",
    "        mask=mask.detach().cpu().numpy()\n",
    "        acc=masked_balanced_accuracy(y,y_pred,mask)\n",
    "        f1=masked_f1_score(y,y_pred,mask)\n",
    "        mcc=masked_mcc(y,y_pred,mask)\n",
    "            \n",
    "        self.log(\"val_acc\", acc, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        self.log(\"val_f1\", f1, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        self.log(\"val_mcc\", mcc, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        \n",
    "        return val_loss\n",
    "     \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        X, y, formula = batch\n",
    "        src, frac = X.squeeze(-1).chunk(2, dim=1)\n",
    "        frac = frac * (1 + (torch.randn_like(frac))*self.fudge)\n",
    "        frac = torch.clamp(frac, 0, 1)\n",
    "        frac[src == 0] = 0\n",
    "        frac = frac / frac.sum(dim=1).unsqueeze(1).repeat(1, frac.shape[-1])\n",
    "        \n",
    "        mask, prediction = self(src, frac)\n",
    "        \n",
    "        prediction = torch.nn.functional.softmax(prediction,dim=2)\n",
    "        y_pred = torch.argmax(prediction,dim=2)\n",
    "        mask = torch.mean (mask,axis=2).long()\n",
    "        y=y.detach().cpu().numpy()\n",
    "        y_pred=y_pred.detach().cpu().numpy()\n",
    "        mask=mask.detach().cpu().numpy()\n",
    "        acc=masked_balanced_accuracy(y,y_pred,mask)\n",
    "        f1=masked_f1_score(y,y_pred,mask)\n",
    "        mcc=masked_mcc(y,y_pred,mask)\n",
    "        \n",
    "        self.log(\"test_acc\", acc, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        self.log(\"test_f1\", f1, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "        self.log(\"test_mcc\", mcc, on_step=False, on_epoch=True, prog_bar=True, logger=True, batch_size=self.batch_size)\n",
    "            \n",
    "        return \n",
    "    \n",
    "    def predict_step(self, batch, batch_idx, dataloader_idx=0):\n",
    "        X, y, formula = batch\n",
    "        src, frac = X.squeeze(-1).chunk(2, dim=1)\n",
    "        frac = frac * (1 + (torch.randn_like(frac))*self.fudge)\n",
    "        frac = torch.clamp(frac, 0, 1)\n",
    "        frac[src == 0] = 0\n",
    "        frac = frac / frac.sum(dim=1).unsqueeze(1).repeat(1, frac.shape[-1])\n",
    "        \n",
    "        mask, prediction = self(src, frac)\n",
    "        prediction = torch.nn.functional.softmax(prediction,dim=2)\n",
    "        mask = torch.mean (mask,axis=2).long()\n",
    "        y_pred = torch.argmax(prediction,dim=2)\n",
    "        y_pred=y_pred.detach().cpu().numpy()\n",
    "        prediction=prediction.detach().cpu().numpy()\n",
    "        mask=mask.detach().cpu().numpy()\n",
    "        \n",
    "        return formula, y_pred, prediction, mask\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "960839e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 35\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "35"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('crabnet/crabnet_config.json','r') as f:\n",
    "        config=json.load(f)\n",
    "\n",
    "L.seed_everything(config['random_seed'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2ba25ece",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 82278/82278 [00:00<00:00, 209553.85formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 9143/9143 [00:00<00:00, 198592.05formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 15 elements in the formula for validation\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 10158/10158 [00:00<00:00, 199008.55formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for testing\n",
      "\n",
      "Model architecture: out_dims, d_model, N, heads\n",
      "3, 512, 4, 4\n",
      "Model size: 15139590 parameters\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\setup.py:176: PossibleUserWarning: GPU available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='gpu', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "Generating EDM: 100%|██████████| 82278/82278 [00:00<00:00, 206798.14formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 9143/9143 [00:00<00:00, 198592.05formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 15 elements in the formula for validation\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 10158/10158 [00:00<00:00, 196952.46formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for testing\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:224: PossibleUserWarning: The dataloader, predict_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eba60ce76ea0424d83356ff34324944c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Model architecture: out_dims, d_model, N, heads\n",
      "3, 512, 4, 4\n",
      "Model size: 15139590 parameters\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\setup.py:176: PossibleUserWarning: GPU available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='gpu', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "Generating EDM: 100%|██████████| 82278/82278 [00:00<00:00, 203841.60formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 9143/9143 [00:00<00:00, 207617.00formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 15 elements in the formula for validation\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 10158/10158 [00:00<00:00, 199006.69formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for testing\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:224: PossibleUserWarning: The dataloader, predict_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc6404bdb5394e1891db25d042d86cd3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Model architecture: out_dims, d_model, N, heads\n",
      "3, 512, 4, 4\n",
      "Model size: 15139590 parameters\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\setup.py:176: PossibleUserWarning: GPU available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='gpu', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "Generating EDM: 100%|██████████| 82278/82278 [00:00<00:00, 203729.69formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 9143/9143 [00:00<00:00, 192113.39formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 15 elements in the formula for validation\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 10158/10158 [00:00<00:00, 195170.59formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for testing\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:224: PossibleUserWarning: The dataloader, predict_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf125e61b56848578b311c920cbdd278",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Model architecture: out_dims, d_model, N, heads\n",
      "3, 512, 4, 4\n",
      "Model size: 15139590 parameters\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\setup.py:176: PossibleUserWarning: GPU available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='gpu', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "Generating EDM: 100%|██████████| 82278/82278 [00:00<00:00, 203008.10formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 9143/9143 [00:00<00:00, 202993.52formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 15 elements in the formula for validation\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 10158/10158 [00:00<00:00, 191232.93formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for testing\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:224: PossibleUserWarning: The dataloader, predict_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be39a870e8634f57b6c8b577b1c6e97f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Model architecture: out_dims, d_model, N, heads\n",
      "3, 512, 4, 4\n",
      "Model size: 15139590 parameters\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\setup.py:176: PossibleUserWarning: GPU available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='gpu', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "Generating EDM: 100%|██████████| 82278/82278 [00:00<00:00, 203013.35formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 9143/9143 [00:00<00:00, 190306.74formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 15 elements in the formula for validation\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating EDM: 100%|██████████| 10158/10158 [00:00<00:00, 205882.52formulae/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading data with up to 16 elements in the formula for testing\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\patykova\\.conda\\envs\\roost\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:224: PossibleUserWarning: The dataloader, predict_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "357240097bc44ac784ff1cfcc416f776",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "disorder_data = CrabNetDataModule(config['train_path'],\n",
    "                                   config['val_path'],\n",
    "                                   config['test_path'],\n",
    "                                   classification = config['classification'])\n",
    "disorder_data.prepare_data()\n",
    "predictloader=disorder_data.predict_dataloader()\n",
    "for x in predictloader:\n",
    "    _,y_true,_=x\n",
    "    datalength=len(y_true)\n",
    "y_true=y_true.detach().cpu().numpy()\n",
    "\n",
    "y_pred_crabnet=np.zeros((5,datalength,16))\n",
    "prediction_crabnet=np.zeros((5,datalength,16,3))\n",
    "\n",
    "for i in range(5):\n",
    "    name='element-disorder-'+str(i)\n",
    "    model = CrabNetLightning.load_from_checkpoint('crabnet_models/crabnet_trained_models/'+name+'.ckpt')\n",
    "    trainer = Trainer()\n",
    "    formula, y_pred, prediction, mask=trainer.predict(model=model,datamodule=disorder_data)[0]\n",
    "    y_pred_crabnet[i,:,:]=y_pred\n",
    "    prediction_crabnet[i,:,:,:]=prediction\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "cafa285a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "y_pred_av=np.zeros((datalength,16))\n",
    "\n",
    "for i in range(datalength):\n",
    "    for j in range(16):\n",
    "        labels=Counter(y_pred_crabnet[:,i,j])\n",
    "        maximum_count_number=max(labels.values())\n",
    "        most_common_labels=[]\n",
    "        for key,val in labels.items():\n",
    "            if(val==maximum_count_number):\n",
    "                most_common_labels.append(key)\n",
    "        if(len(most_common_labels)>1):\n",
    "            probs=np.mean(prediction_crabnet[:,i,j,:],axis=0)\n",
    "            y_pred_av[i,j]=np.argmax(probs)\n",
    "        else:\n",
    "            y_pred_av[i,j]=most_common_labels[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "5649a43c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7889589813362958"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "masked_balanced_accuracy(y_true,y_pred,mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "2cf579df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.7008269344359126, 0.9001472031403337, 0.5672476159158172)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compound_accuracy(y_true,y_pred,mask)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 5
}