From a07853ee15f9719a43322618af7be56805bab738 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 24 Aug 2026 04:58:10 +0800 Subject: [PATCH] feat: reuse relaxation CHGCAR for init MD Coding-Agent: Codex Codex-Version: codex-cli 0.149.0 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- dpgen/data/arginfo.py | 14 +++- dpgen/data/gen.py | 19 ++++++ tests/data/test_init_chgcar.py | 114 +++++++++++++++++++++++++++++++++ 3 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 tests/data/test_init_chgcar.py diff --git a/dpgen/data/arginfo.py b/dpgen/data/arginfo.py index 37c9b1fa9..0d5c94d73 100644 --- a/dpgen/data/arginfo.py +++ b/dpgen/data/arginfo.py @@ -37,7 +37,19 @@ def init_reaction_mdata_arginfo() -> Argument: def init_bulk_vasp_args() -> list[Argument]: - return [] + return [ + Argument( + "reuse_relax_chgcar", + bool, + optional=True, + default=False, + doc=( + "Download each stage-1 relaxation CHGCAR and forward it to all " + "stage-3 VASP MD tasks derived from the same system. The MD " + "INCAR must enable charge-density restart settings as needed." + ), + ) + ] def init_bulk_abacus_args() -> list[Argument]: diff --git a/dpgen/data/gen.py b/dpgen/data/gen.py index 5c10657d2..51704104a 100644 --- a/dpgen/data/gen.py +++ b/dpgen/data/gen.py @@ -871,6 +871,14 @@ def make_vasp_md(jdata, mdata): os.chdir(cwd) for ii in sys_ps: + relax_chgcar = os.path.abspath( + os.path.join(out_dir, global_dirname_02, ii, "CHGCAR") + ) + if jdata.get("reuse_relax_chgcar", False) and not os.path.isfile(relax_chgcar): + raise RuntimeError( + f"file {relax_chgcar} not found; stage-1 VASP relaxation must " + "produce CHGCAR before reuse_relax_chgcar can be enabled" + ) for jj in scale: for kk in range(pert_numb + 1): path_work = path_md @@ -895,6 +903,13 @@ def make_vasp_md(jdata, mdata): os.symlink(os.path.relpath(file_potcar), "POTCAR") except FileExistsError: pass + if jdata.get("reuse_relax_chgcar", False): + # One relaxed charge density seeds every scaled/perturbed MD + # task belonging to the same chemical system. + try: + os.symlink(os.path.relpath(relax_chgcar), "CHGCAR") + except FileExistsError: + pass is_cvasp = False if "cvasp" in mdata["fp_resources"].keys(): @@ -1145,6 +1160,8 @@ def run_vasp_relax(jdata, mdata): user_forward_files = mdata.get("fp" + "_user_forward_files", []) forward_files += [os.path.basename(file) for file in user_forward_files] backward_files = ["OUTCAR", "CONTCAR"] + if jdata.get("reuse_relax_chgcar", False): + backward_files.append("CHGCAR") backward_files += mdata.get("fp" + "_user_backward_files", []) forward_common_files = [] if "cvasp" in mdata["fp_resources"]: @@ -1329,6 +1346,8 @@ def run_vasp_md(jdata, mdata): md_nstep = jdata["md_nstep"] forward_files = ["POSCAR", "INCAR", "POTCAR"] + if jdata.get("reuse_relax_chgcar", False): + forward_files.append("CHGCAR") user_forward_files = mdata.get("fp" + "_user_forward_files", []) forward_files += [os.path.basename(file) for file in user_forward_files] backward_files = ["OUTCAR"] diff --git a/tests/data/test_init_chgcar.py b/tests/data/test_init_chgcar.py new file mode 100644 index 000000000..00a22ce61 --- /dev/null +++ b/tests/data/test_init_chgcar.py @@ -0,0 +1,114 @@ +"""Tests for carrying VASP relaxation charge densities into init MD tasks.""" + +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +from dpgen.data.arginfo import init_bulk_jdata_arginfo +from dpgen.data.gen import make_vasp_md, run_vasp_md, run_vasp_relax + + +class TestInitChgcar(unittest.TestCase): + """Validate local setup and remote staging for CHGCAR reuse.""" + + def setUp(self): + self.tempdir = tempfile.TemporaryDirectory() + self.root = Path(self.tempdir.name) + self.out_dir = self.root / "init" + self.md_incar = self.root / "INCAR.md" + self.potcar = self.root / "POTCAR" + self.md_incar.write_text("ISTART = 0\nICHARG = 1\n") + self.potcar.write_text("pseudopotential\n") + self.jdata = { + "out_dir": str(self.out_dir), + "potcars": [str(self.potcar)], + "scale": [1.0], + "pert_numb": 0, + "md_nstep": 1, + "md_incar": str(self.md_incar), + "reuse_relax_chgcar": True, + } + self.mdata = { + "fp_command": "vasp_std", + "fp_group_size": 1, + "fp_resources": {}, + "fp_machine": {}, + } + + def tearDown(self): + self.tempdir.cleanup() + + def _make_md_source(self): + source = self.out_dir / "01.scale_pert" / "sys-0001" / "scale-1.000" / "000000" + source.mkdir(parents=True) + (source / "POSCAR").write_text("structure\n") + + def test_vasp_schema_accepts_chgcar_reuse(self): + arginfo = init_bulk_jdata_arginfo() + data = { + "stages": [1, 2, 3], + "elements": ["Al"], + "potcars": ["POTCAR"], + "cell_type": "fcc", + "super_cell": [2, 2, 2], + "relax_incar": "INCAR.relax", + "md_incar": "INCAR.md", + "scale": [1.0], + "skip_relax": False, + "pert_numb": 1, + "pert_box": 0.01, + "pert_atom": 0.01, + "md_nstep": 10, + "coll_ndata": 10, + "init_fp_style": "VASP", + "reuse_relax_chgcar": True, + } + + normalized = arginfo.normalize_value(data) + arginfo.check_value(normalized, strict=True) + + def test_make_vasp_md_links_system_relax_chgcar(self): + self._make_md_source() + relax_dir = self.out_dir / "00.place_ele" / "sys-0001" + relax_dir.mkdir(parents=True) + relax_chgcar = relax_dir / "CHGCAR" + relax_chgcar.write_text("charge density\n") + + make_vasp_md(self.jdata, {"fp_resources": {}}) + + task_chgcar = ( + self.out_dir / "02.md" / "sys-0001" / "scale-1.000" / "000000" / "CHGCAR" + ) + self.assertTrue(task_chgcar.is_symlink()) + self.assertTrue(os.path.samefile(task_chgcar, relax_chgcar)) + + def test_make_vasp_md_requires_relax_chgcar(self): + self._make_md_source() + + with self.assertRaisesRegex(RuntimeError, "stage-1 VASP relaxation"): + make_vasp_md(self.jdata, {"fp_resources": {}}) + + @patch("dpgen.data.gen.check_api_version") + @patch("dpgen.data.gen.make_submission") + def test_dispatcher_transfers_chgcar(self, make_submission, _check_api_version): + submission = MagicMock() + make_submission.return_value = submission + + relax_task = self.out_dir / "00.place_ele" / "sys-0001" + relax_task.mkdir(parents=True) + run_vasp_relax(self.jdata, self.mdata) + relax_call = make_submission.call_args.kwargs + self.assertIn("CHGCAR", relax_call["backward_files"]) + + make_submission.reset_mock() + md_task = self.out_dir / "02.md" / "sys-0001" / "scale-1.000" / "000000" + md_task.mkdir(parents=True) + run_vasp_md(self.jdata, self.mdata) + md_call = make_submission.call_args.kwargs + self.assertIn("CHGCAR", md_call["forward_files"]) + + +if __name__ == "__main__": + unittest.main()