From ce621f321f3e4fa44deb48023e04ea5384c27017 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 28 Jul 2026 16:28:19 +0800 Subject: [PATCH 1/5] perf(test): stream dp test so a system need not fit in memory Testing read every frame of a system before looking at any of them. On an LMDB of 38.8 million frames that meant decoding the whole 54 GB set into memory, and `-n` was applied only afterwards, so the default run of 100 frames paid for all of them while the GPU sat idle. Frames are now chosen before anything is read -- atom counts come from the LMDB metadata, so grouping, shuffling and truncation are index arithmetic -- and a system is walked in chunks bounded by atom count, which also lifts the ceiling on how large a dataset can be tested. Chunking is exact rather than an approximation. An MAE is the mean of the absolute errors and an RMSE the root of the mean of their squares, so both are recovered from partial results weighted by the number of elements each was taken over. `merge_weighted_errors` performs that combination, and `weighted_average`, which already combined systems the same way, is now expressed in terms of it. The five test routines had drifted into five shapes of one thing. They are now a single skeleton -- declare the labels, evaluate chunk by chunk, combine, report -- with a tester per model class supplying only what differs. Choosing a tester replaces two parallel isinstance chains, one selecting the routine and one its printer, because the report of a tester drives the per-system and the run-level table alike. That machinery lives in `deepmd/infer/model_test`, beside `model_devi`, which is the same kind of backend-independent analysis driven by a command-line entry point. `deepmd/entrypoints/test.py` keeps system discovery and the run-level report, which returns it to the size of every other entry point. A spin model now reports the virial and the stress. Its magnetic degrees of freedom reach the virial only through the virtual atoms, whose displacement the model already removes, so the virial is with respect to the real atomic positions as for any other energy model; the quantity was computed and then discarded. Deriving the spin tester from the energy one leaves it overriding the force alone, and the exclusions vanish with the branches that carried them. Two defects surface in the same code. A system split into sub-groups, as a mixed-nloc LMDB is, had each group overwrite the detail file of the last, because the append flag tracked the system instead of the group. And `test_wfc` was unreachable while the dispatch still carried its printer, so a wave-function model would have failed on an unbound name. Verified against the previous implementation over the same frames: 144 metrics agree to 8.3e-07, below the 9.98e-07 spread between two runs of one unchanged implementation, which is nondeterministic on this GPU. Chunked and unchunked evaluation agree to 4.6e-07. Tensor detail headers are identical to the previous ones for both model classes, atomic and not, over a range of selected atom counts. The pt, pt_expt and dpmodel suites pass, and a JAX model was trained, frozen and tested end to end. The TensorFlow backend is not built here, so its path was reviewed rather than run. --- deepmd/dpmodel/utils/lmdb_data.py | 283 +++- deepmd/entrypoints/test.py | 1490 +---------------- deepmd/infer/model_test/__init__.py | 111 ++ deepmd/infer/model_test/base.py | 310 ++++ deepmd/infer/model_test/dos.py | 121 ++ deepmd/infer/model_test/ener.py | 703 ++++++++ deepmd/infer/model_test/property.py | 119 ++ deepmd/infer/model_test/tensor.py | 157 ++ deepmd/utils/data.py | 41 + deepmd/utils/weight_avg.py | 66 +- .../tests/common/test_dp_test_ener_split.py | 2 +- source/tests/pt/test_dp_test.py | 12 +- source/tests/pt/test_weighted_avg.py | 62 +- 13 files changed, 1899 insertions(+), 1578 deletions(-) create mode 100644 deepmd/infer/model_test/__init__.py create mode 100644 deepmd/infer/model_test/base.py create mode 100644 deepmd/infer/model_test/dos.py create mode 100644 deepmd/infer/model_test/ener.py create mode 100644 deepmd/infer/model_test/property.py create mode 100644 deepmd/infer/model_test/tensor.py diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index a85f92be87..f84d6c860d 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -25,6 +25,9 @@ from dataclasses import ( dataclass, ) +from itertools import ( + pairwise, +) from pathlib import ( Path, ) @@ -2318,6 +2321,7 @@ def __init__( lmdb_path: str, type_map: list[str] | None = None, shuffle_test: bool = True, + max_frames: float | None = None, **kwargs: Any, ) -> None: self.lmdb_path = str(lmdb_path) @@ -2353,52 +2357,115 @@ def __init__( f"model {self._type_map}, remap={remap}" ) - # Read all frames - self._frames: list[dict[str, Any]] = [] - with self._env.begin() as txn: - for i in range(self.nframes): - key = format(i, self._frame_fmt).encode() - raw = txn.get(key) - if raw is not None: - frame = _remap_keys(_decode_frame(raw)) - # Apply type remapping to atype - if ( - self._type_remap is not None - and "atype" in frame - and isinstance(frame["atype"], np.ndarray) - ): - frame["atype"] = _remap_atom_types( - frame["atype"].reshape(-1), self._type_remap - ) - self._frames.append(frame) - - # Shuffle if requested - if shuffle_test: - rng = np.random.default_rng() - indices = rng.permutation(len(self._frames)) - self._frames = [self._frames[i] for i in indices] - - # Group frames by nloc - self._nloc_groups: dict[int, list[int]] = {} - for idx, frame in enumerate(self._frames): - atype = frame.get("atype") - nloc = len(atype) if isinstance(atype, np.ndarray) else self._natoms - self._nloc_groups.setdefault(nloc, []).append(idx) + # Select the frames to test, without reading any of them. Atom counts + # come from the metadata, so grouping, shuffling and truncation are + # index arithmetic; only the retained frames are ever decoded. + self._nloc_groups = self._select_frames(meta, shuffle_test, max_frames) # Data requirements self._requirements: dict[str, dict[str, Any]] = {} - # Detect PBC: if any frame has a non-zero box + # Detect PBC from the first retained frame. self.pbc = True - if len(self._frames) > 0: - f0 = self._frames[0] - if "box" not in f0: - self.pbc = False - elif isinstance(f0["box"], np.ndarray) and np.allclose(f0["box"], 0.0): + first = next( + (indices[0] for indices in self._nloc_groups.values() if len(indices)), None + ) + probe = self._read_frames([first]) if first is not None else [] + if probe: + box = probe[0].get("box") + if not isinstance(box, np.ndarray) or np.allclose(box, 0.0): self.pbc = False self.mixed_type = True + def _select_frames( + self, + meta: dict, + shuffle_test: bool, + max_frames: float | None, + ) -> dict[int, list[int]]: + """Group the frame indices by atom count, then sample each group. + + Parameters + ---------- + meta : dict + The LMDB metadata. ``frame_nlocs`` gives the atom count of every + frame; without it the frames have to be scanned for it. + shuffle_test : bool + Whether the frames of a group are drawn in random order. + max_frames : int or float or None + Upper bound on the number of frames retained per group. ``None`` + and a non-finite bound retain the whole group. + + Returns + ------- + dict[int, list[int]] + The retained LMDB frame indices of each atom count. + """ + raw_nlocs = meta.get("frame_nlocs") + if _is_encoded_array(raw_nlocs): + nlocs = _decode_array(raw_nlocs).reshape(-1).astype(np.int64) + elif raw_nlocs is not None: + nlocs = np.asarray(raw_nlocs, dtype=np.int64) + else: + nlocs = np.asarray( + _scan_frame_nlocs( + self._env, self.nframes, self._frame_fmt, self._natoms + ), + dtype=np.int64, + ) + + # Sorting once groups every atom count in a single pass, which matters + # for datasets whose frame count reaches into the millions. + order = np.argsort(nlocs, kind="stable") + starts = np.concatenate( + ([0], np.flatnonzero(np.diff(nlocs[order])) + 1, [nlocs.size]) + ) + + keep = ( + None + if max_frames is None or not np.isfinite(max_frames) + else int(max_frames) + ) + rng = np.random.default_rng() + groups: dict[int, list[int]] = {} + for begin, end in pairwise(starts): + indices = order[begin:end] + if shuffle_test: + indices = rng.permutation(indices) + if keep is not None: + indices = indices[:keep] + groups[int(nlocs[order[begin]])] = indices.tolist() + return groups + + def _read_frames(self, frame_indices: list[int]) -> list[dict[str, Any]]: + """Decode the given LMDB frames, applying the type remapping. + + Parameters + ---------- + frame_indices : list[int] + Indices of the frames to read, as keyed in the LMDB. + + Returns + ------- + list[dict[str, Any]] + One decoded frame per index that the LMDB holds. + """ + frames: list[dict[str, Any]] = [] + with self._env.begin() as txn: + for index in frame_indices: + raw = txn.get(format(index, self._frame_fmt).encode()) + if raw is None: + continue + frame = _remap_keys(_decode_frame(raw)) + atype = frame.get("atype") + if self._type_remap is not None and isinstance(atype, np.ndarray): + frame["atype"] = _remap_atom_types( + atype.reshape(-1), self._type_remap + ) + frames.append(frame) + return frames + def __del__(self) -> None: """Release the LMDB environment ref-count on garbage collection.""" path = getattr(self, "lmdb_path", None) @@ -2407,39 +2474,57 @@ def __del__(self) -> None: @property def nloc_groups(self) -> dict[int, list[int]]: - """Nloc → list of frame indices in self._frames.""" + """Nloc → the LMDB frame indices retained for that atom count.""" return self._nloc_groups @staticmethod def _frame_has_data(frame: dict[str, Any], key: str) -> bool: - """Resolve one frame's explicit or inferred ``find_*`` value.""" + """Resolve one frame's explicit or inferred ``find_*`` value. + + The frame may still carry its msgpack payload, so that availability + can be settled without decoding the arrays it describes. + """ find_key = f"find_{key}" if find_key in frame: - return bool(float(np.asarray(frame[find_key]).item())) + return bool(float(np.asarray(_decode_value(frame[find_key])).item())) value = frame.get(key) + if _is_encoded_array(value): + return True return isinstance(value, (np.ndarray, np.generic, int, float, bool)) @property def find_signature_groups( self, ) -> dict[tuple[int, tuple[tuple[str, bool], ...]], list[int]]: - """Group frames by atom count and scalar label availability.""" + """Group the retained frame indices by atom count and label availability. + + Frames that :meth:`_stack_frames` would refuse to stack together land + in different groups, because both settle availability with + :meth:`_frame_has_data`. Only the msgpack payload of each frame is + read, so the grouping does not decode the arrays it separates. + """ groups: dict[tuple[int, tuple[tuple[str, bool], ...]], list[int]] = {} - for index, frame in enumerate(self._frames): - atype = frame.get("atype") - nloc = len(atype) if isinstance(atype, np.ndarray) else self._natoms - signature = tuple( - (f"find_{key}", self._frame_has_data(frame, key)) - for key in _availability_signature_keys(frame, iter(self._requirements)) - ) - groups.setdefault((nloc, signature), []).append(index) + with self._env.begin() as transaction: + for nloc, frame_indices in self._nloc_groups.items(): + for index in frame_indices: + raw = transaction.get(format(index, self._frame_fmt).encode()) + if raw is None: + continue + frame = _remap_keys(msgpack.unpackb(raw, raw=False)) + signature = tuple( + (f"find_{key}", self._frame_has_data(frame, key)) + for key in _availability_signature_keys( + frame, iter(self._requirements) + ) + ) + groups.setdefault((nloc, signature), []).append(index) return groups def get_test_by_indices(self, frame_indices: list[int]) -> dict[str, Any]: """Stack one homogeneous validation group selected by frame index.""" if not frame_indices: raise ValueError("frame_indices must contain at least one frame") - frames = [self._frames[index] for index in frame_indices] + frames = self._read_frames(frame_indices) nlocs = { len(frame["atype"]) for frame in frames @@ -2519,30 +2604,81 @@ def get_test(self, nloc: int | None = None) -> dict[str, Any]: If None and mixed nloc, return the largest group and log a warning. Returns dict matching DeepmdData.get_test() format: """ + frame_indices, natoms = self._resolve_group(nloc) + return self._stack_frames(self._read_frames(frame_indices), natoms) + + def iter_test( + self, + *, + chunk_atoms: int, + numb_test: float = float("inf"), + nloc: int | None = None, + ) -> Iterator[dict[str, Any]]: + """Yield the test frames in chunks, reading each chunk on demand. + + Only the frames of the chunk being yielded are held, which is what + makes a dataset larger than memory testable. + + Parameters + ---------- + chunk_atoms : int + Upper bound on the number of atoms per chunk. A chunk always + carries at least one frame, however many atoms it has. + numb_test : float, optional + Upper bound on the number of frames served. A non-finite bound + serves every frame of the group. + nloc : int or None, optional + Atom count selecting the group, resolved as in :meth:`get_test`. + + Yields + ------ + dict[str, Any] + One chunk of frames, stacked as :meth:`get_test` stacks them. + """ + frame_indices, natoms = self._resolve_group(nloc) + if np.isfinite(numb_test): + frame_indices = frame_indices[: int(numb_test)] + step = max(1, int(chunk_atoms) // max(1, natoms)) + for begin in range(0, len(frame_indices), step): + chunk = frame_indices[begin : begin + step] + yield self._stack_frames(self._read_frames(chunk), natoms) + + def _resolve_group(self, nloc: int | None) -> tuple[list[int], int]: + """Return the retained frame indices and atom count of one group. + + Parameters + ---------- + nloc : int or None + The atom count to select. ``None`` selects the only group when the + dataset is uniform, and the largest group otherwise. + + Returns + ------- + tuple[list[int], int] + The LMDB frame indices of the group and its atom count. + + Raises + ------ + ValueError + If no frame has the requested atom count. + """ if nloc is not None: if nloc not in self._nloc_groups: raise ValueError( f"No frames with nloc={nloc}. Available: {sorted(self._nloc_groups.keys())}" ) - frame_indices = self._nloc_groups[nloc] - natoms = nloc - elif len(self._nloc_groups) == 1: - # Uniform nloc — use all frames + return self._nloc_groups[nloc], nloc + if len(self._nloc_groups) == 1: natoms = next(iter(self._nloc_groups)) - frame_indices = list(range(len(self._frames))) - else: - # Mixed nloc — use the largest group - natoms = max(self._nloc_groups, key=lambda k: len(self._nloc_groups[k])) - frame_indices = self._nloc_groups[natoms] - group_summary = {k: len(v) for k, v in sorted(self._nloc_groups.items())} - log.warning( - f"Mixed-nloc LMDB for dp test: using nloc={natoms} group " - f"({len(frame_indices)} frames). " - f"Available groups: {group_summary}" - ) - - frames = [self._frames[i] for i in frame_indices] - return self._stack_frames(frames, natoms) + return self._nloc_groups[natoms], natoms + natoms = max(self._nloc_groups, key=lambda k: len(self._nloc_groups[k])) + group_summary = {k: len(v) for k, v in sorted(self._nloc_groups.items())} + log.warning( + f"Mixed-nloc LMDB for dp test: using nloc={natoms} group " + f"({len(self._nloc_groups[natoms])} frames). " + f"Available groups: {group_summary}" + ) + return self._nloc_groups[natoms], natoms def _stack_frames( self, frames: list[dict[str, Any]], natoms: int @@ -2684,6 +2820,17 @@ def get_test(self) -> dict[str, Any]: return self._inner.get_test_by_indices(self._frame_indices) return self._inner.get_test(nloc=self._nloc) + def iter_test( + self, + *, + chunk_atoms: int, + numb_test: float = float("inf"), + ) -> Iterator[dict[str, Any]]: + """Yield this group's frames in chunks.""" + return self._inner.iter_test( + chunk_atoms=chunk_atoms, numb_test=numb_test, nloc=self._nloc + ) + def _copy_lmdb_source( src_path: str, diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index 5057e04e0f..066461ce1d 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -1,18 +1,18 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Test trained DeePMD model.""" +"""Command-line entry point of ``dp test``. + +System discovery and the run-level report live here; evaluating a system is +the business of :mod:`deepmd.infer.model_test`. +""" import logging from pathlib import ( Path, ) from typing import ( - TYPE_CHECKING, Any, - NamedTuple, ) -import numpy as np - from deepmd.common import ( j_loader, ) @@ -21,27 +21,11 @@ LmdbTestDataNlocView, is_lmdb, ) -from deepmd.infer.deep_dipole import ( - DeepDipole, -) -from deepmd.infer.deep_dos import ( - DeepDOS, -) from deepmd.infer.deep_eval import ( DeepEval, ) -from deepmd.infer.deep_polar import ( - DeepGlobalPolar, - DeepPolar, -) -from deepmd.infer.deep_pot import ( - DeepPot, -) -from deepmd.infer.deep_property import ( - DeepProperty, -) -from deepmd.infer.deep_wfc import ( - DeepWFC, +from deepmd.infer.model_test import ( + build_tester, ) from deepmd.utils import random as dp_random from deepmd.utils.compat import ( @@ -53,27 +37,10 @@ from deepmd.utils.data_system import ( process_systems, ) -from deepmd.utils.eval_metrics import ( - DP_TEST_HESSIAN_METRIC_KEYS, - DP_TEST_SPIN_WEIGHTED_METRIC_KEYS, - DP_TEST_WEIGHTED_FORCE_METRIC_KEYS, - DP_TEST_WEIGHTED_METRIC_KEYS, - compute_energy_type_metrics, - compute_error_stat, - compute_spin_force_metrics, - compute_weighted_error_stat, - mae, - rmse, -) from deepmd.utils.weight_avg import ( - weighted_average, + merge_weighted_errors, ) -if TYPE_CHECKING: - from deepmd.infer.deep_tensor import ( - DeepTensor, - ) - __all__ = ["test"] log = logging.getLogger(__name__) @@ -178,6 +145,7 @@ def test( # init model dp = DeepEval(model, head=head) + tester = build_tester(dp, atomic=atomic) for cc, system in enumerate(all_sys): log.info("# ---------------output of dp test--------------- ") @@ -190,6 +158,7 @@ def test( system, type_map=tmap, shuffle_test=shuffle_test, + max_frames=numb_test, ) # For mixed-nloc LMDB, test each nloc group separately nloc_keys = sorted(lmdb_data.nloc_groups.keys()) @@ -219,55 +188,20 @@ def test( for data, sys_label in data_items: if sys_label != system: log.info(f"# testing sub-group : {sys_label}") + # Only the very first tested group writes a fresh detail file; a + # system split into sub-groups extends it like any later system. + append_detail = bool(err_coll) - if isinstance(dp, DeepPot): - err = test_ener( - dp, - data, - sys_label, - numb_test, - detail_file, - atomic, - append_detail=(cc != 0), - ) - elif isinstance(dp, DeepDOS): - err = test_dos( - dp, - data, - sys_label, - numb_test, - detail_file, - atomic, - append_detail=(cc != 0), - ) - elif isinstance(dp, DeepProperty): - err = test_property( - dp, - data, - sys_label, - numb_test, - detail_file, - atomic, - append_detail=(cc != 0), - ) - elif isinstance(dp, DeepDipole): - err = test_dipole(dp, data, numb_test, detail_file, atomic) - elif isinstance(dp, DeepPolar): - err = test_polar(dp, data, numb_test, detail_file, atomic=atomic) - elif isinstance( - dp, DeepGlobalPolar - ): # should not appear in this new version - log.warning( - "Global polar model is not currently supported. Please directly use the polar mode and change loss parameters." - ) - err = test_polar( - dp, data, numb_test, detail_file, atomic=False - ) # YWolfeee: downward compatibility + err = tester.run( + data, + sys_label, + numb_test, + detail_file, + append_detail=append_detail, + ) log.info("# ----------------------------------------------- ") err_coll.append(err) - avg_err = weighted_average(err_coll) - # For mixed-nloc LMDB, err_coll may have more entries than all_sys # (one per nloc group per system). Only warn if fewer. if len(err_coll) < len(all_sys): @@ -275,1387 +209,5 @@ def test( log.info("# ----------weighted average of errors----------- ") log.info(f"# number of systems : {len(all_sys)}") - if isinstance(dp, DeepPot): - print_ener_sys_avg(avg_err) - elif isinstance(dp, DeepDOS): - print_dos_sys_avg(avg_err) - elif isinstance(dp, DeepProperty): - print_property_sys_avg(avg_err) - elif isinstance(dp, DeepDipole): - print_dipole_sys_avg(avg_err) - elif isinstance(dp, DeepPolar): - print_polar_sys_avg(avg_err) - elif isinstance(dp, DeepGlobalPolar): - print_polar_sys_avg(avg_err) - elif isinstance(dp, DeepWFC): - print_wfc_sys_avg(avg_err) + tester.log_errors(merge_weighted_errors(err_coll)) log.info("# ----------------------------------------------- ") - - -def save_txt_file( - fname: Path, data: np.ndarray, header: str = "", append: bool = False -) -> None: - """Save numpy array to test file. - - Parameters - ---------- - fname : str - filename - data : np.ndarray - data to save to disk - header : str, optional - header string to use in file, by default "" - append : bool, optional - if true file will be appended instead of overwriting, by default False - """ - flags = "a" if append else "w" - with fname.open(flags, encoding="utf-8") as fp: - np.savetxt(fp, data, header=header) - - -def _reshape_force_by_atom(force_array: np.ndarray, natoms: int) -> np.ndarray: - """Reshape flattened force arrays into `[nframes, natoms, 3]`.""" - return np.reshape(force_array, [-1, natoms, 3]) - - -def _concat_force_rows( - force_blocks: list[np.ndarray], dtype: np.dtype | type[np.generic] -) -> np.ndarray: - """Concatenate per-frame force rows into one 2D array.""" - if not force_blocks: - return np.empty((0, 3), dtype=dtype) - return np.concatenate(force_blocks, axis=0) - - -def _align_spin_force_arrays( - *, - dp: "DeepPot", - atype: np.ndarray, - natoms: int, - prediction_force: np.ndarray, - reference_force: np.ndarray, - prediction_force_mag: np.ndarray | None, - reference_force_mag: np.ndarray | None, - mask_mag: np.ndarray | None, -) -> tuple[np.ndarray, np.ndarray, np.ndarray | None, np.ndarray | None]: - """Align spin force arrays into real-atom and magnetic subsets.""" - prediction_force_by_atom = _reshape_force_by_atom(prediction_force, natoms) - reference_force_by_atom = _reshape_force_by_atom(reference_force, natoms) - if dp.get_ntypes_spin() != 0: # old tf support for spin - ntypes_real = dp.get_ntypes() - dp.get_ntypes_spin() - atype_by_frame = np.reshape(atype, [-1, natoms]) - if atype_by_frame.shape[0] == 1 and prediction_force_by_atom.shape[0] != 1: - atype_by_frame = np.broadcast_to( - atype_by_frame, - (prediction_force_by_atom.shape[0], natoms), - ) - if atype_by_frame.shape[0] != prediction_force_by_atom.shape[0]: - raise ValueError( - "Spin atom types and force arrays must have matching frames." - ) - force_real_prediction_chunks = [] - force_real_reference_chunks = [] - force_magnetic_prediction_chunks = [] - force_magnetic_reference_chunks = [] - for frame_atype, frame_prediction, frame_reference in zip( - atype_by_frame, - prediction_force_by_atom, - reference_force_by_atom, - strict=False, - ): - real_mask = frame_atype < ntypes_real - magnetic_mask = ~real_mask - force_real_prediction_chunks.append(frame_prediction[real_mask]) - force_real_reference_chunks.append(frame_reference[real_mask]) - force_magnetic_prediction_chunks.append(frame_prediction[magnetic_mask]) - force_magnetic_reference_chunks.append(frame_reference[magnetic_mask]) - return ( - _concat_force_rows( - force_real_prediction_chunks, - prediction_force_by_atom.dtype, - ), - _concat_force_rows( - force_real_reference_chunks, - reference_force_by_atom.dtype, - ), - _concat_force_rows( - force_magnetic_prediction_chunks, - prediction_force_by_atom.dtype, - ), - _concat_force_rows( - force_magnetic_reference_chunks, - reference_force_by_atom.dtype, - ), - ) - - force_real_prediction = prediction_force_by_atom.reshape(-1, 3) - force_real_reference = reference_force_by_atom.reshape(-1, 3) - if prediction_force_mag is None or reference_force_mag is None or mask_mag is None: - return force_real_prediction, force_real_reference, None, None - magnetic_mask = mask_mag.reshape(-1).astype(bool) - return ( - force_real_prediction, - force_real_reference, - prediction_force_mag.reshape(-1, 3)[magnetic_mask], - reference_force_mag.reshape(-1, 3)[magnetic_mask], - ) - - -def _write_energy_test_details( - *, - detail_path: Path, - system: str, - natoms: int, - append_detail: bool, - reference_energy: np.ndarray, - prediction_energy: np.ndarray, - reference_force: np.ndarray, - prediction_force: np.ndarray, - reference_virial: np.ndarray | None, - prediction_virial: np.ndarray | None, - out_put_spin: bool, - reference_stress: np.ndarray | None = None, - prediction_stress: np.ndarray | None = None, - reference_force_real: np.ndarray | None = None, - prediction_force_real: np.ndarray | None = None, - reference_force_magnetic: np.ndarray | None = None, - prediction_force_magnetic: np.ndarray | None = None, - reference_hessian: np.ndarray | None = None, - prediction_hessian: np.ndarray | None = None, -) -> None: - """Write energy-type detail outputs after arrays have been aligned.""" - pe = np.concatenate( - ( - np.reshape(reference_energy, [-1, 1]), - np.reshape(prediction_energy, [-1, 1]), - ), - axis=1, - ) - save_txt_file( - detail_path.with_suffix(".e.out"), - pe, - header=f"{system}: data_e pred_e", - append=append_detail, - ) - pe_atom = pe / natoms - save_txt_file( - detail_path.with_suffix(".e_peratom.out"), - pe_atom, - header=f"{system}: data_e pred_e", - append=append_detail, - ) - if not out_put_spin: - pf = np.concatenate( - ( - np.reshape(reference_force, [-1, 3]), - np.reshape(prediction_force, [-1, 3]), - ), - axis=1, - ) - save_txt_file( - detail_path.with_suffix(".f.out"), - pf, - header=f"{system}: data_fx data_fy data_fz pred_fx pred_fy pred_fz", - append=append_detail, - ) - else: - if reference_force_real is None or prediction_force_real is None: - raise ValueError("Spin detail output requires aligned real-atom forces.") - pf_real = np.concatenate( - ( - np.reshape(reference_force_real, [-1, 3]), - np.reshape(prediction_force_real, [-1, 3]), - ), - axis=1, - ) - save_txt_file( - detail_path.with_suffix(".fr.out"), - pf_real, - header=f"{system}: data_fx data_fy data_fz pred_fx pred_fy pred_fz", - append=append_detail, - ) - if (reference_force_magnetic is None) != (prediction_force_magnetic is None): - raise ValueError( - "Spin magnetic detail output requires both reference and prediction forces." - ) - if ( - reference_force_magnetic is not None - and prediction_force_magnetic is not None - ): - pf_mag = np.concatenate( - ( - np.reshape(reference_force_magnetic, [-1, 3]), - np.reshape(prediction_force_magnetic, [-1, 3]), - ), - axis=1, - ) - save_txt_file( - detail_path.with_suffix(".fm.out"), - pf_mag, - header=f"{system}: data_fmx data_fmy data_fmz pred_fmx pred_fmy pred_fmz", - append=append_detail, - ) - if (reference_virial is None) != (prediction_virial is None): - raise ValueError( - "Virial detail output requires both reference and prediction virials." - ) - if reference_virial is not None and prediction_virial is not None: - pv = np.concatenate( - ( - np.reshape(reference_virial, [-1, 9]), - np.reshape(prediction_virial, [-1, 9]), - ), - axis=1, - ) - save_txt_file( - detail_path.with_suffix(".v.out"), - pv, - header=f"{system}: data_vxx data_vxy data_vxz data_vyx data_vyy " - "data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx " - "pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz", - append=append_detail, - ) - pv_atom = pv / natoms - save_txt_file( - detail_path.with_suffix(".v_peratom.out"), - pv_atom, - header=f"{system}: data_vxx data_vxy data_vxz data_vyx data_vyy " - "data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx " - "pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz", - append=append_detail, - ) - if (reference_stress is None) != (prediction_stress is None): - raise ValueError( - "Stress detail output requires both reference and prediction stresses." - ) - if reference_stress is not None and prediction_stress is not None: - ps = np.concatenate( - ( - np.reshape(reference_stress, [-1, 9]), - np.reshape(prediction_stress, [-1, 9]), - ), - axis=1, - ) - save_txt_file( - detail_path.with_suffix(".s.out"), - ps, - header=f"{system} (eV/Å^3): data_sxx data_sxy data_sxz data_syx " - "data_syy data_syz data_szx data_szy data_szz pred_sxx pred_sxy pred_sxz " - "pred_syx pred_syy pred_syz pred_szx pred_szy pred_szz", - append=append_detail, - ) - if reference_hessian is not None and prediction_hessian is not None: - hessian_detail = np.concatenate( - ( - reference_hessian.reshape(-1, 1), - prediction_hessian.reshape(-1, 1), - ), - axis=1, - ) - save_txt_file( - detail_path.with_suffix(".h.out"), - hessian_detail, - header=f"{system}: data_h pred_h (3Na*3Na matrix in row-major order)", - append=append_detail, - ) - - -class _OptionalEnerOutputs(NamedTuple): - """The optional trailing outputs of ``DeepPot.eval`` (``None`` when absent).""" - - atom_energy: "np.ndarray | None" - atom_virial: "np.ndarray | None" - force_mag: "np.ndarray | None" - mask_mag: "np.ndarray | None" - hessian: "np.ndarray | None" - - -def _split_optional_ener_outputs( - ret: tuple, - *, - has_atom_ener: bool, - has_spin: bool, - has_hessian: bool, - numb_test: int, -) -> _OptionalEnerOutputs: - """Split the optional trailing outputs of ``DeepPot.eval``. - - ``DeepPot.eval`` appends its optional outputs after ``(energy, force, - virial)`` in a fixed order: atomic ``(atom_energy, atom_virial)``, then spin - ``(force_mag, mask_mag)``, then ``hessian``. Read them by advancing an index - through the tuple in that same order, so the hessian slot is not confused - with atomic energy/virial or spin outputs when those are also present. - """ - atom_energy = atom_virial = force_mag = mask_mag = hessian = None - idx = 3 - if has_atom_ener: - atom_energy = ret[idx].reshape([numb_test, -1]) - atom_virial = ret[idx + 1].reshape([numb_test, -1]) - idx += 2 - if has_spin: - force_mag = ret[idx].reshape([numb_test, -1]) - mask_mag = ret[idx + 1].reshape([numb_test, -1]) - idx += 2 - if has_hessian: - hessian = ret[idx].reshape([numb_test, -1]) - return _OptionalEnerOutputs(atom_energy, atom_virial, force_mag, mask_mag, hessian) - - -def test_ener( - dp: "DeepPot", - data: DeepmdData, - system: str, - numb_test: int, - detail_file: str | None, - has_atom_ener: bool, - append_detail: bool = False, -) -> dict[str, tuple[float, float]]: - """Test energy type model. - - Parameters - ---------- - dp : DeepPot - instance of deep potential - data : DeepmdData - data container object - system : str - system directory - numb_test : int - munber of tests to do - detail_file : Optional[str] - file where test details will be output - has_atom_ener : bool - whether per atom quantities should be computed - append_detail : bool, optional - if true append output detail file, by default False - - Returns - ------- - dict[str, tuple[float, float]] - weighted-average-ready metric pairs - """ - dict_to_return = {} - - data.add("energy", 1, atomic=False, must=False, high_prec=True) - data.add("force", 3, atomic=True, must=False, high_prec=False) - data.add("atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3) - data.add("virial", 9, atomic=False, must=False, high_prec=False) - if dp.has_efield: - data.add("efield", 3, atomic=True, must=True, high_prec=False) - if has_atom_ener: - data.add("atom_ener", 1, atomic=True, must=True, high_prec=False) - if dp.get_dim_fparam() > 0: - data.add( - "fparam", - dp.get_dim_fparam(), - atomic=False, - must=not dp.has_default_fparam(), - high_prec=False, - ) - if dp.get_dim_aparam() > 0: - data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False) - if dp.has_chg_spin_ebd(): - data.add( - "charge_spin", - 2, - atomic=False, - must=not dp.has_default_chg_spin(), - high_prec=False, - ) - if dp.has_spin: - data.add("spin", 3, atomic=True, must=True, high_prec=False) - data.add("force_mag", 3, atomic=True, must=False, high_prec=False) - if dp.has_hessian: - data.add("hessian", 1, atomic=True, must=True, high_prec=False) - - test_data = data.get_test() - find_energy = test_data.get("find_energy") - find_force = test_data.get("find_force") - find_virial = test_data.get("find_virial") - find_force_mag = test_data.get("find_force_mag") - find_atom_pref = test_data.get("find_atom_pref") - mixed_type = data.mixed_type - natoms = len(test_data["type"][0]) - nframes = test_data["box"].shape[0] - numb_test = min(nframes, numb_test) - - coord = test_data["coord"][:numb_test].reshape([numb_test, -1]) - box = test_data["box"][:numb_test] - if dp.has_efield: - efield = test_data["efield"][:numb_test].reshape([numb_test, -1]) - else: - efield = None - if dp.has_spin: - spin = test_data["spin"][:numb_test].reshape([numb_test, -1]) - else: - spin = None - if not data.pbc: - box = None - if mixed_type: - atype = test_data["type"][:numb_test].reshape([numb_test, -1]) - else: - atype = test_data["type"][0] - if dp.get_dim_fparam() > 0 and test_data["find_fparam"] != 0.0: - fparam = test_data["fparam"][:numb_test] - else: - fparam = None - if dp.get_dim_aparam() > 0: - aparam = test_data["aparam"][:numb_test] - else: - aparam = None - if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 0.0) != 0.0: - charge_spin = test_data["charge_spin"][:numb_test] - else: - charge_spin = None - - ret = dp.eval( - coord, - box, - atype, - fparam=fparam, - aparam=aparam, - atomic=has_atom_ener, - efield=efield, - mixed_type=mixed_type, - spin=spin, - charge_spin=charge_spin, - ) - energy = ret[0] - force = ret[1] - virial = ret[2] - energy = energy.reshape([numb_test, 1]) - force = force.reshape([numb_test, -1]) - virial = virial.reshape([numb_test, 9]) - optional_outputs = _split_optional_ener_outputs( - ret, - has_atom_ener=has_atom_ener, - has_spin=dp.has_spin, - has_hessian=dp.has_hessian, - numb_test=numb_test, - ) - ae = optional_outputs.atom_energy - force_m = optional_outputs.force_mag - mask_mag = optional_outputs.mask_mag - hessian = optional_outputs.hessian - out_put_spin = dp.get_ntypes_spin() != 0 or dp.has_spin - spin_metrics = None - force_r = None - test_force_r = None - test_force_m = None - if out_put_spin: - force_r, test_force_r, force_m, test_force_m = _align_spin_force_arrays( - dp=dp, - atype=atype, - natoms=natoms, - prediction_force=force, - reference_force=test_data["force"][:numb_test], - prediction_force_mag=force_m, - reference_force_mag=( - test_data["force_mag"][:numb_test] if "force_mag" in test_data else None - ), - mask_mag=mask_mag, - ) - if find_force_mag == 1 and (force_m is None or test_force_m is None): - raise RuntimeError( - "Spin magnetic force metrics require magnetic force arrays and mask." - ) - spin_metrics = compute_spin_force_metrics( - force_real_prediction=force_r, - force_real_reference=test_force_r, - force_magnetic_prediction=force_m if find_force_mag == 1 else None, - force_magnetic_reference=test_force_m if find_force_mag == 1 else None, - ) - - energy_metric_input = { - "find_energy": find_energy, - "find_force": find_force if not out_put_spin else 0.0, - "find_virial": find_virial if not out_put_spin else 0.0, - "energy": test_data["energy"][:numb_test], - "force": test_data["force"][:numb_test], - } - energy_metric_prediction = { - "energy": energy, - "force": force, - } - if find_virial == 1 and data.pbc and not out_put_spin: - energy_metric_input["virial"] = test_data["virial"][:numb_test] - energy_metric_prediction["virial"] = virial - shared_metrics = compute_energy_type_metrics( - prediction=energy_metric_prediction, - test_data=energy_metric_input, - natoms=natoms, - has_pbc=data.pbc, - ) - dict_to_return.update( - shared_metrics.as_weighted_average_errors(DP_TEST_WEIGHTED_METRIC_KEYS) - ) - - weighted_force_metrics = None - if find_energy == 1: - if shared_metrics.energy is None or shared_metrics.energy_per_atom is None: - raise RuntimeError("Energy metrics are unavailable for dp test.") - mae_e = shared_metrics.energy.mae - rmse_e = shared_metrics.energy.rmse - mae_ea = shared_metrics.energy_per_atom.mae - rmse_ea = shared_metrics.energy_per_atom.rmse - - if not out_put_spin and find_force == 1: - if shared_metrics.force is None: - raise RuntimeError("Force metrics are unavailable for dp test.") - mae_f = shared_metrics.force.mae - rmse_f = shared_metrics.force.rmse - if find_atom_pref == 1: - weighted_force_metrics = compute_weighted_error_stat( - force, - test_data["force"][:numb_test], - test_data["atom_pref"][:numb_test], - ) - mae_fw = weighted_force_metrics.mae - rmse_fw = weighted_force_metrics.rmse - - prediction_stress = None - reference_stress = None - if data.pbc and not out_put_spin and find_virial == 1: - if shared_metrics.virial is None or shared_metrics.virial_per_atom is None: - raise RuntimeError("Virial metrics are unavailable for dp test.") - mae_v = shared_metrics.virial.mae - rmse_v = shared_metrics.virial.rmse - mae_va = shared_metrics.virial_per_atom.mae - rmse_va = shared_metrics.virial_per_atom.rmse - # Stress sigma = -virial / volume, in eV/Å^3 (tensile-positive convention). - volume = np.abs(np.linalg.det(box.reshape([numb_test, 3, 3]))).reshape( - [numb_test, 1] - ) - prediction_stress = -virial / volume - reference_stress = -test_data["virial"][:numb_test] / volume - stress_metrics = compute_error_stat(prediction_stress, reference_stress) - mae_s = stress_metrics.mae - rmse_s = stress_metrics.rmse - dict_to_return.update( - stress_metrics.as_weighted_average_errors("mae_s", "rmse_s") - ) - - hessian_metrics = None - if dp.has_hessian: - hessian_metrics = compute_error_stat( - hessian, - test_data["hessian"][:numb_test], - ) - mae_h = hessian_metrics.mae - rmse_h = hessian_metrics.rmse - if has_atom_ener: - atomic_energy_metrics = compute_error_stat( - ae.reshape([-1]), - test_data["atom_ener"][:numb_test].reshape([-1]), - ) - mae_ae = atomic_energy_metrics.mae - rmse_ae = atomic_energy_metrics.rmse - if out_put_spin: - if spin_metrics is None or spin_metrics.force_real is None: - raise RuntimeError("Spin force metrics are unavailable for dp test.") - mae_fr = spin_metrics.force_real.mae - rmse_fr = spin_metrics.force_real.rmse - if find_force_mag == 1: - if spin_metrics.force_magnetic is None: - raise RuntimeError("Spin magnetic force metrics are unavailable.") - mae_fm = spin_metrics.force_magnetic.mae - rmse_fm = spin_metrics.force_magnetic.rmse - - log.info(f"# number of test data : {numb_test:d} ") - if find_energy == 1: - log.info(f"Energy MAE : {mae_e:e} eV") - log.info(f"Energy RMSE : {rmse_e:e} eV") - log.info(f"Energy MAE/Natoms : {mae_ea:e} eV") - log.info(f"Energy RMSE/Natoms : {rmse_ea:e} eV") - if not out_put_spin and find_force == 1: - log.info(f"Force MAE : {mae_f:e} eV/Å") - log.info(f"Force RMSE : {rmse_f:e} eV/Å") - if weighted_force_metrics is not None: - log.info(f"Force weighted MAE : {mae_fw:e} eV/Å") - log.info(f"Force weighted RMSE: {rmse_fw:e} eV/Å") - dict_to_return.update( - weighted_force_metrics.as_weighted_average_errors( - *DP_TEST_WEIGHTED_FORCE_METRIC_KEYS - ) - ) - if out_put_spin and find_force == 1: - log.info(f"Force atom MAE : {mae_fr:e} eV/Å") - log.info(f"Force atom RMSE : {rmse_fr:e} eV/Å") - dict_to_return.update( - spin_metrics.as_weighted_average_errors( - {"force_real": DP_TEST_SPIN_WEIGHTED_METRIC_KEYS["force_real"]} - ) - ) - if out_put_spin and find_force_mag == 1: - log.info(f"Force spin MAE : {mae_fm:e} eV/uB") - log.info(f"Force spin RMSE : {rmse_fm:e} eV/uB") - dict_to_return.update( - spin_metrics.as_weighted_average_errors( - {"force_magnetic": DP_TEST_SPIN_WEIGHTED_METRIC_KEYS["force_magnetic"]} - ) - ) - if data.pbc and not out_put_spin and find_virial == 1: - log.info(f"Virial MAE : {mae_v:e} eV") - log.info(f"Virial RMSE : {rmse_v:e} eV") - log.info(f"Virial MAE/Natoms : {mae_va:e} eV") - log.info(f"Virial RMSE/Natoms : {rmse_va:e} eV") - log.info(f"Stress MAE : {mae_s:e} eV/Å^3") - log.info(f"Stress RMSE : {rmse_s:e} eV/Å^3") - if has_atom_ener: - log.info(f"Atomic ener MAE : {mae_ae:e} eV") - log.info(f"Atomic ener RMSE : {rmse_ae:e} eV") - if dp.has_hessian: - log.info(f"Hessian MAE : {mae_h:e} eV/Å^2") - log.info(f"Hessian RMSE : {rmse_h:e} eV/Å^2") - if hessian_metrics is None: - raise RuntimeError("Hessian metrics are unavailable for dp test.") - dict_to_return.update( - hessian_metrics.as_weighted_average_errors(*DP_TEST_HESSIAN_METRIC_KEYS) - ) - - if detail_file is not None: - _write_energy_test_details( - detail_path=Path(detail_file), - system=system, - natoms=natoms, - append_detail=append_detail, - reference_energy=test_data["energy"][:numb_test], - prediction_energy=energy, - reference_force=test_data["force"][:numb_test], - prediction_force=force, - reference_virial=test_data["virial"][:numb_test], - prediction_virial=virial, - reference_stress=reference_stress, - prediction_stress=prediction_stress, - out_put_spin=out_put_spin, - reference_force_real=test_force_r, - prediction_force_real=force_r, - reference_force_magnetic=test_force_m if find_force_mag == 1 else None, - prediction_force_magnetic=force_m - if out_put_spin and find_force_mag == 1 - else None, - reference_hessian=test_data["hessian"][:numb_test] - if dp.has_hessian - else None, - prediction_hessian=hessian if dp.has_hessian else None, - ) - - return dict_to_return - - -def print_ener_sys_avg(avg: dict[str, float]) -> None: - """Print errors summary for energy type potential. - - Parameters - ---------- - avg : np.ndarray - array with summaries - """ - log.info(f"Energy MAE : {avg['mae_e']:e} eV") - log.info(f"Energy RMSE : {avg['rmse_e']:e} eV") - log.info(f"Energy MAE/Natoms : {avg['mae_ea']:e} eV") - log.info(f"Energy RMSE/Natoms : {avg['rmse_ea']:e} eV") - if "rmse_f" in avg: - log.info(f"Force MAE : {avg['mae_f']:e} eV/Å") - log.info(f"Force RMSE : {avg['rmse_f']:e} eV/Å") - if "rmse_fw" in avg: - log.info(f"Force weighted MAE : {avg['mae_fw']:e} eV/Å") - log.info(f"Force weighted RMSE: {avg['rmse_fw']:e} eV/Å") - else: - log.info(f"Force atom MAE : {avg['mae_fr']:e} eV/Å") - log.info(f"Force atom RMSE : {avg['rmse_fr']:e} eV/Å") - if "rmse_fm" in avg: - log.info(f"Force spin MAE : {avg['mae_fm']:e} eV/uB") - log.info(f"Force spin RMSE : {avg['rmse_fm']:e} eV/uB") - if "rmse_v" in avg: - log.info(f"Virial MAE : {avg['mae_v']:e} eV") - log.info(f"Virial RMSE : {avg['rmse_v']:e} eV") - log.info(f"Virial MAE/Natoms : {avg['mae_va']:e} eV") - log.info(f"Virial RMSE/Natoms : {avg['rmse_va']:e} eV") - if "rmse_s" in avg: - log.info(f"Stress MAE : {avg['mae_s']:e} eV/Å^3") - log.info(f"Stress RMSE : {avg['rmse_s']:e} eV/Å^3") - if "rmse_h" in avg: - log.info(f"Hessian MAE : {avg['mae_h']:e} eV/Å^2") - log.info(f"Hessian RMSE : {avg['rmse_h']:e} eV/Å^2") - - -def test_dos( - dp: "DeepDOS", - data: DeepmdData, - system: str, - numb_test: int, - detail_file: str | None, - has_atom_dos: bool, - append_detail: bool = False, -) -> tuple[list[np.ndarray], list[int]]: - """Test DOS type model. - - Parameters - ---------- - dp : DeepDOS - instance of deep potential - data : DeepmdData - data container object - system : str - system directory - numb_test : int - munber of tests to do - detail_file : Optional[str] - file where test details will be output - has_atom_dos : bool - whether per atom quantities should be computed - append_detail : bool, optional - if true append output detail file, by default False - - Returns - ------- - tuple[list[np.ndarray], list[int]] - arrays with results and their shapes - """ - data.add("dos", dp.numb_dos, atomic=False, must=True, high_prec=True) - if has_atom_dos: - data.add("atom_dos", dp.numb_dos, atomic=True, must=False, high_prec=True) - - if dp.get_dim_fparam() > 0: - data.add( - "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False - ) - if dp.get_dim_aparam() > 0: - data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False) - - test_data = data.get_test() - mixed_type = data.mixed_type - natoms = len(test_data["type"][0]) - nframes = test_data["box"].shape[0] - numb_test = min(nframes, numb_test) - - coord = test_data["coord"][:numb_test].reshape([numb_test, -1]) - box = test_data["box"][:numb_test] - - if not data.pbc: - box = None - if mixed_type: - atype = test_data["type"][:numb_test].reshape([numb_test, -1]) - else: - atype = test_data["type"][0] - if dp.get_dim_fparam() > 0: - fparam = test_data["fparam"][:numb_test] - else: - fparam = None - if dp.get_dim_aparam() > 0: - aparam = test_data["aparam"][:numb_test] - else: - aparam = None - - ret = dp.eval( - coord, - box, - atype, - fparam=fparam, - aparam=aparam, - atomic=has_atom_dos, - mixed_type=mixed_type, - ) - dos = ret[0] - - dos = dos.reshape([numb_test, dp.numb_dos]) - - if has_atom_dos: - ados = ret[1] - ados = ados.reshape([numb_test, natoms * dp.numb_dos]) - - diff_dos = dos - test_data["dos"][:numb_test] - mae_dos = mae(diff_dos) - rmse_dos = rmse(diff_dos) - - mae_dosa = mae_dos / natoms - rmse_dosa = rmse_dos / natoms - - if has_atom_dos: - diff_ados = ados - test_data["atom_dos"][:numb_test] - mae_ados = mae(diff_ados) - rmse_ados = rmse(diff_ados) - - log.info(f"# number of test data : {numb_test:d} ") - - log.info(f"DOS MAE : {mae_dos:e} Occupation/eV") - log.info(f"DOS RMSE : {rmse_dos:e} Occupation/eV") - log.info(f"DOS MAE/Natoms : {mae_dosa:e} Occupation/eV") - log.info(f"DOS RMSE/Natoms : {rmse_dosa:e} Occupation/eV") - - if has_atom_dos: - log.info(f"Atomic DOS MAE : {mae_ados:e} Occupation/eV") - log.info(f"Atomic DOS RMSE : {rmse_ados:e} Occupation/eV") - - if detail_file is not None: - detail_path = Path(detail_file) - - for ii in range(numb_test): - test_out = test_data["dos"][ii].reshape(-1, 1) - pred_out = dos[ii].reshape(-1, 1) - - frame_output = np.hstack((test_out, pred_out)) - - save_txt_file( - detail_path.with_suffix(f".dos.out.{ii}"), - frame_output, - header=f"{system} - {ii}: data_dos pred_dos", - append=append_detail, - ) - - if has_atom_dos: - for ii in range(numb_test): - test_out = test_data["atom_dos"][ii].reshape(-1, 1) - pred_out = ados[ii].reshape(-1, 1) - - frame_output = np.hstack((test_out, pred_out)) - - save_txt_file( - detail_path.with_suffix(f".ados.out.{ii}"), - frame_output, - header=f"{system} - {ii}: data_ados pred_ados", - append=append_detail, - ) - - return { - "mae_dos": (mae_dos, dos.size), - "mae_dosa": (mae_dosa, dos.size), - "rmse_dos": (rmse_dos, dos.size), - "rmse_dosa": (rmse_dosa, dos.size), - } - - -def print_dos_sys_avg(avg: dict[str, float]) -> None: - """Print errors summary for DOS type potential. - - Parameters - ---------- - avg : np.ndarray - array with summaries - """ - log.info(f"DOS MAE : {avg['mae_dos']:e} Occupation/eV") - log.info(f"DOS RMSE : {avg['rmse_dos']:e} Occupation/eV") - log.info(f"DOS MAE/Natoms : {avg['mae_dosa']:e} Occupation/eV") - log.info(f"DOS RMSE/Natoms : {avg['rmse_dosa']:e} Occupation/eV") - - -def test_property( - dp: "DeepProperty", - data: DeepmdData, - system: str, - numb_test: int, - detail_file: str | None, - has_atom_property: bool, - append_detail: bool = False, -) -> tuple[list[np.ndarray], list[int]]: - """Test Property type model. - - Parameters - ---------- - dp : DeepProperty - instance of deep potential - data : DeepmdData - data container object - system : str - system directory - numb_test : int - munber of tests to do - detail_file : Optional[str] - file where test details will be output - has_atom_property : bool - whether per atom quantities should be computed - append_detail : bool, optional - if true append output detail file, by default False - - Returns - ------- - tuple[list[np.ndarray], list[int]] - arrays with results and their shapes - """ - var_name = dp.get_var_name() - assert isinstance(var_name, str) - data.add(var_name, dp.task_dim, atomic=False, must=True, high_prec=True) - if has_atom_property: - data.add( - f"atom_{var_name}", - dp.task_dim, - atomic=True, - must=False, - high_prec=True, - ) - - if dp.get_dim_fparam() > 0: - data.add( - "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False - ) - if dp.get_dim_aparam() > 0: - data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False) - - test_data = data.get_test() - mixed_type = data.mixed_type - natoms = len(test_data["type"][0]) - nframes = test_data["box"].shape[0] - numb_test = min(nframes, numb_test) - - coord = test_data["coord"][:numb_test].reshape([numb_test, -1]) - box = test_data["box"][:numb_test] - - if not data.pbc: - box = None - if mixed_type: - atype = test_data["type"][:numb_test].reshape([numb_test, -1]) - else: - atype = test_data["type"][0] - if dp.get_dim_fparam() > 0: - fparam = test_data["fparam"][:numb_test] - else: - fparam = None - if dp.get_dim_aparam() > 0: - aparam = test_data["aparam"][:numb_test] - else: - aparam = None - - ret = dp.eval( - coord, - box, - atype, - fparam=fparam, - aparam=aparam, - atomic=has_atom_property, - mixed_type=mixed_type, - ) - - property = ret[0] - - property = property.reshape([numb_test, dp.task_dim]) - - if has_atom_property: - aproperty = ret[1] - aproperty = aproperty.reshape([numb_test, natoms * dp.task_dim]) - - diff_property = property - test_data[var_name][:numb_test] - mae_property = mae(diff_property) - rmse_property = rmse(diff_property) - - if has_atom_property: - diff_aproperty = aproperty - test_data[f"atom_{var_name}"][:numb_test] - mae_aproperty = mae(diff_aproperty) - rmse_aproperty = rmse(diff_aproperty) - - log.info(f"# number of test data : {numb_test:d} ") - - log.info(f"PROPERTY MAE : {mae_property:e} units") - log.info(f"PROPERTY RMSE : {rmse_property:e} units") - - if has_atom_property: - log.info(f"Atomic PROPERTY MAE : {mae_aproperty:e} units") - log.info(f"Atomic PROPERTY RMSE : {rmse_aproperty:e} units") - - if detail_file is not None: - detail_path = Path(detail_file) - - for ii in range(numb_test): - test_out = test_data[var_name][ii].reshape(-1, 1) - pred_out = property[ii].reshape(-1, 1) - - frame_output = np.hstack((test_out, pred_out)) - - save_txt_file( - detail_path.with_suffix(f".property.out.{ii}"), - frame_output, - header=f"{system} - {ii}: data_property pred_property", - append=append_detail, - ) - - if has_atom_property: - for ii in range(numb_test): - test_out = test_data[f"atom_{var_name}"][ii].reshape(-1, 1) - pred_out = aproperty[ii].reshape(-1, 1) - - frame_output = np.hstack((test_out, pred_out)) - - save_txt_file( - detail_path.with_suffix(f".aproperty.out.{ii}"), - frame_output, - header=f"{system} - {ii}: data_aproperty pred_aproperty", - append=append_detail, - ) - - return { - "mae_property": (mae_property, property.size), - "rmse_property": (rmse_property, property.size), - } - - -def print_property_sys_avg(avg: dict[str, float]) -> None: - """Print errors summary for Property type potential. - - Parameters - ---------- - avg : np.ndarray - array with summaries - """ - log.info(f"PROPERTY MAE : {avg['mae_property']:e} units") - log.info(f"PROPERTY RMSE : {avg['rmse_property']:e} units") - - -def run_test( - dp: "DeepTensor", test_data: dict, numb_test: int, test_sys: DeepmdData -) -> dict: - """Run tests. - - Parameters - ---------- - dp : DeepTensor - instance of deep potential - test_data : dict - dictionary with test data - numb_test : int - munber of tests to do - test_sys : DeepmdData - test system - - Returns - ------- - [type] - [description] - """ - nframes = test_data["box"].shape[0] - numb_test = min(nframes, numb_test) - - coord = test_data["coord"][:numb_test].reshape([numb_test, -1]) - if test_sys.pbc: - box = test_data["box"][:numb_test] - else: - box = None - atype = test_data["type"][0] - prediction = dp.eval(coord, box, atype) - - return prediction.reshape([numb_test, -1]), numb_test, atype - - -def test_wfc( - dp: "DeepWFC", - data: DeepmdData, - numb_test: int, - detail_file: str | None, -) -> tuple[list[np.ndarray], list[int]]: - """Test energy type model. - - Parameters - ---------- - dp : DeepPot - instance of deep potential - data : DeepmdData - data container object - numb_test : int - munber of tests to do - detail_file : Optional[str] - file where test details will be output - - Returns - ------- - tuple[list[np.ndarray], list[int]] - arrays with results and their shapes - """ - data.add( - "wfc", 12, atomic=True, must=True, high_prec=False, type_sel=dp.get_sel_type() - ) - test_data = data.get_test() - wfc, numb_test, _ = run_test(dp, test_data, numb_test, data) - rmse_f = rmse(wfc - test_data["wfc"][:numb_test]) - - log.info(f"# number of test data : {numb_test:d} ") - log.info(f"WFC RMSE : {rmse_f:e}") - - if detail_file is not None: - detail_path = Path(detail_file) - pe = np.concatenate( - ( - np.reshape(test_data["wfc"][:numb_test], [-1, 12]), - np.reshape(wfc, [-1, 12]), - ), - axis=1, - ) - np.savetxt( - detail_path.with_suffix(".out"), - pe, - header="ref_wfc(12 dofs) predicted_wfc(12 dofs)", - ) - return {"rmse": (rmse_f, wfc.size)} - - -def print_wfc_sys_avg(avg: dict) -> None: - """Print errors summary for wfc type potential. - - Parameters - ---------- - avg : np.ndarray - array with summaries - """ - log.info(f"WFC RMSE : {avg['rmse']:e}") - - -def test_polar( - dp: "DeepPolar", - data: DeepmdData, - numb_test: int, - detail_file: str | None, - *, - atomic: bool, -) -> tuple[list[np.ndarray], list[int]]: - """Test energy type model. - - Parameters - ---------- - dp : DeepPot - instance of deep potential - data : DeepmdData - data container object - numb_test : int - munber of tests to do - detail_file : Optional[str] - file where test details will be output - atomic : bool - whether to use glovbal version of polar potential - - Returns - ------- - tuple[list[np.ndarray], list[int]] - arrays with results and their shapes - """ - data.add( - "polarizability" if not atomic else "atomic_polarizability", - 9, - atomic=atomic, - must=True, - high_prec=False, - type_sel=dp.get_sel_type(), - output_natoms_for_type_sel=True, - ) - - test_data = data.get_test() - polar, numb_test, atype = run_test(dp, test_data, numb_test, data) - - sel_type = dp.get_sel_type() - sel_natoms = 0 - for ii in sel_type: - sel_natoms += sum(atype == ii) - - # YWolfeee: do summation in global polar mode - if not atomic: - polar = np.sum(polar.reshape((polar.shape[0], -1, 9)), axis=1) - rmse_f = rmse(polar - test_data["polarizability"][:numb_test]) - rmse_fs = rmse_f / np.sqrt(sel_natoms) - rmse_fa = rmse_f / sel_natoms - else: - sel_mask = np.isin(atype, sel_type) - polar = polar.reshape((polar.shape[0], -1, 9))[:, sel_mask, :].reshape( - (polar.shape[0], -1) - ) - label_polar = ( - test_data["atom_polarizability"][:numb_test] - .reshape((numb_test, -1, 9))[:, sel_mask, :] - .reshape((numb_test, -1)) - ) - rmse_f = rmse(polar - label_polar) - - log.info(f"# number of test data : {numb_test:d} ") - log.info(f"Polarizability RMSE : {rmse_f:e}") - if not atomic: - log.info(f"Polarizability RMSE/sqrtN : {rmse_fs:e}") - log.info(f"Polarizability RMSE/N : {rmse_fa:e}") - log.info("The unit of error is the same as the unit of provided label.") - - if detail_file is not None: - detail_path = Path(detail_file) - - if not atomic: - pe = np.concatenate( - ( - np.reshape(test_data["polarizability"][:numb_test], [-1, 9]), - np.reshape(polar, [-1, 9]), - ), - axis=1, - ) - header_text = ( - "data_pxx data_pxy data_pxz data_pyx data_pyy data_pyz data_pzx " - "data_pzy data_pzz pred_pxx pred_pxy pred_pxz pred_pyx pred_pyy " - "pred_pyz pred_pzx pred_pzy pred_pzz" - ) - else: - pe = np.concatenate( - ( - np.reshape(label_polar, [-1, 9 * sel_natoms]), - np.reshape(polar, [-1, 9 * sel_natoms]), - ), - axis=1, - ) - header_text = [ - f"{letter}{number}" - for number in range(1, sel_natoms + 1) - for letter in [ - "data_pxx", - "data_pxy", - "data_pxz", - "data_pyx", - "data_pyy", - "data_pyz", - "data_pzx", - "data_pzy", - "data_pzz", - ] - ] + [ - f"{letter}{number}" - for number in range(1, sel_natoms + 1) - for letter in [ - "pred_pxx", - "pred_pxy", - "pred_pxz", - "pred_pyx", - "pred_pyy", - "pred_pyz", - "pred_pzx", - "pred_pzy", - "pred_pzz", - ] - ] - header_text = " ".join(header_text) - - np.savetxt( - detail_path.with_suffix(".out"), - pe, - header=header_text, - ) - return {"rmse": (rmse_f, polar.size)} - - -def print_polar_sys_avg(avg: dict) -> None: - """Print errors summary for polar type potential. - - Parameters - ---------- - avg : np.ndarray - array with summaries - """ - log.info(f"Polarizability RMSE : {avg['rmse']:e}") - - -def test_dipole( - dp: "DeepDipole", - data: DeepmdData, - numb_test: int, - detail_file: str | None, - atomic: bool, -) -> tuple[list[np.ndarray], list[int]]: - """Test energy type model. - - Parameters - ---------- - dp : DeepPot - instance of deep potential - data : DeepmdData - data container object - numb_test : int - munber of tests to do - detail_file : Optional[str] - file where test details will be output - atomic : bool - whether atomic dipole is provided - - Returns - ------- - tuple[list[np.ndarray], list[int]] - arrays with results and their shapes - """ - data.add( - "dipole" if not atomic else "atomic_dipole", - 3, - atomic=atomic, - must=True, - high_prec=False, - type_sel=dp.get_sel_type(), - output_natoms_for_type_sel=True, - ) - - test_data = data.get_test() - dipole, numb_test, atype = run_test(dp, test_data, numb_test, data) - - sel_type = dp.get_sel_type() - sel_natoms = 0 - for ii in sel_type: - sel_natoms += sum(atype == ii) - - # do summation in atom dimension - if not atomic: - dipole = np.sum(dipole.reshape((dipole.shape[0], -1, 3)), axis=1) - rmse_f = rmse(dipole - test_data["dipole"][:numb_test]) - rmse_fs = rmse_f / np.sqrt(sel_natoms) - rmse_fa = rmse_f / sel_natoms - else: - sel_mask = np.isin(atype, sel_type) - dipole = dipole.reshape((dipole.shape[0], -1, 3))[:, sel_mask, :].reshape( - (dipole.shape[0], -1) - ) - label_dipole = ( - test_data["atom_dipole"][:numb_test] - .reshape((numb_test, -1, 3))[:, sel_mask, :] - .reshape((numb_test, -1)) - ) - rmse_f = rmse(dipole - label_dipole) - - log.info(f"# number of test data : {numb_test:d}") - log.info(f"Dipole RMSE : {rmse_f:e}") - if not atomic: - log.info(f"Dipole RMSE/sqrtN : {rmse_fs:e}") - log.info(f"Dipole RMSE/N : {rmse_fa:e}") - log.info("The unit of error is the same as the unit of provided label.") - - if detail_file is not None: - detail_path = Path(detail_file) - if not atomic: - pe = np.concatenate( - ( - np.reshape(test_data["dipole"][:numb_test], [-1, 3]), - np.reshape(dipole, [-1, 3]), - ), - axis=1, - ) - header_text = "data_x data_y data_z pred_x pred_y pred_z" - else: - pe = np.concatenate( - ( - np.reshape(label_dipole, [-1, 3 * sel_natoms]), - np.reshape(dipole, [-1, 3 * sel_natoms]), - ), - axis=1, - ) - header_text = [ - f"{letter}{number}" - for number in range(1, sel_natoms + 1) - for letter in ["data_x", "data_y", "data_z"] - ] + [ - f"{letter}{number}" - for number in range(1, sel_natoms + 1) - for letter in ["pred_x", "pred_y", "pred_z"] - ] - header_text = " ".join(header_text) - - np.savetxt( - detail_path.with_suffix(".out"), - pe, - header=header_text, - ) - return {"rmse": (rmse_f, dipole.size)} - - -def print_dipole_sys_avg(avg: dict) -> None: - """Print errors summary for dipole type potential. - - Parameters - ---------- - avg : np.ndarray - array with summaries - """ - log.info(f"Dipole RMSE : {avg['rmse']:e}") diff --git a/deepmd/infer/model_test/__init__.py b/deepmd/infer/model_test/__init__.py new file mode 100644 index 0000000000..8f8b7734b9 --- /dev/null +++ b/deepmd/infer/model_test/__init__.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Evaluation of a trained model against labelled data. + +The machinery behind ``dp test``: a tester walks one system in chunks, +evaluates each chunk and combines the errors, so that a dataset larger than +memory can be tested. :func:`build_tester` selects the tester of a model +class; the command-line entry point only discovers the systems and reports +the run-level average. +""" + +import logging +from typing import ( + Any, +) + +from deepmd.infer.deep_dipole import ( + DeepDipole, +) +from deepmd.infer.deep_dos import ( + DeepDOS, +) +from deepmd.infer.deep_polar import ( + DeepGlobalPolar, + DeepPolar, +) +from deepmd.infer.deep_pot import ( + DeepPot, +) +from deepmd.infer.deep_property import ( + DeepProperty, +) +from deepmd.infer.model_test.base import ( + ChunkContext, + ModelTester, + save_txt_file, + test_chunk_atoms, +) +from deepmd.infer.model_test.dos import ( + DosTester, +) +from deepmd.infer.model_test.ener import ( + EnerTester, + SpinEnerTester, +) +from deepmd.infer.model_test.property import ( + PropertyTester, +) +from deepmd.infer.model_test.tensor import ( + DipoleTester, + PolarTester, + TensorTester, +) + +log = logging.getLogger(__name__) + +__all__ = [ + "ChunkContext", + "DipoleTester", + "DosTester", + "EnerTester", + "ModelTester", + "PolarTester", + "PropertyTester", + "SpinEnerTester", + "TensorTester", + "build_tester", + "save_txt_file", + "test_chunk_atoms", +] + + +def build_tester(dp: Any, *, atomic: bool) -> ModelTester: + """Return the tester of the model class under test. + + Parameters + ---------- + dp : Any + The evaluator of the model under test. + atomic : bool + Whether per-atom quantities are computed. + + Returns + ------- + ModelTester + A tester able to evaluate one system of that model class. + + Raises + ------ + RuntimeError + If no tester covers the model class. + """ + if isinstance(dp, DeepPot): + tester = SpinEnerTester if dp.has_spin else EnerTester + return tester(dp, atomic=atomic) + if isinstance(dp, DeepDOS): + return DosTester(dp, atomic=atomic) + if isinstance(dp, DeepProperty): + return PropertyTester(dp, atomic=atomic) + if isinstance(dp, DeepGlobalPolar): + # A global polar model reports one tensor per frame, which is what a + # polar model does when per-atom output is not requested. + log.warning( + "Global polar model is not currently supported. Please directly " + "use the polar mode and change loss parameters." + ) + return PolarTester(dp, atomic=False) + if isinstance(dp, DeepPolar): + return PolarTester(dp, atomic=atomic) + if isinstance(dp, DeepDipole): + return DipoleTester(dp, atomic=atomic) + raise RuntimeError(f"Testing is not supported for {type(dp).__name__}.") diff --git a/deepmd/infer/model_test/base.py b/deepmd/infer/model_test/base.py new file mode 100644 index 0000000000..de83b60608 --- /dev/null +++ b/deepmd/infer/model_test/base.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Skeleton shared by every model class a ``dp test`` run can evaluate.""" + +import logging +import os +from abc import ( + ABC, + abstractmethod, +) +from collections.abc import ( + Mapping, +) +from dataclasses import ( + dataclass, +) +from pathlib import ( + Path, +) +from typing import ( + Any, + ClassVar, +) + +import numpy as np + +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.weight_avg import ( + merge_weighted_errors, +) + +log = logging.getLogger(__name__) + +__all__ = [ + "ChunkContext", + "ModelTester", + "save_txt_file", + "test_chunk_atoms", +] + +DEFAULT_TEST_CHUNK_ATOMS = 1_000_000 + + +def test_chunk_atoms() -> int: + """Return the number of atoms a test evaluates at once. + + Testing walks a system in chunks so that neither the reference data nor + the predictions of a large system have to be held in full. The bound is + expressed in atoms, as the evaluation batch size is, so that it means the + same amount of work whatever the size of a frame. ``DP_TEST_CHUNK_ATOMS`` + overrides it. + + Returns + ------- + int + The maximum number of atoms per chunk, at least one. + """ + return max(1, int(os.environ.get("DP_TEST_CHUNK_ATOMS", DEFAULT_TEST_CHUNK_ATOMS))) + + +log = logging.getLogger(__name__) + + +def save_txt_file( + fname: Path, data: np.ndarray, header: str = "", append: bool = False +) -> None: + """Save numpy array to test file. + + Parameters + ---------- + fname : str + filename + data : np.ndarray + data to save to disk + header : str, optional + header string to use in file, by default "" + append : bool, optional + if true file will be appended instead of overwriting, by default False + """ + flags = "a" if append else "w" + with fname.open(flags, encoding="utf-8") as fp: + np.savetxt(fp, data, header=header) + + +@dataclass(frozen=True) +class ChunkContext: + """Where one chunk sits within the system being tested. + + Attributes + ---------- + system : str + System label recorded in the detail files. + detail_file : str or None + File the per-frame details are written to, or ``None`` to write none. + append_detail : bool + Whether the details of this chunk extend an existing file rather than + starting one. + frame_offset : int + Index of the first frame of the chunk within its system, which keeps + per-frame detail files numbered consistently across chunks. + """ + + system: str + detail_file: str | None + append_detail: bool + frame_offset: int + + @property + def detail_path(self) -> Path | None: + """The detail file as a path, or ``None`` when details are not kept.""" + return None if self.detail_file is None else Path(self.detail_file) + + +class ModelTester(ABC): + """Evaluate one system of one model class, one chunk at a time. + + A system is walked in chunks so that neither its reference data nor the + predictions over it are ever held in full, which is what makes a dataset + larger than memory testable. The errors of the chunks combine into the + errors of the system exactly, because an MAE and an RMSE are both + recovered from partial results weighted by the number of elements each was + taken over; see :func:`~deepmd.utils.weight_avg.merge_weighted_errors`. + + A subclass supplies only what distinguishes its model class: the labels a + chunk must carry, how a chunk is evaluated, and how the resulting + quantities are named and reported. + + Parameters + ---------- + dp : Any + The evaluator of the model under test. + atomic : bool + Whether per-atom quantities are computed. + """ + + #: ``(quantity, log template)`` pairs, in the order they are reported. A + #: quantity the run did not produce is absent from the errors and is + #: therefore skipped. + report: ClassVar[tuple[tuple[str, str], ...]] = () + #: Line closing the report, if the model class needs one. + report_footer: ClassVar[str | None] = None + #: Quantities reported per system but withheld from the run-level average. + per_system_only: ClassVar[tuple[str, ...]] = () + + def __init__(self, dp: Any, *, atomic: bool) -> None: + self.dp = dp + self.atomic = atomic + + @abstractmethod + def add_data_requirements(self, data: DeepmdData) -> None: + """Declare the labels every chunk of the system must carry. + + Parameters + ---------- + data : DeepmdData + The system about to be tested. + """ + + @abstractmethod + def evaluate_chunk( + self, + data: DeepmdData, + test_data: dict, + context: ChunkContext, + ) -> dict[str, tuple[float, float]]: + """Evaluate one chunk and report the errors over it. + + Parameters + ---------- + data : DeepmdData + The system the chunk was drawn from, consulted for its conventions. + test_data : dict + One chunk of test data, as yielded by ``data.iter_test``. + context : ChunkContext + Where the chunk sits within the system. + + Returns + ------- + dict[str, tuple[float, float]] + The ``(error, weight)`` of every quantity the chunk produced. + """ + + def run( + self, + data: DeepmdData, + system: str, + numb_test: float, + detail_file: str | None, + *, + append_detail: bool = False, + ) -> dict[str, tuple[float, float]]: + """Test one system and report its errors. + + Parameters + ---------- + data : DeepmdData + The system to test. + system : str + System label used in logs and detail files. + numb_test : float + Upper bound on the number of frames tested. A non-finite bound + tests every frame. + detail_file : str, optional + File the per-frame details are written to. + append_detail : bool, optional + Whether the details of this system extend an existing file. + + Returns + ------- + dict[str, tuple[float, float]] + The ``(error, weight)`` of every quantity that takes part in the + run-level average. + + Raises + ------ + RuntimeError + If the system holds no test frame. + """ + self.add_data_requirements(data) + + chunk_errors: list[dict[str, tuple[float, float]]] = [] + frames_tested = 0 + append = append_detail + for chunk in data.iter_test( + chunk_atoms=test_chunk_atoms(), numb_test=numb_test + ): + context = ChunkContext( + system=system, + detail_file=detail_file, + append_detail=append, + frame_offset=frames_tested, + ) + chunk_errors.append(self.evaluate_chunk(data, chunk, context)) + frames_tested += chunk["box"].shape[0] + append = True + + if not chunk_errors: + raise RuntimeError(f"No test frames found in system {system}.") + + errors = merge_weighted_errors(chunk_errors) + log.info(f"# number of test data : {frames_tested:d} ") + self.log_errors(errors) + return { + key: value + for key, value in errors.items() + if key not in self.per_system_only + } + + @classmethod + def log_errors(cls, errors: Mapping[str, tuple[float, float]]) -> None: + """Report the errors of a system or of a whole run. + + Parameters + ---------- + errors : Mapping[str, tuple[float, float]] + The ``(error, weight)`` of each quantity. + """ + for key, template in cls.report: + if key in errors: + log.info(template.format(f"{errors[key][0]:e}")) + if cls.report_footer is not None: + log.info(cls.report_footer) + + +# --------------------------------------------------------------------------- +# Energy models +# --------------------------------------------------------------------------- + + +def _write_per_frame_details( + context: ChunkContext, + *, + suffix: str, + reference: np.ndarray, + prediction: np.ndarray, +) -> None: + """Write one detail file per frame of a chunk. + + Parameters + ---------- + context : ChunkContext + Where the chunk sits within the system, which numbers the files. + suffix : str + Name of the quantity, used in the file suffix and the header. + reference : np.ndarray + Reference values with shape ``(nframes, ...)``. + prediction : np.ndarray + Predicted values with shape ``(nframes, ...)``. + """ + detail_path = context.detail_path + assert detail_path is not None + for index in range(reference.shape[0]): + frame = context.frame_offset + index + save_txt_file( + detail_path.with_suffix(f".{suffix}.out.{frame}"), + np.hstack( + ( + reference[index].reshape(-1, 1), + prediction[index].reshape(-1, 1), + ) + ), + header=f"{context.system} - {frame}: data_{suffix} pred_{suffix}", + append=context.append_detail, + ) + + +# --------------------------------------------------------------------------- +# Tensor models +# --------------------------------------------------------------------------- diff --git a/deepmd/infer/model_test/dos.py b/deepmd/infer/model_test/dos.py new file mode 100644 index 0000000000..239c9fe858 --- /dev/null +++ b/deepmd/infer/model_test/dos.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of density-of-states models.""" + +import logging + +from deepmd.infer.model_test.base import ( + ChunkContext, + ModelTester, + _write_per_frame_details, +) +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.eval_metrics import ( + mae, + rmse, +) + +log = logging.getLogger(__name__) + +__all__ = ["DosTester"] + + +class DosTester(ModelTester): + """Test a model of the electronic density of states.""" + + report = ( + ("mae_dos", "DOS MAE : {} Occupation/eV"), + ("rmse_dos", "DOS RMSE : {} Occupation/eV"), + ("mae_dosa", "DOS MAE/Natoms : {} Occupation/eV"), + ("rmse_dosa", "DOS RMSE/Natoms : {} Occupation/eV"), + ("mae_ados", "Atomic DOS MAE : {} Occupation/eV"), + ("rmse_ados", "Atomic DOS RMSE : {} Occupation/eV"), + ) + per_system_only = ("mae_ados", "rmse_ados") + + def add_data_requirements(self, data: DeepmdData) -> None: + """Declare the labels a density-of-states test consumes.""" + dp = self.dp + data.add("dos", dp.numb_dos, atomic=False, must=True, high_prec=True) + if self.atomic: + data.add("atom_dos", dp.numb_dos, atomic=True, must=True, high_prec=True) + if dp.get_dim_fparam() > 0: + data.add( + "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False + ) + if dp.get_dim_aparam() > 0: + data.add( + "aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False + ) + + def evaluate_chunk( + self, + data: DeepmdData, + test_data: dict, + context: ChunkContext, + ) -> dict[str, tuple[float, float]]: + """Evaluate one chunk of a density-of-states test.""" + dp = self.dp + mixed_type = data.mixed_type + natoms = len(test_data["type"][0]) + nframes = test_data["box"].shape[0] + + coord = test_data["coord"].reshape([nframes, -1]) + box = test_data["box"] if data.pbc else None + if mixed_type: + atype = test_data["type"].reshape([nframes, -1]) + else: + atype = test_data["type"][0] + fparam = test_data["fparam"] if dp.get_dim_fparam() > 0 else None + aparam = test_data["aparam"] if dp.get_dim_aparam() > 0 else None + + ret = dp.eval( + coord, + box, + atype, + fparam=fparam, + aparam=aparam, + atomic=self.atomic, + mixed_type=mixed_type, + ) + dos = ret[0].reshape([nframes, dp.numb_dos]) + + diff_dos = dos - test_data["dos"] + mae_dos = mae(diff_dos) + rmse_dos = rmse(diff_dos) + errors: dict[str, tuple[float, float]] = { + "mae_dos": (mae_dos, dos.size), + "mae_dosa": (mae_dos / natoms, dos.size), + "rmse_dos": (rmse_dos, dos.size), + "rmse_dosa": (rmse_dos / natoms, dos.size), + } + + ados = None + if self.atomic: + ados = ret[1].reshape([nframes, natoms * dp.numb_dos]) + diff_ados = ados - test_data["atom_dos"] + errors["mae_ados"] = (mae(diff_ados), ados.size) + errors["rmse_ados"] = (rmse(diff_ados), ados.size) + + if context.detail_path is not None: + _write_per_frame_details( + context, + suffix="dos", + reference=test_data["dos"], + prediction=dos, + ) + if self.atomic: + _write_per_frame_details( + context, + suffix="ados", + reference=test_data["atom_dos"], + prediction=ados, + ) + + return errors + + +# --------------------------------------------------------------------------- +# Property models +# --------------------------------------------------------------------------- diff --git a/deepmd/infer/model_test/ener.py b/deepmd/infer/model_test/ener.py new file mode 100644 index 0000000000..f51139d056 --- /dev/null +++ b/deepmd/infer/model_test/ener.py @@ -0,0 +1,703 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of energy models, including those carrying spin.""" + +import logging +from dataclasses import ( + dataclass, +) +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, + ClassVar, + NamedTuple, +) + +import numpy as np + +from deepmd.infer.model_test.base import ( + ChunkContext, + ModelTester, + save_txt_file, +) +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.eval_metrics import ( + DP_TEST_HESSIAN_METRIC_KEYS, + DP_TEST_SPIN_WEIGHTED_METRIC_KEYS, + DP_TEST_WEIGHTED_FORCE_METRIC_KEYS, + DP_TEST_WEIGHTED_METRIC_KEYS, + compute_energy_type_metrics, + compute_error_stat, + compute_spin_force_metrics, + compute_weighted_error_stat, +) + +if TYPE_CHECKING: + from deepmd.infer.deep_pot import ( + DeepPot, + ) + +log = logging.getLogger(__name__) + +__all__ = ["EnerTester", "SpinEnerTester"] + + +def _reshape_force_by_atom(force_array: np.ndarray, natoms: int) -> np.ndarray: + """Reshape flattened force arrays into `[nframes, natoms, 3]`.""" + return np.reshape(force_array, [-1, natoms, 3]) + + +def _concat_force_rows( + force_blocks: list[np.ndarray], dtype: np.dtype | type[np.generic] +) -> np.ndarray: + """Concatenate per-frame force rows into one 2D array.""" + if not force_blocks: + return np.empty((0, 3), dtype=dtype) + return np.concatenate(force_blocks, axis=0) + + +def _align_spin_force_arrays( + *, + dp: "DeepPot", + atype: np.ndarray, + natoms: int, + prediction_force: np.ndarray, + reference_force: np.ndarray, + prediction_force_mag: np.ndarray | None, + reference_force_mag: np.ndarray | None, + mask_mag: np.ndarray | None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray | None, np.ndarray | None]: + """Align spin force arrays into real-atom and magnetic subsets.""" + prediction_force_by_atom = _reshape_force_by_atom(prediction_force, natoms) + reference_force_by_atom = _reshape_force_by_atom(reference_force, natoms) + if dp.get_ntypes_spin() != 0: # old tf support for spin + ntypes_real = dp.get_ntypes() - dp.get_ntypes_spin() + atype_by_frame = np.reshape(atype, [-1, natoms]) + if atype_by_frame.shape[0] == 1 and prediction_force_by_atom.shape[0] != 1: + atype_by_frame = np.broadcast_to( + atype_by_frame, + (prediction_force_by_atom.shape[0], natoms), + ) + if atype_by_frame.shape[0] != prediction_force_by_atom.shape[0]: + raise ValueError( + "Spin atom types and force arrays must have matching frames." + ) + force_real_prediction_chunks = [] + force_real_reference_chunks = [] + force_magnetic_prediction_chunks = [] + force_magnetic_reference_chunks = [] + for frame_atype, frame_prediction, frame_reference in zip( + atype_by_frame, + prediction_force_by_atom, + reference_force_by_atom, + strict=False, + ): + real_mask = frame_atype < ntypes_real + magnetic_mask = ~real_mask + force_real_prediction_chunks.append(frame_prediction[real_mask]) + force_real_reference_chunks.append(frame_reference[real_mask]) + force_magnetic_prediction_chunks.append(frame_prediction[magnetic_mask]) + force_magnetic_reference_chunks.append(frame_reference[magnetic_mask]) + return ( + _concat_force_rows( + force_real_prediction_chunks, + prediction_force_by_atom.dtype, + ), + _concat_force_rows( + force_real_reference_chunks, + reference_force_by_atom.dtype, + ), + _concat_force_rows( + force_magnetic_prediction_chunks, + prediction_force_by_atom.dtype, + ), + _concat_force_rows( + force_magnetic_reference_chunks, + reference_force_by_atom.dtype, + ), + ) + + force_real_prediction = prediction_force_by_atom.reshape(-1, 3) + force_real_reference = reference_force_by_atom.reshape(-1, 3) + if prediction_force_mag is None or reference_force_mag is None or mask_mag is None: + return force_real_prediction, force_real_reference, None, None + magnetic_mask = mask_mag.reshape(-1).astype(bool) + return ( + force_real_prediction, + force_real_reference, + prediction_force_mag.reshape(-1, 3)[magnetic_mask], + reference_force_mag.reshape(-1, 3)[magnetic_mask], + ) + + +def _write_energy_test_details( + *, + detail_path: Path, + system: str, + natoms: int, + append_detail: bool, + reference_energy: np.ndarray, + prediction_energy: np.ndarray, + reference_force: np.ndarray, + prediction_force: np.ndarray, + reference_virial: np.ndarray | None, + prediction_virial: np.ndarray | None, + out_put_spin: bool, + reference_stress: np.ndarray | None = None, + prediction_stress: np.ndarray | None = None, + reference_force_real: np.ndarray | None = None, + prediction_force_real: np.ndarray | None = None, + reference_force_magnetic: np.ndarray | None = None, + prediction_force_magnetic: np.ndarray | None = None, + reference_hessian: np.ndarray | None = None, + prediction_hessian: np.ndarray | None = None, +) -> None: + """Write energy-type detail outputs after arrays have been aligned.""" + pe = np.concatenate( + ( + np.reshape(reference_energy, [-1, 1]), + np.reshape(prediction_energy, [-1, 1]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".e.out"), + pe, + header=f"{system}: data_e pred_e", + append=append_detail, + ) + pe_atom = pe / natoms + save_txt_file( + detail_path.with_suffix(".e_peratom.out"), + pe_atom, + header=f"{system}: data_e pred_e", + append=append_detail, + ) + if not out_put_spin: + pf = np.concatenate( + ( + np.reshape(reference_force, [-1, 3]), + np.reshape(prediction_force, [-1, 3]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".f.out"), + pf, + header=f"{system}: data_fx data_fy data_fz pred_fx pred_fy pred_fz", + append=append_detail, + ) + else: + if reference_force_real is None or prediction_force_real is None: + raise ValueError("Spin detail output requires aligned real-atom forces.") + pf_real = np.concatenate( + ( + np.reshape(reference_force_real, [-1, 3]), + np.reshape(prediction_force_real, [-1, 3]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".fr.out"), + pf_real, + header=f"{system}: data_fx data_fy data_fz pred_fx pred_fy pred_fz", + append=append_detail, + ) + if (reference_force_magnetic is None) != (prediction_force_magnetic is None): + raise ValueError( + "Spin magnetic detail output requires both reference and prediction forces." + ) + if ( + reference_force_magnetic is not None + and prediction_force_magnetic is not None + ): + pf_mag = np.concatenate( + ( + np.reshape(reference_force_magnetic, [-1, 3]), + np.reshape(prediction_force_magnetic, [-1, 3]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".fm.out"), + pf_mag, + header=f"{system}: data_fmx data_fmy data_fmz pred_fmx pred_fmy pred_fmz", + append=append_detail, + ) + if (reference_virial is None) != (prediction_virial is None): + raise ValueError( + "Virial detail output requires both reference and prediction virials." + ) + if reference_virial is not None and prediction_virial is not None: + pv = np.concatenate( + ( + np.reshape(reference_virial, [-1, 9]), + np.reshape(prediction_virial, [-1, 9]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".v.out"), + pv, + header=f"{system}: data_vxx data_vxy data_vxz data_vyx data_vyy " + "data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx " + "pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz", + append=append_detail, + ) + pv_atom = pv / natoms + save_txt_file( + detail_path.with_suffix(".v_peratom.out"), + pv_atom, + header=f"{system}: data_vxx data_vxy data_vxz data_vyx data_vyy " + "data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx " + "pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz", + append=append_detail, + ) + if (reference_stress is None) != (prediction_stress is None): + raise ValueError( + "Stress detail output requires both reference and prediction stresses." + ) + if reference_stress is not None and prediction_stress is not None: + ps = np.concatenate( + ( + np.reshape(reference_stress, [-1, 9]), + np.reshape(prediction_stress, [-1, 9]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".s.out"), + ps, + header=f"{system} (eV/Å^3): data_sxx data_sxy data_sxz data_syx " + "data_syy data_syz data_szx data_szy data_szz pred_sxx pred_sxy pred_sxz " + "pred_syx pred_syy pred_syz pred_szx pred_szy pred_szz", + append=append_detail, + ) + if reference_hessian is not None and prediction_hessian is not None: + hessian_detail = np.concatenate( + ( + reference_hessian.reshape(-1, 1), + prediction_hessian.reshape(-1, 1), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".h.out"), + hessian_detail, + header=f"{system}: data_h pred_h (3Na*3Na matrix in row-major order)", + append=append_detail, + ) + + +class _OptionalEnerOutputs(NamedTuple): + """The optional trailing outputs of ``DeepPot.eval`` (``None`` when absent).""" + + atom_energy: "np.ndarray | None" + atom_virial: "np.ndarray | None" + force_mag: "np.ndarray | None" + mask_mag: "np.ndarray | None" + hessian: "np.ndarray | None" + + +def _split_optional_ener_outputs( + ret: tuple, + *, + has_atom_ener: bool, + has_spin: bool, + has_hessian: bool, + numb_test: int, +) -> _OptionalEnerOutputs: + """Split the optional trailing outputs of ``DeepPot.eval``. + + ``DeepPot.eval`` appends its optional outputs after ``(energy, force, + virial)`` in a fixed order: atomic ``(atom_energy, atom_virial)``, then spin + ``(force_mag, mask_mag)``, then ``hessian``. Read them by advancing an index + through the tuple in that same order, so the hessian slot is not confused + with atomic energy/virial or spin outputs when those are also present. + """ + atom_energy = atom_virial = force_mag = mask_mag = hessian = None + idx = 3 + if has_atom_ener: + atom_energy = ret[idx].reshape([numb_test, -1]) + atom_virial = ret[idx + 1].reshape([numb_test, -1]) + idx += 2 + if has_spin: + force_mag = ret[idx].reshape([numb_test, -1]) + mask_mag = ret[idx + 1].reshape([numb_test, -1]) + idx += 2 + if has_hessian: + hessian = ret[idx].reshape([numb_test, -1]) + return _OptionalEnerOutputs(atom_energy, atom_virial, force_mag, mask_mag, hessian) + + +@dataclass(frozen=True) +class _ForceDetails: + """Force arrays a spin detail file records, alongside the shared ones.""" + + reference_real: np.ndarray | None = None + prediction_real: np.ndarray | None = None + reference_magnetic: np.ndarray | None = None + prediction_magnetic: np.ndarray | None = None + + +# --------------------------------------------------------------------------- +# Density of states models +# --------------------------------------------------------------------------- + + +class EnerTester(ModelTester): + """Test an energy model against energies, forces and the virial.""" + + report = ( + ("mae_e", "Energy MAE : {} eV"), + ("rmse_e", "Energy RMSE : {} eV"), + ("mae_ea", "Energy MAE/Natoms : {} eV"), + ("rmse_ea", "Energy RMSE/Natoms : {} eV"), + ("mae_f", "Force MAE : {} eV/Å"), + ("rmse_f", "Force RMSE : {} eV/Å"), + ("mae_fw", "Force weighted MAE : {} eV/Å"), + ("rmse_fw", "Force weighted RMSE: {} eV/Å"), + ("mae_fr", "Force atom MAE : {} eV/Å"), + ("rmse_fr", "Force atom RMSE : {} eV/Å"), + ("mae_fm", "Force spin MAE : {} eV/uB"), + ("rmse_fm", "Force spin RMSE : {} eV/uB"), + ("mae_v", "Virial MAE : {} eV"), + ("rmse_v", "Virial RMSE : {} eV"), + ("mae_va", "Virial MAE/Natoms : {} eV"), + ("rmse_va", "Virial RMSE/Natoms : {} eV"), + ("mae_s", "Stress MAE : {} eV/Å^3"), + ("rmse_s", "Stress RMSE : {} eV/Å^3"), + ("mae_ae", "Atomic ener MAE : {} eV"), + ("rmse_ae", "Atomic ener RMSE : {} eV"), + ("mae_h", "Hessian MAE : {} eV/Å^2"), + ("rmse_h", "Hessian RMSE : {} eV/Å^2"), + ) + per_system_only = ("mae_ae", "rmse_ae") + + #: Whether the force this model class reports is the plain atomic force. + #: A spin model reports a real and a magnetic force instead. + reports_plain_force: ClassVar[bool] = True + + def add_data_requirements(self, data: DeepmdData) -> None: + """Declare the labels an energy test consumes.""" + dp = self.dp + data.add("energy", 1, atomic=False, must=False, high_prec=True) + data.add("force", 3, atomic=True, must=False, high_prec=False) + data.add("atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3) + data.add("virial", 9, atomic=False, must=False, high_prec=False) + if dp.has_efield: + data.add("efield", 3, atomic=True, must=True, high_prec=False) + if self.atomic: + data.add("atom_ener", 1, atomic=True, must=True, high_prec=False) + if dp.get_dim_fparam() > 0: + data.add( + "fparam", + dp.get_dim_fparam(), + atomic=False, + must=not dp.has_default_fparam(), + high_prec=False, + ) + if dp.get_dim_aparam() > 0: + data.add( + "aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False + ) + if dp.has_chg_spin_ebd(): + data.add( + "charge_spin", + 2, + atomic=False, + must=not dp.has_default_chg_spin(), + high_prec=False, + ) + if dp.has_spin: + data.add("spin", 3, atomic=True, must=True, high_prec=False) + data.add("force_mag", 3, atomic=True, must=False, high_prec=False) + if dp.has_hessian: + data.add("hessian", 1, atomic=True, must=True, high_prec=False) + + def evaluate_chunk( + self, + data: DeepmdData, + test_data: dict, + context: ChunkContext, + ) -> dict[str, tuple[float, float]]: + """Evaluate one chunk of an energy test.""" + dp = self.dp + errors: dict[str, tuple[float, float]] = {} + find_energy = test_data.get("find_energy") + find_force = test_data.get("find_force") + find_virial = test_data.get("find_virial") + find_atom_pref = test_data.get("find_atom_pref") + mixed_type = data.mixed_type + natoms = len(test_data["type"][0]) + nframes = test_data["box"].shape[0] + + coord = test_data["coord"].reshape([nframes, -1]) + box = test_data["box"] if data.pbc else None + efield = test_data["efield"].reshape([nframes, -1]) if dp.has_efield else None + spin = test_data["spin"].reshape([nframes, -1]) if dp.has_spin else None + if mixed_type: + atype = test_data["type"].reshape([nframes, -1]) + else: + atype = test_data["type"][0] + fparam = ( + test_data["fparam"] + if dp.get_dim_fparam() > 0 and test_data["find_fparam"] != 0.0 + else None + ) + aparam = test_data["aparam"] if dp.get_dim_aparam() > 0 else None + charge_spin = ( + test_data["charge_spin"] + if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 0.0) != 0.0 + else None + ) + + ret = dp.eval( + coord, + box, + atype, + fparam=fparam, + aparam=aparam, + atomic=self.atomic, + efield=efield, + mixed_type=mixed_type, + spin=spin, + charge_spin=charge_spin, + ) + energy = ret[0].reshape([nframes, 1]) + force = ret[1].reshape([nframes, -1]) + virial = ret[2].reshape([nframes, 9]) + optional_outputs = _split_optional_ener_outputs( + ret, + has_atom_ener=self.atomic, + has_spin=dp.has_spin, + has_hessian=dp.has_hessian, + numb_test=nframes, + ) + + force_details = self.force_errors( + errors, + data=data, + test_data=test_data, + atype=atype, + natoms=natoms, + prediction_force=force, + optional_outputs=optional_outputs, + find_force=find_force, + find_atom_pref=find_atom_pref, + ) + + reports_virial = find_virial == 1 and data.pbc + shared_metrics = compute_energy_type_metrics( + prediction={ + "energy": energy, + "force": force, + **({"virial": virial} if reports_virial else {}), + }, + test_data={ + "find_energy": find_energy, + "find_force": find_force if self.reports_plain_force else 0.0, + "find_virial": find_virial, + "energy": test_data["energy"], + "force": test_data["force"], + **({"virial": test_data["virial"]} if reports_virial else {}), + }, + natoms=natoms, + has_pbc=data.pbc, + ) + errors.update( + shared_metrics.as_weighted_average_errors(DP_TEST_WEIGHTED_METRIC_KEYS) + ) + if find_energy == 1 and ( + shared_metrics.energy is None or shared_metrics.energy_per_atom is None + ): + raise RuntimeError("Energy metrics are unavailable for dp test.") + + prediction_stress = None + reference_stress = None + if reports_virial: + if shared_metrics.virial is None or shared_metrics.virial_per_atom is None: + raise RuntimeError("Virial metrics are unavailable for dp test.") + # Stress sigma = -virial / volume, in eV/Å^3 (tensile-positive + # convention). + volume = np.abs(np.linalg.det(box.reshape([nframes, 3, 3]))).reshape( + [nframes, 1] + ) + prediction_stress = -virial / volume + reference_stress = -test_data["virial"] / volume + errors.update( + compute_error_stat( + prediction_stress, reference_stress + ).as_weighted_average_errors("mae_s", "rmse_s") + ) + + if dp.has_hessian: + errors.update( + compute_error_stat( + optional_outputs.hessian, test_data["hessian"] + ).as_weighted_average_errors(*DP_TEST_HESSIAN_METRIC_KEYS) + ) + if self.atomic: + errors.update( + compute_error_stat( + optional_outputs.atom_energy.reshape([-1]), + test_data["atom_ener"].reshape([-1]), + ).as_weighted_average_errors(*self.per_system_only) + ) + + if context.detail_path is not None: + _write_energy_test_details( + detail_path=context.detail_path, + system=context.system, + natoms=natoms, + append_detail=context.append_detail, + reference_energy=test_data["energy"], + prediction_energy=energy, + reference_force=test_data["force"], + prediction_force=force, + reference_virial=test_data["virial"], + prediction_virial=virial, + reference_stress=reference_stress, + prediction_stress=prediction_stress, + out_put_spin=not self.reports_plain_force, + reference_force_real=force_details.reference_real, + prediction_force_real=force_details.prediction_real, + reference_force_magnetic=force_details.reference_magnetic, + prediction_force_magnetic=force_details.prediction_magnetic, + reference_hessian=test_data["hessian"] if dp.has_hessian else None, + prediction_hessian=optional_outputs.hessian if dp.has_hessian else None, + ) + + return errors + + def force_errors( + self, + errors: dict[str, tuple[float, float]], + *, + data: DeepmdData, + test_data: dict, + atype: np.ndarray, + natoms: int, + prediction_force: np.ndarray, + optional_outputs: "_OptionalEnerOutputs", + find_force: float | None, + find_atom_pref: float | None, + ) -> "_ForceDetails": + """Add the force errors of a chunk and return what the details need. + + The plain atomic force is covered by the shared energy metrics, so only + the optional per-atom weighting is added here. + + Parameters + ---------- + errors : dict[str, tuple[float, float]] + Errors of the chunk, extended in place. + data : DeepmdData + The system the chunk was drawn from. + test_data : dict + The chunk. + atype : np.ndarray + Atom types of the chunk. + natoms : int + Number of atoms per frame. + prediction_force : np.ndarray + Predicted force with shape ``(nframes, natoms * 3)``. + optional_outputs : _OptionalEnerOutputs + The optional trailing outputs of the evaluation. + find_force : float or None + Whether the chunk carries a force label. + find_atom_pref : float or None + Whether the chunk carries per-atom force weights. + + Returns + ------- + _ForceDetails + The force arrays the detail file records, empty for a model whose + forces the shared detail writer already covers. + """ + if find_force == 1 and find_atom_pref == 1: + errors.update( + compute_weighted_error_stat( + prediction_force, + test_data["force"], + test_data["atom_pref"], + ).as_weighted_average_errors(*DP_TEST_WEIGHTED_FORCE_METRIC_KEYS) + ) + return _ForceDetails() + + +class SpinEnerTester(EnerTester): + """Test a spin energy model, whose force splits into real and magnetic. + + Everything else an energy model reports carries over unchanged: the + magnetic degrees of freedom enter the virial only through the virtual + atoms, whose displacement the model removes again, so the virial is with + respect to the real atomic positions as for any other energy model. + """ + + reports_plain_force = False + + def force_errors( + self, + errors: dict[str, tuple[float, float]], + *, + data: DeepmdData, + test_data: dict, + atype: np.ndarray, + natoms: int, + prediction_force: np.ndarray, + optional_outputs: "_OptionalEnerOutputs", + find_force: float | None, + find_atom_pref: float | None, + ) -> "_ForceDetails": + """Add the real and magnetic force errors of a chunk.""" + find_force_mag = test_data.get("find_force_mag") + force_real, reference_real, force_mag, reference_mag = _align_spin_force_arrays( + dp=self.dp, + atype=atype, + natoms=natoms, + prediction_force=prediction_force, + reference_force=test_data["force"], + prediction_force_mag=optional_outputs.force_mag, + reference_force_mag=test_data.get("force_mag"), + mask_mag=optional_outputs.mask_mag, + ) + if find_force_mag == 1 and (force_mag is None or reference_mag is None): + raise RuntimeError( + "Spin magnetic force metrics require magnetic force arrays and mask." + ) + spin_metrics = compute_spin_force_metrics( + force_real_prediction=force_real, + force_real_reference=reference_real, + force_magnetic_prediction=force_mag if find_force_mag == 1 else None, + force_magnetic_reference=reference_mag if find_force_mag == 1 else None, + ) + if spin_metrics.force_real is None: + raise RuntimeError("Spin force metrics are unavailable for dp test.") + if find_force == 1: + errors.update( + spin_metrics.as_weighted_average_errors( + {"force_real": DP_TEST_SPIN_WEIGHTED_METRIC_KEYS["force_real"]} + ) + ) + if find_force_mag == 1: + if spin_metrics.force_magnetic is None: + raise RuntimeError("Spin magnetic force metrics are unavailable.") + errors.update( + spin_metrics.as_weighted_average_errors( + { + "force_magnetic": DP_TEST_SPIN_WEIGHTED_METRIC_KEYS[ + "force_magnetic" + ] + } + ) + ) + return _ForceDetails( + reference_real=reference_real, + prediction_real=force_real, + reference_magnetic=reference_mag if find_force_mag == 1 else None, + prediction_magnetic=force_mag if find_force_mag == 1 else None, + ) diff --git a/deepmd/infer/model_test/property.py b/deepmd/infer/model_test/property.py new file mode 100644 index 0000000000..d16c51910e --- /dev/null +++ b/deepmd/infer/model_test/property.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of models fitting an arbitrary per-frame property.""" + +import logging + +from deepmd.infer.model_test.base import ( + ChunkContext, + ModelTester, + _write_per_frame_details, +) +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.eval_metrics import ( + mae, + rmse, +) + +log = logging.getLogger(__name__) + +__all__ = ["PropertyTester"] + + +class PropertyTester(ModelTester): + """Test a model of an arbitrary per-frame property.""" + + report = ( + ("mae_property", "PROPERTY MAE : {} units"), + ("rmse_property", "PROPERTY RMSE : {} units"), + ("mae_aproperty", "Atomic PROPERTY MAE : {} units"), + ("rmse_aproperty", "Atomic PROPERTY RMSE : {} units"), + ) + per_system_only = ("mae_aproperty", "rmse_aproperty") + + def add_data_requirements(self, data: DeepmdData) -> None: + """Declare the labels a property test consumes.""" + dp = self.dp + var_name = dp.get_var_name() + assert isinstance(var_name, str) + data.add(var_name, dp.task_dim, atomic=False, must=True, high_prec=True) + if self.atomic: + data.add( + f"atom_{var_name}", + dp.task_dim, + atomic=True, + must=False, + high_prec=True, + ) + if dp.get_dim_fparam() > 0: + data.add( + "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False + ) + if dp.get_dim_aparam() > 0: + data.add( + "aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False + ) + + def evaluate_chunk( + self, + data: DeepmdData, + test_data: dict, + context: ChunkContext, + ) -> dict[str, tuple[float, float]]: + """Evaluate one chunk of a property test.""" + dp = self.dp + var_name = dp.get_var_name() + mixed_type = data.mixed_type + natoms = len(test_data["type"][0]) + nframes = test_data["box"].shape[0] + + coord = test_data["coord"].reshape([nframes, -1]) + box = test_data["box"] if data.pbc else None + if mixed_type: + atype = test_data["type"].reshape([nframes, -1]) + else: + atype = test_data["type"][0] + fparam = test_data["fparam"] if dp.get_dim_fparam() > 0 else None + aparam = test_data["aparam"] if dp.get_dim_aparam() > 0 else None + + ret = dp.eval( + coord, + box, + atype, + fparam=fparam, + aparam=aparam, + atomic=self.atomic, + mixed_type=mixed_type, + ) + prediction = ret[0].reshape([nframes, dp.task_dim]) + + diff = prediction - test_data[var_name] + errors: dict[str, tuple[float, float]] = { + "mae_property": (mae(diff), prediction.size), + "rmse_property": (rmse(diff), prediction.size), + } + + atom_prediction = None + if self.atomic: + atom_prediction = ret[1].reshape([nframes, natoms * dp.task_dim]) + atom_diff = atom_prediction - test_data[f"atom_{var_name}"] + errors["mae_aproperty"] = (mae(atom_diff), atom_prediction.size) + errors["rmse_aproperty"] = (rmse(atom_diff), atom_prediction.size) + + if context.detail_path is not None: + _write_per_frame_details( + context, + suffix="property", + reference=test_data[var_name], + prediction=prediction, + ) + if self.atomic: + _write_per_frame_details( + context, + suffix="aproperty", + reference=test_data[f"atom_{var_name}"], + prediction=atom_prediction, + ) + + return errors diff --git a/deepmd/infer/model_test/tensor.py b/deepmd/infer/model_test/tensor.py new file mode 100644 index 0000000000..3b8ea305b2 --- /dev/null +++ b/deepmd/infer/model_test/tensor.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of atomic tensor models, such as dipole and polarizability.""" + +import logging +from typing import ( + ClassVar, +) + +import numpy as np + +from deepmd.infer.model_test.base import ( + ChunkContext, + ModelTester, + save_txt_file, +) +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.eval_metrics import ( + rmse, +) + +log = logging.getLogger(__name__) + +__all__ = ["DipoleTester", "PolarTester", "TensorTester"] + + +class TensorTester(ModelTester): + """Test a model of an atomic tensor, summed over atoms unless per-atom. + + A tensor model is evaluated over the atoms its selected types cover, so + both the reported error and the detail layout follow from the number of + such atoms. + """ + + #: Label of the per-frame quantity, and of its per-atom counterpart. + label: ClassVar[str] + atomic_label: ClassVar[str] + #: Number of components of the tensor. + ndof: ClassVar[int] + #: Per-component names used in the detail header. + components: ClassVar[tuple[str, ...]] + + def add_data_requirements(self, data: DeepmdData) -> None: + """Declare the label a tensor test consumes.""" + data.add( + self.atomic_label if self.atomic else self.label, + self.ndof, + atomic=self.atomic, + must=True, + high_prec=False, + type_sel=self.dp.get_sel_type(), + output_natoms_for_type_sel=True, + ) + + def evaluate_chunk( + self, + data: DeepmdData, + test_data: dict, + context: ChunkContext, + ) -> dict[str, tuple[float, float]]: + """Evaluate one chunk of a tensor test.""" + nframes = test_data["box"].shape[0] + coord = test_data["coord"].reshape([nframes, -1]) + box = test_data["box"] if data.pbc else None + atype = test_data["type"][0] + prediction = self.dp.eval(coord, box, atype).reshape([nframes, -1]) + + sel_type = self.dp.get_sel_type() + sel_natoms = int(sum(sum(atype == ii) for ii in sel_type)) + + if self.atomic: + sel_mask = np.isin(atype, sel_type) + prediction = prediction.reshape((nframes, -1, self.ndof))[ + :, sel_mask, : + ].reshape((nframes, -1)) + reference = ( + test_data[self.atomic_label] + .reshape((nframes, -1, self.ndof))[:, sel_mask, :] + .reshape((nframes, -1)) + ) + else: + prediction = np.sum(prediction.reshape((nframes, -1, self.ndof)), axis=1) + reference = test_data[self.label] + + rmse_tensor = rmse(prediction - reference) + errors: dict[str, tuple[float, float]] = { + "rmse": (rmse_tensor, prediction.size) + } + if not self.atomic: + errors["rmse_sqrtn"] = (rmse_tensor / np.sqrt(sel_natoms), prediction.size) + errors["rmse_n"] = (rmse_tensor / sel_natoms, prediction.size) + + if context.detail_path is not None: + width = self.ndof * (sel_natoms if self.atomic else 1) + save_txt_file( + context.detail_path.with_suffix(".out"), + np.concatenate( + ( + np.reshape(reference, [-1, width]), + np.reshape(prediction, [-1, width]), + ), + axis=1, + ), + header=self._detail_header(sel_natoms), + append=context.append_detail, + ) + + return errors + + def _detail_header(self, sel_natoms: int) -> str: + """Return the detail-file header for the layout in use.""" + if not self.atomic: + return " ".join( + [f"data_{name}" for name in self.components] + + [f"pred_{name}" for name in self.components] + ) + return " ".join( + [ + f"{prefix}_{name}{number}" + for prefix in ("data", "pred") + for number in range(1, sel_natoms + 1) + for name in self.components + ] + ) + + +class DipoleTester(TensorTester): + """Test a dipole model.""" + + label = "dipole" + atomic_label = "atomic_dipole" + ndof = 3 + components = ("x", "y", "z") + report = ( + ("rmse", "Dipole RMSE : {}"), + ("rmse_sqrtn", "Dipole RMSE/sqrtN : {}"), + ("rmse_n", "Dipole RMSE/N : {}"), + ) + report_footer = "The unit of error is the same as the unit of provided label." + per_system_only = ("rmse_sqrtn", "rmse_n") + + +class PolarTester(TensorTester): + """Test a polarizability model.""" + + label = "polarizability" + atomic_label = "atomic_polarizability" + ndof = 9 + components = ("pxx", "pxy", "pxz", "pyx", "pyy", "pyz", "pzx", "pzy", "pzz") + report = ( + ("rmse", "Polarizability RMSE : {}"), + ("rmse_sqrtn", "Polarizability RMSE/sqrtN : {}"), + ("rmse_n", "Polarizability RMSE/N : {}"), + ) + report_footer = "The unit of error is the same as the unit of provided label." + per_system_only = ("rmse_sqrtn", "rmse_n") diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 0db472dca9..30ff347540 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -5,6 +5,9 @@ import copy import functools import logging +from collections.abc import ( + Iterator, +) from concurrent.futures import ( ThreadPoolExecutor, as_completed, @@ -338,6 +341,44 @@ def get_test(self, ntests: int = -1) -> dict: self.modifier.modify_data(ret, self) return ret + def iter_test( + self, + *, + chunk_atoms: int, + numb_test: float = float("inf"), + ) -> Iterator[dict]: + """Yield the test data in chunks of at most ``chunk_atoms`` atoms. + + A set is loaded as a whole, so chunking here bounds what a consumer + holds at once rather than what is read. + + Parameters + ---------- + chunk_atoms : int + Upper bound on the number of atoms per chunk. A chunk always + carries at least one frame, however many atoms it has. + numb_test : float, optional + Upper bound on the number of frames served. A non-finite bound + serves the whole test set. + + Yields + ------ + dict + One chunk of test data, keyed as :meth:`get_test`. + """ + if not hasattr(self, "test_set"): + self._load_test_set(self.shuffle_test) + total = int(self.test_set["type"].shape[0]) + if np.isfinite(numb_test): + total = min(total, int(numb_test)) + step = max(1, int(chunk_atoms) // max(1, self.natoms)) + for begin in range(0, total, step): + idx = np.arange(begin, min(begin + step, total), dtype=np.int64) + chunk = self._get_subdata(self.test_set, idx=idx) + if self.modifier is not None: + self.modifier.modify_data(chunk, self) + yield chunk + def get_ntypes(self) -> int: """Number of atom types in the system.""" if self.type_map is not None: diff --git a/deepmd/utils/weight_avg.py b/deepmd/utils/weight_avg.py index 8328be5fcf..d2932fe4a4 100644 --- a/deepmd/utils/weight_avg.py +++ b/deepmd/utils/weight_avg.py @@ -6,24 +6,37 @@ import numpy as np -def weighted_average(errors: list[dict[str, tuple[float, float]]]) -> dict: - """Compute weighted average of prediction errors (MAE or RMSE) for model. +def merge_weighted_errors( + errors: list[dict[str, tuple[float, float]]], +) -> dict[str, tuple[float, float]]: + """Combine prediction errors, keeping the weight they were combined over. + + An MAE is the mean of the absolute errors and an RMSE the root of the mean + of their squares, so both are recovered exactly from the partial results by + weighting the mean, respectively the squared mean, by the number of + elements each was taken over. Combining partial results is therefore + equivalent to evaluating the whole set at once, which lets a caller + evaluate in chunks. Parameters ---------- errors : list[dict[str, tuple[float, float]]] - List: the error of systems - Dict: the error of quantities, name given by the key - str: the name of the quantity, must starts with 'mae' or 'rmse' - Tuple: (error, weight) + One ``{quantity: (error, weight)}`` mapping per partial result. A + quantity name starts with ``mae`` or ``rmse``. Returns ------- - Dict - weighted averages + dict[str, tuple[float, float]] + The combined ``(error, weight)`` of every quantity, itself suitable as + one partial result of a further combination. + + Raises + ------ + RuntimeError + If a quantity name identifies neither an MAE nor an RMSE. """ - sum_err = defaultdict(float) - sum_siz = defaultdict(int) + sum_err: dict[str, float] = defaultdict(float) + sum_siz: dict[str, float] = defaultdict(float) for err in errors: for kk, (ee, ss) in err.items(): if kk.startswith("mae"): @@ -33,11 +46,28 @@ def weighted_average(errors: list[dict[str, tuple[float, float]]]) -> dict: else: raise RuntimeError("unknown error type") sum_siz[kk] += ss - for kk in sum_err.keys(): - if kk.startswith("mae"): - sum_err[kk] = sum_err[kk] / sum_siz[kk] - elif kk.startswith("rmse"): - sum_err[kk] = np.sqrt(sum_err[kk] / sum_siz[kk]) - else: - raise RuntimeError("unknown error type") - return sum_err + merged: dict[str, tuple[float, float]] = {} + for kk, total in sum_err.items(): + weight = sum_siz[kk] + mean = total / weight + merged[kk] = (mean if kk.startswith("mae") else float(np.sqrt(mean)), weight) + return merged + + +def weighted_average(errors: list[dict[str, tuple[float, float]]]) -> dict: + """Compute weighted average of prediction errors (MAE or RMSE) for model. + + Parameters + ---------- + errors : list[dict[str, tuple[float, float]]] + List: the error of systems + Dict: the error of quantities, name given by the key + str: the name of the quantity, must starts with 'mae' or 'rmse' + Tuple: (error, weight) + + Returns + ------- + Dict + weighted averages + """ + return {kk: value for kk, (value, _) in merge_weighted_errors(errors).items()} diff --git a/source/tests/common/test_dp_test_ener_split.py b/source/tests/common/test_dp_test_ener_split.py index eb339b3af8..d98bf9b4aa 100644 --- a/source/tests/common/test_dp_test_ener_split.py +++ b/source/tests/common/test_dp_test_ener_split.py @@ -13,7 +13,7 @@ import numpy as np -from deepmd.entrypoints.test import ( +from deepmd.infer.model_test.ener import ( _split_optional_ener_outputs, ) diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py index 725de06b15..131f07149f 100644 --- a/source/tests/pt/test_dp_test.py +++ b/source/tests/pt/test_dp_test.py @@ -15,10 +15,12 @@ import torch from deepmd.entrypoints.test import test as dp_test -from deepmd.entrypoints.test import test_ener as dp_test_ener from deepmd.infer.deep_eval import ( DeepEval, ) +from deepmd.infer.model_test import ( + build_tester, +) from deepmd.pt.entrypoints.main import ( get_trainer, ) @@ -335,13 +337,11 @@ def test_force_weight(self) -> None: type_map=dp.get_type_map(), sort_atoms=False, ) - err = dp_test_ener( - dp, + err = build_tester(dp, atomic=False).run( data, self.system_dir, numb_test=1, detail_file=None, - has_atom_ener=False, ) test_data = data.get_test() coord = test_data["coord"].reshape([1, -1]) @@ -424,13 +424,11 @@ def test_stress(self) -> None: sort_atoms=False, ) numb_test = 1 - err = dp_test_ener( - dp, + err = build_tester(dp, atomic=False).run( data, self.system_dir, numb_test=numb_test, detail_file=self.detail_file, - has_atom_ener=False, ) os.unlink(tmp_model_path) diff --git a/source/tests/pt/test_weighted_avg.py b/source/tests/pt/test_weighted_avg.py index cbaa5e3692..d24a79657d 100644 --- a/source/tests/pt/test_weighted_avg.py +++ b/source/tests/pt/test_weighted_avg.py @@ -14,10 +14,12 @@ import numpy as np import torch -from deepmd.entrypoints.test import test_ener as dp_test_ener from deepmd.infer.deep_eval import ( DeepEval, ) +from deepmd.infer.model_test import ( + build_tester, +) from deepmd.pt.entrypoints.main import ( get_trainer, ) @@ -64,13 +66,11 @@ def test_dp_test_ener_without_spin(self) -> None: type_map=dp.get_type_map(), sort_atoms=False, ) - err = dp_test_ener( - dp, + err = build_tester(dp, atomic=False).run( data, system, numb_test=1, detail_file=None, - has_atom_ener=False, ) self.assertIn("mae_e", err, "'mae_e' key is missing in the result") self.assertNotIn( @@ -92,13 +92,11 @@ def test_dp_test_ener_with_multisys_and_with_virial(self) -> None: sort_atoms=False, ) err = [] - err_novirial = dp_test_ener( - dp, + err_novirial = build_tester(dp, atomic=False).run( data, system, numb_test=1, detail_file=None, - has_atom_ener=False, ) err.append(err_novirial) ener_nv, weight_nv = err_novirial["mae_e"] @@ -115,13 +113,11 @@ def test_dp_test_ener_with_multisys_and_with_virial(self) -> None: type_map=dp.get_type_map(), sort_atoms=False, ) - err_virial = dp_test_ener( - dp, + err_virial = build_tester(dp, atomic=False).run( data, system, numb_test=1, detail_file=None, - has_atom_ener=False, ) self.assertIn("mae_e", err_virial, "'mae_e' key is missing in the result") @@ -204,24 +200,53 @@ def test_dp_test_ener_with_spin(self) -> None: sort_atoms=False, ) - err = dp_test_ener( - dp, + err = build_tester(dp, atomic=False).run( data, system, numb_test=1, detail_file=None, - has_atom_ener=False, ) self.assertIn("mae_e", err, "'mae_e' key is missing in the result") self.assertIn("mae_fm", err, "'mae_fm' key is missing in the result") + # A spin model reports a real and a magnetic force in place of the + # plain one, and this system carries no virial label. self.assertNotIn( - "mae_v", err, "'mae_v' key should not be present in the result" + "mae_f", err, "'mae_f' key should not be present in the result" ) self.assertNotIn( - "mae_f", err, "'mae_f' key should not be present in the result" + "mae_v", err, "'mae_v' key should not be present in the result" ) os.unlink(self.tmp_model.name) + def test_dp_test_ener_with_spin_and_with_virial(self) -> None: + # The magnetic degrees of freedom enter the virial only through the + # virtual atoms, whose displacement the model removes again, so a spin + # model reports the virial and the stress like any energy model. + dp = DeepEval(self.tmp_model.name, head="PyTorch") + system = self.config["training"]["validation_data"]["systems"][0] + np.save( + os.path.join(system, "set.000", "virial.npy"), + np.ones([1, 9], dtype=np.float64), + ) + data = DeepmdData( + sys_path=system, + set_prefix="set", + shuffle_test=False, + type_map=dp.get_type_map(), + sort_atoms=False, + ) + + err = build_tester(dp, atomic=False).run( + data, + system, + numb_test=1, + detail_file=None, + ) + self.assertIn("mae_fm", err, "'mae_fm' key is missing in the result") + for key in ("mae_v", "rmse_v", "mae_va", "rmse_va", "mae_s", "rmse_s"): + self.assertIn(key, err, f"'{key}' key is missing in the result") + os.unlink(self.tmp_model.name) + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): @@ -232,6 +257,13 @@ def tearDown(self) -> None: os.remove(f) if f in ["stat_files"]: shutil.rmtree(f) + virial_path_fake = os.path.join( + self.config["training"]["validation_data"]["systems"][0], + "set.000", + "virial.npy", + ) + if os.path.exists(virial_path_fake): + os.remove(virial_path_fake) if __name__ == "__main__": From 5ea2b6605f0b30acb8f01a9ea33f7943af17c5c4 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 28 Jul 2026 16:37:35 +0800 Subject: [PATCH 2/5] fix(dpmodel): build the neighbor representation the atomic model declares The output-bias forward wrapper is the only caller of an atomic model that starts from raw coordinates, so it constructs the neighbor input itself. It always built a fixed-capacity neighbor list sized by `get_sel()`, even though the model already declares through `uses_graph_lower()` which representation it consumes, and already implements both `forward_common_atomic` and `forward_common_atomic_graph`. Every caller that starts from an extended input honours that declaration; this one did not. A graph-native model reports no finite neighbor capacity, so sizing a dense list from `get_sel()` is not merely wasteful there: the allocation is unbounded and the index array alone reaches tens of gigabytes on a few hundred atoms. Fine-tuning ran out of memory at `change_out_bias`, which drives this wrapper over sampled training frames. The wrapper now builds a carry-all `NeighborGraph` when the model is graph-native and keeps the neighbor list otherwise. Pair exclusion stays a build-time transform on both routes, folded into `edge_mask` by the graph builder. The graph route works on a flat node axis, so its result is restored to the per-frame layout the dense route returns and the contract with `compute_output_stats` is unchanged. The routing tests a capability, not a descriptor, so any graph-native model is covered. Backends whose atomic models never report `uses_graph_lower` keep their own dense wrappers untouched. --- .../dpmodel/atomic_model/base_atomic_model.py | 91 +++++++++++++------ 1 file changed, 65 insertions(+), 26 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index f2f2218443..891da33016 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -805,9 +805,22 @@ def _store_out_stat( self.out_std = out_std_data def _get_forward_wrapper_func(self) -> Callable[..., dict[str, np.ndarray]]: - """Get a forward wrapper of the atomic model for output bias calculation.""" + """Get a forward wrapper of the atomic model for output bias calculation. + + The wrapper starts from raw coordinates and therefore has to construct + the neighbor representation itself. It builds the one this model + declares through :meth:`uses_graph_lower`: a carry-all + ``NeighborGraph`` for graph-native models, whose neighbor count follows + the geometry, or the fixed-capacity neighbor list sized by + :meth:`get_sel` otherwise. Sizing a dense list from ``get_sel`` is not + merely wasteful for a graph-native model -- such a model reports no + finite capacity, so the allocation is unbounded. + """ import array_api_compat + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + ) from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, ) @@ -841,31 +854,57 @@ def model_forward( if charge_spin is not None: charge_spin = xp.asarray(charge_spin, device=device) - ( - extended_coord, - extended_atype, - mapping, - nlist, - ) = extend_input_and_build_neighbor_list( - coord, - atype, - self.get_rcut(), - self.get_sel(), - mixed_types=self.mixed_types(), - box=box, - # exclusion is a nlist-BUILD transform (decision #18/A4); - # forward_common_atomic consumes a pre-excluded nlist. - pair_excl=self.pair_excl, - ) - atomic_ret = self.forward_common_atomic( - extended_coord, - extended_atype, - nlist, - mapping=mapping, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - ) + if self.uses_graph_lower(): + nframes, nloc = atype.shape + # Pair exclusion is a neighbor-BUILD transform (decision + # #18/A4) on both routes; the graph builder folds it into + # ``edge_mask``. + graph = build_neighbor_graph( + coord, + atype, + box, + self.get_rcut(), + pair_excl=self.pair_excl, + ) + atomic_ret = self.forward_common_atomic_graph( + graph, + xp.reshape(atype, (-1,)), + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + # The graph route works on a flat node axis; restore the + # per-frame layout the dense route returns. + atomic_ret = { + kk: xp.reshape(vv, (nframes, nloc, *vv.shape[1:])) + for kk, vv in atomic_ret.items() + } + else: + ( + extended_coord, + extended_atype, + mapping, + nlist, + ) = extend_input_and_build_neighbor_list( + coord, + atype, + self.get_rcut(), + self.get_sel(), + mixed_types=self.mixed_types(), + box=box, + # exclusion is a nlist-BUILD transform (decision #18/A4); + # forward_common_atomic consumes a pre-excluded nlist. + pair_excl=self.pair_excl, + ) + atomic_ret = self.forward_common_atomic( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) # Convert outputs back to numpy arrays return {kk: to_numpy_array(vv) for kk, vv in atomic_ret.items()} From 01fe8fdb91ce5ebd0e75b6fb4b8a21f2ce49d7e5 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 2 Aug 2026 19:08:00 +0800 Subject: [PATCH 3/5] fix(test): handle empty and grouped tensor outputs --- deepmd/infer/model_test/base.py | 28 ++++++++++++++++++++-------- deepmd/infer/model_test/tensor.py | 16 ++++++++++------ 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/deepmd/infer/model_test/base.py b/deepmd/infer/model_test/base.py index e916b738be..b5457e3b2f 100644 --- a/deepmd/infer/model_test/base.py +++ b/deepmd/infer/model_test/base.py @@ -102,8 +102,8 @@ class ChunkContext: per-frame detail files numbered consistently across chunks. detail_group : int Zero-based test-group index. The first group keeps the historical - per-frame filenames; later groups include this index to avoid mixing - frames from different systems or LMDB subgroups. + detail filenames; later groups include this index to avoid mixing + data from different systems or LMDB subgroups. """ system: str @@ -213,8 +213,8 @@ def run( append_detail : bool, optional Whether the details of this system extend an existing file. detail_group : int, optional - Zero-based test-group index used to disambiguate per-frame detail - filenames across systems and LMDB subgroups. + Zero-based test-group index used to disambiguate detail filenames + across systems and LMDB subgroups. Returns ------- @@ -274,6 +274,20 @@ def log_errors(cls, errors: Mapping[str, tuple[float, float]]) -> None: log.info(cls.report_footer) +def _detail_output_path( + context: ChunkContext, + suffix: str, + *, + frame: int | None = None, +) -> Path: + """Build a detail path unique to its test group and optional frame.""" + detail_path = context.detail_path + assert detail_path is not None + group = f".{context.detail_group}" if context.detail_group else "" + frame_suffix = f".{frame}" if frame is not None else "" + return detail_path.with_suffix(f"{suffix}{group}{frame_suffix}") + + def _write_per_frame_details( context: ChunkContext, *, @@ -294,13 +308,11 @@ def _write_per_frame_details( prediction : np.ndarray Predicted values with shape ``(nframes, ...)``. """ - detail_path = context.detail_path - assert detail_path is not None + assert context.detail_path is not None for index in range(reference.shape[0]): frame = context.frame_offset + index - group = f".{context.detail_group}" if context.detail_group else "" save_txt_file( - detail_path.with_suffix(f".{suffix}.out{group}.{frame}"), + _detail_output_path(context, f".{suffix}.out", frame=frame), np.hstack( ( reference[index].reshape(-1, 1), diff --git a/deepmd/infer/model_test/tensor.py b/deepmd/infer/model_test/tensor.py index 27456d2b5a..fd858923c9 100644 --- a/deepmd/infer/model_test/tensor.py +++ b/deepmd/infer/model_test/tensor.py @@ -10,6 +10,7 @@ from deepmd.infer.model_test.base import ( ChunkContext, ModelTester, + _detail_output_path, save_txt_file, ) from deepmd.utils.data import ( @@ -27,7 +28,8 @@ class TensorTester(ModelTester): A tensor model is evaluated over the atoms its selected types cover, so both the reported error and the detail layout follow from the number of - such atoms. + such atoms. If a chunk contains none, a global tensor still has a defined + RMSE, but its atom-normalized errors and all atomic metrics are omitted. """ #: Label of the per-frame quantity, and of its per-atom counterpart. @@ -63,11 +65,10 @@ def evaluate_chunk( atype = test_data["type"][0] prediction = self.dp.eval(coord, box, atype).reshape([nframes, -1]) - sel_type = self.dp.get_sel_type() - sel_natoms = int(sum(sum(atype == ii) for ii in sel_type)) + sel_mask = np.isin(atype, self.dp.get_sel_type()) + sel_natoms = int(np.count_nonzero(sel_mask)) if self.atomic: - sel_mask = np.isin(atype, sel_type) prediction = prediction.reshape((nframes, -1, self.ndof))[ :, sel_mask, : ].reshape((nframes, -1)) @@ -80,18 +81,21 @@ def evaluate_chunk( prediction = np.sum(prediction.reshape((nframes, -1, self.ndof)), axis=1) reference = test_data[self.label] + if self.atomic and sel_natoms == 0: + return {} + rmse_tensor = rmse(prediction - reference) errors: dict[str, tuple[float, float]] = { "rmse": (rmse_tensor, prediction.size) } - if not self.atomic: + if not self.atomic and sel_natoms: errors["rmse_sqrtn"] = (rmse_tensor / np.sqrt(sel_natoms), prediction.size) errors["rmse_n"] = (rmse_tensor / sel_natoms, prediction.size) if context.detail_path is not None: width = self.ndof * (sel_natoms if self.atomic else 1) save_txt_file( - context.detail_path.with_suffix(".out"), + _detail_output_path(context, ".out"), np.concatenate( ( np.reshape(reference, [-1, width]), From b014db81c00eca6b5dcf2103405ae4d12f4c6d57 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 2 Aug 2026 19:22:36 +0800 Subject: [PATCH 4/5] fix(test): limit inherited dp test checks to one frame --- source/tests/pt/test_dp_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py index a66a470d76..4826048bce 100644 --- a/source/tests/pt/test_dp_test.py +++ b/source/tests/pt/test_dp_test.py @@ -43,7 +43,7 @@ class DPTest: def _run_dp_test( - self, use_input_json: bool, numb_test: int = 0, use_train: bool = False + self, use_input_json: bool, numb_test: int = 1, use_train: bool = False ) -> None: trainer = get_trainer(deepcopy(self.config)) with torch.device("cpu"): From 2d530bf8821cd1ec7e36b8af18c909f1fbc96cfe Mon Sep 17 00:00:00 2001 From: OutisLi Date: Mon, 3 Aug 2026 12:27:45 +0800 Subject: [PATCH 5/5] refactor(test): remove obsolete size collection --- deepmd/entrypoints/test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index 197ef4f42a..d9972b967e 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -137,7 +137,6 @@ def test( if len(all_sys) == 0: raise RuntimeError("Did not find valid system") err_coll = [] - siz_coll = [] # init random seed if rand_seed is not None: