Refactor data module

mace/calculators/mace.py

@@ -182,10 +182,18 @@ def __init__(
             [int(z) for z in self.models[0].atomic_numbers]
         )
         self.charges_key = charges_key
+
         try:
-            self.heads = self.models[0].heads
+            self.available_heads = self.models[0].heads
         except AttributeError:
-            self.heads = ["Default"]
+            self.available_heads = ["Default"]
+        self.head = kwargs.get("head", "Default")
+        assert (
+            self.head in self.available_heads
+        ), f"specified head {self.head}, but model available model heads are {self.available_heads}"
+
+        print("using head", self.head, "out of", self.available_heads)
+
         model_dtype = get_model_dtype(self.models[0])
         if default_dtype == "":
             print(
@@ -235,11 +243,19 @@ def _create_result_tensors(
         return dict_of_tensors
 
     def _atoms_to_batch(self, atoms):
-        config = data.config_from_atoms(atoms, charges_key=self.charges_key)
+        keyspec = data.KeySpecification(
+            info_keys={}, arrays_keys={"charges": self.charges_key}
+        )
+        config = data.config_from_atoms(
+            atoms, key_specification=keyspec, head_name=self.head
+        )
         data_loader = torch_geometric.dataloader.DataLoader(
             dataset=[
                 data.AtomicData.from_config(
-                    config, z_table=self.z_table, cutoff=self.r_max, heads=self.heads
+                    config,
+                    z_table=self.z_table,
+                    cutoff=self.r_max,
+                    heads=self.available_heads,
                 )
             ],
             batch_size=1,

mace/cli/preprocess_data.py

@@ -17,6 +17,7 @@
 import tqdm
 
 from mace import data, tools
+from mace.data import KeySpecification, update_keyspec_from_kwargs
 from mace.data.utils import save_configurations_as_HDF5
 from mace.modules import compute_statistics
 from mace.tools import torch_geometric
@@ -144,6 +145,10 @@ def run(args: argparse.Namespace):
     new hdf5 file that is ready for training with on-the-fly dataloading
     """
 
+    # currently support only command line property_key syntax
+    args.key_specification = KeySpecification()
+    update_keyspec_from_kwargs(args.key_specification, vars(args))
+
     # Setup
     tools.set_seeds(args.seed)
     random.seed(args.seed)
@@ -177,12 +182,8 @@ def run(args: argparse.Namespace):
         config_type_weights=config_type_weights,
         test_path=args.test_file,
         seed=args.seed,
-        energy_key=args.energy_key,
-        forces_key=args.forces_key,
-        stress_key=args.stress_key,
-        virials_key=args.virials_key,
-        dipole_key=args.dipole_key,
-        charges_key=args.charges_key,
+        key_specification=args.key_specification,
+        head_name=None,
     )
 
     # Atomic number table

mace/cli/run_train.py

@@ -24,6 +24,7 @@
 import mace
 from mace import data, tools
 from mace.calculators.foundations_models import mace_mp, mace_off
+from mace.data import KeySpecification, update_keyspec_from_kwargs
 from mace.tools import torch_geometric
 from mace.tools.model_script_utils import configure_model
 from mace.tools.multihead_tools import (
@@ -70,6 +71,10 @@ def run(args: argparse.Namespace) -> None:
     tag = tools.get_tag(name=args.name, seed=args.seed)
     args, input_log_messages = tools.check_args(args)
 
+    # default keyspec to update using heads dictionary
+    args.key_specification = KeySpecification()
+    update_keyspec_from_kwargs(args.key_specification, vars(args))
+
     if args.device == "xpu":
         try:
             import intel_extension_for_pytorch as ipex
@@ -153,12 +158,26 @@ def run(args: argparse.Namespace) -> None:
 
     if args.heads is not None:
         args.heads = ast.literal_eval(args.heads)
+        for _, head_dict in args.heads.items():
+            # priority is global args < head property_key values < head info_keys+arrays_keys
+            head_keyspec = deepcopy(args.key_specification)
+            update_keyspec_from_kwargs(head_keyspec, head_dict)
+            head_keyspec.update(
+                info_keys=head_dict.get("info_keys", {}),
+                arrays_keys=head_dict.get("arrays_keys", {}),
+            )
+            head_dict["key_specification"] = head_keyspec
     else:
         args.heads = prepare_default_head(args)
 
     logging.info("===========LOADING INPUT DATA===========")
     heads = list(args.heads.keys())
     logging.info(f"Using heads: {heads}")
+    logging.info("Using the key specifications to parse data:")
+    for name, head_dict in args.heads.items():
+        head_keyspec = head_dict["key_specification"]
+        logging.info(f"{name}: {head_keyspec}")
+
     head_configs: List[HeadConfig] = []
     for head, head_args in args.heads.items():
         logging.info(f"=============    Processing head {head}     ===========")
@@ -187,7 +206,6 @@ def run(args: argparse.Namespace) -> None:
                 head_config.atomic_energies_dict = ast.literal_eval(
                     statistics["atomic_energies"]
                 )
-
         # Data preparation
         if check_path_ase_read(head_config.train_file):
             if head_config.valid_file is not None:
@@ -205,12 +223,7 @@ def run(args: argparse.Namespace) -> None:
                 config_type_weights=config_type_weights,
                 test_path=head_config.test_file,
                 seed=args.seed,
-                energy_key=head_config.energy_key,
-                forces_key=head_config.forces_key,
-                stress_key=head_config.stress_key,
-                virials_key=head_config.virials_key,
-                dipole_key=head_config.dipole_key,
-                charges_key=head_config.charges_key,
+                key_specification=head_config.key_specification,
                 head_name=head_config.head_name,
                 keep_isolated_atoms=head_config.keep_isolated_atoms,
             )
@@ -251,14 +264,21 @@ def run(args: argparse.Namespace) -> None:
                 "Using foundation model for multiheads finetuning with Materials Project data"
             )
             heads = list(dict.fromkeys(["pt_head"] + heads))
+            mp_keyspec = KeySpecification()
+            update_keyspec_from_kwargs(mp_keyspec, vars(args))
+            mp_keyspec.update(
+                info_keys={"energy": "energy", "stress": "stress"},
+                arrays_keys={"forces": "forces"},
+            )
             head_config_pt = HeadConfig(
                 head_name="pt_head",
                 E0s="foundation",
                 statistics_file=args.statistics_file,
+                key_specification=mp_keyspec,
                 compute_avg_num_neighbors=False,
                 avg_num_neighbors=model_foundation.interactions[0].avg_num_neighbors,
             )
-            collections = assemble_mp_data(args, tag, head_configs)
+            collections = assemble_mp_data(args, tag, head_configs, head_config_pt)
             head_config_pt.collections = collections
             head_config_pt.train_file = f"mp_finetuning-{tag}.xyz"
             head_configs.append(head_config_pt)
@@ -275,12 +295,7 @@ def run(args: argparse.Namespace) -> None:
                 config_type_weights=None,
                 test_path=None,
                 seed=args.seed,
-                energy_key=args.energy_key,
-                forces_key=args.forces_key,
-                stress_key=args.stress_key,
-                virials_key=args.virials_key,
-                dipole_key=args.dipole_key,
-                charges_key=args.charges_key,
+                key_specification=args.key_specification,
                 head_name="pt_head",
                 keep_isolated_atoms=args.keep_isolated_atoms,
             )
@@ -292,12 +307,7 @@ def run(args: argparse.Namespace) -> None:
                 statistics_file=args.statistics_file,
                 valid_fraction=args.valid_fraction,
                 config_type_weights=None,
-                energy_key=args.energy_key,
-                forces_key=args.forces_key,
-                stress_key=args.stress_key,
-                virials_key=args.virials_key,
-                dipole_key=args.dipole_key,
-                charges_key=args.charges_key,
+                key_specification=args.key_specification,
                 keep_isolated_atoms=args.keep_isolated_atoms,
                 collections=collections,
                 avg_num_neighbors=model_foundation.interactions[0].avg_num_neighbors,

mace/data/__init__.py

@@ -4,6 +4,7 @@
 from .utils import (
     Configuration,
     Configurations,
+    KeySpecification,
     compute_average_E0s,
     config_from_atoms,
     config_from_atoms_list,
@@ -13,6 +14,7 @@
     save_configurations_as_HDF5,
     save_dataset_as_HDF5,
     test_config_types,
+    update_keyspec_from_kwargs,
 )
 
 __all__ = [
@@ -31,4 +33,6 @@
     "dataset_from_sharded_hdf5",
     "save_AtomicData_to_HDF5",
     "save_configurations_as_HDF5",
+    "KeySpecification",
+    "update_keyspec_from_kwargs",
 ]