deploy: 9a30dfe

_sources/tutorials/advanced/embedding-monkeypatch.md → _downloads/f926bfac3d372dde77d6e6dc68532a9c/embedding-monkeypatch.py

@@ -1,27 +1,9 @@
----
-jupytext:
-  text_representation:
-    extension: .md
-    format_name: myst
-    format_version: 0.13
-    jupytext_version: 1.16.1
-kernelspec:
-  display_name: Python 3 (ipykernel)
-  language: python
-  name: python3
----
-
-Embedding monkeypatching
--------------------------------
-
-
-```{code-cell} ipython3
 from ocpmodels.models.gemnet_oc.gemnet_oc import GemNetOC
 
 import torch
 from ocpmodels.common.utils import conditional_grad, scatter_det
-    
-    
+
+
 @conditional_grad(torch.enable_grad())
 def newforward(self, data):
     pos = data.pos
@@ -160,14 +142,12 @@ def newforward(self, data):
         out["energy"] = E_t.squeeze(1)  # (num_molecules)
 
     # This is the section I adapted from abishek's code
-    if hasattr(self, 'return_embedding') and self.return_embedding:
+    if hasattr(self, "return_embedding") and self.return_embedding:
         nMolecules = (torch.max(batch) + 1).item()
 
         # This seems to be an earlier block
-        out["h sum"] = scatter_det(
-            h, batch, dim=0, dim_size=nMolecules, reduce="add"
-        )
-        
+        out["h sum"] = scatter_det(h, batch, dim=0, dim_size=nMolecules, reduce="add")
+
         # These embedding are closer to energy output
         out["x_E sum"] = scatter_det(
             x_E, batch, dim=0, dim_size=nMolecules, reduce="add"
@@ -176,11 +156,10 @@ def newforward(self, data):
         # This is an embedding related to forces.
         # Something seems off on I couldn't do the same thing as above with scatter_det.
         out["x_F sum"] = torch.sum(x_F, axis=0)[None, :]
-        
 
         # tuples with nMolecules tensors of size nAtoms x embedding_size.
         out["x_E"] = x_E.split(data.natoms.tolist(), dim=0)
-        
+
         out["h"] = h.split(data.natoms.tolist(), dim=0)
 
         return out
@@ -190,10 +169,10 @@ def newforward(self, data):
         else:
             return out["energy"]
 
+
 GemNetOC.forward = newforward
-```
 
-```{code-cell} ipython3
+# +
 from ocpmodels.common.relaxation.ase_utils import OCPCalculator
 from ocpmodels.datasets import data_list_collater
 
@@ -212,21 +191,14 @@ def embed(self, atoms):
         with torch.no_grad():
             out = self.trainer.model([batch_list])
 
-    if (
-        self.trainer.normalizers is not None
-        and "target" in self.trainer.normalizers
-    ):
-        out["energy"] = self.trainer.normalizers["target"].denorm(
-            out["energy"]
-        )
-        out["forces"] = self.trainer.normalizers["grad_target"].denorm(
-            out["forces"]
-        )
+    if self.trainer.normalizers is not None and "target" in self.trainer.normalizers:
+        out["energy"] = self.trainer.normalizers["target"].denorm(out["energy"])
+        out["forces"] = self.trainer.normalizers["grad_target"].denorm(out["forces"])
 
     if self.trainer.ema:
         self.trainer.ema.restore()
 
     return out
 
+
 OCPCalculator.embed = embed
-```

_sources/tutorials/NRR/NRR_example.md

@@ -121,7 +121,10 @@ You need to provide the calculator with a path to a model checkpoint file. That
 Running the model with BFGS prints at each relaxation step is a lot to print. So we will just run one to demonstrate what happens on each iteration.
 
 ```{code-cell} ipython3
-checkpoint_path = "../escn_l6_m3_lay20_all_md_s2ef.pt"
+from ocpmodels.common.model_registry import model_name_to_local_file
+
+checkpoint_path = model_name_to_local_file('eSCN-L6-M3-Lay20All+MD', local_cache='/tmp/ocp_checkpoints/')
+
 os.makedirs(f"data/{bulk_src_id}_{adsorbate_smiles_h}", exist_ok=True)
 
 # Define the calculator

_sources/tutorials/advanced/embeddings.md

@@ -32,13 +32,12 @@ We provide 5 different kinds of embeddings:
 4. 'x_E sum' - summed over atoms
 5. 'x_F sum' - This is related to the forces
 
-In principle other models could be adapted in a similar way. See [embedding-monkeypatch.ipynb](./embedding-monkeypatch) for details on the patch. We simply run this notebook below to load it.
+In principle other models could be adapted in a similar way. See [embedding-monkeypatch.ipynb](./embedding-monkeypatch.py) for details on the patch. We simply run this notebook below to load it.
 
 The OCP project is still under active development, and it is not yet clear what the best way to access these embeddings are, so this code is not yet part of the main development branch. This code was adapted from a branch at https://github.com/Open-Catalyst-Project/ocp/blob/gnoc-embeddings.
 
 ```{code-cell} ipython3
