From 64fc8e45c5cccfdc855f75464bb3ef038fa6e1fb Mon Sep 17 00:00:00 2001 From: Carlo Asam Date: Tue, 3 Feb 2026 21:49:31 -0800 Subject: [PATCH 1/3] dpa3 inclusion framework calculating forces and energy --- maple/function/calculator/deepmd/__init__.py | 3 + .../calculator/deepmd/_dpa3_calculator.py | 86 +++++++++++++++++++ maple/function/calculator/set_calculator.py | 5 ++ 3 files changed, 94 insertions(+) create mode 100644 maple/function/calculator/deepmd/__init__.py create mode 100644 maple/function/calculator/deepmd/_dpa3_calculator.py diff --git a/maple/function/calculator/deepmd/__init__.py b/maple/function/calculator/deepmd/__init__.py new file mode 100644 index 0000000..63f9a4e --- /dev/null +++ b/maple/function/calculator/deepmd/__init__.py @@ -0,0 +1,3 @@ +from ._dpa3_calculator import DPA3Calculator + +__all__ = ["DPA3Calculator"] diff --git a/maple/function/calculator/deepmd/_dpa3_calculator.py b/maple/function/calculator/deepmd/_dpa3_calculator.py new file mode 100644 index 0000000..89e7320 --- /dev/null +++ b/maple/function/calculator/deepmd/_dpa3_calculator.py @@ -0,0 +1,86 @@ +""" +DPA3 / DeepMD-Kit calculator (PyTorch backend). +Uses the unified DP ASE calculator from deepmd.calculator; converts eV -> Hartree. +""" +import os +from typing import Literal + +from ase.calculators.calculator import all_changes +from ..calculator_base import CalcABC + +EV2HARTREE = 1.0 / 27.211386245988 + + +def _get_dp_calculator(): + """Import DeepMD-Kit DP calculator (works with .pt/.pth DPA3 models).""" + try: + from deepmd.calculator import DP + return DP + except ImportError: + try: + from deepmd.pt.utils.ase_calc import DPCalculator as DP + return DP + except ImportError: + raise ImportError( + "DeepMD-Kit is required for DPA3. Install with: pip install deepmd-kit" + ) from None + + +class DPA3Calculator(CalcABC): + """ASE calculator wrapping DeepMD-Kit DP (DPA3 and other Deep Potential models).""" + + implemented_properties = ["energy", "forces", "free_energy"] + + def __init__( + self, + device, + model: str = "dpa3", + overwrite: bool = False, + implicit: Literal["gbsa", "none"] = "gbsa", + solvent: str = "none", + ): + super().__init__() + DP = _get_dp_calculator() + + model_dir = os.path.dirname(os.path.realpath(__file__)) + model_dir = os.path.dirname(model_dir) + model_dir = os.path.dirname(model_dir) + default_path = os.path.join(model_dir, "model", f"{model}.pt") + if os.path.isfile(default_path): + model_path = default_path + elif os.path.isdir(model): + model_path = model + elif os.path.isfile(model): + model_path = model + else: + model_path = default_path + + self._dp = DP(model=model_path) + self.device = device + self.overwrite = overwrite + self.implicit_solv_init(implicit=implicit, solvent=solvent) + + def calculate( + self, atoms=None, properties=("energy", "forces"), system_changes=all_changes + ): + super().calculate(atoms, properties, system_changes) + + self._dp.calculate(atoms, properties=properties, system_changes=system_changes) + + energy_eV = self._dp.results.get("energy", 0.0) + energy = energy_eV * EV2HARTREE + + if self.solvent_correction: + energy = energy + self.implicit_solv_energy(atoms).item() + + self.results["energy"] = energy if isinstance(energy, float) else energy.item() + self.results["free_energy"] = self.results["energy"] + + if "forces" in properties: + forces = self._dp.results.get("forces") + if forces is not None: + forces = forces * EV2HARTREE # eV/Š-> Hartree/Š+ if self.solvent_correction: + _, solvent_force = self.implicit_solv_energy_and_force(atoms) + forces = forces + solvent_force.cpu().numpy() + self.results["forces"] = forces diff --git a/maple/function/calculator/set_calculator.py b/maple/function/calculator/set_calculator.py index 895d361..73660e5 100644 --- a/maple/function/calculator/set_calculator.py +++ b/maple/function/calculator/set_calculator.py @@ -22,6 +22,7 @@ 'uma', 'maceomol', 'aimnet2nse', + 'dpa3', ] # Model name to filename mapping for HuggingFace download @@ -145,6 +146,10 @@ def set_calculator(self) -> ase.calculators.calculator.Calculator: from .mace._mace_general_calculator import MACEModelCalculator calculator = MACEModelCalculator(model=self.model, device=self.device, implicit = self.implicit, solvent = self.solvent) return calculator + elif self.model == 'dpa3': + from .deepmd._dpa3_calculator import DPA3Calculator + calculator = DPA3Calculator(model=self.model, device=self.device, implicit=self.implicit, solvent=self.solvent) + return calculator else: raise ValueError(f"Model '{self.model}' is not implemented yet.") From a552ae1047ea44e95d372de6bc4229f0174f70cd Mon Sep 17 00:00:00 2001 From: Carlo Asam Date: Tue, 3 Feb 2026 21:53:55 -0800 Subject: [PATCH 2/3] dpa3 model inp file, needs weights in model folder --- maple/function/read/command_control.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/maple/function/read/command_control.py b/maple/function/read/command_control.py index fb40397..e26fc69 100644 --- a/maple/function/read/command_control.py +++ b/maple/function/read/command_control.py @@ -13,7 +13,8 @@ class CommandControl: SUPPORTED_MODELS = { "ani2x", "ani1x", "ani1ccx", "ani1xnr", "maceoff23s", "maceoff23m", "maceoff23l", - "egret", "aimnet2", "uma", "maceomol", "aimnet2nse" + "egret", "aimnet2", "uma", "maceomol", "aimnet2nse", + "dpa3", } SUPPORTED_TASKS = {"sp", "opt", "ts", "scan", "freq", "irc"} From a058de3c20414475b5dacb55be7367f14dc31f68 Mon Sep 17 00:00:00 2001 From: Carlo Asam Date: Tue, 3 Feb 2026 21:58:16 -0800 Subject: [PATCH 3/3] chgnet implemented, pretrained model uploads automatically --- maple/function/calculator/chgnet/__init__.py | 3 + .../calculator/chgnet/_chgnet_calculator.py | 72 +++++++++++++++++++ maple/function/calculator/set_calculator.py | 5 ++ maple/function/read/command_control.py | 2 +- 4 files changed, 81 insertions(+), 1 deletion(-) create mode 100644 maple/function/calculator/chgnet/__init__.py create mode 100644 maple/function/calculator/chgnet/_chgnet_calculator.py diff --git a/maple/function/calculator/chgnet/__init__.py b/maple/function/calculator/chgnet/__init__.py new file mode 100644 index 0000000..c06b6f9 --- /dev/null +++ b/maple/function/calculator/chgnet/__init__.py @@ -0,0 +1,3 @@ +from ._chgnet_calculator import CHGNetCalc + +__all__ = ["CHGNetCalc"] diff --git a/maple/function/calculator/chgnet/_chgnet_calculator.py b/maple/function/calculator/chgnet/_chgnet_calculator.py new file mode 100644 index 0000000..851fd4a --- /dev/null +++ b/maple/function/calculator/chgnet/_chgnet_calculator.py @@ -0,0 +1,72 @@ +""" +CHGNet calculator wrapper. Uses pretrained CHGNet via chgnet package; converts eV -> Hartree. +""" +from typing import Literal + +from ase.calculators.calculator import all_changes +from ..calculator_base import CalcABC + +EV2HARTREE = 1.0 / 27.211386245988 + + +def _get_chgnet_calculator(device): + """Build CHGNet ASE calculator (pretrained model, no local file).""" + try: + from chgnet.model import CHGNet + from chgnet.model import CHGNetCalculator + except ImportError: + raise ImportError( + "CHGNet is required. Install with: pip install chgnet" + ) from None + chgnet = CHGNet.load() + return CHGNetCalculator(potential=chgnet) + + +class CHGNetCalc(CalcABC): + """ASE calculator wrapping CHGNet (charge-informed graph neural network potential).""" + + implemented_properties = ["energy", "forces", "free_energy", "stress"] + + def __init__( + self, + device, + model: str = "chgnet", + overwrite: bool = False, + implicit: Literal["gbsa", "none"] = "gbsa", + solvent: str = "none", + ): + super().__init__() + self._chg = _get_chgnet_calculator(device) + self.device = device + self.overwrite = overwrite + self.implicit_solv_init(implicit=implicit, solvent=solvent) + + def calculate( + self, atoms=None, properties=("energy", "forces"), system_changes=all_changes + ): + super().calculate(atoms, properties, system_changes) + + self._chg.calculate(atoms, properties=properties, system_changes=system_changes) + + energy_eV = self._chg.results.get("energy", 0.0) + energy = energy_eV * EV2HARTREE + + if self.solvent_correction: + energy = energy + self.implicit_solv_energy(atoms).item() + + self.results["energy"] = energy if isinstance(energy, float) else energy.item() + self.results["free_energy"] = self.results["energy"] + + if "forces" in properties: + forces = self._chg.results.get("forces") + if forces is not None: + forces = forces * EV2HARTREE # eV/Š-> Hartree/Š+ if self.solvent_correction: + _, solvent_force = self.implicit_solv_energy_and_force(atoms) + forces = forces + solvent_force.cpu().numpy() + self.results["forces"] = forces + + if "stress" in properties: + stress = self._chg.results.get("stress") + if stress is not None: + self.results["stress"] = stress * EV2HARTREE # eV/ų -> Hartree/ų diff --git a/maple/function/calculator/set_calculator.py b/maple/function/calculator/set_calculator.py index 73660e5..031cbb0 100644 --- a/maple/function/calculator/set_calculator.py +++ b/maple/function/calculator/set_calculator.py @@ -23,6 +23,7 @@ 'maceomol', 'aimnet2nse', 'dpa3', + 'chgnet', ] # Model name to filename mapping for HuggingFace download @@ -150,6 +151,10 @@ def set_calculator(self) -> ase.calculators.calculator.Calculator: from .deepmd._dpa3_calculator import DPA3Calculator calculator = DPA3Calculator(model=self.model, device=self.device, implicit=self.implicit, solvent=self.solvent) return calculator + elif self.model == 'chgnet': + from .chgnet._chgnet_calculator import CHGNetCalc + calculator = CHGNetCalc(model=self.model, device=self.device, implicit=self.implicit, solvent=self.solvent) + return calculator else: raise ValueError(f"Model '{self.model}' is not implemented yet.") diff --git a/maple/function/read/command_control.py b/maple/function/read/command_control.py index e26fc69..a14244a 100644 --- a/maple/function/read/command_control.py +++ b/maple/function/read/command_control.py @@ -14,7 +14,7 @@ class CommandControl: "ani2x", "ani1x", "ani1ccx", "ani1xnr", "maceoff23s", "maceoff23m", "maceoff23l", "egret", "aimnet2", "uma", "maceomol", "aimnet2nse", - "dpa3", + "dpa3", "chgnet", } SUPPORTED_TASKS = {"sp", "opt", "ts", "scan", "freq", "irc"}