Support electronic temperature (#258)

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

- **New Features**
- Introduced a new argument for handling electronic temperature
configurations, enhancing simulation flexibility.
- Added the capability to specify optional output paths, improving
output management.
- Expanded functionality for retrieving configurations with new
temperature parameters.
- Enhanced exploration scheduling with elemental temperature
configurations.

- **Bug Fixes**
- Improved data validation processes to ensure consistency across
trajectories and outputs.

- **Documentation**
- Updated documentation to reflect new parameters and functionalities
across multiple modules.

- **Tests**
- Added unit tests to validate the functionality of temperature setup
and related methods.
<!-- 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>

dpgen2/entrypoint/args.py

@@ -529,6 +529,7 @@ def input_args():
     doc_valid_data_prefix = "The prefix of validation data systems"
     doc_valid_sys = "The validation data systems"
     doc_valid_data_uri = "The URI of validation data"
+    doc_use_ele_temp = "Whether to use electronic temperature, 0 for no, 1 for frame temperature, and 2 for atomic temperature"
     doc_multi_valid_data = (
         "The validation data for multitask, it should be a dict, whose keys are task names and each value is a dict"
         "containing fields `prefix` and `sys` for initial data of each task"
@@ -612,6 +613,13 @@ def input_args():
             default=None,
             doc=doc_valid_data_uri,
         ),
+        Argument(
+            "use_ele_temp",
+            int,
+            optional=True,
+            default=0,
+            doc=doc_use_ele_temp,
+        ),
         Argument(
             "multi_valid_data",
             dict,

dpgen2/entrypoint/submit.py

@@ -370,11 +370,12 @@ def make_lmp_naive_exploration_scheduler(config):
     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)
+    render = TrajRenderLammps(nopbc=output_nopbc, use_ele_temp=use_ele_temp)
     # selector
     selector = ConfSelectorFrames(
         render,
@@ -625,6 +626,9 @@ def workflow_concurrent_learning(
     else:
         init_models = None
 
+    if config["inputs"]["use_ele_temp"]:
+        explore_config["use_ele_temp"] = config["inputs"]["use_ele_temp"]
+
     optional_parameter = make_optional_parameter(
         config["inputs"]["mixed_type"],
     )

dpgen2/exploration/render/traj_render.py

@@ -52,6 +52,7 @@ def get_confs(
         id_selected: List[List[int]],
         type_map: Optional[List[str]] = None,
         conf_filters: Optional["ConfFilters"] = None,
+        optional_outputs: Optional[List[Path]] = None,
     ) -> dpdata.MultiSystems:
         r"""Get configurations from trajectory by selection.
 
@@ -64,6 +65,10 @@ def get_confs(
             from the ii-th trajectory. id_selected[ii] may be an empty list.
         type_map : List[str]
             The type map.
+        conf_filters : ConfFilters
+            Configuration filters
+        optional_outputs : List[Path]
+            Optional outputs of the exploration
 
         Returns
         -------

dpgen2/exploration/render/traj_render_lammps.py

@@ -1,3 +1,4 @@
+import json
 from pathlib import (
     Path,
 )
@@ -12,6 +13,10 @@
 import dpdata
 import numpy as np
 
+from dpgen2.utils import (
+    setup_ele_temp,
+)
+
 from ..deviation import (
     DeviManager,
     DeviManagerStd,
@@ -30,8 +35,10 @@ class TrajRenderLammps(TrajRender):
     def __init__(
         self,
         nopbc: bool = False,
+        use_ele_temp: int = 0,
     ):
         self.nopbc = nopbc
+        self.use_ele_temp = use_ele_temp
 
     def get_model_devi(
         self,
@@ -57,20 +64,54 @@ def _load_one_model_devi(self, fname, model_devi):
         model_devi.add(DeviManager.MIN_DEVI_F, dd[:, 5])
         model_devi.add(DeviManager.AVG_DEVI_F, dd[:, 6])
 
+    def get_ele_temp(self, optional_outputs):
+        ele_temp = []
+        for ii in range(len(optional_outputs)):
+            with open(optional_outputs[ii], "r") as f:
+                data = json.load(f)
+            if self.use_ele_temp:
+                ele_temp.append(data["ele_temp"])
+        if self.use_ele_temp:
+            if self.use_ele_temp == 1:
+                setup_ele_temp(False)
+            elif self.use_ele_temp == 2:
+                setup_ele_temp(True)
+            else:
+                raise ValueError(
+                    "Invalid value for 'use_ele_temp': %s", self.use_ele_temp
+                )
+        return ele_temp
+
+    def set_ele_temp(self, system, ele_temp):
+        if self.use_ele_temp == 1 and ele_temp:
+            system.data["fparam"] = np.tile(ele_temp, [len(system), 1])
+        elif self.use_ele_temp == 2 and ele_temp:
+            system.data["aparam"] = np.tile(
+                ele_temp, [len(system), system.get_natoms(), 1]
+            )
+
     def get_confs(
         self,
         trajs: List[Path],
         id_selected: List[List[int]],
         type_map: Optional[List[str]] = None,
         conf_filters: Optional["ConfFilters"] = None,
+        optional_outputs: Optional[List[Path]] = None,
     ) -> dpdata.MultiSystems:
         ntraj = len(trajs)
+        ele_temp = None
+        if optional_outputs:
+            assert ntraj == len(optional_outputs)
+            ele_temp = self.get_ele_temp(optional_outputs)
+
         traj_fmt = "lammps/dump"
         ms = dpdata.MultiSystems(type_map=type_map)
         for ii in range(ntraj):
             if len(id_selected[ii]) > 0:
                 ss = dpdata.System(trajs[ii], fmt=traj_fmt, type_map=type_map)
                 ss.nopbc = self.nopbc
+                if ele_temp:
+                    self.set_ele_temp(ss, ele_temp[ii])
                 ss = ss.sub_system(id_selected[ii])
                 ms.append(ss)
         if conf_filters is not None:

dpgen2/exploration/selector/conf_selector.py

@@ -32,5 +32,6 @@ def select(
         trajs: List[Path],
         model_devis: List[Path],
         type_map: Optional[List[str]] = None,
+        optional_outputs: Optional[List[Path]] = None,
     ) -> Tuple[List[Path], ExplorationReport]:
         pass

dpgen2/exploration/selector/conf_selector_frame.py

@@ -55,6 +55,7 @@ def select(
         trajs: List[Path],
         model_devis: List[Path],
         type_map: Optional[List[str]] = None,
+        optional_outputs: Optional[List[Path]] = None,
     ) -> Tuple[List[Path], ExplorationReport]:
         """Select configurations
 
@@ -69,6 +70,8 @@ def select(
             where `md` stands for model deviation, v for virial and f for force
         type_map : List[str]
             The `type_map` of the systems
+        optional_outputs : List[Path]
+            Optional outputs of the exploration
 
         Returns
         -------
@@ -88,7 +91,11 @@ def select(
         id_cand_list = self.report.get_candidate_ids(self.max_numb_sel)
 
         ms = self.traj_render.get_confs(
-            trajs, id_cand_list, type_map, self.conf_filters
+            trajs,
+            id_cand_list,
+            type_map,
+            self.conf_filters,
+            optional_outputs,
         )
 
         out_path = Path("confs")

dpgen2/fp/prep_fp.py

@@ -30,6 +30,7 @@
 )
 from dpgen2.utils import (
     set_directory,
+    setup_ele_temp,
 )
 
 
@@ -114,6 +115,10 @@ def execute(
         counter = 0
         # loop over list of MultiSystems
         for mm in confs:
+            if len(list(mm.rglob("fparam.npy"))) > 0:
+                setup_ele_temp(False)
+            if len(list(mm.rglob("aparam.npy"))) > 0:
+                setup_ele_temp(True)
             ms = dpdata.MultiSystems(type_map=type_map)
             ms.from_deepmd_npy(mm, labeled=False)  # type: ignore
             # loop over Systems in MultiSystems

dpgen2/fp/vasp.py

@@ -1,4 +1,7 @@
+import json
 import logging
+import os
+import re
 from pathlib import (
     Path,
 )
@@ -33,8 +36,9 @@
     fp_default_log_name,
     fp_default_out_data_name,
 )
-from dpgen2.utils.run_command import (
+from dpgen2.utils import (
     run_command,
+    setup_ele_temp,
 )
 
 from .prep_fp import (
@@ -55,7 +59,71 @@
 vasp_kp_name = "KPOINTS"
 
 
+def clean_lines(string_list, remove_empty_lines=True):
+    """[migrated from pymatgen]
+    Strips whitespace, carriage returns and empty lines from a list of strings.
+
+    Args:
+        string_list: List of strings
+        remove_empty_lines: Set to True to skip lines which are empty after
+            stripping.
+
+    Returns:
+        List of clean strings with no whitespaces.
+    """
+    for s in string_list:
+        clean_s = s
+        if "#" in s:
+            ind = s.index("#")
+            clean_s = s[:ind]
+        clean_s = clean_s.strip()
+        if (not remove_empty_lines) or clean_s != "":
+            yield clean_s
+
+
+def loads_incar(incar: str):
+    lines = list(clean_lines(incar.splitlines()))
+    params = {}
+    for line in lines:
+        for sline in line.split(";"):
+            m = re.match(r"(\w+)\s*=\s*(.*)", sline.strip())
+            if m:
+                key = m.group(1).strip()
+                val = m.group(2).strip()
+                params[key.upper()] = val
+    return params
+
+
+def dumps_incar(params: dict):
+    incar = "\n".join([key + " = " + str(val) for key, val in params.items()]) + "\n"
+    return incar
+
+
 class PrepVasp(PrepFp):
+    def set_ele_temp(self, conf_frame, incar):
+        use_ele_temp = 0
+        ele_temp = None
+        if "fparam" in conf_frame.data:
+            use_ele_temp = 1
+            ele_temp = conf_frame.data["fparam"][0][0]
+        if "aparam" in conf_frame.data:
+            use_ele_temp = 2
+            ele_temp = conf_frame.data["aparam"][0][0][0]
+        if ele_temp:
+            import scipy.constants as pc
+
+            params = loads_incar(incar)
+            params["ISMEAR"] = -1
+            params["SIGMA"] = ele_temp * pc.Boltzmann / pc.electron_volt
+            incar = dumps_incar(params)
+            data = {
+                "use_ele_temp": use_ele_temp,
+                "ele_temp": ele_temp,
+            }
+            with open("job.json", "w") as f:
+                json.dump(data, f, indent=4)
+        return incar
+
     def prep_task(
         self,
         conf_frame: dpdata.System,
@@ -72,7 +140,10 @@ def prep_task(
         """
 
         conf_frame.to("vasp/poscar", vasp_conf_name)
-        Path(vasp_input_name).write_text(vasp_inputs.incar_template)
+        incar = vasp_inputs.incar_template
+        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
         tmp_frame = dpdata.System(vasp_conf_name, fmt="vasp/poscar")
         Path(vasp_pot_name).write_text(vasp_inputs.make_potcar(tmp_frame["atom_names"]))
@@ -100,7 +171,21 @@ def optional_input_files(self) -> List[str]:
             A list of optional input files names.
 
         """
-        return []
+        return ["job.json"]
+
+    def set_ele_temp(self, system):
+        if os.path.exists("job.json"):
+            with open("job.json", "r") as f:
+                data = json.load(f)
+            if "use_ele_temp" in data and "ele_temp" in data:
+                if data["use_ele_temp"] == 1:
+                    setup_ele_temp(False)
+                    system.data["fparam"] = np.tile(data["ele_temp"], [1, 1])
+                elif data["use_ele_temp"] == 2:
+                    setup_ele_temp(True)
+                    system.data["aparam"] = np.tile(
+                        data["ele_temp"], [1, system.get_natoms(), 1]
+                    )
 
     def run_task(
         self,
@@ -141,6 +226,7 @@ def run_task(
             raise TransientError("vasp failed")
         # convert the output to deepmd/npy format
         sys = dpdata.LabeledSystem("OUTCAR")
+        self.set_ele_temp(sys)
         sys.to("deepmd/npy", out_name)
         return out_name, log_name
 

dpgen2/op/collect_data.py

@@ -19,6 +19,10 @@
     Parameter,
 )
 
+from dpgen2.utils import (
+    setup_ele_temp,
+)
+
 
 class CollectData(OP):
     """Collect labeled data and add to the iteration dataset.
@@ -91,6 +95,10 @@ def execute(
 
         ms = dpdata.MultiSystems(type_map=type_map)
         for ii in labeled_data:
+            if len(list(ii.rglob("fparam.npy"))) > 0:
+                setup_ele_temp(False)
+            if len(list(ii.rglob("aparam.npy"))) > 0:
+                setup_ele_temp(True)
             ss = dpdata.LabeledSystem(ii, fmt="deepmd/npy")
             ms.append(ss)