update mace ,chgnet io

HojeChun · HojeChun · commit 7385124d9833 · 2025-08-26T23:00:34.000-04:00
diff --git a/nff/io/chgnet.py b/nff/io/chgnet.py
@@ -1,15 +1,131 @@
 """Convert NFF Dataset to CHGNet StructureData"""
 
-from typing import Dict, List
-
+from typing import Dict, List, TYPE_CHECKING
+import functools
+import random
 import torch
-from chgnet.data.dataset import StructureData
+import numpy as np
 from pymatgen.core.structure import Structure
 from pymatgen.io.ase import AseAtomsAdaptor
 
 from nff.data import Dataset
 from nff.io import AtomsBatch
 from nff.utils.cuda import batch_detach, detach
+from chgnet.graph import CrystalGraph, CrystalGraphConverter
+from collections.abc import Sequence
+
+datatype = torch.float32
+
+class StructureData(Dataset):
+    """A simple torch Dataset of structures."""
+
+    def __init__(
+        self,
+        structures: list[Structure],
+        energies: list[float],
+        forces: list[Sequence[Sequence[float]]],
+        stresses: list[Sequence[Sequence[float]]] | None = None,
+        magmoms: list[Sequence[Sequence[float]]] | None = None,
+        structure_ids: list[str] | None = None,
+        graph_converter: CrystalGraphConverter | None = None,
+    ) -> None:
+        """Initialize the dataset.
+
+        Args:
+            structures (list[dict]): pymatgen Structure objects.
+            energies (list[float]): [data_size, 1]
+            forces (list[list[float]]): [data_size, n_atoms, 3]
+            stresses (list[list[float]], optional): [data_size, 3, 3]
+            magmoms (list[list[float]], optional): [data_size, n_atoms, 1]
+            structure_ids (list[str], optional): a list of ids to track the structures
+            graph_converter (CrystalGraphConverter, optional): Converts the structures
+                to graphs. If None, it will be set to CHGNet 0.3.0 converter
+                with AtomGraph cutoff = 6A.
+
+        Raises:
+            RuntimeError: if the length of structures and labels (energies, forces,
+                stresses, magmoms) are not equal.
+        """
+        for idx, struct in enumerate(structures):
+            if not isinstance(struct, Structure):
+                raise ValueError(f"{idx} is not a pymatgen Structure object: {struct}")
+        for name in "energies forces stresses magmoms structure_ids".split():
+            labels = locals()[name]
+            if labels is not None and len(labels) != len(structures):
+                raise RuntimeError(
+                    f"Inconsistent number of structures and labels: "
+                    f"{len(structures)=}, len({name})={len(labels)}"
+                )
+        self.structures = structures
+        self.energies = energies
+        self.forces = forces
+        self.stresses = stresses
+        self.magmoms = magmoms
+        self.structure_ids = structure_ids
+        self.keys = np.arange(len(structures))
+        random.shuffle(self.keys)
+        self.graph_converter = graph_converter or CrystalGraphConverter(
+            atom_graph_cutoff=6, bond_graph_cutoff=3
+        )
+        self.failed_idx: list[int] = []
+        self.failed_graph_id: dict[str, str] = {}
+
+    def __len__(self) -> int:
+        """Get the number of structures in this dataset."""
+        return len(self.keys)
+
+    @functools.cache  # Cache loaded structures
+    def __getitem__(self, idx: int) -> tuple[CrystalGraph, dict]:
+        """Get one graph for a structure in this dataset.
+
+        Args:
+            idx (int): Index of the structure
+
+        Returns:
+            crystal_graph (CrystalGraph): graph of the crystal structure
+            targets (dict): list of targets. i.e. energy, force, stress
+        """
+        if idx not in self.failed_idx:
+            graph_id = self.keys[idx]
+            try:
+                struct = self.structures[graph_id]
+                if self.structure_ids is not None:
+                    mp_id = self.structure_ids[graph_id]
+                else:
+                    mp_id = graph_id
+                crystal_graph = self.graph_converter(
+                    struct, graph_id=graph_id, mp_id=mp_id
+                )
+                targets = {
+                    "e": torch.tensor(self.energies[graph_id], dtype=datatype),
+                    "f": torch.tensor(self.forces[graph_id], dtype=datatype),
+                }
+                if self.stresses is not None:
+                    # Convert VASP stress
+                    targets["s"] = torch.tensor(
+                        self.stresses[graph_id], dtype=datatype
+                    ) * (-0.1)
+                if self.magmoms is not None:
+                    mag = self.magmoms[graph_id]
+                    # use absolute value for magnetic moments
+                    if mag is None:
+                        targets["m"] = None
+                    else:
+                        targets["m"] = torch.abs(torch.tensor(mag, dtype=datatype))
+
+                return crystal_graph, targets
+
+            # Omit structures with isolated atoms. Return another randomly selected
+            # structure
+            except Exception:
+                struct = self.structures[graph_id]
+                self.failed_graph_id[graph_id] = struct.composition.formula
+                self.failed_idx.append(idx)
+                idx = random.randint(0, len(self) - 1)
+                return self.__getitem__(idx)
+        else:
+            idx = random.randint(0, len(self) - 1)
+            return self.__getitem__(idx)
 
 
 def convert_nff_to_chgnet_structure_data(
diff --git a/nff/io/mace.py b/nff/io/mace.py
@@ -1,6 +1,7 @@
 import logging
 import os
 import urllib
+from pathlib import Path
 from collections.abc import Iterable
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -17,6 +18,7 @@
 from e3nn import o3
 from nff.data import Dataset
 from nff.utils.cuda import detach
+from mace.calculators.foundations_models import download_mace_mp_checkpoint, mace_mp_names
 
 # get the path to NFF models dir, which is the parent directory of this file
 module_dir = os.path.abspath(os.path.join(os.path.abspath(__file__), "..", "..", "..", "models"))
@@ -37,7 +39,7 @@ def _check_non_zero(std):
     return std
 
 
-def get_mace_mp_model_path(model: Optional[str] = None, supress_print=True) -> str:
+def get_mace_foundtion_model_path(model: Optional[str] = None, supress_print=True) -> str:
     """Get the default MACE MP model. Replicated from the MACE codebase,
     Copyright (c) 2022 ACEsuit/mace and licensed under the MIT license.
 
@@ -52,32 +54,17 @@ def get_mace_mp_model_path(model: Optional[str] = None, supress_print=True) -> s
     Returns:
         str: path to the model
     """
-    if model in (None, "medium") and os.path.isfile(LOCAL_MODEL_PATH):
-        model_path = LOCAL_MODEL_PATH
+    try:
+        if model in mace_mp_names or str(model).startswith("https:"):
+            model_path = download_mace_mp_checkpoint(model)
+        else:
+            if not Path(model).exists():
+                raise FileNotFoundError(f"{model} not found locally")
+            model_path = model
         if not supress_print:
-            print(f"Using local medium Materials Project MACE model for MACECalculator {model}")
-    elif model in (None, "small", "medium", "large") or str(model).startswith("https:"):
-        try:
-            checkpoint_url = (
-                MACE_URLS.get(model, MACE_URLS["medium"]) if model in (None, "small", "medium", "large") else model
-            )
-            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 "_")
-            model_path = f"{cache_dir}/{checkpoint_url_name}"
-            if not os.path.isfile(model_path):
-                os.makedirs(cache_dir, exist_ok=True)
-                # download and save to disk
-                urllib.request.urlretrieve(checkpoint_url, model_path)
-                if not supress_print:
-                    print(f"Downloading MACE model from {checkpoint_url!r}")
-                    print(f"Cached MACE model to {model_path}")
-            if not supress_print:
-                msg = f"Loading Materials Project MACE with {model_path}"
-                print(msg)
-        except Exception as exc:
-            raise RuntimeError("Model download failed and no local model found") from exc
-    else:
-        raise RuntimeError("Model download failed and no local model found")
+            print(f"Using MACE mdoel with {model_path}")
+    except Exception as exc:
+        raise RuntimeError("Model download failed or no local model found") from exc
 
     return model_path
 
diff --git a/nff/nn/models/mace.py b/nff/nn/models/mace.py
@@ -25,7 +25,7 @@
     NffBatch,
     get_atomic_number_table_from_zs,
     get_init_kwargs_from_model,
-    get_mace_mp_model_path,
+    get_mace_foundtion_model_path,
 )
 
 
@@ -293,8 +293,8 @@ def load_foundations(
         Returns:
             NffScaleMACE: NffScaleMACE foundational model.
         """
-        mace_model_path = get_mace_mp_model_path(model)
-        mace_model = torch.load(mace_model_path, map_location=map_location)
+        mace_model_path = get_mace_foundtion_model_path(model)
+        mace_model = torch.load(mace_model_path, map_location=map_location, weights_only=False)
         init_params = get_init_kwargs_from_model(mace_model)
         model_dtype = get_model_dtype(mace_model)
         if default_dtype == "":
