Updating the AIMNet package

pyproject.toml

@@ -93,7 +93,7 @@ plugins = ['pydantic.mypy']
 
 [[tool.mypy.overrides]]
 module = [
-  "aimnet2calc", "aimnet2calc.*",
+  "aimnet", "aimnet.*",
   "torch", "torch.*",
   "fairchem.core", "fairchem.core.*",
   "ase", "ase.*",

requirements/aimnet2.txt

@@ -1,7 +1 @@
---index-url https://download.pytorch.org/whl/cpu
---extra-index-url https://pypi.org/simple
---find-links https://data.pyg.org/whl/torch-2.8.0+cpu.html
-
-torch==2.8.0
-torch-cluster
-aimnet2calc @ https://github.com/isayevlab/AIMNet2/archive/refs/heads/main.zip
+aimnet>=0.1.0

src/oet/calculator/aimnet2.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Calculator for using AIMNet2 (https://github.com/isayevlab/AIMNet2),
+Calculator for using AIMNet2 (https://github.com/isayevlab/aimnetcentral),
 compatible with ORCA's ExtTool interface.
 
 Provides
@@ -13,7 +13,6 @@
 
 import shutil
 import sys
-import warnings
 from argparse import ArgumentParser
 from pathlib import Path
 from typing import Any
@@ -25,20 +24,16 @@
 from oet.core.misc import ENERGY_CONVERSION, LENGTH_CONVERSION, xyzfile_to_at_coord
 
 try:
-    # Suppress PySisyphus missing warning
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        from aimnet2calc import AIMNet2Calculator
-        from aimnet2calc.models import get_model_path, model_registry_aliases
-except ImportError:
+    from aimnet.calculators import AIMNet2Calculator
+    from aimnet.calculators.model_registry import get_model_path, load_model_registry
+except ImportError as err:
     print(
-        "[MISSING] Required module aimnet2calc not found.\n"
+        f"[MISSING] Required module aimnet not found: {err}.\n"
         "Please install the packages in the virtual environment.\n"
         "Therefore, activate the venv, got to the orca-external-tools "
         "main directory and use pip install -r requirements/aimnet2.txt"
     )
     sys.exit(1)
-
 try:
     import torch
 except ImportError as e:
@@ -65,6 +60,7 @@ class Aimnet2Calc(BaseCalc):
         "As": 33,
         "Se": 34,
         "Br": 35,
+        "Pd": 46,
         "I": 53,
     }
 
@@ -138,8 +134,14 @@ def get_model_file(model: str, model_dir: str) -> Path:
             return model_path
         # `model` must be the name of a model
         else:
-            # check aliases
-            model_file = model_registry_aliases.get(model, model)
+            # Check aliases
+            # First, check if the model is available in the registry. If not, assume it is a
+            # filename and treat it as is.
+            model_registry = load_model_registry()
+            if model in model_registry["aliases"]:
+                model_file = model_registry["aliases"][model]
+            else:
+                model_file = model
             # add jpt extension if not already present
             if not model_file.endswith(".jpt"):
                 model_file += ".jpt"
@@ -213,11 +215,14 @@ def extend_parser(cls, parser: ArgumentParser) -> None:
             "--model",
             type=str,
             dest="model",
-            default="aimnet2_wb97m",
-            help="The AIMNet2 model name or file name or absolute path. "
-            "If an absolute path is given, the file must exist. "
-            "Otherwise, it will be downloaded to DIR if necessary. "
-            'Default: "aimnet2_wb97m".',
+            default="aimnet2",
+            help="The AIMNet2 model name, file name, or absolute path. "
+            "If an absolute path is given, the file must exist locally. "
+            "Otherwise, the model will be downloaded to DIR if needed. "
+            "Note that different models are available. For example, `aimnet2nse` "
+            "is designed for open-shell/radical systems, such as those occurring in "
+            "transition states during bond breaking or formation. "
+            'Default: "aimnet2".',
         )
         parser.add_argument(
             "-p",

src/oet/calculator/uma.py

@@ -316,9 +316,9 @@ def calc(
             # Check if the model files are locally available. If not, subsequent errors will occur
             # as they cannot be downloaded.
             print(
-                    "WARNING: Offline mode selected, but no model files were detected. "
-                    "This will likely cause subsequent errors."
-                )
+                "WARNING: Offline mode selected, but no model files were detected. "
+                "This will likely cause subsequent errors."
+            )
 
         # setup calculator if not already set
         # this is important as usage on a server would otherwise cause

src/oet/server_client/server.py

@@ -20,6 +20,7 @@
 import importlib
 import io
 import logging
+import multiprocessing as mp
 import os
 import signal
 import sys
@@ -31,7 +32,6 @@
 from collections import OrderedDict
 from collections.abc import Mapping, Sequence
 from concurrent.futures import BrokenExecutor, ProcessPoolExecutor
-import multiprocessing as mp
 from contextlib import redirect_stdout
 from pathlib import Path
 from types import FrameType

tests/aimnet2/test_aiment2_client.py

@@ -71,17 +71,17 @@ def test_H2O_engrad(self):
         )
         run_aimnet2(input_file, output_file)
         expected_num_atoms = 3
-        expected_energy = -76.47682538331
+        expected_energy = -7.647682644970e+01
         expected_gradients = [
-            -0.01020942255855,
-            -0.007558935321867,
-            0.005339920055121,
-            0.003577792551368,
-            0.009023879654706,
-            0.001832913258113,
-            0.0066316309385,
-            -0.001464945613407,
-            -0.007172833196819,
+            -1.020941790193e-02,
+            -7.558942306787e-03,
+            5.339907016605e-03,
+            3.577796975151e-03,
+            9.023884311318e-03,
+            1.832915237173e-03,
+            6.631625350565e-03,
+            -1.464942586608e-03,
+            -7.172828074545e-03
         ]
 
         try:
@@ -109,14 +109,14 @@ def test_OH_anion_eng_grad(self):
         )
         run_aimnet2(input_file, output_file)
         expected_num_atoms = 2
