diff --git a/deepmd/main.py b/deepmd/main.py index 1206b2fa0e..d0b45e3260 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -668,6 +668,15 @@ def main_parser() -> argparse.ArgumentParser: default=None, help="The training script of the input frozen model", ) + parser_compress.add_argument( + "--recompute-min-nbor-dist", + action="store_true", + help="(Supported backend: PyTorch Exportable) Ignore the minimal neighbor " + "distance saved in the model and recompute it from the training data. " + "Requires -t,--training-script. Only applies to models whose compression " + "consumes the statistic; it is ignored for a model that tabulates over an " + "analytically bounded domain", + ) parser_compress.add_argument( "--head", "--model-branch", diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 9aeae29ca3..016489aa00 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/deepmd/pt_expt/entrypoints/compress.py @@ -2,6 +2,9 @@ """Compress a pt_expt model (.pte) by tabulating embedding nets.""" import logging +from typing import ( + Any, +) from deepmd.pt_expt.utils.serialization import ( deserialize_to_file, @@ -11,6 +14,67 @@ log = logging.getLogger(__name__) +def _read_saved_min_nbor_dist(model: Any, model_dict: dict) -> tuple[float | None, str]: + """Read the stored minimal neighbor distance and where it was read from. + + ``@variables`` is the cross-backend location of this value: it is written + by :mod:`deepmd.pt.utils.serialization` and read back by the PyTorch and + Paddle backends, so a ``.pt2`` produced by ``dp convert-backend`` carries + the value there rather than inside the serialized model dict. + + Returns + ------- + float or None + The stored minimal neighbor distance, None if the model has none. + str + Human-readable description of where the value was read from. + """ + min_nbor_dist = model.get_min_nbor_dist() + if min_nbor_dist is not None: + return float(min_nbor_dist), "the model" + min_nbor_dist = model_dict.get("min_nbor_dist") + if min_nbor_dist is not None: + return float(min_nbor_dist), "the model file" + min_nbor_dist = (model_dict.get("@variables") or {}).get("min_nbor_dist") + if min_nbor_dist is not None: + return float(min_nbor_dist), "the model file (@variables)" + return None, "" + + +def _measure_min_nbor_dist(training_script: str) -> float: + """Measure the minimal neighbor distance from the training data. + + This is a dense all-pairs pass over every frame, so callers run it only + when the tabulation actually consumes the statistic. + + Returns + ------- + float + The shortest neighbor distance in the training data. + """ + from deepmd.common import ( + j_loader, + ) + from deepmd.pt_expt.utils.update_sel import ( + UpdateSel, + ) + from deepmd.utils.compat import ( + update_deepmd_input, + ) + from deepmd.utils.data_system import ( + get_data, + ) + + jdata = update_deepmd_input(j_loader(training_script)) + train_data = get_data( + jdata["training"]["training_data"], + 0, + jdata["model"].get("type_map", None), + None, + ) + return UpdateSel().get_min_nbor_dist(train_data) + + def enable_compression( input_file: str, output: str, @@ -18,6 +82,7 @@ def enable_compression( extrapolate: int = 5, check_frequency: int = -1, training_script: str | None = None, + recompute_min_nbor_dist: bool = False, ) -> None: """Compress a .pte model by tabulating embedding nets. @@ -36,6 +101,12 @@ def enable_compression( training_script : str or None Path to training script, used to compute min_nbor_dist if not stored in the model. + recompute_min_nbor_dist : bool + Ignore the min_nbor_dist stored in the model and recompute it from + the training data. Requires training_script. Only applies to models + whose compression consumes the statistic; it is ignored, with a + warning, for a model that tabulates over an analytically bounded + domain. """ from deepmd.pt_expt.model.model import ( BaseModel, @@ -45,55 +116,59 @@ def enable_compression( model_dict = serialize_from_file(input_file) model = BaseModel.deserialize(model_dict["model"]) - # 2. Get or compute min_nbor_dist. Measuring it is a dense all-pairs pass - # over the training data, so it is only run for models that tabulate - # from the shortest observed distance. - min_nbor_dist = model.get_min_nbor_dist() - if min_nbor_dist is None: - min_nbor_dist = model_dict.get("min_nbor_dist") - if min_nbor_dist is None and not model.compression_needs_min_nbor_dist(): - log.info( - "The model tabulates over an analytically bounded domain; " - "skipping the neighbor statistics." - ) - min_nbor_dist = 0.0 - if min_nbor_dist is None: - log.info( - "Minimal neighbor distance is not saved in the model, " - "compute it from the training data." - ) - if training_script is None: - raise ValueError( - "The model does not have a minimum neighbor distance, " - "so the training script and data must be provided " - "(via -t,--training-script)." + # 2. Resolve min_nbor_dist, branching on whether compression consumes it. + # Measuring it is a dense all-pairs pass over the training data, so a + # model whose tabulation domain is fixed analytically never scans. + saved_min_nbor_dist, source = _read_saved_min_nbor_dist(model, model_dict) + if model.compression_needs_min_nbor_dist(): + if recompute_min_nbor_dist: + log.info( + "Recompute the minimal neighbor distance from the training data, " + "ignoring the one saved in the model." ) - from deepmd.common import ( - j_loader, - ) - from deepmd.pt_expt.utils.update_sel import ( - UpdateSel, - ) - from deepmd.utils.compat import ( - update_deepmd_input, - ) - from deepmd.utils.data_system import ( - get_data, - ) - - jdata = j_loader(training_script) - jdata = update_deepmd_input(jdata) - type_map = jdata["model"].get("type_map", None) - train_data = get_data( - jdata["training"]["training_data"], - 0, - type_map, - None, - ) - update_sel = UpdateSel() - min_nbor_dist = update_sel.get_min_nbor_dist(train_data) + if training_script is None: + raise ValueError( + "Recomputing the minimal neighbor distance requires the " + "training script and data (via -t,--training-script)." + ) + min_nbor_dist = _measure_min_nbor_dist(training_script) + elif saved_min_nbor_dist is not None: + min_nbor_dist = saved_min_nbor_dist + log.info(f"Minimal neighbor distance read from {source}: {min_nbor_dist:f}") + else: + log.info( + "Minimal neighbor distance is not saved in the model, " + "compute it from the training data." + ) + if training_script is None: + raise ValueError( + "The model does not have a minimum neighbor distance, " + "so the training script and data must be provided " + "(via -t,--training-script)." + ) + min_nbor_dist = _measure_min_nbor_dist(training_script) + else: + if recompute_min_nbor_dist: + log.warning( + "--recompute-min-nbor-dist does not apply to this model: its " + "tabulation domain is fixed analytically, so no neighbor " + "statistics are measured. Ignoring the flag." + ) + else: + log.info( + "The model tabulates over an analytically bounded domain; " + "skipping the neighbor statistics." + ) + # Carry a measured distance through untouched; leave it absent when + # the model has none, rather than inventing one. + min_nbor_dist = saved_min_nbor_dist + if min_nbor_dist is not None: + log.info(f"Minimal neighbor distance read from {source}: {min_nbor_dist:f}") - model.min_nbor_dist = min_nbor_dist + # ``enable_compression`` reads the value off the model, so a descriptor + # that ignores it still needs a number there. That placeholder stays local + # to the call: afterwards the model carries a measured distance or none. + model.min_nbor_dist = 0.0 if min_nbor_dist is None else min_nbor_dist # 3. Enable compression (also ensures fake ops are registered now that # the C++ custom op library is loaded via enable_compression imports) @@ -110,6 +185,7 @@ def enable_compression( stride * 10, check_frequency, ) + model.min_nbor_dist = min_nbor_dist # 4. Serialize the compressed model dict (includes tabulated data) compressed_model_dict = model.serialize() @@ -130,11 +206,18 @@ def enable_compression( ) model_def_script = model_dict.get("model_def_script") + # Only a measured distance is metadata; a model that never measured one + # leaves the field out rather than carrying a placeholder. + measured = {} if min_nbor_dist is None else {"min_nbor_dist": float(min_nbor_dist)} if output.endswith(".pt2") and model_uses_graph_lower(model): log.info("Re-exporting compressed graph...") deserialize_to_file( output, - {"model": compressed_model_dict, "model_def_script": model_def_script}, + { + "model": compressed_model_dict, + "model_def_script": model_def_script, + **measured, + }, lower_kind="auto", ) else: @@ -145,7 +228,7 @@ def enable_compression( model_json_override={ "model": compressed_model_dict, "model_def_script": model_def_script, - "min_nbor_dist": float(min_nbor_dist), + **measured, }, ) log.info("Compressed model saved to %s", output) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index e4ff2f813d..64294d86a5 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -970,6 +970,7 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None: extrapolate=FLAGS.extrapolate, check_frequency=FLAGS.frequency, training_script=FLAGS.training_script, + recompute_min_nbor_dist=FLAGS.recompute_min_nbor_dist, ) else: raise RuntimeError( diff --git a/deepmd/pt_expt/model/spin_model.py b/deepmd/pt_expt/model/spin_model.py index 83cc8ecba4..b4e9a8bcd0 100644 --- a/deepmd/pt_expt/model/spin_model.py +++ b/deepmd/pt_expt/model/spin_model.py @@ -48,6 +48,20 @@ def __getattr__(self, name: str) -> Any: return getattr(backbone, name) raise AttributeError(name) + @property + def min_nbor_dist(self) -> float | None: + """Minimal neighbor distance, stored on the backbone model. + + ``__getattr__`` only delegates reads, so without this property an + assignment would land on the wrapper while ``get_min_nbor_dist`` and + ``enable_compression`` keep reading the backbone. + """ + return self.backbone_model.min_nbor_dist + + @min_nbor_dist.setter + def min_nbor_dist(self, value: float | None) -> None: + self.backbone_model.min_nbor_dist = value + def forward_common_lower_exportable( self, extended_coord: torch.Tensor, diff --git a/deepmd/pt_expt/utils/auto_batch_size.py b/deepmd/pt_expt/utils/auto_batch_size.py new file mode 100644 index 0000000000..f1b540d302 --- /dev/null +++ b/deepmd/pt_expt/utils/auto_batch_size.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from deepmd.pt.utils.auto_batch_size import AutoBatchSize as AutoBatchSizeBase +from deepmd.pt_expt.utils import ( + env, +) + + +class AutoBatchSize(AutoBatchSizeBase): + """Auto batch size following the device pt_expt runs on. + + ``env.DEVICE`` is CPU whenever ``DEVICE=cpu`` is set, even on a CUDA host. + Growing the batch there risks a host OOM, which the CUDA-OOM handler + cannot recover from, so the growth policy follows the selected device + rather than CUDA availability. The module is held rather than the value + so that the check reads the device in effect at call time. + """ + + def is_gpu_available(self) -> bool: + """Check if the selected device is a GPU. + + Returns + ------- + bool + True if pt_expt runs on a CUDA device + """ + return env.DEVICE.type == "cuda" diff --git a/deepmd/pt_expt/utils/neighbor_stat.py b/deepmd/pt_expt/utils/neighbor_stat.py index 850bf86def..2b09ae6a02 100644 --- a/deepmd/pt_expt/utils/neighbor_stat.py +++ b/deepmd/pt_expt/utils/neighbor_stat.py @@ -10,6 +10,9 @@ from deepmd.pt_expt.common import ( torch_module, ) +from deepmd.pt_expt.utils.auto_batch_size import ( + AutoBatchSize, +) from deepmd.pt_expt.utils.env import ( DEVICE, GLOBAL_PT_FLOAT_PRECISION, @@ -28,6 +31,12 @@ class NeighborStatOP(NeighborStatOPDP): class NeighborStat(BaseNeighborStat): """Neighbor statistics using torch on DEVICE. + The statistics are evaluated one frame batch at a time. The intermediate + tensor is of shape ``[nframes, nloc, nall, 3]``, so processing a whole set + at once needs hundreds of GiB for a large set; :class:`AutoBatchSize` + keeps the batch within the available device memory, as the PyTorch, + Paddle, JAX and TensorFlow backends already do. + Parameters ---------- ntypes : int @@ -46,6 +55,7 @@ def __init__( ) -> None: super().__init__(ntypes, rcut, mixed_type) self.op = NeighborStatOP(ntypes, rcut, mixed_type) + self.auto_batch_size = AutoBatchSize() def iterator( self, data: DeepmdDataSystem @@ -65,7 +75,10 @@ def iterator( for jj in data.data_systems[ii].dirs: data_set = data.data_systems[ii] data_set_data = data_set._load_set(jj) - minrr2, max_nnei = self._execute( + minrr2, max_nnei = self.auto_batch_size.execute_all( + self._execute, + data_set_data["coord"].shape[0], + data_set.get_natoms(), data_set_data["coord"], data_set_data["type"], data_set_data["box"] if data_set.pbc else None, diff --git a/source/tests/consistent/test_neighbor_stat.py b/source/tests/consistent/test_neighbor_stat.py index 9c9f97045b..34c6ca0e22 100644 --- a/source/tests/consistent/test_neighbor_stat.py +++ b/source/tests/consistent/test_neighbor_stat.py @@ -16,6 +16,7 @@ INSTALLED_JAX, INSTALLED_PD, INSTALLED_PT, + INSTALLED_PT_EXPT, INSTALLED_TF, ) @@ -82,6 +83,10 @@ def test_neighbor_stat_tf(self) -> None: def test_neighbor_stat_pt(self) -> None: self.run_neighbor_stat("pytorch") + @unittest.skipUnless(INSTALLED_PT_EXPT, "PyTorch Exportable is not installed") + def test_neighbor_stat_pt_expt(self) -> None: + self.run_neighbor_stat("pt-expt") + def test_neighbor_stat_dp(self) -> None: self.run_neighbor_stat("numpy") diff --git a/source/tests/pt_expt/model/test_model_compression.py b/source/tests/pt_expt/model/test_model_compression.py index 95a66b528d..2c84e54049 100644 --- a/source/tests/pt_expt/model/test_model_compression.py +++ b/source/tests/pt_expt/model/test_model_compression.py @@ -187,6 +187,35 @@ def test_min_nbor_dist_roundtrip(self) -> None: finally: os.unlink(frozen_path) + def test_compress_min_nbor_dist_from_variables(self) -> None: + """Test that compress recovers min_nbor_dist from @variables and keeps it. + + ``dp convert-backend`` stores the value under ``@variables``. Compress + must read it from there — without a training script it would otherwise + raise — and carry it into the compressed archive. + """ + from deepmd.pt_expt.entrypoints.compress import ( + enable_compression, + ) + + md = self._make_model() + md.eval() + + model_data = {"model": md.serialize(), "@variables": {"min_nbor_dist": 0.5}} + with tempfile.NamedTemporaryFile(suffix=".pte", delete=False) as f: + frozen_path = f.name + with tempfile.NamedTemporaryFile(suffix=".pte", delete=False) as f: + compressed_path = f.name + try: + deserialize_to_file(frozen_path, model_data) + enable_compression(input_file=frozen_path, output=compressed_path) + compressed_data = serialize_from_file(compressed_path) + self.assertAlmostEqual(compressed_data["min_nbor_dist"], 0.5) + finally: + os.unlink(frozen_path) + if os.path.exists(compressed_path): + os.unlink(compressed_path) + def test_compress_state_serialized(self) -> None: """Test that compression state persists through serialize/deserialize. diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py new file mode 100644 index 0000000000..b7a9dd4e74 --- /dev/null +++ b/source/tests/pt_expt/test_compress_min_nbor_dist.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for reading the stored minimal neighbor distance in pt_expt compress.""" + +import logging +from typing import ( + Any, +) + +import pytest + +from deepmd.main import ( + parse_args, +) +from deepmd.pt_expt.entrypoints import compress as compress_mod + + +class _FakeModel: + """Stand-in exposing only what the compression entry point calls.""" + + def __init__( + self, + min_nbor_dist: float | None, + *, + needs_min_nbor_dist: bool = True, + ) -> None: + self._min_nbor_dist = min_nbor_dist + self._needs_min_nbor_dist = needs_min_nbor_dist + + def get_min_nbor_dist(self) -> float | None: + return self._min_nbor_dist + + def compression_needs_min_nbor_dist(self) -> bool: + return self._needs_min_nbor_dist + + +@pytest.mark.parametrize( + ("model_min_nbor_dist", "model_dict", "expected"), + [ + # the model buffer wins over both metadata locations + ( + 1.5, + {"min_nbor_dist": 9.0, "@variables": {"min_nbor_dist": 8.0}}, + (1.5, "the model"), + ), + # the top-level key, written by pt_expt's own compress + ( + None, + {"min_nbor_dist": 2.5, "@variables": {"min_nbor_dist": 8.0}}, + (2.5, "the model file"), + ), + # "@variables", the cross-backend location that dp convert-backend writes + ( + None, + {"@variables": {"min_nbor_dist": 3.5}}, + (3.5, "the model file (@variables)"), + ), + # nothing stored anywhere + (None, {}, (None, "")), + (None, {"@variables": {}}, (None, "")), + (None, {"@variables": None}, (None, "")), + ], +) +def test_read_saved_min_nbor_dist( + model_min_nbor_dist: float | None, + model_dict: dict[str, Any], + expected: tuple[float | None, str], +) -> None: + """Every storage location is honored, in order of precedence.""" + got = compress_mod._read_saved_min_nbor_dist( + _FakeModel(model_min_nbor_dist), model_dict + ) + assert got == expected + + +def test_recompute_min_nbor_dist_flag_defaults_to_false() -> None: + """The compress parser exposes the flag and leaves it off by default.""" + args = ["--pt-expt", "compress", "-i", "in.pt2", "-o", "out.pt2"] + assert parse_args(args).recompute_min_nbor_dist is False + assert parse_args([*args, "--recompute-min-nbor-dist"]).recompute_min_nbor_dist + + +def test_auto_batch_size_follows_the_selected_device() -> None: + """Growth is allowed only when pt_expt actually runs on a GPU. + + ``DEVICE`` is CPU whenever ``DEVICE=cpu`` is set, even on a CUDA host, + where an expanding batch would risk an unrecoverable host OOM. + """ + from deepmd.pt_expt.utils.auto_batch_size import ( + AutoBatchSize, + ) + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + assert AutoBatchSize().is_gpu_available() == (DEVICE.type == "cuda") + + +def test_auto_batch_size_ignores_cuda_when_cpu_is_selected(monkeypatch: Any) -> None: + """A CUDA host with ``DEVICE=cpu`` must still pin the CPU policy. + + This is the combination the policy exists for, and the only one where + following ``torch.cuda.is_available()`` instead of the selected device + gives a different answer. + """ + import torch + + from deepmd.pt_expt.utils import ( + env, + ) + from deepmd.pt_expt.utils.auto_batch_size import ( + AutoBatchSize, + ) + + monkeypatch.setattr(env, "DEVICE", torch.device("cpu")) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + assert AutoBatchSize(silent=True).is_gpu_available() is False + + +def test_graph_lower_output_keeps_min_nbor_dist(monkeypatch: Any) -> None: + """The graph-lower branch carries the value into the exported archive. + + ``model.serialize()`` does not include the runtime buffer, so without this + the compressed artifact loses a value that was just recovered. + """ + import deepmd.pt_expt.model.graph_lower as graph_lower_mod + import deepmd.pt_expt.model.model as model_mod + import deepmd.pt_expt.utils.tabulate_ops as tabulate_ops_mod + + class _Model(_FakeModel): + def enable_compression(self, *args: Any, **kwargs: Any) -> None: + pass + + def serialize(self) -> dict: + return {"type": "compressed"} + + captured: dict[str, Any] = {} + + monkeypatch.setattr( + compress_mod, + "serialize_from_file", + lambda _: {"model": {}, "@variables": {"min_nbor_dist": 1.25}}, + ) + monkeypatch.setattr( + model_mod.BaseModel, "deserialize", staticmethod(lambda _: _Model(None)) + ) + monkeypatch.setattr(tabulate_ops_mod, "ensure_fake_registered", lambda: None) + monkeypatch.setattr(graph_lower_mod, "model_uses_graph_lower", lambda _: True) + monkeypatch.setattr( + compress_mod, + "deserialize_to_file", + lambda output, data, **kwargs: captured.update(data=data), + ) + + compress_mod.enable_compression(input_file="in.pt2", output="out.pt2") + + assert captured["data"]["min_nbor_dist"] == 1.25 + + +def test_spin_model_min_nbor_dist_reaches_the_backbone() -> None: + """Assigning on the wrapper must reach the model that gets compressed. + + ``SpinModel.__getattr__`` only delegates reads, while both + ``get_min_nbor_dist`` and ``enable_compression`` run on the backbone. + """ + from deepmd.pt_expt.model.spin_model import ( + SpinModel, + ) + + class _Backbone: + min_nbor_dist: float | None = None + + class _Wrapper: + # the property under test, on a stand-in that needs no torch module + min_nbor_dist = SpinModel.min_nbor_dist + + def __init__(self) -> None: + self.backbone_model = _Backbone() + + wrapper = _Wrapper() + wrapper.min_nbor_dist = 1.25 + assert wrapper.backbone_model.min_nbor_dist == 1.25 + assert wrapper.min_nbor_dist == 1.25 + + +def _run_compression( + monkeypatch: Any, + *, + saved: float | None, + needs: bool, + flag: bool, + training_script: str | None = "train.json", +) -> tuple[dict, list, Any]: + """Compress a stand-in model and report what reached the archive. + + Returns + ------- + dict + The data handed to ``deserialize_to_file``. + list + One entry per training-data scan, so an empty list proves none ran. + Any + The model, to inspect what the run left on it. + """ + import deepmd.pt_expt.model.graph_lower as graph_lower_mod + import deepmd.pt_expt.model.model as model_mod + import deepmd.pt_expt.utils.tabulate_ops as tabulate_ops_mod + + class _Model(_FakeModel): + def enable_compression(self, *args: Any, **kwargs: Any) -> None: + pass + + def serialize(self) -> dict: + return {"type": "compressed"} + + model = _Model(None, needs_min_nbor_dist=needs) + model_dict: dict[str, Any] = {"model": {}} + if saved is not None: + model_dict["@variables"] = {"min_nbor_dist": saved} + + captured: dict[str, Any] = {} + scans: list[str] = [] + + monkeypatch.setattr(compress_mod, "serialize_from_file", lambda _: model_dict) + monkeypatch.setattr( + model_mod.BaseModel, "deserialize", staticmethod(lambda _: model) + ) + monkeypatch.setattr(tabulate_ops_mod, "ensure_fake_registered", lambda: None) + monkeypatch.setattr(graph_lower_mod, "model_uses_graph_lower", lambda _: True) + monkeypatch.setattr( + compress_mod, + "deserialize_to_file", + lambda output, data, **kwargs: captured.update(data=data), + ) + monkeypatch.setattr( + compress_mod, + "_measure_min_nbor_dist", + lambda script: (scans.append(script), 9.99)[1], + ) + + compress_mod.enable_compression( + input_file="in.pt2", + output="out.pt2", + training_script=training_script, + recompute_min_nbor_dist=flag, + ) + return captured["data"], scans, model + + +@pytest.mark.parametrize("flag", [False, True]) +def test_unused_statistic_keeps_a_saved_distance(monkeypatch: Any, flag: bool) -> None: + """A model that ignores the statistic must not lose a measured one. + + The recompute flag is inapplicable here, so it may not replace a real + distance with the placeholder the compression ABI is given. + """ + data, scans, model = _run_compression( + monkeypatch, saved=0.75, needs=False, flag=flag + ) + + assert data["min_nbor_dist"] == 0.75 + assert scans == [] + assert model.min_nbor_dist == 0.75 + + +@pytest.mark.parametrize("flag", [False, True]) +def test_unused_statistic_leaves_an_absent_distance_absent( + monkeypatch: Any, flag: bool +) -> None: + """Never measured means absent, not zero. + + ``enable_compression`` still needs a number, but that placeholder stays + local to the call instead of being recorded as a measurement. + """ + data, scans, model = _run_compression( + monkeypatch, saved=None, needs=False, flag=flag + ) + + assert "min_nbor_dist" not in data + assert scans == [] + assert model.min_nbor_dist is None + + +def test_unused_statistic_warns_the_flag_is_inapplicable( + monkeypatch: Any, caplog: Any +) -> None: + """The flag is ignored loudly, not silently.""" + with caplog.at_level(logging.WARNING): + _run_compression(monkeypatch, saved=0.75, needs=False, flag=True) + + assert "--recompute-min-nbor-dist does not apply" in caplog.text + + +@pytest.mark.parametrize( + ("saved", "flag", "expected", "scanned"), + [ + # a saved value is reused, so no pass over the training data + (1.25, False, 1.25, False), + # an explicit request measures again and wins over the saved value + (1.25, True, 9.99, True), + # nothing saved leaves measuring as the only option + (None, False, 9.99, True), + # an explicit request with nothing saved measures just the same + (None, True, 9.99, True), + ], +) +def test_consumed_statistic_paths( + monkeypatch: Any, + saved: float | None, + flag: bool, + expected: float, + scanned: bool, +) -> None: + """A model whose tabulation consumes the statistic is unaffected.""" + data, scans, _ = _run_compression(monkeypatch, saved=saved, needs=True, flag=flag) + + assert data["min_nbor_dist"] == expected + assert bool(scans) is scanned + + +@pytest.mark.parametrize( + ("flag", "message"), + [ + (True, "Recomputing the minimal neighbor distance requires"), + (False, "does not have a minimum neighbor distance"), + ], +) +def test_measuring_requires_a_training_script( + monkeypatch: Any, flag: bool, message: str +) -> None: + """Both routes into measuring keep their training-script validation.""" + with pytest.raises(ValueError, match=message): + _run_compression( + monkeypatch, + saved=None if not flag else 1.25, + needs=True, + flag=flag, + training_script=None, + ) diff --git a/source/tests/pt_expt/utils/test_neighbor_stat.py b/source/tests/pt_expt/utils/test_neighbor_stat.py new file mode 100644 index 0000000000..2a0af3d210 --- /dev/null +++ b/source/tests/pt_expt/utils/test_neighbor_stat.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the batched neighbor statistics of the pt_expt backend.""" + +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.pt_expt.utils.neighbor_stat import ( + NeighborStat, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) + +NATOMS = 27 +NFRAMES = 8 +SET_SIZE = 4 + +# A 3x3x3 grid spaced 1.0 apart in a 3.0 box: the shortest distance is +# exactly 1.0 and, just past it, every atom sees its six axis-aligned +# images. Used for the analytical check. +UNIFORM_SPACING = 1.0 +UNIFORM_RCUT = 1.001 + +# For the batching comparison the cutoff also admits the second shell +# (sqrt(2) * spacing), which is what makes the neighbor count depend on +# the spacing of each individual frame. +MIXED_RCUT = 1.5 +SPARSE_SPACING = 1.2 +DENSE_SPACING = 1.0 +CLOSE_PAIR_DISTANCE = 0.5 + + +def _grid(spacing: float) -> np.ndarray: + """Return a 3x3x3 grid of ``NATOMS`` positions with the given spacing.""" + axis = np.arange(3) * spacing + return np.vstack([g.ravel() for g in np.meshgrid(*([axis] * 3), indexing="ij")]).T + + +def _write_system(path: Path, coords: np.ndarray, boxes: np.ndarray) -> None: + """Write frames as a multi-set system in the DeePMD npy layout.""" + nframes = coords.shape[0] + path.mkdir(parents=True, exist_ok=True) + (path / "type.raw").write_text("\n".join(["0"] * NATOMS) + "\n") + (path / "type_map.raw").write_text("TYPE\n") + for index, start in enumerate(range(0, nframes, SET_SIZE)): + subset = path / f"set.{index:03d}" + subset.mkdir() + frames = slice(start, start + SET_SIZE) + nsub = coords[frames].shape[0] + np.save(subset / "coord.npy", coords[frames].reshape(nsub, NATOMS * 3)) + np.save(subset / "box.npy", boxes[frames]) + np.save(subset / "energy.npy", np.zeros(nsub)) + np.save(subset / "force.npy", np.zeros((nsub, NATOMS * 3))) + + +def _box(spacing: float) -> np.ndarray: + return (np.eye(3) * spacing * 3).ravel() + + +def _uniform_frames() -> tuple[np.ndarray, np.ndarray]: + """Identical frames, so the statistics are known analytically.""" + coords = np.repeat(_grid(UNIFORM_SPACING)[np.newaxis], NFRAMES, axis=0) + boxes = np.tile(_box(UNIFORM_SPACING), (NFRAMES, 1)) + return coords, boxes + + +def _mixed_frames() -> tuple[np.ndarray, np.ndarray]: + """Frames whose two extrema are contributed by different sets. + + The first set holds a frame with one pair pulled close, which owns the + global minimum distance but barely changes any neighbor count. The last + set holds a denser frame, whose second shell falls inside the cutoff and + owns the global maximum neighbor count. A reduction that keeps only one + batch therefore cannot report both. + """ + sparse = _grid(SPARSE_SPACING) + close_pair = sparse.copy() + # index 9 is (spacing, 0, 0); move atom 0 to sit CLOSE_PAIR_DISTANCE from it + close_pair[0] = sparse[9] - np.array([CLOSE_PAIR_DISTANCE, 0.0, 0.0]) + dense = _grid(DENSE_SPACING) + + coords = np.repeat(sparse[np.newaxis], NFRAMES, axis=0) + boxes = np.tile(_box(SPARSE_SPACING), (NFRAMES, 1)) + coords[0] = close_pair + coords[-1] = dense + boxes[-1] = _box(DENSE_SPACING) + return coords, boxes + + +def _stat_with_batch_size( + monkeypatch: Any, system: str, rcut: float, batch_atoms: int +) -> tuple[tuple[float, np.ndarray], list[int]]: + """Run the statistics at a pinned batch size. + + ``DP_INFER_BATCH_SIZE`` is read when :class:`AutoBatchSize` is built, so + it has to be set before the ``NeighborStat`` is constructed. + + Returns + ------- + tuple + The ``(min_nbor_dist, max_nbor_size)`` that ``get_stat`` produced. + list + The number of frames handed to each ``_execute`` call, one entry per + call, so the caller can tell how the work was split. + """ + frames_per_call: list[int] = [] + # A context keeps the patched ``_execute`` from leaking into a later run, + # which would otherwise chain the spies and mix their records. + with monkeypatch.context() as patch: + patch.setenv("DP_INFER_BATCH_SIZE", str(batch_atoms)) + original = NeighborStat._execute + + def spy( + self: NeighborStat, + coord: np.ndarray, + atype: np.ndarray, + cell: np.ndarray | None, + ) -> tuple[np.ndarray, np.ndarray]: + frames_per_call.append(coord.shape[0]) + return original(self, coord, atype, cell) + + patch.setattr(NeighborStat, "_execute", spy) + + data = DeepmdDataSystem( + systems=[system], + batch_size=1, + test_size=1, + rcut=rcut, + type_map=["TYPE"], + ) + data.get_batch() + stat = NeighborStat(data.get_ntypes(), rcut, mixed_type=False).get_stat(data) + return stat, frames_per_call + + +def _single_frame_stat( + monkeypatch: Any, path: Path, coord: np.ndarray, box: np.ndarray +) -> tuple[float, np.ndarray]: + """Statistics of one frame on its own, as a reference extremum.""" + _write_system(path, coord[np.newaxis], box[np.newaxis]) + stat, _ = _stat_with_batch_size(monkeypatch, str(path), MIXED_RCUT, NATOMS) + return stat + + +def test_batching_splits_the_work_and_keeps_the_result( + monkeypatch: Any, tmp_path: Path +) -> None: + """Batching must change how the work is split, not what it produces. + + Sending a whole set at once allocates an ``[nframes, nloc, nall, 3]`` + tensor, which is why the statistics are batched. Both the split and the + reduction across batches are asserted: the frames are built so that the + two extrema come from different sets, so keeping only one batch would + get one of them wrong. + """ + coords, boxes = _mixed_frames() + system = tmp_path / "system" + _write_system(system, coords, boxes) + + whole, whole_calls = _stat_with_batch_size( + monkeypatch, str(system), MIXED_RCUT, NATOMS * NFRAMES + ) + batched, batched_calls = _stat_with_batch_size( + monkeypatch, str(system), MIXED_RCUT, NATOMS + ) + + # one call per set when everything fits, one call per frame when it does not + assert whole_calls == [SET_SIZE] * (NFRAMES // SET_SIZE) + assert batched_calls == [1] * NFRAMES + assert len(batched_calls) > len(whole_calls) + + # the frame in the first set owns the shortest distance, the one in the + # last set the largest neighbor count + first = _single_frame_stat(monkeypatch, tmp_path / "first", coords[0], boxes[0]) + last = _single_frame_stat(monkeypatch, tmp_path / "last", coords[-1], boxes[-1]) + assert first[0] < last[0] + assert first[1].max() < last[1].max() + + for stat in (whole, batched): + assert stat[0] == first[0] + np.testing.assert_array_equal(stat[1], last[1]) + + +def test_batched_statistics_are_correct(monkeypatch: Any, tmp_path: Path) -> None: + """The reduction over batches reproduces the known geometry. + + On a unit grid the shortest distance is 1.0, and just past it each atom + sees its six axis-aligned images. + """ + coords, boxes = _uniform_frames() + system = tmp_path / "system" + _write_system(system, coords, boxes) + + (min_nbor_dist, max_nbor_size), _ = _stat_with_batch_size( + monkeypatch, str(system), UNIFORM_RCUT, NATOMS + ) + + assert min_nbor_dist == UNIFORM_SPACING + np.testing.assert_array_equal(max_nbor_size, [6])