fix pyright errors (#154)

Signed-off-by: zjgemi <liuxin_zijian@163.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>

dpgen2/conf/alloy_conf.py

@@ -166,7 +166,7 @@ def __init__(
             sys = sys.replicate(replicate)
         # set atom types
         self.ntypes = len(type_map)
-        self.natoms = sum(sys["atom_numbs"])
+        self.natoms = sum(sys["atom_numbs"])  # type: ignore
         sys.data["atom_names"] = type_map
         sys.data["atom_numbs"] = [0] * self.ntypes
         sys.data["atom_numbs"][0] = self.natoms
@@ -222,7 +222,7 @@ def generate_systems(
         concentration: Union[List[List[float]], List[float], None] = None,
         cell_pert_frac: float = 0.0,
         atom_pert_dist: float = 0.0,
-    ) -> List[str]:
+    ) -> List[dpdata.System]:
         """
         Parameters
         ----------

dpgen2/entrypoint/submit.py

@@ -16,6 +16,7 @@
 
 import dpdata
 from dflow import (
+    ArgoStep,
     InputArtifact,
     InputParameter,
     Inputs,
@@ -592,7 +593,7 @@ def copy_scheduler_plans(
 
 def submit_concurrent_learning(
     wf_config,
-    reuse_step: Optional[List[Step]] = None,
+    reuse_step: Optional[List[ArgoStep]] = None,
     old_style: bool = False,
     replace_scheduler: bool = False,
     no_submission: bool = False,

dpgen2/exploration/selector/conf_filter.py

@@ -58,8 +58,8 @@ def add(
     def check(
         self,
         conf: dpdata.System,
-    ) -> bool:
-        natoms = sum(conf["atom_numbs"])
+    ) -> dpdata.System:
+        natoms = sum(conf["atom_numbs"])  # type: ignore
         selected_idx = np.arange(conf.get_nframes())
         for ff in self._filters:
             fsel = np.where(

dpgen2/exploration/selector/conf_selector_frame.py

@@ -93,6 +93,6 @@ def select(
 
         out_path = Path("confs")
         out_path.mkdir(exist_ok=True)
-        ms.to_deepmd_npy(out_path)
+        ms.to_deepmd_npy(out_path)  # type: ignore
 
         return [out_path], copy.deepcopy(self.report)

dpgen2/flow/dpgen_loop.py

@@ -156,7 +156,7 @@ class ConcurrentLearningLoop(Steps):
     def __init__(
         self,
         name: str,
-        block_op: OPTemplate,
+        block_op: ConcurrentLearningBlock,
         step_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
     ):
@@ -240,7 +240,7 @@ class ConcurrentLearning(Steps):
     def __init__(
         self,
         name: str,
-        block_op: OPTemplate,
+        block_op: ConcurrentLearningBlock,
         step_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
     ):

dpgen2/fp/deepmd.py

@@ -137,7 +137,7 @@ def _get_dp_model(self, teacher_model_path: BinaryFileInput):
         from deepmd.infer import DeepPot  # type: ignore
 
         teacher_model_path.save_as_file(deepmd_teacher_model)
-        dp = DeepPot(deepmd_teacher_model)
+        dp = DeepPot(Path(deepmd_teacher_model))
 
         type_map_teacher = dp.get_type_map()
 
@@ -148,7 +148,7 @@ def _prep_input(self, type_map_teacher):
         ss = dpdata.System(deepmd_input_path, fmt="deepmd/npy")
         conf_type_map = ss["atom_names"]
 
-        if not set(conf_type_map).issubset(set(type_map_teacher)):
+        if not set(conf_type_map).issubset(set(type_map_teacher)):  # type: ignore
             err_message = (
                 f"the type map of system ({conf_type_map}) is not subset of "
                 + f"the type map of the teacher model ({type_map_teacher})."
@@ -157,7 +157,7 @@ def _prep_input(self, type_map_teacher):
 
         # make sure the order of elements in sys_type_map
         # is the same as that in type_map_teacher
-        temp_type_map = [ele for ele in type_map_teacher if ele in set(conf_type_map)]
+        temp_type_map = [ele for ele in type_map_teacher if ele in set(conf_type_map)]  # type: ignore
 
         ss.apply_type_map(temp_type_map)
         ss.to("deepmd/npy", deepmd_temp_path)
@@ -179,8 +179,8 @@ def _dp_infer(self, dp, type_map_teacher, out_name):
 
         nframe = ss.get_nframes()
         coord = ss["coords"]
-        cell = None if ss.nopbc else ss["cells"].reshape([nframe, -1])
-        atype = ss["atom_types"].tolist()
+        cell = None if ss.nopbc else ss["cells"].reshape([nframe, -1])  # type: ignore
+        atype = ss["atom_types"].tolist()  # type: ignore
 
         energy, force, virial_force = dp.eval(coord, cell, atype)
 

dpgen2/fp/prep_fp.py

@@ -115,7 +115,7 @@ def execute(
         # loop over list of MultiSystems
         for mm in confs:
             ms = dpdata.MultiSystems(type_map=type_map)
-            ms.from_deepmd_npy(mm, labeled=False)
+            ms.from_deepmd_npy(mm, labeled=False)  # type: ignore
             # loop over Systems in MultiSystems
             for ii in range(len(ms)):
                 ss = ms[ii]

dpgen2/fp/vasp.py

@@ -75,7 +75,7 @@ def prep_task(
         # 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"]))
-        Path(vasp_kp_name).write_text(vasp_inputs.make_kpoints(conf_frame["cells"][0]))
+        Path(vasp_kp_name).write_text(vasp_inputs.make_kpoints(conf_frame["cells"][0]))  # type: ignore
 
 
 class RunVasp(RunFp):

dpgen2/op/collect_data.py

@@ -98,9 +98,9 @@ def execute(
         # if ms.get_nframes() == 0, ms.to_deepmd_npy would not make the dir Path(name)
         Path(name).mkdir()
         if mixed_type:
-            ms.to_deepmd_npy_mixed(name)
+            ms.to_deepmd_npy_mixed(name)  # type: ignore
         else:
-            ms.to_deepmd_npy(name)
+            ms.to_deepmd_npy(name)  # type: ignore
         iter_data.append(Path(name))
 
         return OPIO(

dpgen2/op/run_dp_train.py

@@ -425,9 +425,9 @@ def _get_data_size_of_all_systems(data_dirs):
 def _get_data_size_of_mult_sys(data_dir, mixed_type=False):
     ms = dpdata.MultiSystems()
     if mixed_type:
-        ms.from_deepmd_npy_mixed(data_dir)
+        ms.from_deepmd_npy_mixed(data_dir)  # type: ignore
     else:
-        ms.from_deepmd_npy(data_dir)
+        ms.from_deepmd_npy(data_dir)  # type: ignore
     return ms.get_nframes()
 
 

dpgen2/superop/block.py

@@ -11,6 +11,7 @@
     List,
     Optional,
     Set,
+    Type,
 )
 
 from dflow import (
@@ -47,6 +48,16 @@
 )
 from dpgen2.utils.step_config import normalize as normalize_step_dict
 
+from .prep_run_dp_train import (
+    PrepRunDPTrain,
+)
+from .prep_run_fp import (
+    PrepRunFp,
+)
+from .prep_run_lmp import (
+    PrepRunLmp,
+)
+
 block_default_optional_parameter = {
     "data_mixed_type": False,
 }
@@ -68,11 +79,11 @@ class ConcurrentLearningBlock(Steps):
     def __init__(
         self,
         name: str,
-        prep_run_dp_train_op: OPTemplate,
-        prep_run_lmp_op: OPTemplate,
-        select_confs_op: OP,
-        prep_run_fp_op: OPTemplate,
-        collect_data_op: OP,
+        prep_run_dp_train_op: PrepRunDPTrain,
+        prep_run_lmp_op: PrepRunLmp,
+        select_confs_op: Type[OP],
+        prep_run_fp_op: PrepRunFp,
+        collect_data_op: Type[OP],
         select_confs_config: dict = normalize_step_dict({}),
         collect_data_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
@@ -172,9 +183,9 @@ def _block_cl(
     name: str,
     prep_run_dp_train_op: OPTemplate,
     prep_run_lmp_op: OPTemplate,
-    select_confs_op: OP,
+    select_confs_op: Type[OP],
     prep_run_fp_op: OPTemplate,
-    collect_data_op: OP,
+    collect_data_op: Type[OP],
     select_confs_config: dict = normalize_step_dict({}),
     collect_data_config: dict = normalize_step_dict({}),
     upload_python_packages: Optional[List[os.PathLike]] = None,

dpgen2/superop/prep_run_dp_train.py

@@ -9,6 +9,7 @@
     List,
     Optional,
     Set,
+    Type,
 )
 
 from dflow import (
@@ -39,6 +40,9 @@
 from dpgen2.constants import (
     train_index_pattern,
 )
+from dpgen2.op import (
+    RunDPTrain,
+)
 from dpgen2.utils.step_config import (
     init_executor,
 )
@@ -49,8 +53,8 @@ class PrepRunDPTrain(Steps):
     def __init__(
         self,
         name: str,
-        prep_train_op: OP,
-        run_train_op: OP,
+        prep_train_op: Type[OP],
+        run_train_op: Type[RunDPTrain],
         prep_config: dict = normalize_step_dict({}),
         run_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
@@ -131,8 +135,8 @@ def keys(self):
 def _prep_run_dp_train(
     train_steps,
     step_keys,
-    prep_train_op: OP,
-    run_train_op: OP,
+    prep_train_op: Type[OP],
+    run_train_op: Type[OP],
     prep_config: dict = normalize_step_dict({}),
     run_config: dict = normalize_step_dict({}),
     upload_python_packages: Optional[List[os.PathLike]] = None,

dpgen2/superop/prep_run_fp.py

@@ -9,6 +9,7 @@
     List,
     Optional,
     Set,
+    Type,
 )
 
 from dflow import (
@@ -49,8 +50,8 @@ class PrepRunFp(Steps):
     def __init__(
         self,
         name: str,
-        prep_op: OP,
-        run_op: OP,
+        prep_op: Type[OP],
+        run_op: Type[OP],
         prep_config: dict = normalize_step_dict({}),
         run_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
@@ -124,8 +125,8 @@ def keys(self):
 def _prep_run_fp(
     prep_run_steps,
     step_keys,
-    prep_op: OP,
-    run_op: OP,
+    prep_op: Type[OP],
+    run_op: Type[OP],
     prep_config: dict = normalize_step_dict({}),
     run_config: dict = normalize_step_dict({}),
     upload_python_packages: Optional[List[os.PathLike]] = None,

dpgen2/superop/prep_run_lmp.py

@@ -9,6 +9,7 @@
     List,
     Optional,
     Set,
+    Type,
 )
 
 from dflow import (
@@ -49,8 +50,8 @@ class PrepRunLmp(Steps):
     def __init__(
         self,
         name: str,
-        prep_op: OP,
-        run_op: OP,
+        prep_op: Type[OP],
+        run_op: Type[OP],
         prep_config: dict = normalize_step_dict({}),
         run_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
@@ -128,8 +129,8 @@ def keys(self):
 def _prep_run_lmp(
     prep_run_steps,
     step_keys,
-    prep_op: OP,
-    run_op: OP,
+    prep_op: Type[OP],
+    run_op: Type[OP],
     prep_config: dict = normalize_step_dict({}),
     run_config: dict = normalize_step_dict({}),
     upload_python_packages: Optional[List[os.PathLike]] = None,

tests/test_prep_run_dp_labeling.py

@@ -151,7 +151,7 @@ def test_get_dp_model(self):
             _dp, _type_map = run_deepmd._get_dp_model(self.teacher_model)
             self.assertTrue(_dp is dp)
             self.assertEqual(_type_map, ["H", "C"])
-            deepmd.infer.DeepPot.assert_called_once_with(deepmd_teacher_model)
+            deepmd.infer.DeepPot.assert_called_once_with(Path(deepmd_teacher_model))
             self.assertFalse(Path(deepmd_teacher_model).is_file())
 
     def test_dp_infer(self):