adding new notebook for using fairchem models with NEBs without CatTS…

…unami enumeration (#764)

* adding new notebook for using fairchem models with NEBs

* adding md tutorials

* blocking code cells that arent needed or take too long

docs/tutorials/cattsunami_walkthrough.md → docs/tutorials/cattsunami_tutorial.md

@@ -4,14 +4,14 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.16.1
+    jupytext_version: 1.16.3
 kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
   name: python3
 ---
 
-# CatTSunami tutorial
+# CatTSunami Tutorial
 
 ```{code-cell} ipython3
 ---
@@ -31,20 +31,17 @@ import matplotlib.pyplot as plt
 from fairchem.applications.cattsunami.core.autoframe import AutoFrameDissociation
 from fairchem.applications.cattsunami.core import OCPNEB
 from ase.io import read
-from IPython.display import Image
 
-# Optional
-# from x3dase.x3d import X3D
+#Optional
+from IPython.display import Image
+from x3dase.x3d import X3D
 
-# Set random seed
+#Set random seed
 import numpy as np
 np.random.seed(22)
 ```
 
-## Do enumerations in an AdsorbML style for CH dissociation on Ru (001)
-
-To start, we generate placements for the reactant and product species on the surface. We utilize the random placement approach which was developed for AdsorbML, and use an OCP model to relax our placements on the surface. These placements and their ML-determined energies are used as input to the CatTSunami automatic NEB frame generation approach.
-
+## Do enumerations in an AdsorbML style
 
 ```{code-cell} ipython3
 ---
@@ -54,16 +51,31 @@ tags: ["skip-execution"]
 reaction = Reaction(reaction_str_from_db="*CH -> *C + *H",
                     reaction_db_path=DISSOCIATION_REACTION_DB_PATH,
                     adsorbate_db_path = ADSORBATE_PKL_PATH)
+```
 
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
 # Instantiate our adsorbate class for the reactant and product
 reactant = Adsorbate(adsorbate_id_from_db=reaction.reactant1_idx, adsorbate_db_path=ADSORBATE_PKL_PATH)
 product1 = Adsorbate(adsorbate_id_from_db=reaction.product1_idx, adsorbate_db_path=ADSORBATE_PKL_PATH)
 product2 = Adsorbate(adsorbate_id_from_db=reaction.product2_idx, adsorbate_db_path=ADSORBATE_PKL_PATH)
+```
 
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
 # Grab the bulk and cut the slab we are interested in
 bulk = Bulk(bulk_src_id_from_db="mp-33", bulk_db_path=BULK_PKL_PATH)
 slab = Slab.from_bulk_get_specific_millers(bulk = bulk, specific_millers=(0,0,1))
+```
 
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
 # Perform site enumeration
 # For AdsorbML num_sites = 100, but we use 5 here for brevity. This should be increased for practical use.
 reactant_configs = AdsorbateSlabConfig(slab = slab[0], adsorbate = reactant,
@@ -89,16 +101,6 @@ cpu = True
 calc = OCPCalculator(checkpoint_path = checkpoint_path, cpu = cpu)
 ```
 
-### Run ML local relaxations:
-
-There are 2 options for how to do this.
- 1. Using `OCPCalculator` as the calculator within the ASE framework
- 2. By writing objects to lmdb and relaxing them using `main.py` in the ocp repo
-
-(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).
-
-You need to provide the calculator with a path to a model checkpoint file. That can be downloaded [here](../core/model_checkpoints)
-
 ```{code-cell} ipython3
 ---
 tags: ["skip-execution"]
@@ -110,15 +112,25 @@ for config in reactant_configs:
     opt = BFGS(config)
     opt.run(fmax = 0.05, steps=200)
     reactant_energies.append(config.get_potential_energy())
+```
 
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
 # Relax the product systems
 product1_energies = []
 for config in product1_configs:
     config.calc = calc
     opt = BFGS(config)
     opt.run(fmax = 0.05, steps=200)
     product1_energies.append(config.get_potential_energy())
+```
 
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
 product2_energies = []
 for config in product2_configs:
     config.calc = calc
@@ -128,8 +140,13 @@ for config in product2_configs:
 ```
 
 ## Enumerate NEBs
-Here we use the class we created to handle automatic generation of NEB frames to create frames using the structures we just relaxed as input.
-![dissociation_scheme](https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/applications/cattsunami/tutorial/dissociation_scheme.png)
+
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
+Image(filename="dissociation_scheme.png")
+```
 
 ```{code-cell} ipython3
 ---
@@ -146,7 +163,12 @@ af = AutoFrameDissociation(
             r_product2_max=3, #r3 in the above fig
             r_product2_min=1, #r2 in the above fig
 )
+```
 
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
 nframes = 10
 frame_sets, mapping_idxs = af.get_neb_frames(calc,
                                n_frames = nframes,
@@ -156,7 +178,6 @@ frame_sets, mapping_idxs = af.get_neb_frames(calc,
 ```
 
 ## Run NEBs
-Here we use the custom child class we created to run NEB relaxations using ML. The class we created allows the frame relaxations to be batched, improving efficiency.
 
 ```{code-cell} ipython3
 ---
@@ -189,7 +210,7 @@ tags: ["skip-execution"]
 #         conv = optimizer.run(fmax=fmax, steps=300)
 #         if conv:
 #             converged_idxs.append(idx)
-
+            
 # print(converged_idxs)
 ```
 
@@ -217,14 +238,13 @@ if conv:
     conv = optimizer.run(fmax=fmax, steps=300)
 ```
 
-## (Optional) Visualize the results
+## Visualize the results
 
 ```{code-cell} ipython3
 ---
 tags: ["skip-execution"]
 ---
-idx_of_interest = 0
-optimized_neb = read(f"n2_dissoc_on_Ru_{idx_of_interest}.traj", ":")[-1*nframes:]
+optimized_neb = read(f"ch_dissoc_on_Ru_{converged_idxs[0]}.traj", ":")[-1*nframes:]
 ```
 
 ```{code-cell} ipython3
@@ -235,7 +255,12 @@ es  = []
 for frame in optimized_neb:
     frame.set_calculator(calc)
     es.append(frame.get_potential_energy())
+```
 
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
 # Plot the reaction coordinate
 
 es = [e - es[0] for e in es]

docs/tutorials/fairchem_models_for_nebs.md

@@ -0,0 +1,162 @@
+---
+jupytext:
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+    jupytext_version: 1.16.3
+kernelspec:
+  display_name: Python 3 (ipykernel)
+  language: python
+  name: python3
+---
+
+# Tutorial for using Fair Chemistry models to relax NEBs
+
+```{code-cell} ipython3
+from ase.optimize import BFGS
+from ase.io import read
+
+from fairchem.applications.cattsunami.core.autoframe import interpolate
+from fairchem.applications.cattsunami.core import OCPNEB
+from fairchem.core.models.model_registry import model_name_to_local_file
+
+#Optional
+from x3dase.x3d import X3D
+import matplotlib.pyplot as plt
+from pathlib import Path
+import os
+```
+
+## Set up inputs
+
+Shown here are the values used consistently throughout the paper.
+
+```{code-cell} ipython3
+fmax = 0.05 # [eV / ang]
+delta_fmax_climb = 0.4 # this means that when the fmax is below 0.45 eV/Ang climbing image will be turned on
+k = 1 # you may adjust this value as you see fit
+cpu = True # set to False if you have a GPU
+
+
+# NOTE: Change the checkpoint path to locally downloaded files as needed
+checkpoint_path = model_name_to_local_file('EquiformerV2-31M-S2EF-OC20-All+MD', local_cache='/tmp/ocp_checkpoints/')
+```
+
+## If you have your own set of NEB frames
+
+```{code-cell} ipython3
+"""
+Load your frames (change to the appropriate loading method)
+The approach uses ase, so you must provide a list of ase.Atoms objects
+with the appropriate constraints.
+"""
+path_ = Path(__file__).parents[2]
+path_ = os.path.join(path_, "src", "fairchem", "applications", "cattsunami", "tutorial", "sample_traj.traj")
+frame_set = read(path_, ":")[0:10] # Change to the path to your atoms of the frame set
+```
+
+```{code-cell} ipython3
+neb = OCPNEB(
+    frame_set,
+    checkpoint_path=checkpoint_path,
+    k=k,
+    batch_size=8, # If you get a memory error, try reducing this to 4
+    cpu = cpu,
+)
+optimizer = BFGS(
+    neb,
+    trajectory=f"your-neb.traj",
+)
+conv = optimizer.run(fmax=fmax + delta_fmax_climb, steps=200)
+if conv:
+    neb.climb = True
+    conv = optimizer.run(fmax=fmax, steps=300)
+```
+
+## If you have a proposed initial and final frame
+
+You may use the `interpolate` function we implemented which is very similar to idpp but not sensative to periodic boundary crossings. Alternatively you can adopt whatever interpolation scheme you prefer. The `interpolate` function lacks some of the extra protections implemented in the `interpolate_and_correct_frames` which is used in the CatTSunami enumeration workflow. Care should be taken to ensure the results are reasonable.
+
+IMPORTANT NOTES: 
+1. Make sure the indices in the initial and final frame map to the same atoms
+2. Ensure you have the proper constraints on subsurface atoms
+
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
+"""
+Load your initial and frames (change to the appropriate loading method)
+The approach uses ase, so you must provide ase.Atoms objects
+with the appropriate constraints (i.e. fixed subsurface atoms).
+"""
+initial_frame = read("path-to-your-initial-atoms.traj")
+final_frame = read("path-to-your-final-atoms.traj")
+num_frames = 10 # you may change this to whatever you like
+```
+
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
+frame_set = interpolate(initial_frame, final_frame, num_frames)
+
+neb = OCPNEB(
+    frame_set,
+    checkpoint_path=checkpoint_path,
+    k=k,
+    batch_size=8, # If you get a memory error, try reducing this to 4
+    cpu = cpu,
+)
+optimizer = BFGS(
+    neb,
+    trajectory=f"your-neb.traj",
+)
+conv = optimizer.run(fmax=fmax + delta_fmax_climb, steps=200)
+if conv:
+    neb.climb = True
+    conv = optimizer.run(fmax=fmax, steps=300)
+```
+
+## Visualize the results
+
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
+optimized_neb = read(f"your-neb.traj", ":")[-1*nframes:]
+```
+
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
+es  = []
+for frame in optimized_neb:
+    frame.set_calculator(calc)
+    es.append(frame.get_potential_energy())
+```
+
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
+# Plot the reaction coordinate
+
+es = [e - es[0] for e in es]
+plt.plot(es)
+plt.xlabel("frame number")
+plt.ylabel("relative energy [eV]")
+plt.title(f"Ea = {max(es):1.2f} eV")
+plt.savefig("reaction_coordinate.png")
+```
+
+```{code-cell} ipython3
+---
+tags: ["skip-execution"]
+---
+# Make an interative html file of the optimized neb trajectory
+x3d = X3D(optimized_neb)
+x3d.write("your-neb.html")
+```

src/fairchem/applications/AdsorbML/tutorials/adsorbml_walkthrough.ipynb

@@ -7,7 +7,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from ocpmodels.common.relaxation.ase_utils import OCPCalculator\n",
+    "from faichem.core.common.relaxation.ase_utils import OCPCalculator\n",
     "import ase.io\n",
     "from ase.optimize import BFGS\n",
     "\n",