Add torchani models (prebuild only)

ipsuite/models/__init__.pyi

@@ -4,5 +4,15 @@ from .ensemble import EnsembleModel
 from .gap import GAP
 from .mace_model import MACE
 from .nequip import Nequip
+from .torchani import TorchAni
 
-__all__ = ["MLModel", "EnsembleModel", "GAP", "Nequip", "MACE", "Apax", "ApaxEnsemble"]
+__all__ = [
+    "MLModel",
+    "EnsembleModel",
+    "GAP",
+    "Nequip",
+    "MACE",
+    "Apax",
+    "ApaxEnsemble",
+    "TorchAni",
+]

ipsuite/models/torchani.py

@@ -0,0 +1,23 @@
+import ase.calculators.calculator
+import torchani
+import zntrack
+
+from ipsuite.models import MLModel
+
+_models = {
+    "ANI2x": torchani.models.ANI2x,
+    "ANI1ccx": torchani.models.ANI1ccx,
+    "ANI1x": torchani.models.ANI1x,
+}
+
+
+class TorchAni(MLModel):
+    """TorchAni model implementation."""
+
+    model_name: str = zntrack.zn.params("ANI2x")
+
+    def run(self) -> None:
+        pass
+
+    def get_calculator(self, **kwargs) -> ase.calculators.calculator.Calculator:
+        return _models[self.model_name].ase()

pyproject.toml

@@ -43,6 +43,7 @@ apax = { git = "https://github.com/apax-hub/apax.git", rev = "dev", optional = t
 rdkit = "^2023"
 lazy-loader = "^0.2"
 znflow = "^0.1.11"
+torchani = {version = "^2.2.3", optional = true}
 
 [tool.poetry.group.allegro.dependencies]
 mir-allegro = { git = "https://github.com/mir-group/allegro.git" }
@@ -73,6 +74,7 @@ comparison = ["h5py", "tensorflow", "llvmlite", "numba", "dscribe"]
 gap = ["quippy-ase", "xmltodict"]
 nequip = ["nequip"]
 apax = ["apax"]
+torchani = ["torchani"]
 
 [build-system]
 requires = ["poetry-core"]