add fp_style custom

dpgen/generator/arginfo.py

@@ -789,6 +789,35 @@ def fp_style_amber_diff_args() -> List[Argument]:
     ]
 
 
+def fp_style_custom_args() -> List[Argument]:
+    """Arguments for FP style custom.
+
+    Returns
+    -------
+    list[dargs.Argument]
+        list of Gaussian fp style arguments
+    """
+    doc_fp_params_custom = "Parameters for FP calculation."
+    doc_input_fmt = "Input dpdata format of the custom FP code. Such format should only need the first argument as the file name."
+    doc_output_fmt = "Output dpata format of the custom FP code. Such format should only need the first argument as the file name."
+    doc_input_fn = "Input file name of the custom FP code."
+    doc_output_fn = "Output file name of the custom FP code."
+    return [
+        Argument(
+            "fp_params",
+            dict,
+            optional=False,
+            doc=doc_fp_params_custom,
+            sub_fields=[
+                Argument("input_fmt", str, optional=False, doc=doc_input_fmt),
+                Argument("input_fn", str, optional=False, doc=doc_input_fn),
+                Argument("output_fmt", str, optional=False, doc=doc_output_fmt),
+                Argument("output_fn", str, optional=False, doc=doc_output_fn),
+            ],
+        ),
+    ]
+
+
 def fp_style_variant_type_args() -> Variant:
     doc_fp_style = "Software for First Principles."
     doc_amber_diff = (
@@ -798,6 +827,11 @@ def fp_style_variant_type_args() -> Variant:
         "The command argument in the machine file should be path to sander. "
         "One should also install dpamber and make it visible in the PATH."
     )
+    doc_custom = (
+        "Custom FP code. You need to provide the input and output file format and name. "
+        "The command argument in the machine file should be the script to run custom FP codes. "
+        "The extra forward and backward files can be defined in the machine file."
+    )
 
     return Variant(
         "fp_style",
@@ -812,6 +846,7 @@ def fp_style_variant_type_args() -> Variant:
             ),
             Argument("pwmat", dict, [], doc="TODO: add doc"),
             Argument("pwscf", dict, [], doc="TODO: add doc"),
+            Argument("custom", dict, fp_style_custom_args(), doc=doc_custom),
         ],
         optional=False,
         doc=doc_fp_style,

dpgen/generator/run.py

@@ -23,6 +23,7 @@
 import warnings
 from collections import Counter
 from collections.abc import Iterable
+from pathlib import Path
 from typing import List
 
 import dpdata
@@ -88,7 +89,13 @@
     write_incar_dict,
 )
 from dpgen.remote.decide_machine import convert_mdata
-from dpgen.util import convert_training_data_to_hdf5, expand_sys_str, normalize, sepline
+from dpgen.util import (
+    convert_training_data_to_hdf5,
+    expand_sys_str,
+    normalize,
+    sepline,
+    set_directory,
+)
 
 from .arginfo import run_jdata_arginfo
 
@@ -3534,6 +3541,30 @@ def make_fp_amber_diff(iter_index: int, jdata: dict):
     os.chdir(cwd)
 
 
