"""
Configure MLIP calculators.
Similar in spirit to matcalc and quacc approaches
- https://github.com/materialsvirtuallab/matcalc
- https://github.com/Quantum-Accelerators/quacc.git
"""
from __future__ import annotations
from os import environ
from pathlib import Path
from sys import platform
from typing import TYPE_CHECKING, Any, get_args
from warnings import warn
from ase import units
from ase.calculators.mixing import SumCalculator
from janus_core.helpers.janus_types import Architectures, Devices, PathLike
from janus_core.helpers.utils import none_to_dict
if TYPE_CHECKING:
from ase.calculators.calculator import Calculator
import torch
[docs]
def _set_model(
model: PathLike | None = None,
kwargs: dict[str, Any] | None = None,
) -> PathLike | torch.nn.Module | None:
"""
Set `model`.
Parameters
----------
model
Path to MLIP file.
kwargs
Dictionary of additional keyword arguments passed to the selected calculator.
Returns
-------
PathLike | torch.nn.Module | None
Name of MLIP model, or path to MLIP model file or loaded model.
"""
(kwargs,) = none_to_dict(kwargs)
# kwargs that may be used for `model` for different MLIPs
# Note: "model" for chgnet (but not mace_mp or mace_off) and "potential" may refer
# to loaded PyTorch models
model_kwargs = {
"model_paths",
"potential",
"path",
"model_name",
"checkpoint_path",
"predict_unit",
}
present = kwargs.keys() & model_kwargs
# Use model if specified, but check not also specified via kwargs
if model and present:
raise ValueError(
f"`model` cannot be used in combination with any keys in {model_kwargs}"
)
if len(present) > 1:
# Check at most one suitable kwarg is specified
raise ValueError(f"Only one of keyword in {model_kwargs} can be specified.")
if present:
# Set model from kwargs if any are specified
model = kwargs.pop(present.pop())
# Convert to path if file/directory exists
if isinstance(model, Path | str) and Path(model).expanduser().exists():
return Path(model).expanduser()
return model
[docs]
def _set_no_weights_only_load():
"""Set environment variable to fix models for torch 2.6."""
environ.setdefault("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1")
[docs]
def add_dispersion(
calc: Calculator,
device: Devices = "cpu",
dtype: torch.dtype | None = None,
**kwargs,
) -> SumCalculator:
"""
Add D3 dispersion calculator to existing calculator.
Parameters
----------
calc
Calculator to add D3 correction to.
device
Device to run calculator on. Default is "cpu".
dtype
Calculation precision. Default is current torch dtype.
**kwargs
Additional keyword arguments passed to `TorchDFTD3Calculator`.
Returns
-------
SumCalculator
Configured calculator with D3 dispersion correction added.
"""
try:
from torch import get_default_dtype
from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
except ImportError as err:
raise ImportError("Please install the d3 extra.") from err
dtype = dtype if dtype else get_default_dtype()
d3_calc = TorchDFTD3Calculator(
device=device,
dtype=dtype,
**kwargs,
)
sum_calc = SumCalculator([calc, d3_calc])
# Copy calculator parameters to make more accessible
sum_calc.parameters = calc.parameters
if "arch" in sum_calc.parameters:
sum_calc.parameters["arch"] = sum_calc.parameters["arch"] + "_d3"
return sum_calc
[docs]
def choose_calculator(
arch: Architectures,
device: Devices = "cpu",
model: PathLike | None = None,
dispersion: bool = False,
dispersion_kwargs: dict[str, Any] | None = None,
**kwargs,
) -> Calculator:
"""
Choose MLIP calculator to configure.
Parameters
----------
arch
MLIP architecture.
device
Device to run calculator on. Default is "cpu".
model
MLIP model label, path to model, or loaded model. Default is `None`.
dispersion
Whether to add D3 dispersion.
dispersion_kwargs
Additional keyword arguments for `TorchDFTD3Calculator`. Defaults for mace_mp
are taken from mace_mp's defaults.
**kwargs
Additional keyword arguments passed to the selected calculator.
Returns
-------
Calculator
Configured MLIP calculator.
Raises
------
ModuleNotFoundError
MLIP module not correctly been installed.
ValueError
Invalid architecture specified.
"""
dispersion_kwargs = dispersion_kwargs if dispersion_kwargs else {}
model = _set_model(model, kwargs)
if device not in get_args(Devices):
raise ValueError(f"`device` must be one of: {get_args(Devices)}")
# Fix torch 2.6 (must be before MLIP modules are loaded)
_set_no_weights_only_load()
match arch:
case "mace":
from mace import __version__
from mace.calculators import MACECalculator
# No default `model`
if model is None:
raise ValueError(
f"Please specify `model`, as there is no default model for {arch}"
)
# Default to float64 precision
kwargs.setdefault("default_dtype", "float64")
calculator = MACECalculator(model_paths=model, device=device, **kwargs)
case "mace_mp":
from mace import __version__
from mace.calculators import mace_mp
# Default to "small" model and float64 precision
model = model if model else "small"
kwargs.setdefault("default_dtype", "float64")
# Set mace_mp dispersion defaults
dispersion_kwargs.setdefault("damping", kwargs.pop("damping", "bj"))
dispersion_kwargs.setdefault("xc", kwargs.pop("dispersion_xc", "pbe"))
dispersion_kwargs.setdefault(
"cutoff", kwargs.pop("dispersion_cutoff", 40.0 * units.Bohr)
)
calculator = mace_mp(model=model, device=device, **kwargs)
case "mace_off":
from mace import __version__
from mace.calculators import mace_off
# Default to "small" model and float64 precision
model = model if model else "small"
kwargs.setdefault("default_dtype", "float64")
calculator = mace_off(model=model, device=device, **kwargs)
case "mace_omol":
from mace import __version__
from mace.calculators import mace_omol
# Default to "extra_large" model and float64 precision
model = model if model else "extra_large"
kwargs.setdefault("default_dtype", "float64")
calculator = mace_omol(model=model, device=device, **kwargs)
case "mace_polar":
from mace import __version__
from mace.calculators import mace_polar
# Default to "polar-1-m" model and float64 precision
model = model if model else "polar-1-m"
kwargs.setdefault("default_dtype", "float64")
calculator = mace_polar(model=model, device=device, **kwargs)
case "chgnet":
from chgnet import __version__
from chgnet.model.dynamics import CHGNetCalculator
from chgnet.model.model import CHGNet
import torch
# Set before loading to avoid type mismatches
torch.set_default_dtype(torch.float32)
# Use loaded model (from kwargs) if specified
# Otherwise, load the model if given a path, else use a default model
match model:
case CHGNet():
loaded_model = model
model = "loaded_CHGNet"
case Path():
loaded_model = CHGNet.from_file(model)
case str():
loaded_model = CHGNet.load(model_name=model, use_device=device)
case _:
model = __version__
loaded_model = None
calculator = CHGNetCalculator(
model=loaded_model, use_device=device, **kwargs
)
case "sevennet":
from sevenn import __version__
from sevenn.sevennet_calculator import SevenNetCalculator
import torch
# Set before loading model to avoid type mismatches
torch.set_default_dtype(torch.float32)
match model:
case Path() | str():
model = str(model)
case _:
model = "SevenNet-0_11July2024"
kwargs.setdefault("file_type", "checkpoint")
kwargs.setdefault("sevennet_config", None)
calculator = SevenNetCalculator(model=model, device=device, **kwargs)
case "nequip":
from nequip import __version__
from nequip.ase import NequIPCalculator
# No default `model`
if model is None:
raise ValueError(
f"Please specify `model`, as there is no default model for {arch}"
)
model = str(model)
calculator = NequIPCalculator.from_compiled_model(
compile_path=model, device=device, **kwargs
)
case "deepmd" | "dpa3":
from deepmd import __version__
from deepmd.calculator import DP
if arch == "dpa3":
warn(
"`dpa3` has been deprecated. Please use `deepmd`.",
FutureWarning,
stacklevel=2,
)
# No default `model`
if model is None:
# Pretrained models registered in deepmd-kit. See
# https://docs.deepmodeling.com/projects/deepmd/en/latest/model/pretrained.html
raise ValueError(
"Please specify `model`, as there is no "
f"default model for {arch} "
"e.g. a model downloaded with `dp pretrained download DPA-3.3-1M`"
)
model = str(model)
calculator = DP(model=model, **kwargs)
case "orb":
from orb_models import __version__
from orb_models.forcefield import pretrained
from orb_models.forcefield.inference.calculator import ORBCalculator
from orb_models.forcefield.inference.d3_model import D3SumModel
from orb_models.forcefield.models.conservative_regressor import (
ConservativeForcefieldRegressor,
)
from orb_models.forcefield.models.direct_regressor import (
DirectForcefieldRegressor,
)
model_func = None
# Set by `model_func`, if a pre-trained model label is used
atoms_adapter = kwargs.pop("atoms_adapter", None)
# Options for the calculator, rather than loading the model
calc_kwargs = {
key: kwargs.pop(key)
for key in (
"edge_method",
"max_num_neighbors",
"half_supercell",
"directory",
)
if key in kwargs
}
match model:
case DirectForcefieldRegressor():
loaded_model = model
model = "loaded_DirectForcefieldRegressor"
case ConservativeForcefieldRegressor():
loaded_model = model
model = "loaded_ConservativeForcefieldRegressor"
case D3SumModel():
loaded_model = model
model = "loaded_D3SumModel"
case str() if hasattr(pretrained, model.replace("-", "_")):
model_func = getattr(pretrained, model.replace("-", "_"))
case None:
# Default model
model = "orb_v3_conservative_20_omat"
model_func = getattr(pretrained, model)
case _:
raise ValueError(
"`model` must be a `DirectForcefieldRegressor`, "
"`ConservativeForcefieldRegressor`, `D3SumModel`, or "
"pre-trained model label (e.g. 'orb-v2'), or `None` (uses "
"default, orb_v3_conservative_20_omat)"
)
if model_func:
loaded_model, atoms_adapter = model_func(device=device, **kwargs)
calculator = ORBCalculator(
model=loaded_model,
atoms_adapter=atoms_adapter,
device=device,
**calc_kwargs,
)
case "mattersim":
from mattersim import __version__
from mattersim.forcefield import MatterSimCalculator
from torch.nn import Module
potential = None
match model:
case Module():
potential = model
model = "loaded_Module"
case Path() | str():
model = str(model)
case None:
model = "mattersim-v1.0.0-5M"
calculator = MatterSimCalculator(
potential=potential, load_path=model, device=device, **kwargs
)
case "grace":
from tensorpotential import __version__
from tensorpotential.calculator import grace_fm
# Default model
model = model if model else "GRACE-2L-OMAT"
if isinstance(model, Path):
model = str(model)
calculator = grace_fm(model, **kwargs)
case "upet":
from upet import __version__
from upet.calculator import UPETCalculator
# Default model
model = model if model else "pet-mad-s"
checkpoint_path = None
# Pass checkpoint (local/URL) via `checkpoint_path`
if model.endswith(".ckpt"):
checkpoint_path = model
model = None
calculator = UPETCalculator(
model=model, checkpoint_path=checkpoint_path, device=device, **kwargs
)
case "fairchem":
from fairchem.core import FAIRChemCalculator, __version__, pretrained_mlip
from fairchem.core.units.mlip_unit import MLIPPredictUnit
# fairchem's "default" inference settings compile the model, which
# requires a C++ compiler, typically unavailable on Windows
inference_settings = kwargs.pop(
"inference_settings", "batch" if platform == "win32" else "default"
)
match model:
case MLIPPredictUnit():
predict_unit = model
model = "loaded_Module"
case Path() | str():
predict_unit = pretrained_mlip.get_predict_unit(
model_name=model,
device=device,
inference_settings=inference_settings,
)
case None:
model = "uma-m-1p1"
predict_unit = pretrained_mlip.get_predict_unit(
model_name=model,
device=device,
inference_settings=inference_settings,
)
kwargs.setdefault("task_name", "omat")
calculator = FAIRChemCalculator(
predict_unit=predict_unit,
**kwargs,
)
case _:
raise ValueError(
f"Unrecognized {arch=}. Suported architectures "
f"are {', '.join(Architectures.__args__)}"
)
if isinstance(model, Path):
model = model.as_posix()
calculator.parameters["version"] = __version__
calculator.parameters["arch"] = arch
calculator.parameters["model"] = str(model)
if dispersion:
return add_dispersion(calc=calculator, device=device, **dispersion_kwargs)
return calculator
[docs]
def check_calculator(calc: Calculator, attribute: str) -> None:
"""
Ensure calculator has ability to calculate properties.
If the calculator is a SumCalculator that inlcudes the TorchDFTD3Calculator, this
also sets the relevant function so that the MLIP component of the calculator is
used for properties unrelated to dispersion.
Parameters
----------
calc
ASE Calculator to check.
attribute
Attribute to check calculator for.
"""
# If dispersion added to MLIP calculator, use only MLIP calculator for calculation
if (
isinstance(calc, SumCalculator)
and len(calc.mixer.calcs) == 2
and calc.mixer.calcs[1].name == "TorchDFTD3Calculator"
and hasattr(calc.mixer.calcs[0], attribute)
):
setattr(calc, attribute, getattr(calc.mixer.calcs[0], attribute))
if not hasattr(calc, attribute) or not callable(getattr(calc, attribute)):
raise NotImplementedError(
f"The attached calculator does not currently support {attribute}"
)