fix type hints errors; add type checker (#92)

* fix type hints errors; add type checker

This commit uses microsoft/pyright as a type checker. Several errors have been fixed.
Sometimes the type of a varible is too "dynamic". It will be helpful to use `assert` to narrow types. Finally, one can add `# type: ignore` to skip type checker of a line.

* remove optional from `incar_template_name` and `potcar_names`

* fix tests

* fix tests; fix docstring

* fix tests

.github/workflows/pyright.yml

@@ -0,0 +1,15 @@
+on:
+  push:
+  pull_request:
+name: Type checker
+jobs:
+  test:
+    name: pyright
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+    - run: pip install -e .
+    - uses: jakebailey/pyright-action@v1

dpgen2/entrypoint/download.py

@@ -28,6 +28,7 @@ def download(
     if wf_keys is None:
         wf_keys = wf.query_keys_of_steps()
 
+    assert wf_keys is not None
     for kk in wf_keys:
         download_dpgen2_artifacts(wf, kk, prefix=prefix)
         logging.info(f'step {kk} downloaded')

dpgen2/entrypoint/submit.py

@@ -102,16 +102,16 @@ def make_concurrent_learning_op (
         train_style : str = 'dp',
         explore_style : str = 'lmp',
         fp_style : str = 'vasp',
-        prep_train_config : str = default_config,
-        run_train_config : str = default_config,
-        prep_explore_config : str = default_config,
-        run_explore_config : str = default_config,
-        prep_fp_config : str = default_config,
-        run_fp_config : str = default_config,
-        select_confs_config : str = default_config,
-        collect_data_config : str = default_config,
-        cl_step_config : str = default_config,
-        upload_python_package : bool = None,
+        prep_train_config : dict = default_config,
+        run_train_config : dict = default_config,
+        prep_explore_config : dict = default_config,
+        run_explore_config : dict = default_config,
+        prep_fp_config : dict = default_config,
+        run_fp_config : dict = default_config,
+        select_confs_config : dict = default_config,
+        collect_data_config : dict = default_config,
+        cl_step_config : dict = default_config,
+        upload_python_package : Optional[List[os.PathLike]] = None,
 ):
     if train_style == 'dp':
         prep_run_train_op = PrepRunDPTrain(
@@ -277,6 +277,8 @@ def get_kspacing_kgamma_from_incar(
 ):
     with open(fname) as fp:
         lines = fp.readlines()
+    ks = None
+    kg = None
     for ii in lines:
         if 'KSPACING' in ii:
             ks = float(ii.split('=')[1])
@@ -287,12 +289,13 @@ def get_kspacing_kgamma_from_incar(
                 kg = False
             else:
                 raise RuntimeError(f"invalid kgamma value {ii.split('=')[1]}")
+    assert ks is not None and kg is not None
     return ks, kg
 
 
 def workflow_concurrent_learning(
         config : Dict,
-        old_style : Optional[bool] = False,
+        old_style : bool = False,
 ):
     default_config = normalize_step_dict(config.get('default_config', {})) if old_style else config['default_step_config']
 

dpgen2/entrypoint/watch.py

@@ -24,7 +24,7 @@
 
 def update_finished_steps(
         wf,
-        finished_keys : List[str] = None,
+        finished_keys : Optional[List[str]] = None,
         download : Optional[bool] = False,
         watching_keys : Optional[List[str]] = None,
         prefix : Optional[str] = None,
@@ -51,7 +51,7 @@ def watch(
         workflow_id,
         wf_config : Optional[Dict] = {},
         watching_keys : Optional[List] = default_watching_keys,
-        frequency : Optional[float] = 600.,
+        frequency : float = 600.,
         download : Optional[bool] = False,
         prefix : Optional[str] = None,
 ):

dpgen2/exploration/report/trajs_report.py

@@ -3,6 +3,7 @@
 from . import ExplorationReport
 from typing import (
     List,
+    Optional,
     Tuple,
 )
 from dflow.python import FatalError
@@ -92,7 +93,7 @@ def candidate_ratio(
 
     def get_candidates(
             self,
-            max_nframes : int = None,
+            max_nframes : Optional[int] = None,
     )->List[Tuple[int,int]]:
         """
         Get candidates. If number of candidates is larger than `max_nframes`, 

dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py

@@ -1,5 +1,6 @@
 from typing import (
     List,
+    Optional,
     Tuple,
 )
 from dflow.python import (
@@ -17,7 +18,7 @@ def __init__(
             stage : ExplorationStage,
             selector : ConfSelector,
             conv_accuracy : float = 0.9,
-            max_numb_iter : int = None,
+            max_numb_iter : Optional[int] = None,
             fatal_at_max : bool = True,
     ):
         self.stage = stage
@@ -42,9 +43,9 @@ def reached_max_iteration(self):
 
     def plan_next_iteration(
             self,
-            report : ExplorationReport = None,
-            trajs : List[Path] = None,
-    ) -> Tuple[bool, ExplorationTaskGroup, ConfSelector] :
+            report : Optional[ExplorationReport] = None,
+            trajs : Optional[List[Path]] = None,
+    ) -> Tuple[bool, Optional[ExplorationTaskGroup], Optional[ConfSelector]] :
         if report is None:
             stg_complete = False
             self.conv = stg_complete

dpgen2/exploration/scheduler/scheduler.py

@@ -2,6 +2,7 @@
 
 from typing import (
     List,
+    Optional,
     Tuple,
 )
 from dflow.python import (
@@ -76,9 +77,9 @@ def complete(self):
 
     def plan_next_iteration(
             self,
-            report : ExplorationReport = None,
-            trajs : List[Path] = None,
-    ) -> Tuple[bool, ExplorationTaskGroup, ConfSelector] :
+            report : Optional[ExplorationReport] = None,
+            trajs : Optional[List[Path]] = None,
+    ) -> Tuple[bool, Optional[ExplorationTaskGroup], Optional[ConfSelector]] :
         """
         Make the plan for the next DPGEN iteration.
 
@@ -202,7 +203,9 @@ def print_convergence(self):
         for iidx in range(len(accu)):
             if stage_idx[iidx] != prev_stg_idx:
                 if prev_stg_idx >= 0:
-                    ret.append(self._print_prev_summary(prev_stg_idx))
+                    _summary = self._print_prev_summary(prev_stg_idx)
+                    assert _summary is not None
+                    ret.append(_summary)
                 ret.append(f'# Stage {stage_idx[iidx]:4d}  ' + '-'*20)
                 prev_stg_idx = stage_idx[iidx]
             ret.append(' ' + fmt_str % (
@@ -213,6 +216,8 @@ def print_convergence(self):
             ))
         if self.complete():
             if prev_stg_idx >= 0:
-                ret.append(self._print_prev_summary(prev_stg_idx))
+                _summary = self._print_prev_summary(prev_stg_idx)
+                assert _summary is not None
+                ret.append(_summary)
                 ret.append(f'# All stages converged')
         return '\n'.join(ret + [''])

dpgen2/exploration/selector/conf_filter.py

@@ -7,9 +7,9 @@ class ConfFilter(ABC):
     @abstractmethod
     def check (
             self,
-            coords : np.array,
-            cell: np.array,
-            atom_types : np.array,
+            coords : np.ndarray,
+            cell: np.ndarray,
+            atom_types : np.ndarray,
             nopbc: bool,
     ) -> bool :
         """Check if the configuration is valid.

dpgen2/exploration/selector/conf_selector.py

@@ -1,6 +1,6 @@
 import dpdata
 from abc import ABC, abstractmethod
-from typing import Tuple, List, Set
+from typing import Optional, Tuple, List, Set
 from pathlib import Path
 from . import (
     ConfFilters,
@@ -17,7 +17,7 @@ def select (
             trajs : List[Path],
             model_devis : List[Path],
             traj_fmt : str = 'deepmd/npy',
-            type_map : List[str] = None,
+            type_map : Optional[List[str]] = None,
     ) -> Tuple[List[ Path ], ExplorationReport]:
         pass
 

dpgen2/exploration/selector/conf_selector_frame.py

@@ -3,6 +3,7 @@
 from collections import Counter
 from typing import (
     List,
+    Optional,
     Tuple,
 )
 from pathlib import Path
@@ -25,8 +26,8 @@ class ConfSelectorLammpsFrames(ConfSelector):
     def __init__(
             self,
             trust_level,
-            max_numb_sel : int = None,
-            conf_filters : ConfFilters = None,
+            max_numb_sel : Optional[int] = None,
+            conf_filters : Optional[ConfFilters] = None,
     ):
         self.trust_level = trust_level
         self.max_numb_sel = max_numb_sel
@@ -38,7 +39,7 @@ def select (
             trajs : List[Path],
             model_devis : List[Path],
             traj_fmt : str = 'lammps/dump',
-            type_map : List[str] = None,
+            type_map : Optional[List[str]] = None,
     ) -> Tuple[List[ Path ], ExplorationReport]:
         """Select configurations
 
@@ -136,6 +137,6 @@ def _load_traj(
     @staticmethod
     def _load_model_devi(
             fname : Path,
-    ) -> Tuple[np.array, np.array] : 
+    ) -> Tuple[np.ndarray, np.ndarray] : 
         dd = np.loadtxt(fname)
         return dd[:,4], dd[:,1]

dpgen2/exploration/task/conf_sampling_task_group.py

@@ -1,6 +1,7 @@
 import itertools, random
 from typing import (
     List,
+    Optional,
 )
 from . import (
     ExplorationTask,
@@ -22,7 +23,7 @@ def __init__(
     def set_conf(
             self,
             conf_list : List[str],
-            n_sample : int = None,
+            n_sample : Optional[int] = None,
             random_sample : bool = False,
     ):
         """

dpgen2/exploration/task/lmp/lmp_input.py

@@ -2,6 +2,7 @@
 import numpy as np
 from typing import (
     List,
+    Optional,
 )
 from distutils.version import (
     LooseVersion,
@@ -10,6 +11,9 @@
     lmp_traj_name,
 )
 
+import dpdata
+import scipy.constants as pc
+
 def _sample_sphere() :
     while True:
         vv = np.array([np.random.normal(), np.random.normal(), np.random.normal()])
@@ -24,19 +28,19 @@ def make_lmp_input(
         graphs : List[str],
         nsteps : int,
         dt : float,
-        neidelay : int,
+        neidelay : Optional[int],
         trj_freq : int,
         mass_map : List[float],
         temp : float, 
         tau_t : float = 0.1,
-        pres : float= None,
+        pres : Optional[float]= None,
         tau_p : float = 0.5,
         use_clusters : bool = False,        
-        relative_f_epsilon : float = None,
-        relative_v_epsilon : float = None,
-        pka_e : float = None,
-        ele_temp_f : float = None,
-        ele_temp_a : float = None,
+        relative_f_epsilon : Optional[float] = None,
+        relative_v_epsilon : Optional[float] = None,
+        pka_e : Optional[float] = None,
+        ele_temp_f : Optional[float] = None,
+        ele_temp_a : Optional[float] = None,
         nopbc : bool = False,
         max_seed : int = 1000000,
         deepmd_version = '2.0', 
@@ -115,7 +119,7 @@ def make_lmp_input(
         sys_data = sys.data
         pka_mass = mass_map[sys_data['atom_types'][0] - 1]
         pka_vn = pka_e * pc.electron_volt / \
-                 (0.5 * pka_mass * 1e-3 / pc.Avogadro * (pc.angstrom / pc.pico) ** 2)
+                 (0.5 * pka_mass * 1e-3 / pc.Avogadro * (pc.angstrom / pc.pico) ** 2) # type: ignore
         pka_vn = np.sqrt(pka_vn)
         print(pka_vn)
         pka_vec = _sample_sphere()

dpgen2/exploration/task/lmp_template_task_group.py

@@ -1,6 +1,7 @@
 import itertools, random
 from typing import (
     List,
+    Optional,
 )
 from pathlib import Path
 from . import (
@@ -30,7 +31,7 @@ def set_lmp(
             self,
             numb_models : int,
             lmp_template_fname : str,
-            plm_template_fname : str = None,
+            plm_template_fname : Optional[str] = None,
             revisions : dict = {},
             traj_freq : int = 10,
     )->None:
@@ -92,7 +93,7 @@ def _make_lmp_task(
             self,
             conf: str,
             lmp_cont : str,
-            plm_cont : str = None,
+            plm_cont : Optional[str] = None,
     )->ExplorationTask:
         task = ExplorationTask()
         task\

dpgen2/exploration/task/npt_task_group.py

@@ -1,6 +1,7 @@
 import itertools, random
 from typing import (
     List,
+    Optional,
 )
 from . import (
     ExplorationTask,
@@ -26,21 +27,21 @@ def set_md(
             numb_models,
             mass_map,
             temps : List[float],
-            press : List[float] = None,
+            press : Optional[List[float]] = None,
             ens : str = 'npt',
             dt : float = 0.001,
             nsteps : int = 1000,
             trj_freq : int = 10,
             tau_t : float = 0.1,
             tau_p : float = 0.5,
-            pka_e : float = None,
-            neidelay : int = None,
+            pka_e : Optional[float] = None,
+            neidelay : Optional[int] = None,
             no_pbc : bool = False,
             use_clusters : bool = False,
-            relative_f_epsilon : float = None,
-            relative_v_epsilon : float = None,
-            ele_temp_f : float = None,
-            ele_temp_a : float = None,
+            relative_f_epsilon : Optional[float] = None,
+            relative_v_epsilon : Optional[float] = None,
+            ele_temp_f : Optional[float] = None,
+            ele_temp_a : Optional[float] = None,
     ):
         """
         Set MD parameters
@@ -95,7 +96,7 @@ def _make_lmp_task(
             self,
             conf : str,
             tt : float,
-            pp : float,
+            pp : Optional[float],
     ) -> ExplorationTask:
         task = ExplorationTask()
         task\

dpgen2/exploration/task/task.py

@@ -122,7 +122,8 @@ def __init__(
 class FooTaskGroup(ExplorationTaskGroup):
     def __init__(self, numb_task):
         super().__init__()
-        self.tlist = []
+        # TODO: confirm the following is correct
+        self.tlist = ExplorationTaskGroup()
         for ii in range(numb_task):
             self.tlist.add_task( 
                 FooTask(f'conf.{ii}', f'this is conf.{ii}',