-%run embedding-monkeypatch.ipynb
-%run ../ocp-tutorial.ipynb
+import embedding-monkeypatch
 ```
 
 # A diagnostic example

_sources/tutorials/advanced/fine-tuning-in-python.md

@@ -62,6 +62,9 @@ calc = OCPCalculator(checkpoint=checkpoint, trainer='forces', cpu=False)
 
 ```{code-cell} ipython3
 ! rm -fr train.db test.db val.db
+
+from ocpmodels.common.tutorial_utils import train_test_val_split
+
 train, test, val = train_test_val_split('../fine-tuning/oxides.db')
 train, test, val
 ```
@@ -71,6 +74,8 @@ train, test, val
 We start by making the config.yml. We build this from the calculator checkpoint.
 
 ```{code-cell} ipython3
+from ocpmodels.common.tutorial_utils import generate_yml_config
+
 yml = generate_yml_config(checkpoint, 'config.yml',
                    delete=['slurm', 'cmd', 'logger', 'task', 'model_attributes',
                            'optim.loss_force', # the checkpoint setting causes an error

_sources/tutorials/advanced/mass-inference.md

@@ -33,20 +33,25 @@ You can retrieve the dataset below. In this notebook we learn how to do "mass in
 You have to choose a checkpoint to start with. The newer checkpoints may require too much memory for this environment. 
 
 ```{code-cell} ipython3
-%run ../ocp-tutorial.ipynb
-list_checkpoints()
+from ocpmodels.common.model_registry import MODEL_REGISTRY
+print(MODEL_REGISTRY.keys())
 
 ```
 
 ```{code-cell} ipython3
-checkpoint = get_checkpoint('GemNet-dT OC22')
-checkpoint
+from ocpmodels.common.model_registry import model_name_to_local_file
+
+checkpoint_path = model_name_to_local_file('GemNet-dT OC22', local_cache='/tmp/ocp_checkpoints/')
+checkpoint_path
+
 ```
 
 We have to update our configuration yml file with the dataset. It is necessary to specify the train and test set for some reason. 
 
 ```{code-cell} ipython3
-yml = generate_yml_config(checkpoint, 'config.yml',
+
+from ocpmodels.common.tutorial_utils import generate_yml_config
+yml = generate_yml_config(checkpoint_path, 'config.yml',
                    delete=['cmd', 'logger', 'task', 'model_attributes',
                            'dataset', 'slurm'],
                    update={'amp': True,
@@ -75,7 +80,7 @@ It is a good idea to redirect the output to a file. If the output gets too large
 %%capture inference
 import time
 t0 = time.time()
-! python {ocp_main()} --mode predict --config-yml {yml} --checkpoint {checkpoint} --amp
+! python {ocp_main()} --mode predict --config-yml {yml} --checkpoint {checkpoint_path} --amp
 print(f'Elapsed time = {time.time() - t0:1.1f} seconds')
 ```
 
@@ -138,7 +143,7 @@ We include this here just to show that:
 
 ```{code-cell} ipython3
 from ocpmodels.common.relaxation.ase_utils import OCPCalculator
-calc = OCPCalculator(checkpoint=os.path.expanduser(checkpoint), cpu=False)
+calc = OCPCalculator(checkpoint=os.path.expanduser(checkpoint_path), cpu=False)
 ```
 
 ```{code-cell} ipython3

autoapi/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/data_parallel/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/distutils/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/flags/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/gp_utils/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/hpo_utils/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/logger/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/model_registry/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/registry/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/relaxation/ase_utils/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/relaxation/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/relaxation/ml_relaxation/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/relaxation/optimizers/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/relaxation/optimizers/lbfgs_torch/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/transforms/index.html

@@ -227,7 +227,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/tutorial_utils/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/typing/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/common/utils/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/datasets/_utils/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/datasets/ase_datasets/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/datasets/embeddings/atomic_radii/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/datasets/embeddings/continuous_embeddings/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/datasets/embeddings/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>
 
 

autoapi/ocpmodels/datasets/embeddings/khot_embeddings/index.html

@@ -225,7 +225,6 @@
 
 
 
-<li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embedding-monkeypatch.html">Embedding monkeypatching</a></li>
 <li class="toctree-l2"><a class="reference internal" href="../../../../../tutorials/advanced/embeddings.html">Working with embeddings</a></li>