support specific files by glob

mace/cli/preprocess_data.py

@@ -22,7 +22,7 @@
 from mace.modules import compute_statistics
 from mace.tools import torch_geometric
 from mace.tools.scripts_utils import get_atomic_energies, get_dataset_from_xyz
-from mace.tools.utils import AtomicNumberTable
+from mace.tools.utils import AtomicNumberTable, expand_glob
 
 
 def compute_stats_target(
@@ -176,11 +176,11 @@ def run(args: argparse.Namespace):
     # Data preparation
     collections, atomic_energies_dict = get_dataset_from_xyz(
         work_dir=args.work_dir,
-        train_path=args.train_file,
-        valid_path=args.valid_file,
+        train_path=expand_glob(args.train_file),
+        valid_path=expand_glob(args.valid_file),
         valid_fraction=args.valid_fraction,
         config_type_weights=config_type_weights,
-        test_path=args.test_file,
+        test_path=expand_glob(args.test_file),
         seed=args.seed,
         key_specification=args.key_specification,
         head_name="",

mace/tools/utils.py

@@ -9,10 +9,11 @@
 import os
 import sys
 from pathlib import Path
-from typing import Any, Dict, Iterable, Optional, Sequence, Union
+from typing import Any, Dict, Iterable, Optional, Sequence, Union, List
 
 import numpy as np
 import torch
+from glob import glob
 
 from .torch_tools import to_numpy
 
@@ -205,3 +206,19 @@ def filter_nonzero_weight(
 
     quantity_l[-1] = filtered_q
     return 1.0
+
+
+def expand_glob(file_path: Optional[Union[str, List[str]]]):
+    if file_path is None:
+        return file_path
+    if isinstance(file_path, str):
+        file_path = [file_path]
+    if not isinstance(file_path, list):
+        return file_path
+    expanded_paths = []
+    for path in file_path:
+        _files = glob(path)
+        if not _files:
+            raise FileNotFoundError(f"No files matched the pattern: {path}")
+        expanded_paths.extend(_files)
+    return expanded_paths