Merge branch 'master' into feature/add_cp2k

fpop/abacus.py

@@ -517,7 +517,7 @@ class PrepAbacus(PrepFp):
     def prep_task(
             self,
             conf_frame,
-            abacus_inputs: AbacusInputs,
+            inputs: AbacusInputs,
             prepare_image_config: Optional[Dict] = None,
             optional_input: Optional[Dict] = None,
             optional_artifact: Optional[Dict] = None,
@@ -539,13 +539,13 @@ def prep_task(
         """
 
         element_list = conf_frame['atom_names']
-        pp, orb = abacus_inputs.write_pporb(element_list)
-        dpks = abacus_inputs.write_deepks()
-        mass = abacus_inputs.get_mass(element_list)
+        pp, orb = inputs.write_pporb(element_list)
+        dpks = inputs.write_deepks()
+        mass = inputs.get_mass(element_list)
         conf_frame.to('abacus/stru', 'STRU', pp_file=pp,numerical_orbital=orb,numerical_descriptor=dpks,mass=mass)
 
-        abacus_inputs.write_input("INPUT")
-        abacus_inputs.write_kpt("KPT")
+        inputs.write_input("INPUT")
+        inputs.write_kpt("KPT")
 
         if optional_artifact:
             for file_name, file_path in optional_artifact.items():
@@ -636,7 +636,7 @@ def run_task(
         if run_image_config:
             kwargs.update(run_image_config)
             kwargs.pop("command", None)
-        ret, out, err = run_command(command, raise_error=False, **kwargs)
+        ret, out, err = run_command(command, raise_error=False, **kwargs) # type: ignore
         if ret != 0:
             raise TransientError(
                 "abacus failed\n", "out msg", out, "\n", "err msg", err, "\n"

fpop/prep_fp.py

@@ -25,7 +25,7 @@
 class PrepFp(OP, ABC):
     r"""Prepares the working directories for first-principles (FP) tasks.
 
-    A list of (same length as ip["confs"]) working directories
+    A list of (same length as op_in["confs"]) working directories
     containing all files needed to start FP tasks will be
     created. The paths of the directories will be returned as
     `op["task_paths"]`. The identities of the tasks are returned as
@@ -86,13 +86,13 @@ def prep_task(
     @OP.exec_sign_check
     def execute(
             self,
-            ip : OPIO,
+            op_in : OPIO,
     ) -> OPIO:
         r"""Execute the OP.
 
         Parameters
         ----------
-        ip : dict
+        op_in : dict
             Input dict with components:
 
             - `prep_image_config` : (`dict`) It defines the parameters of the process of preparing FP tasks.
@@ -118,12 +118,12 @@ def execute(
         """
         import dpdata
 
-        inputs = ip['inputs']
-        confs = ip['confs']
-        type_map = ip['type_map']
-        prepare_image_config = ip["prep_image_config"]
-        optional_artifact = ip["optional_artifact"]
-        optional_input = ip["optional_input"]
+        inputs = op_in['inputs']
+        confs = op_in['confs']
+        type_map = op_in['type_map']
+        prepare_image_config = op_in["prep_image_config"]
+        optional_artifact = op_in["optional_artifact"]
+        optional_input = op_in["optional_input"]
         try:
             conf_format = optional_input["conf_format"]
         except:

fpop/preprun_fp.py

@@ -1,3 +1,4 @@
+from abc import ABCMeta
 from dflow import (
     InputParameter,
     OutputParameter,
@@ -25,7 +26,7 @@
 
 from dflow.plugins.dispatcher import DispatcherExecutor
 import os,sys
-from typing import Optional, Set, List
+from typing import Optional, Set, List, Union
 from pathlib import Path
 from fpop.utils.step_config import (
     init_executor,
@@ -35,16 +36,16 @@ class PrepRunFp(Steps):
     def __init__(
         self,
         name : str,
-        prep_op : OP,
-        run_op : OP,
+        prep_op : Union[OP, ABCMeta],
+        run_op : Union[OP, ABCMeta],
         prep_image : str,
         run_image : str,
         prep_template_config : Optional[dict] = None,
         prep_step_config : Optional[dict] = None,
         run_template_config : Optional[dict] = None,
         run_slice_config : Optional[dict] = None,
         run_step_config : Optional[dict] = None,
-        upload_python_packages : Optional[List[str]] = None,
+        upload_python_packages : Optional[Union[List[Path], List[str]]] = None,
     ):
         self._input_parameters = {
             "inputs" : InputParameter(),
@@ -112,16 +113,16 @@ def keys(self):
 def _prep_run_fp(
         prep_run_steps,
         step_keys,
-        prep_op : OP,
-        run_op : OP,
+        prep_op : Union[OP, ABCMeta],
+        run_op : Union[OP, ABCMeta],
         prep_image,
         run_image,
         prep_template_config : Optional[dict] = None,
         prep_step_config : Optional[dict] = None,
         run_template_config : Optional[dict] = None,
         run_slice_config : Optional[dict] = None,
         run_step_config : Optional[dict] = None,
-        upload_python_packages : Optional[List[str]] = None,
+        upload_python_packages : Optional[Union[List[Path], List[str]]] = None,
 ):
     if not prep_template_config: prep_template_config = {}
     if not prep_step_config: prep_step_config = {}
@@ -144,7 +145,7 @@ def _prep_run_fp(
             output_artifact_archive={
                 "task_paths": None
             },
-            python_packages = upload_python_packages,
+            python_packages = upload_python_packages, # type: ignore
             image = prep_image,
             **prep_template_config,
         ),
@@ -175,7 +176,7 @@ def _prep_run_fp(
                 output_artifact = ["backward_dir"],
                 **run_slice_config,
             ),
-            python_packages = upload_python_packages,
+            python_packages = upload_python_packages, # type: ignore
             image = run_image,
             **run_template_config,
         ),

fpop/run_fp.py

@@ -109,12 +109,12 @@ def run_task(
     @OP.exec_sign_check
     def execute(
         self,
-        ip: OPIO,
+        op_in: OPIO,
     ) -> OPIO:
         r'''Execute the OP.
         Parameters
         ----------
-        ip : dict
+        op_in : dict
             Input dict with components:
             - `task_name`: (`str`) The name of task.
             - `task_path`: (`Artifact(Path)`) The path that contains all input files prepareed by `PrepFp`.
@@ -140,19 +140,19 @@ def execute(
         FatalError
             When mandatory files are not found.
         '''
-        run_image_config = ip["run_image_config"]
-        backward_dir_name = ip["backward_dir_name"] 
-        log_name = ip["log_name"] 
-        backward_list = ip["backward_list"]
-        optional_input = ip["optional_input"]
-        task_name = ip["task_name"]
-        task_path = ip["task_path"]
+        run_image_config = op_in["run_image_config"]
+        backward_dir_name = op_in["backward_dir_name"] 
+        log_name = op_in["log_name"] 
+        backward_list = op_in["backward_list"]
+        optional_input = op_in["optional_input"]
+        task_name = op_in["task_name"]
+        task_path = op_in["task_path"]
         input_files = self.input_files(task_path)
         input_files = [(Path(task_path) / ii).resolve() for ii in input_files]
         work_dir = Path(task_name)
         opt_input_files = []
-        if ip["optional_artifact"]:
-            for ss,vv in ip["optional_artifact"].items():
+        if op_in["optional_artifact"]:
+            for ss,vv in op_in["optional_artifact"].items():
                 opt_input_files.append(ss)
         opt_input_files = [(Path(task_path) / ii).resolve() for ii in opt_input_files]
 

fpop/vasp.py

@@ -155,7 +155,7 @@ class PrepVasp(PrepFp):
     def prep_task(
             self,
             conf_frame: dpdata.System,
-            vasp_inputs: VaspInputs,
+            inputs: VaspInputs,
             prepare_image_config: Optional[Dict] = None,
             optional_input: Optional[Dict] = None,
             optional_artifact: Optional[Dict] = None,
@@ -186,15 +186,15 @@ def prep_task(
 
         conf_frame.to('vasp/poscar', 'POSCAR')
         Path('INCAR').write_text(
-            vasp_inputs.incar_template
+            inputs.incar_template
         )
         # fix the case when some element have 0 atom, e.g. H0O2
         tmp_frame = dpdata.System('POSCAR', fmt='vasp/poscar')
         Path('POTCAR').write_text(
-            vasp_inputs.make_potcar(tmp_frame['atom_names'])
+            inputs.make_potcar(tmp_frame['atom_names'])
         )
         Path('KPOINTS').write_text(
-            vasp_inputs.make_kpoints(conf_frame['cells'][0])
+            inputs.make_kpoints(conf_frame['cells'][0])
         )
 
         if optional_artifact:
@@ -257,7 +257,7 @@ def run_task(
         if run_image_config:
             kwargs.update(run_image_config)
             kwargs.pop("command", None)
-        ret, out, err = run_command(command, raise_error=False, **kwargs)
+        ret, out, err = run_command(command, raise_error=False, **kwargs) # type: ignore
         if ret != 0:
             raise TransientError(
                 "vasp failed\n", "out msg", out, "\n", "err msg", err, "\n"

tests/constants.py

@@ -72,7 +72,7 @@ def dump_conf_from_poscar(
         confs = []
         for ii in range(len(conf_list)):
             ls = dpdata.System("POSCAR_%d"%ii, fmt="vasp/poscar")
-            ls.to_deepmd_npy("data.%03d"%ii)
+            ls.to_deepmd_npy("data.%03d"%ii) # type: ignore
             confs.append("data.%03d"%ii)
             os.remove("POSCAR_%d"%ii)
         return confs

tests/mocked_ops.py

@@ -20,16 +20,16 @@ class MockedRunVasp(RunVasp):
     @OP.exec_sign_check
     def execute(
             self,
-            ip : OPIO,
+            op_in : OPIO,
     ) -> OPIO:
-        task_name = ip["task_name"]
-        task_path = ip["task_path"]
+        task_name = op_in["task_name"]
+        task_path = op_in["task_path"]
 
-        assert(ip['task_path'].is_dir())
-        assert(re.match('task.[0-9][0-9][0-9][0-9][0-9][0-9]', ip['task_name']))
-        assert(ip['task_name'] in str(ip['task_path']))
+        assert(op_in['task_path'].is_dir())
+        assert(re.match('task.[0-9][0-9][0-9][0-9][0-9][0-9]', op_in['task_name']))
+        assert(op_in['task_name'] in str(op_in['task_path']))
         for ii in ['INCAR','POSCAR','POTCAR','KPOINTS']:
-            assert((ip['task_path']/ii).is_file())
+            assert((op_in['task_path']/ii).is_file())
 
         work_dir = Path(task_name)
 
@@ -42,16 +42,16 @@ def execute(
             if not Path(Path(ii).name).exists():
                 Path(Path(ii).name).symlink_to(ii)
 
-        backward_dir = Path(ip['backward_dir_name'])
+        backward_dir = Path(op_in['backward_dir_name'])
         backward_dir.mkdir(exist_ok=True, parents=True)
-        log = Path(Path(ip["backward_dir_name"])/ip["log_name"])
+        log = Path(Path(op_in["backward_dir_name"])/op_in["log_name"])
 
         fc=[]
         for ii in ['POSCAR','INCAR']:
             fc.append(Path(ii).read_text())
         log.write_text('\n'.join(fc))
 
-        for ii in ip["backward_list"]:
+        for ii in op_in["backward_list"]:
             fc = []
             fc.append(f'This is {ii} which users need in {task_name}')
             fc.append(Path('KPOINTS').read_text())
@@ -73,7 +73,7 @@ def run_task(
         run_image_config: Optional[Dict]=None,
         optional_input: Optional[Dict]=None,
     ):
-        pass
+        return "mock"
 
 class TestInputFiles2(RunFp):
     def input_files(self, task_path) -> List[str]:
@@ -87,4 +87,4 @@ def run_task(
         run_image_config: Optional[Dict]=None,
         optional_input: Optional[Dict]=None,
     ):
-        pass
+        return "mock"

tests/test_prep_run_vasp.py

@@ -73,7 +73,7 @@ def dump_conf_from_poscar(
         confs = []
         for ii in range(len(conf_list)):
             ls = dpdata.System("POSCAR_%d"%ii, fmt="vasp/poscar")
-            ls.to_deepmd_npy("data.%03d"%ii)
+            ls.to_deepmd_npy("data.%03d"%ii) # type: ignore
             confs.append("data.%03d")
             os.remove("POSCAR_%d"%ii)
         return confs