Split data of last iteration to training data and validation data (#293)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added an option to automatically split the last iteration's training
data into separate training and validation sets using a configurable
ratio.

- **Bug Fixes**
- Fixed an issue ensuring temperature settings are correctly applied
when preparing VASP input files.
- Improved validation data loading logic for multitask and single-task
workflows.

- **Tests**
- Added tests to verify the correct splitting of datasets into training
and validation subsets.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: zjgemi <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Han Wang <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

Author: 4 people

dpgen2/entrypoint/submit.py

@@ -533,7 +533,7 @@ def workflow_concurrent_learning(
     else:
         if config["inputs"]["valid_data_uri"] is not None:
             valid_data = get_artifact_from_uri(config["inputs"]["valid_data_uri"])
-        elif config["inputs"]["valid_data_prefix"] is not None:
+        elif config["inputs"]["valid_data_sys"] is not None:
             valid_data_prefix = config["inputs"]["valid_data_prefix"]
             valid_data = config["inputs"]["valid_data_sys"]
             valid_data = get_systems_from_data(valid_data, valid_data_prefix)

dpgen2/fp/vasp.py

@@ -141,7 +141,7 @@ def prep_task(
 
         conf_frame.to("vasp/poscar", vasp_conf_name)
         incar = vasp_inputs.incar_template
-        self.set_ele_temp(conf_frame, incar)
+        incar = self.set_ele_temp(conf_frame, incar)
 
         Path(vasp_input_name).write_text(incar)
         # fix the case when some element have 0 atom, e.g. H0O2

dpgen2/op/run_dp_train.py

@@ -1,7 +1,9 @@
 import glob
 import json
 import logging
+import math
 import os
+import random
 import shutil
 from pathlib import (
     Path,
@@ -197,6 +199,12 @@ def execute(
         valid_data = ip["valid_data"]
         iter_data_old_exp = _expand_all_multi_sys_to_sys(iter_data[:-1])
         iter_data_new_exp = _expand_all_multi_sys_to_sys(iter_data[-1:])
+        if config["split_last_iter_valid_ratio"] is not None:
+            train_systems, valid_systems = split_valid(
+                iter_data_new_exp, config["split_last_iter_valid_ratio"]
+            )
+            iter_data_new_exp = train_systems
+            valid_data = append_valid_data(config, valid_data, valid_systems)
         iter_data_exp = iter_data_old_exp + iter_data_new_exp
         work_dir = Path(task_name)
         init_model_with_finetune = config["init_model_with_finetune"]
@@ -517,6 +525,9 @@ def training_args():
         doc_head = "Head to use in the multitask training"
         doc_init_model_with_finetune = "Use finetune for init model"
         doc_train_args = "Extra arguments for dp train"
+        doc_split_last_iter_valid_ratio = (
+            "Ratio of valid data if split data of last iter"
+        )
         return [
             Argument(
                 "command",
@@ -618,6 +629,13 @@ def training_args():
                 default="",
                 doc=doc_train_args,
             ),
+            Argument(
+                "split_last_iter_valid_ratio",
+                float,
+                optional=True,
+                default=None,
+                doc=doc_split_last_iter_valid_ratio,
+            ),
         ]
 
     @staticmethod
@@ -672,4 +690,75 @@ def _expand_all_multi_sys_to_sys(list_multi_sys):
     return all_sys_dirs
 
 
+def split_valid(systems: List[str], valid_ratio: float):
+    train_systems = []
+    valid_systems = []
+    for system in systems:
+        d = dpdata.MultiSystems()
+        mixed_type = len(glob.glob("%s/*/real_atom_types.npy" % system)) > 0
+        if mixed_type:
+            d.load_systems_from_file(system, fmt="deepmd/npy/mixed")
+        else:
+            k = dpdata.LabeledSystem(system, fmt="deepmd/npy")
+            d.append(k)
+
+        train_multi_systems = dpdata.MultiSystems()
+        valid_multi_systems = dpdata.MultiSystems()
+        for s in d:
+            nvalid = math.floor(len(s) * valid_ratio)
+            if random.random() < len(s) * valid_ratio - nvalid:
+                nvalid += 1
+            valid_indices = random.sample(range(len(s)), nvalid)
+            train_indices = list(set(range(len(s))).difference(valid_indices))
+            if len(valid_indices) > 0:
+                valid_multi_systems.append(s.sub_system(valid_indices))
+            if len(train_indices) > 0:
+                train_multi_systems.append(s.sub_system(train_indices))
+
+        if len(train_multi_systems) > 0:
+            target = "train_data/" + system
+            if mixed_type:
+                # The multisystem is loaded from one dir, thus we can safely keep one dir
+                train_multi_systems.to_deepmd_npy_mixed("%s.tmp" % target)  # type: ignore
+                fs = os.listdir("%s.tmp" % target)
+                assert len(fs) == 1
+                os.rename(os.path.join("%s.tmp" % target, fs[0]), target)
+                os.rmdir("%s.tmp" % target)
+            else:
+                train_multi_systems[0].to_deepmd_npy(target)  # type: ignore
+            train_systems.append(os.path.abspath(target))
+
+        if len(valid_multi_systems) > 0:
+            target = "valid_data/" + system
+            if mixed_type:
+                # The multisystem is loaded from one dir, thus we can safely keep one dir
+                valid_multi_systems.to_deepmd_npy_mixed("%s.tmp" % target)  # type: ignore
+                fs = os.listdir("%s.tmp" % target)
+                assert len(fs) == 1
+                os.rename(os.path.join("%s.tmp" % target, fs[0]), target)
+                os.rmdir("%s.tmp" % target)
+            else:
+                valid_multi_systems[0].to_deepmd_npy(target)  # type: ignore
+            valid_systems.append(os.path.abspath(target))
+
+    return train_systems, valid_systems
+
+
+def append_valid_data(config, valid_data, valid_systems):
+    if not valid_systems:
+        return valid_data
+    if config["multitask"]:
+        head = config["head"]
+        if not valid_data:
+            valid_data = {}
+        if head not in valid_data:
+            valid_data[head] = []
+        valid_data[head] += valid_systems
+    else:
+        if not valid_data:
+            valid_data = []
+        valid_data += valid_systems
+    return valid_data
+
+
 config_args = RunDPTrain.training_args

tests/op/test_run_dp_train.py

@@ -7,6 +7,7 @@
     Path,
 )
 
+import dpdata
 import numpy as np
 from dflow.python import (
     OP,
@@ -37,6 +38,7 @@
     RunDPTrain,
     _get_data_size_of_all_mult_sys,
     _make_train_command,
+    split_valid,
 )
 
 # isort: on
@@ -942,3 +944,44 @@ def test_exec_v2_empty_dir(self, mocked_run):
         with open(out["script"]) as fp:
             jdata = json.load(fp)
             self.assertDictEqual(jdata, self.expected_odict_v2)
+
+
+class TestSplitValid(unittest.TestCase):
+    def setUp(self):
+        s = fake_system(10, 1)
+        s.to_deepmd_npy("fake_data")
+        ms = fake_multi_sys([10, 20], [1, 2])
+        ms.to_deepmd_npy_mixed("fake_mixed_data")
+
+    def test_split_valid(self):
+        train_systems, valid_systems = split_valid(["fake_data"], 0.1)
+        self.assertEqual(len(train_systems), 1)
+        s = dpdata.LabeledSystem(train_systems[0], fmt="deepmd/npy")
+        self.assertEqual(len(s), 9)
+        self.assertEqual(len(valid_systems), 1)
+        s = dpdata.LabeledSystem(valid_systems[0], fmt="deepmd/npy")
+        self.assertEqual(len(s), 1)
+
+    def test_split_valid_mixed(self):
+        train_systems, valid_systems = split_valid(
+            ["fake_mixed_data/1", "fake_mixed_data/2"], 0.1
+        )
+        self.assertEqual(len(train_systems), 2)
+        ms = dpdata.MultiSystems()
+        ms.load_systems_from_file(train_systems[0], fmt="deepmd/npy/mixed")
+        self.assertEqual(len(ms[0]), 9)
+        ms = dpdata.MultiSystems()
+        ms.load_systems_from_file(train_systems[1], fmt="deepmd/npy/mixed")
+        self.assertEqual(len(ms[0]), 18)
+        self.assertEqual(len(valid_systems), 2)
+        ms = dpdata.MultiSystems()
+        ms.load_systems_from_file(valid_systems[0], fmt="deepmd/npy/mixed")
+        self.assertEqual(len(ms[0]), 1)
+        ms = dpdata.MultiSystems()
+        ms.load_systems_from_file(valid_systems[1], fmt="deepmd/npy/mixed")
+        self.assertEqual(len(ms[0]), 2)
+
+    def tearDown(self):
+        for f in ["fake_data", "fake_mixed_data", "train_data", "valid_data"]:
+            if os.path.exists(f):
+                shutil.rmtree(f)