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

@@ -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 }}
 

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

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 <https://github.com/isayevlab/AIMNet2>"""
+
+    _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})

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())

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"]

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)
 

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"}),
 ]