fix: fix the pyright check

dpgen2/exploration/task/calypso/caly_input.py

@@ -242,7 +242,7 @@ def make_calypso_input(
     name_of_atoms: List[str],
     atomic_number: List[int],
     numb_of_atoms: List[int],
-    distance_of_ions: List[List[float]],
+    distance_of_ions,
     pop_size: int = 30,
     max_step: int = 5,
     system_name: str = "CALYPSO",

dpgen2/exploration/task/make_task_group_from_config.py

@@ -555,4 +555,4 @@ def make_lmp_task_group_from_config(
 
 
 if __name__ == "__main__":
-    print(normalize({"type": "lmp-md"}))
+    print(lmp_normalize({"type": "lmp-md"}))

dpgen2/op/prep_run_dp_optim.py

@@ -169,7 +169,7 @@ def execute(
                     target = traj_results_dir.joinpath(str(cnt_num) + "." + traj.name)
                     shutil.copyfile(traj, target)
 
-            elif finished != "false":
+            else:
                 optim_results_dir = Path("optim_results_dir")
                 optim_results_dir.mkdir(parents=True, exist_ok=True)
                 traj_results_dir = Path("traj_results")

dpgen2/op/run_caly_model_devi.py

@@ -124,8 +124,8 @@ def execute(
                         f.write(dump_str)
                         pbc = np.all(atoms.get_pbc())
                         coord = atoms.get_positions().reshape(1, -1)
-                        cell = atoms.get_cell().reshape(1, -1) if pbc else None
-                        atype = [type_map.index(atom.symbol) for atom in atoms]
+                        cell = atoms.get_cell().array.reshape(1, -1) if pbc else None
+                        atype = [type_map.index(atom.symbol) for atom in atoms]  # type: ignore
                         devi = calc_model_devi(coord, cell, atype, graphs)[0]
                         devi[0] = tcount
                         Devis.append(devi)
@@ -189,21 +189,21 @@ def atoms2lmpdump(atoms, struc_idx, type_map):
     dump_str += "%20.10f %20.10f %20.10f\n" % (zlo_bound, zhi_bound, yz)
     dump_str += "ITEM: ATOMS id type x y z fx fy fz\n"
     for idx, atom in enumerate(new_atoms):
-        type_id = type_map.index(atom.symbol) + 1
+        type_id = type_map.index(atom.symbol) + 1  # type: ignore
         dump_str += "%5d %5d" % (idx + 1, type_id)
         dump_str += "%20.10f %20.10f %20.10f" % (
-            atom.position[0],
-            atom.position[1],
-            atom.position[2],
+            atom.position[0],  # type: ignore
+            atom.position[1],  # type: ignore
+            atom.position[2],  # type: ignore
         )
         dump_str += "%20.10f %20.10f %20.10f\n" % (0, 0, 0)
     # dump_str = dump_str.strip("\n")
     return dump_str
 
 
-def parse_traj(traj_file) -> Union[None, List[Atoms]]:
+def parse_traj(traj_file):
     # optimization will at least return one structures in traj file
-    trajs = read(traj_file, index=":", format="traj")
+    trajs: List[Atoms] = read(traj_file, index=":", format="traj")  # type: ignore
 
     numb_traj = len(trajs)
     assert numb_traj >= 1, "traj file is broken."
@@ -216,6 +216,7 @@ def parse_traj(traj_file) -> Union[None, List[Atoms]]:
     dis_mtx[row, col] = np.nan
     is_reasonable = np.nanmin(dis_mtx) > 0.6
 
+    selected_traj: Union[List[Atoms], None] = None
     if is_reasonable:
         if len(trajs) >= 20:
             selected_traj = [trajs[iii] for iii in [4, 9, -10, -5, -1]]
@@ -236,7 +237,7 @@ def parse_traj(traj_file) -> Union[None, List[Atoms]]:
     return selected_traj
 
 
-def write_model_devi_out(devi: np.ndarray, fname: str, header: str = ""):
+def write_model_devi_out(devi: np.ndarray, fname: Union[str, Path], header: str = ""):
     assert devi.shape[1] == 7
     header = "%s\n%10s" % (header, "step")
     for item in "vf":

dpgen2/superop/prep_run_calypso.py

@@ -57,7 +57,7 @@ def __init__(
         self,
         name: str,
         prep_caly_input_op: Type[OP],
-        caly_evo_step_op: Type[CalyEvoStep],
+        caly_evo_step_op: OPTemplate,
         run_caly_model_devi_op: Type[OP],
         prep_config: dict = normalize_step_dict({}),
         run_config: dict = normalize_step_dict({}),

dpgen2/superop/prep_run_lmp.py

@@ -60,7 +60,7 @@ def __init__(
             "block_id": InputParameter(type=str, value=""),
             "explore_config": InputParameter(),
             "expl_task_grp": InputParameter(),
-            "type_map": InputParameter(value=""),
+            "type_map": InputParameter(value=[]),
         }
         self._input_artifacts = {
             "models": InputArtifact(),