Upload folder using huggingface_hub

equiformer_calculator.py

@@ -0,0 +1,185 @@
+from typing import Dict, List, Optional, Tuple
+from omegaconf import ListConfig
+import yaml
+import os
+import torch
+from torch_geometric.data import Batch
+from torch_geometric.data import Data as TGData
+from torch_geometric.loader import DataLoader as TGDataLoader
+
+from nets.equiformer_v2.equiformer_v2_oc20 import EquiformerV2_OC20
+from nets.prediction_utils import compute_extra_props
+from ocpmodels.common.relaxation.ase_utils import (
+    batch_to_atoms,
+    ase_atoms_to_torch_geometric,
+)
+from ocpmodels.datasets import data_list_collater
+from ocpmodels.preprocessing import AtomsToGraphs
+
+from ase.calculators.calculator import Calculator
+from ase import Atoms
+
+
+def get_model(config_path):
+    with open(config_path, "r") as file:
+        config = yaml.safe_load(file)
+    model_config = config["model"]
+    print("model_config", model_config)
+    return EquiformerV2_OC20(**model_config), model_config
+
+
+class EquiformerCalculator:
+    def __init__(
+        self,
+        checkpoint_path: Optional[str] = None,
+        device: Optional[torch.device] = None,
+        **kwargs,
+    ):
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        project_root = os.path.dirname(os.path.dirname(__file__))
+
+        config_path = os.path.join(project_root, "configs/equiformer_v2.yaml")
+        self.model, self.model_config = get_model(config_path)
+
+        if checkpoint_path is None:
+            checkpoint_path = os.path.join(project_root, "ckpt/eqv2.ckpt")
+        state_dict = torch.load(checkpoint_path, weights_only=True)["state_dict"]
+        state_dict = {k.replace("potential.", ""): v for k, v in state_dict.items()}
+        self.model.load_state_dict(state_dict, strict=False)
+
+        self.model.eval()
+        self.model.to(device)
+
+        # ocpmodels/common/relaxation/ase_utils.py
+        self.a2g = AtomsToGraphs(
+            max_neigh=self.model.max_neighbors,
+            radius=self.model.cutoff,
+            r_energy=False,
+            r_forces=False,
+            r_distances=False,
+            r_edges=False,
+            r_pbc=True,
+        )
+
+    def predict(self, batch):
+        """Predict one or multiple batches"""
+        batch = batch.to(self.model.device)
+        batch = compute_extra_props(batch, pos_require_grad=False)
+        energy, forces, eigenpred = self.model.forward(batch, eigen=True)
+        return energy, forces, eigenpred
+
+    def get_forces(self, batch):
+        """Get forces from the model"""
+        batch = batch.to(self.model.device)
+        batch = compute_extra_props(batch, pos_require_grad=False)
+        _, forces, _ = self.model.forward(batch, eigen=False)
+        return forces
+
+    def get_energy(self, batch):
+        """Get energy from the model"""
+        batch = batch.to(self.model.device)
+        batch = compute_extra_props(batch, pos_require_grad=False)
+        energy, _, _ = self.model.forward(batch, eigen=False)
+        return energy
+
+    def predict_with_hessian(self, batch):
+        """Predict one batch with autodiff Hessian"""
+        B = batch.batch.max() + 1
+        assert B == 1, "Only one batch is supported for Hessian prediction"
+
+        batch = batch.to(self.model.device)
+
+        # Prepare batch with extra properties
+        batch = compute_extra_props(batch, pos_require_grad=True)
+
+        # Run prediction
+        with torch.enable_grad():
+            energy, forces, eigenpred = self.model.forward(batch, eigen=True)
+
+        # 3D coordinates -> 3N^2 Hessian elements
+        N = batch.pos.shape[0]
+        forces = forces.reshape(-1)
+        num_elements = forces.shape[0]
+
+        def get_vjp(v):
+            return torch.autograd.grad(
+                outputs=-1 * forces,
+                inputs=batch.pos,
+                grad_outputs=v,
+                retain_graph=True,
+                create_graph=False,
+                allow_unused=False,
+            )
+
+        I_N = torch.eye(num_elements, device=forces.device)
+        hessian = torch.vmap(get_vjp, in_dims=0, out_dims=0, chunk_size=None)(I_N)[0]
+        hessian = hessian.view(N * 3, N * 3)
+
+        eigenvalues, eigenvectors = torch.linalg.eigh(hessian)
+        smallest_eigenvals = eigenvalues[:2]
+        smallest_eigenvecs = eigenvectors[:, :2]
+        eigenvalues = smallest_eigenvals
+        eigenvectors = smallest_eigenvecs.T.view(2, N, 3)
+        return energy, forces, hessian, eigenvalues, eigenvectors, eigenpred
+
+    def predict_gad(self, batch):
+        """
+        Gentlest Ascent Dynamics (GAD)
+        dx/dt = -∇V(x) + 2(∇V, v(x))v(x)
+        = F + 2(-F, v(x))v(x)
+        since F=-∇V(x)
+        where v(x) is the eigenvector of the Hessian with the smallest eigenvalue.
+        """
+        B = batch.batch.max() + 1
+        energy, forces, eigenpred = self.predict(batch)
+        v = eigenpred["eigvec_1"].reshape(B, -1)
+        # normalize eigenvector
+        v = v / torch.norm(v, dim=1, keepdim=True)
+        forces = forces.reshape(B, -1)
+        # −∇V(x) + 2(∇V, v(x))v(x)
+        gad = forces + 2 * torch.einsum("bi,bi->b", -forces, v) * v
+        out = {
+            "energy": energy,
+            "forces": forces,
+        }
+        out.update(eigenpred)
+        return gad, out
+
+    def predict_gad_with_hessian(self, batch):
+        energy, forces, hessian, eigenvalues, eigenvectors, eigenpred = (
+            self.predict_with_hessian(batch)
+        )
+        v = eigenvectors[0].reshape(-1)  # N*3
+        v = v / torch.norm(v, dim=0, keepdim=True)
+        forces = forces.reshape(-1)  # N*3
+        # Diagnostic prints
+        dot_product = torch.dot(-forces, v)
+        gad = forces + (2 * dot_product * v)
+
+        out = {
+            "energy": energy,
+            "forces": forces,
+            "hessian": hessian,
+            "eigenvalues": eigenvalues,
+            "eigenvectors": eigenvectors,
+        }
+        out.update(eigenpred)
+        return gad, out
+
+    def find_transitionstate_with_gad_from_hessian(self, batch):
+        """Integrate the equations of motion of the GAD vector field."""
+        raise NotImplementedError("Not implemented")
+
+    def ase_to_batch(self, atoms: Atoms):
+        # Call base class to set atoms attribute
+        Calculator.calculate(self, atoms)
+
+        # ocpmodels/common/relaxation/ase_utils.py
+        # data_object = self.a2g.convert(atoms)
+        # batch = data_list_collater([data_object], otf_graph=True)
+
+        # Convert ASE atoms to torch_geometric format
+        batch = ase_atoms_to_torch_geometric(atoms)
+        return batch

example_inference.py

@@ -68,9 +68,10 @@ def predict_gad(batch, potential):
     B = batch.batch.max() + 1
     energy, forces, eigenpred = predict(batch, potential)
     v = eigenpred["eigvec_1"].reshape(B, -1)
+    v = v / torch.norm(v, dim=1, keepdim=True)
     forces = forces.reshape(B, -1)
     # −∇V(x) + 2(∇V, v(x))v(x)
-    gad = -forces + 2 * torch.einsum("bi,bi->b", forces, v) * v
+    gad = forces + 2 * torch.einsum("bi,bi->b", -forces, v) * v
     return gad
 
 
@@ -79,9 +80,10 @@ def predict_gad_with_hessian(batch, potential):
         predict_with_hessian(batch, potential)
     )
     v = eigenvectors[0].reshape(-1)  # N*3
+    v = v / torch.norm(v, dim=0, keepdim=True)
     forces = forces.reshape(-1)  # N*3
     # −∇V(x) + 2(∇V, v(x))v(x)
-    gad = -forces + 2 * torch.einsum("i,i->", forces, v) * v
+    gad = forces + 2 * torch.einsum("i,i->", -forces, v) * v
     return gad