-        expected_energy = -75.82629634884
+        expected_energy = -7.582629740302e+01
         expected_gradients = [
-            -0.000485832511913,
-            -0.001563806785271,
-            -0.0004455488233361,
-            0.000485832511913,
-            0.001563805621117,
-            0.0004455488233361,
+            -4.858186002821e-04,
+            -1.563774305396e-03,
+            -4.455401503947e-04,
+            4.858186002821e-04,
+            1.563771977089e-03,
+            4.455401503947e-04
         ]
 
         try:
@@ -144,14 +144,14 @@ def test_OH_rad_eng_grad(self):
         )
         run_aimnet2(input_file, output_file)
         expected_num_atoms = 2
-        expected_energy = -75.68258695326
+        expected_energy = -7.568258800204e+01
         expected_gradients = [
-            -3.78393149e-03,
-            -1.21797854e-02,
-            -3.47019313e-03,
-            3.78393149e-03,
-            1.21797854e-02,
-            3.47019313e-03,
+            -3.783911699429e-03,
+            -1.217970903963e-02,
+            -3.470176830888e-03,
+            3.783911699429e-03,
+            1.217970997095e-02,
+            3.470176830888e-03
         ]
 
         try:

tests/aimnet2/test_aiment2_standalone.py

@@ -34,17 +34,17 @@ def test_H2O_engrad(self):
         )
         run_aimnet2(input_file, output_file)
         expected_num_atoms = 3
-        expected_energy = -76.47682538331
+        expected_energy = -7.647682644970e+01
         expected_gradients = [
-            -0.01020942255855,
-            -0.007558935321867,
-            0.005339920055121,
-            0.003577792551368,
-            0.009023879654706,
-            0.001832913258113,
-            0.0066316309385,
-            -0.001464945613407,
-            -0.007172833196819,
+            -1.020941790193e-02,
+            -7.558942306787e-03,
+            5.339907016605e-03,
+            3.577796975151e-03,
+            9.023884311318e-03,
+            1.832915237173e-03,
+            6.631625350565e-03,
+            -1.464942586608e-03,
+            -7.172828074545e-03
         ]
 
         try:
@@ -72,14 +72,14 @@ def test_OH_anion_eng_grad(self):
         )
         run_aimnet2(input_file, output_file)
         expected_num_atoms = 2
-        expected_energy = -75.82629634884
+        expected_energy = -7.582629740302e+01
         expected_gradients = [
-            -0.000485832511913,
-            -0.001563806785271,
-            -0.0004455488233361,
-            0.000485832511913,
-            0.001563805621117,
-            0.0004455488233361,
+            -4.858186002821e-04,
+            -1.563774305396e-03,
+            -4.455401503947e-04,
+            4.858186002821e-04,
+            1.563771977089e-03,
+            4.455401503947e-04
         ]
 
         try:
@@ -107,14 +107,14 @@ def test_OH_rad_eng_grad(self):
         )
         run_aimnet2(input_file, output_file)
         expected_num_atoms = 2
-        expected_energy = -75.68258695326
+        expected_energy = -7.568258800204e+01
         expected_gradients = [
-            -3.78393149e-03,
-            -1.21797854e-02,
-            -3.47019313e-03,
-            3.78393149e-03,
-            1.21797854e-02,
-            3.47019313e-03,
+            -3.783911699429e-03,
+            -1.217970903963e-02,
+            -3.470176830888e-03,
+            3.783911699429e-03,
+            1.217970997095e-02,
+            3.470176830888e-03
         ]
 
         try: