notebooks/denoiser/demo.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "62bc0cf4-5fd3-49c4-9f16-df1b5676bf41",
   "metadata": {},
   "source": [
    "# Graph network model denoiser for denoising thermal perturbation\n",
    "\n",
    "This notebook will guide you through training and applying a graph network model for denoising thermal perturbation in crystal structures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fae635ac-e4cd-4054-932d-64731b851448",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import ase.io"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d1d99b0-454f-49d8-a57c-4916b84ba2a6",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a28de9a3-ad7c-4949-b35b-420f30509881",
   "metadata": {},
   "source": [
    "### Create lightning datamodule and module instances\n",
    "\n",
    "A lot of details are hidden away for your convenience using PyTorch Lightning. They include:\n",
    "- How the data is loaded, processed, batched, and so on.\n",
    "- Denoiser model definition.\n",
    "- How the model is trained, which optimizer, etc.\n",
    "\n",
    "These implementations can be found at `./lit/` if you are interested."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7ef4cdb6-deef-4a58-9799-fd3b6f38e954",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lit.datamodules import PeriodicStructureDataModule\n",
    "from lit.modules import LitEquivariantNoiseNet\n",
    "\n",
    "# Read ideal structures\n",
    "file_list = [\n",
    "    './data/Cu/ideal-bcc.dump',\n",
    "    './data/Cu/ideal-fcc.dump',\n",
    "    './data/Cu/ideal-hcp.dump'\n",
    "]\n",
    "\n",
    "# Datamodule\n",
    "datamodule = PeriodicStructureDataModule(\n",
    "    file_list    = file_list, \n",
    "    large_cutoff = 5.0,  # a large cutoff to precompute long bonds/edges\n",
    "    duplicate    = 256,\n",
    "    batch_size   = 8,\n",
    "    num_workers  = 4,\n",
    ")\n",
    "\n",
    "# New model instance\n",
    "# noise_net = LitEquivariantNoiseNet(\n",
    "#     num_species   = 1, \n",
    "#     node_dim      = 64, \n",
    "#     ff_dim        = 64, \n",
    "#     init_edge_dim = 1, \n",
    "#     edge_dim      = 64, \n",
    "#     num_heads     = 4, \n",
    "#     num_layers    = 4,\n",
    "#     sigma_max     = 0.3,  # maximum level of noise to be applied to ideal crystals \n",
    "#     cutoff        = 3.2,  # a smaller cutoff for downselecting the precomputed edges\n",
    "#     learn_rate    = 1e-4, \n",
    "# )\n",
    "\n",
    "# Load model weights from a saved checkpoint\n",
    "noise_net = LitEquivariantNoiseNet.load_from_checkpoint(\n",
    "    './lit_logs/equivariant-noise-net/version_0/checkpoints/epoch=1149-step=100000.ckpt'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b9c36d9-b4e6-4f4e-adb6-f73a9d34ee7e",
   "metadata": {},
   "source": [
    "### Start training session\n",
    "\n",
    "Skip this section is you loaded a pre-trained model ready for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1294cf8d-3154-4839-8256-a1bc0d1a8694",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightning as L\n",
    "from lightning.pytorch.callbacks import TQDMProgressBar\n",
    "from lightning.pytorch.loggers import TensorBoardLogger\n",
    "\n",
    "trainer = L.Trainer(\n",
    "    max_steps = 100_000,\n",
    "    logger    = TensorBoardLogger(save_dir='./lit_logs/', name='equivariant-noise-net'),\n",
    "    callbacks = [TQDMProgressBar(refresh_rate=10)],\n",
    ")\n",
    "\n",
    "trainer.fit(\n",
    "    noise_net, datamodule,\n",
    "    # ckpt_path = './lit_logs/...',  # specify this if you are resuming from a previous training session\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36b9488a-b26f-4878-86a1-4e843adf7527",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30f99573-a1d7-48eb-9b8b-af73390760dc",
   "metadata": {},
   "source": [
    "### Denoise test data\n",
    "\n",
    "In the case that the model was trained with Cu systems, it is still applicable to systems of other elements. You just have to scale the input data to match the atomic radius or the interatomic distance of the ideal Cu lattices. I used the following scaling factors:\n",
    "- In-house data\n",
    "    - Cu: 1.0\n",
    "    - Ta: 0.82\n",
    "- Benchmark data from Chung et al. (*\"Data-centric framework for crystal structure identification in atomistic simulations using machine learning\"*)\n",
    "    - Al (FCC): 0.9\n",
    "    - Ar (FCC): 0.65\n",
    "    - Fe (BCC): 0.9\n",
    "    - Li (BCC): 0.75\n",
    "    - Mg (HCP): 0.75\n",
    "    - Ti (HCP): 0.8"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d468e69a-8ac9-4581-81f2-9a6bb47c0897",
   "metadata": {},
   "source": [
    "### Helper function for denoising"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6718bf47-5f7a-4e93-8be1-347e454e794a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import LabelEncoder\n",
    "from tqdm.notebook import trange\n",
    "from graphite.nn import periodic_radius_graph\n",
    "\n",
    "@torch.no_grad()\n",
    "def denoise_snapshot(atoms, model, scale=1.0, steps=8):\n",
    "    # Convert to PyG format\n",
    "    x = LabelEncoder().fit_transform(atoms.numbers)\n",
    "    species = torch.tensor(x,                   dtype=torch.long)\n",
    "    pos     = torch.tensor(atoms.positions,     dtype=torch.float)\n",
    "    cell    = torch.tensor(atoms.cell.tolist(), dtype=torch.float)\n",
    "\n",
    "    # Scale\n",
    "    pos  *= scale\n",
    "    cell *= scale\n",
    "    \n",
    "    # Denoising trajectory\n",
    "    pos_traj = [atoms.positions]    \n",
    "    for _ in trange(steps):\n",
    "        edge_index, edge_vec = periodic_radius_graph(pos, r=3.2, cell=cell)\n",
    "        edge_len = torch.linalg.norm(edge_vec, dim=1, keepdim=True)\n",
    "        _, disp = model(species, edge_index, edge_attr=edge_len, edge_vec=edge_vec)\n",
    "        pos -= disp\n",
    "        pos_traj.append(pos.clone().numpy() / scale)\n",
    "    \n",
    "    return pos_traj"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a035dcd0-d8d8-4575-b579-f5b365556663",
   "metadata": {},
   "source": [
    "### Denoise a perturbed FCC snapshot\n",
    "\n",
    "You can also apply the model to other structures. Just change `TEST_FNAME` to a different structure file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "90a14335-0b86-4041-9a81-22af4308dd26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "09129a7e0b8840d7b04c9487e6b9d8bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "from scipy.spatial.transform import Rotation\n",
    "\n",
    "TEST_FNAME = './data/Cu/3400K-fcc.dump'\n",
    "path = Path(TEST_FNAME)\n",
    "\n",
    "# Read perturbed data\n",
    "noisy_atoms = ase.io.read(path)\n",
    "\n",
    "# Apply rotation to test equivariance\n",
    "R = Rotation.random().as_matrix()\n",
    "noisy_atoms.positions @= R\n",
    "noisy_atoms.cell      @= R\n",
    "\n",
    "# Get denoising trajectory (positions only)\n",
    "noise_net.eval(); noise_net.to('cpu')\n",
    "pos_traj = denoise_snapshot(noisy_atoms, noise_net.model, scale=1.0, steps=8)\n",
    "\n",
    "# Construct the denoising trajectory including other information (chemical symbols, cell dimenisons, etc.)\n",
    "denoising_traj = [\n",
    "    ase.Atoms(\n",
    "        symbols   = noisy_atoms.get_chemical_symbols(),\n",
    "        positions = pos,\n",
    "        cell      = noisy_atoms.cell,\n",
    "        pbc       = True\n",
    "    )\n",
    "    for pos in pos_traj\n",
    "]\n",
    "\n",
    "# Wrap the atoms into the simulaiton box\n",
    "for atoms in denoising_traj: atoms.wrap()\n",
    "\n",
    "# Save the denoising steps\n",
    "ase.io.write(path.with_suffix('.denoised_.extxyz'), denoising_traj)"
   ]
  }
 ],
 "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}