tutorials/adsorbml_walkthrough.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cee85286-fcfe-452b-9017-7c0bce8168bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ocpmodels.common.relaxation.ase_utils import OCPCalculator\n",
    "import ase.io\n",
    "from ase.optimize import BFGS\n",
    "\n",
    "from ocdata.core import Adsorbate, AdsorbateSlabConfig, Bulk, Slab\n",
    "import os\n",
    "from glob import glob\n",
    "import pandas as pd\n",
    "from ocdata.utils import DetectTrajAnomaly\n",
    "from ocdata.utils.vasp import write_vasp_input_files\n",
    "\n",
    "# Optional - see below\n",
    "import numpy as np\n",
    "from dscribe.descriptors import SOAP\n",
    "from scipy.spatial.distance import pdist, squareform\n",
    "from x3dase.visualize import view_x3d_n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a774b7f5-e662-4959-9202-660d7bcaaacd",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Enumerate the adsorbate-slab configurations to run relaxations on"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6b10e112-25be-43e2-8f7d-0cda02187746",
   "metadata": {},
   "source": [
    "Be sure to set the path to the bulk and adsorbate pickle files in `ocdata/configs/paths.py` or pass the paths as an argument. The database pickles can be found in `ocdata/databases/pkls`. AdsorbML incorporates random placement, which is especially useful for more complicated adsorbates which may have many degrees of freedom. I have opted sample a few random placements and a few heuristic. Here I am using *CO on copper (1,1,1) as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6c7132ce-9536-4b5f-a7f7-cf42a528b312",
   "metadata": {},
   "outputs": [],
   "source": [
    "bulk_src_id = \"mp-30\"\n",
    "adsorbate_smiles = \"*CO\"\n",
    "\n",
    "bulk = Bulk(bulk_src_id_from_db = bulk_src_id, bulk_db_path = \"your-path-here.pkl\")\n",
    "adsorbate = Adsorbate(adsorbate_smiles_from_db=adsorbate_smiles, adsorbate_db_path = \"your-path-here.pkl\")\n",
    "slabs = Slab.from_bulk_get_specific_millers(bulk = bulk, specific_millers=(1,1,1))\n",
    "\n",
    "# There may be multiple slabs with this miller index.\n",
    "# For demonstrative purposes we will take the first entry.\n",
    "slab = slabs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d834adae-9d08-4e6f-9ada-23fc7c6b67d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform heuristic placements\n",
    "heuristic_adslabs = AdsorbateSlabConfig(slabs[0], adsorbate, mode=\"heuristic\")\n",
    "\n",
    "# Perform random placements\n",
    "# (for AdsorbML we use `num_sites = 100` but we will use 4 for brevity here)\n",
    "random_adslabs = AdsorbateSlabConfig(slabs[0], adsorbate, mode=\"random_site_heuristic_placement\", num_sites = 4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "31df1a82-46b1-4ab1-8c2c-e4983fe31747",
   "metadata": {},
   "source": [
    "# Run ML relaxations:\n",
    "\n",
    "There are 2 options for how to do this.\n",
    " 1. Using `OCPCalculator` as the calculator within the ASE framework\n",
    " 2. By writing objects to lmdb and relaxing them using `main.py` in the ocp repo\n",
    " \n",
    "(1) is really only adequate for small stuff and it is what I will show here, but if you plan to run many relaxations, you should definitely use (2). More details about writing lmdbs has been provided [here](https://github.com/Open-Catalyst-Project/ocp/blob/main/tutorials/lmdb_dataset_creation.ipynb) - follow the IS2RS/IS2RE instructions. And more information about running relaxations once the lmdb has been written is [here](https://github.com/Open-Catalyst-Project/ocp/blob/main/TRAIN.md#initial-structure-to-relaxed-structure-is2rs).\n",
    "\n",
    "You need to provide the calculator with a path to a model checkpoint file. That can be downloaded [here](https://github.com/Open-Catalyst-Project/ocp/blob/main/MODELS.md)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fe297959-7de2-4e54-bca6-562b230de12c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Unrecognized arguments: ['symmetric_edge_symmetrization']\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "amp: false\n",
      "cmd:\n",
      "  checkpoint_dir: /private/home/brookwander/AdsorbML/tutorials/checkpoints/2023-06-07-22-15-44\n",
      "  commit: 819e11d\n",
      "  identifier: ''\n",
      "  logs_dir: /private/home/brookwander/AdsorbML/tutorials/logs/tensorboard/2023-06-07-22-15-44\n",
      "  print_every: 100\n",
      "  results_dir: /private/home/brookwander/AdsorbML/tutorials/results/2023-06-07-22-15-44\n",
      "  seed: null\n",
      "  timestamp_id: 2023-06-07-22-15-44\n",
      "dataset: null\n",
      "gpus: 1\n",
      "logger: tensorboard\n",
      "model: gemnet_oc\n",
      "model_attributes:\n",
      "  activation: silu\n",
      "  atom_edge_interaction: true\n",
      "  atom_interaction: true\n",
      "  cbf:\n",
      "    name: spherical_harmonics\n",
      "  cutoff: 12.0\n",
      "  cutoff_aeaint: 12.0\n",
      "  cutoff_aint: 12.0\n",
      "  cutoff_qint: 12.0\n",
      "  direct_forces: true\n",
      "  edge_atom_interaction: true\n",
      "  emb_size_aint_in: 64\n",
      "  emb_size_aint_out: 64\n",
      "  emb_size_atom: 256\n",
      "  emb_size_cbf: 16\n",
      "  emb_size_edge: 1024\n",
      "  emb_size_quad_in: 64\n",
      "  emb_size_quad_out: 32\n",
      "  emb_size_rbf: 32\n",
      "  emb_size_sbf: 64\n",
      "  emb_size_trip_in: 64\n",
      "  emb_size_trip_out: 128\n",
      "  envelope:\n",
      "    exponent: 5\n",
      "    name: polynomial\n",
      "  extensive: true\n",
      "  forces_coupled: false\n",
      "  max_neighbors: 30\n",
      "  max_neighbors_aeaint: 20\n",
      "  max_neighbors_aint: 1000\n",
      "  max_neighbors_qint: 8\n",
      "  num_after_skip: 2\n",
      "  num_atom: 3\n",
      "  num_atom_emb_layers: 2\n",
      "  num_before_skip: 2\n",
      "  num_blocks: 6\n",
      "  num_concat: 4\n",
      "  num_global_out_layers: 2\n",
      "  num_output_afteratom: 3\n",
      "  num_radial: 128\n",
      "  num_spherical: 7\n",
      "  otf_graph: true\n",
      "  output_init: HeOrthogonal\n",
      "  qint_tags:\n",
      "  - 1\n",
      "  - 2\n",
      "  quad_interaction: true\n",
      "  rbf:\n",
      "    name: gaussian\n",
      "  regress_forces: true\n",
      "  sbf:\n",
      "    name: legendre_outer\n",
      "  symmetric_edge_symmetrization: false\n",
      "noddp: false\n",
      "optim:\n",
      "  batch_size: 4\n",
      "  clip_grad_norm: 10\n",
      "  ema_decay: 0.999\n",
      "  energy_coefficient: 1\n",
      "  eval_batch_size: 4\n",
      "  eval_every: 5000\n",
      "  factor: 0.8\n",
      "  force_coefficient: 100\n",
      "  load_balancing: atoms\n",
      "  loss_energy: mae\n",
      "  loss_force: l2mae\n",
      "  lr_initial: 0.0002\n",
      "  max_epochs: 80\n",
      "  mode: min\n",
      "  num_workers: 2\n",
      "  optimizer: AdamW\n",
      "  optimizer_params:\n",
      "    amsgrad: true\n",
      "  patience: 3\n",
      "  scheduler: ReduceLROnPlateau\n",
      "  weight_decay: 0\n",
      "slurm:\n",
      "  additional_parameters:\n",
      "    constraint: volta32gb\n",
      "  cpus_per_task: 3\n",
      "  folder: /private/home/anuroops/workspace/ocp_internal/gemnet/ocp/logs/47397573\n",
      "  gpus_per_node: 8\n",
      "  job_id: '47397573'\n",
      "  job_name: gemnet_base_allmd_8x8x4_nb6_ed1024_qin64_tout128_rbf32_sbf64_ncat4_lr2e-4\n",
      "  mem: 480GB\n",
      "  nodes: 8\n",
      "  ntasks_per_node: 8\n",
      "  partition: ocp,learnfair\n",
      "  time: 4320\n",
      "task:\n",
      "  dataset: trajectory_lmdb\n",
      "  description: Regressing to energies and forces for DFT trajectories from OCP\n",
      "  eval_on_free_atoms: true\n",
      "  grad_input: atomic forces\n",
      "  labels:\n",
      "  - potential energy\n",
      "  metric: mae\n",
      "  primary_metric: forces_mae\n",
      "  train_on_free_atoms: true\n",
      "  type: regression\n",
      "trainer: forces\n",
      "\n",
      "      Step     Time          Energy         fmax\n",
      "BFGS:    0 22:15:44       -0.561423        1.9519\n",
      "BFGS:    1 22:15:45       -0.558931        2.3817\n",
      "BFGS:    2 22:15:45       -0.595281        0.7013\n",
      "BFGS:    3 22:15:45       -0.609127        0.6193\n",
      "BFGS:    4 22:15:45       -0.636346        1.1445\n",
      "BFGS:    5 22:15:46       -0.647304        0.6850\n",
      "BFGS:    6 22:15:46       -0.651209        0.1758\n",
      "BFGS:    7 22:15:46       -0.653599        0.0947\n",
      "BFGS:    8 22:15:46       -0.654467        0.2220\n",
      "BFGS:    9 22:15:47       -0.655311        0.2275\n",
      "BFGS:   10 22:15:47       -0.655538        0.1208\n",
      "BFGS:   11 22:15:47       -0.655465        0.0255\n",
      "      Step     Time          Energy         fmax\n",
      "BFGS:    0 22:15:47       -0.430243        1.3932\n",
      "BFGS:    1 22:15:48       -0.435630        1.7549\n",
      "BFGS:    2 22:15:48       -0.470890        1.0809\n",
      "BFGS:    3 22:15:48       -0.566138        1.2374\n",
      "BFGS:    4 22:15:48       -0.569227        0.4316\n",
      "BFGS:    5 22:15:49       -0.571576        0.1568\n",
      "BFGS:    6 22:15:49       -0.572382        0.3412\n",
      "BFGS:    7 22:15:49       -0.571792        0.4606\n",
      "BFGS:    8 22:15:49       -0.575012        0.3174\n",
      "BFGS:    9 22:15:50       -0.573991        0.0981\n",
      "BFGS:   10 22:15:50       -0.573509        0.1604\n",
      "BFGS:   11 22:15:50       -0.573198        0.3230\n",
      "BFGS:   12 22:15:50       -0.582042        0.3933\n",
      "BFGS:   13 22:15:51       -0.588403        0.2653\n",
      "BFGS:   14 22:15:51       -0.592158        0.0948\n",
      "BFGS:   15 22:15:51       -0.593196        0.0509\n",
      "BFGS:   16 22:15:51       -0.593280        0.1092\n",
      "BFGS:   17 22:15:52       -0.588864        0.1371\n",
      "BFGS:   18 22:15:52       -0.587637        0.0867\n",
      "BFGS:   19 22:15:52       -0.584854        0.0319\n",
      "      Step     Time          Energy         fmax\n",
      "BFGS:    0 22:15:52       -0.523496        1.9444\n",
      "BFGS:    1 22:15:53       -0.527777        2.3635\n",
      "BFGS:    2 22:15:53       -0.564942        0.7029\n",
      "BFGS:    3 22:15:53       -0.579909        0.6344\n",
      "BFGS:    4 22:15:53       -0.611283        1.1946\n",
      "BFGS:    5 22:15:54       -0.625374        0.7463\n",
      "BFGS:    6 22:15:54       -0.628869        0.2049\n",
      "BFGS:    7 22:15:54       -0.633914        0.0966\n",
      "BFGS:    8 22:15:54       -0.632737        0.2104\n",
      "BFGS:    9 22:15:54       -0.632730        0.2323\n",
      "BFGS:   10 22:15:55       -0.636210        0.1323\n",
      "BFGS:   11 22:15:55       -0.635213        0.0343\n",
      "      Step     Time          Energy         fmax\n",
      "BFGS:    0 22:15:55       -0.267540        1.0896\n",
      "BFGS:    1 22:15:55       -0.306573        1.4670\n",
      "BFGS:    2 22:15:56       -0.417836        1.8248\n",
      "BFGS:    3 22:15:56       -0.477398        1.4147\n",
      "BFGS:    4 22:15:56       -0.531581        0.6218\n",
      "BFGS:    5 22:15:56       -0.540693        0.2894\n",
      "BFGS:    6 22:15:57       -0.541306        0.1435\n",
      "BFGS:    7 22:15:57       -0.543003        0.1314\n",
      "BFGS:    8 22:15:57       -0.563062        0.1333\n",
      "BFGS:    9 22:15:57       -0.564199        0.0919\n",
      "BFGS:   10 22:15:58       -0.567372        0.0903\n",
      "BFGS:   11 22:15:58       -0.573318        0.1011\n",
      "BFGS:   12 22:15:58       -0.583518        0.1063\n",
      "BFGS:   13 22:15:58       -0.584335        0.0603\n",
      "BFGS:   14 22:15:59       -0.580738        0.0354\n",
      "      Step     Time          Energy         fmax\n",
      "BFGS:    0 22:15:59       -0.459526        1.5999\n",
      "BFGS:    1 22:15:59       -0.462524        1.9972\n",
      "BFGS:    2 22:15:59       -0.494281        0.9750\n",
      "BFGS:    3 22:15:59       -0.548273        1.3185\n",
      "BFGS:    4 22:16:00       -0.591102        1.2766\n",
      "BFGS:    5 22:16:00       -0.600462        0.4214\n",
      "BFGS:    6 22:16:00       -0.602835        0.1933\n",
      "BFGS:    7 22:16:00       -0.604258        0.3799\n",
      "BFGS:    8 22:16:01       -0.608175        0.5353\n",
      "BFGS:    9 22:16:01       -0.609217        0.4061\n",
      "BFGS:   10 22:16:01       -0.607468        0.1399\n",
      "BFGS:   11 22:16:01       -0.607435        0.2059\n",
      "BFGS:   12 22:16:02       -0.611716        0.4337\n",
      "BFGS:   13 22:16:02       -0.614226        0.6105\n",
      "BFGS:   14 22:16:02       -0.620421        0.5321\n",
      "BFGS:   15 22:16:02       -0.621105        0.2439\n",
      "BFGS:   16 22:16:03       -0.624506        0.1582\n",
      "BFGS:   17 22:16:03       -0.623408        0.3073\n",
      "BFGS:   18 22:16:03       -0.620581        0.5686\n",
      "BFGS:   19 22:16:03       -0.619284        0.6212\n",
      "BFGS:   20 22:16:04       -0.622535        0.4151\n",
      "BFGS:   21 22:16:04       -0.625551        0.1114\n",
      "BFGS:   22 22:16:04       -0.626730        0.2276\n",
      "BFGS:   23 22:16:04       -0.625358        0.4183\n",
      "BFGS:   24 22:16:04       -0.626929        0.4653\n",
      "BFGS:   25 22:16:05       -0.629821        0.3058\n",
      "BFGS:   26 22:16:05       -0.632398        0.1476\n",
      "BFGS:   27 22:16:05       -0.635386        0.1480\n",
      "BFGS:   28 22:16:05       -0.636952        0.1874\n",
      "BFGS:   29 22:16:06       -0.642222        0.2479\n",
      "BFGS:   30 22:16:06       -0.640195        0.1642\n",
      "BFGS:   31 22:16:06       -0.637421        0.0496\n",
      "      Step     Time          Energy         fmax\n",
      "BFGS:    0 22:16:06       -0.297759        1.0291\n",
      "BFGS:    1 22:16:07       -0.319552        1.2143\n",
      "BFGS:    2 22:16:07       -0.391521        1.6311\n",
      "BFGS:    3 22:16:07       -0.143413        4.5038\n",
      "BFGS:    4 22:16:07       -0.467872        0.5046\n",
      "BFGS:    5 22:16:08       -0.484418        0.3222\n",
      "BFGS:    6 22:16:08       -0.496631        0.4339\n",
      "BFGS:    7 22:16:08       -0.499850        0.3310\n",
      "BFGS:    8 22:16:08       -0.510275        0.1918\n",
      "BFGS:    9 22:16:09       -0.525447        0.1974\n",
      "BFGS:   10 22:16:09       -0.533154        0.2342\n",
      "BFGS:   11 22:16:09       -0.534695        0.1489\n",
      "BFGS:   12 22:16:09       -0.535248        0.1343\n",
      "BFGS:   13 22:16:10       -0.537551        0.2158\n",
      "BFGS:   14 22:16:10       -0.544007        0.2638\n",
      "BFGS:   15 22:16:10       -0.543970        0.2368\n",
      "BFGS:   16 22:16:10       -0.546522        0.1934\n",
      "BFGS:   17 22:16:11       -0.547154        0.1538\n",
      "BFGS:   18 22:16:11       -0.549221        0.3359\n",
      "BFGS:   19 22:16:11       -0.552973        0.6825\n",
      "BFGS:   20 22:16:11       -0.559761        0.9237\n",
      "BFGS:   21 22:16:11       -0.568866        0.6173\n",
      "BFGS:   22 22:16:12       -0.583400        0.2377\n",
      "BFGS:   23 22:16:12       -0.586051        0.2631\n",
      "BFGS:   24 22:16:12       -0.597416        0.2508\n",
      "BFGS:   25 22:16:12       -0.610889        0.3062\n",
      "BFGS:   26 22:16:13       -0.616256        0.2844\n",
      "BFGS:   27 22:16:13       -0.618282        0.1322\n",
      "BFGS:   28 22:16:13       -0.621477        0.1828\n",
      "BFGS:   29 22:16:13       -0.627896        0.1580\n",
      "BFGS:   30 22:16:14       -0.629583        0.2922\n",
      "BFGS:   31 22:16:14       -0.628426        0.3300\n",
      "BFGS:   32 22:16:14       -0.628647        0.2546\n",
      "BFGS:   33 22:16:14       -0.630497        0.0922\n",
      "BFGS:   34 22:16:15       -0.633954        0.2232\n",
      "BFGS:   35 22:16:15       -0.636784        0.2473\n",
      "BFGS:   36 22:16:15       -0.641608        0.1797\n",
      "BFGS:   37 22:16:15       -0.646288        0.0961\n",
      "BFGS:   38 22:16:16       -0.649047        0.0955\n",
      "BFGS:   39 22:16:16       -0.648605        0.1767\n",
      "BFGS:   40 22:16:16       -0.649710        0.1442\n",
      "BFGS:   41 22:16:16       -0.651890        0.0666\n",
      "BFGS:   42 22:16:17       -0.651770        0.0679\n",
      "BFGS:   43 22:16:17       -0.653087        0.1280\n",
      "BFGS:   44 22:16:17       -0.653415        0.1278\n",
      "BFGS:   45 22:16:17       -0.653022        0.0801\n",
      "BFGS:   46 22:16:18       -0.654412        0.0381\n",
      "      Step     Time          Energy         fmax\n",
      "BFGS:    0 22:16:18       -0.288048        0.9128\n",
      "BFGS:    1 22:16:18       -0.295804        1.0280\n",
      "BFGS:    2 22:16:18       -0.412738        1.9463\n",
      "BFGS:    3 22:16:18       -0.366151        2.7266\n",
      "BFGS:    4 22:16:19       -0.470306        0.5755\n",
      "BFGS:    5 22:16:19       -0.477315        0.2972\n",
      "BFGS:    6 22:16:19       -0.483320        0.3629\n",
      "BFGS:    7 22:16:19       -0.487669        0.4059\n",
      "BFGS:    8 22:16:20       -0.508777        0.3446\n",
      "BFGS:    9 22:16:20       -0.516416        0.1635\n",
      "BFGS:   10 22:16:20       -0.520069        0.1174\n",
      "BFGS:   11 22:16:20       -0.523108        0.1247\n",
      "BFGS:   12 22:16:21       -0.526440        0.1328\n",
      "BFGS:   13 22:16:21       -0.528147        0.1551\n",
      "BFGS:   14 22:16:21       -0.532637        0.1109\n",
      "BFGS:   15 22:16:21       -0.532983        0.0563\n",
      "BFGS:   16 22:16:22       -0.532751        0.0491\n",
      "      Step     Time          Energy         fmax\n",
      "BFGS:    0 22:16:22       -0.238480        0.8653\n",
      "BFGS:    1 22:16:22       -0.256584        0.8340\n",
      "BFGS:    2 22:16:22       -0.403654        2.6614\n",
      "BFGS:    3 22:16:23       -0.456160        1.7490\n",
      "BFGS:    4 22:16:23       -0.478113        0.3854\n",
      "BFGS:    5 22:16:23       -0.478790        0.4265\n",
      "BFGS:    6 22:16:23       -0.484233        0.3568\n",
      "BFGS:    7 22:16:24       -0.498288        1.0200\n",
      "BFGS:    8 22:16:24       -0.510282        0.6611\n",
      "BFGS:    9 22:16:24       -0.519586        0.2115\n",
      "BFGS:   10 22:16:24       -0.521503        0.1278\n",
      "BFGS:   11 22:16:25       -0.526684        0.0722\n",
      "BFGS:   12 22:16:25       -0.528006        0.2383\n",
      "BFGS:   13 22:16:25       -0.529157        0.1802\n",
      "BFGS:   14 22:16:25       -0.528227        0.1073\n",
      "BFGS:   15 22:16:26       -0.527684        0.0753\n",
      "BFGS:   16 22:16:26       -0.527903        0.1274\n",
      "BFGS:   17 22:16:26       -0.526584        0.2011\n",
      "BFGS:   18 22:16:26       -0.524105        0.2224\n",
      "BFGS:   19 22:16:26       -0.524389        0.1359\n",
      "BFGS:   20 22:16:27       -0.525183        0.0855\n",
      "BFGS:   21 22:16:27       -0.528138        0.1352\n",
      "BFGS:   22 22:16:27       -0.530999        0.2716\n",
      "BFGS:   23 22:16:27       -0.533321        0.3112\n",
      "BFGS:   24 22:16:28       -0.534824        0.1513\n",
      "BFGS:   25 22:16:28       -0.536266        0.0583\n",
      "BFGS:   26 22:16:28       -0.537404        0.0909\n",
      "BFGS:   27 22:16:28       -0.540191        0.1603\n",
      "BFGS:   28 22:16:29       -0.535117        0.2099\n",
      "BFGS:   29 22:16:29       -0.533477        0.1828\n",
      "BFGS:   30 22:16:29       -0.533067        0.0712\n",
      "BFGS:   31 22:16:29       -0.532191        0.0651\n",
      "BFGS:   32 22:16:30       -0.542638        0.3142\n",
      "BFGS:   33 22:16:30       -0.545449        0.3568\n",
      "BFGS:   34 22:16:30       -0.554024        0.3208\n",
      "BFGS:   35 22:16:30       -0.568983        0.3860\n",
      "BFGS:   36 22:16:31       -0.567062        0.2222\n",
      "BFGS:   37 22:16:31       -0.569164        0.1600\n",
      "BFGS:   38 22:16:31       -0.575710        0.3226\n",
      "BFGS:   39 22:16:31       -0.575801        0.2419\n",
      "BFGS:   40 22:16:32       -0.578145        0.1278\n",
      "BFGS:   41 22:16:32       -0.579855        0.0793\n",
      "BFGS:   42 22:16:32       -0.581149        0.0645\n",
      "BFGS:   43 22:16:32       -0.579463        0.0590\n",
      "BFGS:   44 22:16:33       -0.580303        0.0319\n"
     ]
    }
   ],
   "source": [
    "checkpoint_path = \"your-path-here.pt\"\n",
    "os.makedirs(f\"data/{bulk}_{adsorbate}\", exist_ok=True)\n",
    "\n",
    "# Define the calculator\n",
    "calc = OCPCalculator(checkpoint=checkpoint_path) # if you have a gpu, add `cpu=False` to speed up calculations\n",
    "\n",
    "adslabs = [*heuristic_adslabs.atoms_list, *random_adslabs.atoms_list]\n",
    "# Set up the calculator\n",
    "for idx, adslab in enumerate(adslabs):\n",
    "    adslab.calc = calc\n",
    "    opt = BFGS(adslab, trajectory=f\"data/{bulk}_{adsorbate}/{idx}.traj\")\n",
    "    opt.run(fmax=0.05, steps=100) # For the AdsorbML results we used fmax = 0.02 and steps = 300, but we will use less strict values for brevity."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5bd48f41-52d9-475f-9b11-ec48fc0dc529",
   "metadata": {},
   "source": [
    "# Parse the trajectories and post-process\n",
    "\n",
    "As a post-processing step we check to see if:\n",
    "1. the adsorbate desorbed\n",
    "2. the adsorbate disassociated\n",
    "3. the adsorbate intercalated\n",
    "4. the surface has changed\n",
    "\n",
    "We check these because they effect our referencing scheme and may result in erroneous energies. For (4), the relaxed surface should really be supplied as well. It will be necessary when correcting the SP / RX energies later. Since we don't have it here, we will ommit supplying it, and the detector will instead compare the initial and final slab from the adsorbate-slab relaxation trajectory. If a relaxed slab is provided, the detector will compare it and the slab after the adsorbate-slab relaxation. The latter is more correct! Note: for the results in the AdsorbML paper, we did not check if the adsorbate was intercalated (`is_adsorbate_intercalated()`) because it is a new addition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0d005711-c702-4362-b930-f5bb8a6bbc55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iterate over trajs to extract results\n",
    "results = []\n",
    "for file in glob(f\"data/{bulk}_{adsorbate}/*.traj\"):\n",
    "    rx_id = file.split(\"/\")[-1].split(\".\")[0]\n",
    "    traj = ase.io.read(file, \":\")\n",
    "    \n",
    "    # Check to see if the trajectory is anomolous\n",
    "    initial_atoms = traj[0]\n",
    "    final_atoms = traj[-1]\n",
    "    atom_tags = initial_atoms.get_tags()\n",
    "    detector = DetectTrajAnomaly(initial_atoms, final_atoms, atom_tags)\n",
    "    anom = (\n",
    "        detector.is_adsorbate_dissociated()\n",
    "        or detector.is_adsorbate_desorbed()\n",
    "        or detector.has_surface_changed()\n",
    "        or detector.is_adsorbate_intercalated()\n",
    "    )\n",
    "    rx_energy = traj[-1].get_potential_energy()\n",
    "    results.append({\"relaxation_idx\": rx_id, \"relaxed_atoms\": traj[-1],\n",
    "                    \"relaxed_energy_ml\": rx_energy, \"anomolous\": anom})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ad52f34b-a885-4950-b040-b782acf0b899",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>relaxation_idx</th>\n",
       "      <th>relaxed_atoms</th>\n",
       "      <th>relaxed_energy_ml</th>\n",
       "      <th>anomolous</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.655465</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.584854</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.635213</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.580738</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.637421</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>5</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.654412</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>6</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.532751</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>7</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.580303</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  relaxation_idx                                      relaxed_atoms  \\\n",
       "0              0  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "1              1  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "2              2  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "3              3  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "4              4  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "5              5  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "6              6  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "7              7  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "\n",
       "   relaxed_energy_ml  anomolous  \n",
       "0          -0.655465      False  \n",
       "1          -0.584854      False  \n",
       "2          -0.635213      False  \n",
       "3          -0.580738      False  \n",
       "4          -0.637421      False  \n",
       "5          -0.654412      False  \n",
       "6          -0.532751      False  \n",
       "7          -0.580303      False  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame(results)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3cec27e4-ea87-468b-8088-389ba7635035",
   "metadata": {},
   "outputs": [],
   "source": [
    "#scrap anomalies\n",
    "df = df[~df.anomolous].copy().reset_index()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c1b71df2-8742-4789-bbbf-1654a5a174e4",
   "metadata": {},
   "source": [
    "# (Optional) Deduplicate structures\n",
    "We may have enumerated very similar structures or structures may have relaxed to the same configuration. For this reason, it is advantageous to cull systems if they are very similar. This results in marginal improvements in the recall metrics we calculated for AdsorbML, so it wasnt implemented there. It is, however, a good way to prevent wasteful VASP calculations. You can also imagine that if we would have enumerated 1000 configs per slab adsorbate combo rather than 100 for AdsorbML, it is more likely that having redundant systems would reduce performance, so its a good thing to keep in mind. This may be done by eye for a small number of systems, but with many systems it is easier to use an automated approach. Here is an example of one such approach, which uses a SOAP descriptor to find similar systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cd356d23-4b5e-4aaf-8f04-eb99ab96394f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract the configs and their energies\n",
    "def deduplicate(configs_for_deduplication: list,\n",
    "                adsorbate_binding_index: int,\n",
    "                cosine_similarity = 1e-3,\n",
    "               ):\n",
    "    \"\"\"\n",
    "    A function that may be used to deduplicate similar structures.\n",
    "    Among duplicate entries, the one with the lowest energy will be kept.\n",
    "    \n",
    "    Args:\n",
    "        configs_for_deduplication: a list of ML relaxed adsorbate-\n",
    "            surface configurations.\n",
    "        cosine_similarity: The cosine simularity value above which,\n",
    "            configurations are considered duplicate.\n",
    "            \n",
    "    Returns:\n",
    "        (list): the indices of configs which should be kept as non-duplicate\n",
    "    \"\"\"\n",
    "    \n",
    "    energies_for_deduplication = np.array([atoms.get_potential_energy() for atoms in configs_for_deduplication])\n",
    "    # Instantiate the soap descriptor\n",
    "    soap = SOAP(\n",
    "        species=np.unique(configs_for_deduplication[0].get_chemical_symbols()),\n",
    "        r_cut = 2.0,\n",
    "        n_max=6,\n",
    "        l_max=3,\n",
    "        periodic=True,\n",
    "    )\n",
    "    #Figure out which index cooresponds to \n",
    "    ads_len = list(configs_for_deduplication[0].get_tags()).count(2)\n",
    "    position_idx = -1*(ads_len-adsorbate_binding_index)\n",
    "    # Iterate over the systems to get the SOAP vectors\n",
    "    soap_desc = []\n",
    "    for config in configs_for_deduplication:\n",
    "        soap_ex = soap.create(config, centers=[position_idx])\n",
    "        soap_desc.extend(soap_ex)\n",
    "\n",
    "    soap_descs = np.vstack(soap_desc)\n",
    "\n",
    "    #Use euclidean distance to assess similarity\n",
    "    distance = squareform(pdist(soap_descs, metric=\"cosine\"))\n",
    "\n",
    "    bool_matrix = np.where(distance <= cosine_similarity, 1, 0)\n",
    "    # For configs that are found to be similar, just keep the lowest energy one\n",
    "    idxs_to_keep = []\n",
    "    pass_idxs = []\n",
    "    for idx, row in enumerate(bool_matrix):\n",
    "        if idx in pass_idxs:\n",
    "            continue\n",
    "            \n",
    "        elif sum(row) == 1:\n",
    "            idxs_to_keep.append(idx)\n",
    "        else:\n",
    "            same_idxs = [row_idx for row_idx, val in enumerate(row) if val == 1]\n",
    "            pass_idxs.extend(same_idxs)\n",
    "            # Pick the one with the lowest energy by ML\n",
    "            min_e = min(energies_for_deduplication[same_idxs])\n",
    "            idxs_to_keep.append(list(energies_for_deduplication).index(min_e))\n",
    "    return idxs_to_keep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a1dfa417-02fe-4c66-b06a-aeab56f702e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "configs_for_deduplication =  df.relaxed_atoms.tolist()\n",
    "idxs_to_keep = deduplicate(configs_for_deduplication, adsorbate.binding_indices[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "729be6f0-22b9-4d10-860b-39d7d8209ec2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 1, 3, 6]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"1000px\"\n",
       "            height=\"500px\"\n",
       "            src=\"x3dase.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7f6ca76bdc70>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Flip through your configurations to check them out (and make sure deduplication looks good)\n",
    "print(idxs_to_keep)\n",
    "view_x3d_n(configs_for_deduplication[2].repeat((2,2,1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1fd2c12e-d355-4213-b100-81a4054cbbcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.iloc[idxs_to_keep]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "04b87513-8397-4216-a567-ee81a25cef41",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The lowest 5 energies are: [-0.655 -0.585 -0.581 -0.533]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>index</th>\n",
       "      <th>relaxation_idx</th>\n",
       "      <th>relaxed_atoms</th>\n",
       "      <th>relaxed_energy_ml</th>\n",
       "      <th>anomolous</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.655465</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.584854</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.580738</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>6</td>\n",
       "      <td>6</td>\n",
       "      <td>(Atom('Cu', [1.3000465215529715, 2.25174662753...</td>\n",
       "      <td>-0.532751</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   index relaxation_idx                                      relaxed_atoms  \\\n",
       "0      0              0  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "1      1              1  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "3      3              3  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "6      6              6  (Atom('Cu', [1.3000465215529715, 2.25174662753...   \n",
       "\n",
       "   relaxed_energy_ml  anomolous  \n",
       "0          -0.655465      False  \n",
       "1          -0.584854      False  \n",
       "3          -0.580738      False  \n",
       "6          -0.532751      False  "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "low_e_values = np.round(df.sort_values(by = \"relaxed_energy_ml\").relaxed_energy_ml.tolist()[0:5],3)\n",
    "print(f\"The lowest 5 energies are: {low_e_values}\")\n",
    "df"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fe926290-3bec-406c-83d3-fcb2eee76c0a",
   "metadata": {},
   "source": [
    "# Write VASP input files\n",
    "\n",
    "This assumes you have access to VASP pseudopotentials. The default VASP flags (which are equivalent to those used to make OC20) are located in `ocdata.utils.vasp`. Alternatively, you may pass your own vasp flags to the `write_vasp_input_files` function as `vasp_flags`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5baeaa19-bc79-424c-8d1f-9b4a066a486c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grab the 5 systems with the lowest energy\n",
    "configs_for_dft = df.sort_values(by = \"relaxed_energy_ml\").relaxed_atoms.tolist()[0:5]\n",
    "config_idxs = df.sort_values(by = \"relaxed_energy_ml\").relaxation_idx.tolist()[0:5]\n",
    "\n",
    "# Write the inputs\n",
    "for idx, config in enumerate(configs_for_dft):\n",
    "    os.mkdir(f\"data/{config_idxs[idx]}\")\n",
    "    write_vasp_input_files(config, outdir = f\"data/{config_idxs[idx]}/\")"
   ]
  }
 ],
 "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}