feat: reuse relaxation CHGCAR for init MD

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]:

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"]

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()