From ce4ab3e6deb980611cc8c397b444c7a9dfaed883 Mon Sep 17 00:00:00 2001 From: Xinzijian Liu Date: Wed, 4 Sep 2024 15:32:14 +0800 Subject: [PATCH] Support electronic temperature (#258) ## 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. --------- Signed-off-by: zjgemi Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- dpgen2/entrypoint/args.py | 8 ++ dpgen2/entrypoint/submit.py | 6 +- dpgen2/exploration/render/traj_render.py | 5 + .../exploration/render/traj_render_lammps.py | 41 ++++++++ dpgen2/exploration/selector/conf_selector.py | 1 + .../selector/conf_selector_frame.py | 9 +- dpgen2/fp/prep_fp.py | 5 + dpgen2/fp/vasp.py | 92 +++++++++++++++++- dpgen2/op/collect_data.py | 8 ++ dpgen2/op/run_lmp.py | 40 ++++++++ dpgen2/op/select_confs.py | 40 ++++++-- dpgen2/superop/block.py | 3 + dpgen2/superop/prep_run_lmp.py | 12 ++- dpgen2/utils/__init__.py | 3 + dpgen2/utils/setup_ele_temp.py | 33 +++++++ pyproject.toml | 2 +- tests/exploration/test_traj_render_lammps.py | 62 ++++++++++++ tests/fp/test_vasp.py | 95 +++++++++++++++++++ tests/mocked_ops.py | 2 + tests/op/test_run_lmp.py | 23 +++++ tests/test_select_confs.py | 26 ++++- tests/utils/test_ele_temp.py | 66 +++++++++++++ 22 files changed, 565 insertions(+), 17 deletions(-) create mode 100644 dpgen2/utils/setup_ele_temp.py create mode 100644 tests/exploration/test_traj_render_lammps.py create mode 100644 tests/utils/test_ele_temp.py diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index 645e8c6b..4ba69ea2 100644 --- a/dpgen2/entrypoint/args.py +++ b/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, diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index ce414cfb..a6911ef0 100644 --- a/dpgen2/entrypoint/submit.py +++ b/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"], ) diff --git a/dpgen2/exploration/render/traj_render.py b/dpgen2/exploration/render/traj_render.py index 91703b77..eb7296b6 100644 --- a/dpgen2/exploration/render/traj_render.py +++ b/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 ------- diff --git a/dpgen2/exploration/render/traj_render_lammps.py b/dpgen2/exploration/render/traj_render_lammps.py index 89a0f7da..d51ec040 100644 --- a/dpgen2/exploration/render/traj_render_lammps.py +++ b/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: diff --git a/dpgen2/exploration/selector/conf_selector.py b/dpgen2/exploration/selector/conf_selector.py index 3f1061bb..df00afd3 100644 --- a/dpgen2/exploration/selector/conf_selector.py +++ b/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 diff --git a/dpgen2/exploration/selector/conf_selector_frame.py b/dpgen2/exploration/selector/conf_selector_frame.py index cacd9ed5..74eee689 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/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") diff --git a/dpgen2/fp/prep_fp.py b/dpgen2/fp/prep_fp.py index b908278e..a43f582a 100644 --- a/dpgen2/fp/prep_fp.py +++ b/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 diff --git a/dpgen2/fp/vasp.py b/dpgen2/fp/vasp.py index 4dfe6fd2..a8ab6b88 100644 --- a/dpgen2/fp/vasp.py +++ b/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 diff --git a/dpgen2/op/collect_data.py b/dpgen2/op/collect_data.py index 838a695d..b1e057e5 100644 --- a/dpgen2/op/collect_data.py +++ b/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) diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 2f60631d..2822a325 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -79,6 +79,7 @@ def get_output_sign(cls): "traj": Artifact(Path), "model_devi": Artifact(Path), "plm_output": Artifact(Path, optional=True), + "optional_output": Artifact(Path, optional=True), } ) @@ -183,6 +184,16 @@ def execute( ) raise TransientError("lmp failed") + ele_temp = None + if config.get("use_ele_temp", 0): + ele_temp = get_ele_temp(lmp_log_name) + if ele_temp is not None: + data = { + "ele_temp": ele_temp, + } + with open("job.json", "w") as f: + json.dump(data, f, indent=4) + ret_dict = { "log": work_dir / lmp_log_name, "traj": work_dir / lmp_traj_name, @@ -194,6 +205,8 @@ def execute( else {} ) ret_dict.update(plm_output) + if ele_temp is not None: + ret_dict["optional_output"] = work_dir / "job.json" return OPIO(ret_dict) @@ -203,6 +216,7 @@ def lmp_args(): doc_teacher_model = "The teacher model in `Knowledge Distillation`" doc_shuffle_models = "Randomly pick a model from the group of models to drive theexploration MD simulation" doc_head = "Select a head from multitask" + doc_use_ele_temp = "Whether to use electronic temperature, 0 for no, 1 for frame temperature, and 2 for atomic temperature" return [ Argument("command", str, optional=True, default="lmp", doc=doc_lmp_cmd), Argument( @@ -219,6 +233,10 @@ def lmp_args(): default=False, doc=doc_shuffle_models, ), + Argument("head", str, optional=True, default=None, doc=doc_head), + Argument( + "use_ele_temp", int, optional=True, default=0, doc=doc_use_ele_temp + ), Argument( "model_frozen_head", str, optional=True, default=None, doc=doc_head ), @@ -293,6 +311,28 @@ def find_only_one_key(lmp_lines, key, raise_not_found=True): return found[0] +def get_ele_temp(lmp_log_name): + with open(lmp_log_name, encoding="utf8") as f: + lmp_log_lines = f.readlines() + + for line in lmp_log_lines: + fields = line.split() + if fields[:2] == ["pair_style", "deepmd"]: + if "fparam" in fields: + # for rendering variables + try: + return float(fields[fields.index("fparam") + 1]) + except Exception: + pass + if "aparam" in fields: + try: + return float(fields[fields.index("aparam") + 1]) + except Exception: + pass + + return None + + def freeze_model(input_model, frozen_model, head=None): freeze_args = "-o %s" % frozen_model if head is not None: diff --git a/dpgen2/op/select_confs.py b/dpgen2/op/select_confs.py index 502aad91..fea0bf59 100644 --- a/dpgen2/op/select_confs.py +++ b/dpgen2/op/select_confs.py @@ -37,6 +37,7 @@ def get_input_sign(cls): "type_map": List[str], "trajs": Artifact(List[Path]), "model_devis": Artifact(List[Path]), + "optional_outputs": Artifact(List[Path], optional=True), } ) @@ -80,12 +81,16 @@ def execute( trajs = ip["trajs"] model_devis = ip["model_devis"] - trajs, model_devis = SelectConfs.validate_trajs(trajs, model_devis) + optional_outputs = ip["optional_outputs"] + trajs, model_devis, optional_outputs = SelectConfs.validate_trajs( + trajs, model_devis, optional_outputs + ) confs, report = conf_selector.select( trajs, model_devis, type_map=type_map, + optional_outputs=optional_outputs, ) return OPIO( @@ -99,18 +104,41 @@ def execute( def validate_trajs( trajs, model_devis, + optional_outputs=None, ): ntrajs = len(trajs) if ntrajs != len(model_devis): raise FatalError( "length of trajs list is not equal to the " "model_devis list" ) + if optional_outputs and ntrajs != len(optional_outputs): + raise FatalError( + "length of trajs list is not equal to the " "optional_output list" + ) rett = [] retm = [] - for tt, mm in zip(trajs, model_devis): # type: ignore - if (tt is None and mm is not None) or (tt is not None and mm is None): - raise FatalError("trajs frame is {tt} while model_devis frame is {mm}") - elif tt is not None and mm is not None: + reto = [] + for i in range(ntrajs): + tt = trajs[i] + mm = model_devis[i] + if tt is not None and mm is not None: rett.append(tt) retm.append(mm) - return rett, retm + if optional_outputs: + oo = optional_outputs[i] + if oo is not None: + reto.append(oo) + else: + raise FatalError( + f"trajs frame is {tt} while optional_outputs frame is {oo}" + ) + elif tt is None and mm is None: + if optional_outputs: + oo = optional_outputs[i] + if oo is not None: + raise FatalError( + f"trajs frame is {tt} while optional_outputs frame is {oo}" + ) + else: + raise FatalError(f"trajs frame is {tt} while model_devis frame is {mm}") + return rett, retm, reto diff --git a/dpgen2/superop/block.py b/dpgen2/superop/block.py index eea566d2..0e39ab38 100644 --- a/dpgen2/superop/block.py +++ b/dpgen2/superop/block.py @@ -266,6 +266,9 @@ def _block_cl( artifacts={ "trajs": prep_run_explore.outputs.artifacts["trajs"], "model_devis": prep_run_explore.outputs.artifacts["model_devis"], + "optional_outputs": prep_run_explore.outputs.artifacts["optional_outputs"] + if "optional_outputs" in prep_run_explore.outputs.artifacts + else None, }, key=step_keys["select-confs"], executor=select_confs_executor, diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py index 6b4de94a..48d4028a 100644 --- a/dpgen2/superop/prep_run_lmp.py +++ b/dpgen2/superop/prep_run_lmp.py @@ -75,6 +75,7 @@ def __init__( "trajs": OutputArtifact(), "model_devis": OutputArtifact(), "plm_output": OutputArtifact(), + "optional_outputs": OutputArtifact(), } super().__init__( @@ -172,7 +173,13 @@ def _prep_run_lmp( "int('{{item}}')", input_parameter=["task_name"], input_artifact=["task_path"], - output_artifact=["log", "traj", "model_devi", "plm_output"], + output_artifact=[ + "log", + "traj", + "model_devi", + "plm_output", + "optional_output", + ], **template_slice_config, ), python_packages=upload_python_packages, @@ -207,5 +214,8 @@ def _prep_run_lmp( prep_run_steps.outputs.artifacts["plm_output"]._from = run_lmp.outputs.artifacts[ "plm_output" ] + prep_run_steps.outputs.artifacts[ + "optional_outputs" + ]._from = run_lmp.outputs.artifacts["optional_output"] return prep_run_steps diff --git a/dpgen2/utils/__init__.py b/dpgen2/utils/__init__.py index 6eebbe98..b7824f57 100644 --- a/dpgen2/utils/__init__.py +++ b/dpgen2/utils/__init__.py @@ -34,6 +34,9 @@ from .run_command import ( run_command, ) +from .setup_ele_temp import ( + setup_ele_temp, +) from .step_config import gen_doc as gen_doc_step_dict from .step_config import ( init_executor, diff --git a/dpgen2/utils/setup_ele_temp.py b/dpgen2/utils/setup_ele_temp.py new file mode 100644 index 00000000..7a218886 --- /dev/null +++ b/dpgen2/utils/setup_ele_temp.py @@ -0,0 +1,33 @@ +import dpdata +import numpy as np +from dpdata.data_type import ( + Axis, + DataType, +) + + +def setup_ele_temp(atomic: bool): + """Set electronic temperature as required input data. + + Parameters + ---------- + atomic : bool + Whether to use atomic temperature or frame temperature + """ + if atomic: + ele_temp_data_type = DataType( + "aparam", + np.ndarray, + shape=(Axis.NFRAMES, Axis.NATOMS, 1), + required=False, + ) + else: + ele_temp_data_type = DataType( + "fparam", + np.ndarray, + shape=(Axis.NFRAMES, 1), + required=False, + ) + + dpdata.System.register_data_type(ele_temp_data_type) + dpdata.LabeledSystem.register_data_type(ele_temp_data_type) diff --git a/pyproject.toml b/pyproject.toml index 33bdd90c..e327dec3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ classifiers = [ ] dependencies = [ 'numpy', - 'dpdata', + 'dpdata>=0.2.20', 'pydflow>=1.6.57', 'dargs>=0.3.1', 'scipy', diff --git a/tests/exploration/test_traj_render_lammps.py b/tests/exploration/test_traj_render_lammps.py new file mode 100644 index 00000000..3659ffd8 --- /dev/null +++ b/tests/exploration/test_traj_render_lammps.py @@ -0,0 +1,62 @@ +import json +import os +import unittest + +import dpdata +import numpy as np + +# isort: off +from .context import ( + dpgen2, +) +from dpgen2.exploration.render import TrajRenderLammps + +# isort: on + + +class TestTrajRenderLammps(unittest.TestCase): + def test_use_ele_temp_1(self): + with open("job.json", "w") as f: + json.dump({"ele_temp": 6.6}, f) + traj_render = TrajRenderLammps(use_ele_temp=1) + ele_temp = traj_render.get_ele_temp(["job.json"]) + self.assertEqual(ele_temp, [6.6]) + + system = dpdata.System( + data={ + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": np.zeros(1, dtype=int), + "cells": np.eye(3).reshape(1, 3, 3), + "coords": np.zeros((1, 1, 3)), + "orig": np.zeros(3), + "nopbc": True, + } + ) + traj_render.set_ele_temp(system, ele_temp[0]) + np.testing.assert_array_almost_equal(system.data["fparam"], np.array([[6.6]])) + + def test_use_ele_temp_2(self): + with open("job.json", "w") as f: + json.dump({"ele_temp": 6.6}, f) + traj_render = TrajRenderLammps(use_ele_temp=2) + ele_temp = traj_render.get_ele_temp(["job.json"]) + self.assertEqual(ele_temp, [6.6]) + + system = dpdata.System( + data={ + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": np.zeros(1, dtype=int), + "cells": np.eye(3).reshape(1, 3, 3), + "coords": np.zeros((1, 1, 3)), + "orig": np.zeros(3), + "nopbc": True, + } + ) + traj_render.set_ele_temp(system, ele_temp[0]) + np.testing.assert_array_almost_equal(system.data["aparam"], np.array([[[6.6]]])) + + def tearDown(self): + if os.path.exists("job.json"): + os.remove("job.json") diff --git a/tests/fp/test_vasp.py b/tests/fp/test_vasp.py index f0c25e9b..8d9bee25 100644 --- a/tests/fp/test_vasp.py +++ b/tests/fp/test_vasp.py @@ -17,12 +17,44 @@ dpgen2, ) from dpgen2.fp.vasp import ( + PrepVasp, + RunVasp, VaspInputs, + dumps_incar, + loads_incar, make_kspacing_kpoints, ) +from dpgen2.utils import ( + setup_ele_temp, +) # isort: on +INCAR = """ +PREC=A +ENCUT=600 +ISYM=0 +ALGO=fast +EDIFF=1.000000e-06 +LREAL=A +NPAR=1 +KPAR=1 + +NELMIN=4 +ISIF=2 +ISMEAR=1 +SIGMA=1.000000 +IBRION=-1 + +NSW=0 + +LWAVE=F +LCHARG=F +PSTRESS=0 + +KSPACING=0.160000 +KGAMMA=.FALSE.""" + class TestVASPInputs(unittest.TestCase): def setUp(self): @@ -119,3 +151,66 @@ def test_vasp_input_kp(self): ss = dpdata.System("POSCAR") kps = vi.make_kpoints(ss["cells"][0]) self.assertEqual(ref, kps) + + +class TestIncar(unittest.TestCase): + def test_loads_dumps_incar(self): + incar = INCAR + params = loads_incar(incar) + self.assertEqual(len(params), 19) + self.assertEqual(params["PREC"], "A") + self.assertEqual(params["EDIFF"], "1.000000e-06") + self.assertEqual(params["IBRION"], "-1") + self.assertEqual(params["KGAMMA"], ".FALSE.") + new_incar = dumps_incar(params) + new_params = loads_incar(new_incar) + self.assertEqual(params, new_params) + + +class TestPrepVasp(unittest.TestCase): + def test_set_ele_temp(self): + setup_ele_temp(False) + frame = dpdata.System( + data={ + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": np.zeros(1, dtype=int), + "cells": np.eye(3).reshape(1, 3, 3), + "coords": np.zeros((1, 1, 3)), + "orig": np.zeros(3), + "nopbc": True, + "fparam": np.array([[6.6]]), + } + ) + op = PrepVasp() + incar = INCAR + incar = op.set_ele_temp(frame, incar) + params = loads_incar(incar) + self.assertEqual(int(params["ISMEAR"]), -1) + self.assertAlmostEqual(float(params["SIGMA"]), 0.0005687439953015817) + + +class TestRunVasp(unittest.TestCase): + def test_set_ele_temp(self): + with open("job.json", "w") as f: + json.dump({"use_ele_temp": 1, "ele_temp": 6.6}, f) + system = dpdata.LabeledSystem( + data={ + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": np.zeros(1, dtype=int), + "cells": np.eye(3).reshape(1, 3, 3), + "coords": np.zeros((1, 1, 3)), + "orig": np.zeros(3), + "energies": np.zeros(1), + "forces": np.zeros((1, 1, 3)), + "nopbc": True, + } + ) + op = RunVasp() + op.set_ele_temp(system) + np.testing.assert_array_almost_equal(system.data["fparam"], [[6.6]]) + + def tearDown(self): + if os.path.exists("job.json"): + os.remove("job.json") diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py index fbd2308c..9cd13c00 100644 --- a/tests/mocked_ops.py +++ b/tests/mocked_ops.py @@ -19,6 +19,7 @@ ) from typing import ( List, + Optional, Tuple, ) @@ -864,6 +865,7 @@ def select( trajs: List[Path], model_devis: List[Path], type_map: List[str] = None, + optional_outputs: Optional[List[Path]] = None, ) -> Tuple[List[Path], ExplorationReport]: confs = [] if len(trajs) == mocked_numb_lmp_tasks: diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index de885ce7..b727fb76 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -34,6 +34,7 @@ ) from dpgen2.op.run_lmp import ( RunLmp, + get_ele_temp, set_models, ) from dpgen2.utils import ( @@ -263,3 +264,25 @@ def test_failed_no_matching(self): input_name.write_text(lmp_config) with self.assertRaises(RuntimeError) as re: set_models(input_name, self.model_names) + + +class TestGetEleTemp(unittest.TestCase): + def test_get_ele_temp_none(self): + with open("log", "w") as f: + f.write( + "pair_style deepmd model.000.pb model.001.pb model.002.pb model.003.pb model.004.pb out_freq 10 out_file model_devi.out" + ) + ele_temp = get_ele_temp("log") + self.assertIsNone(ele_temp) + + def test_get_ele_temp(self): + with open("log", "w") as f: + f.write( + "pair_style deepmd model.000.pb model.001.pb model.002.pb model.003.pb model.004.pb out_freq 10 out_file model_devi.out fparam 6.6" + ) + ele_temp = get_ele_temp("log") + self.assertEqual(ele_temp, 6.6) + + def tearDown(self): + if os.path.exists("log"): + os.remove("log") diff --git a/tests/test_select_confs.py b/tests/test_select_confs.py index c8840eaf..491d42f7 100644 --- a/tests/test_select_confs.py +++ b/tests/test_select_confs.py @@ -129,21 +129,39 @@ def test(self): def test_validate_trajs(self): trajs = ["foo", "bar", None, "tar"] model_devis = ["zar", "par", None, "mar"] - trajs, model_devis = SelectConfs.validate_trajs(trajs, model_devis) + trajs, model_devis, _ = SelectConfs.validate_trajs(trajs, model_devis, None) self.assertEqual(trajs, ["foo", "bar", "tar"]) self.assertEqual(model_devis, ["zar", "par", "mar"]) trajs = ["foo", "bar", None, "tar"] model_devis = ["zar", "par", None] with self.assertRaises(FatalError) as context: - trajs, model_devis = SelectConfs.validate_trajs(trajs, model_devis) + trajs, model_devis, _ = SelectConfs.validate_trajs(trajs, model_devis, None) trajs = ["foo", "bar"] model_devis = ["zar", None] with self.assertRaises(FatalError) as context: - trajs, model_devis = SelectConfs.validate_trajs(trajs, model_devis) + trajs, model_devis, _ = SelectConfs.validate_trajs(trajs, model_devis, None) trajs = ["foo", None] model_devis = ["zar", "par"] with self.assertRaises(FatalError) as context: - trajs, model_devis = SelectConfs.validate_trajs(trajs, model_devis) + trajs, model_devis, _ = SelectConfs.validate_trajs(trajs, model_devis, None) + + trajs = ["foo", "bar", None, "tar"] + model_devis = ["zar", "par", None, "mar"] + optional_outputs = ["dar", "far", None, "gar"] + trajs, model_devis, optional_outputs = SelectConfs.validate_trajs( + trajs, model_devis, optional_outputs + ) + self.assertEqual(trajs, ["foo", "bar", "tar"]) + self.assertEqual(model_devis, ["zar", "par", "mar"]) + self.assertEqual(optional_outputs, ["dar", "far", "gar"]) + + trajs = ["foo", "bar"] + model_devis = ["zar", "par"] + optional_outputs = ["dar", None] + with self.assertRaises(FatalError) as context: + trajs, model_devis, optional_outputs = SelectConfs.validate_trajs( + trajs, model_devis, optional_outputs + ) diff --git a/tests/utils/test_ele_temp.py b/tests/utils/test_ele_temp.py new file mode 100644 index 00000000..6a7de791 --- /dev/null +++ b/tests/utils/test_ele_temp.py @@ -0,0 +1,66 @@ +import glob +import os +import shutil +import unittest + +import dpdata +import numpy as np + +# isort: off +from .context import ( + dpgen2, +) +from dpgen2.utils import ( + setup_ele_temp, +) + +# isort: on + + +class TestSetupEleTemp(unittest.TestCase): + def test_setup_ele_temp_unlabeled(self): + system = dpdata.System( + data={ + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": np.zeros(1, dtype=int), + "cells": np.eye(3).reshape(1, 3, 3), + "coords": np.zeros((1, 1, 3)), + "orig": np.zeros(3), + "nopbc": True, + } + ) + setup_ele_temp(False) + system.data["fparam"] = np.array([[1.0]]) + system.to_deepmd_npy("ele_temp_data") + self.assertEqual(len(glob.glob("ele_temp_data/*/fparam.npy")), 1) + new_system = dpdata.System("ele_temp_data", fmt="deepmd/npy") + self.assertTrue("fparam" in new_system.data) + + def test_setup_ele_temp_mixed(self): + system = dpdata.System( + data={ + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": np.zeros(1, dtype=int), + "cells": np.eye(3).reshape(1, 3, 3), + "coords": np.zeros((1, 1, 3)), + "orig": np.zeros(3), + "nopbc": True, + } + ) + setup_ele_temp(True) + system.data["aparam"] = np.array([[[1.0]]]) + system.to_deepmd_npy_mixed("ele_temp_mixed_data") + self.assertEqual(len(glob.glob("ele_temp_mixed_data/*/aparam.npy")), 1) + ms = dpdata.MultiSystems() + ms.load_systems_from_file( + "ele_temp_mixed_data", fmt="deepmd/npy/mixed", labeled=False + ) + self.assertTrue("aparam" in ms[0].data) + + def tearDown(self): + if os.path.exists("ele_temp_data"): + shutil.rmtree("ele_temp_data") + if os.path.exists("ele_temp_mixed_data"): + shutil.rmtree("ele_temp_mixed_data")