notebooks/dielectric_tensor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This notebook aims to showcase the application of our model for predicting equivariant dielectric tensors of inorganic materials.\n",
    "\n",
    "PFP (see details at https://tech.preferred.jp/en/blog/development-of-universal-neural-network-for-materials-discovery/) is a neural network potential model that incorporates scalar, equivariant vector, and tensor features in each node. We demonstrate that the equivariant features learned in the pre-trained PFP can be effectively utilized to predict other tensorial properties, yielding promising accuracy even with limited available data. In this notebook, we utilize the dielectric constants (~6.6k data) extracted from the Materials Project (https://next-gen.materialsproject.org/) as an illustrative example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Model training\n",
    "\n",
    "A specifically-designed equivariant readout module is utilized for the dielectric tensor. This module takes into account the scalar, vector, and tensor features of nodes, as well as the scalar and vector features of edges, which are extracted from the pre-trained PFP. By combining and processing these inputs, the module predicts the equivariant 3 by 3 dielectric constants."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/nfs-mnj-archive-02/user/pe_zmao/micromamba/envs/pfp/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "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 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",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning) \n",
    "\n",
    "from utils.training import DielectricModule\n",
    "from data.dataset import JSONDataset\n",
    "from models.equivariant_model import GatedEquivariantModel\n",
    "from confidential.models import CombinedModel\n",
    "from confidential.utils import load_pfp, collate_fn_dict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the configuration for model training. This configuration is used to train a model for the electronic contribution of dielectric tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Seed': 3,\n",
       " 'Train': {'target': 'electronic',\n",
       "  'num_workers': 12,\n",
       "  'dataset': '../data/mp_dielectric.json',\n",
       "  'batch': 64,\n",
       "  'epoch': 5,\n",
       "  'patience': 200,\n",
       "  'lr': 0.0001,\n",
       "  'accelerator': 'gpu',\n",
       "  'device': [3],\n",
       "  'save_path': '../confidential/checkpoints/',\n",
       "  'gradient_clip': 2.0},\n",
       " 'Model': {'pfp_layer': 3,\n",
       "  'train_pfp': False,\n",
       "  'latent_feat': 64,\n",
       "  'n_gate_layers': 2,\n",
       "  'dropout_rate': 0.0,\n",
       "  'residual': True,\n",
       "  'gate_sigmoid': True,\n",
       "  'mlp_layer': 3,\n",
       "  'integrate_es_ev': True,\n",
       "  'integrate_nv_nt': True,\n",
       "  'apply_mask': False}}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('../scripts/train_config.yaml','r') as f:\n",
    "    configs = yaml.safe_load(f)\n",
    "\n",
    "# dimensionality of each feature of PFP\n",
    "ns_feat = configs['Model'].pop('ns_feat')   # node scalar feature\n",
    "nv_feat = configs['Model'].pop('nv_feat')   # node vector feature\n",
    "nt_feat = configs['Model'].pop('nt_feat')   # node tensor feature\n",
    "es_feat = configs['Model'].pop('es_feat')   # edge scalar feature\n",
    "ev_feat = configs['Model'].pop('ev_feat')   # edge vector feature\n",
    "\n",
    "configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_config = configs.pop(\"Train\")\n",
    "model_config = configs.pop(\"Model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the random seed for code reproducibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fa8cd646250>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seed = configs.pop(\"Seed\")\n",
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare data for model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6648 data are prepared for training!\n"
     ]
    }
   ],
   "source": [
    "dataset_file = train_config.pop(\"dataset\")\n",
    "target = train_config.pop(\"target\")\n",
    "dataset = JSONDataset(json_file = dataset_file,target = target)\n",
    "\n",
    "keys = sorted(dataset.keys)\n",
    "print(f\"{len(keys)} data are prepared for training!\")\n",
    "\n",
    "# split dataset\n",
    "train_keys, test_keys = train_test_split(keys, test_size=0.1, random_state=seed)\n",
    "train_keys, val_keys = train_test_split(train_keys, test_size=0.1/0.9, random_state=seed)\n",
    "\n",
    "# prepare dataloader\n",
    "train_set = JSONDataset(dataset_file, target = target, keys = train_keys)\n",
    "val_set = JSONDataset(dataset_file, target = target, keys=val_keys)\n",
    "test_set = JSONDataset(dataset_file, target = target, keys=test_keys)\n",
    "\n",
    "batch_size = train_config.pop(\"batch\")\n",
    "num_workers = train_config.pop(\"num_workers\")\n",
    "train_loader = DataLoader(train_set, batch_size=batch_size, collate_fn=collate_fn_dict, shuffle=True, num_workers=num_workers, pin_memory=True)\n",
    "val_loader = DataLoader(val_set, batch_size=batch_size, collate_fn=collate_fn_dict, shuffle=False, num_workers=num_workers, pin_memory=True)\n",
    "test_loader = DataLoader(test_set, batch_size=batch_size, collate_fn=collate_fn_dict, shuffle=False, num_workers=num_workers, pin_memory=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the pre-trained PFP model and freeze its parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pfp_layer = model_config.pop(\"pfp_layer\")\n",
    "pfp_wrapped = load_pfp(load_parameters=True, return_layer=pfp_layer)\n",
    "\n",
    "train_pfp = model_config.pop(\"train_pfp\")\n",
    "if train_pfp:\n",
    "    # Initalize the parameters if PFP is also trained\n",
    "    pfp_wrapped.pfp.reset_parameters()\n",
    "else:\n",
    "    # Otherwise, freeze parameters in PFP\n",
    "    for param in pfp_wrapped.parameters():\n",
    "        param.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initalize the equivariant readout model and combine it with the pre-trained PFP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorial_model = GatedEquivariantModel(\n",
    "    ns_feat=ns_feat,\n",
    "    nv_feat=nv_feat,\n",
    "    nt_feat=nt_feat,\n",
    "    es_feat=es_feat,\n",
    "    ev_feat=ev_feat,\n",
    "    **model_config)\n",
    "model = CombinedModel(pfp_wrapped, tensorial_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the Trainer in the pytorch lightning package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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",
      "/mnt/nfs-mnj-archive-02/user/pe_zmao/micromamba/envs/pfp/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py:67: UserWarning: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `pytorch_lightning` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default\n",
      "  warning_cache.warn(\n"
     ]
    }
   ],
   "source": [
    "lr = train_config.pop(\"lr\")\n",
    "pl_module = DielectricModule(model, learning_rate=lr, optimizer='adam')\n",
    "outdir = train_config.pop(\"save_path\")\n",
    "\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    save_top_k=1,\n",
    "    monitor=\"val_loss\",\n",
    "    mode=\"min\",\n",
    "    dirpath=f\"{outdir}/pl_checkpoints/\",\n",
    "    filename=\"eps-{epoch:02d}-{val_loss:.2f}\",\n",
    ")\n",
    "\n",
    "patience = train_config.pop(\"patience\")\n",
    "earlystopping_callback = EarlyStopping(\"val_loss\", mode=\"min\", patience=patience)\n",
    "\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\")\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=epoch, \n",
    "    accelerator=accelerator, \n",
    "    devices=devices, \n",
    "    callbacks=[checkpoint_callback, earlystopping_callback],\n",
    "    gradient_clip_val=clip_val,\n",
    "    enable_model_summary=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train the model. For showcase purposes, we only train for 5 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [4,3,0,7,5,6,2,1]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4: 100%|██████████| 84/84 [00:07<00:00, 10.95it/s, v_num=0, val_loss=25.40, val_trace_loss=6.870, val_diag_loss=25.20, val_off_loss=0.075, val_tensor_loss=25.40, train_loss=39.70, train_trace_loss=10.60, train_diag_loss=39.10, train_off_loss=0.303, train_tensor_loss=39.70, lr=0.0001] "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=5` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4: 100%|██████████| 84/84 [00:07<00:00, 10.66it/s, v_num=0, val_loss=25.40, val_trace_loss=6.870, val_diag_loss=25.20, val_off_loss=0.075, val_tensor_loss=25.40, train_loss=39.70, train_trace_loss=10.60, train_diag_loss=39.10, train_off_loss=0.303, train_tensor_loss=39.70, lr=0.0001]\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(model=pl_module, train_dataloaders=train_loader, val_dataloaders=val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Evaluation of model performance on the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Restoring states from the checkpoint path at ../scripts/pl_checkpoints/eps-epoch=79-val_loss=11.27.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [4,3,0,7,5,6,2,1]\n",
      "Loaded model weights from the checkpoint at ../scripts/pl_checkpoints/eps-epoch=79-val_loss=11.27.ckpt\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing DataLoader 0: 100%|██████████| 11/11 [00:00<00:00, 16.43it/s]\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "          loss              0.4318067133426666\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'loss': 0.4318067133426666}]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test(model=pl_module, dataloaders=test_loader, ckpt_path='../scripts/pl_checkpoints/eps-epoch=79-val_loss=11.27.ckpt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Predict dielectric tensors based on structures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 3.3504e+00, -4.8901e-08,  2.2176e-07],\n",
      "         [-4.8901e-08,  3.3504e+00, -9.7048e-08],\n",
      "         [ 2.2176e-07, -9.7048e-08,  3.3882e+00]]])\n"
     ]
    }
   ],
   "source": [
    "pl_module = DielectricModule.load_from_checkpoint(\n",
    "    '../scripts/pl_checkpoints/eps-epoch=79-val_loss=11.27.ckpt',\n",
    "    map_location = \"cuda:0\",\n",
    "    model = model,\n",
    ")\n",
    "pred = pl_module.predict_atoms_from_file('../scripts/test_structures/Na14Mn2O9.cif')\n",
    "\n",
    "print(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pfp",
   "language": "python",
   "name": "pfp"
  },
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}