diff --git a/.gitignore b/.gitignore index 82f9275..80307c0 100644 --- a/.gitignore +++ b/.gitignore @@ -160,3 +160,9 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ +*.bas4 +*.bas2 +*.bas3 +*.bas5 +*.bas0 +*.bas1 diff --git a/README.md b/README.md index 83b3135..79cabf6 100644 --- a/README.md +++ b/README.md @@ -75,6 +75,12 @@ performance flags (`--compile`, `--nb-threshold`, `--ensemble-member`), long-range Coulomb (`--coulomb-method`, `--coulomb-cutoff`), DFT-D3 flags, GPU server deployment notes, and complete ORCA-input examples. +### MACE options + +`oet_mace` and the `oet_server mace` server expose model-selection +and performance flags specific to MACE. See +[`readmes/mace.md`](readmes/mace.md) for further details. + ## Interface All scripts must be executable as: diff --git a/install.py b/install.py index 651beab..4174b8b 100755 --- a/install.py +++ b/install.py @@ -9,7 +9,7 @@ # Available extras -EXTRAS = ["uma", "aimnet2", "mlatom"] +EXTRAS = ["aimnet2", "mace", "mlatom", "uma"] # Minimal python interpreter required by the base class (currently 3.10) minimal_python_version = (3, 10) diff --git a/pyproject.toml b/pyproject.toml index 4052799..47c3c96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ packages = ["src/oet"] oet_aenet = "oet.calculator.aenet:main" oet_aimnet2 = "oet.calculator.aimnet2:main" oet_gxtb = "oet.calculator.gxtb:main" +oet_mace = "oet.calculator.mace:main" oet_mlatom = "oet.calculator.mlatom:main" oet_mopac = "oet.calculator.mopac:main" oet_uma = "oet.calculator.uma:main" @@ -101,6 +102,7 @@ module = [ "fairchem.core", "fairchem.core.*", "ase", "ase.*", "huggingface_hub", "huggingface_hub.*", + "mace", "mace.*", ] ignore_missing_imports = true diff --git a/readmes/mace.md b/readmes/mace.md new file mode 100644 index 0000000..259fe2a --- /dev/null +++ b/readmes/mace.md @@ -0,0 +1,20 @@ +# MACE ExtTool for ORCA + +Wrapper around MACE calculators to use ORCA's `otool_external` interface. + +- Suites: `mace-mp` (Materials Project) and `mace-omol` (OMOL foundation model) +- Extras: `dispersion` (MP only), `default_dtype` (`float32` or `float64`), optional `device`, and `head` for advanced MP heads. + +## Arguments + +| Category | Argument | Values / Description | Default | +|---|---|---|---| +| Common | `-s, --suite` | `mp` or `omol` | `omol` | +| Common | `-m, --model` | Model spec or local path. MP examples: `medium-mpa-0`, `medium`, `small`; OMOL: `extra_large` or local path | — | +| Common | `--default-dtype` | `float32` (MD speed) or `float64` (optimization accuracy) | — | +| Common | `--device` | Compute device, e.g. `cpu`, `cuda` | Auto / framework default | +| MP only | `--dispersion` | Enable D3 dispersion correction | Off | +| MP only | `--damping` | Advanced dispersion damping controls | — | +| MP only | `--dispersion-xc` | Exchange-correlation functional for dispersion | — | +| MP only | `--dispersion-cutoff` | Dispersion cutoff distance | — | +| MP only | `--head` | MACE head selection for multi-head variants | — | \ No newline at end of file diff --git a/requirements/mace.txt b/requirements/mace.txt new file mode 100644 index 0000000..1d58135 --- /dev/null +++ b/requirements/mace.txt @@ -0,0 +1 @@ +mace-torch>=0.3.14 diff --git a/src/oet/calculator/mace.py b/src/oet/calculator/mace.py new file mode 100755 index 0000000..2992e37 --- /dev/null +++ b/src/oet/calculator/mace.py @@ -0,0 +1,332 @@ +#!/usr/bin/env python3 +""" +MACE wrapper for ORCA's ExtTool interface. +""" + +import sys +import warnings +from argparse import ArgumentParser +from typing import Any + +from oet.core.base_calc import BaseCalc, CalculationData +from oet.core.misc import ENERGY_CONVERSION, LENGTH_CONVERSION, xyzfile_to_at_coord + +try: + with warnings.catch_warnings(): + # Catch the warning that the `weights_only` bool is not set during `torch.load`. This will not affect the calculation. + warnings.filterwarnings( + "ignore", + message="Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected", + ) + from mace.calculators.foundations_models import mace_mp, mace_omol + from mace.calculators.mace import MACECalculator +except ImportError as e: + print( + f"[MISSING] Required module mace not found: {e}.\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/mace.txt" + ) + sys.exit(1) + +try: + import torch +except ImportError as e: + print("[MISSING] torch not found:", e) + sys.exit(1) + +try: + from ase import Atoms +except ImportError as e: + print("[MISSING] ase not found:", e) + sys.exit(1) + + +class MaceCalc(BaseCalc): + # MACE calculator used for computations + _calc: MACECalculator | None = None + + def set_calculator( + self, + suite: str, + model: str, + dispersion: bool, + damping: str, + dispersion_xc: str, + dispersion_cutoff: float, + device: str, + default_dtype: str, + head: str, + force: bool = False, + ) -> None: + """ + Prepare a MACE calculator object that is compatible with the ASE calculator object to compute energy and gradient, if not done already. + + Parameters + ---------- + suite: str + Suite for the calculator. + model: str + Model/model path to be used. + dispersion: bool + Whether to use a dispersion model or not. + damping: str + Dispersion damping. + dispersion_xc: str + XC functional for the dispersion model. + dispersion_cutoff: str + Cutoff radius for the dispersion model. + device: str + Device (cuda or cpu). + default_dtype: str + Precision type. + head: str + The head of the MACE model. + force: bool, default = False + Force re-initialization of the calculator, even if already initialized. + """ + if self._calc and not force: + return + match suite: + case "mp": + kwargs = dict( + model=model, + device=device, + default_dtype=default_dtype or "float32", + dispersion=dispersion, + damping=damping, + dispersion_xc=dispersion_xc, + dispersion_cutoff=(dispersion_cutoff * LENGTH_CONVERSION["Ang"]), + ) + if head: + kwargs["head"] = head + calc = mace_mp(**kwargs) + # No dispersion applicable for omol + case "omol": + calc = mace_omol( + model=model, + device=device, + default_dtype=default_dtype or "float64", + ) + case _: + raise RuntimeError(f"The suite {suite} is not supported.") + self._calc = calc + + def get_calculator(self) -> MACECalculator | None: + """ + Returns the current MACE calculator + """ + return self._calc + + @classmethod + def extend_parser(cls, parser: ArgumentParser) -> None: + """ + Add Mace parsing options. + + Parameters + ---------- + parser: ArgumentParser + Parser that should be extended + """ + parser.add_argument( + "-s", + "--suite", + choices=["mp", "omol", "mace-mp", "mace-omol"], + default="omol", + help="Select MACE suite: mp/mace-mp or omol/mace-omol. Default: omol", + ) + # Settings the default model to None uses the MACE default. + parser.add_argument( + "-m", + "--model", + type=str, + default=None, + help="Model spec or local path. MP: small/medium/large/medium-mpa-0/... OMOL: extra_large or path.", + ) + parser.add_argument( + "--default-dtype", + choices=["float32", "float64"], + default="float64", + help="Default float precision (recommended: use float32 only for MD).", + ) + parser.add_argument( + "-d", + "--device", + type=str, + default="cpu", + metavar="DEVICE", + dest="device", + choices=(device_choices := ["cpu", "cuda"]), + help="Device to perform the calculation on. " + "Options: " + ", ".join(device_choices) + ". " + "Default: cpu. ", + ) + # MP-specific extras + parser.add_argument( + "--dispersion", + action="store_true", + help=( + "Enable D3 dispersion. Only applicable for the MP suite. " + "This requires the installation of torch_dftd to the virtual " + "environment and that liblzma was installed when setting up " + "the environment." + ) + ) + parser.add_argument( + "--damping", + type=str, + default="bj", + choices=(damp_choices := ["zero", "bj", "zerom", "bjm"]), + help="D3 damping. " + "Options: " + ", ".join(damp_choices) + "." + "Only applicable for the MP suite. Default BJ for MP suite. ", + ) + parser.add_argument( + "--dispersion-xc", + type=str, + default="pbe", + help="XC functional for D3. Only applicable for the MP suite. Default PBE for MP suite.", + ) + parser.add_argument( + "--dispersion-cutoff", + type=float, + default=40.0, + help="Cutoff radius for D3 in Bohr. " + "Only applicable for the MP suite. " + "Default: 40 Bohr for MP suite.", + ) + parser.add_argument( + "--head", + type=str, + default=None, + help="Advanced option: Select MACE head. Only applicable for the MP suite.", + ) + + def run_mace( + self, + atom_types: list[str], + coordinates: list[tuple[float, float, float]], + calc_data: CalculationData, + ) -> tuple[float, list[float]]: + """ + Runs a Mace calculation. + + Parameters + ---------- + atom_types : list[str] + List of element symbols (e.g., ["O", "H", "H"]) + coordinates : list[tuple[float, float, float]] + List of (x, y, z) coordinates + calc_data: CalculationData + Object with calculation data for the run + + Returns + ------- + float + The computed energy (Eh) + list[float] + Flattened gradient vector (Eh/Bohr), if computed, otherwise empty + """ + + # set the number of threads + torch.set_num_threads(calc_data.ncores) + + # make ase atoms object for calculation + atoms = Atoms(symbols=atom_types, positions=coordinates) + atoms.info = {"charge": calc_data.charge, "spin": calc_data.mult} + atoms.calc = self._calc + + # Get the energy and gradients + energy = atoms.get_potential_energy() / ENERGY_CONVERSION["eV"] + gradient = [] + try: + forces = atoms.get_forces() + # Convert forces to gradient (-1) and unit conversion + fac = -LENGTH_CONVERSION["Ang"] / ENERGY_CONVERSION["eV"] + gradient = (fac * forces).flatten().tolist() + except Exception: + # forces may not be available + pass + + return energy, gradient + + def calc( + self, + calc_data: CalculationData, + args_parsed: dict[str, Any], + args_not_parsed: list[str], + ) -> tuple[float, list[float]]: + """ + Routine for calculating energy and optional gradient. + + Parameters + ---------- + calc_data: CalculationData + Object with calculation data for the run + args_parsed: dict[str, Any] + Arguments parsed as defined in extend_parser + args_not_parsed: list[str] + Arguments not parsed so far + + Returns + ------- + float + The computed energy (Eh) + list[float] + Flattened gradient vector (Eh/Bohr), if computed, otherwise empty + """ + + suite = args_parsed["suite"] + if suite.startswith("mace"): + suite = suite.split("-", 1)[1] + dispersion = bool(args_parsed["dispersion"]) + dispersion_xc = args_parsed["dispersion_xc"] + dispersion_cutoff = args_parsed["dispersion_cutoff"] + if dispersion and suite != "mp": + print( + "WARNING: Dispersion flag recognized, but MP suite not used. Ignoring all options related to dispersion." + ) + + # Setup calculator if not already set. + # This is important as usage on a server would otherwise cause an initialization with every call so that nothing is gained. + # Catch the warning that the `weights_only` bool is not set during `torch.load`. This will not affect the calculation. + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected", + ) + self.set_calculator( + suite=suite, + model=args_parsed["model"], + dispersion=dispersion, + damping=args_parsed["damping"], + dispersion_xc=dispersion_xc, + dispersion_cutoff=dispersion_cutoff, + device=args_parsed["device"], + default_dtype=args_parsed["default_dtype"], + head=args_parsed["head"], + ) + + # process the XYZ file + atom_types, coordinates = xyzfile_to_at_coord(calc_data.xyzfile) + + # run mace + energy, gradient = self.run_mace( + atom_types=atom_types, coordinates=coordinates, calc_data=calc_data + ) + + return energy, gradient + + +def main() -> None: + """ + Main routine for execution. + """ + calculator = MaceCalc() + inputfile, args, args_not_parsed = calculator.parse_args() + calculator.run(inputfile=inputfile, args_parsed=args, args_not_parsed=args_not_parsed) + + +# Python entry point +if __name__ == "__main__": + main() diff --git a/src/oet/core/base_calc.py b/src/oet/core/base_calc.py index 373cd01..2f4cc09 100644 --- a/src/oet/core/base_calc.py +++ b/src/oet/core/base_calc.py @@ -39,6 +39,7 @@ "aenet": ("oet.calculator.aenet", "AenetCalc"), "aimnet2": ("oet.calculator.aimnet2", "Aimnet2Calc"), "gxtb": ("oet.calculator.gxtb", "GxtbCalc"), + "mace": ("oet.calculator.mace", "MaceCalc"), "mlatom": ("oet.calculator.mlatom", "MlatomCalc"), "mopac": ("oet.calculator.mopac", "MopacCalc"), "uma": ("oet.calculator.uma", "UmaCalc"), diff --git a/tests/g-xtb/test_gxtb.py b/tests/g-xtb/test_gxtb.py index 709d008..5aba57a 100644 --- a/tests/g-xtb/test_gxtb.py +++ b/tests/g-xtb/test_gxtb.py @@ -13,7 +13,7 @@ # Path to the scripts, adjust if needed. gxtb_script_path = ROOT_DIR / "../../bin/oet_gxtb" -# Leave uma_executable_path empty, if gxtb from system path should be called +# Leave gxtb_executable_path empty, if gxtb from system path should be called gxtb_executable_path = "" diff --git a/tests/mace/test_mace_client.py b/tests/mace/test_mace_client.py new file mode 100644 index 0000000..2260ae0 --- /dev/null +++ b/tests/mace/test_mace_client.py @@ -0,0 +1,179 @@ +import os +import signal +import subprocess +import time +import unittest + +from oet import ROOT_DIR +from oet.core.test_utilities import ( + OH, + WATER, + get_filenames, + read_result_file, + run_wrapper, + write_input_file, + write_xyz_file, +) + +# Path to the script, adjust if needed. +mace_server_path = ROOT_DIR / "../../bin/oet_server" +mace_client_path = ROOT_DIR / "../../bin/oet_client" +# Default maximum time (in sec) to download the model files if not present +timeout = 600 +# Default ID and port of server. Change if needed +id_port = "127.0.0.1:9000" + + +def run_mace(inputfile: str, output_file: str, args: list[str]) -> None: + # Run the wrapper with an increased timeout as loading the MACE model files might take a while + args.extend(["--bind", id_port]) + run_wrapper( + inputfile=inputfile, + script_path=mace_client_path, + outfile=output_file, + timeout=30, + args=args, + ) + + +class MACETests(unittest.TestCase): + @classmethod + def setUpClass(cls): + """ + Test starting the server + """ + print("Starting the server. A detailed server log can be found on file server.out") + with open("server.out", "a") as f: + cls.server = subprocess.Popen( + [mace_server_path, "mace", "--bind", id_port, "--nthreads", "2"], + stdout=f, + stderr=subprocess.STDOUT, + preexec_fn=os.setsid, + ) + # Wait a little to make sure it is setup + time.sleep(5) + + @classmethod + def tearDownClass(cls): + """ + Shut the server at the end + """ + print("Killing the server.") + os.killpg(os.getpgid(cls.server.pid), signal.SIGTERM) + cls.server.wait(timeout=10) + + def test_H2O_engrad(self): + xyz_file, input_file, engrad_out, output_file = get_filenames("H2O") + + write_xyz_file(xyz_file, WATER) + write_input_file( + filename=input_file, + xyz_filename=xyz_file, + charge=0, + multiplicity=1, + ncores=2, + do_gradient=1, + ) + args = ["-s", "mace-mp", "-m", "medium", "--head", "mh0"] + run_mace(input_file, output_file, args) + expected_num_atoms = 3 + expected_energy = -5.203530407103e-01 + expected_gradients = [ + 2.897306550196e-03, + 2.144325522181e-03, + -1.515869870431e-03, + -7.767454655073e-04, + -3.378720118971e-03, + -1.222818329722e-03, + -2.120561084688e-03, + 1.234394596790e-03, + 2.738688200153e-03, + ] + + try: + num_atoms, energy, gradients = read_result_file(engrad_out) + except Exception as e: + raise FileNotFoundError( + f"Error wrapper outputfile not found. Check {output_file} for details" + ) from e + + self.assertEqual(num_atoms, expected_num_atoms) + self.assertAlmostEqual(energy, expected_energy, places=7) + for g1, g2 in zip(gradients, expected_gradients): + self.assertAlmostEqual(g1, g2, places=7) + + def test_OH_anion_eng_grad(self): + xyz_file, input_file, engrad_out, output_file = get_filenames("OH_anion") + write_xyz_file(xyz_file, OH) + write_input_file( + filename=input_file, + xyz_filename=xyz_file, + charge=-1, + multiplicity=1, + ncores=2, + do_gradient=1, + ) + args = ["-s", "omol"] + run_mace(input_file, output_file, args) + expected_num_atoms = 2 + expected_energy = -7.580657311911e01 + expected_gradients = [ + -1.082263844645e-03, + -3.483610415645e-03, + -9.925303607707e-04, + 1.082263844645e-03, + 3.483610415645e-03, + 9.925303607707e-04, + ] + + try: + num_atoms, energy, gradients = read_result_file(engrad_out) + except Exception as e: + raise FileNotFoundError( + f"Error wrapper outputfile not found. Check {output_file} for details" + ) from e + + self.assertEqual(num_atoms, expected_num_atoms) + self.assertAlmostEqual(energy, expected_energy, places=7) + for g1, g2 in zip(gradients, expected_gradients): + self.assertAlmostEqual(g1, g2, places=7) + + def test_OH_rad_eng_grad(self): + xyz_file, input_file, engrad_out, output_file = get_filenames("OH_rad") + write_xyz_file(xyz_file, OH) + write_input_file( + filename=input_file, + xyz_filename=xyz_file, + charge=0, + multiplicity=2, + ncores=2, + do_gradient=1, + ) + # Test the defaults (no arguments) + run_mace(input_file, output_file, args=[]) + expected_num_atoms = 2 + expected_energy = -7.574239878105e01 + expected_gradients = [ + 1.056554797887e-03, + 3.400856384066e-03, + 9.689529220712e-04, + -1.056554797887e-03, + -3.400856384066e-03, + -9.689529220712e-04, + ] + + try: + num_atoms, energy, gradients = read_result_file(engrad_out) + except Exception as e: + raise FileNotFoundError( + f"Error wrapper outputfile not found. Check {output_file} for details" + ) from e + + self.assertEqual(num_atoms, expected_num_atoms) + self.assertAlmostEqual(energy, expected_energy, places=7) + for g1, g2 in zip(gradients, expected_gradients): + self.assertAlmostEqual(g1, g2, places=7) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/mace/test_mace_standalone.py b/tests/mace/test_mace_standalone.py new file mode 100644 index 0000000..5d353b6 --- /dev/null +++ b/tests/mace/test_mace_standalone.py @@ -0,0 +1,143 @@ +import unittest + +from oet import ROOT_DIR +from oet.core.test_utilities import ( + OH, + WATER, + get_filenames, + read_result_file, + run_wrapper, + write_input_file, + write_xyz_file, +) + +# Path to the script, adjust if needed. +mace_script_path = ROOT_DIR / "../../bin/oet_mace" + + +def run_mace(inputfile: str, output_file: str, args: list[str] | None = None) -> None: + # Run the wrapper with an increased timeout as loading the MACE model files might take a while + run_wrapper( + inputfile=inputfile, + script_path=mace_script_path, + outfile=output_file, + timeout=30, + args=args, + ) + + +class MACETests(unittest.TestCase): + def test_H2O_engrad(self): + xyz_file, input_file, engrad_out, output_file = get_filenames("H2O") + + write_xyz_file(xyz_file, WATER) + write_input_file( + filename=input_file, + xyz_filename=xyz_file, + charge=0, + multiplicity=1, + ncores=2, + do_gradient=1, + ) + args = ["-s", "mace-mp", "-m", "medium", "--head", "mh0"] + run_mace(input_file, output_file, args) + expected_num_atoms = 3 + expected_energy = -5.203530407103e-01 + expected_gradients = [ + 2.897306550196e-03, + 2.144325522181e-03, + -1.515869870431e-03, + -7.767454655073e-04, + -3.378720118971e-03, + -1.222818329722e-03, + -2.120561084688e-03, + 1.234394596790e-03, + 2.738688200153e-03, + ] + + try: + num_atoms, energy, gradients = read_result_file(engrad_out) + except Exception as e: + raise FileNotFoundError( + f"Error wrapper outputfile not found. Check {output_file} for details" + ) from e + + self.assertEqual(num_atoms, expected_num_atoms) + self.assertAlmostEqual(energy, expected_energy, places=7) + for g1, g2 in zip(gradients, expected_gradients): + self.assertAlmostEqual(g1, g2, places=7) + + def test_OH_anion_eng_grad(self): + xyz_file, input_file, engrad_out, output_file = get_filenames("OH_anion") + write_xyz_file(xyz_file, OH) + write_input_file( + filename=input_file, + xyz_filename=xyz_file, + charge=-1, + multiplicity=1, + ncores=2, + do_gradient=1, + ) + args = ["-s", "omol"] + run_mace(input_file, output_file, args) + expected_num_atoms = 2 + expected_energy = -7.580657311911e01 + expected_gradients = [ + -1.082263844645e-03, + -3.483610415645e-03, + -9.925303607707e-04, + 1.082263844645e-03, + 3.483610415645e-03, + 9.925303607707e-04, + ] + + try: + num_atoms, energy, gradients = read_result_file(engrad_out) + except Exception as e: + raise FileNotFoundError( + f"Error wrapper outputfile not found. Check {output_file} for details" + ) from e + + self.assertEqual(num_atoms, expected_num_atoms) + self.assertAlmostEqual(energy, expected_energy, places=7) + for g1, g2 in zip(gradients, expected_gradients): + self.assertAlmostEqual(g1, g2, places=7) + + def test_OH_rad_eng_grad(self): + xyz_file, input_file, engrad_out, output_file = get_filenames("OH_rad") + write_xyz_file(xyz_file, OH) + write_input_file( + filename=input_file, + xyz_filename=xyz_file, + charge=0, + multiplicity=2, + ncores=2, + do_gradient=1, + ) + run_mace(input_file, output_file) + expected_num_atoms = 2 + expected_energy = -7.574239878105e01 + expected_gradients = [ + 1.056554797887e-03, + 3.400856384066e-03, + 9.689529220712e-04, + -1.056554797887e-03, + -3.400856384066e-03, + -9.689529220712e-04, + ] + + try: + num_atoms, energy, gradients = read_result_file(engrad_out) + except Exception as e: + raise FileNotFoundError( + f"Error wrapper outputfile not found. Check {output_file} for details" + ) from e + + self.assertEqual(num_atoms, expected_num_atoms) + self.assertAlmostEqual(energy, expected_energy, places=7) + for g1, g2 in zip(gradients, expected_gradients): + self.assertAlmostEqual(g1, g2, places=7) + + +if __name__ == "__main__": + unittest.main()