+def make_fp_custom(iter_index, jdata):
+    """Make input file for customized FP style.
+
+    Convert the POSCAR file to custom format.
+
+    Parameters
+    ----------
+    iter_index : int
+        iter index
+    jdata : dict
+        Run parameters.
+    """
+    work_path = os.path.join(make_iter_name(iter_index), fp_name)
+    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
+    fp_params = jdata["fp_params"]
+    input_fn = fp_params["input_fn"]
+    input_fmt = fp_params["input_fmt"]
+
+    for ii in fp_tasks:
+        with set_directory(Path(ii)):
+            system = dpdata.System("POSCAR", fmt="vasp/poscar")
+            system.to(input_fmt, input_fn)
+
+
 def make_fp(iter_index, jdata, mdata):
     """Select the candidate strutures and make the input file of FP calculation.
 
@@ -3581,6 +3612,8 @@ def make_fp_calculation(iter_index, jdata, mdata):
         make_fp_pwmat(iter_index, jdata)
     elif fp_style == "amber/diff":
         make_fp_amber_diff(iter_index, jdata)
+    elif fp_style == "custom":
+        make_fp_custom(iter_index, jdata)
     else:
         raise RuntimeError("unsupported fp style")
     # Copy user defined forward_files
@@ -3890,6 +3923,19 @@ def run_fp(iter_index, jdata, mdata):
             log_file="output",
             forward_common_files=forward_common_files,
         )
+    elif fp_style == "custom":
+        fp_params = jdata["fp_params"]
+        forward_files = [fp_params["input_fn"]]
+        backward_files = [fp_params["output_fn"]]
+        run_fp_inner(
+            iter_index,
+            jdata,
+            mdata,
+            forward_files,
+            backward_files,
+            None,
+            log_file="output",
+        )
     else:
         raise RuntimeError("unsupported fp style")
 
@@ -4354,6 +4400,51 @@ def post_fp_amber_diff(iter_index, jdata):
         all_sys.to_deepmd_npy(sys_data_path, set_size=len(sys_output), prec=np.float64)
 
 
+def post_fp_custom(iter_index, jdata):
+    """Post fp for custom fp. Collect data from user-defined `output_fn`.
+
+    Parameters
+    ----------
+    iter_index : int
+        The index of the current iteration.
+    jdata : dict
+        The parameter data.
+    """
+    model_devi_jobs = jdata["model_devi_jobs"]
+    assert iter_index < len(model_devi_jobs)
+
+    iter_name = make_iter_name(iter_index)
+    work_path = os.path.join(iter_name, fp_name)
+    fp_tasks = glob.glob(os.path.join(work_path, "task.*"))
+    fp_tasks.sort()
+    if len(fp_tasks) == 0:
+        return
+
+    system_index = []
+    for ii in fp_tasks:
+        system_index.append(os.path.basename(ii).split(".")[1])
+    system_index.sort()
+    set_tmp = set(system_index)
+    system_index = list(set_tmp)
+    system_index.sort()
+
+    fp_params = jdata["fp_params"]
+    output_fn = fp_params["output_fn"]
+    output_fmt = fp_params["output_fmt"]
+
+    for ss in system_index:
+        sys_output = glob.glob(os.path.join(work_path, "task.%s.*" % ss))
+        sys_output.sort()
+        all_sys = dpdata.MultiSystems(type_map=jdata["type_map"])
+        for oo in sys_output:
+            if os.path.exists(os.path.join(oo, output_fn)):
+                sys = dpdata.LabeledSystem(os.path.join(oo, output_fn), fmt=output_fmt)
+                all_sys.append(sys)
+        sys_data_path = os.path.join(work_path, "data.%s" % ss)
+        all_sys.to_deepmd_raw(sys_data_path)
+        all_sys.to_deepmd_npy(sys_data_path, set_size=len(sys_output), prec=np.float64)
+
+
 def post_fp(iter_index, jdata):
     fp_style = jdata["fp_style"]
     if fp_style == "vasp":
@@ -4372,6 +4463,8 @@ def post_fp(iter_index, jdata):
         post_fp_pwmat(iter_index, jdata)
     elif fp_style == "amber/diff":
         post_fp_amber_diff(iter_index, jdata)
+    elif fp_style == "custom":
+        post_fp_custom(iter_index, jdata)
     else:
         raise RuntimeError("unsupported fp style")
     post_fp_check_fail(iter_index, jdata)

dpgen/simplify/arginfo.py

@@ -8,6 +8,7 @@
     data_args,
     fp_style_abacus_args,
     fp_style_cp2k_args,
+    fp_style_custom_args,
     fp_style_gaussian_args,
     fp_style_siesta_args,
     fp_style_vasp_args,
@@ -66,6 +67,11 @@ def fp_style_variant_type_args() -> Variant:
     doc_fp_style_none = "No fp."
     doc_fp_style_vasp = "VASP."
     doc_fp_style_gaussian = "Gaussian. The command should be set as `g16 < input`."
+    doc_custom = (
+        "Custom FP code. You need to provide the input and output file format and name. "
+        "The command argument in the machine file should be the script to run custom FP codes. "
+        "The extra forward and backward files can be defined in the machine file."
+    )
 
     return Variant(
         "fp_style",
@@ -85,6 +91,7 @@ def fp_style_variant_type_args() -> Variant:
             # ),
             Argument("pwmat", dict, [], doc="TODO: add doc"),
             Argument("pwscf", dict, [], doc="TODO: add doc"),
+            Argument("custom", dict, fp_style_custom_args(), doc=doc_custom),
         ],
         optional=True,
         default_tag="none",

dpgen/util.py

@@ -1,6 +1,9 @@
 #!/usr/bin/env python
 import json
 import os
+from contextlib import (
+    contextmanager,
+)
 from pathlib import Path
 from typing import List, Union
 
@@ -145,3 +148,30 @@ def convert_training_data_to_hdf5(input_files: List[str], h5_file: str):
                 group = f.create_group(str(pp))
                 s = dpdata.LabeledSystem(ii, fmt="deepmd/npy")
                 s.to("deepmd/hdf5", group)
+
+
+@contextmanager
+def set_directory(path: Path):
+    """Sets the current working path within the context.
+
+    Parameters
+    ----------
+    path : Path
+        The path to the cwd
+
+    Yields
+    ------
+    None
+
+    Examples
+    --------
+    >>> with set_directory("some_path"):
+    ...    do_something()
+    """
+    cwd = Path().absolute()
+    path.mkdir(exist_ok=True, parents=True)
+    try:
+        os.chdir(path)
+        yield
+    finally:
+        os.chdir(cwd)

tests/generator/context.py

@@ -40,6 +40,7 @@
 param_diy_abacus_post_file = "param-methane-abacus-diy.json"
 param_amber_file = "param-amber.json"
 param_multiple_trust_file = "param-mg-vasp-multi-trust.json"
+param_custom_fp_file = "param-custom-fp.json"
 
 
 def my_file_cmp(test, f0, f1):

tests/generator/param-custom-fp.json

@@ -0,0 +1,107 @@
+{    
+    "type_map":		["C", "H", "N"],
+    "mass_map":		[16, 2, 14],
+
+    "init_data_prefix":	"/home/linfengz/SCR/wanghan/deepgen.pyridine/init",
+    "init_data_sys":	["Pyridine-I",
+			 "Pyridine-II"
+			],
+    "init_batch_size":	[1, 1],
+    "sys_configs":	[
+	["/home/linfengz/SCR/wanghan/data/pyridine/pyI.POSCAR.01x01x01/01.scale_pert/sys-0080-0080-0016/scale-1.000/00009?/POSCAR"],
+	["/home/linfengz/SCR/wanghan/data/pyridine/pyI.POSCAR.01x01x01/01.scale_pert/sys-0080-0080-0016/scale-1.000/0000[7-8]?/POSCAR"],
+	["/home/linfengz/SCR/wanghan/data/pyridine/pyI.POSCAR.01x01x01/01.scale_pert/sys-0080-0080-0016/scale-1.000/0000[5-6]?/POSCAR"],
+	["/home/linfengz/SCR/wanghan/data/pyridine/pyI.POSCAR.01x01x01/01.scale_pert/sys-0080-0080-0016/scale-1.000/0000[0-4]?/POSCAR"],
+	["/home/linfengz/SCR/wanghan/data/pyridine/pyII.POSCAR.01x01x01/01.scale_pert/sys-0080-0080-0016/scale-1.000/00009?/POSCAR"],
+	["/home/linfengz/SCR/wanghan/data/pyridine/pyII.POSCAR.01x01x01/01.scale_pert/sys-0080-0080-0016/scale-1.000/0000[7-8]?/POSCAR"],
+	["/home/linfengz/SCR/wanghan/data/pyridine/pyII.POSCAR.01x01x01/01.scale_pert/sys-0080-0080-0016/scale-1.000/0000[5-6]?/POSCAR"],
+	["/home/linfengz/SCR/wanghan/data/pyridine/pyII.POSCAR.01x01x01/01.scale_pert/sys-0080-0080-0016/scale-1.000/0000[0-4]?/POSCAR"]
+    ],
+    "_comment":		"0  1  2  3",
+    "_comment":		"4  5  6  7",
+    "sys_batch_size":	[1, 1, 1, 1,
+			 1, 1, 1, 1
+			],
+
+    "_comment":		" 00.train ",
+    "numb_models":	4,
+    "default_training_param" : {
+	"_comment": " model parameters",
+	"use_smooth":		true,
+	"sel_a":		[81, 81, 20],
+	"rcut_smth":		0.50,
+	"rcut":			6.50,
+	"filter_neuron":	[25, 50, 100],
+	"filter_resnet_dt":	false,
+	"n_axis_neuron":	12,
+	"n_neuron":		[240, 240, 240],
+	"resnet_dt":		true,
+	"coord_norm":		true,
+	"type_fitting_net":	false,
+
+	"_comment": " traing controls",
+	"systems":		[],
+	"set_prefix":		"set",    
+	"stop_batch":		400000,
+	"batch_size":		1,
+	"start_lr":		0.002,
+	"decay_steps":		2000,
+	"decay_rate":		0.95,
+	"seed":			0,
+
+	"start_pref_e":		0.02,
+	"limit_pref_e":		2,
+	"start_pref_f":		1000,
+	"limit_pref_f":		1,
+	"start_pref_v":		0.0,
+	"limit_pref_v":		0.0,
+
+	"_comment": " display and restart",
+	"_comment": " frequencies counted in batch",
+	"disp_file":		"lcurve.out",
+	"disp_freq":		2000,
+	"numb_test":		10,
+	"save_freq":		20000,
+	"save_ckpt":		"model.ckpt",
+	"load_ckpt":		"model.ckpt",
+	"disp_training":	true,
+	"time_training":	true,
+	"profiling":		false,
+	"profiling_file":	"timeline.json",
+
+	"_comment":		"that's all"
+    },
+
+    "_comment":		" 01.model_devi ",
+    "_comment": "model_devi_skip: the first x of the recorded frames",
+    "model_devi_dt":		0.001,
+    "model_devi_skip":		0,
+    "model_devi_f_trust_lo":	0.050,
+    "model_devi_f_trust_hi":	0.150,
+    "model_devi_clean_traj":	false,
+    "model_devi_jobs":	[
+	{"sys_idx": [0,4], "temps": [  50], "press": [1e0,1e1,1e2,1e3,1e4,2e4,4e4], "trj_freq": 10, "nsteps": 1000, "ensemble": "npt", "_idx": "00"},
+	{"sys_idx": [1,5], "temps": [  50], "press": [1e0,1e1,1e2,1e3,1e4,2e4,4e4], "trj_freq": 10, "nsteps": 1000, "ensemble": "npt", "_idx": "01"},
+	{"sys_idx": [0,4], "temps": [  50], "press": [1e0,1e1,1e2,1e3,1e4,2e4,4e4], "trj_freq": 10, "nsteps": 1000, "ensemble": "npt", "_idx": "02"},
+	{"sys_idx": [1,5], "temps": [  50], "press": [1e0,1e1,1e2,1e3,1e4,2e4,4e4], "trj_freq": 10, "nsteps": 1000, "ensemble": "npt", "_idx": "03"},
+	{"sys_idx": [0,4], "temps": [ 100], "press": [1e0,1e1,1e2,1e3,1e4,2e4,4e4], "trj_freq": 10, "nsteps": 1000, "ensemble": "npt", "_idx": "04"},
+	{"sys_idx": [1,5], "temps": [ 100], "press": [1e0,1e1,1e2,1e3,1e4,2e4,4e4], "trj_freq": 10, "nsteps": 1000, "ensemble": "npt", "_idx": "05"},
+	{"sys_idx": [0,4], "temps": [ 100], "press": [1e0,1e1,1e2,1e3,1e4,2e4,4e4], "trj_freq": 10, "nsteps": 1000, "ensemble": "npt", "_idx": "06"},
+	{"sys_idx": [1,5], "temps": [ 100], "press": [1e0,1e1,1e2,1e3,1e4,2e4,4e4], "trj_freq": 10, "nsteps": 1000, "ensemble": "npt", "_idx": "07"}
+    ],
+
+    "_comment":		" 02.fp ",
+    "fp_style":		"custom",
+    "shuffle_poscar":	false,
+    "fp_task_max":	100,
+    "fp_task_min":	10,
+    "fp_pp_path":	".",
+    "fp_pp_files":	[],
+    "fp_params":	{
+		"input_fn": "input.h5",
+		"input_fmt": "deepmd/hdf5",
+		"output_fn": "output.h5",
+		"output_fmt": "deepmd/hdf5"
+    },
+    "_comment":		" that's all "
+}