diff --git a/docs/input.md b/docs/input.md index 6fc03bd1..bd71a2ea 100644 --- a/docs/input.md +++ b/docs/input.md @@ -127,7 +127,39 @@ The {dargs:argument}`"configurations"` provides the - Automatic alloy configuration generator. See {ref}`the detailed doc` for the allowed keys. - Configurations load from files. See {ref}`the detailed doc` for the allowed keys. -The {dargs:argument}`"stages"` defines the exploration stages. It is of type `list[list[dict]]`. The outer `list` enumerate the exploration stages, the inner list enumerate the task groups of the stage. Each `dict` defines a stage. See {ref}`the full documentation of the task group` for writting task groups. +The {dargs:argument}`"stages"` defines the exploration stages. The legacy form is `list[list[dict]]`: the outer list enumerates stages and each inner list contains that stage's task groups. See {ref}`the full documentation of the task group` for writing task groups. + +When different stages need different convergence or labeling limits, a stage may instead be a dictionary with a required `task_groups` list and optional `convergence`, `max_numb_iter`, `fatal_at_max`, and `task_max` overrides. Missing overrides inherit the global values from `explore` and `fp`: + +```json +"stages": [ + { + "task_groups": [ + { + "type": "lmp-md", + "ensemble": "nvt", + "nsteps": 50, + "temps": [50], + "trj_freq": 10, + "conf_idx": [0], + "n_sample": 3 + } + ], + "convergence": { + "type": "adaptive-lower", + "rate_candi_f": 0.15, + "level_f_hi": 5.0, + "n_checked_steps": 3, + "conv_tolerance": 0.005 + }, + "max_numb_iter": 2, + "fatal_at_max": false, + "task_max": 4000 + } +] +``` + +Training loss endpoints can also be specialized without changing the template script. `train.config.finetune_end_pref_e/f/v` apply to the initial finetuning run, while `train.config.init_model_end_pref_e/f/v` apply to subsequent init-model training. An unset value keeps the corresponding `limit_pref_e/f/v` from the template. The {dargs:argument}`"n_sample"` tells the number of confgiruations randomly sampled from the set picked by {dargs:argument}`"conf_idx"` from {dargs:argument}`"configurations"` for each exploration task. All configurations has the equal possibility to be sampled. The default value of `"n_sample"` is `null`, in this case all picked configurations are sampled. In the example, we have 3 samples for stage 0 task group 0 and 2 thermodynamic states (NVT, T=50 and 100K), then the task group has 3x2=6 NVT DPMD tasks. diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index df11ff7f..19aed5f2 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -206,9 +206,10 @@ def lmp_args(): doc_convergence = "The method of convergence check." doc_configuration = "A list of initial configurations." doc_stages = ( - "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. " - "The outer list provides the enumeration of the exploration stages. " - "Then each stage is defined by a list of exploration task groups. " + "The definition of exploration stages. A legacy stage is a list of " + "exploration task groups. A stage dictionary stores those groups under " + "`task_groups` and may override `convergence`, `max_numb_iter`, " + "`fatal_at_max`, and `task_max` for that stage. " "Each task group is described in :ref:`the task group definition` " ) doc_filters = "A list of configuration filters" @@ -249,7 +250,7 @@ def lmp_args(): doc=doc_configuration, alias=["configuration"], ), - Argument("stages", List[List[dict]], optional=False, doc=doc_stages), + Argument("stages", list, optional=False, doc=doc_stages), Argument( "filters", list, @@ -299,9 +300,10 @@ def caly_args(): doc_convergence = "The method of convergence check." doc_configuration = "A list of initial configurations." doc_stages = ( - "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. " - "The outer list provides the enumeration of the exploration stages. " - "Then each stage is defined by a list of exploration task groups. " + "The definition of exploration stages. A legacy stage is a list of " + "exploration task groups. A stage dictionary stores those groups under " + "`task_groups` and may override `convergence`, `max_numb_iter`, " + "`fatal_at_max`, and `task_max` for that stage. " "Each task group is described in :ref:`the task group definition` " ) doc_filters = "A list of configuration filters" @@ -342,7 +344,7 @@ def caly_args(): doc=doc_configuration, alias=["configuration"], ), - Argument("stages", List[List[dict]], optional=False, doc=doc_stages), + Argument("stages", list, optional=False, doc=doc_stages), Argument( "filters", list, @@ -400,9 +402,10 @@ def diffcsp_args(): doc_output_nopbc = "Remove pbc of the output configurations" doc_convergence = "The method of convergence check." doc_stages = ( - "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. " - "The outer list provides the enumeration of the exploration stages. " - "Then each stage is defined by a list of exploration task groups. " + "The definition of exploration stages. A legacy stage is a list of " + "exploration task groups. A stage dictionary stores those groups under " + "`task_groups` and may override `convergence`, `max_numb_iter`, " + "`fatal_at_max`, and `task_max` for that stage. " "Each task group is described in :ref:`the task group definition` " ) doc_filters = "A list of configuration filters" @@ -432,7 +435,7 @@ def diffcsp_args(): optional=False, doc=doc_convergence, ), - Argument("stages", List[List[dict]], optional=False, doc=doc_stages), + Argument("stages", list, optional=False, doc=doc_stages), Argument( "filters", list, diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 9e07374f..67435100 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -298,36 +298,99 @@ def get_conf_filters(config): return conf_filters +def _normalize_exploration_stage(stage_config, explore_config, fp_task_max): + """Resolve task groups and effective controls for one exploration stage. + + Legacy configurations represent a stage directly as a list of task-group + dictionaries. Metadata-bearing stage dictionaries keep that list under + ``task_groups`` and override only the controls that need to differ from the + global exploration and FP defaults. + """ + if isinstance(stage_config, list): + task_groups = stage_config + stage_options = {} + elif isinstance(stage_config, dict) and "task_groups" in stage_config: + stage_options = deepcopy(stage_config) + task_groups = stage_options.pop("task_groups") + elif isinstance(stage_config, dict): + # Keep accepting a single task-group dictionary as one legacy stage. + task_groups = [stage_config] + stage_options = {} + else: + raise TypeError("each exploration stage must be a list or dictionary") + + if isinstance(task_groups, dict): + task_groups = [task_groups] + if not isinstance(task_groups, list) or not all( + isinstance(task_group, dict) for task_group in task_groups + ): + raise TypeError("stage task_groups must be a list of dictionaries") + + convergence = deepcopy( + stage_options.get("convergence", explore_config["convergence"]) + ) + max_numb_iter = stage_options.get("max_numb_iter", explore_config["max_numb_iter"]) + fatal_at_max = stage_options.get("fatal_at_max", explore_config["fatal_at_max"]) + stage_task_max = stage_options.get("task_max", fp_task_max) + if not isinstance(convergence, dict): + raise TypeError("stage convergence must be a dictionary") + if max_numb_iter is not None and not isinstance(max_numb_iter, int): + raise TypeError("stage max_numb_iter must be an integer or null") + if not isinstance(fatal_at_max, bool): + raise TypeError("stage fatal_at_max must be a boolean") + if stage_task_max is not None and not isinstance(stage_task_max, int): + raise TypeError("stage task_max must be an integer or null") + return task_groups, convergence, max_numb_iter, fatal_at_max, stage_task_max + + +def _make_stage_selector( + convergence, + render, + fp_task_max, + conf_filters, +): + """Build an independent report and selector for an exploration stage.""" + conv_style = convergence.pop("type") + report = conv_styles[conv_style](**convergence) + return ConfSelectorFrames( + render, + report, + fp_task_max, + conf_filters, + ) + + def make_naive_exploration_scheduler_without_conf(config, explore_style): model_devi_jobs = config["explore"]["stages"] fp_task_max = config["fp"]["task_max"] - max_numb_iter = config["explore"]["max_numb_iter"] - fatal_at_max = config["explore"]["fatal_at_max"] - convergence = config["explore"]["convergence"] output_nopbc = config["explore"]["output_nopbc"] conf_filters = get_conf_filters(config["explore"]["filters"]) scheduler = ExplorationScheduler() - # report - conv_style = convergence.pop("type") - report = conv_styles[conv_style](**convergence) # trajectory render, the format of the output trajs are assumed to be lammps/dump render = TrajRenderLammps(nopbc=output_nopbc) - # selector - selector = ConfSelectorFrames( - render, - report, - fp_task_max, - conf_filters, - ) - for job_ in model_devi_jobs: - if not isinstance(job_, list): - job = [job_] - else: - job = job_ + for stage_config in model_devi_jobs: + ( + job, + convergence, + max_numb_iter, + fatal_at_max, + stage_task_max, + ) = _normalize_exploration_stage( + stage_config, + config["explore"], + fp_task_max, + ) + selector = _make_stage_selector( + convergence, + render, + stage_task_max, + conf_filters, + ) # stage stage = ExplorationStage() for jj in job: + jj = deepcopy(jj) if "calypso" in explore_style: jconf = caly_normalize(jj) # make task group @@ -361,39 +424,41 @@ def make_lmp_naive_exploration_scheduler(config): type_map = config["inputs"]["type_map"] numb_models = config["train"]["numb_models"] fp_task_max = config["fp"]["task_max"] - max_numb_iter = config["explore"]["max_numb_iter"] - fatal_at_max = config["explore"]["fatal_at_max"] - convergence = config["explore"]["convergence"] output_nopbc = config["explore"]["output_nopbc"] conf_filters = get_conf_filters(config["explore"]["filters"]) use_ele_temp = config["inputs"]["use_ele_temp"] scheduler = ExplorationScheduler() - # report - conv_style = convergence.pop("type") - report = conv_styles[conv_style](**convergence) render = TrajRenderLammps(nopbc=output_nopbc, use_ele_temp=use_ele_temp) - # selector - selector = ConfSelectorFrames( - render, - report, - fp_task_max, - conf_filters, - ) sys_configs_lmp = [] for sys_config in sys_configs: + sys_config = deepcopy(sys_config) conf_style = sys_config.pop("type") generator = conf_styles[conf_style](**sys_config) sys_configs_lmp.append(generator.get_file_content(type_map)) - for job_ in model_devi_jobs: - if not isinstance(job_, list): - job = [job_] - else: - job = job_ + for stage_config in model_devi_jobs: + ( + job, + convergence, + max_numb_iter, + fatal_at_max, + stage_task_max, + ) = _normalize_exploration_stage( + stage_config, + config["explore"], + fp_task_max, + ) + selector = _make_stage_selector( + convergence, + render, + stage_task_max, + conf_filters, + ) # stage stage = ExplorationStage() for jj in job: + jj = deepcopy(jj) jconf = normalize_lmp_task_group_config(jj) n_sample = jconf.pop("n_sample") ## ignore the expansion of sys_idx diff --git a/dpgen2/op/run_dp_train.py b/dpgen2/op/run_dp_train.py index c0cf2d4d..e6b641b0 100644 --- a/dpgen2/op/run_dp_train.py +++ b/dpgen2/op/run_dp_train.py @@ -265,7 +265,11 @@ def execute( valid_data, ) train_dict = RunDPTrain.write_other_to_input_script( - train_dict, config, do_init_model, major_version + train_dict, + config, + do_init_model, + major_version, + finetune_mode, ) if RunDPTrain.skip_training( @@ -430,21 +434,21 @@ def write_other_to_input_script( config, do_init_model, major_version: str = "1", + finetune_mode: str = "no", ): odict = idict.copy() odict["training"]["disp_file"] = "lcurve.out" + loss_sections = ( + [value for value in odict["loss_dict"].values() if isinstance(value, dict)] + if "loss_dict" in odict + else [odict["loss"]] + ) if do_init_model: odict["learning_rate"]["start_lr"] = config["init_model_start_lr"] - if "loss_dict" in odict: - for v in odict["loss_dict"].values(): - if isinstance(v, dict): - v["start_pref_e"] = config["init_model_start_pref_e"] - v["start_pref_f"] = config["init_model_start_pref_f"] - v["start_pref_v"] = config["init_model_start_pref_v"] - else: - odict["loss"]["start_pref_e"] = config["init_model_start_pref_e"] - odict["loss"]["start_pref_f"] = config["init_model_start_pref_f"] - odict["loss"]["start_pref_v"] = config["init_model_start_pref_v"] + for loss in loss_sections: + loss["start_pref_e"] = config["init_model_start_pref_e"] + loss["start_pref_f"] = config["init_model_start_pref_f"] + loss["start_pref_v"] = config["init_model_start_pref_v"] if major_version == "1": odict["training"]["stop_batch"] = config["init_model_numb_steps"] elif major_version == "2": @@ -453,6 +457,16 @@ def write_other_to_input_script( raise RuntimeError( "unsupported DeePMD-kit major version", major_version ) + + # Finetuning and later init-model training can need different terminal + # loss weights. Unset values deliberately preserve the template script. + end_pref_prefix = "finetune" if finetune_mode == "finetune" else "init_model" + if finetune_mode == "finetune" or do_init_model: + for loss in loss_sections: + for component in ("e", "f", "v"): + value = config[f"{end_pref_prefix}_end_pref_{component}"] + if value is not None: + loss[f"limit_pref_{component}"] = value return odict @staticmethod @@ -536,6 +550,30 @@ def training_args(): doc_init_model_start_pref_v = ( "The start virial prefactor in loss when init-model" ) + doc_init_model_end_pref_e = ( + "The ending energy prefactor in loss for init-model training. " + "Keep the template value when unset." + ) + doc_init_model_end_pref_f = ( + "The ending force prefactor in loss for init-model training. " + "Keep the template value when unset." + ) + doc_init_model_end_pref_v = ( + "The ending virial prefactor in loss for init-model training. " + "Keep the template value when unset." + ) + doc_finetune_end_pref_e = ( + "The ending energy prefactor in loss for the initial finetuning. " + "Keep the template value when unset." + ) + doc_finetune_end_pref_f = ( + "The ending force prefactor in loss for the initial finetuning. " + "Keep the template value when unset." + ) + doc_finetune_end_pref_v = ( + "The ending virial prefactor in loss for the initial finetuning. " + "Keep the template value when unset." + ) doc_finetune_args = "Extra arguments for finetuning" doc_multitask = "Do multitask training" doc_head = "Head to use in the multitask training" @@ -610,6 +648,48 @@ def training_args(): default=0.0, doc=doc_init_model_start_pref_v, ), + Argument( + "init_model_end_pref_e", + float, + optional=True, + default=None, + doc=doc_init_model_end_pref_e, + ), + Argument( + "init_model_end_pref_f", + float, + optional=True, + default=None, + doc=doc_init_model_end_pref_f, + ), + Argument( + "init_model_end_pref_v", + float, + optional=True, + default=None, + doc=doc_init_model_end_pref_v, + ), + Argument( + "finetune_end_pref_e", + float, + optional=True, + default=None, + doc=doc_finetune_end_pref_e, + ), + Argument( + "finetune_end_pref_f", + float, + optional=True, + default=None, + doc=doc_finetune_end_pref_f, + ), + Argument( + "finetune_end_pref_v", + float, + optional=True, + default=None, + doc=doc_finetune_end_pref_v, + ), Argument( "init_model_with_finetune", bool, diff --git a/tests/entrypoint/test_submit.py b/tests/entrypoint/test_submit.py index bc551a2c..051b9971 100644 --- a/tests/entrypoint/test_submit.py +++ b/tests/entrypoint/test_submit.py @@ -5,6 +5,9 @@ import tempfile import textwrap import unittest +from copy import ( + deepcopy, +) from pathlib import ( Path, ) @@ -20,6 +23,7 @@ copy_scheduler_plans, expand_idx, get_resubmit_keys, + make_lmp_naive_exploration_scheduler, print_list_steps, submit_concurrent_learning, update_reuse_step_scheduler, @@ -380,6 +384,48 @@ def test(self): remove_executor_if_debug(wf_config) submit_concurrent_learning(wf_config, no_submission=True) + def test_stage_specific_controls(self): + from dpgen2.entrypoint.args import ( + normalize, + ) + + wf_config = normalize(json.loads(input_std)) + task_groups = wf_config["explore"]["stages"][0] + wf_config["explore"]["stages"] = [ + { + "task_groups": deepcopy(task_groups), + "convergence": { + "type": "fixed-levels", + "level_f_lo": 0.1, + "level_f_hi": 0.2, + "conv_accuracy": 0.8, + }, + "max_numb_iter": 2, + "fatal_at_max": True, + "task_max": 7, + }, + { + "task_groups": [deepcopy(task_groups[0])], + "max_numb_iter": 9, + "fatal_at_max": False, + }, + ] + + scheduler = make_lmp_naive_exploration_scheduler(wf_config) + first, second = scheduler.stage_schedulers + self.assertEqual(first.max_numb_iter, 2) + self.assertTrue(first.fatal_at_max) + self.assertEqual(first.selector.max_numb_sel, 7) + self.assertAlmostEqual(first.selector.report.level_f_lo, 0.1) + self.assertAlmostEqual(first.selector.report.level_f_hi, 0.2) + self.assertEqual(second.max_numb_iter, 9) + self.assertFalse(second.fatal_at_max) + self.assertEqual(second.selector.max_numb_sel, wf_config["fp"]["task_max"]) + self.assertAlmostEqual(second.selector.report.level_f_lo, 0.05) + self.assertAlmostEqual(second.selector.report.level_f_hi, 0.5) + self.assertIsNot(first.selector, second.selector) + self.assertIsNot(first.selector.report, second.selector.report) + def remove_executor_if_debug(conf): from dflow.config import ( diff --git a/tests/op/test_run_dp_train.py b/tests/op/test_run_dp_train.py index 45ba950c..06ebea96 100644 --- a/tests/op/test_run_dp_train.py +++ b/tests/op/test_run_dp_train.py @@ -3,6 +3,9 @@ import os import shutil import unittest +from copy import ( + deepcopy, +) from pathlib import ( Path, ) @@ -254,6 +257,12 @@ def test_normalize_config(self): self.assertAlmostEqual(config["init_model_start_pref_e"], 0.1) self.assertAlmostEqual(config["init_model_start_pref_f"], 100) self.assertAlmostEqual(config["init_model_start_pref_v"], 0.0) + self.assertIsNone(config["init_model_end_pref_e"]) + self.assertIsNone(config["init_model_end_pref_f"]) + self.assertIsNone(config["init_model_end_pref_v"]) + self.assertIsNone(config["finetune_end_pref_e"]) + self.assertIsNone(config["finetune_end_pref_f"]) + self.assertIsNone(config["finetune_end_pref_v"]) def test_get_size_of_all_mult_sys(self): cc = _get_data_size_of_all_mult_sys(self.iter_data) @@ -483,6 +492,41 @@ def test_update_input_dict_v2(self): ) self.assertDictEqual(odict, self.expected_odict_v2) + def test_stage_specific_end_preferences(self): + config = self.config.copy() + config.update( + { + "init_model_end_pref_e": 0.2, + "init_model_end_pref_f": 2.0, + "init_model_end_pref_v": 0.02, + "finetune_end_pref_e": 0.3, + "finetune_end_pref_f": 3.0, + "finetune_end_pref_v": 0.03, + } + ) + + init_model_dict = RunDPTrain.write_other_to_input_script( + deepcopy(self.idict_v2), + config, + True, + major_version="2", + finetune_mode="no", + ) + self.assertAlmostEqual(init_model_dict["loss"]["limit_pref_e"], 0.2) + self.assertAlmostEqual(init_model_dict["loss"]["limit_pref_f"], 2.0) + self.assertAlmostEqual(init_model_dict["loss"]["limit_pref_v"], 0.02) + + finetune_dict = RunDPTrain.write_other_to_input_script( + deepcopy(self.idict_v2), + config, + False, + major_version="2", + finetune_mode="finetune", + ) + self.assertAlmostEqual(finetune_dict["loss"]["limit_pref_e"], 0.3) + self.assertAlmostEqual(finetune_dict["loss"]["limit_pref_f"], 3.0) + self.assertAlmostEqual(finetune_dict["loss"]["limit_pref_v"], 0.03) + @patch("dpgen2.op.run_dp_train.run_command") def test_exec_v1(self, mocked_run): mocked_run.side_effect = [(0, "foo\n", ""), (0, "bar\n", "")]