check arguments with 2D lists

dpgen/arginfo.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import Callable, Tuple
 
 from dargs import Argument
 
@@ -43,3 +43,35 @@ def general_mdata_arginfo(name: str, tasks: Tuple[str]) -> Argument:
             )
         )
     return Argument(name, dict, sub_fields=sub_fields, doc=doc_run_mdata)
+
+
+def check_nd_list(dimesion: int = 2) -> Callable:
+    """Return a method to check if the input is a nd list.
+
+    Parameters
+    ----------
+    dimesion : int, default=2
+        dimension of the array
+
+    Returns
+    -------
+    callable
+        check function
+    """
+
+    def check(value, dimension=dimesion):
+        if value is None:
+            # do not check null
+            return True
+        if dimension:
+            if not isinstance(value, list):
+                return False
+        if dimension > 1:
+            if not all(check(v, dimension=dimesion - 1) for v in value):
+                return False
+        return True
+
+    return check
+
+
+errmsg_nd_list = "Must be a %d-dimension list."

dpgen/generator/arginfo.py

@@ -3,7 +3,7 @@
 
 from dargs import Argument, Variant
 
-from dpgen.arginfo import general_mdata_arginfo
+from dpgen.arginfo import check_nd_list, errmsg_nd_list, general_mdata_arginfo
 
 
 def run_mdata_arginfo() -> Argument:
@@ -41,7 +41,7 @@ def data_args() -> List[Argument]:
     doc_sys_format = "Format of sys_configs."
     doc_init_batch_size = "Each number is the batch_size of corresponding system for training in init_data_sys. One recommended rule for setting the sys_batch_size and init_batch_size is that batch_size mutiply number of atoms ot the stucture should be larger than 32. If set to auto, batch size will be 32 divided by number of atoms. This argument will not override the mixed batch size in `default_training_param`."
     doc_sys_configs_prefix = "Prefix of sys_configs."
-    doc_sys_configs = "Containing directories of structures to be explored in iterations.Wildcard characters are supported here."
+    doc_sys_configs = "2D list. Containing directories of structures to be explored in iterations for each system. Wildcard characters are supported here."
     doc_sys_batch_size = "Each number is the batch_size for training of corresponding system in sys_configs. If set to auto, batch size will be 32 divided by number of atoms. This argument will not override the mixed batch size in `default_training_param`."
 
     return [
@@ -54,7 +54,14 @@ def data_args() -> List[Argument]:
             "init_batch_size", [list, str], optional=True, doc=doc_init_batch_size
         ),
         Argument("sys_configs_prefix", str, optional=True, doc=doc_sys_configs_prefix),
-        Argument("sys_configs", list, optional=False, doc=doc_sys_configs),
+        Argument(
+            "sys_configs",
+            list,
+            optional=False,
+            doc=doc_sys_configs,
+            extra_check=check_nd_list(2),
+            extra_check_errmsg=errmsg_nd_list % 2,
+        ),
         Argument("sys_batch_size", list, optional=True, doc=doc_sys_batch_size),
     ]
 
@@ -466,7 +473,7 @@ def model_devi_amber_args() -> List[Argument]:
         "List of ints. The number of steps to run. Each number maps to a system."
     )
     doc_r = (
-        "3D or 4D list of floats. Constrict values for the enhanced sampling. "
+        "2D or 3D list of floats. Constrict values for the enhanced sampling. "
         "The first dimension maps to systems. "
         "The second dimension maps to confs in each system. The third dimension is the "
         "constrict value. It can be a single float for 1D or list of floats for nD."
@@ -502,7 +509,14 @@ def model_devi_amber_args() -> List[Argument]:
         Argument("qm_region", list, optional=False, doc=doc_qm_region),
         Argument("qm_charge", list, optional=False, doc=doc_qm_charge),
         Argument("nsteps", list, optional=False, doc=doc_nsteps),
-        Argument("r", list, optional=False, doc=doc_r),
+        Argument(
+            "r",
+            list,
+            optional=False,
+            doc=doc_r,
+            extra_check=check_nd_list(2),
+            extra_check_errmsg=errmsg_nd_list % 2,
+        ),
         Argument("disang_prefix", str, optional=True, doc=doc_disang_prefix),
         Argument("disang", list, optional=False, doc=doc_disang),
         # post model devi args

pyproject.toml

@@ -29,7 +29,7 @@ dependencies = [
     'GromacsWrapper>=0.8.0',
     'dpdispatcher>=0.3.11',
     'netCDF4',
-    'dargs>=0.2.9',
+    'dargs>=0.3.5',
     'h5py',
     'pymatgen-analysis-defects',
     'openbabel-wheel',