From d681f22cb8b64812440e5062be39382f982affa2 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 12 Jul 2026 09:23:27 +0800 Subject: [PATCH 1/6] refactor(lmp): share MPI pair test driver Consolidate the common pair_deepmd and pair_deepspin MPI setup while preserving their existing script entry points and model-specific configuration. Authored by OpenClaw (model: gpt-5.6-sol) Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lmp/tests/mpi_pair_deepmd.py | 81 ++++++++++++++++++++ source/lmp/tests/run_mpi_pair_deepmd.py | 71 ++--------------- source/lmp/tests/run_mpi_pair_deepmd_spin.py | 71 ++--------------- 3 files changed, 97 insertions(+), 126 deletions(-) create mode 100644 source/lmp/tests/mpi_pair_deepmd.py diff --git a/source/lmp/tests/mpi_pair_deepmd.py b/source/lmp/tests/mpi_pair_deepmd.py new file mode 100644 index 0000000000..0b10422855 --- /dev/null +++ b/source/lmp/tests/mpi_pair_deepmd.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared MPI driver for DeepMD and DeepSpin LAMMPS pair-style tests.""" + +import argparse +from typing import ( + NamedTuple, +) + +import numpy as np +from lammps import ( + PyLammps, +) +from mpi4py import ( + MPI, +) + + +class PairStyleConfig(NamedTuple): + """LAMMPS commands that differ between the DeepMD and DeepSpin runners.""" + + atom_style: str + masses: tuple[str, str] + pair_style: str + + +def run_mpi_pair_deepmd(config: PairStyleConfig) -> None: + """Run the common two-rank model-deviation scenario. + + The public runner scripts remain separate because their model and data-file + contracts differ. Keeping those scripts as wrappers also preserves their + command-line entry points for pytest and external build tooling. + """ + parser = argparse.ArgumentParser() + parser.add_argument("DATAFILE", type=str) + parser.add_argument("PBFILE", type=str) + parser.add_argument("PBFILE2", type=str) + parser.add_argument("MD_FILE", type=str) + parser.add_argument("OUTPUT", type=str) + parser.add_argument("--balance", action="store_true") + parser.add_argument("--nopbc", action="store_true") + args = parser.parse_args() + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + lammps = PyLammps() + if args.balance: + # 4 and 2 atoms + lammps.processors("2 1 1") + else: + # 6 and 0 atoms + lammps.processors("1 2 1") + lammps.units("metal") + if args.nopbc: + lammps.boundary("f f f") + else: + lammps.boundary("p p p") + lammps.atom_style(config.atom_style) + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(args.DATAFILE) + lammps.mass(f"1 {config.masses[0]}") + lammps.mass(f"2 {config.masses[1]}") + lammps.timestep(0.0005) + lammps.fix("1 all nve") + + relative = 1.0 + lammps.pair_style( + f"{config.pair_style} {args.PBFILE} {args.PBFILE2} " + f"out_file {args.MD_FILE} out_freq 1 atomic relative {relative}" + ) + lammps.pair_coeff("* *") + lammps.run(0) + if rank == 0: + pe = lammps.eval("pe") + np.savetxt(args.OUTPUT, np.array([pe])) + + # LAMMPS owns MPI resources, so its destructor must run before finalization. + # Changing this order can make the destructor call MPI after MPI_Finalize. + del lammps + MPI.Finalize() diff --git a/source/lmp/tests/run_mpi_pair_deepmd.py b/source/lmp/tests/run_mpi_pair_deepmd.py index 21119ab0d3..3257b122dd 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd.py +++ b/source/lmp/tests/run_mpi_pair_deepmd.py @@ -1,68 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Use mpi4py to run a LAMMPS pair_deepmd + model deviation (atomic, relative) task.""" +"""Run the pair_deepmd MPI model-deviation test scenario.""" -import argparse - -import numpy as np -from lammps import ( - PyLammps, -) -from mpi4py import ( - MPI, +from mpi_pair_deepmd import ( + PairStyleConfig, + run_mpi_pair_deepmd, ) -comm = MPI.COMM_WORLD -rank = comm.Get_rank() - -parser = argparse.ArgumentParser() -parser.add_argument("DATAFILE", type=str) -parser.add_argument("PBFILE", type=str) -parser.add_argument("PBFILE2", type=str) -parser.add_argument("MD_FILE", type=str) -parser.add_argument("OUTPUT", type=str) -parser.add_argument("--balance", action="store_true") -parser.add_argument("--nopbc", action="store_true") -args = parser.parse_args() -data_file = args.DATAFILE -pb_file = args.PBFILE -pb_file2 = args.PBFILE2 -md_file = args.MD_FILE -output = args.OUTPUT -balance = args.balance - -lammps = PyLammps() -if balance: - # 4 and 2 atoms - lammps.processors("2 1 1") -else: - # 6 and 0 atoms - lammps.processors("1 2 1") -lammps.units("metal") -if args.nopbc: - lammps.boundary("f f f") -else: - lammps.boundary("p p p") -lammps.atom_style("atomic") -lammps.neighbor("2.0 bin") -lammps.neigh_modify("every 10 delay 0 check no") -lammps.read_data(data_file) -lammps.mass("1 16") -lammps.mass("2 2") -lammps.timestep(0.0005) -lammps.fix("1 all nve") - -relative = 1.0 -lammps.pair_style( - f"deepmd {pb_file} {pb_file2} out_file {md_file} out_freq 1 atomic relative {relative}" -) -lammps.pair_coeff("* *") -lammps.run(0) -if rank == 0: - pe = lammps.eval("pe") - arr = [pe] - np.savetxt(output, np.array(arr)) -# Tear down LAMMPS before MPI.Finalize() to avoid MPI-after-Finalize -# in the LAMMPS destructor. See run_mpi_pair_deepmd_spin_dpa3_pt2.py. -del lammps -MPI.Finalize() +if __name__ == "__main__": + run_mpi_pair_deepmd( + PairStyleConfig(atom_style="atomic", masses=("16", "2"), pair_style="deepmd") + ) diff --git a/source/lmp/tests/run_mpi_pair_deepmd_spin.py b/source/lmp/tests/run_mpi_pair_deepmd_spin.py index 7e87d3e5a1..c53ccaa4ed 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd_spin.py +++ b/source/lmp/tests/run_mpi_pair_deepmd_spin.py @@ -1,68 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Use mpi4py to run a LAMMPS pair_deepmd + model deviation (atomic, relative) task.""" +"""Run the pair_deepspin MPI model-deviation test scenario.""" -import argparse - -import numpy as np -from lammps import ( - PyLammps, -) -from mpi4py import ( - MPI, +from mpi_pair_deepmd import ( + PairStyleConfig, + run_mpi_pair_deepmd, ) -comm = MPI.COMM_WORLD -rank = comm.Get_rank() - -parser = argparse.ArgumentParser() -parser.add_argument("DATAFILE", type=str) -parser.add_argument("PBFILE", type=str) -parser.add_argument("PBFILE2", type=str) -parser.add_argument("MD_FILE", type=str) -parser.add_argument("OUTPUT", type=str) -parser.add_argument("--balance", action="store_true") -parser.add_argument("--nopbc", action="store_true") -args = parser.parse_args() -data_file = args.DATAFILE -pb_file = args.PBFILE -pb_file2 = args.PBFILE2 -md_file = args.MD_FILE -output = args.OUTPUT -balance = args.balance - -lammps = PyLammps() -if balance: - # 4 and 2 atoms - lammps.processors("2 1 1") -else: - # 6 and 0 atoms - lammps.processors("1 2 1") -lammps.units("metal") -if args.nopbc: - lammps.boundary("f f f") -else: - lammps.boundary("p p p") -lammps.atom_style("spin") -lammps.neighbor("2.0 bin") -lammps.neigh_modify("every 10 delay 0 check no") -lammps.read_data(data_file) -lammps.mass("1 58") -lammps.mass("2 16") -lammps.timestep(0.0005) -lammps.fix("1 all nve") - -relative = 1.0 -lammps.pair_style( - f"deepspin {pb_file} {pb_file2} out_file {md_file} out_freq 1 atomic relative {relative}" -) -lammps.pair_coeff("* *") -lammps.run(0) -if rank == 0: - pe = lammps.eval("pe") - arr = [pe] - np.savetxt(output, np.array(arr)) -# Tear down LAMMPS before MPI.Finalize() to avoid MPI-after-Finalize -# in the LAMMPS destructor. See run_mpi_pair_deepmd_spin_dpa3_pt2.py. -del lammps -MPI.Finalize() +if __name__ == "__main__": + run_mpi_pair_deepmd( + PairStyleConfig(atom_style="spin", masses=("58", "16"), pair_style="deepspin") + ) From cbe6ae83e5a951736f34b02b2afc49bc2e314946 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 12 Jul 2026 06:29:01 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/lmp/tests/run_mpi_pair_deepmd.py | 1 - source/lmp/tests/run_mpi_pair_deepmd_spin.py | 1 - 2 files changed, 2 deletions(-) diff --git a/source/lmp/tests/run_mpi_pair_deepmd.py b/source/lmp/tests/run_mpi_pair_deepmd.py index 3257b122dd..3617dcbd6c 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd.py +++ b/source/lmp/tests/run_mpi_pair_deepmd.py @@ -6,7 +6,6 @@ run_mpi_pair_deepmd, ) - if __name__ == "__main__": run_mpi_pair_deepmd( PairStyleConfig(atom_style="atomic", masses=("16", "2"), pair_style="deepmd") diff --git a/source/lmp/tests/run_mpi_pair_deepmd_spin.py b/source/lmp/tests/run_mpi_pair_deepmd_spin.py index c53ccaa4ed..ca022133d3 100644 --- a/source/lmp/tests/run_mpi_pair_deepmd_spin.py +++ b/source/lmp/tests/run_mpi_pair_deepmd_spin.py @@ -6,7 +6,6 @@ run_mpi_pair_deepmd, ) - if __name__ == "__main__": run_mpi_pair_deepmd( PairStyleConfig(atom_style="spin", masses=("58", "16"), pair_style="deepspin") From c5b2ad03282c5a32d4e502c3600818d3d9adacdb Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 12 Jul 2026 18:53:49 +0800 Subject: [PATCH 3/6] refactor(lmp): share test infrastructure Centralize atomic and spin LAMMPS construction, common water fixture lifecycle, and MPI subprocess parsing across backend and model-format tests. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lmp/tests/lammps_test_utils.py | 221 ++++++++++++++++++ source/lmp/tests/test_deeptensor.py | 34 +-- source/lmp/tests/test_lammps.py | 89 ++----- source/lmp/tests/test_lammps_3types.py | 17 +- .../lmp/tests/test_lammps_dpa1_graph_pt2.py | 18 +- source/lmp/tests/test_lammps_dpa2_pt.py | 88 ++----- source/lmp/tests/test_lammps_dpa2_pt_nopbc.py | 89 ++----- source/lmp/tests/test_lammps_dpa3_pt2.py | 102 ++------ .../lmp/tests/test_lammps_dpa3_pt2_nopbc.py | 35 +-- source/lmp/tests/test_lammps_dpa4_pt2.py | 39 +--- source/lmp/tests/test_lammps_dpa_jax.py | 90 ++----- source/lmp/tests/test_lammps_faparam.py | 32 +-- source/lmp/tests/test_lammps_faparam_pt.py | 15 +- source/lmp/tests/test_lammps_faparam_pt2.py | 15 +- .../tests/test_lammps_fparam_from_fix_dedn.py | 14 +- source/lmp/tests/test_lammps_jax.py | 88 ++----- .../lmp/tests/test_lammps_model_devi_pt2.py | 18 +- source/lmp/tests/test_lammps_pd.py | 88 ++----- source/lmp/tests/test_lammps_pt.py | 88 ++----- source/lmp/tests/test_lammps_pt2.py | 34 +-- source/lmp/tests/test_lammps_spin.py | 55 ++--- source/lmp/tests/test_lammps_spin_dpa3_pt2.py | 72 ++---- source/lmp/tests/test_lammps_spin_nopbc.py | 56 ++--- source/lmp/tests/test_lammps_spin_nopbc_pt.py | 55 ++--- .../lmp/tests/test_lammps_spin_nopbc_pt2.py | 25 +- source/lmp/tests/test_lammps_spin_pt.py | 55 ++--- source/lmp/tests/test_lammps_spin_pt2.py | 25 +- 27 files changed, 494 insertions(+), 1063 deletions(-) create mode 100644 source/lmp/tests/lammps_test_utils.py diff --git a/source/lmp/tests/lammps_test_utils.py b/source/lmp/tests/lammps_test_utils.py new file mode 100644 index 0000000000..983b106383 --- /dev/null +++ b/source/lmp/tests/lammps_test_utils.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared infrastructure for the LAMMPS Python tests. + +Model paths, expected values, and scenario-specific assertions stay in their +test modules. This module owns setup that must remain identical across model +formats and backends, so changes to the LAMMPS test system have one source of +truth. +""" + +from __future__ import ( + annotations, +) + +import os +import subprocess as sp +import sys +import tempfile +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import constants +import numpy as np +import pytest +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data, +) + + +def require_backend(environment_variable: str, backend_name: str) -> None: + """Skip the current module when its compiled backend is unavailable.""" + if os.environ.get(environment_variable, "1") != "1": + pytest.skip(f"Skip test because {backend_name} support is not enabled.") + + +def remove_test_files(*paths: Path) -> None: + """Remove generated test files, tolerating partial setup and prior cleanup.""" + for path in paths: + path.unlink(missing_ok=True) + + +def write_water_data_variants( + box: np.ndarray, + coord: np.ndarray, + type_oh: np.ndarray, + type_ho: np.ndarray, + data_file: Path, + type_map_file: Path, + si_file: Path, +) -> None: + """Write the standard metal, type-map, and SI water test fixtures.""" + write_lmp_data(box, coord, type_oh, data_file) + write_lmp_data(box, coord, type_ho, type_map_file) + write_lmp_data( + box * constants.dist_metal2si, + coord * constants.dist_metal2si, + type_oh, + si_file, + ) + + +def make_atomic_lammps( + data_file: Path, + units: str = "metal", + *, + boundary: str = "p p p", + atom_map: str | None = None, + masses: tuple[float, ...] = (16, 2), +) -> PyLammps: + """Create the standard two-type atomic LAMMPS test system. + + ``atom_map="no"`` deliberately omits ``atom_modify`` because LAMMPS + rejects ``atom_modify map no``; this preserves the no-map failure-path + tests used by the graph-model fixtures. + """ + if units not in {"metal", "real", "si"}: + raise ValueError("units should be metal, real, or si") + + lammps = PyLammps() + lammps.units(units) + lammps.boundary(boundary) + lammps.atom_style("atomic") + if atom_map is not None and atom_map != "no": + lammps.atom_modify(f"map {atom_map}") + lammps.neighbor("2.0e-10 bin" if units == "si" else "2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + for atom_type, mass in enumerate(masses, start=1): + if units == "si": + lammps.mass(f"{atom_type} {mass * constants.mass_metal2si:.10e}") + else: + lammps.mass(f"{atom_type} {mass:g}") + lammps.timestep({"metal": 0.0005, "real": 0.5, "si": 5e-16}[units]) + lammps.fix("1 all nve") + return lammps + + +def make_spin_lammps( + data_file: Path, + units: str = "metal", + *, + boundary: str = "p p p", +) -> PyLammps: + """Create the standard two-type DeepSpin LAMMPS test system.""" + if units != "metal": + raise ValueError("units for spin should be metal") + + lammps = PyLammps() + lammps.units(units) + lammps.boundary(boundary) + lammps.atom_style("spin") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 58") + lammps.mass("2 16") + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +def run_mpi_pair_runner( + runner: Path, + data_file: Path, + model_file: Path, + *, + nprocs: int = 2, + processors: str | None = None, + extra_args: list[str] | None = None, + runner_args: list[str] | None = None, + output_columns: tuple[tuple[str, int], ...] = (("forces", 3), ("virials", 9)), + capture: bool = False, +) -> dict[str, Any]: + """Invoke a DPA MPI runner and parse its energy/per-atom output. + + The runner output contract is one energy line followed by a rectangular + per-atom table. ``output_columns`` names and slices that table while each + model-specific wrapper retains its own defaults and explanatory docstring. + """ + with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: + output_path = Path(f.name) + try: + argv = [ + "mpirun", + "-n", + str(nprocs), + sys.executable, + str(runner), + str(data_file.resolve()), + str(model_file.resolve()), + str(output_path), + ] + if processors is not None: + argv.extend(["--processors", processors]) + elif nprocs == 1: + argv.extend(["--processors", "1 1 1"]) + if extra_args: + argv.extend(extra_args) + if runner_args: + argv.extend(runner_args) + if capture: + proc = sp.run(argv, capture_output=True, text=True) + return { + "returncode": proc.returncode, + "stdout": proc.stdout, + "stderr": proc.stderr, + } + + sp.check_call(argv) + lines = output_path.read_text().strip().splitlines() + rows = np.array( + [list(map(float, line.split())) for line in lines[1:]], + dtype=np.float64, + ) + result: dict[str, Any] = {"pe": float(lines[0])} + start = 0 + for name, width in output_columns: + result[name] = rows[:, start : start + width] + start += width + if rows.shape[1] != start: + raise ValueError( + f"MPI runner produced {rows.shape[1]} columns; expected {start}" + ) + return result + finally: + output_path.unlink(missing_ok=True) + + +def run_mpi_model_deviation( + runner: Path, + data_file: Path, + model_file: Path, + second_model_file: Path, + deviation_file: Path, + *, + extra_args: list[str] | None = None, +) -> float: + """Run the two-rank model-deviation driver and return rank-zero energy.""" + with tempfile.NamedTemporaryFile() as output: + argv = [ + "mpirun", + "-n", + "2", + sys.executable, + str(runner), + str(data_file), + str(model_file), + str(second_model_file), + str(deviation_file), + output.name, + ] + if extra_args: + argv.extend(extra_args) + sp.check_call(argv) + return float(np.loadtxt(output.name, ndmin=1)[0]) diff --git a/source/lmp/tests/test_deeptensor.py b/source/lmp/tests/test_deeptensor.py index c0f9a2d7f7..37eabf8e41 100644 --- a/source/lmp/tests/test_deeptensor.py +++ b/source/lmp/tests/test_deeptensor.py @@ -10,6 +10,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from model_convert import ( ensure_converted_pb, ) @@ -83,36 +86,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units) @pytest.fixture diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py index eca782365e..6ec1fedb3f 100644 --- a/source/lmp/tests/test_lammps.py +++ b/source/lmp/tests/test_lammps.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import importlib -import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -15,12 +11,16 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, + remove_test_files, + require_backend, + run_mpi_model_deviation, + write_water_data_variants, +) from model_convert import ( ensure_converted_pb, ) -from write_lmp_data import ( - write_lmp_data, -) pbtxt_file = Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot.pbtxt" pbtxt_file2 = ( @@ -225,59 +225,20 @@ def setup_module() -> None: - if os.environ.get("ENABLE_TENSORFLOW", "1") != "1": - pytest.skip( - "Skip test because TensorFlow support is not enabled.", - ) + require_backend("ENABLE_TENSORFLOW", "TensorFlow") ensure_converted_pb(pbtxt_file, pb_file) ensure_converted_pb(pbtxt_file2, pb_file2) - - write_lmp_data(box, coord, type_OH, data_file) - write_lmp_data(box, coord, type_HO, data_type_map_file) - write_lmp_data( - box * constants.dist_metal2si, - coord * constants.dist_metal2si, - type_OH, - data_file_si, + write_water_data_variants( + box, coord, type_OH, type_HO, data_file, data_type_map_file, data_file_si ) def teardown_module() -> None: - os.remove(data_file) - os.remove(data_type_map_file) + remove_test_files(data_file, data_type_map_file, data_file_si, md_file) def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units) @pytest.fixture @@ -732,24 +693,14 @@ def test_pair_deepmd_si(lammps_si) -> None: [(["--balance"],), ([],)], ) def test_pair_deepmd_mpi(balance_args: list) -> None: - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_3types.py b/source/lmp/tests/test_lammps_3types.py index a3cef29a62..6354388314 100644 --- a/source/lmp/tests/test_lammps_3types.py +++ b/source/lmp/tests/test_lammps_3types.py @@ -9,6 +9,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from model_convert import ( ensure_converted_pb, ) @@ -264,19 +267,7 @@ def teardown_module() -> None: def _lammps(data_file) -> PyLammps: - lammps = PyLammps() - lammps.units("metal") - lammps.boundary("p p p") - lammps.atom_style("atomic") - lammps.neighbor("2.0 bin") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - lammps.mass("1 16") - lammps.mass("2 2") - lammps.mass("3 16") - lammps.timestep(0.0005) - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, masses=(16, 2, 16)) @pytest.fixture diff --git a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py index ab898e15f6..f3f3e0bf59 100644 --- a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py +++ b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py @@ -37,6 +37,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from write_lmp_data import ( write_lmp_data, ) @@ -112,20 +115,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal", atom_map: str = "yes") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if atom_map != "no": - lammps.atom_modify(f"map {atom_map}") - lammps.neighbor("2.0 bin") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - lammps.mass("1 16") - lammps.mass("2 2") - lammps.timestep(0.0005) - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, atom_map=atom_map) @pytest.fixture diff --git a/source/lmp/tests/test_lammps_dpa2_pt.py b/source/lmp/tests/test_lammps_dpa2_pt.py index 860a62b631..29a5f8c917 100644 --- a/source/lmp/tests/test_lammps_dpa2_pt.py +++ b/source/lmp/tests/test_lammps_dpa2_pt.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -18,12 +15,16 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, + remove_test_files, + require_backend, + run_mpi_model_deviation, + write_water_data_variants, +) from model_convert import ( ensure_converted_pb, ) -from write_lmp_data import ( - write_lmp_data, -) pbtxt_file2 = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot-1.pbtxt" @@ -152,59 +153,20 @@ def setup_module() -> None: - if os.environ.get("ENABLE_PYTORCH", "1") != "1": - pytest.skip( - "Skip test because PyTorch support is not enabled.", - ) + require_backend("ENABLE_PYTORCH", "PyTorch") if os.environ.get("ENABLE_TENSORFLOW", "1") == "1": ensure_converted_pb(pbtxt_file2, pb_file2) - - write_lmp_data(box, coord, type_OH, data_file) - write_lmp_data(box, coord, type_HO, data_type_map_file) - write_lmp_data( - box * constants.dist_metal2si, - coord * constants.dist_metal2si, - type_OH, - data_file_si, + write_water_data_variants( + box, coord, type_OH, type_HO, data_file, data_type_map_file, data_file_si ) def teardown_module() -> None: - os.remove(data_file) - os.remove(data_type_map_file) + remove_test_files(data_file, data_type_map_file, data_file_si, md_file) def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units) @pytest.fixture @@ -662,24 +624,14 @@ def test_pair_deepmd_si(lammps_si) -> None: reason="Skip test because TensorFlow support is not enabled.", ) def test_pair_deepmd_mpi(balance_args: list) -> None: - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_dpa2_pt_nopbc.py b/source/lmp/tests/test_lammps_dpa2_pt_nopbc.py index 3b8b60ea46..4f30227418 100644 --- a/source/lmp/tests/test_lammps_dpa2_pt_nopbc.py +++ b/source/lmp/tests/test_lammps_dpa2_pt_nopbc.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -18,12 +15,16 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, + remove_test_files, + require_backend, + run_mpi_model_deviation, + write_water_data_variants, +) from model_convert import ( ensure_converted_pb, ) -from write_lmp_data import ( - write_lmp_data, -) pbtxt_file2 = Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot.pbtxt" pb_file = Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa2.pth" @@ -148,59 +149,20 @@ def setup_module() -> None: - if os.environ.get("ENABLE_PYTORCH", "1") != "1": - pytest.skip( - "Skip test because PyTorch support is not enabled.", - ) + require_backend("ENABLE_PYTORCH", "PyTorch") if os.environ.get("ENABLE_TENSORFLOW", "1") == "1": ensure_converted_pb(pbtxt_file2, pb_file2) - - write_lmp_data(box, coord, type_OH, data_file) - write_lmp_data(box, coord, type_HO, data_type_map_file) - write_lmp_data( - box * constants.dist_metal2si, - coord * constants.dist_metal2si, - type_OH, - data_file_si, + write_water_data_variants( + box, coord, type_OH, type_HO, data_file, data_type_map_file, data_file_si ) def teardown_module() -> None: - os.remove(data_file) - os.remove(data_type_map_file) + remove_test_files(data_file, data_type_map_file, data_file_si, md_file) def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("f f f") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, boundary="f f f") @pytest.fixture @@ -648,25 +610,14 @@ def test_pair_deepmd_si(lammps_si) -> None: reason="Skip test because TensorFlow support is not enabled.", ) def test_pair_deepmd_mpi(balance_args: list) -> None: - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - "--nopbc", - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=[*balance_args, "--nopbc"], + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_dpa3_pt2.py b/source/lmp/tests/test_lammps_dpa3_pt2.py index a55d80561c..7445965898 100644 --- a/source/lmp/tests/test_lammps_dpa3_pt2.py +++ b/source/lmp/tests/test_lammps_dpa3_pt2.py @@ -8,9 +8,6 @@ import importlib.util import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -24,6 +21,10 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, + run_mpi_pair_runner, +) from write_lmp_data import ( write_lmp_data, ) @@ -181,41 +182,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal", atom_map: str = "yes") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - # LAMMPS rejects ``atom_modify map no``; the supported way to leave - # the atom-map disabled is to simply omit the command (default for - # ``atom_style atomic``). - if atom_map != "no": - lammps.atom_modify(f"map {atom_map}") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, atom_map=atom_map) @pytest.fixture @@ -439,55 +406,16 @@ def _run_mpi_subprocess( (``"2 1 1"``); used by the ``test_*_decomposition`` variants to exercise 2D / 3D processor grids (Px*Py*Pz must equal nprocs). """ - if data_path is None: - data_path = data_file - if pb_path is None: - pb_path = pb_file_mpi - with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: - out_path = f.name - try: - argv = [ - "mpirun", - "-n", - str(nprocs), - sys.executable, - str(Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py"), - str(data_path.resolve()), - str(pb_path.resolve()), - out_path, - ] - if processors is not None: - argv.extend(["--processors", processors]) - elif nprocs == 1: - argv.extend(["--processors", "1 1 1"]) - if extra_args: - argv.extend(extra_args) - if runner_args: - argv.extend(runner_args) - if capture: - # Return raw process info instead of parsing output — used by - # tests that expect the subprocess to fail (the fail-fast cases). - proc = sp.run(argv, capture_output=True, text=True) - return { - "returncode": proc.returncode, - "stdout": proc.stdout, - "stderr": proc.stderr, - } - sp.check_call(argv) - with open(out_path) as fh: - lines = fh.read().strip().splitlines() - pe = float(lines[0]) - rows = np.array( - [list(map(float, line.split())) for line in lines[1:]], - dtype=np.float64, - ) - # Each row is (3 force) + (9 virial); see runner script. - forces = rows[:, :3] - virials = rows[:, 3:] - return {"pe": pe, "forces": forces, "virials": virials} - finally: - if os.path.exists(out_path): - os.remove(out_path) + return run_mpi_pair_runner( + Path(__file__).parent / "run_mpi_pair_deepmd_dpa3_pt2.py", + data_path or data_file, + pb_path or pb_file_mpi, + nprocs=nprocs, + processors=processors, + extra_args=extra_args, + runner_args=runner_args, + capture=capture, + ) @pytest.mark.skipif( diff --git a/source/lmp/tests/test_lammps_dpa3_pt2_nopbc.py b/source/lmp/tests/test_lammps_dpa3_pt2_nopbc.py index e856b7aaf4..35c2f5483b 100644 --- a/source/lmp/tests/test_lammps_dpa3_pt2_nopbc.py +++ b/source/lmp/tests/test_lammps_dpa3_pt2_nopbc.py @@ -19,6 +19,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from write_lmp_data import ( write_lmp_data, ) @@ -81,37 +84,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("f f f") - lammps.atom_style("atomic") - lammps.atom_modify("map yes") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, boundary="f f f", atom_map="yes") @pytest.fixture diff --git a/source/lmp/tests/test_lammps_dpa4_pt2.py b/source/lmp/tests/test_lammps_dpa4_pt2.py index d68680a996..bf4fcf2940 100644 --- a/source/lmp/tests/test_lammps_dpa4_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_pt2.py @@ -52,6 +52,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from write_lmp_data import ( write_lmp_data, ) @@ -117,41 +120,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal", atom_map: str = "yes") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - # LAMMPS rejects ``atom_modify map no``; the supported way to leave - # the atom-map disabled is to simply omit the command (default for - # ``atom_style atomic``). - if atom_map != "no": - lammps.atom_modify(f"map {atom_map}") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, atom_map=atom_map) @pytest.fixture diff --git a/source/lmp/tests/test_lammps_dpa_jax.py b/source/lmp/tests/test_lammps_dpa_jax.py index d4d63dd340..6fc27139f1 100644 --- a/source/lmp/tests/test_lammps_dpa_jax.py +++ b/source/lmp/tests/test_lammps_dpa_jax.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -15,12 +12,16 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, + remove_test_files, + require_backend, + run_mpi_model_deviation, + write_water_data_variants, +) from model_convert import ( ensure_converted_pb, ) -from write_lmp_data import ( - write_lmp_data, -) pbtxt_file2 = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot-1.pbtxt" @@ -226,61 +227,20 @@ def setup_module(): - if os.environ.get("ENABLE_JAX", "1") != "1": - pytest.skip( - "Skip test because JAX support is not enabled.", - ) + require_backend("ENABLE_JAX", "JAX") if os.environ.get("ENABLE_TENSORFLOW", "1") == "1": ensure_converted_pb(pbtxt_file2, pb_file2) - - write_lmp_data(box, coord, type_OH, data_file) - write_lmp_data(box, coord, type_HO, data_type_map_file) - write_lmp_data( - box * constants.dist_metal2si, - coord * constants.dist_metal2si, - type_OH, - data_file_si, + write_water_data_variants( + box, coord, type_OH, type_HO, data_file, data_type_map_file, data_file_si ) def teardown_module(): - os.remove(data_file) - os.remove(data_type_map_file) + remove_test_files(data_file, data_type_map_file, data_file_si, md_file) def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - # Requires for DPA-2 - lammps.atom_modify("map yes") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, atom_map="yes") @pytest.fixture @@ -725,24 +685,14 @@ def test_pair_deepmd_si(lammps_si): ) @pytest.mark.skip("MPI is not supported") def test_pair_deepmd_mpi(balance_args: list): - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_faparam.py b/source/lmp/tests/test_lammps_faparam.py index 857acc9f17..6579427aae 100644 --- a/source/lmp/tests/test_lammps_faparam.py +++ b/source/lmp/tests/test_lammps_faparam.py @@ -12,6 +12,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from model_convert import ( ensure_converted_pb, ) @@ -150,34 +153,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, masses=(16,)) @pytest.fixture diff --git a/source/lmp/tests/test_lammps_faparam_pt.py b/source/lmp/tests/test_lammps_faparam_pt.py index 6a22282f43..7a88e30daf 100644 --- a/source/lmp/tests/test_lammps_faparam_pt.py +++ b/source/lmp/tests/test_lammps_faparam_pt.py @@ -11,6 +11,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from write_lmp_data import ( write_lmp_data, ) @@ -86,17 +89,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - lammps.neighbor("2.0 bin") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - lammps.mass("1 16") - lammps.timestep(0.0005) - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, masses=(16,)) @pytest.fixture diff --git a/source/lmp/tests/test_lammps_faparam_pt2.py b/source/lmp/tests/test_lammps_faparam_pt2.py index 721afb0357..7ca73a7673 100644 --- a/source/lmp/tests/test_lammps_faparam_pt2.py +++ b/source/lmp/tests/test_lammps_faparam_pt2.py @@ -14,6 +14,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from write_lmp_data import ( write_lmp_data, ) @@ -76,17 +79,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - lammps.neighbor("2.0 bin") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - lammps.mass("1 16") - lammps.timestep(0.0005) - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units, masses=(16,)) @pytest.fixture diff --git a/source/lmp/tests/test_lammps_fparam_from_fix_dedn.py b/source/lmp/tests/test_lammps_fparam_from_fix_dedn.py index c929380051..e7c788812d 100644 --- a/source/lmp/tests/test_lammps_fparam_from_fix_dedn.py +++ b/source/lmp/tests/test_lammps_fparam_from_fix_dedn.py @@ -11,6 +11,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from model_convert import ( ensure_converted_pb, ) @@ -54,16 +57,7 @@ def teardown_module() -> None: def _lammps(fp_value, units="metal") -> PyLammps: """Build a LAMMPS instance configured for frame-parameter derivative tests.""" - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - lammps.neighbor("2.0 bin") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - lammps.mass("1 16") - lammps.timestep(0.0005) - lammps.fix("1 all nve") + lammps = make_atomic_lammps(data_file, units, masses=(16,)) lammps.variable("fp equal " + str(fp_value)) lammps.variable("dummy equal 0.0") lammps.fix("fpfix all ave/time 1 1 1 v_dummy v_fp") diff --git a/source/lmp/tests/test_lammps_jax.py b/source/lmp/tests/test_lammps_jax.py index dcc20a1e85..5cdfc594a6 100644 --- a/source/lmp/tests/test_lammps_jax.py +++ b/source/lmp/tests/test_lammps_jax.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -15,12 +12,16 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, + remove_test_files, + require_backend, + run_mpi_model_deviation, + write_water_data_variants, +) from model_convert import ( ensure_converted_pb, ) -from write_lmp_data import ( - write_lmp_data, -) pbtxt_file2 = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot-1.pbtxt" @@ -226,59 +227,20 @@ def setup_module(): - if os.environ.get("ENABLE_JAX", "1") != "1": - pytest.skip( - "Skip test because JAX support is not enabled.", - ) + require_backend("ENABLE_JAX", "JAX") if os.environ.get("ENABLE_TENSORFLOW", "1") == "1": ensure_converted_pb(pbtxt_file2, pb_file2) - - write_lmp_data(box, coord, type_OH, data_file) - write_lmp_data(box, coord, type_HO, data_type_map_file) - write_lmp_data( - box * constants.dist_metal2si, - coord * constants.dist_metal2si, - type_OH, - data_file_si, + write_water_data_variants( + box, coord, type_OH, type_HO, data_file, data_type_map_file, data_file_si ) def teardown_module(): - os.remove(data_file) - os.remove(data_type_map_file) + remove_test_files(data_file, data_type_map_file, data_file_si, md_file) def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units) @pytest.fixture @@ -726,24 +688,14 @@ def test_pair_deepmd_si(lammps_si): reason="Skip test because TensorFlow support is not enabled.", ) def test_pair_deepmd_mpi(balance_args: list): - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_model_devi_pt2.py b/source/lmp/tests/test_lammps_model_devi_pt2.py index e802182e85..bb8649ebdb 100644 --- a/source/lmp/tests/test_lammps_model_devi_pt2.py +++ b/source/lmp/tests/test_lammps_model_devi_pt2.py @@ -17,6 +17,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from write_lmp_data import ( write_lmp_data, ) @@ -121,20 +124,7 @@ def teardown_module() -> None: def _lammps(data_file_path=data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - lammps.neighbor("2.0 bin") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file_path.resolve()) - lammps.mass("1 16") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file_path, units, masses=(16,)) @pytest.fixture() diff --git a/source/lmp/tests/test_lammps_pd.py b/source/lmp/tests/test_lammps_pd.py index 2c83a65b80..9debfc7050 100644 --- a/source/lmp/tests/test_lammps_pd.py +++ b/source/lmp/tests/test_lammps_pd.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -15,12 +12,16 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, + remove_test_files, + require_backend, + run_mpi_model_deviation, + write_water_data_variants, +) from model_convert import ( ensure_converted_pb, ) -from write_lmp_data import ( - write_lmp_data, -) pbtxt_file2 = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot-1.pbtxt" @@ -227,59 +228,20 @@ def setup_module(): - if os.environ.get("ENABLE_PADDLE", "1") != "1": - pytest.skip( - "Skip test because Paddle support is not enabled.", - ) + require_backend("ENABLE_PADDLE", "Paddle") if os.environ.get("ENABLE_TENSORFLOW", "1") == "1": ensure_converted_pb(pbtxt_file2, pb_file2) - - write_lmp_data(box, coord, type_OH, data_file) - write_lmp_data(box, coord, type_HO, data_type_map_file) - write_lmp_data( - box * constants.dist_metal2si, - coord * constants.dist_metal2si, - type_OH, - data_file_si, + write_water_data_variants( + box, coord, type_OH, type_HO, data_file, data_type_map_file, data_file_si ) def teardown_module(): - os.remove(data_file) - os.remove(data_type_map_file) + remove_test_files(data_file, data_type_map_file, data_file_si, md_file) def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units) @pytest.fixture @@ -756,24 +718,14 @@ def test_pair_deepmd_si(lammps_si): reason="Skip test because TensorFlow support is not enabled.", ) def test_pair_deepmd_mpi(balance_args: list): - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_pt.py b/source/lmp/tests/test_lammps_pt.py index a95cc4dbaa..8337c822dc 100644 --- a/source/lmp/tests/test_lammps_pt.py +++ b/source/lmp/tests/test_lammps_pt.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -15,12 +12,16 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, + remove_test_files, + require_backend, + run_mpi_model_deviation, + write_water_data_variants, +) from model_convert import ( ensure_converted_pb, ) -from write_lmp_data import ( - write_lmp_data, -) pbtxt_file2 = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot-1.pbtxt" @@ -224,59 +225,20 @@ def setup_module() -> None: - if os.environ.get("ENABLE_PYTORCH", "1") != "1": - pytest.skip( - "Skip test because PyTorch support is not enabled.", - ) + require_backend("ENABLE_PYTORCH", "PyTorch") if os.environ.get("ENABLE_TENSORFLOW", "1") == "1": ensure_converted_pb(pbtxt_file2, pb_file2) - - write_lmp_data(box, coord, type_OH, data_file) - write_lmp_data(box, coord, type_HO, data_type_map_file) - write_lmp_data( - box * constants.dist_metal2si, - coord * constants.dist_metal2si, - type_OH, - data_file_si, + write_water_data_variants( + box, coord, type_OH, type_HO, data_file, data_type_map_file, data_file_si ) def teardown_module() -> None: - os.remove(data_file) - os.remove(data_type_map_file) + remove_test_files(data_file, data_type_map_file, data_file_si, md_file) def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units) @pytest.fixture @@ -724,24 +686,14 @@ def test_pair_deepmd_si(lammps_si) -> None: reason="Skip test because TensorFlow support is not enabled.", ) def test_pair_deepmd_mpi(balance_args: list) -> None: - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_pt2.py b/source/lmp/tests/test_lammps_pt2.py index a800e81838..c4609749bd 100644 --- a/source/lmp/tests/test_lammps_pt2.py +++ b/source/lmp/tests/test_lammps_pt2.py @@ -16,6 +16,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_atomic_lammps, +) from write_lmp_data import ( write_lmp_data, ) @@ -156,36 +159,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("atomic") - if units == "metal" or units == "real": - lammps.neighbor("2.0 bin") - elif units == "si": - lammps.neighbor("2.0e-10 bin") - else: - raise ValueError("units should be metal, real, or si") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal" or units == "real": - lammps.mass("1 16") - lammps.mass("2 2") - elif units == "si": - lammps.mass("1 %.10e" % (16 * constants.mass_metal2si)) - lammps.mass("2 %.10e" % (2 * constants.mass_metal2si)) - else: - raise ValueError("units should be metal, real, or si") - if units == "metal": - lammps.timestep(0.0005) - elif units == "real": - lammps.timestep(0.5) - elif units == "si": - lammps.timestep(5e-16) - else: - raise ValueError("units should be metal, real, or si") - lammps.fix("1 all nve") - return lammps + return make_atomic_lammps(data_file, units) @pytest.fixture diff --git a/source/lmp/tests/test_lammps_spin.py b/source/lmp/tests/test_lammps_spin.py index a5210e2390..b99f7b9b56 100644 --- a/source/lmp/tests/test_lammps_spin.py +++ b/source/lmp/tests/test_lammps_spin.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -15,6 +12,10 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_spin_lammps, + run_mpi_model_deviation, +) from model_convert import ( ensure_converted_pb, ) @@ -229,27 +230,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("spin") - if units == "metal": - lammps.neighbor("2.0 bin") - else: - raise ValueError("units for spin should be metal") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal": - lammps.mass("1 58") - lammps.mass("2 16") - else: - raise ValueError("units for spin should be metal") - if units == "metal": - lammps.timestep(0.0005) - else: - raise ValueError("units for spin should be metal") - lammps.fix("1 all nve") - return lammps + return make_spin_lammps(data_file, units) @pytest.fixture @@ -415,24 +396,14 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: [(["--balance"],), ([],)], ) def test_pair_deepmd_mpi(balance_args: list) -> None: - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd_spin.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd_spin.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_spin_dpa3_pt2.py b/source/lmp/tests/test_lammps_spin_dpa3_pt2.py index d11d29ed34..5ad30dcf29 100644 --- a/source/lmp/tests/test_lammps_spin_dpa3_pt2.py +++ b/source/lmp/tests/test_lammps_spin_dpa3_pt2.py @@ -37,15 +37,15 @@ import importlib.util import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) import numpy as np import pytest +from lammps_test_utils import ( + run_mpi_pair_runner, +) from write_lmp_data import ( write_lmp_data_spin, ) @@ -158,61 +158,17 @@ def _run_mpi_subprocess( fixtures. ``runner_args`` flows additional flags (e.g. ``--pair-coeff``, ``--mass3``) to the subprocess runner. """ - if data_path is None: - data_path = data_file - if pb_path is None: - pb_path = pb_file_mpi - with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: - out_path = f.name - try: - argv = [ - "mpirun", - "-n", - str(nprocs), - sys.executable, - str(Path(__file__).parent / "run_mpi_pair_deepmd_spin_dpa3_pt2.py"), - str(data_path.resolve()), - str(pb_path.resolve()), - out_path, - ] - if processors is not None: - argv.extend(["--processors", processors]) - elif nprocs == 1: - argv.extend(["--processors", "1 1 1"]) - if extra_args: - argv.extend(extra_args) - if runner_args: - argv.extend(runner_args) - if capture: - # Used by fail-fast tests: return raw subprocess info instead of - # parsing output (the subprocess is expected to exit non-zero). - proc = sp.run(argv, capture_output=True, text=True) - return { - "returncode": proc.returncode, - "stdout": proc.stdout, - "stderr": proc.stderr, - } - sp.check_call(argv) - with open(out_path) as fh: - lines = fh.read().strip().splitlines() - pe = float(lines[0]) - rows = np.array( - [list(map(float, line.split())) for line in lines[1:]], - dtype=np.float64, - ) - # Each row: 3 force + 3 force_mag + 9 virial = 15 cols (see runner). - forces = rows[:, :3] - force_mag = rows[:, 3:6] - virials = rows[:, 6:] - return { - "pe": pe, - "forces": forces, - "force_mag": force_mag, - "virials": virials, - } - finally: - if os.path.exists(out_path): - os.remove(out_path) + return run_mpi_pair_runner( + Path(__file__).parent / "run_mpi_pair_deepmd_spin_dpa3_pt2.py", + data_path or data_file, + pb_path or pb_file_mpi, + nprocs=nprocs, + processors=processors, + extra_args=extra_args, + runner_args=runner_args, + output_columns=(("forces", 3), ("force_mag", 3), ("virials", 9)), + capture=capture, + ) @pytest.mark.skipif( diff --git a/source/lmp/tests/test_lammps_spin_nopbc.py b/source/lmp/tests/test_lammps_spin_nopbc.py index 309d1591cc..a7b2ee1e26 100644 --- a/source/lmp/tests/test_lammps_spin_nopbc.py +++ b/source/lmp/tests/test_lammps_spin_nopbc.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -14,6 +11,10 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_spin_lammps, + run_mpi_model_deviation, +) from model_convert import ( ensure_converted_pb, ) @@ -228,27 +229,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("f f f") - lammps.atom_style("spin") - if units == "metal": - lammps.neighbor("2.0 bin") - else: - raise ValueError("units for spin should be metal") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal": - lammps.mass("1 58") - lammps.mass("2 16") - else: - raise ValueError("units for spin should be metal") - if units == "metal": - lammps.timestep(0.0005) - else: - raise ValueError("units for spin should be metal") - lammps.fix("1 all nve") - return lammps + return make_spin_lammps(data_file, units, boundary="f f f") @pytest.fixture @@ -375,25 +356,14 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: [(["--balance"],), ([],)], ) def test_pair_deepmd_mpi(balance_args: list) -> None: - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd_spin.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - "--nopbc", - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd_spin.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=[*balance_args, "--nopbc"], + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_spin_nopbc_pt.py b/source/lmp/tests/test_lammps_spin_nopbc_pt.py index 86d74f68c4..9b25f9a0f2 100644 --- a/source/lmp/tests/test_lammps_spin_nopbc_pt.py +++ b/source/lmp/tests/test_lammps_spin_nopbc_pt.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -14,6 +11,10 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_spin_lammps, + run_mpi_model_deviation, +) from model_convert import ( ensure_converted_pb, ) @@ -106,27 +107,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("f f f") - lammps.atom_style("spin") - if units == "metal": - lammps.neighbor("2.0 bin") - else: - raise ValueError("units for spin should be metal") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal": - lammps.mass("1 58") - lammps.mass("2 16") - else: - raise ValueError("units for spin should be metal") - if units == "metal": - lammps.timestep(0.0005) - else: - raise ValueError("units for spin should be metal") - lammps.fix("1 all nve") - return lammps + return make_spin_lammps(data_file, units, boundary="f f f") @pytest.fixture @@ -222,24 +203,14 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: reason="Skip test because TensorFlow support is not enabled.", ) def test_pair_deepmd_mpi(balance_args: list) -> None: - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd_spin.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd_spin.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_spin_nopbc_pt2.py b/source/lmp/tests/test_lammps_spin_nopbc_pt2.py index 56952fdc0d..fdc08284e4 100644 --- a/source/lmp/tests/test_lammps_spin_nopbc_pt2.py +++ b/source/lmp/tests/test_lammps_spin_nopbc_pt2.py @@ -9,6 +9,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_spin_lammps, +) from write_lmp_data import ( write_lmp_data_spin, ) @@ -95,27 +98,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("f f f") - lammps.atom_style("spin") - if units == "metal": - lammps.neighbor("2.0 bin") - else: - raise ValueError("units for spin should be metal") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal": - lammps.mass("1 58") - lammps.mass("2 16") - else: - raise ValueError("units for spin should be metal") - if units == "metal": - lammps.timestep(0.0005) - else: - raise ValueError("units for spin should be metal") - lammps.fix("1 all nve") - return lammps + return make_spin_lammps(data_file, units, boundary="f f f") @pytest.fixture diff --git a/source/lmp/tests/test_lammps_spin_pt.py b/source/lmp/tests/test_lammps_spin_pt.py index ff9fc6ecb8..522ad57b76 100644 --- a/source/lmp/tests/test_lammps_spin_pt.py +++ b/source/lmp/tests/test_lammps_spin_pt.py @@ -2,9 +2,6 @@ import importlib import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -15,6 +12,10 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_spin_lammps, + run_mpi_model_deviation, +) from model_convert import ( ensure_converted_pb, ) @@ -197,27 +198,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("spin") - if units == "metal": - lammps.neighbor("2.0 bin") - else: - raise ValueError("units for spin should be metal") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal": - lammps.mass("1 58") - lammps.mass("2 16") - else: - raise ValueError("units for spin should be metal") - if units == "metal": - lammps.timestep(0.0005) - else: - raise ValueError("units for spin should be metal") - lammps.fix("1 all nve") - return lammps + return make_spin_lammps(data_file, units) @pytest.fixture @@ -354,24 +335,14 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: reason="Skip test because TensorFlow support is not enabled.", ) def test_pair_deepmd_mpi(balance_args: list) -> None: - with tempfile.NamedTemporaryFile() as f: - sp.check_call( - [ - "mpirun", - "-n", - "2", - sys.executable, - Path(__file__).parent / "run_mpi_pair_deepmd_spin.py", - data_file, - pb_file, - pb_file2, - md_file, - f.name, - *balance_args, - ] - ) - arr = np.loadtxt(f.name, ndmin=1) - pe = arr[0] + pe = run_mpi_model_deviation( + Path(__file__).parent / "run_mpi_pair_deepmd_spin.py", + data_file, + pb_file, + pb_file2, + md_file, + extra_args=balance_args, + ) relative = 1.0 assert pe == pytest.approx(expected_e) diff --git a/source/lmp/tests/test_lammps_spin_pt2.py b/source/lmp/tests/test_lammps_spin_pt2.py index 6a6dc50933..5460d6de5a 100644 --- a/source/lmp/tests/test_lammps_spin_pt2.py +++ b/source/lmp/tests/test_lammps_spin_pt2.py @@ -10,6 +10,9 @@ from lammps import ( PyLammps, ) +from lammps_test_utils import ( + make_spin_lammps, +) from write_lmp_data import ( write_lmp_data_spin, ) @@ -145,27 +148,7 @@ def teardown_module() -> None: def _lammps(data_file, units="metal") -> PyLammps: - lammps = PyLammps() - lammps.units(units) - lammps.boundary("p p p") - lammps.atom_style("spin") - if units == "metal": - lammps.neighbor("2.0 bin") - else: - raise ValueError("units for spin should be metal") - lammps.neigh_modify("every 10 delay 0 check no") - lammps.read_data(data_file.resolve()) - if units == "metal": - lammps.mass("1 58") - lammps.mass("2 16") - else: - raise ValueError("units for spin should be metal") - if units == "metal": - lammps.timestep(0.0005) - else: - raise ValueError("units for spin should be metal") - lammps.fix("1 all nve") - return lammps + return make_spin_lammps(data_file, units) @pytest.fixture From 2fde99e2e7f5335265fdea4f4e7a2f2b3e35ce6c Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 12 Jul 2026 20:40:58 +0800 Subject: [PATCH 4/6] fix(lmp): preserve nopbc in spin MPI test Document the capture-mode return contract for the shared MPI runner. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lmp/tests/lammps_test_utils.py | 3 +++ source/lmp/tests/test_lammps_spin_nopbc_pt.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/source/lmp/tests/lammps_test_utils.py b/source/lmp/tests/lammps_test_utils.py index 983b106383..e3f33f5cb1 100644 --- a/source/lmp/tests/lammps_test_utils.py +++ b/source/lmp/tests/lammps_test_utils.py @@ -142,6 +142,9 @@ def run_mpi_pair_runner( The runner output contract is one energy line followed by a rectangular per-atom table. ``output_columns`` names and slices that table while each model-specific wrapper retains its own defaults and explanatory docstring. + + If ``capture`` is true, skip parsing and return the subprocess result as + ``{"returncode": int, "stdout": str, "stderr": str}``. """ with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: output_path = Path(f.name) diff --git a/source/lmp/tests/test_lammps_spin_nopbc_pt.py b/source/lmp/tests/test_lammps_spin_nopbc_pt.py index 9b25f9a0f2..64967fc042 100644 --- a/source/lmp/tests/test_lammps_spin_nopbc_pt.py +++ b/source/lmp/tests/test_lammps_spin_nopbc_pt.py @@ -209,7 +209,7 @@ def test_pair_deepmd_mpi(balance_args: list) -> None: pb_file, pb_file2, md_file, - extra_args=balance_args, + extra_args=[*balance_args, "--nopbc"], ) relative = 1.0 From ee53139b064ce0aab715392b70b0408683cc42bb Mon Sep 17 00:00:00 2001 From: "A bot of @njzjz" <48687836+njzjz-bot@users.noreply.github.com> Date: Mon, 13 Jul 2026 12:40:04 +0800 Subject: [PATCH 5/6] refactor(lmp): reuse shared DPA1 MPI runner Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .../lmp/tests/test_lammps_dpa1_graph_pt2.py | 56 ++++--------------- 1 file changed, 11 insertions(+), 45 deletions(-) diff --git a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py index f3f3e0bf59..3549ecfd8b 100644 --- a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py +++ b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py @@ -21,9 +21,6 @@ import importlib.util import os import shutil -import subprocess as sp -import sys -import tempfile from pathlib import ( Path, ) @@ -39,6 +36,7 @@ ) from lammps_test_utils import ( make_atomic_lammps, + run_mpi_pair_runner, ) from write_lmp_data import ( write_lmp_data, @@ -184,53 +182,21 @@ def _run_mpi_subprocess( processors: str | None = None, runner_args: list[str] | None = None, ) -> dict: - """Invoke the (backend-agnostic) DPA3 MPI runner under - ``mpirun -n `` against the dpa1 graph .pt2 and return - ``{"pe": float, "forces": (n, 3), "virials": (n, 9)}``. + """Invoke the shared DPA MPI runner against the DPA1 graph model. ``nprocs == 1`` forces ``--processors 1 1 1`` so the C++ side sees ``nprocs == 1`` and routes to the single-rank graph path — a same-archive reference for the multi-rank comparison. """ - if data_path is None: - data_path = data_file - with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: - out_path = f.name - try: - argv = [ - "mpirun", - "-n", - str(nprocs), - sys.executable, - str(mpi_runner), - str(data_path.resolve()), - str(pb_file.resolve()), - out_path, - ] - if processors is not None: - argv.extend(["--processors", processors]) - elif nprocs == 1: - argv.extend(["--processors", "1 1 1"]) - if extra_args: - argv.extend(extra_args) - if runner_args: - argv.extend(runner_args) - sp.check_call(argv) - with open(out_path) as fh: - lines = fh.read().strip().splitlines() - pe = float(lines[0]) - rows = np.array( - [list(map(float, line.split())) for line in lines[1:]], - dtype=np.float64, - ) - forces = rows[:, :3] - virials = rows[:, 3:] - return {"pe": pe, "forces": forces, "virials": virials} - finally: - if os.path.exists(out_path): - os.remove(out_path) - - + return run_mpi_pair_runner( + mpi_runner, + data_path or data_file, + pb_file, + nprocs=nprocs, + processors=processors, + extra_args=extra_args, + runner_args=runner_args, + ) @pytest.mark.skipif( shutil.which("mpirun") is None, reason="MPI is not installed on this system" ) From 4f330eba4cbca97b7cb4bfb242f7da78c2370596 Mon Sep 17 00:00:00 2001 From: "A bot of @njzjz" <48687836+njzjz-bot@users.noreply.github.com> Date: Mon, 13 Jul 2026 12:41:23 +0800 Subject: [PATCH 6/6] style(lmp): format shared runner call Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lmp/tests/test_lammps_dpa1_graph_pt2.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py index 3549ecfd8b..09e8c10738 100644 --- a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py +++ b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py @@ -197,6 +197,8 @@ def _run_mpi_subprocess( extra_args=extra_args, runner_args=runner_args, ) + + @pytest.mark.skipif( shutil.which("mpirun") is None, reason="MPI is not installed on this system" )