renamed cont_dim to container_ndim

Author: tclose

pydra/compose/base/task.py

@@ -160,7 +160,7 @@ class Task(ty.Generic[OutputsType]):
     # The following fields are used to store split/combine state information
     _splitter = attrs.field(default=None, init=False, repr=False)
     _combiner = attrs.field(default=None, init=False, repr=False)
-    _cont_dim = attrs.field(default=None, init=False, repr=False)
+    _container_ndim = attrs.field(default=None, init=False, repr=False)
     _hashes = attrs.field(default=None, init=False, eq=False, repr=False)
 
     RESERVED_FIELD_NAMES = ("split", "combine")
@@ -265,7 +265,7 @@ def split(
         splitter: ty.Union[str, ty.List[str], ty.Tuple[str, ...], None] = None,
         /,
         overwrite: bool = False,
-        cont_dim: ty.Optional[dict] = None,
+        container_ndim: ty.Optional[dict] = None,
         **inputs,
     ) -> Self:
         """
@@ -279,9 +279,9 @@ def split(
             then the fields to split are taken from the keyword-arg names.
         overwrite : bool, optional
             whether to overwrite an existing split on the node, by default False
-        cont_dim : dict, optional
+        container_ndim : dict, optional
             Container dimensions for specific inputs, used in the splitter.
-            If input name is not in cont_dim, it is assumed that the input values has
+            If input name is not in container_ndim, it is assumed that the input values has
             a container dimension of 1, so only the most outer dim will be used for splitting.
         **inputs
             fields to split over, will be automatically wrapped in a StateArray object
@@ -321,7 +321,7 @@ def split(
         else:
             # If no splitter is provided, use the names of the inputs as combinatorial splitter
             split_names = splitter = list(inputs)
-        for field_name in cont_dim or []:
+        for field_name in container_ndim or []:
             if field_name not in split_names:
                 raise ValueError(
                     f"Container dimension for {field_name} is provided but the field "
@@ -342,7 +342,7 @@ def split(
             split_inputs[name] = split_val
         split_def = attrs.evolve(self, **split_inputs)
         split_def._splitter = splitter
-        split_def._cont_dim = cont_dim
+        split_def._container_ndim = container_ndim
         return split_def
 
     def combine(
@@ -613,8 +613,8 @@ def bytes_repr_task(obj: Task, cache: Cache) -> ty.Iterator[bytes]:
     yield hash_single(obj._splitter, cache)
     yield b",_combiner="
     yield hash_single(obj._combiner, cache)
-    yield b",_cont_dim="
-    yield hash_single(obj._cont_dim, cache)
+    yield b",_container_ndim="
+    yield hash_single(obj._container_ndim, cache)
     yield b",_xor="
     yield hash_single(obj._xor, cache)
     yield b")"

pydra/compose/tests/test_workflow_run.py

@@ -1524,7 +1524,7 @@ def test_wf_ndstinner_5(worker: str, tmp_path: Path):
     """workflow with 3 tasks,
     the second task has two inputs and inner splitter from one of the input,
     (inner input come from the first task that has its own splitter,
-    there is a inner_cont_dim)
+    there is a inner_container_ndim)
     the third task has no new splitter
     """
 

pydra/engine/node.py

@@ -173,25 +173,25 @@ def _check_if_outputs_have_been_used(self, msg):
             )
 
     def _set_state(self) -> None:
-        # Add node name to state's splitter, combiner and cont_dim loaded from the def
+        # Add node name to state's splitter, combiner and container_ndim loaded from the def
         splitter = deepcopy(self._task._splitter)  # these can be modified in state
         combiner = deepcopy(self._task._combiner)  # these can be modified in state
-        cont_dim = {}
+        container_ndim = {}
         if splitter:
             splitter = add_name_splitter(splitter, self.name)
         if combiner:
             combiner = add_name_combiner(combiner, self.name)
-        if self._task._cont_dim:
-            for key, val in self._task._cont_dim.items():
-                cont_dim[f"{self.name}.{key}"] = val
+        if self._task._container_ndim:
+            for key, val in self._task._container_ndim.items():
+                container_ndim[f"{self.name}.{key}"] = val
         other_states = self._get_upstream_states()
         if splitter or combiner or other_states:
             self._state = State(
                 self.name,
                 splitter=splitter,
                 other_states=other_states,
                 combiner=combiner,
-                cont_dim=cont_dim,
+                container_ndim=container_ndim,
             )
             if combiner:
                 if not_split := [
@@ -218,7 +218,7 @@ def _get_upstream_states(self) -> dict[str, tuple["State", list[str]]]:
                 node: Node = val._node
                 # variables that are part of inner splitters should be treated as a containers
                 if node.state and f"{node.name}.{val._field}" in node.state.splitter:
-                    node.state._inner_cont_dim[f"{node.name}.{val._field}"] = 1
+                    node.state._inner_container_ndim[f"{node.name}.{val._field}"] = 1
                 # adding task_name: (task.state, [a field from the connection]
                 if node.name not in upstream_states:
                     upstream_states[node.name] = (node.state, [val._field])

pydra/engine/state.py

@@ -85,7 +85,7 @@ def __init__(
         name,
         splitter=None,
         combiner=None,
-        cont_dim=None,
+        container_ndim=None,
         other_states=None,
     ):
         """
