diff --git a/pyproject.toml b/pyproject.toml index 687ea88..6943fcc 100644 --- a/pyproject.toml +++ b/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.*", diff --git a/requirements/aimnet2.txt b/requirements/aimnet2.txt index 6ee3daf..278698e 100644 --- a/requirements/aimnet2.txt +++ b/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 diff --git a/src/oet/calculator/aimnet2.py b/src/oet/calculator/aimnet2.py index 0f609a2..c6100b4 100755 --- a/src/oet/calculator/aimnet2.py +++ b/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", diff --git a/src/oet/calculator/uma.py b/src/oet/calculator/uma.py index c326156..429ccb7 100755 --- a/src/oet/calculator/uma.py +++ b/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 diff --git a/src/oet/server_client/server.py b/src/oet/server_client/server.py index 13e7b8d..3c637d7 100755 --- a/src/oet/server_client/server.py +++ b/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 diff --git a/tests/aimnet2/test_aiment2_client.py b/tests/aimnet2/test_aiment2_client.py index 8244011..c83c4cc 100644 --- a/tests/aimnet2/test_aiment2_client.py +++ b/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: diff --git a/tests/aimnet2/test_aiment2_standalone.py b/tests/aimnet2/test_aiment2_standalone.py index 2c88e3a..88baed1 100644 --- a/tests/aimnet2/test_aiment2_standalone.py +++ b/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: