diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index f2f2218443..26979d0e91 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,64 @@ 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=( + xp.reshape( + aparam, + (nframes * nloc, self.get_dim_aparam()), + ) + if aparam is not None + else None + ), + 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()} diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index a85f92be87..e642d25093 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, ) @@ -43,6 +46,7 @@ GLOBAL_ENER_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION, ) +from deepmd.utils import random as dp_random from deepmd.utils.data import ( DataRequirementItem, ) @@ -177,10 +181,15 @@ def _decode_frame( def _remap_keys(frame: dict[str, Any]) -> dict[str, Any]: - """Remap LMDB key names to DeePMD convention, pass through unknown keys.""" + """Remap LMDB key names to the canonical in-memory DeePMD convention.""" out = {} for k, v in frame.items(): - out[_KEY_REMAP.get(k, k)] = v + key = _KEY_REMAP.get(k, k) + if key.startswith("find_atomic_"): + key = "find_atom_" + key.removeprefix("find_atomic_") + elif key.startswith("atomic_"): + key = "atom_" + key.removeprefix("atomic_") + out[key] = v return out @@ -2318,6 +2327,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 +2363,114 @@ 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) + ) + groups: dict[int, list[int]] = {} + for begin, end in pairwise(starts): + indices = order[begin:end].copy() + if shuffle_test: + dp_random.shuffle(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 +2479,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 +2609,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 @@ -2605,6 +2746,8 @@ def _stack_frames( f"LMDB validation group mixes find_{key} values {availability}" ) has_key = availability[0] + if not has_key and req_info.get("must", False): + raise RuntimeError(f"Required LMDB test-data field {key!r} is missing.") result[f"find_{key}"] = 1.0 if has_key else 0.0 # Get repeat factor from registered requirements @@ -2684,6 +2827,26 @@ 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.""" + if self._frame_indices is not None: + frame_indices = self._frame_indices + if np.isfinite(numb_test): + frame_indices = frame_indices[: int(numb_test)] + step = max(1, int(chunk_atoms) // max(1, self._nloc)) + return ( + self._inner.get_test_by_indices(frame_indices[begin : begin + step]) + for begin in range(0, len(frame_indices), step) + ) + 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..d9972b967e 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__) @@ -170,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: @@ -178,6 +144,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 +157,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 +187,22 @@ def test( for data, sys_label in data_items: if sys_label != system: log.info(f"# testing sub-group : {sys_label}") - - 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 + # Only the very first tested group writes a fresh detail file; a + # system split into sub-groups extends it like any later system. + detail_group = len(err_coll) + append_detail = bool(detail_group) + + err = tester.run( + data, + sys_label, + numb_test, + detail_file, + append_detail=append_detail, + detail_group=detail_group, + ) 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 +210,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..4afd626d0e --- /dev/null +++ b/deepmd/infer/model_test/__init__.py @@ -0,0 +1,113 @@ +# 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. Lazy data sources such as LMDB +therefore need not fit in memory; ordinary ``DeepmdData`` systems materialize +the complete test system before it is sliced into evaluation chunks. +: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): + has_spin = dp.has_spin or dp.get_ntypes_spin() != 0 + tester = SpinEnerTester if 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..b5457e3b2f --- /dev/null +++ b/deepmd/infer/model_test/base.py @@ -0,0 +1,324 @@ +# 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 bounds the number of atoms sent to the evaluator at once. Lazy + data sources such as LMDB also decode only this many atoms, while ordinary + ``DeepmdData`` systems materialize the full test system before iteration. + 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))) + + +def save_txt_file( + fname: Path, data: np.ndarray, header: str = "", append: bool = False +) -> None: + """Save numpy array to test file. + + Parameters + ---------- + fname : Path + File to write. + 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 + """ + write_header = ( + "" if append and fname.exists() and fname.stat().st_size > 0 else header + ) + flags = "a" if append else "w" + with fname.open(flags, encoding="utf-8") as fp: + np.savetxt(fp, data, header=write_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. + detail_group : int + Zero-based test-group index. The first group keeps the historical + detail filenames; later groups include this index to avoid mixing + data from different systems or LMDB subgroups. + """ + + system: str + detail_file: str | None + append_detail: bool + frame_offset: int + detail_group: int = 0 + + @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 evaluation chunks. Lazy data sources such as LMDB + decode only the current chunk; ordinary ``DeepmdData`` systems materialize + the complete test system first. 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, + detail_group: int = 0, + ) -> 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. + detail_group : int, optional + Zero-based test-group index used to disambiguate detail filenames + across systems and LMDB subgroups. + + 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, + detail_group=detail_group, + ) + 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) + + +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, + *, + 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, ...)``. + """ + assert context.detail_path is not None + for index in range(reference.shape[0]): + frame = context.frame_offset + index + save_txt_file( + _detail_output_path(context, f".{suffix}.out", frame=frame), + np.hstack( + ( + reference[index].reshape(-1, 1), + prediction[index].reshape(-1, 1), + ) + ), + header=f"{context.system} - {frame}: data_{suffix} pred_{suffix}", + append=False, + ) diff --git a/deepmd/infer/model_test/dos.py b/deepmd/infer/model_test/dos.py new file mode 100644 index 0000000000..afe7d3bc5d --- /dev/null +++ b/deepmd/infer/model_test/dos.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of density-of-states models.""" + +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, +) + +__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 diff --git a/deepmd/infer/model_test/ener.py b/deepmd/infer/model_test/ener.py new file mode 100644 index 0000000000..6c02b1566e --- /dev/null +++ b/deepmd/infer/model_test/ener.py @@ -0,0 +1,695 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of energy models, including those carrying spin.""" + +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, + ) + +__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 + + +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..af2813ccdc --- /dev/null +++ b/deepmd/infer/model_test/property.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of models fitting an arbitrary per-frame property.""" + +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, +) + +__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=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 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..fd858923c9 --- /dev/null +++ b/deepmd/infer/model_test/tensor.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Testing of atomic tensor models, such as dipole and polarizability.""" + +from typing import ( + ClassVar, +) + +import numpy as np + +from deepmd.infer.model_test.base import ( + ChunkContext, + ModelTester, + _detail_output_path, + save_txt_file, +) +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.eval_metrics import ( + rmse, +) + +__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. 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. + 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_mask = np.isin(atype, self.dp.get_sel_type()) + sel_natoms = int(np.count_nonzero(sel_mask)) + + if self.atomic: + 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] + + 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 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( + _detail_output_path(context, ".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 = "atom_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 = "atom_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 c7a682c921..4bde3a8f9e 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -7,6 +7,7 @@ import logging from collections.abc import ( Iterable, + Iterator, ) from concurrent.futures import ( ThreadPoolExecutor, @@ -350,6 +351,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. + + The complete test system is materialized first, so chunking here bounds + the evaluation batch and its predictions rather than source-data memory. + + 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/dpmodel/test_dpa4_call_graph.py b/source/tests/common/dpmodel/test_dpa4_call_graph.py index 9896049594..c2410f2e5c 100644 --- a/source/tests/common/dpmodel/test_dpa4_call_graph.py +++ b/source/tests/common/dpmodel/test_dpa4_call_graph.py @@ -468,6 +468,31 @@ def test_dpa4_ener_fitting_call_graph_matches_dense() -> None: np.testing.assert_allclose(got.reshape(ref.shape), ref, rtol=1e-12, atol=1e-14) +def test_output_stat_graph_wrapper_flattens_aparam() -> None: + """The output-stat graph route passes atomic parameters on the flat node axis.""" + descriptor = make_descriptor() + fitting = SeZMEnergyFittingNet( + ntypes=3, + dim_descrpt=descriptor.get_dim_out(), + neuron=[16], + numb_aparam=2, + precision="float64", + seed=5, + ) + model = DPAtomicModel(descriptor, fitting, type_map=["A", "B", "C"]) + coord, atype, _ = make_inputs() + aparam = np.random.default_rng(17).normal(size=(*atype.shape, 2)) + + output = model._get_forward_wrapper_func()( + coord, + atype, + box=None, + aparam=aparam, + ) + + assert output["energy"].shape == (*atype.shape, 1) + + def make_bridging_descriptor(seed: int = 99) -> DescrptDPA4: """A SFPG-bridging variant of ``make_message_sensitive_descriptor()``. diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index f184efc53b..939d2b8a70 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -45,6 +45,7 @@ is_lmdb, make_neighbor_stat_data, ) +from deepmd.utils import random as dp_random from deepmd.utils.data import ( DataRequirementItem, ) @@ -771,6 +772,39 @@ def test_test_data_nloc_view(self): else: self.assertEqual(actual_value, expected_value) + def test_test_data_nloc_view_iterates_selected_frames(self): + """A label-availability view streams only its selected frame indices.""" + td = LmdbTestData(self._lmdb_path, type_map=self._type_map, shuffle_test=False) + selected = td.nloc_groups[9][::2] + view = LmdbTestDataNlocView(td, 9, frame_indices=selected) + + chunks = list(view.iter_test(chunk_atoms=9)) + + self.assertEqual(len(chunks), len(selected)) + expected = td.get_test_by_indices(selected) + np.testing.assert_array_equal( + np.concatenate([chunk["coord"] for chunk in chunks]), + expected["coord"], + ) + + def test_test_data_shuffle_uses_global_seed(self): + """The CLI random seed makes LMDB test-frame shuffling reproducible.""" + self.addCleanup(dp_random.seed, None) + dp_random.seed(123) + first = LmdbTestData( + self._lmdb_path, + type_map=self._type_map, + shuffle_test=True, + ) + dp_random.seed(123) + second = LmdbTestData( + self._lmdb_path, + type_map=self._type_map, + shuffle_test=True, + ) + + self.assertEqual(first.nloc_groups, second.nloc_groups) + def test_test_data_get_test_default_mixed(self): td = LmdbTestData(self._lmdb_path, type_map=self._type_map, shuffle_test=False) td.add("energy", 1, atomic=False, must=False, high_prec=True) @@ -1633,6 +1667,37 @@ def test_testdata_missing_key_not_found(self): self.assertEqual(result.get("find_atom_pref", 0.0), 0.0) tmpdir.cleanup() + def test_testdata_required_key_must_exist(self): + """A required LMDB test label cannot be replaced by a default value.""" + td = LmdbTestData(self._lmdb_path, type_map=self._type_map, shuffle_test=False) + td.add("missing_label", 1, atomic=False, must=True) + + with self.assertRaisesRegex(RuntimeError, "missing_label"): + td.get_test() + + def test_testdata_normalizes_atomic_label_prefix(self): + """LMDB atomic_* labels use the same in-memory atom_* keys as NPY data.""" + tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(tmpdir.cleanup) + path = _create_lmdb_with_extra_keys( + f"{tmpdir.name}/atomic-label.lmdb", + nframes=2, + natoms=self._natoms, + extra_keys={ + "atomic_dipole": (lambda n: (n, 3), np.float64), + }, + ) + td = LmdbTestData(path, type_map=self._type_map, shuffle_test=False) + td.add("atom_dipole", 3, atomic=True, must=True) + + result = td.get_test() + + self.assertEqual(result["find_atom_dipole"], 1.0) + self.assertEqual( + result["atom_dipole"].shape, + (2, self._natoms * 3), + ) + class _StalledPool: """A pool whose decoder exited, leaving its submissions unfinished. diff --git a/source/tests/common/test_dp_test_ener_split.py b/source/tests/common/test_dp_test_ener_split.py index eb339b3af8..a95abf6a49 100644 --- a/source/tests/common/test_dp_test_ener_split.py +++ b/source/tests/common/test_dp_test_ener_split.py @@ -10,10 +10,20 @@ """ import unittest +from types import ( + SimpleNamespace, +) import numpy as np -from deepmd.entrypoints.test import ( +from deepmd.infer.deep_pot import ( + DeepPot, +) +from deepmd.infer.model_test import ( + build_tester, +) +from deepmd.infer.model_test.ener import ( + SpinEnerTester, _split_optional_ener_outputs, ) @@ -80,5 +90,16 @@ def test_atomic_and_spin_no_hessian(self) -> None: self.assertIsNone(out.hessian) +def test_legacy_spin_model_dispatch() -> None: + """Legacy TF spin metadata selects the spin-aware energy tester.""" + dp = object.__new__(DeepPot) + dp.deep_eval = SimpleNamespace( + get_has_spin=lambda: False, + get_ntypes_spin=lambda: 1, + ) + + assert isinstance(build_tester(dp, atomic=False), SpinEnerTester) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py index 725de06b15..4826048bce 100644 --- a/source/tests/pt/test_dp_test.py +++ b/source/tests/pt/test_dp_test.py @@ -10,15 +10,20 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) import numpy as np 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, ) @@ -38,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"): @@ -335,13 +340,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]) @@ -402,6 +405,9 @@ def _prepare_virial_system(self) -> str: tmp_dir = tempfile.mkdtemp() shutil.copytree(src, tmp_dir, dirs_exist_ok=True) set_dir = Path(tmp_dir) / "set.000" + for data_file in set_dir.glob("*.npy"): + values = np.load(data_file) + np.save(data_file, np.repeat(values, 3, axis=0)) nframes = np.load(set_dir / "box.npy").shape[0] rng = np.random.default_rng(0) np.save(set_dir / "virial.npy", rng.standard_normal((nframes, 9))) @@ -416,26 +422,56 @@ def test_stress(self) -> None: os.close(tmp_fd) torch.jit.save(model, tmp_model_path) dp = DeepEval(tmp_model_path) - data = DeepmdData( - self.system_dir, - set_prefix="set", - shuffle_test=False, - type_map=dp.get_type_map(), - sort_atoms=False, - ) - numb_test = 1 - err = dp_test_ener( - dp, - data, - self.system_dir, - numb_test=numb_test, - detail_file=self.detail_file, - has_atom_ener=False, - ) + + def run_test(detail_file: str) -> tuple[dict, DeepmdData]: + data = DeepmdData( + self.system_dir, + set_prefix="set", + shuffle_test=False, + type_map=dp.get_type_map(), + sort_atoms=False, + ) + errors = build_tester(dp, atomic=False).run( + data, + self.system_dir, + numb_test=3, + detail_file=detail_file, + ) + return errors, data + + single_detail = f"{self.detail_file}_single_chunk" + with mock.patch.dict(os.environ, clear=False): + os.environ.pop("DP_TEST_CHUNK_ATOMS", None) + single_err, _ = run_test(single_detail) + with mock.patch.dict( + os.environ, + {"DP_TEST_CHUNK_ATOMS": "192"}, + clear=False, + ): + err, data = run_test(self.detail_file) os.unlink(tmp_model_path) + self.assertEqual(single_err.keys(), err.keys()) + for key in err: + np.testing.assert_allclose(single_err[key][0], err[key][0]) + self.assertEqual(single_err[key][1], err[key][1]) + for suffix in ("e", "f", "v", "s"): + single_path = Path(f"{single_detail}.{suffix}.out") + chunked_path = Path(f"{self.detail_file}.{suffix}.out") + np.testing.assert_allclose( + np.loadtxt(single_path, ndmin=2), + np.loadtxt(chunked_path, ndmin=2), + ) + self.assertEqual( + sum( + line.startswith("#") + for line in chunked_path.read_text().splitlines() + ), + 1, + ) + test_data = data.get_test() - box = test_data["box"][:numb_test].reshape(-1, 3, 3) + box = test_data["box"][:3].reshape(-1, 3, 3) volume = np.abs(np.linalg.det(box)).reshape(-1, 1) stress_out = np.loadtxt(self.detail_file + ".s.out", ndmin=2) @@ -465,13 +501,38 @@ def setUp(self) -> None: self.config = json.load(f) self.config["training"]["numb_steps"] = 1 self.config["training"]["save_freq"] = 1 - data_file = [str(Path(__file__).parent / "property/single")] + self.system_tmpdir = tempfile.TemporaryDirectory() + shutil.copytree( + Path(__file__).parent / "property/single", + self.system_tmpdir.name, + dirs_exist_ok=True, + ) + set_dir = Path(self.system_tmpdir.name) / "set.000000" + global_property = np.load(set_dir / "band_property.npy") + atom_types = np.load(set_dir / "real_atom_types.npy") + np.save( + set_dir / "atom_band_property.npy", + np.zeros( + ( + global_property.shape[0], + atom_types.shape[1], + global_property.shape[1], + ), + dtype=global_property.dtype, + ), + ) + data_file = [self.system_tmpdir.name] self.config["training"]["training_data"]["systems"] = data_file self.config["training"]["validation_data"]["systems"] = data_file self.config["model"] = deepcopy(model_property) self.config["model"]["type_map"] = [ self.config["model"]["type_map"][i] for i in [1, 0, 3, 2] ] + self.datafile = "test_dp_test_property_systems.txt" + Path(self.datafile).write_text( + "\n".join(data_file * 2) + "\n", + encoding="utf-8", + ) self.input_json = "test_dp_test_property.json" with open(self.input_json, "w") as fp: json.dump(self.config, fp, indent=4) @@ -488,8 +549,8 @@ def test_dp_test_1_frame(self) -> None: torch.jit.save(model, tmp_model_path) dp_test( model=tmp_model_path, - system=self.config["training"]["validation_data"]["systems"][0], - datafile=None, + system=None, + datafile=self.datafile, set_prefix="set", numb_test=0, rand_seed=None, @@ -499,6 +560,7 @@ def test_dp_test_1_frame(self) -> None: ) os.unlink(tmp_model_path) pred_property = np.loadtxt(self.detail_file + ".property.out.0")[:, 1] + self.assertTrue(os.path.exists(self.detail_file + ".property.out.1.0")) np.testing.assert_almost_equal( pred_property, to_numpy_array(result[model.get_var_name()])[0], @@ -510,10 +572,11 @@ def tearDown(self) -> None: os.remove(f) if f.startswith(self.detail_file): os.remove(f) - if f in ["lcurve.out", self.input_json]: + if f in ["lcurve.out", self.input_json, self.datafile]: os.remove(f) if f in ["stat_files"]: shutil.rmtree(f) + self.system_tmpdir.cleanup() if __name__ == "__main__": diff --git a/source/tests/pt/test_weighted_avg.py b/source/tests/pt/test_weighted_avg.py index cbaa5e3692..97a178e230 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, ) @@ -42,7 +44,13 @@ def setUp(self) -> None: self.config = json.load(f) self.config["training"]["numb_steps"] = 1 self.config["training"]["save_freq"] = 1 - data_file = [str(Path(__file__).parent / "water/data/single")] + self.system_tmpdir = tempfile.TemporaryDirectory() + shutil.copytree( + Path(__file__).parent / "water/data/single", + self.system_tmpdir.name, + dirs_exist_ok=True, + ) + data_file = [self.system_tmpdir.name] self.config["training"]["training_data"]["systems"] = data_file self.config["training"]["validation_data"]["systems"] = data_file self.config["model"] = deepcopy(model_se_e2_a) @@ -64,13 +72,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,21 +98,15 @@ 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"] - virial_path_fake = os.path.join( - self.config["training"]["validation_data"]["systems"][0], - "set.000", - "virial.npy", - ) + virial_path_fake = os.path.join(system, "set.000", "virial.npy") np.save(virial_path_fake, np.ones([1, 9], dtype=np.float64)) data = DeepmdData( sys_path=system, @@ -115,13 +115,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") @@ -151,8 +149,6 @@ def test_dp_test_ener_with_multisys_and_with_virial(self) -> None: f"Expected mae_e in avg_err to be {mae_e_expected} but got {avg_err['mae_e']}", ) - os.unlink(self.tmp_model.name) - def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): @@ -163,13 +159,10 @@ 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) + self.tmp_model.close() + if os.path.exists(self.tmp_model.name): + os.unlink(self.tmp_model.name) + self.system_tmpdir.cleanup() class Test_testener_spin(unittest.TestCase): @@ -180,7 +173,13 @@ def setUp(self) -> None: self.config = json.load(f) self.config["training"]["numb_steps"] = 1 self.config["training"]["save_freq"] = 1 - data_file = [str(Path(__file__).parent / "NiO/data/single")] + self.system_tmpdir = tempfile.TemporaryDirectory() + shutil.copytree( + Path(__file__).parent / "NiO/data/single", + self.system_tmpdir.name, + dirs_exist_ok=True, + ) + data_file = [self.system_tmpdir.name] self.config["training"]["training_data"]["systems"] = data_file self.config["training"]["validation_data"]["systems"] = data_file self.config["model"] = deepcopy(model_spin) @@ -204,23 +203,50 @@ 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") def tearDown(self) -> None: for f in os.listdir("."): @@ -232,6 +258,10 @@ def tearDown(self) -> None: os.remove(f) if f in ["stat_files"]: shutil.rmtree(f) + self.tmp_model.close() + if os.path.exists(self.tmp_model.name): + os.unlink(self.tmp_model.name) + self.system_tmpdir.cleanup() if __name__ == "__main__":