diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index f0132348f1..56e1839d06 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5360,12 +5360,12 @@ def training_data_args() -> list[ doc_systems = ( "The data systems for training. " "This key can be a list or a str. " - "When provided as a string, it can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. " - "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories." - ) - doc_patterns = ( - "The customized patterns used in `rglob` to collect all training systems. " + "Each value can be a system directory path (containing 'type.raw'), a parent directory path to recursively search for system subdirectories, or an explicitly named labeled '.xyz' or '.extxyz' file. " + "Extended-XYZ files are read with dpdata and transparently cached as DeePMD NumPy systems. Every frame must contain species, positions, total energy, and atomic forces; virial or ASE-style stress is also required when virial loss is enabled. " + "Files containing heterogeneous atom counts or compositions are split into fixed-shape systems in first-occurrence order. Partially periodic PBC cannot be represented by traditional DeePMD NumPy systems and is rejected. If such a file expands into multiple systems, list-valued batch_size, sys_probs, and indexed auto_prob blocks are rejected as ambiguous. " + "Lists may contain multiple extended-XYZ files and may mix them with existing DeePMD system directories." ) + doc_patterns = "The customized patterns used in `rglob` to collect DeePMD system directories. Explicit '.xyz' and '.extxyz' files are not discovered or filtered by these patterns. " doc_batch_size = f'This key can be \n\n\ - list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ - int: all {link_sys} use the same batch size.\n\n\ @@ -5458,12 +5458,12 @@ def validation_data_args() -> list[ doc_systems = ( "The data systems for validation. " "This key can be a list or a str. " - "When provided as a string, it can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. " - "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories." - ) - doc_patterns = ( - "The customized patterns used in `rglob` to collect all validation systems. " + "Each value can be a system directory path (containing 'type.raw'), a parent directory path to recursively search for system subdirectories, or an explicitly named labeled '.xyz' or '.extxyz' file. " + "Extended-XYZ files are read with dpdata and transparently cached as DeePMD NumPy systems. Every frame must contain species, positions, total energy, and atomic forces; virial or ASE-style stress is also required when virial loss is enabled. " + "Files containing heterogeneous atom counts or compositions are split into fixed-shape systems in first-occurrence order. Partially periodic PBC cannot be represented by traditional DeePMD NumPy systems and is rejected. If such a file expands into multiple systems, list-valued batch_size, sys_probs, and indexed auto_prob blocks are rejected as ambiguous. " + "Lists may contain multiple extended-XYZ files and may mix them with existing DeePMD system directories." ) + doc_patterns = "The customized patterns used in `rglob` to collect DeePMD system directories. Explicit '.xyz' and '.extxyz' files are not discovered or filtered by these patterns. " doc_batch_size = f'This key can be \n\n\ - list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ - int: all {link_sys} use the same batch size.\n\n\ @@ -6543,7 +6543,14 @@ def normalize( _check_dpa3_chg_spin_migration(data) validate_no_multitask_lora(data, multi_task=multi_task) - return data + # External training files are materialized only after schema validation. + # Every backend then receives ordinary DeePMD system paths through its + # existing data-loader and neighbor-statistics code. + from deepmd.utils.data_conversion import ( + normalize_extxyz_training_data, + ) + + return normalize_extxyz_training_data(data, multi_task=multi_task) if __name__ == "__main__": diff --git a/deepmd/utils/data_conversion.py b/deepmd/utils/data_conversion.py new file mode 100644 index 0000000000..a95f4da4d5 --- /dev/null +++ b/deepmd/utils/data_conversion.py @@ -0,0 +1,553 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-independent conversion of external training data.""" + +from __future__ import ( + annotations, +) + +import hashlib +import json +import logging +import os +import shlex +import shutil +import tempfile +import time +import uuid +from collections import ( + OrderedDict, +) +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import dpdata + +log = logging.getLogger(__name__) + +_EXTXYZ_SUFFIXES = frozenset({".xyz", ".extxyz"}) +_CACHE_ENV = "DEEPMD_EXTXYZ_CACHE" +_CACHE_MANIFEST = ".deepmd_extxyz_cache.json" +_CONVERTER_SCHEMA_VERSION = 1 +_SET_SIZE = 2000 +_STRESS_SIGN = -1 +_LOCK_TIMEOUT = 600.0 +_STALE_LOCK_AGE = 24 * 60 * 60 + + +def is_extxyz_path(path: str | os.PathLike[str]) -> bool: + """Return whether an explicit path names an extxyz input file.""" + source = Path(path) + return source.suffix.lower() in _EXTXYZ_SUFFIXES and not source.is_dir() + + +def _is_lmdb_path(path: str | os.PathLike[str]) -> bool: + """Match the existing LMDB path predicate without importing dpmodel.""" + source = Path(path) + return str(path).endswith(".lmdb") or (source / "data.mdb").is_file() + + +def _cache_root() -> Path: + configured = os.environ.get(_CACHE_ENV) + if configured: + return Path(configured).expanduser().resolve() + return (Path(tempfile.gettempdir()) / "deepmd-kit" / "extxyz").resolve() + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _fingerprint(source: Path, dpdata_version: str) -> dict[str, Any]: + return { + "converter_schema_version": _CONVERTER_SCHEMA_VERSION, + "conversion_settings": { + "format": "extxyz", + "output_format": "deepmd/npy", + "set_size": _SET_SIZE, + "stress_sign": _STRESS_SIGN, + }, + "dpdata_version": dpdata_version, + "source_path": str(source), + "source_sha256": _file_sha256(source), + } + + +def _fingerprint_digest(fingerprint: dict[str, Any]) -> str: + encoded = json.dumps(fingerprint, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _read_manifest(root: Path) -> dict[str, Any] | None: + try: + with (root / _CACHE_MANIFEST).open(encoding="utf-8") as stream: + manifest = json.load(stream) + except (FileNotFoundError, json.JSONDecodeError, OSError): + return None + if not isinstance(manifest, dict): + return None + return manifest + + +def _manifest_system_paths(root: Path, manifest: dict[str, Any]) -> list[str] | None: + entries = manifest.get("systems") + if not isinstance(entries, list) or not entries: + return None + + resolved_root = root.resolve() + result = [] + for entry in entries: + if not isinstance(entry, dict) or not isinstance(entry.get("path"), str): + return None + system = (root / entry["path"]).resolve() + if not system.is_relative_to(resolved_root): + return None + if not (system / "type.raw").is_file(): + return None + if not any(path.is_dir() for path in system.glob("set.*")): + return None + result.append(str(system)) + return result + + +def expand_extxyz_cache(path: str | os.PathLike[str]) -> list[str] | None: + """Expand a materialized extxyz cache in deterministic manifest order.""" + root = Path(path) + manifest = _read_manifest(root) + if manifest is None: + return None + return _manifest_system_paths(root, manifest) + + +def _valid_cache( + target: Path, fingerprint: dict[str, Any] +) -> tuple[dict[str, Any], list[str]] | None: + manifest = _read_manifest(target) + if manifest is None or manifest.get("fingerprint") != fingerprint: + return None + systems = _manifest_system_paths(target, manifest) + if systems is None: + return None + return manifest, systems + + +def _labels(system: Any) -> list[str]: + """Return DeePMD label names for fields produced by dpdata.""" + labels = [] + if "energies" in system.data: + labels.append("energy") + # dpdata uses plural dictionary keys. The singular names recorded here + # are the canonical DeePMD dataset names (energy.npy, force.npy, etc.). + if "forces" in system.data: + labels.append("force") + if "virials" in system.data: + labels.append("virial") + return labels + + +def _frame_error(source: Path, frame_index: int, exc: Exception) -> ValueError: + detail = str(exc) + lowered = detail.lower() + if "energy" in lowered or "energies" in lowered: + problem = "is missing a usable total-energy label" + elif "force" in lowered or "forces" in lowered: + problem = "is missing usable atomic-force labels" + else: + problem = "is not a valid labeled extended-XYZ frame" + return ValueError( + f"Frame {frame_index + 1} in extxyz file '{source}' {problem}. " + "Each frame must contain atomic species, positions, total energy, and " + f"atomic forces. dpdata reported: {detail}" + ) + + +def _validate_representable_pbc(source: Path) -> None: + """Reject per-axis PBC that traditional DeePMD systems cannot encode.""" + false_values = {"f", "false", "0"} + true_values = {"t", "true", "1"} + frame_index = 0 + with source.open(encoding="utf-8") as stream: + while line := stream.readline(): + try: + natoms = int(line.strip()) + except ValueError: + continue + header = stream.readline() + try: + fields = shlex.split(header) + except ValueError: + fields = [] # Let dpdata report malformed extxyz syntax. + for field in fields: + key, separator, value = field.partition("=") + if not separator or key.casefold() != "pbc": + continue + flags = value.split() + normalized = [flag.casefold() for flag in flags] + recognized = all( + flag in false_values or flag in true_values for flag in normalized + ) + if not recognized or len(flags) not in {1, 3}: + raise ValueError( + f"Frame {frame_index + 1} in extxyz file '{source}' has " + f"an invalid pbc field: '{value}'. Use one or three " + "boolean values." + ) + periodic = [flag in true_values for flag in normalized] + if any(periodic) and not all(periodic): + raise ValueError( + f"Frame {frame_index + 1} in extxyz file '{source}' is " + f"partially periodic (pbc='{value}'). Traditional DeePMD " + "NumPy systems can represent only all-periodic or " + "all-nonperiodic boundary conditions." + ) + break + for _ in range(natoms): + stream.readline() + frame_index += 1 + + +def _read_extxyz_systems(source: Path) -> list[dpdata.LabeledSystem]: + """Read every frame through dpdata and form fixed-shape DeePMD systems.""" + _validate_representable_pbc(source) + + # dpdata exposes formats through its plugin registry. Using the registered + # extxyz reader here keeps aliases, units, and stress conversion in dpdata. + from dpdata.format import ( + Format, + ) + + format_class = Format.get_formats().get("extxyz") + if format_class is None: # pragma: no cover - guaranteed by the dependency + raise RuntimeError( + f"dpdata {dpdata.__version__} does not provide its extxyz reader" + ) + + groups: OrderedDict[tuple[Any, ...], dpdata.LabeledSystem] = OrderedDict() + frames = iter( + format_class().from_multi_systems(str(source), stress_sign=_STRESS_SIGN) + ) + frame_index = 0 + while True: + try: + frame_data = next(frames) + except StopIteration: + break + except Exception as exc: + raise _frame_error(source, frame_index, exc) from exc + + try: + frame = dpdata.LabeledSystem(data=frame_data) + except Exception as exc: + raise _frame_error(source, frame_index, exc) from exc + + labels = tuple(_labels(frame)) + if "energy" not in labels: + raise _frame_error( + source, frame_index, ValueError("energies not found in data") + ) + if "force" not in labels: + raise _frame_error( + source, frame_index, ValueError("forces not found in data") + ) + + # Canonicalizing type-map names lets dpdata append compatible + # compositions while retaining the first frame's atom order. + frame.sort_atom_names() + key = ( + frame.uniq_formula, + bool(frame.data.get("nopbc", False)), + labels, + ) + if key not in groups: + groups[key] = frame + else: + try: + groups[key].append(frame) + except Exception as exc: + raise ValueError( + f"Frame {frame_index + 1} in extxyz file '{source}' cannot " + "be combined safely with frames of the same composition. " + f"dpdata reported: {exc}" + ) from exc + frame_index += 1 + + if frame_index == 0: + raise ValueError(f"Extxyz file '{source}' contains no frames.") + return list(groups.values()) + + +def _write_cache( + temporary: Path, + source: Path, + fingerprint: dict[str, Any], +) -> dict[str, Any]: + systems = _read_extxyz_systems(source) + entries = [] + for index, system in enumerate(systems): + relative = f"system.{index:03d}" + system.to("deepmd/npy", str(temporary / relative), set_size=_SET_SIZE) + entries.append( + { + "atom_names": list(system.data["atom_names"]), + "atom_numbs": [int(value) for value in system.data["atom_numbs"]], + "frames": int(system.get_nframes()), + "labels": _labels(system), + "nopbc": bool(system.data.get("nopbc", False)), + "path": relative, + } + ) + + manifest = { + "fingerprint": fingerprint, + "source": str(source), + "systems": entries, + } + with (temporary / _CACHE_MANIFEST).open("w", encoding="utf-8") as stream: + json.dump(manifest, stream, indent=2, sort_keys=True) + stream.write("\n") + return manifest + + +def _remove_invalid_cache(path: Path, cache_root: Path) -> None: + if not path.exists(): + return + if path.resolve().parent != cache_root.resolve(): + raise RuntimeError(f"Refusing to remove cache outside '{cache_root}': {path}") + shutil.rmtree(path) + + +def materialize_extxyz( + path: str | os.PathLike[str], +) -> tuple[str, dict[str, Any]]: + """Convert an extxyz file to an atomically published DeePMD NumPy cache.""" + source = Path(path).expanduser().resolve() + if not source.is_file(): + raise FileNotFoundError(f"Extxyz training-data file does not exist: '{source}'") + + fingerprint = _fingerprint(source, dpdata.__version__) + digest = _fingerprint_digest(fingerprint) + cache_root = _cache_root() + cache_root.mkdir(parents=True, exist_ok=True) + target = cache_root / digest + + cached = _valid_cache(target, fingerprint) + if cached is not None: + return str(target), cached[0] + + lock = cache_root / f".{digest}.lock" + deadline = time.monotonic() + _LOCK_TIMEOUT + while True: + try: + lock.mkdir() + break + except FileExistsError: + cached = _valid_cache(target, fingerprint) + if cached is not None: + return str(target), cached[0] + try: + lock_age = time.time() - lock.stat().st_mtime + if lock_age > _STALE_LOCK_AGE: + lock.rmdir() + continue + except FileNotFoundError: + continue + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for extxyz cache creation for '{source}'. " + f"If no conversion is running, remove stale lock '{lock}'." + ) from None + time.sleep(0.1) + + temporary = cache_root / f".{digest}.tmp-{os.getpid()}-{uuid.uuid4().hex}" + try: + cached = _valid_cache(target, fingerprint) + if cached is not None: + return str(target), cached[0] + + _remove_invalid_cache(target, cache_root) + temporary.mkdir() + log.info("Converting extxyz training data %s with dpdata", source) + manifest = _write_cache(temporary, source, fingerprint) + if _file_sha256(source) != fingerprint["source_sha256"]: + raise RuntimeError( + f"Extxyz training-data file '{source}' changed while it was " + "being converted. Retry after the file is no longer being written." + ) + temporary.replace(target) + return str(target), manifest + finally: + if temporary.exists(): + shutil.rmtree(temporary) + try: + lock.rmdir() + except FileNotFoundError: + pass + + +def _active(loss: dict[str, Any], *keys: str) -> bool: + return any(float(loss.get(key, 0.0)) != 0.0 for key in keys) + + +def _validate_loss_labels( + source: str, manifest: dict[str, Any], loss: dict[str, Any] +) -> None: + loss_type = loss.get("type", "ener") + if loss_type not in {"ener", "dens"}: + raise ValueError( + f"Extxyz training data '{source}' currently supports energy-model " + f"losses only; configured loss type is '{loss_type}'." + ) + + unsupported = { + "atomic energy": ("start_pref_ae", "limit_pref_ae"), + "atomic preference": ("start_pref_pf", "limit_pref_pf"), + "Hessian": ("start_pref_h", "limit_pref_h"), + "generalized force": ("start_pref_gf", "limit_pref_gf"), + } + for label, keys in unsupported.items(): + if _active(loss, *keys): + raise ValueError( + f"The configured loss requires {label} labels, but dpdata's " + f"extxyz reader does not convert them for '{source}'." + ) + + required = {"energy", "force"} + if _active(loss, "start_pref_v", "limit_pref_v"): + required.add("virial") + for index, system in enumerate(manifest["systems"]): + missing = sorted(required.difference(system["labels"])) + if missing: + missing_text = ", ".join(missing) + raise ValueError( + f"Extxyz file '{source}' is missing required {missing_text} " + f"label(s) in converted system {index}; the configured loss " + "uses those labels. Supply virial or ASE-style stress when " + "virial loss is enabled." + ) + + +def _normalize_dataset( + dataset: dict[str, Any], loss: dict[str, Any], location: str +) -> None: + systems_value = dataset.get("systems") + if isinstance(systems_value, str): + systems = [systems_value] + scalar = True + elif isinstance(systems_value, list): + systems = systems_value + scalar = False + else: + return + + extxyz_indices = [ + index for index, path in enumerate(systems) if is_extxyz_path(path) + ] + if not extxyz_indices: + return + + if any(_is_lmdb_path(path) for path in systems): + raise ValueError( + f"{location}/systems cannot mix extxyz and LMDB inputs. LMDB is " + "supported only as a single systems path; convert the inputs to one " + "representation first." + ) + + normalized = list(systems) + expands_to_multiple = False + for index in extxyz_indices: + source = systems[index] + cache, manifest = materialize_extxyz(source) + _validate_loss_labels(source, manifest, loss) + normalized[index] = cache + expands_to_multiple |= len(manifest["systems"]) > 1 + + if expands_to_multiple: + if isinstance(dataset.get("batch_size"), list): + raise ValueError( + f"{location}/batch_size cannot be a list when one extxyz file " + "expands into multiple fixed-shape systems. Use a scalar or " + "'auto' batch size." + ) + if dataset.get("sys_probs") is not None: + raise ValueError( + f"{location}/sys_probs is ambiguous when one extxyz file expands " + "into multiple fixed-shape systems. Omit it and use auto_prob." + ) + if ";" in dataset.get("auto_prob", ""): + raise ValueError( + f"{location}/auto_prob cannot use indexed blocks when one extxyz " + "file expands into multiple fixed-shape systems." + ) + + dataset["systems"] = normalized[0] if scalar else normalized + + +def normalize_extxyz_training_data( + data: dict[str, Any], *, multi_task: bool = False +) -> dict[str, Any]: + """Materialize explicit extxyz paths on a normalized configuration copy.""" + training = data.get("training") + if not isinstance(training, dict): + return data + + if multi_task: + data_dict = training.get("data_dict", {}) + bindings = [ + ( + task_data, + data.get("loss_dict", {}).get(task, {"type": "ener"}), + f"training/data_dict/{task}", + ) + for task, task_data in data_dict.items() + ] + else: + bindings = [(training, data.get("loss", {"type": "ener"}), "training")] + + if not any( + is_extxyz_path(path) + for task_data, _, _ in bindings + for name in ("training_data", "validation_data") + if isinstance(task_data.get(name), dict) + for path in ( + [task_data[name]["systems"]] + if isinstance(task_data[name].get("systems"), str) + else task_data[name].get("systems", []) + ) + ): + return data + + result = deepcopy(data) + if multi_task: + result_bindings = [ + ( + task_data, + result.get("loss_dict", {}).get(task, {"type": "ener"}), + f"training/data_dict/{task}", + ) + for task, task_data in result["training"].get("data_dict", {}).items() + ] + else: + result_bindings = [ + ( + result["training"], + result.get("loss", {"type": "ener"}), + "training", + ) + ] + + for task_data, loss, location in result_bindings: + for name in ("training_data", "validation_data"): + dataset = task_data.get(name) + if isinstance(dataset, dict): + _normalize_dataset(dataset, loss, f"{location}/{name}") + return result diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index cea270bd37..b1e340e82c 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -24,6 +24,11 @@ DataRequirementItem, DeepmdData, ) +from deepmd.utils.data_conversion import ( + expand_extxyz_cache, + is_extxyz_path, + materialize_extxyz, +) from deepmd.utils.out_stat import ( compute_stats_from_redu, ) @@ -851,6 +856,7 @@ def process_systems( If it is a single directory, search for all the systems in the directory. If it is a list, each item in the list is treated as a directory to search. + If it is an explicit extxyz file, materialize and expand its ordered cache. If it is a single LMDB path, return it directly without expansion. Check if the systems are valid. @@ -888,7 +894,20 @@ def process_systems( # Iterate over the search_paths list and apply expansion logic to each path result_systems = [] for path in search_paths: - if patterns is None: + cached_paths = expand_extxyz_cache(path) + if cached_paths is not None: + # A cache manifest records the deterministic order of systems split + # from one heterogeneous extxyz input. rglob_patterns deliberately + # does not filter explicitly named extxyz files. + expanded_paths = cached_paths + elif is_extxyz_path(path): + cache_root, _ = materialize_extxyz(path) + expanded_paths = expand_extxyz_cache(cache_root) + if expanded_paths is None: # pragma: no cover - cache is validated + raise RuntimeError( + f"Extxyz cache for '{path}' was created without a valid manifest" + ) + elif patterns is None: expanded_paths = expand_sys_str(path) else: expanded_paths = rglob_sys_str(path, patterns) diff --git a/doc/data/data-conv.md b/doc/data/data-conv.md index 30be98bcfe..2e45eea01e 100644 --- a/doc/data/data-conv.md +++ b/doc/data/data-conv.md @@ -1,6 +1,6 @@ # Formats of a system -Two binary formats, NumPy and HDF5, are supported for training. The raw format is not directly supported, but a tool is provided to convert data from the raw format to the NumPy format. +The native on-disk training formats are NumPy and HDF5. Labeled extended-XYZ files can also be listed directly in `training_data/systems` or `validation_data/systems`; DeePMD-kit transparently converts and caches them as NumPy systems before initializing the existing data loader. The raw format is not directly supported, but a tool is provided to convert data from the raw format to the NumPy format. ## NumPy format diff --git a/doc/data/dpdata.md b/doc/data/dpdata.md index 63fe4f39c3..6af843c56c 100644 --- a/doc/data/dpdata.md +++ b/doc/data/dpdata.md @@ -1,12 +1,8 @@ # Prepare data with dpdata -One can use a convenient tool [`dpdata`](https://github.com/deepmodeling/dpdata) to convert data directly from the output of first principle packages to the DeePMD-kit format. +DeePMD-kit includes [`dpdata`](https://github.com/deepmodeling/dpdata) and uses it to convert data from first-principles packages to the DeePMD-kit format. -To install one can execute - -```bash -pip install dpdata -``` +Labeled `.xyz` and `.extxyz` files are accepted directly by the `systems` key in training and validation configurations. That path is converted automatically; the explicit Python conversion shown below remains useful for other source formats and standalone data preparation. An example of converting data [VASP](https://www.vasp.at/) data in `OUTCAR` format to DeePMD-kit data can be found at diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 31609c09b8..3b61fe3924 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -53,8 +53,39 @@ Other training parameters are given in the {ref}`training ` section. The sections {ref}`training_data ` and {ref}`validation_data ` give the training dataset and validation dataset, respectively. Taking the training dataset for example, the keys are explained below: - {ref}`systems ` provide paths of the training data systems. DeePMD-kit allows you to provide multiple systems with different numbers of atoms. This key can be a `list` or a `str`. - - `str`: {ref}`systems ` should be a valid path. It can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. - - `list`: {ref}`systems ` gives a list of paths. Each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories. + - `str`: {ref}`systems ` should be a valid path. It can be a system directory path (containing `type.raw`), a parent directory path to recursively search for all system subdirectories, or an explicitly named labeled `.xyz` or `.extxyz` file. + - `list`: {ref}`systems ` gives a list of paths. Each item can use any of the forms accepted for `str`, so multiple extended-XYZ files and mixtures of extended-XYZ files and existing DeePMD systems are supported. + +### Labeled extended-XYZ datasets + +An extended-XYZ file can be used directly for either training or validation; no separate conversion command or modification of `input.json` is needed: + +```json +{ + "training": { + "training_data": { + "systems": [ + "data/train_part_1.extxyz", + "data/train_part_2.xyz" + ], + "batch_size": "auto" + }, + "validation_data": { + "systems": "data/validation.extxyz", + "batch_size": "auto" + } + } +} +``` + +DeePMD-kit uses `dpdata` to read every frame and to write a transparent DeePMD NumPy cache. Species and positions are read from the standard `Properties` declaration. Atomic forces may use `force` or the common `forces` property name. Total energy accepts the aliases supported by `dpdata`. A `virial`/`virials` field is retained directly; an ASE-style `stress`/`stresses` field is converted using `virial = -volume * stress`. Stress may contain nine row-major tensor components or six components in ASE Voigt order (`xx yy zz yz xz xy`). In the absence of explicit unit metadata, positions and cells are interpreted as angstrom, energy and virial as eV, forces as eV/angstrom, and stress as eV/angstrom^3. Unit metadata supported by `dpdata` is converted to those DeePMD units. + +Every frame must contain total energy and atomic forces. A coordinate-only XYZ file therefore fails with a label-specific error. If the configured loss has a nonzero virial prefactor, every frame must additionally contain either virial or usable stress data. + +Frames with different atom counts, compositions, periodicity, or available label sets are split deterministically into ordinary fixed-shape DeePMD NumPy systems in first-occurrence order. Traditional DeePMD NumPy systems distinguish only all-periodic from all-nonperiodic data, so a partially periodic field such as `pbc="T T F"` is rejected explicitly. A heterogeneous file cannot be combined with list-valued `batch_size`, explicit `sys_probs`, or indexed `auto_prob` blocks because one input path then maps to multiple internal systems; use scalar/`"auto"` batching and automatic probabilities instead. + +Converted data is cached under the platform temporary directory. The cache key includes the canonical source path, complete source-file hash, conversion settings, and `dpdata` version, so editing the source creates a new cache. Publication is atomic and concurrent launches coordinate through a per-key lock. Set `DEEPMD_EXTXYZ_CACHE` to use a different cache root. `rglob_patterns` continues to control only directory discovery: arbitrary `.xyz` files below a parent directory are never discovered implicitly, and an explicitly listed extended-XYZ file is never filtered by those patterns. + - At each training step, DeePMD-kit randomly picks {ref}`batch_size ` frame(s) from one of the systems. The probability of using a system is by default in proportion to the number of batches in the system. More options are available for automatically determining the probability of using systems. One can set the key {ref}`auto_prob ` to - `"prob_uniform"` all systems are used with the same probability. - `"prob_sys_size"` the probability of using a system is proportional to its size (number of frames). diff --git a/pyproject.toml b/pyproject.toml index 3b64d3ec1b..5512906cc2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ dependencies = [ 'array-api-compat', 'lmdb', 'msgpack', + 'dpdata>=1.1.0', ] requires-python = ">=3.10" keywords = ["deepmd"] @@ -79,7 +80,6 @@ repository = "https://github.com/deepmodeling/deepmd-kit" # which can be read by the build backend. [tool.deepmd_build_backend.optional-dependencies] test = [ - "dpdata>=0.2.7", # ASE issue: https://gitlab.com/ase/ase/-/merge_requests/2843 # fixed in 3.23.0 "ase>=3.23.0", @@ -97,7 +97,6 @@ test = [ ] dpa-adapt = [ "scikit-learn", - "dpdata", "torch", "ase", "rdkit", diff --git a/source/tests/common/test_extxyz_training_data.py b/source/tests/common/test_extxyz_training_data.py new file mode 100644 index 0000000000..552f6bc218 --- /dev/null +++ b/source/tests/common/test_extxyz_training_data.py @@ -0,0 +1,571 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for labeled extended-XYZ training and validation inputs.""" + +from __future__ import ( + annotations, +) + +import copy +import json +import shutil +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import dpdata +import numpy as np +import pytest +from scipy.constants import ( + electron_volt, +) + +from deepmd.utils.data_conversion import ( + expand_extxyz_cache, + materialize_extxyz, + normalize_extxyz_training_data, +) + +_DEFAULT_FORCES = object() +_ENERGY_FORCE_LOSS = { + "type": "ener", + "start_pref_e": 1.0, + "limit_pref_e": 1.0, + "start_pref_f": 1.0, + "limit_pref_f": 1.0, + "start_pref_v": 0.0, + "limit_pref_v": 0.0, +} + + +@pytest.fixture(autouse=True) +def _isolated_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DEEPMD_EXTXYZ_CACHE", str(tmp_path / "extxyz-cache")) + + +def _frame( + *, + species: tuple[str, ...] = ("H", "O"), + positions: np.ndarray | None = None, + forces: Any = _DEFAULT_FORCES, + energy: float | None = -1.25, + cell: np.ndarray | None = None, + pbc: bool = True, + tensor_name: str | None = None, + tensor: np.ndarray | None = None, + energy_unit: str | None = None, + force_unit: str | None = None, + stress_unit: str | None = None, +) -> str: + natoms = len(species) + if positions is None: + positions = np.arange(natoms * 3, dtype=float).reshape(natoms, 3) / 10.0 + if forces is _DEFAULT_FORCES: + forces = -np.arange(1, natoms * 3 + 1, dtype=float).reshape(natoms, 3) / 10.0 + if cell is None: + cell = np.diag([2.0, 3.0, 4.0]) + + properties = "Properties=species:S:1:pos:R:3" + if forces is not None: + properties += ":forces:R:3" + fields = [ + f'Lattice="{" ".join(str(value) for value in cell.reshape(-1))}"', + properties, + f'pbc="{"T T T" if pbc else "F F F"}"', + ] + if energy is not None: + fields.append(f"energy={energy}") + if tensor_name is not None and tensor is not None: + values = " ".join(str(value) for value in np.asarray(tensor).reshape(-1)) + fields.append(f'{tensor_name}="{values}"') + if energy_unit is not None: + fields.append(f"energy_unit={energy_unit}") + if force_unit is not None: + fields.append(f"force_unit={force_unit}") + if stress_unit is not None: + fields.append(f"stress_unit={stress_unit}") + + atom_lines = [] + for index, name in enumerate(species): + values = [*positions[index]] + if forces is not None: + values.extend(forces[index]) + atom_lines.append(f"{name} " + " ".join(str(value) for value in values)) + return f"{natoms}\n{' '.join(fields)}\n" + "\n".join(atom_lines) + "\n" + + +def _write(path: Path, *frames: str) -> Path: + path.write_text("".join(frames), encoding="utf-8") + return path + + +def _converted(source: Path) -> tuple[Path, dict[str, Any]]: + cache, manifest = materialize_extxyz(source) + systems = expand_extxyz_cache(cache) + assert systems is not None + assert len(systems) == 1 + return Path(systems[0]), manifest + + +def _load_system(path: Path) -> dpdata.LabeledSystem: + return dpdata.LabeledSystem(str(path), fmt="deepmd/npy") + + +def _partial_config( + training_systems: str | list[str], + validation_systems: str | list[str] | None = None, + *, + loss: dict[str, Any] | None = None, + batch_size: str | int | list[int] = "auto", +) -> dict[str, Any]: + training: dict[str, Any] = { + "training_data": { + "systems": training_systems, + "batch_size": batch_size, + } + } + if validation_systems is not None: + training["validation_data"] = { + "systems": validation_systems, + "batch_size": batch_size, + } + return { + "loss": copy.deepcopy(loss or _ENERGY_FORCE_LOSS), + "training": training, + } + + +def _energies(paths: list[str]) -> list[float]: + return [float(_load_system(Path(path)).data["energies"][0]) for path in paths] + + +def test_periodic_energy_force_cell_and_pbc_round_trip(tmp_path: Path) -> None: + positions = np.array([[0.1, 0.2, 0.3], [1.1, 1.2, 1.3]]) + forces = np.array([[0.4, 0.5, 0.6], [-0.4, -0.5, -0.6]]) + cell = np.array([[2.0, 0.1, 0.2], [0.0, 3.0, 0.3], [0.0, 0.0, 4.0]]) + source = _write( + tmp_path / "periodic.extxyz", + _frame(positions=positions, forces=forces, energy=-3.5, cell=cell), + _frame(positions=positions + 0.25, forces=forces * 2, energy=-2.5, cell=cell), + ) + + system_path, manifest = _converted(source) + system = _load_system(system_path) + + np.testing.assert_allclose(system.data["energies"], [-3.5, -2.5]) + np.testing.assert_allclose(system.data["coords"], [positions, positions + 0.25]) + np.testing.assert_allclose(system.data["forces"], [forces, forces * 2]) + np.testing.assert_allclose(system.data["cells"], [cell, cell]) + assert not system.nopbc + assert manifest["systems"][0]["labels"] == ["energy", "force"] + + +def test_nonperiodic_pbc_survives_conversion(tmp_path: Path) -> None: + source = _write(tmp_path / "nonperiodic.extxyz", _frame(pbc=False)) + + system_path, manifest = _converted(source) + system = _load_system(system_path) + + assert system.nopbc + assert manifest["systems"][0]["nopbc"] + + +def test_partially_periodic_pbc_is_rejected(tmp_path: Path) -> None: + source = _write( + tmp_path / "partial-pbc.extxyz", + _frame().replace('pbc="T T T"', 'pbc="T T F"'), + ) + + with pytest.raises(ValueError, match="partially periodic"): + materialize_extxyz(source) + + +def test_explicit_virial_round_trip(tmp_path: Path) -> None: + virial = np.arange(1.0, 10.0).reshape(3, 3) + source = _write( + tmp_path / "virial.xyz", + _frame(tensor_name="virial", tensor=virial), + ) + + system_path, manifest = _converted(source) + system = _load_system(system_path) + + np.testing.assert_allclose(system.data["virials"], [virial]) + assert manifest["systems"][0]["labels"] == ["energy", "force", "virial"] + + +@pytest.mark.parametrize( + ("stress", "expected_matrix"), + [ + ( + np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), + np.array([[1.0, 6.0, 5.0], [6.0, 2.0, 4.0], [5.0, 4.0, 3.0]]), + ), + ( + np.arange(1.0, 10.0), + np.arange(1.0, 10.0).reshape(3, 3), + ), + ], +) +def test_ase_stress_to_virial_sign_units_and_ordering( + tmp_path: Path, stress: np.ndarray, expected_matrix: np.ndarray +) -> None: + cell = np.diag([2.0, 3.0, 4.0]) + source = _write( + tmp_path / f"stress-{stress.size}.extxyz", + _frame( + cell=cell, + tensor_name="stress", + tensor=stress, + stress_unit="GPa", + ), + ) + + system_path, _ = _converted(source) + virial = _load_system(system_path).data["virials"][0] + gpa_to_ev_per_angstrom3 = 1e9 * 1e-30 / electron_volt + expected = -abs(np.linalg.det(cell)) * expected_matrix * gpa_to_ev_per_angstrom3 + np.testing.assert_allclose(virial, expected) + + +def test_energy_and_force_units_are_converted(tmp_path: Path) -> None: + source = _write( + tmp_path / "units.extxyz", + _frame( + energy=1.0, + forces=np.ones((2, 3)), + energy_unit="hartree", + force_unit="hartree/bohr", + ), + ) + + system_path, _ = _converted(source) + system = _load_system(system_path) + + # Independent CODATA values in DeePMD's eV and eV/angstrom units. + np.testing.assert_allclose(system.data["energies"], [27.211386245988], rtol=1e-10) + np.testing.assert_allclose(system.data["forces"], 51.4220674763, rtol=1e-10) + + +def test_multiple_relative_and_absolute_inputs_preserve_order( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + first = _write(tmp_path / "first.extxyz", _frame(energy=-1.0)) + second = _write(tmp_path / "second.xyz", _frame(energy=-2.0)) + validation = _write(tmp_path / "validation.extxyz", _frame(energy=-3.0)) + monkeypatch.chdir(tmp_path) + original = _partial_config( + [first.name, str(second.resolve())], + validation.name, + ) + snapshot = copy.deepcopy(original) + + normalized = normalize_extxyz_training_data(original) + + assert original == snapshot + from deepmd.utils.data_system import ( + process_systems, + ) + + training_paths = process_systems(normalized["training"]["training_data"]["systems"]) + validation_paths = process_systems( + normalized["training"]["validation_data"]["systems"] + ) + assert _energies(training_paths) == [-1.0, -2.0] + assert _energies(validation_paths) == [-3.0] + + +def test_existing_numpy_raw_hdf5_parent_lmdb_and_mixed_inputs_are_unchanged( + tmp_path: Path, +) -> None: + from deepmd.utils.data_system import ( + process_systems, + ) + from deepmd.utils.path import ( + DPH5Path, + ) + + source = _write(tmp_path / "new.extxyz", _frame(energy=-2.0)) + generated, _ = _converted(source) + existing = tmp_path / "existing" + shutil.copytree(generated, existing) + parent = tmp_path / "parent" + parent.mkdir() + nested = parent / "nested" + shutil.copytree(generated, nested) + _write(parent / "not-discovered.xyz", _frame(energy=-99.0)) + system = _load_system(generated) + raw = tmp_path / "raw-system" + system.to("deepmd/raw", str(raw)) + hdf5 = tmp_path / "system.hdf5" + system.to("deepmd/hdf5", str(hdf5)) + + assert process_systems(str(existing)) == [str(existing)] + assert process_systems(str(raw)) == [str(raw)] + assert process_systems(str(hdf5)) == [f"{hdf5}#/"] + assert process_systems(str(parent)) == [str(nested)] + assert process_systems(str(tmp_path / "dataset.lmdb")) == [ + str(tmp_path / "dataset.lmdb") + ] + mixed = process_systems([str(existing), str(source)]) + assert mixed[0] == str(existing) + assert _energies(mixed) == [-2.0, -2.0] + + # DPH5Path caches read handles globally; release the synthetic fixture on + # Windows so pytest can remove its temporary directory. + DPH5Path._load_h5py(str(hdf5), "r").close() + DPH5Path._load_h5py.cache_clear() + DPH5Path._file_keys.cache_clear() + + +def test_rglob_patterns_do_not_discover_or_filter_explicit_extxyz( + tmp_path: Path, +) -> None: + from deepmd.utils.data_system import ( + process_systems, + ) + + source = _write(tmp_path / "explicit.extxyz", _frame(energy=-4.0)) + assert process_systems(str(tmp_path), patterns=["*.xyz", "*.extxyz"]) == [] + explicit = process_systems(str(source), patterns=["does-not-match"]) + assert _energies(explicit) == [-4.0] + + +@pytest.mark.parametrize( + ("frame", "message"), + [ + (_frame(energy=None), "energy"), + (_frame(forces=None), "force"), + ], +) +def test_missing_required_energy_or_forces_fails_clearly( + tmp_path: Path, frame: str, message: str +) -> None: + source = _write(tmp_path / f"missing-{message}.extxyz", frame) + with pytest.raises(ValueError, match=message): + materialize_extxyz(source) + + +def test_coordinate_only_xyz_fails_clearly(tmp_path: Path) -> None: + source = _write( + tmp_path / "coordinates.xyz", + "2\ncoordinate-only XYZ\nH 0 0 0\nO 1 1 1\n", + ) + with pytest.raises(ValueError, match=r"species.*positions.*energy.*forces"): + materialize_extxyz(source) + + +def test_missing_virial_fails_when_virial_loss_is_enabled(tmp_path: Path) -> None: + source = _write(tmp_path / "no-virial.extxyz", _frame()) + loss = copy.deepcopy(_ENERGY_FORCE_LOSS) + loss["start_pref_v"] = 1.0 + loss["limit_pref_v"] = 1.0 + + with pytest.raises(ValueError, match="missing required virial"): + normalize_extxyz_training_data(_partial_config(str(source), loss=loss)) + + +def test_heterogeneous_frames_split_deterministically(tmp_path: Path) -> None: + source = _write( + tmp_path / "heterogeneous.extxyz", + _frame(species=("H", "H"), energy=-1.0), + _frame(species=("H", "O", "H"), energy=-2.0), + _frame(species=("H", "H"), energy=-3.0), + ) + + cache, manifest = materialize_extxyz(source) + paths = expand_extxyz_cache(cache) + assert paths is not None + assert [entry["frames"] for entry in manifest["systems"]] == [2, 1] + assert [entry["atom_numbs"] for entry in manifest["systems"]] == [[2], [2, 1]] + assert _load_system(Path(paths[0])).data["energies"].tolist() == [-1.0, -3.0] + assert _load_system(Path(paths[1])).data["energies"].tolist() == [-2.0] + + +@pytest.mark.parametrize( + ("setting", "value", "message"), + [ + ("batch_size", [1], "batch_size"), + ("sys_probs", [1.0], "sys_probs"), + ("auto_prob", "prob_sys_size;0:1:1.0", "auto_prob"), + ], +) +def test_heterogeneous_expansion_rejects_ambiguous_per_system_settings( + tmp_path: Path, setting: str, value: Any, message: str +) -> None: + source = _write( + tmp_path / f"heterogeneous-{setting}.extxyz", + _frame(species=("H", "H")), + _frame(species=("H", "O", "H")), + ) + config = _partial_config(str(source)) + config["training"]["training_data"][setting] = value + + with pytest.raises(ValueError, match=message): + normalize_extxyz_training_data(config) + + +def test_atom_layout_changes_do_not_corrupt_coordinates_or_forces( + tmp_path: Path, +) -> None: + first_positions = np.array([[1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]) + first_forces = np.array([[10.0, 0.0, 0.0], [20.0, 0.0, 0.0]]) + second_positions = np.array([[3.0, 0.0, 0.0], [4.0, 0.0, 0.0]]) + second_forces = np.array([[30.0, 0.0, 0.0], [40.0, 0.0, 0.0]]) + source = _write( + tmp_path / "layouts.extxyz", + _frame( + species=("H", "O"), + positions=first_positions, + forces=first_forces, + energy=-1.0, + ), + _frame( + species=("O", "H"), + positions=second_positions, + forces=second_forces, + energy=-2.0, + ), + ) + + system_path, manifest = _converted(source) + system = _load_system(system_path) + assert manifest["systems"][0]["frames"] == 2 + for frame_index in range(2): + coord_force_pairs = sorted( + zip( + system.data["coords"][frame_index, :, 0], + system.data["forces"][frame_index, :, 0], + strict=True, + ) + ) + expected = [[(1.0, 10.0), (2.0, 20.0)], [(3.0, 30.0), (4.0, 40.0)]] + assert coord_force_pairs == expected[frame_index] + + +def test_cache_reuse_and_content_invalidation(tmp_path: Path) -> None: + source = _write(tmp_path / "cache.extxyz", _frame(energy=-1.0)) + first_cache, first_manifest = materialize_extxyz(source) + reused_cache, reused_manifest = materialize_extxyz(source) + assert reused_cache == first_cache + assert reused_manifest == first_manifest + + _write(source, _frame(energy=-2.0)) + second_cache, _ = materialize_extxyz(source) + assert second_cache != first_cache + paths = expand_extxyz_cache(second_cache) + assert paths is not None + assert _energies(paths) == [-2.0] + + +def test_normalization_does_not_modify_parsed_input_or_json_file( + tmp_path: Path, +) -> None: + source = _write(tmp_path / "source.extxyz", _frame()) + input_file = tmp_path / "input.json" + config = _partial_config(str(source), str(source)) + input_file.write_text(json.dumps(config, indent=2), encoding="utf-8") + before = input_file.read_bytes() + parsed = json.loads(input_file.read_text(encoding="utf-8")) + snapshot = copy.deepcopy(parsed) + + normalized = normalize_extxyz_training_data(parsed) + + assert parsed == snapshot + assert input_file.read_bytes() == before + assert normalized is not parsed + assert normalized["training"]["training_data"]["systems"] != str(source) + + +def test_schema_documents_extxyz_paths() -> None: + from deepmd.utils.argcheck import ( + training_data_args, + validation_data_args, + ) + + assert ".extxyz" in training_data_args()["systems"].doc + assert "dpdata" in training_data_args()["systems"].doc + assert ".extxyz" in validation_data_args()["systems"].doc + + +def test_normalize_and_initialize_training_and_validation_data( + tmp_path: Path, +) -> None: + from deepmd.dpmodel.loss.ener import ( + EnergyLoss, + ) + from deepmd.utils.argcheck import ( + normalize, + ) + from deepmd.utils.data_system import ( + get_data, + process_systems, + ) + + stress = np.array([1.0, 2.0, 3.0, 0.4, 0.5, 0.6]) + virial = np.arange(1.0, 10.0).reshape(3, 3) + training_source = _write( + tmp_path / "train.extxyz", + _frame(energy=-1.0, tensor_name="stress", tensor=stress), + ) + validation_source = _write( + tmp_path / "validation.extxyz", + _frame(energy=-2.0, tensor_name="virial", tensor=virial), + ) + repository_root = Path(__file__).resolve().parents[3] + config = json.loads( + (repository_root / "examples" / "water" / "se_e2_a" / "input.json").read_text( + encoding="utf-8" + ) + ) + config["training"]["training_data"] = { + "systems": [str(training_source)], + "batch_size": 1, + } + config["training"]["validation_data"] = { + "systems": [str(validation_source)], + "batch_size": 1, + "numb_btch": 1, + } + config["loss"]["start_pref_v"] = 1.0 + config["loss"]["limit_pref_v"] = 1.0 + original = copy.deepcopy(config) + + normalized = normalize(config) + + assert config == original + train_params = normalized["training"]["training_data"] + validation_params = normalized["training"]["validation_data"] + train_paths = process_systems(train_params["systems"]) + validation_paths = process_systems(validation_params["systems"]) + # Neighbor-statistics and loader construction both call process_systems; + # repeated expansion must resolve to exactly the same cache paths. + assert process_systems(train_params["systems"]) == train_paths + assert process_systems(validation_params["systems"]) == validation_paths + + train_data = get_data(train_params, 6.0, ["O", "H"], None) + validation_data = get_data(validation_params, 6.0, ["O", "H"], None) + assert train_data.system_dirs == train_paths + assert validation_data.system_dirs == validation_paths + + loss = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=1.0, + limit_pref_v=1.0, + ) + for data, expected_energy in ( + (train_data, -1.0), + (validation_data, -2.0), + ): + data.add_data_requirements(loss.label_requirement) + batch = data.get_batch(0) + np.testing.assert_allclose(batch["energy"].reshape(-1), [expected_energy]) + assert np.all(batch["find_energy"] == 1.0) + assert np.all(batch["find_force"] == 1.0) + assert np.all(batch["find_virial"] == 1.0)