@@ -109,8 +109,8 @@ def __init__(
         self.splitter = splitter
         # temporary combiner
         self.combiner = combiner
-        self.cont_dim = cont_dim or {}
-        self._inner_cont_dim = {}
+        self.container_ndim = container_ndim or {}
+        self._inner_container_ndim = {}
         self._inputs_ind = None
         # if other_states, the connections have to be updated
         if self.other_states:
@@ -377,12 +377,12 @@ def prev_state_splitter_rpn_compact(self):
         return self._prev_state_splitter_rpn_compact
 
     @property
-    def cont_dim_all(self):
-        # adding inner_cont_dim to the general container_dimension provided by the users
-        cont_dim_all = deepcopy(self.cont_dim)
-        for k, v in self._inner_cont_dim.items():
-            cont_dim_all[k] = cont_dim_all.get(k, 1) + v
-        return cont_dim_all
+    def container_ndim_all(self):
+        # adding inner_container_ndim to the general container_dimension provided by the users
+        container_ndim_all = deepcopy(self.container_ndim)
+        for k, v in self._inner_container_ndim.items():
+            container_ndim_all[k] = container_ndim_all.get(k, 1) + v
+        return container_ndim_all
 
     @property
     def combiner(self):
@@ -869,7 +869,7 @@ def combiner_validation(self):
     def prepare_states(
         self,
         inputs: dict[str, ty.Any],
-        cont_dim: dict[str, int] | None = None,
+        container_ndim: dict[str, int] | None = None,
     ):
         """
         Prepare a full list of state indices and state values.
@@ -885,13 +885,13 @@ def prepare_states(
         self.combiner_validation()
         self.set_input_groups()
         self.inputs = inputs
-        if cont_dim is not None:
-            self.cont_dim = cont_dim
+        if container_ndim is not None:
+            self.container_ndim = container_ndim
         if self.other_states:
             st: State
             for nm, (st, _) in self.other_states.items():
                 self.inputs.update(st.inputs)
-                self.cont_dim.update(st.cont_dim_all)
+                self.container_ndim.update(st.container_ndim_all)
 
         self.prepare_states_ind()
         self.prepare_states_val()
@@ -995,7 +995,9 @@ def prepare_states_combined_ind(self, elements_to_remove_comb):
     def prepare_states_val(self):
         """Evaluate states values having states indices."""
         self.states_val = list(
-            map_splits(self.states_ind, self.inputs, cont_dim=self.cont_dim_all)
+            map_splits(
+                self.states_ind, self.inputs, container_ndim=self.container_ndim_all
+            )
         )
         return self.states_val
 
@@ -1165,8 +1167,8 @@ def _processing_terms(self, term, previous_states_ind):
             var_ind, new_keys = previous_states_ind[term]
             shape = (len(var_ind),)
         else:
-            cont_dim = self.cont_dim_all.get(term, 1)
-            shape = input_shape(self.inputs[term], cont_dim=cont_dim)
+            container_ndim = self.container_ndim_all.get(term, 1)
+            shape = input_shape(self.inputs[term], container_ndim=container_ndim)
             var_ind = range(reduce(lambda x, y: x * y, shape))
             new_keys = [term]
             # checking if the term is in inner_inputs
@@ -1186,7 +1188,8 @@ def _processing_terms(self, term, previous_states_ind):
     def _single_op_splits(self, op_single):
         """splits function if splitter is a singleton"""
         shape = input_shape(
-            self.inputs[op_single], cont_dim=self.cont_dim_all.get(op_single, 1)
+            self.inputs[op_single],
+            container_ndim=self.container_ndim_all.get(op_single, 1),
         )
         val_ind = range(reduce(lambda x, y: x * y, shape))
         if op_single in self.inner_inputs:
@@ -1211,8 +1214,8 @@ def _single_op_splits(self, op_single):
     def _get_element(self, value: ty.Any, field_name: str, ind: int) -> ty.Any:
         """
         Extracting element of the inputs taking into account
-        container dimension of the specific element that can be set in self.state.cont_dim.
-        If input name is not in cont_dim, it is assumed that the input values has
+        container dimension of the specific element that can be set in self.state.container_ndim.
+        If input name is not in container_ndim, it is assumed that the input values has
         a container dimension of 1, so only the most outer dim will be used for splitting.
 
         Parameters
@@ -1229,11 +1232,11 @@ def _get_element(self, value: ty.Any, field_name: str, ind: int) -> ty.Any:
         Any
             specific element of the input field
         """
-        if f"{self.name}.{field_name}" in self.cont_dim_all:
+        if f"{self.name}.{field_name}" in self.container_ndim_all:
             return list(
                 flatten(
                     ensure_list(value),
-                    max_depth=self.cont_dim_all[f"{self.name}.{field_name}"],
+                    max_depth=self.container_ndim_all[f"{self.name}.{field_name}"],
                 )
             )[ind]
         else:
@@ -1600,15 +1603,15 @@ def iter_splits(iterable, keys):
         yield dict(zip(keys, list(flatten(iter, max_depth=1000))))
 
 
-def input_shape(inp, cont_dim=1):
+def input_shape(inp, container_ndim=1):
     """Get input shape, depends on the container dimension, if not specify it is assumed to be 1"""
     # TODO: have to be changed for inner splitter (sometimes different length)
-    cont_dim -= 1
+    container_ndim -= 1
     shape = [len(inp)]
     last_shape = None
     for value in inp:
-        if isinstance(value, list) and cont_dim > 0:
-            cur_shape = input_shape(value, cont_dim)
+        if isinstance(value, list) and container_ndim > 0:
+            cur_shape = input_shape(value, container_ndim)
             if last_shape is None:
                 last_shape = cur_shape
             elif last_shape != cur_shape:
@@ -1828,13 +1831,15 @@ def combine_final_groups(combiner, groups, groups_stack, keys):
     return keys_final, groups_final_map, groups_stack_final, combiner_all
 
 
-def map_splits(split_iter, inputs, cont_dim=None):
+def map_splits(split_iter, inputs, container_ndim=None):
     """generate a dictionary of inputs prescribed by the splitter."""
-    if cont_dim is None:
-        cont_dim = {}
+    if container_ndim is None:
+        container_ndim = {}
     for split in split_iter:
         yield {
-            k: list(flatten(ensure_list(inputs[k]), max_depth=cont_dim.get(k, None)))[v]
+            k: list(
+                flatten(ensure_list(inputs[k]), max_depth=container_ndim.get(k, None))
+            )[v]
             for k, v in split.items()
         }
 

pydra/engine/submitter.py

@@ -215,7 +215,7 @@ def __call__(
                 name="outer_split",
                 splitter=deepcopy(task._splitter),
                 combiner=deepcopy(task._combiner),
-                cont_dim=deepcopy(task._cont_dim),
+                container_ndim=deepcopy(task._container_ndim),
             )
 
             def wrap_type(tp):