diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index b787738d76..2148fe65b7 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -217,25 +217,12 @@ def make_stat_input( if frame_mask is not None: stat_data = select_batch_frames(stat_data, frame_mask) accepted_batches += 1 - if ( - "find_fparam" in stat_data - and "fparam" in stat_data - and stat_data["find_fparam"] == 0.0 - ): - # for model using default fparam - stat_data.pop("fparam") - stat_data.pop("find_fparam") - for dd in stat_data: - if stat_data[dd] is None: - sys_stat[dd] = None - elif isinstance(stat_data[dd], torch.Tensor): - if dd not in sys_stat: - sys_stat[dd] = [] - sys_stat[dd].append(stat_data[dd]) - elif isinstance(stat_data[dd], np.float32): - sys_stat[dd] = stat_data[dd] - else: - pass + _append_stat_data(sys_stat, stat_data) + _append_missing_type_frames( + sys_stat, + datasets[system_index], + filter_enabled=filter_enabled, + ) if not sys_stat: if min_pair_dist > 0.0: @@ -263,6 +250,129 @@ def make_stat_input( return lst +def _append_stat_data(sys_stat: dict[str, Any], stat_data: dict[str, Any]) -> None: + """Append one statistics batch to the per-system accumulator.""" + if ( + "find_fparam" in stat_data + and "fparam" in stat_data + and stat_data["find_fparam"] == 0.0 + ): + # for model using default fparam + stat_data.pop("fparam") + stat_data.pop("find_fparam") + for dd, value in stat_data.items(): + if value is None: + sys_stat[dd] = None + elif isinstance(value, torch.Tensor): + if dd not in sys_stat: + sys_stat[dd] = [] + sys_stat[dd].append(value) + elif isinstance(value, np.float32): + sys_stat[dd] = value + + +def _append_missing_type_frames( + sys_stat: dict[str, Any], dataset: Any, filter_enabled: bool = False +) -> None: + """Add representative mixed-type frames for atom types missed by sampling. + + Global output statistics solve a per-type linear regression from the sampled + frame compositions. In mixed-type datasets a random small sample can miss a + type that exists elsewhere in the dataset, making that type's bias + unconstrained. We therefore append a minimal set of real frames, one by one, + until every type present in the full system is also present in the statistics + sample. Non-mixed systems have a fixed composition, so the initially sampled + frames already cover the system-level types. + """ + if "real_natoms_vec" not in sys_stat or sys_stat["real_natoms_vec"] is None: + return + if not hasattr(dataset, "data_system"): + return + sampled_natoms_vec = sys_stat["real_natoms_vec"] + if len(sampled_natoms_vec) == 0: + return + sampled_counts = torch.cat(sampled_natoms_vec, dim=0)[:, 2:].sum(dim=0) + dataset_counts, candidate_frames = _mixed_type_coverage(dataset) + if dataset_counts is None or candidate_frames is None: + return + + missing_types = np.flatnonzero((dataset_counts > 0) & (sampled_counts.numpy() == 0)) + if len(missing_types) == 0: + return + + # Import lazily to keep this utility independent from dataloader import time. + from deepmd.pt.utils.dataloader import ( + collate_batch, + ) + + used_frames: set[int] = set() + while len(missing_types) > 0: + extra_batch: dict[str, Any] | None = None + for type_i in missing_types: + for frame_idx in candidate_frames[int(type_i)]: + if frame_idx in used_frames: + continue + used_frames.add(frame_idx) + # Reuse the dataset and collate path so augmented frames have + # exactly the same tensor layout as DataLoader batches. + with torch.device("cpu"): + candidate_batch = collate_batch([dataset[frame_idx]]) + frame_mask = pair_margin_frame_mask(candidate_batch, filter_enabled) + if frame_mask is not None and not torch.any(frame_mask): + continue + if frame_mask is not None: + candidate_batch = select_batch_frames(candidate_batch, frame_mask) + extra_batch = candidate_batch + break + if extra_batch is not None: + break + if extra_batch is None: + log.warning( + "No frame containing atom types %s satisfies " + "filter_enabled=%s; statistics will not cover these types.", + missing_types.tolist(), + filter_enabled, + ) + break + _append_stat_data(sys_stat, extra_batch) + sampled_counts += extra_batch["real_natoms_vec"][:, 2:].sum(dim=0) + missing_types = np.flatnonzero( + (dataset_counts > 0) & (sampled_counts.numpy() == 0) + ) + + +def _mixed_type_coverage( + dataset: Any, +) -> tuple[np.ndarray | None, list[list[int]] | None]: + """Return full-dataset type counts and candidate frames for every type.""" + data_system = dataset.data_system + if not getattr(data_system, "mixed_type", False): + return None, None + ntypes = data_system.get_ntypes() + counts = np.zeros(ntypes, dtype=np.int64) + candidate_frames: list[list[int]] = [[] for _ in range(ntypes)] + frame_offset = 0 + for set_dir, frame_end in zip( + data_system.dirs, data_system.prefix_sum, strict=True + ): + type_path = set_dir / "real_atom_types.npy" + real_type = type_path.load_numpy() + if getattr(data_system, "enforce_type_map", False): + real_type = real_type.copy() + real_mask = real_type >= 0 + # Negative types denote padding, not indices into the type map. + real_type[real_mask] = data_system.type_idx_map[real_type[real_mask]] + real_type = real_type.reshape(frame_end - frame_offset, data_system.natoms) + for type_i in range(ntypes): + frame_hits = np.flatnonzero((real_type == type_i).any(axis=1)) + counts[type_i] += int((real_type == type_i).sum()) + candidate_frames[type_i].extend( + frame_offset + int(frame_idx) for frame_idx in frame_hits + ) + frame_offset = frame_end + return counts, candidate_frames + + def _full_pass_loader(dataloader: Any) -> Any: """Return a loader that covers the dataset once, whatever sampler it uses. diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index ac927d0e4f..2a7d930010 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -159,6 +159,7 @@ def collect_batches( if dd == "natoms_vec": stat_data[dd] = stat_data[dd].astype(np.int32) sys_stat[dd].append(stat_data[dd]) + _append_missing_type_frames(data, ii, sys_stat) for dd in sys_stat: if merge_sys: for bb in sys_stat[dd]: @@ -168,6 +169,202 @@ def collect_batches( return all_stat +def _append_missing_type_frames( + data: Any, + sys_idx: int, + sys_stat: dict[str, list[Any]], + *, + normalized: bool = False, + filter_enabled: bool = False, +) -> None: + """Append representative mixed-type frames for types missed by sampling. + + Energy/output bias statistics regress one bias per atom type from the sampled + frame compositions. Mixed-type systems can contain types that do not appear + in the small random statistics sample. When that happens, append the first acceptable + frame containing each missing type so the regression is constrained for every + type that exists in the underlying system. Standard (non-mixed) systems have + fixed composition and do not need augmentation. + """ + if "real_natoms_vec" not in sys_stat or not hasattr(data, "data_systems"): + return + mixed_systems = getattr(data, "mixed_systems", False) + if mixed_systems: + dataset_counts, candidate_frames = _mixed_system_coverage(data) + else: + data_system = data.data_systems[sys_idx] + dataset_counts, frames_for_type = _mixed_type_coverage(data_system) + if dataset_counts is None or frames_for_type is None: + return + candidate_frames = [ + [(sys_idx, frame_idx) for frame_idx in frames] for frames in frames_for_type + ] + if dataset_counts is None or candidate_frames is None: + return + sampled_counts = np.concatenate(sys_stat["real_natoms_vec"], axis=0)[:, 2:].sum( + axis=0 + ) + missing_types = np.flatnonzero((dataset_counts > 0) & (sampled_counts == 0)) + if len(missing_types) == 0: + return + + used_frames: set[tuple[int, int]] = set() + while len(missing_types) > 0: + appended = False + for type_i in missing_types: + for system_idx, frame_idx in candidate_frames[int(type_i)]: + candidate = (system_idx, frame_idx) + if candidate in used_frames: + continue + used_frames.add(candidate) + extra_batch = _get_representative_batch( + data, + system_idx, + frame_idx, + merge_mixed=mixed_systems, + ) + # The generic statistics path now consumes canonical batches; + # the TF collect_batches entry point still needs legacy keys. + if normalized: + extra_batch = normalize_batch(extra_batch) + frame_mask = pair_margin_frame_mask(extra_batch, filter_enabled) + if frame_mask is not None and not frame_mask.all(): + continue + if not _append_compatible_batch(sys_stat, extra_batch): + continue + sampled_counts += ( + extra_batch["real_natoms_vec"].reshape(1, -1)[:, 2:].sum(axis=0) + ) + appended = True + break + if appended: + break + if not appended: + if mixed_systems: + log.warning( + "Cannot add shape-compatible representative frames for atom " + "types %s in mixed-system statistics.", + missing_types.tolist(), + ) + break + missing_types = np.flatnonzero((dataset_counts > 0) & (sampled_counts == 0)) + + +def _get_representative_batch( + data: Any, + system_idx: int, + frame_idx: int, + *, + merge_mixed: bool, +) -> dict[str, Any]: + """Load one frame and give it the same layout as an ordinary batch.""" + data_system = data.data_systems[system_idx] + extra_batch = data_system.get_single_frame(frame_idx, num_worker=1) + extra_batch["natoms_vec"] = data.natoms_vec[system_idx].astype(np.int32) + extra_batch["default_mesh"] = data.default_mesh[system_idx] + for key, value in list(extra_batch.items()): + if ( + key not in {"natoms_vec", "default_mesh"} + and isinstance(value, np.ndarray) + and value.ndim >= 1 + ): + extra_batch[key] = value.reshape((1, *value.shape)) + if not merge_mixed: + return extra_batch + + # Reuse the production mixed-batch merger so atomic labels and input arrays + # follow exactly the same padding rules as randomly sampled mixed batches. + merged_batch = data._merge_batch_data([extra_batch]) + if "real_natoms_vec" in extra_batch: + # _merge_batch_data derives this field from fixed system composition; + # preserve the actual per-frame composition for mixed-type systems. + merged_batch["real_natoms_vec"] = extra_batch["real_natoms_vec"] + return merged_batch + + +def _append_compatible_batch( + sys_stat: dict[str, list[Any]], extra_batch: dict[str, Any] +) -> bool: + """Append a batch without changing the sampled schema or array widths.""" + selected: dict[str, Any] = {} + for key, value in extra_batch.items(): + if key not in sys_stat: + continue + reference = sys_stat[key][0] + if isinstance(reference, np.ndarray) and isinstance(value, np.ndarray): + reference_shape = ( + reference.shape[1:] if reference.ndim >= 2 else reference.shape + ) + value_shape = value.shape[1:] if value.ndim >= 2 else value.shape + if reference_shape != value_shape: + return False + selected[key] = value + if "real_natoms_vec" not in selected: + return False + for key, value in selected.items(): + sys_stat[key].append(value) + return True + + +def _mixed_system_coverage( + data: Any, +) -> tuple[np.ndarray | None, list[list[tuple[int, int]]] | None]: + """Aggregate type counts and representative frames across mixed systems.""" + if not hasattr(data, "data_systems") or len(data.data_systems) == 0: + return None, None + ntypes = int(data.get_ntypes()) + counts = np.zeros(ntypes, dtype=np.int64) + candidate_frames: list[list[tuple[int, int]]] = [[] for _ in range(ntypes)] + for system_idx, data_system in enumerate(data.data_systems): + system_counts, frames_for_type = _mixed_type_coverage(data_system) + if system_counts is None or frames_for_type is None: + system_counts = np.asarray(data.natoms_vec[system_idx][2:], dtype=np.int64) + # Fixed composition does not imply identical geometry: a later + # frame may pass the pair-margin filter when the first one fails. + nframes = getattr(data_system, "prefix_sum", [1])[-1] + frames_for_type = [ + list(range(nframes)) if count > 0 else [] for count in system_counts + ] + counts[: len(system_counts)] += system_counts + for type_i, frames in enumerate(frames_for_type): + candidate_frames[type_i].extend( + (system_idx, frame_idx) for frame_idx in frames + ) + return counts, candidate_frames + + +def _mixed_type_coverage( + data_system: Any, +) -> tuple[np.ndarray | None, list[list[int]] | None]: + """Return type counts and all candidates, in on-disk order. + + A representative can be rejected by the geometry filter or batch schema; + retaining later frames lets statistics cover each type that has valid data. + """ + if not getattr(data_system, "mixed_type", False): + return None, None + ntypes = data_system.get_ntypes() + counts = np.zeros(ntypes, dtype=np.int64) + frames_for_type: list[list[int]] = [[] for _ in range(ntypes)] + frame_offset = 0 + for set_dir, frame_end in zip( + data_system.dirs, data_system.prefix_sum, strict=True + ): + real_type = (set_dir / "real_atom_types.npy").load_numpy() + if getattr(data_system, "enforce_type_map", False): + real_type = real_type.copy() + real_mask = real_type >= 0 + # Negative types denote padding, not indices into the type map. + real_type[real_mask] = data_system.type_idx_map[real_type[real_mask]] + real_type = real_type.reshape(frame_end - frame_offset, data_system.natoms) + for type_i in range(ntypes): + frame_hits = np.flatnonzero((real_type == type_i).any(axis=1)) + counts[type_i] += int((real_type == type_i).sum()) + frames_for_type[type_i].extend((frame_hits + frame_offset).tolist()) + frame_offset = frame_end + return counts, frames_for_type + + def make_stat_input( data: Any, nbatches: int, @@ -218,11 +415,19 @@ def make_stat_input( ) log.info(f"Skipping statistical system {ii}: {reason}.") continue + sys_stat = {key: [batch[key] for batch in batches] for key in batches[0]} + _append_missing_type_frames( + data, + ii, + sys_stat, + normalized=True, + filter_enabled=min_pair_dist > 0.0, + ) # Arrays of two or more dimensions lead with the frame axis; every # other entry (find_* flags, an absent box) is constant per system. lst.append( { - key: np.concatenate([batch[key] for batch in batches], axis=0) + key: np.concatenate(sys_stat[key], axis=0) if isinstance(value, np.ndarray) and value.ndim >= 2 else value for key, value in batches[0].items() diff --git a/doc/conf.py b/doc/conf.py index 5753b2afed..78b3bc2a85 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -142,6 +142,10 @@ # This pattern also affects html_static_path and html_extra_path. exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] +# External inventories are optional cross-reference metadata. Limit connection +# and read-inactivity waits so Sphinx can warn and continue if a host stalls. +# This is not a deadline for the full download or documentation build. +intersphinx_timeout = 30 intersphinx_mapping = { "numpy": ("https://docs.scipy.org/doc/numpy/", None), "python": ("https://docs.python.org/", None), diff --git a/source/tests/common/test_model_stat.py b/source/tests/common/test_model_stat.py new file mode 100644 index 0000000000..97612d3307 --- /dev/null +++ b/source/tests/common/test_model_stat.py @@ -0,0 +1,257 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for backend-agnostic statistics sampling helpers.""" + +import unittest + +import numpy as np + +from deepmd.utils.model_stat import ( + make_stat_input, +) + + +class _FakeTypePath: + """Fake path object that returns in-memory atom types.""" + + def __init__(self, real_types: np.ndarray) -> None: + self.real_types = real_types + + def load_numpy(self) -> np.ndarray: + """Return the stored atom-type array.""" + return self.real_types + + +class _FakeSetDir: + """Fake set directory exposing ``real_atom_types.npy``.""" + + def __init__(self, real_types: np.ndarray) -> None: + self.real_types = real_types + + def __truediv__(self, name: str) -> _FakeTypePath: + """Return a fake path for the atom-type file.""" + assert name == "real_atom_types.npy" + return _FakeTypePath(self.real_types) + + +class _FakeMixedDataSystem: + """Minimal mixed-type data system for stat sampling tests.""" + + mixed_type = True + enforce_type_map = False + natoms = 2 + dirs: list[_FakeSetDir] + prefix_sum: list[int] + + def __init__(self, real_types: np.ndarray | None = None) -> None: + if real_types is None: + real_types = np.array([[0, -1], [1, -1]], dtype=np.int32) + self.dirs = [_FakeSetDir(real_types)] + self.prefix_sum = [len(real_types)] + + def get_ntypes(self) -> int: + """Return the number of real atom types.""" + return 2 + + def get_single_frame(self, index: int, num_worker: int = 1) -> dict: + """Return the representative frame containing the missing type.""" + assert index == 1 + return { + "coord": np.zeros((6,), dtype=np.float32), + "type": np.array([1, -1], dtype=np.int32), + "atype": np.array([1, -1], dtype=np.int32), + "box": np.eye(3, dtype=np.float32).reshape(-1), + "real_natoms_vec": np.array([2, 2, 0, 1], dtype=np.int32), + "find_energy": np.float32(1.0), + "energy": np.array([1.0], dtype=np.float64), + "fid": index, + } + + +class _FakeMixedData: + """Minimal multi-system data wrapper for ``make_stat_input``.""" + + mixed_systems = False + natoms_vec: list[np.ndarray] + default_mesh: list[np.ndarray] + + def __init__(self, system: _FakeMixedDataSystem | None = None) -> None: + self.data_systems = [_FakeMixedDataSystem() if system is None else system] + self.natoms_vec = [np.array([2, 2, 1, 0], dtype=np.int32)] + self.default_mesh = [np.array([], dtype=np.int32)] + + def get_nsystems(self) -> int: + """Return the number of systems.""" + return 1 + + def get_batch(self, sys_idx: int | None = None) -> dict: + """Return the initially sampled batch that misses one type.""" + assert sys_idx == 0 + return { + "coord": np.zeros((1, 6), dtype=np.float32), + "type": np.array([[0, -1]], dtype=np.int32), + "box": np.eye(3, dtype=np.float32).reshape(1, 9), + "real_natoms_vec": np.array([[2, 2, 1, 0]], dtype=np.int32), + "natoms_vec": np.array([2, 2, 1, 0], dtype=np.int32), + "default_mesh": np.array([], dtype=np.int32), + "find_energy": np.float32(1.0), + "energy": np.array([[0.0]], dtype=np.float64), + } + + +class _FilteredMixedDataSystem(_FakeMixedDataSystem): + """Two rare-type frames, of which only the later frame is valid.""" + + def __init__(self) -> None: + super().__init__(np.array([[0, -1], [1, -1], [1, -1]], dtype=np.int32)) + self.read_frames = [] + + def get_single_frame(self, index: int, num_worker: int = 1) -> dict: + self.read_frames.append(index) + frame = super().get_single_frame(1, num_worker) + frame["pair_margin"] = np.array([0.2 if index == 1 else 2.0]) + frame["find_pair_margin"] = np.float32(1.0) + return frame + + +class _FilteredMixedData(_FakeMixedData): + """Keep the initial type-zero sample valid while filtering candidates.""" + + def __init__(self) -> None: + super().__init__(_FilteredMixedDataSystem()) + + def get_nbatches(self) -> list[int]: + return [3] + + def get_batch(self, sys_idx: int | None = None) -> dict: + batch = super().get_batch(sys_idx) + batch["pair_margin"] = np.array([[2.0]]) + batch["find_pair_margin"] = np.float32(1.0) + return batch + + +class _FakeFixedDataSystem: + """One-frame fixed-composition system used by mixed-batch tests.""" + + mixed_type = False + + def __init__(self, atom_type: int) -> None: + self.atom_type = atom_type + + def get_single_frame(self, index: int, num_worker: int = 1) -> dict: + """Return the only frame in this fixed-composition system.""" + assert index == 0 + return { + "coord": np.full((6,), self.atom_type, dtype=np.float32), + "type": np.full((2,), self.atom_type, dtype=np.int32), + "atype": np.full((2,), self.atom_type, dtype=np.int32), + "box": np.eye(3, dtype=np.float32).reshape(-1), + "find_energy": np.float32(1.0), + "energy": np.array([float(self.atom_type)], dtype=np.float64), + "fid": index, + } + + +class _FakeMixedSystemsData: + """Mixed-batch wrapper whose random sample always misses the rare system.""" + + mixed_systems = True + + def __init__(self) -> None: + self.data_systems = [_FakeFixedDataSystem(0), _FakeFixedDataSystem(1)] + self.natoms_vec = [ + np.array([2, 2, 2, 0], dtype=np.int32), + np.array([2, 2, 0, 2], dtype=np.int32), + ] + self.default_mesh = [ + np.array([], dtype=np.int32), + np.array([], dtype=np.int32), + ] + + def get_nsystems(self) -> int: + """Return the number of underlying systems.""" + return 2 + + def get_ntypes(self) -> int: + """Return the global number of atom types.""" + return 2 + + def get_batch(self, sys_idx: int | None = None) -> dict: + """Mimic a low-probability rare system by always sampling type zero.""" + raw_batch = self.data_systems[0].get_single_frame(0) + raw_batch = { + key: value.reshape((1, *value.shape)) + if isinstance(value, np.ndarray) and value.ndim >= 1 + else value + for key, value in raw_batch.items() + if key not in {"atype", "fid"} + } + raw_batch["natoms_vec"] = self.natoms_vec[0] + raw_batch["default_mesh"] = self.default_mesh[0] + return self._merge_batch_data([raw_batch]) + + def _merge_batch_data(self, batch_data: list[dict]) -> dict: + """Merge same-sized frames using the production mixed-batch schema.""" + return { + "natoms_vec": np.array([2, 2, 2, 0], dtype=np.int32), + "real_natoms_vec": np.vstack([batch["natoms_vec"] for batch in batch_data]), + "type": np.concatenate([batch["type"] for batch in batch_data]), + "default_mesh": np.array([], dtype=np.int32), + "coord": np.concatenate([batch["coord"] for batch in batch_data]), + "box": np.concatenate([batch["box"] for batch in batch_data]), + "find_energy": batch_data[0]["find_energy"], + "energy": np.concatenate([batch["energy"] for batch in batch_data]), + } + + +class TestModelStatSamplingCoverage(unittest.TestCase): + """Mixed-type make_stat_input should cover types beyond initial batches.""" + + def test_make_stat_input_appends_missing_mixed_type_frame(self) -> None: + """Append a representative frame when the first batch misses a type.""" + data = _FakeMixedData() + ordinary_batch = data.get_batch(sys_idx=0) + self.assertIn("type", ordinary_batch) + self.assertNotIn("atype", ordinary_batch) + self.assertNotIn("fid", ordinary_batch) + + raw_frame = data.data_systems[0].get_single_frame(1) + self.assertIn("type", raw_frame) + self.assertIn("atype", raw_frame) + self.assertIn("fid", raw_frame) + + sampled = make_stat_input(data, nbatches=1) + + self.assertEqual(len(sampled), 1) + counts = sampled[0]["real_natoms_vec"][:, 2:].sum(axis=0) + self.assertTrue(np.all(counts > 0)) + self.assertEqual(sampled[0]["atype"].shape, (2, 2)) + self.assertEqual(sampled[0]["energy"].shape[0], 2) + self.assertNotIn("fid", sampled[0]) + + def test_filtered_statistics_try_later_rare_type_frames(self) -> None: + """A rejected first representative must not hide a valid rare type.""" + for threshold, expected_reads in [(0.0, [1]), (0.5, [1, 2])]: + with self.subTest(min_pair_dist=threshold): + data = _FilteredMixedData() + sampled = make_stat_input(data, nbatches=1, min_pair_dist=threshold) + counts = sampled[0]["real_natoms_vec"][:, 2:].sum(axis=0) + np.testing.assert_array_equal(counts, [1, 1]) + self.assertEqual(data.data_systems[0].read_frames, expected_reads) + if threshold: + self.assertTrue(np.all(sampled[0]["pair_margin"] >= 1.0)) + + def test_make_stat_input_covers_rare_type_in_mixed_systems(self) -> None: + """Cover a low-probability system in ``batch_size: mixed`` sampling.""" + sampled = make_stat_input(_FakeMixedSystemsData(), nbatches=1) + + self.assertEqual(len(sampled), 2) + for system in sampled: + counts = system["real_natoms_vec"][:, 2:].sum(axis=0) + self.assertTrue(np.all(counts > 0)) + self.assertEqual(system["atype"].shape, (2, 2)) + self.assertEqual(system["energy"].shape, (2, 1)) + self.assertNotIn("fid", system) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_observed_type.py b/source/tests/pt/test_observed_type.py index a2c7e37a73..f9388cd2d4 100644 --- a/source/tests/pt/test_observed_type.py +++ b/source/tests/pt/test_observed_type.py @@ -17,10 +17,15 @@ import numpy as np import torch +from deepmd.pt.utils.dataloader import ( + collate_batch, +) from deepmd.pt.utils.stat import ( + _append_missing_type_frames, _restore_observed_type_from_file, _save_observed_type_to_file, collect_observed_types, + make_stat_input, ) from deepmd.utils.path import ( DPPath, @@ -72,6 +77,164 @@ def test_out_of_range_index_ignored(self) -> None: self.assertEqual(result, ["O"]) +class _FakeTypePath: + """Fake path object that returns in-memory atom types.""" + + def __init__(self, real_types: np.ndarray) -> None: + self.real_types = real_types + + def load_numpy(self) -> np.ndarray: + """Return the stored atom-type array.""" + return self.real_types + + +class _FakeSetDir: + """Fake set directory exposing ``real_atom_types.npy``.""" + + def __init__(self, real_types: np.ndarray) -> None: + self.real_types = real_types + + def __truediv__(self, name: str) -> _FakeTypePath: + """Return a fake path for the atom-type file.""" + assert name == "real_atom_types.npy" + return _FakeTypePath(self.real_types) + + +class _FakeMixedDataSystem: + """Minimal mixed-type data system for stat sampling tests.""" + + mixed_type = True + enforce_type_map = False + natoms = 2 + dirs: list[_FakeSetDir] + prefix_sum: list[int] + + def __init__(self) -> None: + self.dirs = [_FakeSetDir(np.array([[0, -1], [1, -1], [1, -1]], dtype=np.int32))] + self.prefix_sum = [3] + + def get_ntypes(self) -> int: + """Return the number of real atom types.""" + return 2 + + +class _FakeMixedDataset: + """Minimal PyTorch dataset wrapper for mixed-type stat sampling.""" + + data_system: _FakeMixedDataSystem + + def __init__(self, min_pair_distances: tuple[float, ...] = (1.0, 1.0, 1.0)) -> None: + self.data_system = _FakeMixedDataSystem() + self.min_pair_distances = min_pair_distances + + def __getitem__(self, index: int) -> dict: + """Return the representative frame containing the missing type.""" + atom_type = 0 if index == 0 else 1 + return { + "coord": np.full((6,), index, dtype=np.float32), + "atype": np.array([atom_type, -1], dtype=np.int32), + "box": np.eye(3, dtype=np.float32).reshape(-1), + "real_natoms_vec": np.array( + [2, 2, int(atom_type == 0), int(atom_type == 1)], dtype=np.int32 + ), + "pair_margin": np.array( + [self.min_pair_distances[index] / 0.5], dtype=np.float64 + ), + "find_pair_margin": np.float32(1.0), + } + + +class TestStatSamplingCoverage(unittest.TestCase): + """Mixed-type statistics samples should cover all dataset types.""" + + def test_make_stat_input_covers_later_valid_rare_type(self) -> None: + """The public caller supplements its bounded sample after filtering.""" + dataset = _FakeMixedDataset((1.0, 0.1, 0.8)) + with torch.device("cpu"): + # The one-batch budget initially observes only type 0. Type 1's + # first representative is rejected; its later frame must survive + # supplementation, concatenation, and final device conversion. + dataloader = [collate_batch([dataset[index]]) for index in range(3)] + for options in ( + {"min_pair_dist": 0.5}, + {"filter_enabled_by_system": [True]}, + ): + with self.subTest(options=options): + sampled = make_stat_input( + [dataset], [dataloader], nbatches=1, **options + ) + self.assertEqual(len(sampled), 1) + result = sampled[0] + torch.testing.assert_close( + result["atype"].cpu(), + torch.tensor([[0, -1], [1, -1]], dtype=torch.int32), + ) + torch.testing.assert_close( + result["coord"].cpu(), + torch.stack([torch.zeros(6), torch.full((6,), 2.0)]), + ) + self.assertTrue(torch.all(result["pair_margin"] >= 1.0)) + self.assertTrue( + torch.all(result["real_natoms_vec"][:, 2:].sum(dim=0) > 0) + ) + + def test_append_missing_mixed_type_frame(self) -> None: + """Append a representative frame under an explicit CPU device.""" + with torch.device("cpu"): + sys_stat = { + "coord": [torch.zeros((1, 6), dtype=torch.float32)], + "atype": [torch.tensor([[0, -1]], dtype=torch.int32)], + "box": [torch.eye(3, dtype=torch.float32).reshape(1, 9)], + "real_natoms_vec": [torch.tensor([[2, 2, 1, 0]], dtype=torch.int32)], + } + + _append_missing_type_frames(sys_stat, _FakeMixedDataset()) + + self.assertEqual(len(sys_stat["real_natoms_vec"]), 2) + sampled_counts = torch.cat(sys_stat["real_natoms_vec"], dim=0)[:, 2:].sum( + dim=0 + ) + self.assertTrue(torch.all(sampled_counts > 0)) + + def test_append_missing_type_uses_later_valid_distance_frame(self) -> None: + """Skip a too-close rare-type frame and append a later valid one.""" + with torch.device("cpu"): + sys_stat = { + "coord": [torch.zeros((1, 6), dtype=torch.float32)], + "atype": [torch.tensor([[0, -1]], dtype=torch.int32)], + "box": [torch.eye(3, dtype=torch.float32).reshape(1, 9)], + "real_natoms_vec": [torch.tensor([[2, 2, 1, 0]], dtype=torch.int32)], + } + + _append_missing_type_frames( + sys_stat, + _FakeMixedDataset((1.0, 0.1, 0.8)), + filter_enabled=True, + ) + + self.assertEqual(len(sys_stat["real_natoms_vec"]), 2) + torch.testing.assert_close(sys_stat["coord"][-1], torch.full((1, 6), 2.0)) + + def test_append_missing_type_warns_when_all_frames_are_too_close(self) -> None: + """Leave a type uncovered when no representative passes the threshold.""" + with torch.device("cpu"): + sys_stat = { + "coord": [torch.zeros((1, 6), dtype=torch.float32)], + "atype": [torch.tensor([[0, -1]], dtype=torch.int32)], + "box": [torch.eye(3, dtype=torch.float32).reshape(1, 9)], + "real_natoms_vec": [torch.tensor([[2, 2, 1, 0]], dtype=torch.int32)], + } + + with self.assertLogs("deepmd.pt.utils.stat", level="WARNING"): + _append_missing_type_frames( + sys_stat, + _FakeMixedDataset((1.0, 0.1, 0.2)), + filter_enabled=True, + ) + + self.assertEqual(len(sys_stat["real_natoms_vec"]), 1) + + class TestObservedTypeStatFile(unittest.TestCase): """Test stat file save/load round-trip for observed_type."""