From ae291d81350d135e45b578a948523d74cd011cf0 Mon Sep 17 00:00:00 2001
From: Yuan Chiang <cyrusyc@berkeley.edu>
Date: Wed, 22 Jan 2025 16:06:50 -0800
Subject: [PATCH] add zbl test

---
 mlip_arena/models/__init__.py       |  5 ++--
 mlip_arena/models/classicals/zbl.py | 16 ++++---------
 tests/test_internal_calculators.py  | 36 +++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 15 deletions(-)
 create mode 100644 tests/test_internal_calculators.py

diff --git a/mlip_arena/models/__init__.py b/mlip_arena/models/__init__.py
index b21dc02..400f2be 100644
--- a/mlip_arena/models/__init__.py
+++ b/mlip_arena/models/__init__.py
@@ -39,9 +39,6 @@
 
 MLIPEnum = Enum("MLIPEnum", MLIPMap)
 
-# https://github.com/pytorch/pytorch/blob/3cbc8c54fd37eb590e2a9206aecf3ab568b3e63c/torch/_dynamo/config.py#L534
-# torch._dynamo.config.compiled_autograd = True
-
 class MLIP(
     nn.Module,
     PyTorchModelHubMixin,
@@ -49,6 +46,8 @@ class MLIP(
 ):
     def __init__(self, model: nn.Module) -> None:
         super().__init__()
+        # https://github.com/pytorch/pytorch/blob/3cbc8c54fd37eb590e2a9206aecf3ab568b3e63c/torch/_dynamo/config.py#L534
+        # torch._dynamo.config.compiled_autograd = True
         # self.model = torch.compile(model)
         self.model = model
 
diff --git a/mlip_arena/models/classicals/zbl.py b/mlip_arena/models/classicals/zbl.py
index 97a411c..5821539 100644
--- a/mlip_arena/models/classicals/zbl.py
+++ b/mlip_arena/models/classicals/zbl.py
@@ -13,23 +13,15 @@
 class ZBL(nn.Module):
     """Ziegler-Biersack-Littmark (ZBL) screened nuclear repulsion"""
 
-    # name: str
-    # implemented_properties: list[str] = ["energy", "forces", "stress"]
-
     def __init__(
         self,
         trianable: bool = False,
         **kwargs,
     ) -> None:
         nn.Module.__init__(self, **kwargs)
-        # Calculator.__init__(
-        #     self, restart=restart, atoms=atoms, directory=directory, **calculator_kwargs
-        # )
-
+        
         torch.set_default_dtype(torch.double)
 
-        # self.cutoff = torch.tensor(cutoff, dtype=torch.get_default_dtype())
-
         self.a = torch.nn.parameter.Parameter(
             torch.tensor(
                 [0.18175, 0.50986, 0.28022, 0.02817], dtype=torch.get_default_dtype()
@@ -142,7 +134,7 @@ def envelope(self, r: torch.Tensor, rc: torch.Tensor, p: int = 6):
         return y
 
     def _get_derivatives(self, energy: torch.Tensor, data: Data):
-        egrad, fij = torch.autograd.grad(
+        egradi, egradij = torch.autograd.grad(
             outputs=[energy],  # TODO: generalized derivatives
             inputs=[data.positions, data.vij],  # TODO: generalized derivatives
             grad_outputs=[torch.ones_like(energy)],
@@ -152,7 +144,7 @@ def _get_derivatives(self, energy: torch.Tensor, data: Data):
         )
 
         volume = torch.det(data.cell)  # (batch,)
-        rfaxy = torch.einsum("ax,ay->axy", data.vij, fij)
+        rfaxy = torch.einsum("ax,ay->axy", data.vij, -egradij)
 
         edge_batch = data.batch[data.edge_index[0]]
 
@@ -162,7 +154,7 @@ def _get_derivatives(self, energy: torch.Tensor, data: Data):
             / volume.view(-1, 1)
         )
 
-        return -egrad, stress
+        return -egradi, stress
 
     def forward(
         self,
diff --git a/tests/test_internal_calculators.py b/tests/test_internal_calculators.py
new file mode 100644
index 0000000..76447f7
--- /dev/null
+++ b/tests/test_internal_calculators.py
@@ -0,0 +1,36 @@
+import numpy as np
+from mlip_arena.models import MLIPCalculator
+from mlip_arena.models.classicals.zbl import ZBL
+
+from ase.build import bulk
+
+
+def test_zbl():
+    calc = MLIPCalculator(model=ZBL(), cutoff=6.0)
+
+    energies = []
+    forces = []
+    stresses = []
+
+    lattice_constants = [1, 3, 5, 7]
+
+    for a in lattice_constants:
+        atoms = bulk("Cu", "fcc", a=a) * (2, 2, 2)
+        atoms.calc = calc
+
+        energies.append(atoms.get_potential_energy())
+        forces.append(atoms.get_forces())
+        stresses.append(atoms.get_stress(voigt=False))
+
+    # test energy monotonicity
+    assert all(np.diff(energies) <= 0)
+
+    # test force vectors are all zeros due to symmetry
+    for f in forces:
+        assert np.allclose(f, 0)
+
+    # test trace of stress is monotonically increasing (less negative) and zero beyond cutoff
+    traces = [np.trace(s) for s in stresses]
+
+    assert all(np.diff(traces) >= 0)
+    assert np.allclose(stresses[-1], 0)