Add files via upload

mace-bench/3rdparty/mace/mace/__init__.py

+import os
+
+from .__version__ import __version__
+
+os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"

mace-bench/3rdparty/mace/mace/__version__.py

+__version__ = "0.3.13"
+
+__all__ = ["__version__"]

mace-bench/3rdparty/mace/mace/calculators/__init__.py

+from .foundations_models import mace_anicc, mace_mp, mace_off
+from .lammps_mace import LAMMPS_MACE
+from .mace import MACECalculator
+
+__all__ = [
+    "MACECalculator",
+    "LAMMPS_MACE",
+    "mace_mp",
+    "mace_off",
+    "mace_anicc",
+]

mace-bench/3rdparty/mace/mace/calculators/foundations_models.py

+import os
+import urllib.request
+from pathlib import Path
+from typing import Union
+
+import torch
+from ase import units
+from ase.calculators.mixing import SumCalculator
+
+from .mace import MACECalculator
+
+module_dir = os.path.dirname(__file__)
+local_model_path = os.path.join(
+    module_dir, "foundations_models/mace-mpa-0-medium.model"
+)
+
+
+def download_mace_mp_checkpoint(model: Union[str, Path] = None) -> str:
+    """
+    Downloads or locates the MACE-MP checkpoint file.
+
+    Args:
+        model (str, optional): Path to the model or size specification.
+            Defaults to None which uses the medium model.
+
+    Returns:
+        str: Path to the downloaded (or cached, if previously loaded) checkpoint file.
+    """
+    if model in (None, "medium-mpa-0") and os.path.isfile(local_model_path):
+        return local_model_path
+
+    urls = {
+        "small": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model",
+        "medium": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-03-mace-128-L1_epoch-199.model",
+        "large": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/MACE_MPtrj_2022.9.model",
+        "small-0b": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b/mace_agnesi_small.model",
+        "medium-0b": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b/mace_agnesi_medium.model",
+        "small-0b2": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b2/mace-small-density-agnesi-stress.model",
+        "medium-0b2": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b2/mace-medium-density-agnesi-stress.model",
+        "large-0b2": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b2/mace-large-density-agnesi-stress.model",
+        "medium-0b3": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b3/mace-mp-0b3-medium.model",
+        "medium-mpa-0": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mpa_0/mace-mpa-0-medium.model",
+        "medium-omat-0": "https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/mace-omat-0-medium.model",
+        "mace-matpes-pbe-0": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_matpes_0/MACE-matpes-pbe-omat-ft.model",
+        "mace-matpes-r2scan-0": "https://github.com/ACEsuit/mace-foundations/releases/download/mace_matpes_0/MACE-matpes-r2scan-omat-ft.model",
+    }
+
+    checkpoint_url = (
+        urls.get(model, urls["medium-mpa-0"])
+        if model
+        in (
+            None,
+            "small",
+            "medium",
+            "large",
+            "small-0b",
+            "medium-0b",
+            "small-0b2",
+            "medium-0b2",
+            "large-0b2",
+            "medium-0b3",
+            "medium-mpa-0",
+            "medium-omat-0",
+        )
+        else model
+    )
+
+    if checkpoint_url == urls["medium-mpa-0"]:
+        print(
+            "Using medium MPA-0 model as default MACE-MP model, to use previous (before 3.10) default model please specify 'medium' as model argument"
+        )
+    ASL_checkpoint_urls = {
+        urls["medium-omat-0"],
+        urls["mace-matpes-pbe-0"],
+        urls["mace-matpes-r2scan-0"],
+    }
+    if checkpoint_url in ASL_checkpoint_urls:
+        print(
+            "Using model under Academic Software License (ASL) license, see https://github.com/gabor1/ASL \n To use this model you accept the terms of the license."
+        )
+
+    cache_dir = os.path.expanduser("~/.cache/mace")
+    checkpoint_url_name = "".join(
+        c for c in os.path.basename(checkpoint_url) if c.isalnum() or c in "_"
+    )
+    cached_model_path = f"{cache_dir}/{checkpoint_url_name}"
+
+    if not os.path.isfile(cached_model_path):
+        os.makedirs(cache_dir, exist_ok=True)
+        print(f"Downloading MACE model from {checkpoint_url!r}")
+        _, http_msg = urllib.request.urlretrieve(checkpoint_url, cached_model_path)
+        if "Content-Type: text/html" in http_msg:
+            raise RuntimeError(
+                f"Model download failed, please check the URL {checkpoint_url}"
+            )
+        print(f"Cached MACE model to {cached_model_path}")
+
+    return cached_model_path
+
+
+def mace_mp(
+    model: Union[str, Path] = None,
+    device: str = "",
+    default_dtype: str = "float32",
+    dispersion: bool = False,
+    damping: str = "bj",  # choices: ["zero", "bj", "zerom", "bjm"]
+    dispersion_xc: str = "pbe",
+    dispersion_cutoff: float = 40.0 * units.Bohr,
+    return_raw_model: bool = False,
+    **kwargs,
+) -> MACECalculator:
+    """
+    Constructs a MACECalculator with a pretrained model based on the Materials Project (89 elements).
+    The model is released under the MIT license. See https://github.com/ACEsuit/mace-mp for all models.
+    Note:
+        If you are using this function, please cite the relevant paper for the Materials Project,
+        any paper associated with the MACE model, and also the following:
+        - MACE-MP by Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M. Elena,
+            Dávid P. Kovács, Janosh Riebesell, et al., 2023, arXiv:2401.00096
+        - MACE-Universal by Yuan Chiang, 2023, Hugging Face, Revision e5ebd9b,
+            DOI: 10.57967/hf/1202, URL: https://huggingface.co/cyrusyc/mace-universal
+        - Matbench Discovery by Janosh Riebesell, Rhys EA Goodall, Philipp Benner, Yuan Chiang,
+            Alpha A Lee, Anubhav Jain, Kristin A Persson, 2023, arXiv:2308.14920
+
+    Args:
+        model (str, optional): Path to the model. Defaults to None which first checks for
+            a local model and then downloads the default model from figshare. Specify "small",
+            "medium" or "large" to download a smaller or larger model from figshare.
+        device (str, optional): Device to use for the model. Defaults to "cuda" if available.
+        default_dtype (str, optional): Default dtype for the model. Defaults to "float32".
+        dispersion (bool, optional): Whether to use D3 dispersion corrections. Defaults to False.
+        damping (str): The damping function associated with the D3 correction. Defaults to "bj" for D3(BJ).
+        dispersion_xc (str, optional): Exchange-correlation functional for D3 dispersion corrections.
+        dispersion_cutoff (float, optional): Cutoff radius in Bohr for D3 dispersion corrections.
+        return_raw_model (bool, optional): Whether to return the raw model or an ASE calculator. Defaults to False.
+        **kwargs: Passed to MACECalculator and TorchDFTD3Calculator.
+
+    Returns:
+        MACECalculator: trained on the MPtrj dataset (unless model otherwise specified).
+    """
+    try:
+        if model in (
+            None,
+            "small",
+            "medium",
+            "large",
+            "medium-mpa-0",
+            "small-0b",
+            "medium-0b",
+            "small-0b2",
+            "medium-0b2",
+            "medium-0b3",
+            "large-0b2",
+            "medium-omat-0",
+        ) or str(model).startswith("https:"):
+            model_path = download_mace_mp_checkpoint(model)
+            print(f"Using Materials Project MACE for MACECalculator with {model_path}")
+        else:
+            if not Path(model).exists():
+                raise FileNotFoundError(f"{model} not found locally")
+            model_path = model
+    except Exception as exc:
+        raise RuntimeError("Model download failed and no local model found") from exc
+
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    if default_dtype == "float64":
+        print(
+            "Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization."
+        )
+    if default_dtype == "float32":
+        print(
+            "Using float32 for MACECalculator, which is faster but less accurate. Recommended for MD. Use float64 for geometry optimization."
+        )
+
+    if return_raw_model:
+        return torch.load(model_path, map_location=device)
+
+    mace_calc = MACECalculator(
+        model_paths=model_path, device=device, default_dtype=default_dtype, **kwargs
+    )
+
+    if not dispersion:
+        return mace_calc
+
+    try:
+        from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
+    except ImportError as exc:
+        raise RuntimeError(
+            "Please install torch-dftd to use dispersion corrections (see https://github.com/pfnet-research/torch-dftd)"
+        ) from exc
+
+    print("Using TorchDFTD3Calculator for D3 dispersion corrections")
+    dtype = torch.float32 if default_dtype == "float32" else torch.float64
+    d3_calc = TorchDFTD3Calculator(
+        device=device,
+        damping=damping,
+        dtype=dtype,
+        xc=dispersion_xc,
+        cutoff=dispersion_cutoff,
+        **kwargs,
+    )
+
+    return SumCalculator([mace_calc, d3_calc])
+
+
+def mace_off(
+    model: Union[str, Path] = None,
+    device: str = "",
+    default_dtype: str = "float64",
+    return_raw_model: bool = False,
+    **kwargs,
+) -> MACECalculator:
+    """
+    Constructs a MACECalculator with a pretrained model based on the MACE-OFF23 models.
+    The model is released under the ASL license.
+    Note:
+        If you are using this function, please cite the relevant paper by Kovacs et.al., arXiv:2312.15211
+
+    Args:
+        model (str, optional): Path to the model. Defaults to None which first checks for
+            a local model and then downloads the default medium model from https://github.com/ACEsuit/mace-off.
+            Specify "small", "medium" or "large" to download a smaller or larger model.
+        device (str, optional): Device to use for the model. Defaults to "cuda".
+        default_dtype (str, optional): Default dtype for the model. Defaults to "float64".
+        return_raw_model (bool, optional): Whether to return the raw model or an ASE calculator. Defaults to False.
+        **kwargs: Passed to MACECalculator.
+
+    Returns:
+        MACECalculator: trained on the MACE-OFF23 dataset
+    """
+    try:
+        if model in (None, "small", "medium", "large") or str(model).startswith(
+            "https:"
+        ):
+            urls = dict(
+                small="https://github.com/ACEsuit/mace-off/blob/main/mace_off23/MACE-OFF23_small.model?raw=true",
+                medium="https://github.com/ACEsuit/mace-off/raw/main/mace_off23/MACE-OFF23_medium.model?raw=true",
+                large="https://github.com/ACEsuit/mace-off/blob/main/mace_off23/MACE-OFF23_large.model?raw=true",
+            )
+            checkpoint_url = (
+                urls.get(model, urls["medium"])
+                if model in (None, "small", "medium", "large")
+                else model
+            )
+            cache_dir = os.path.expanduser("~/.cache/mace")
+            checkpoint_url_name = os.path.basename(checkpoint_url).split("?")[0]
+            cached_model_path = f"{cache_dir}/{checkpoint_url_name}"
+            if not os.path.isfile(cached_model_path):
+                os.makedirs(cache_dir, exist_ok=True)
+                # download and save to disk
+                print(f"Downloading MACE model from {checkpoint_url!r}")
+                print(
+                    "The model is distributed under the Academic Software License (ASL) license, see https://github.com/gabor1/ASL \n To use the model you accept the terms of the license."
+                )
+                print(
+                    "ASL is based on the Gnu Public License, but does not permit commercial use"
+                )
+                urllib.request.urlretrieve(checkpoint_url, cached_model_path)
+                print(f"Cached MACE model to {cached_model_path}")
+            model = cached_model_path
+            msg = f"Using MACE-OFF23 MODEL for MACECalculator with {model}"
+            print(msg)
+        else:
+            if not Path(model).exists():
+                raise FileNotFoundError(f"{model} not found locally")
+    except Exception as exc:
+        raise RuntimeError("Model download failed and no local model found") from exc
+
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+    if return_raw_model:
+        return torch.load(model, map_location=device)
+
+    if default_dtype == "float64":
+        print(
+            "Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization."
+        )
+    if default_dtype == "float32":
+        print(
+            "Using float32 for MACECalculator, which is faster but less accurate. Recommended for MD. Use float64 for geometry optimization."
+        )
+    mace_calc = MACECalculator(
+        model_paths=model, device=device, default_dtype=default_dtype, **kwargs
+    )
+    return mace_calc
+
+
+def mace_anicc(
+    device: str = "cuda",
+    model_path: str = None,
+    return_raw_model: bool = False,
+) -> MACECalculator:
+    """
+    Constructs a MACECalculator with a pretrained model based on the ANI (H, C, N, O).
+    The model is released under the MIT license.
+    Note:
+        If you are using this function, please cite the relevant paper associated with the MACE model, ANI dataset, and also the following:
+        - "Evaluation of the MACE Force Field Architecture by Dávid Péter Kovács, Ilyes Batatia, Eszter Sára Arany, and Gábor Csányi, The Journal of Chemical Physics, 2023, URL: https://doi.org/10.1063/5.0155322
+    """
+    if model_path is None:
+        model_path = os.path.join(
+            module_dir, "foundations_models/ani500k_large_CC.model"
+        )
+        print(
+            "Using ANI couple cluster model for MACECalculator, see https://doi.org/10.1063/5.0155322"
+        )
+
+    if not os.path.exists(model_path):
+        model_dir = os.path.dirname(model_path)
+        os.makedirs(model_dir, exist_ok=True)
+
+        # Download the model
+        print(f"Model not found at {model_path}. Downloading...")
+        model_url = "https://github.com/ACEsuit/mace/raw/main/mace/calculators/foundations_models/ani500k_large_CC.model"
+
+        try:
+
+            def report_progress(block_num, block_size, total_size):
+                downloaded = block_num * block_size
+                percent = min(100, downloaded * 100 / total_size)
+                if total_size > 0:
+                    print(
+                        f"\rDownloading model: {percent:.1f}% ({downloaded / 1024 / 1024:.1f} MB / {total_size / 1024 / 1024:.1f} MB)",
+                        end="",
+                    )
+
+            urllib.request.urlretrieve(
+                model_url, model_path, reporthook=report_progress
+            )
+            print("\nDownload complete!")
+
+        except Exception as e:
+            raise RuntimeError(f"Failed to download model: {e}") from e
+
+    if return_raw_model:
+        return torch.load(model_path, map_location=device)
+    return MACECalculator(
+        model_paths=model_path, device=device, default_dtype="float64"
+    )

