aimnet 0.0.1__py3-none-any.whl

aimnet/base.py

@@ -0,0 +1,41 @@
+from typing import ClassVar, Dict, Final
+
+import torch
+from torch import Tensor, nn
+
+from aimnet import nbops
+
+
+class AIMNet2Base(nn.Module):  # pylint: disable=abstract-method
+    """Base class for AIMNet2 models. Implements pre-processing data:
+    converting to right dtype and device, setting nb mode, calculating masks.
+    """
+
+    _required_keys: Final = ["coord", "numbers", "charge"]
+    _required_keys_dtype: Final = [torch.float32, torch.int64, torch.float32]
+    _optional_keys: Final = ["mult", "nbmat", "nbmat_lr", "mol_idx", "shifts", "cell"]
+    _optional_keys_dtype: Final = [torch.float32, torch.int64, torch.int64, torch.int64, torch.float32, torch.float32]
+    __constants__: ClassVar = ["_required_keys", "_required_keys_dtype", "_optional_keys", "_optional_keys_dtype"]
+
+    def __init__(self):
+        super().__init__()
+
+    def _prepare_dtype(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        for k, d in zip(self._required_keys, self._required_keys_dtype):
+            assert k in data, f"Key {k} is required"
+            data[k] = data[k].to(d)
+        for k, d in zip(self._optional_keys, self._optional_keys_dtype):
+            if k in data:
+                data[k] = data[k].to(d)
+        return data
+
+    def prepare_input(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Some sommon operations"""
+        data = self._prepare_dtype(data)
+        data = nbops.set_nb_mode(data)
+        data = nbops.calc_masks(data)
+
+        assert data["charge"].ndim == 1, "Charge should be 1D tensor."
+        if "mult" in data:
+            assert data["mult"].ndim == 1, "Mult should be 1D tensor."
+        return data

aimnet/calculators/__init__.py

@@ -0,0 +1,15 @@
+import importlib.util
+
+from .calculator import AIMNet2Calculator
+
+__all__ = ["AIMNet2Calculator"]
+
+if importlib.util.find_spec("ase") is not None:
+    from .aimnet2ase import AIMNet2ASE
+
+    __all__.append(AIMNet2ASE)  # type: ignore
+
+if importlib.util.find_spec("pysisyphus") is not None:
+    from .aimnet2pysis import AIMNet2Pysis
+
+    __all__.append(AIMNet2Pysis)  # type: ignore

aimnet/calculators/aimnet2ase.py

@@ -0,0 +1,98 @@
+import numpy as np
+import torch
+
+try:
+    from ase.calculators.calculator import Calculator, all_changes  # type: ignore
+except ImportError:
+    raise ImportError("ASE is not installed. Please install ASE to use this module.") from None
+
+from .calculator import AIMNet2Calculator
+
+
+class AIMNet2ASE(Calculator):
+    from typing import ClassVar
+
+    implemented_properties: ClassVar[list[str]] = ["energy", "forces", "free_energy", "charges", "stress"]
+
+    def __init__(self, base_calc: AIMNet2Calculator | str = "aimnet2", charge=0, mult=1):
+        super().__init__()
+        if isinstance(base_calc, str):
+            base_calc = AIMNet2Calculator(base_calc)
+        self.base_calc = base_calc
+        self.charge = charge
+        self.mult = mult
+        self.reset()
+        # list of implemented species
+        if hasattr(base_calc, "implemented_species"):
+            self.implemented_species = base_calc.implemented_species.cpu().numpy()  # type: ignore
+        else:
+            self.implemented_species = None
+
+    def reset(self):
+        super().reset()
+        self._t_numbers = None
+        self._t_charge = None
+        self._t_mult = None
+        self.charge = 0.0
+        self.mult = 1.0
+
+    def set_atoms(self, atoms):
+        if self.implemented_species is not None and not np.in1d(atoms.numbers, self.implemented_species).all():
+            raise ValueError("Some species are not implemented in the AIMNet2Calculator")
+        self.reset()
+        self.atoms = atoms
+
+    def set_charge(self, charge):
+        self.charge = charge
+        self._t_charge = None
+        self.update_tensors()
+
+    def set_mult(self, mult):
+        self.mult = mult
+        self._t_mult = None
+        self.update_tensors()
+
+    def update_tensors(self):
+        if self._t_numbers is None:
+            self._t_numbers = torch.tensor(self.atoms.numbers, dtype=torch.int64, device=self.base_calc.device)
+        if self._t_charge is None:
+            self._t_charge = torch.tensor(self.charge, dtype=torch.float32, device=self.base_calc.device)
+        if self._t_mult is None:
+            self._t_mult = torch.tensor(self.mult, dtype=torch.float32, device=self.base_calc.device)
+
+    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+        if properties is None:
+            properties = ["energy"]
+        super().calculate(atoms, properties, system_changes)
+        self.update_tensors()
+
+        cell = self.atoms.cell.array if self.atoms.cell is not None and self.atoms.pbc.any() else None
+
+        _in = {
+            "coord": torch.tensor(self.atoms.positions, dtype=torch.float32, device=self.base_calc.device),
+            "numbers": self._t_numbers,
+            "charge": self._t_charge,
+            "mult": self._t_mult,
+        }
+
+        _unsqueezed = False
+        if cell is not None:
+            _in["cell"] = cell
+        else:
+            for k, v in _in.items():
+                _in[k] = v.unsqueeze(0)
+            _unsqueezed = True
+
+        results = self.base_calc(_in, forces="forces" in properties, stress="stress" in properties)
+
+        for k, v in results.items():
+            if _unsqueezed:
+                v = v.squeeze(0)
+            results[k] = v.detach().cpu().numpy()  # type: ignore
+
+        self.results["energy"] = results["energy"]
+        self.results["charges"] = results["charges"]
+        if "forces" in properties:
+            self.results["forces"] = results["forces"]
+        if "stress" in properties:
+            self.results["stress"] = results["stress"]

aimnet/calculators/aimnet2pysis.py

@@ -0,0 +1,76 @@
+from typing import ClassVar
+
+import torch
+
+from .calculator import AIMNet2Calculator
+
+try:
+    import pysisyphus.run  # type: ignore
+    from pysisyphus.calculators.Calculator import Calculator  # type: ignore
+    from pysisyphus.constants import ANG2BOHR, AU2EV, BOHR2ANG  # type: ignore
+    from pysisyphus.elem_data import ATOMIC_NUMBERS  # type: ignore
+except ImportError:
+    raise ImportError("Pysisyphus is not installed. Please install Pysisyphus to use this module.") from None
+
+EV2AU = 1 / AU2EV
+
+
+class AIMNet2Pysis(Calculator):
+    implemented_properties: ClassVar = ["energy", "forces", "free_energy", "charges", "stress"]
+
+    def __init__(self, model: AIMNet2Calculator | str = "aimnet2", charge=0, mult=1, **kwargs):
+        super().__init__(charge=charge, mult=mult, **kwargs)
+        if isinstance(model, str):
+            model = AIMNet2Calculator(model)
+        self.model = model
+
+    def _prepere_input(self, atoms, coord):
+        device = self.model.device
+        numbers = torch.as_tensor([ATOMIC_NUMBERS[a.lower()] for a in atoms], device=device)
+        coord = torch.as_tensor(coord, dtype=torch.float, device=device).view(-1, 3) * BOHR2ANG
+        charge = torch.as_tensor([self.charge], dtype=torch.float, device=device)
+        mult = torch.as_tensor([self.mult], dtype=torch.float, device=device)
+        return {"coord": coord, "numbers": numbers, "charge": charge, "mult": mult}
+
+    @staticmethod
+    def _results_get_energy(results):
+        return results["energy"].item() * EV2AU
+
+    @staticmethod
+    def _results_get_forces(results):
+        return (results["forces"].detach() * (EV2AU / ANG2BOHR)).flatten().to(torch.double).cpu().numpy()
+
+    @staticmethod
+    def _results_get_hessian(results):
+        return (
+            (results["hessian"].flatten(0, 1).flatten(-2, -1) * (EV2AU / ANG2BOHR / ANG2BOHR))
+            .to(torch.double)
+            .cpu()
+            .numpy()
+        )
+
+    def get_energy(self, atoms, coords):
+        _in = self._prepere_input(atoms, coords)
+        res = self.model(_in)
+        energy = self._results_get_energy(res)
+        return {"energy": energy}
+
+    def get_forces(self, atoms, coords):
+        _in = self._prepere_input(atoms, coords)
+        res = self.model(_in, forces=True)
+        energy = self._results_get_energy(res)
+        forces = self._results_get_forces(res)
+        return {"energy": energy, "forces": forces}
+
+    def get_hessian(self, atoms, coords):
+        _in = self._prepere_input(atoms, coords)
+        res = self.model(_in, forces=True, hessian=True)
+        energy = self._results_get_energy(res)
+        forces = self._results_get_forces(res)
+        hessian = self._results_get_hessian(res)
+        return {"energy": energy, "forces": forces, "hessian": hessian}
+
+
+def run_pysis():
+    pysisyphus.run.CALC_DICT["aimnet"] = AIMNet2Pysis
+    pysisyphus.run.run()

aimnet/calculators/calculator.py

@@ -0,0 +1,320 @@
+import warnings
+from typing import Any, ClassVar, Dict, Literal
+
+import torch
+from torch import Tensor, nn
+
+from .model_registry import get_model_path
+from .nbmat import TooManyNeighborsError, calc_nbmat
+
+
+class AIMNet2Calculator:
+    """Genegic AIMNet2 calculator
+    A helper class to load AIMNet2 models and perform inference.
+    """
+
+    keys_in: ClassVar[Dict[str, torch.dtype]] = {"coord": torch.float, "numbers": torch.int, "charge": torch.float}
+    keys_in_optional: ClassVar[Dict[str, torch.dtype]] = {
+        "mult": torch.float,
+        "mol_idx": torch.int,
+        "nbmat": torch.int,
+        "nbmat_lr": torch.int,
+        "nb_pad_mask": torch.bool,
+        "nb_pad_mask_lr": torch.bool,
+        "shifts": torch.float,
+        "shifts_lr": torch.float,
+        "cell": torch.float,
+    }
+    keys_out: ClassVar[list[str]] = ["energy", "charges", "forces", "hessian", "stress"]
+    atom_feature_keys: ClassVar[list[str]] = ["coord", "numbers", "charges", "forces"]
+
+    def __init__(self, model: str | nn.Module = "aimnet2", nb_threshold: int = 320):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        if isinstance(model, str):
+            p = get_model_path(model)
+            self.model = torch.jit.load(p, map_location=self.device)
+        elif isinstance(model, nn.Module):
+            self.model = model.to(self.device)
+        else:
+            raise TypeError("Invalid model type/name.")
+
+        self.cutoff = self.model.cutoff
+        self.lr = hasattr(self.model, "cutoff_lr")
+        self.cutoff_lr = getattr(self.model, "cutoff_lr", float("inf"))
+        self.max_density = 0.2
+        self.nb_threshold = nb_threshold
+
+        # indicator if input was flattened
+        self._batch = None
+        self._max_mol_size: int = 0
+        # placeholder for tensors that require grad
+        self._saved_for_grad = {}
+        # set flag of current Coulomb model
+        coul_methods = {getattr(mod, "method", None) for mod in iter_lrcoulomb_mods(self.model)}
+        if len(coul_methods) > 1:
+            raise ValueError("Multiple Coulomb modules found.")
+        if len(coul_methods):
+            self._coulomb_method = coul_methods.pop()
+        else:
+            self._coulomb_method = None
+
+    def __call__(self, *args, **kwargs):
+        return self.eval(*args, **kwargs)
+
+    def set_lrcoulomb_method(
+        self, method: Literal["simple", "dsf", "ewald"], cutoff: float = 15.0, dsf_alpha: float = 0.2
+    ):
+        if method not in ("simple", "dsf", "ewald"):
+            raise ValueError(f"Invalid method: {method}")
+        for mod in iter_lrcoulomb_mods(self.model):
+            mod.method = method  # type: ignore
+            if method == "simple":
+                self.cutoff_lr = float("inf")
+            elif method == "dsf":
+                self.cutoff_lr = cutoff
+                mod.dsf_alpha = dsf_alpha  # type: ignore
+                mod.dsf_rc = cutoff  # type: ignore
+            elif method == "ewald":
+                # current implementaion of Ewald does not use nb mat
+                self.cutoff_lr = cutoff
+        self._coulomb_method = method
+
+    def eval(self, data: Dict[str, Any], forces=False, stress=False, hessian=False) -> Dict[str, Tensor]:
+        data = self.prepare_input(data)
+        if hessian and "mol_idx" in data and data["mol_idx"][-1] > 0:
+            raise NotImplementedError("Hessian calculation is not supported for multiple molecules")
+        data = self.set_grad_tensors(data, forces=forces, stress=stress, hessian=hessian)
+        with torch.jit.optimized_execution(False):  # type: ignore
+            data = self.model(data)
+        data = self.get_derivatives(data, forces=forces, stress=stress, hessian=hessian)
+        data = self.process_output(data)
+        return data
+
+    def prepare_input(self, data: Dict[str, Any]) -> Dict[str, Tensor]:
+        data = self.to_input_tensors(data)
+        data = self.mol_flatten(data)
+        if data.get("cell") is not None:
+            if data["mol_idx"][-1] > 0:
+                raise NotImplementedError("PBC with multiple molecules is not implemented yet.")
+            if self._coulomb_method == "simple":
+                warnings.warn("Switching to DSF Coulomb for PBC", stacklevel=1)
+                self.set_lrcoulomb_method("dsf")
+        if data["coord"].ndim == 2:
+            data = self.make_nbmat(data)
+            data = self.pad_input(data)
+        return data
+
+    def process_output(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        if data["coord"].ndim == 2:
+            data = self.unpad_output(data)
+        data = self.mol_unflatten(data)
+        data = self.keep_only(data)
+        return data
+
+    def to_input_tensors(self, data: Dict[str, Any]) -> Dict[str, Tensor]:
+        ret = {}
+        for k in self.keys_in:
+            if k not in data:
+                raise KeyError(f"Missing key {k} in the input data")
+            # always detach !!
+            ret[k] = torch.as_tensor(data[k], device=self.device, dtype=self.keys_in[k]).detach()
+        for k in self.keys_in_optional:
+            if k in data and data[k] is not None:
+                ret[k] = torch.as_tensor(data[k], device=self.device, dtype=self.keys_in_optional[k]).detach()
+        # convert any scalar tensors to shape (1,) tensors
+        for k, v in ret.items():
+            if v.ndim == 0:
+                ret[k] = v.unsqueeze(0)
+        return ret
+
+    def mol_flatten(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        """Flatten the input data for multiple molecules.
+        Will not flatten for batched input and molecule size below threshold.
+        """
+        ndim = data["coord"].ndim
+        if ndim == 2:
+            # single molecule or already flattened
+            self._batch = None
+            if "mol_idx" not in data:
+                data["mol_idx"] = torch.zeros(data["coord"].shape[0], dtype=torch.long, device=self.device)
+                self._max_mol_size = data["coord"].shape[0]
+            elif data["mol_idx"][-1] == 0:
+                self._max_mol_size = len(data["mol_idx"])
+            else:
+                self._max_mol_size = data["mol_idx"].unique(return_counts=True)[1].max().item()
+
+        elif ndim == 3:
+            # batched input
+            B, N = data["coord"].shape[:2]
+            if self.nb_threshold > N or self.device == "cpu":
+                self._batch = B
+                data["mol_idx"] = torch.repeat_interleave(
+                    torch.arange(0, B, device=self.device), torch.full((B,), N, device=self.device)
+                )
+                for k, v in data.items():
+                    if k in self.atom_feature_keys:
+                        data[k] = v.flatten(0, 1)
+            else:
+                self._batch = None
+            self._max_mol_size = N
+        return data
+
+    def mol_unflatten(self, data: Dict[str, Tensor], batch=None) -> Dict[str, Tensor]:
+        batch = batch if batch is not None else self._batch
+        if batch is not None:
+            for k, v in data.items():
+                if k in self.atom_feature_keys:
+                    data[k] = v.view(batch, -1, *v.shape[1:])
+        return data
+
+    def make_nbmat(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        assert self._max_mol_size > 0, "Molecule size is not set"
+
+        if "cell" in data and data["cell"] is not None:
+            data["coord"] = move_coord_to_cell(data["coord"], data["cell"])
+            cell = data["cell"]
+        else:
+            cell = None
+
+        while True:
+            try:
+                maxnb1 = calc_max_nb(self.cutoff, self.max_density)
+                maxnb2 = calc_max_nb(self.cutoff_lr, self.max_density) if self.lr else None  # type: ignore
+                if cell is None:
+                    maxnb1 = min(maxnb1, self._max_mol_size)
+                    maxnb2 = min(maxnb2, self._max_mol_size) if self.lr else None  # type: ignore
+                maxnb = (maxnb1, maxnb2)
+                nbmat1, nbmat2, shifts1, shifts2 = calc_nbmat(
+                    data["coord"],
+                    (self.cutoff, self.cutoff_lr),
+                    maxnb,  # type: ignore
+                    cell,
+                    data.get("mol_idx"),  # type: ignore
+                )
+                break
+            except TooManyNeighborsError:
+                self.max_density *= 1.2
+                assert self.max_density <= 4, "Something went wrong in nbmat calculation"
+        data["nbmat"] = nbmat1
+        if self.lr:
+            assert nbmat2 is not None
+            data["nbmat_lr"] = nbmat2
+        if cell is not None:
+            assert shifts1 is not None
+            data["shifts"] = shifts1
+            if self.lr:
+                assert shifts2 is not None
+                data["shifts_lr"] = shifts2
+        return data
+
+    def pad_input(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        N = data["nbmat"].shape[0]
+        data["mol_idx"] = maybe_pad_dim0(data["mol_idx"], N, value=data["mol_idx"][-1].item())
+        for k in ("coord", "numbers"):
+            if k in data:
+                data[k] = maybe_pad_dim0(data[k], N)
+        return data
+
+    def unpad_output(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        N = data["nbmat"].shape[0] - 1
+        for k, v in data.items():
+            if k in self.atom_feature_keys:
+                data[k] = maybe_unpad_dim0(v, N)
+        return data
+
+    def set_grad_tensors(self, data: Dict[str, Tensor], forces=False, stress=False, hessian=False) -> Dict[str, Tensor]:
+        self._saved_for_grad = {}
+        if forces or hessian:
+            data["coord"].requires_grad_(True)
+            self._saved_for_grad["coord"] = data["coord"]
+        if stress:
+            assert "cell" in data and data["cell"] is not None, "Stress calculation requires cell"
+            scaling = torch.eye(3, requires_grad=True, dtype=data["cell"].dtype, device=data["cell"].device)
+            data["coord"] = data["coord"] @ scaling
+            data["cell"] = data["cell"] @ scaling
+            self._saved_for_grad["scaling"] = scaling
+        return data
+
+    def keep_only(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        ret = {}
+        for k, v in data.items():
+            if k in self.keys_out or (k.endswith("_std") and k[:-4] in self.keys_out):
+                ret[k] = v
+        return ret
+
+    def get_derivatives(self, data: Dict[str, Tensor], forces=False, stress=False, hessian=False) -> Dict[str, Tensor]:
+        training = getattr(self.model, "training", False)
+        _create_graph = hessian or training
+        x = []
+        if hessian:
+            forces = True
+        if forces and ("forces" not in data or (_create_graph and not data["forces"].requires_grad)):
+            forces = True
+            x.append(self._saved_for_grad["coord"])
+        if stress:
+            x.append(self._saved_for_grad["scaling"])
+        if x:
+            tot_energy = data["energy"].sum()
+            deriv = torch.autograd.grad(tot_energy, x, create_graph=_create_graph)
+            if forces:
+                data["forces"] = -deriv[0]
+            if stress:
+                dedc = deriv[0] if not forces else deriv[1]
+                data["stress"] = dedc / data["cell"].detach().det().abs()
+        if hessian:
+            data["hessian"] = self.calculate_hessian(data["forces"], self._saved_for_grad["coord"])
+        return data
+
+    @staticmethod
+    def calculate_hessian(forces: Tensor, coord: Tensor) -> Tensor:
+        # here forces have shape (N, 3) and coord has shape (N+1, 3)
+        # return hessian with shape (N, 3, N, 3)
+        hessian = -torch.stack([
+            torch.autograd.grad(_f, coord, retain_graph=True)[0] for _f in forces.flatten().unbind()
+        ]).view(-1, 3, coord.shape[0], 3)[:-1, :, :-1, :]
+        return hessian
+
+
+def maybe_pad_dim0(a: Tensor, N: int, value=0.0) -> Tensor:
+    _shape_diff = N - a.shape[0]
+    assert _shape_diff == 0 or _shape_diff == 1, "Invalid shape"
+    if _shape_diff == 1:
+        a = pad_dim0(a, value=value)
+    return a
+
+
+def pad_dim0(a: Tensor, value=0.0) -> Tensor:
+    shapes = [0] * ((a.ndim - 1) * 2) + [0, 1]
+    a = torch.nn.functional.pad(a, shapes, mode="constant", value=value)
+    return a
+
+
+def maybe_unpad_dim0(a: Tensor, N: int) -> Tensor:
+    _shape_diff = a.shape[0] - N
+    assert _shape_diff == 0 or _shape_diff == 1, "Invalid shape"
+    if _shape_diff == 1:
+        a = a[:-1]
+    return a
+
+
+def move_coord_to_cell(coord, cell):
+    coord_f = coord @ cell.inverse()
+    coord_f = coord_f % 1
+    return coord_f @ cell
+
+
+def _named_children_rec(module):
+    if isinstance(module, torch.nn.Module):
+        for name, child in module.named_children():
+            yield name, child
+            yield from _named_children_rec(child)
+
+
+def iter_lrcoulomb_mods(model):
+    for name, module in _named_children_rec(model):
+        if name == "lrcoulomb":
+            yield module
+
+
+def calc_max_nb(cutoff: float, density: float = 0.2) -> int | float:
+    return int(density * 4 / 3 * 3.14159 * cutoff**3) if cutoff < float("inf") else float("inf")

aimnet/calculators/model_registry.py

@@ -0,0 +1,60 @@
+import logging
+import os
+from typing import Dict, Optional
+
+import click
+import requests
+import yaml
+
+logging.basicConfig(level=logging.INFO)
+
+
+def load_model_registry(registry_file: Optional[str] = None) -> Dict[str, str]:
+    registry_file = registry_file or os.path.join(os.path.dirname(__file__), "model_registry.yaml")
+    with open(os.path.join(os.path.dirname(__file__), "model_registry.yaml")) as f:
+        return yaml.load(f, Loader=yaml.SafeLoader)
+
+
+def create_assets_dir():
+    os.makedirs(os.path.join(os.path.dirname(__file__), "assets"), exist_ok=True)
+
+
+def get_registry_model_path(model_name: str) -> str:
+    model_registry = load_model_registry()
+    create_assets_dir()
+    if model_name in model_registry["aliases"]:
+        model_name = model_registry["aliases"][model_name]  # type: ignore
+    if model_name not in model_registry["models"]:
+        raise ValueError(f"Model {model_name} not found in the registry.")
+    cfg = model_registry["models"][model_name]  # type: ignore
+    model_path = _maybe_download_asset(**cfg)  # type: ignore
+    return model_path
+
+
+def _maybe_download_asset(file: str, url: str) -> str:
+    filename = os.path.join(os.path.dirname(__file__), "assets", file)
+    if not os.path.exists(filename):
+        print(f"Downloading {url} -> {filename}")
+        with open(filename, "wb") as f:
+            response = requests.get(url, timeout=60)
+            f.write(response.content)
+    return filename
+
+
+def get_model_path(s: str) -> str:
+    # direct file path
+    if os.path.isfile(s):
+        print("Found model file:", s)
+    else:
+        s = get_registry_model_path(s)
+    return s
+
+
+@click.command(short_help="Clear assets directory.")
+def clear_assets():
+    from glob import glob
+
+    for fil in glob(os.path.join(os.path.dirname(__file__), "assets", "*")):
+        if os.path.isfile(fil):
+            logging.warn(f"Removing {fil}")
+            os.remove(fil)

aimnet/calculators/model_registry.yaml

@@ -0,0 +1,33 @@
+# map file name to url
+models:
+    aimnet2_wb97m_d3_0:
+        file: aimnet2_wb97m_d3_0.jpt
+        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_wb97m_d3_0.jpt
+    aimnet2_wb97m_d3_1:
+        file: aimnet2_wb97m_d3_1.jpt
+        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_wb97m_d3_1.jpt
+    aimnet2_wb97m_d3_2:
+        file: aimnet2_wb97m_d3_2.jpt
+        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_wb97m_d3_2.jpt
+    aimnet2_wb97m_d3_3:
+        file: aimnet2_wb97m_d3_3.jpt
+        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_wb97m_d3_3.jpt
+    aimnet2_b973c_d3_0:
+        file: aimnet2_b973c_d3_0.jpt
+        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_b973c_d3_0.jpt
+    aimnet2_b973c_d3_1:
+        file: aimnet2_b973c_d3_1.jpt
+        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_b973c_d3_1.jpt
+    aimnet2_b973c_d3_2:
+        file: aimnet2_b973c_d3_2.jpt
+        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_b973c_d3_2.jpt
+    aimnet2_b973c_d3_3:
+        file: aimnet2_b973c_d3_3.jpt
+        url: https://storage.googleapis.com/aimnetcentral/AIMNet2/aimnet2_b973c_d3_3.jpt
+
+# map model alias to file name
+aliases:
+    aimnet2: aimnet2_wb97m_d3_0
+    aimnet2_wb97m: aimnet2_wb97m_d3_0
+    aimnet2_b973c: aimnet2_b973c_d3_0
+    aimnet2_qr: aimnet2_qr_v0