From 6b259169c75b3ba58fdb675d703c67259e31ebc2 Mon Sep 17 00:00:00 2001 From: Josh Horton Date: Thu, 2 May 2024 07:18:52 +0100 Subject: [PATCH] Add AIMNET2 harness (#443) * add pyaimnet2 harness * lint and fix type hint * fix using check * add aimnet2 to CI testing * fix general harness tests * fix gradient tests * use conda package * update test * lint * add PR feedback * fix tests --- .github/workflows/CI.yml | 6 ++ devtools/conda-envs/aimnet2.yaml | 19 ++++ qcengine/programs/aimnet2.py | 122 +++++++++++++++++++++++ qcengine/programs/base.py | 2 + qcengine/programs/tests/test_programs.py | 47 +++++++++ qcengine/testing.py | 1 + qcengine/tests/test_harness_canonical.py | 3 +- 7 files changed, 199 insertions(+), 1 deletion(-) create mode 100644 devtools/conda-envs/aimnet2.yaml create mode 100644 qcengine/programs/aimnet2.py diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index f028fd89..e6e84e27 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -104,6 +104,12 @@ jobs: runs-on: ubuntu-latest pytest: "" + - conda-env: aimnet2 + python-version: 3.11 + label: AIMNET2 + runs-on: ubuntu-latest + pytest: "" + name: "🐍 ${{ matrix.cfg.python-version }} • ${{ matrix.cfg.label }} • ${{ matrix.cfg.runs-on }}" runs-on: ${{ matrix.cfg.runs-on }} diff --git a/devtools/conda-envs/aimnet2.yaml b/devtools/conda-envs/aimnet2.yaml new file mode 100644 index 00000000..511d7a0d --- /dev/null +++ b/devtools/conda-envs/aimnet2.yaml @@ -0,0 +1,19 @@ +name: test +channels: + - conda-forge +dependencies: + + # Core + - python + - pip + - pyyaml + - py-cpuinfo + - psutil + - qcelemental >=0.12.0 + - pydantic>=1.0.0 + + # Testing + - pytest + - pytest-cov + - codecov + - pyaimnet2 diff --git a/qcengine/programs/aimnet2.py b/qcengine/programs/aimnet2.py new file mode 100644 index 00000000..4f91b595 --- /dev/null +++ b/qcengine/programs/aimnet2.py @@ -0,0 +1,122 @@ +from typing import TYPE_CHECKING, Dict +from qcengine.programs.model import ProgramHarness +from qcelemental.util import safe_version, which_import +from qcelemental.models import AtomicResult, Provenance +from qcengine.exceptions import InputError + +if TYPE_CHECKING: + from qcelemental.models import AtomicInput + from qcengine.config import TaskConfig + + +class AIMNET2Harness(ProgramHarness): + """A harness to run AIMNET2 models """ + + _CACHE = {} + + _defaults = { + "name": "AIMNET2", + "scratch": False, + "thread_safe": True, + "thread_parallel": False, + "node_parallel": False, + "managed_memory": False, + } + + version_cache: Dict[str, str] = {} + + @staticmethod + def found(raise_error: bool = False) -> bool: + return which_import( + "pyaimnet2", + return_bool=True, + raise_error=raise_error, + raise_msg="Please install via `pip install git+https://github.com/jthorton/AIMNet2.git@main`", + ) + + def get_version(self) -> str: + self.found(raise_error=True) + + which_prog = which_import("pyaimnet2") + if which_prog not in self.version_cache: + import pyaimnet2 + + self.version_cache[which_prog] = safe_version(pyaimnet2.__version__) + + return self.version_cache[which_prog] + + def load_model(self, name: str): + model_name = name.lower() + if model_name in self._CACHE: + return self._CACHE[model_name] + + from pyaimnet2 import load_model + + model = load_model(model_name=model_name) + self._CACHE[model_name] = model + return self._CACHE[model_name] + + def compute(self, input_data: "AtomicInput", config: "TaskConfig"): + self.found(raise_error=True) + import torch + from qcengine.units import ureg + + # check we can run on the set of elements + known_elements = {"H", "B", "C", "N", "O", "F", "Si", "P", "S", "Cl", "As", "Se", "Br", "I"} + target_elements = set(input_data.molecule.symbols) + + unknown_elements = target_elements - known_elements + if unknown_elements: + raise InputError(f"AIMNET2 model {input_data.model.method} does not support elements {unknown_elements}.") + + method = input_data.model.method + # load the model using the method as the file name + model = self.load_model(name=method) + + # build the required input data + aimnet_input = { + "coord": torch.tensor( + [input_data.molecule.geometry * ureg.conversion_factor("bohr", "angstrom")], + dtype=torch.float64, + device="cpu", + ), + "numbers": torch.tensor([input_data.molecule.atomic_numbers], dtype=torch.long, device="cpu"), + "charge": torch.tensor([input_data.molecule.molecular_charge], dtype=torch.float64, device="cpu"), + } + + if input_data.driver == "gradient": + aimnet_input["coord"].requires_grad_(True) + out = model(aimnet_input) + + ret_data = { + "success": False, + "properties": { + "return_energy": out["energy"].item() * ureg.conversion_factor("eV", "hartree"), + "return_gradient": ( + -1.0 * out["forces"][0].detach().numpy() * ureg.conversion_factor("eV / angstrom", "hartree / bohr") + ), + "calcinfo_natom": len(input_data.molecule.atomic_numbers), + }, + "extras": input_data.extras.copy(), + } + # update with calculated extras + ret_data["extras"]["aimnet2"] = { + "charges": out["charges"].detach()[0].cpu().numpy(), + "ensemble_charges_std": out["charges_std"].detach()[0].cpu().numpy(), + "ensemble_energy_std": out["energy_std"].item(), + "ensemble_forces_std": out["forces_std"].detach()[0].cpu().numpy(), + } + if input_data.driver == "energy": + ret_data["return_result"] = ret_data["properties"]["return_energy"] + elif input_data.driver == "gradient": + ret_data["return_result"] = ret_data["properties"]["return_gradient"] + else: + raise InputError( + f"AIMNET2 can only compute energy and gradients driver methods. Requested {input_data.driver} not supported." + ) + + ret_data["provenance"] = Provenance(creator="pyaimnet2", version=self.get_version(), routine="load_model") + + ret_data["success"] = True + + return AtomicResult(**{**input_data.dict(), **ret_data}) diff --git a/qcengine/programs/base.py b/qcengine/programs/base.py index bf2fb155..c72687f4 100644 --- a/qcengine/programs/base.py +++ b/qcengine/programs/base.py @@ -29,6 +29,7 @@ from .turbomole import TurbomoleHarness from .xtb import XTBHarness from .mace import MACEHarness +from .aimnet2 import AIMNET2Harness __all__ = ["register_program", "get_program", "list_all_programs", "list_available_programs"] @@ -127,6 +128,7 @@ def list_available_programs() -> Set[str]: # AI register_program(TorchANIHarness()) register_program(MACEHarness()) +register_program(AIMNET2Harness()) # Molecular Mechanics register_program(RDKitHarness()) diff --git a/qcengine/programs/tests/test_programs.py b/qcengine/programs/tests/test_programs.py index f5ae4abc..dae4dcd7 100644 --- a/qcengine/programs/tests/test_programs.py +++ b/qcengine/programs/tests/test_programs.py @@ -413,3 +413,50 @@ def test_mace_gradient(): result = qcng.compute(atomic_input, "mace") assert result.success assert pytest.approx(result.return_result) == expected_result + + +@using("aimnet2") +@pytest.mark.parametrize( + "model, expected_energy", + [ + pytest.param("b973c", -76.39604306960972, id="b973c"), + pytest.param("wb97m-d3", -76.47412023758551, id="wb97m-d3"), + ], +) +def test_aimnet2_energy(model, expected_energy): + """Test computing the energies of water with two aimnet2 models.""" + + water = qcng.get_molecule("water") + atomic_input = AtomicInput(molecule=water, model={"method": model, "basis": None}, driver="energy") + + result = qcng.compute(atomic_input, "aimnet2") + assert result.success + assert pytest.approx(result.return_result) == expected_energy + assert "charges" in result.extras["aimnet2"] + assert "ensemble_charges_std" in result.extras["aimnet2"] + assert "ensemble_forces_std" in result.extras["aimnet2"] + + +@using("aimnet2") +def test_aimnet2_gradient(): + """Test computing the gradient of water using one aimnet2 model.""" + + water = qcng.get_molecule("water") + atomic_input = AtomicInput(molecule=water, model={"method": "wb97m-d3", "basis": None}, driver="gradient") + + result = qcng.compute(atomic_input, "aimnet2") + assert result.success + # make sure the gradient is now the return result + assert np.allclose( + result.return_result, + np.array( + [ + [-0.0, 2.6080331227973375e-09, -0.04097248986363411], + [-0.0, -0.029529934749007225, 0.020486244931817055], + [-0.0, 0.029529931023716927, 0.020486244931817055], + ] + ), + ) + assert pytest.approx(result.properties.return_energy) == -76.47412023758551 + # make sure the other properties were also saved + assert "charges" in result.extras["aimnet2"] diff --git a/qcengine/testing.py b/qcengine/testing.py index f20aaff2..c7a3ecf1 100644 --- a/qcengine/testing.py +++ b/qcengine/testing.py @@ -186,6 +186,7 @@ def get_job(self): "xtb": which_import("xtb", return_bool=True), "mrchem": is_program_new_enough("mrchem", "1.0.0"), "mace": is_program_new_enough("mace", "0.3.2"), + "aimnet2": which_import("pyaimnet2", return_bool=True), } _programs["openmm"] = _programs["rdkit"] and which_import(".openmm", package="simtk", return_bool=True) diff --git a/qcengine/tests/test_harness_canonical.py b/qcengine/tests/test_harness_canonical.py index a0aecba8..92d31cda 100644 --- a/qcengine/tests/test_harness_canonical.py +++ b/qcengine/tests/test_harness_canonical.py @@ -33,7 +33,8 @@ ("cfour", {"method": "hf", "basis": "6-31G"}, {}), ("gamess", {"method": "hf", "basis": "n31"}, {"basis__NGAUSS": 6}), ("mctc-gcp", {"method": "dft/sv"}, {}), - ("mace", {"method": "small"}, {}) + ("mace", {"method": "small"}, {}), + ("aimnet2", {"method": "b973c"}, {}), # add as programs available # ("terachem", {"method": "bad"}), ]