mace-bench/3rdparty/mace/mace/calculators/lammps_mace.py

+from typing import Dict, List, Optional
+
+import torch
+from e3nn.util.jit import compile_mode
+
+from mace.tools.scatter import scatter_sum
+
+
+@compile_mode("script")
+class LAMMPS_MACE(torch.nn.Module):
+    def __init__(self, model, **kwargs):
+        super().__init__()
+        self.model = model
+        self.register_buffer("atomic_numbers", model.atomic_numbers)
+        self.register_buffer("r_max", model.r_max)
+        self.register_buffer("num_interactions", model.num_interactions)
+        if not hasattr(model, "heads"):
+            model.heads = [None]
+        self.register_buffer(
+            "head",
+            torch.tensor(
+                self.model.heads.index(kwargs.get("head", self.model.heads[-1])),
+                dtype=torch.long,
+            ).unsqueeze(0),
+        )
+
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def forward(
+        self,
+        data: Dict[str, torch.Tensor],
+        local_or_ghost: torch.Tensor,
+        compute_virials: bool = False,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        num_graphs = data["ptr"].numel() - 1
+        compute_displacement = False
+        if compute_virials:
+            compute_displacement = True
+        data["head"] = self.head
+        out = self.model(
+            data,
+            training=False,
+            compute_force=False,
+            compute_virials=False,
+            compute_stress=False,
+            compute_displacement=compute_displacement,
+        )
+        node_energy = out["node_energy"]
+        if node_energy is None:
+            return {
+                "total_energy_local": None,
+                "node_energy": None,
+                "forces": None,
+                "virials": None,
+            }
+        positions = data["positions"]
+        displacement = out["displacement"]
+        forces: Optional[torch.Tensor] = torch.zeros_like(positions)
+        virials: Optional[torch.Tensor] = torch.zeros_like(data["cell"])
+        # accumulate energies of local atoms
+        node_energy_local = node_energy * local_or_ghost
+        total_energy_local = scatter_sum(
+            src=node_energy_local, index=data["batch"], dim=-1, dim_size=num_graphs
+        )
+        # compute partial forces and (possibly) partial virials
+        grad_outputs: List[Optional[torch.Tensor]] = [
+            torch.ones_like(total_energy_local)
+        ]
+        if compute_virials and displacement is not None:
+            forces, virials = torch.autograd.grad(
+                outputs=[total_energy_local],
+                inputs=[positions, displacement],
+                grad_outputs=grad_outputs,
+                retain_graph=False,
+                create_graph=False,
+                allow_unused=True,
+            )
+            if forces is not None:
+                forces = -1 * forces
+            else:
+                forces = torch.zeros_like(positions)
+            if virials is not None:
+                virials = -1 * virials
+            else:
+                virials = torch.zeros_like(displacement)
+        else:
+            forces = torch.autograd.grad(
+                outputs=[total_energy_local],
+                inputs=[positions],
+                grad_outputs=grad_outputs,
+                retain_graph=False,
+                create_graph=False,
+                allow_unused=True,
+            )[0]
+            if forces is not None:
+                forces = -1 * forces
+            else:
+                forces = torch.zeros_like(positions)
+        return {
+            "total_energy_local": total_energy_local,
+            "node_energy": node_energy,
+            "forces": forces,
+            "virials": virials,
+        }

mace-bench/3rdparty/mace/mace/calculators/lammps_mliap_mace.py

+import logging
+import os
+import sys
+import time
+from contextlib import contextmanager
+from typing import Dict, Tuple
+
+import torch
+from ase.data import chemical_symbols
+from e3nn.util.jit import compile_mode
+
+try:
+    from lammps.mliap.mliap_unified_abc import MLIAPUnified
+except ImportError:
+
+    class MLIAPUnified:
+        def __init__(self):
+            pass
+
+
+class MACELammpsConfig:
+    """Configuration settings for MACE-LAMMPS integration."""
+
+    def __init__(self):
+        self.debug_time = self._get_env_bool("MACE_TIME", False)
+        self.debug_profile = self._get_env_bool("MACE_PROFILE", False)
+        self.profile_start_step = int(os.environ.get("MACE_PROFILE_START", "5"))
+        self.profile_end_step = int(os.environ.get("MACE_PROFILE_END", "10"))
+        self.allow_cpu = self._get_env_bool("MACE_ALLOW_CPU", False)
+        self.force_cpu = self._get_env_bool("MACE_FORCE_CPU", False)
+
+    @staticmethod
+    def _get_env_bool(var_name: str, default: bool) -> bool:
+        return os.environ.get(var_name, str(default)).lower() in (
+            "true",
+            "1",
+            "t",
+            "yes",
+        )
+
+
+@contextmanager
+def timer(name: str, enabled: bool = True):
+    """Context manager for timing code blocks."""
+    if not enabled:
+        yield
+        return
+
+    start = time.perf_counter()
+    try:
+        yield
+    finally:
+        elapsed = time.perf_counter() - start
+        logging.info(f"Timer - {name}: {elapsed*1000:.3f} ms")
+
+
+@compile_mode("script")
+class MACEEdgeForcesWrapper(torch.nn.Module):
+    """Wrapper that adds per-pair force computation to a MACE model."""
+
+    def __init__(self, model: torch.nn.Module, **kwargs):
+        super().__init__()
+        self.model = model
+        self.register_buffer("atomic_numbers", model.atomic_numbers)
+        self.register_buffer("r_max", model.r_max)
+        self.register_buffer("num_interactions", model.num_interactions)
+
+        if not hasattr(model, "heads"):
+            model.heads = ["Default"]
+
+        head_name = kwargs.get("head", model.heads[-1])
+        head_idx = model.heads.index(head_name)
+        self.register_buffer("head", torch.tensor([head_idx], dtype=torch.long))
+
+        for p in self.model.parameters():
+            p.requires_grad = False
+
+    def forward(
+        self, data: Dict[str, torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Compute energies and per-pair forces."""
+        data["head"] = self.head
+
+        out = self.model(
+            data,
+            training=False,
+            compute_force=False,
+            compute_virials=False,
+            compute_stress=False,
+            compute_displacement=False,
+            compute_edge_forces=True,
+            lammps_mliap=True,
+        )
+
+        node_energy = out["node_energy"]
+        pair_forces = out["edge_forces"]
+        total_energy = out["energy"][0]
+
+        if pair_forces is None:
+            pair_forces = torch.zeros_like(data["vectors"])
+
+        return total_energy, node_energy, pair_forces
+
+
+class LAMMPS_MLIAP_MACE(MLIAPUnified):
+    """MACE integration for LAMMPS using the MLIAP interface."""
+
+    def __init__(self, model, **kwargs):
+        super().__init__()
+        self.config = MACELammpsConfig()
+        self.model = MACEEdgeForcesWrapper(model, **kwargs)
+        self.element_types = [chemical_symbols[s] for s in model.atomic_numbers]
+        self.num_species = len(self.element_types)
+        self.rcutfac = 0.5 * float(model.r_max)
+        self.ndescriptors = 1
+        self.nparams = 1
+        self.dtype = model.r_max.dtype
+        self.device = "cpu"
+        self.initialized = False
+        self.step = 0
+
+    def _initialize_device(self, data):
+        using_kokkos = "kokkos" in data.__class__.__module__.lower()
+
+        if using_kokkos and not self.config.force_cpu:
+            device = torch.as_tensor(data.elems).device
+            if device.type == "cpu" and not self.config.allow_cpu:
+                raise ValueError(
+                    "GPU requested but tensor is on CPU. Set MACE_ALLOW_CPU=true to allow CPU computation."
+                )
+        else:
+            device = torch.device("cpu")
+
+        self.device = device
+        self.model = self.model.to(device)
+        logging.info(f"MACE model initialized on device: {device}")
+        self.initialized = True
+
+    def compute_forces(self, data):
+        natoms = data.nlocal
+        ntotal = data.ntotal
+        nghosts = ntotal - natoms
+        npairs = data.npairs
+        species = torch.as_tensor(data.elems, dtype=torch.int64)
+
+        if not self.initialized:
+            self._initialize_device(data)
+
+        self.step += 1
+        self._manage_profiling()
+
+        if natoms == 0 or npairs <= 1:
+            return
+
+        with timer("total_step", enabled=self.config.debug_time):
+            with timer("prepare_batch", enabled=self.config.debug_time):
+                batch = self._prepare_batch(data, natoms, nghosts, species)
+
+            with timer("model_forward", enabled=self.config.debug_time):
+                _, atom_energies, pair_forces = self.model(batch)
+
+                if self.device.type != "cpu":
+                    torch.cuda.synchronize()
+
+            with timer("update_lammps", enabled=self.config.debug_time):
+                self._update_lammps_data(data, atom_energies, pair_forces, natoms)
+
+    def _prepare_batch(self, data, natoms, nghosts, species):
+        """Prepare the input batch for the MACE model."""
+        return {
+            "vectors": torch.as_tensor(data.rij).to(self.dtype).to(self.device),
+            "node_attrs": torch.nn.functional.one_hot(
+                species.to(self.device), num_classes=self.num_species
+            ).to(self.dtype),
+            "edge_index": torch.stack(
+                [
+                    torch.as_tensor(data.pair_j, dtype=torch.int64).to(self.device),
+                    torch.as_tensor(data.pair_i, dtype=torch.int64).to(self.device),
+                ],
+                dim=0,
+            ),
+            "batch": torch.zeros(natoms, dtype=torch.int64, device=self.device),
+            "lammps_class": data,
+            "natoms": (natoms, nghosts),
+        }
+
+    def _update_lammps_data(self, data, atom_energies, pair_forces, natoms):
+        """Update LAMMPS data structures with computed energies and forces."""
+        if self.dtype == torch.float32:
+            pair_forces = pair_forces.double()
+        eatoms = torch.as_tensor(data.eatoms)
+        eatoms.copy_(atom_energies[:natoms])
+        data.energy = torch.sum(atom_energies[:natoms])
+        data.update_pair_forces_gpu(pair_forces)
+
+    def _manage_profiling(self):
+        if not self.config.debug_profile:
+            return
+
+        if self.step == self.config.profile_start_step:
+            logging.info(f"Starting CUDA profiler at step {self.step}")
+            torch.cuda.profiler.start()
+
+        if self.step == self.config.profile_end_step:
+            logging.info(f"Stopping CUDA profiler at step {self.step}")
+            torch.cuda.profiler.stop()
+            logging.info("Profiling complete. Exiting.")
+            sys.exit()
+
+    def compute_descriptors(self, data):
+        pass
+
+    def compute_gradients(self, data):
+        pass

mace-bench/3rdparty/mace/mace/py.typed

+

mace-bench/3rdparty/mace/mace/tools/__init__.py

+from .arg_parser import build_default_arg_parser, build_preprocess_arg_parser
+from .arg_parser_tools import check_args
+from .cg import U_matrix_real
+from .checkpoint import CheckpointHandler, CheckpointIO, CheckpointState
+from .default_keys import DefaultKeys
+from .finetuning_utils import load_foundations, load_foundations_elements
+from .torch_tools import (
+    TensorDict,
+    cartesian_to_spherical,
+    count_parameters,
+    init_device,
+    init_wandb,
+    set_default_dtype,
+    set_seeds,
+    spherical_to_cartesian,
+    to_numpy,
+    to_one_hot,
+    voigt_to_matrix,
+)
+from .train import SWAContainer, evaluate, train
+from .utils import (
+    AtomicNumberTable,
+    MetricsLogger,
+    atomic_numbers_to_indices,
+    compute_c,
+    compute_mae,
+    compute_q95,
+    compute_rel_mae,
+    compute_rel_rmse,
+    compute_rmse,
+    get_atomic_number_table_from_zs,
+    get_tag,
+    setup_logger,
+)
+
+__all__ = [
+    "TensorDict",
+    "AtomicNumberTable",
+    "atomic_numbers_to_indices",
+    "to_numpy",
+    "to_one_hot",
+    "build_default_arg_parser",
+    "check_args",
+    "DefaultKeys",
+    "set_seeds",
+    "init_device",
+    "setup_logger",
+    "get_tag",
+    "count_parameters",
+    "MetricsLogger",
+    "get_atomic_number_table_from_zs",
+    "train",
+    "evaluate",
+    "SWAContainer",
+    "CheckpointHandler",
+    "CheckpointIO",
+    "CheckpointState",
+    "set_default_dtype",
+    "compute_mae",
+    "compute_rel_mae",
+    "compute_rmse",
+    "compute_rel_rmse",
+    "compute_q95",
+    "compute_c",
+    "U_matrix_real",
+    "spherical_to_cartesian",
+    "cartesian_to_spherical",
+    "voigt_to_matrix",
+    "init_wandb",
+    "load_foundations",
+    "load_foundations_elements",
+    "build_preprocess_arg_parser",
+]