From 29bcbc2c8068404a4ac79db1378c5cb97816e751 Mon Sep 17 00:00:00 2001
From: Bryant Li <94406497+BryantLi-BLI@users.noreply.github.com>
Date: Tue, 4 Jun 2024 13:24:49 -0700
Subject: [PATCH] fixed pyacemdmaker

---
 src/atomate2/forcefields/__init__.py | 1 +
 src/atomate2/forcefields/md.py       | 9 ++-------
 src/atomate2/forcefields/utils.py    | 5 +++++
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/atomate2/forcefields/__init__.py b/src/atomate2/forcefields/__init__.py
index a76400755f..74ec2672dc 100644
--- a/src/atomate2/forcefields/__init__.py
+++ b/src/atomate2/forcefields/__init__.py
@@ -12,3 +12,4 @@ class MLFF(Enum):  # TODO inherit from StrEnum when 3.11+
     CHGNet = "CHGNet"
     Forcefield = "Forcefield"  # default placeholder option
     Nequip = "Nequip"
+    Pyace = "Pyace"
diff --git a/src/atomate2/forcefields/md.py b/src/atomate2/forcefields/md.py
index a8b2be96ae..9e6369a056 100644
--- a/src/atomate2/forcefields/md.py
+++ b/src/atomate2/forcefields/md.py
@@ -397,10 +397,5 @@ def _calculator(self) -> Calculator:
 class PyACEMDMaker(ForceFieldMDMaker):
     """Perform an MD run with PACE."""
 
-    name: str = "Py-ACE MD"
-    force_field_name: str = "Py-ACE"
-
-    def _calculator(self) -> Calculator:
-        import pyace
-
-        return pyace.PyACECalculator(**self.calculator_kwargs)
+    name: str = f"{MLFF.Pyace} MD"
+    force_field_name: str = f"{MLFF.Pyace}"
diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py
index ec65b5d282..897cd1a279 100644
--- a/src/atomate2/forcefields/utils.py
+++ b/src/atomate2/forcefields/utils.py
@@ -439,6 +439,11 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
 
             calculator = NequIPCalculator.from_deployed_model(**kwargs)
 
+        elif calculator_name == MLFF.Pyace:
+            import pyace
+
+            calculator = pyace.PyACECalculator(**kwargs)
+
     elif isinstance(calculator_meta, dict):
         calc_cls = MontyDecoder().decode(json.dumps(calculator_meta))
         calculator = calc_cls(**kwargs)