save

src/graph_pes/models/__init__.py

@@ -28,7 +28,7 @@
 ALL_MODELS: list[type[GraphPESModel]] = [
     globals()[model]
     for model in __all__
-    if model not in ["FixedOffset", "LearnableOffset"]
+    if model not in ["FixedOffset", "LearnableOffset", "load_model"]
 ]
 
 

src/graph_pes/models/e3nn/nequip.py

@@ -3,6 +3,7 @@
 import warnings
 
 import e3nn
+import e3nn.util.jit
 import torch
 from e3nn import o3
 from graph_pes.core import GraphPESModel
@@ -321,6 +322,7 @@ def __call__(
         )
 
 
+@e3nn.util.jit.compile_mode("script")
 class NequIP(GraphPESModel):
     def __init__(
         self,

src/graph_pes/models/pairwise.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from typing import Optional
 
 import torch
 from jaxtyping import Float
@@ -141,7 +142,12 @@ def sigma(self):
 
     # don't use Z_i and Z_j, but include them for consistency with the
     # abstract method
-    def interaction(self, r: torch.Tensor, Z_i=None, Z_j=None):
+    def interaction(
+        self,
+        r: torch.Tensor,
+        Z_i: Optional[torch.Tensor] = None,
+        Z_j: Optional[torch.Tensor] = None,
+    ):
         """
         Evaluate the pair potential.
 

src/graph_pes/nn.py

@@ -471,15 +471,17 @@ def register_Zs(self, Zs: list[int]):
 
     def forward(self, Z: Tensor) -> Tensor:
         internal_idx = self.Z_to_idx[Z]
-        try:
-            return torch.nn.functional.one_hot(internal_idx, self.n_elements)
-        except IndexError:
-            raise ValueError(
-                f"Unknown atomic number: {sorted(set(Z.tolist()))}. "
-                f"Expected {self.n_elements} elements. "
-                "Did you forget to call `register_Zs`?"
-            ) from None
-
+        return torch.nn.functional.one_hot(internal_idx, self.n_elements)
+        # try:
+        #     return torch.nn.functional.one_hot(internal_idx, self.n_elements)
+        # except IndexError:
+        #     raise ValueError(
+        #         f"Unknown atomic number: {sorted(set(Z.tolist()))}. "
+        #         f"Expected {self.n_elements} elements. "
+        #         "Did you forget to call `register_Zs`?"
+        #     ) from None
+
+    @torch.jit.unused
     @property
     def registered_elements(self) -> list[str]:
         return [chemical_symbols[Z] for Z in self.Z_to_idx if Z <= MAX_Z]