From 1eedc8ca56e16b7473f763944b6e56728ceba95a Mon Sep 17 00:00:00 2001
From: BowenD-UCB <84425382+BowenD-UCB@users.noreply.github.com>
Date: Thu, 24 Oct 2024 17:49:36 -0700
Subject: [PATCH] added site_energy observable in MD

---
 chgnet/model/dynamics.py | 15 +++++++++++++--
 tests/test_md.py         |  6 ++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/chgnet/model/dynamics.py b/chgnet/model/dynamics.py
index c1a8601f..9e7fa96c 100644
--- a/chgnet/model/dynamics.py
+++ b/chgnet/model/dynamics.py
@@ -51,7 +51,7 @@
 class CHGNetCalculator(Calculator):
     """CHGNet Calculator for ASE applications."""
 
-    implemented_properties = ("energy", "forces", "stress", "magmoms")
+    implemented_properties = ("energy", "forces", "stress", "magmoms", "energies")
 
     def __init__(
         self,
@@ -61,6 +61,7 @@ def __init__(
         check_cuda_mem: bool = False,
         stress_weight: float | None = 1 / 160.21766208,
         on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
+        return_site_energies: bool = False,
         **kwargs,
     ) -> None:
         """Provide a CHGNet instance to calculate various atomic properties using ASE.
@@ -80,6 +81,7 @@ def __init__(
             on_isolated_atoms ('ignore' | 'warn' | 'error'): how to handle Structures
                 with isolated atoms.
                 Default = 'warn'
+            return_site_energies (bool): whether to return the energy of each atom
             **kwargs: Passed to the Calculator parent class.
         """
         super().__init__(**kwargs)
@@ -95,6 +97,7 @@ def __init__(
             self.model = model.to(self.device)
         self.model.graph_converter.set_isolated_atom_response(on_isolated_atoms)
         self.stress_weight = stress_weight
+        self.return_site_energies = return_site_energies
         print(f"CHGNet will run on {self.device}")
 
     @classmethod
@@ -143,7 +146,10 @@ def calculate(
         structure = AseAtomsAdaptor.get_structure(atoms)
         graph = self.model.graph_converter(structure)
         model_prediction = self.model.predict_graph(
-            graph.to(self.device), task="efsm", return_crystal_feas=True
+            graph.to(self.device),
+            task="efsm",
+            return_crystal_feas=True,
+            return_site_energies=self.return_site_energies,
         )
 
         # Convert Result
@@ -156,6 +162,8 @@ def calculate(
             stress=model_prediction["s"] * self.stress_weight,
             crystal_fea=model_prediction["crystal_fea"],
         )
+        if self.return_site_energies:
+            self.results.update(energies=model_prediction["site_energies"])
 
 
 class StructOptimizer:
@@ -430,6 +438,7 @@ def __init__(
         crystal_feas_logfile: str | None = None,
         append_trajectory: bool = False,
         on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
+        return_site_energies: bool = False,
         use_device: str | None = None,
     ) -> None:
         """Initialize the MD class.
@@ -494,6 +503,7 @@ def __init__(
             on_isolated_atoms ('ignore' | 'warn' | 'error'): how to handle Structures
                 with isolated atoms.
                 Default = 'warn'
+            return_site_energies (bool): whether to return the energy of each atom
             use_device (str): the device for the MD run
                 Default = None
         """
@@ -517,6 +527,7 @@ def __init__(
                 model=model,
                 use_device=use_device,
                 on_isolated_atoms=on_isolated_atoms,
+                return_site_energies=return_site_energies,
             )
 
         if taut is None:
diff --git a/tests/test_md.py b/tests/test_md.py
index b97115b0..ec62f632 100644
--- a/tests/test_md.py
+++ b/tests/test_md.py
@@ -7,6 +7,7 @@
 import numpy as np
 import pytest
 from ase import Atoms
+from ase.io.trajectory import Trajectory
 from ase.md.npt import NPT
 from ase.md.nptberendsen import Inhomogeneous_NPTBerendsen
 from ase.md.nvtberendsen import NVTBerendsen
@@ -75,6 +76,7 @@ def test_md_nvt_berendsen(
         trajectory="md_out.traj",
         logfile="md_out.log",
         loginterval=10,
+        return_site_energies=True,
     )
     md.run(100)
 
@@ -102,6 +104,10 @@ def test_md_nvt_berendsen(
     )
     assert_allclose(logs, ref, rtol=2.1e-3, atol=1e-8)
 
+    traj = Trajectory("md_out.traj")
+    assert isinstance(traj[0].get_potential_energy(), float)
+    assert isinstance(traj[0].get_potential_energies(), np.ndarray)
+
 
 def test_md_nve(tmp_path: Path, monkeypatch: MonkeyPatch):
     monkeypatch.chdir(tmp_path)  # run MD in temporary directory