adding cond_dim option to splitter

Author: djarecka

pydra/engine/helpers_state.py

@@ -362,14 +362,15 @@ def iter_splits(iterable, keys):
         yield dict(zip(keys, list(flatten(iter, max_depth=1000))))
 
 
-def input_shape(in1):
+def input_shape(inp, cont_dim=1):
     """Get input shape."""
     # TODO: have to be changed for inner splitter (sometimes different length)
-    shape = [len(in1)]
+    cont_dim -= 1
+    shape = [len(inp)]
     last_shape = None
-    for value in in1:
-        if isinstance(value, list):
-            cur_shape = input_shape(value)
+    for value in inp:
+        if isinstance(value, list) and cont_dim > 0:
+            cur_shape = input_shape(value, cont_dim)
             if last_shape is None:
                 last_shape = cur_shape
             elif last_shape != cur_shape:
@@ -383,11 +384,34 @@ def input_shape(in1):
     return tuple(shape)
 
 
-def splits(splitter_rpn, inputs, inner_inputs=None):
-    """Split process as specified by an rpn splitter, from left to right."""
+def splits(splitter_rpn, inputs, inner_inputs=None, cont_dim=None):
+    """
+    Splits input variable as specified by splitter
+
+    Parameters
+    ----------
+    splitter_rpn : list
+        splitter in RPN notation
+    inputs: dict
+        input variables
+    inner_inputs: dict, optional
+        inner input specification
+
+
+    Returns
+    -------
+    splitter : list
+        each element contains indices for inputs
+    keys: list
+        names of input variables
+
+    """
+
     stack = []
     keys = []
-    shapes_var = {}
+    if cont_dim is None:
+        cont_dim = {}
+    # analysing states from connected tasks if inner_inputs
     if inner_inputs:
         previous_states_ind = {
             "_{}".format(v.name): (v.ind_l_final, v.keys_final)
@@ -407,9 +431,9 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
             op_single,
             inputs,
             inner_inputs,
-            shapes_var,
             previous_states_ind,
             keys_fromLeftSpl,
+            cont_dim=cont_dim,
         )
 
     terms = {}
@@ -418,7 +442,11 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
     shape = {}
     # iterating splitter_rpn
     for token in splitter_rpn:
-        if token in [".", "*"]:
+        if token not in [".", "*"]:  # token is one of the input var
+            # adding variable to the stack
+            stack.append(token)
+        else:
+            # removing Right and Left var from the stack
             terms["R"] = stack.pop()
             terms["L"] = stack.pop()
             # checking if terms are strings, shapes, etc.
@@ -429,10 +457,14 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
                         trm_val[lr] = previous_states_ind[term][0]
                         shape[lr] = (len(trm_val[lr]),)
                     else:
-                        shape[lr] = input_shape(inputs[term])
+                        if term in cont_dim:
+                            shape[lr] = input_shape(
+                                inputs[term], cont_dim=cont_dim[term]
+                            )
+                        else:
+                            shape[lr] = input_shape(inputs[term])
                         trm_val[lr] = range(reduce(lambda x, y: x * y, shape[lr]))
                     trm_str[lr] = True
-                    shapes_var[term] = shape[lr]
                 else:
                     trm_val[lr], shape[lr] = term
                     trm_str[lr] = False
@@ -447,6 +479,7 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
                     )
                 newshape = shape["R"]
             if token == "*":
+                # TODO: pomyslec
                 newshape = tuple(list(shape["L"]) + list(shape["R"]))
 
             # creating list with keys
@@ -466,7 +499,6 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
             elif trm_str["R"]:
                 keys = keys + new_keys["R"]
 
-            #
             newtrm_val = {}
             for lr in ["R", "L"]:
                 # TODO: rewrite once I have more tests
@@ -491,13 +523,11 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
 
             pushval = (op[token](newtrm_val["L"], newtrm_val["R"]), newshape)
             stack.append(pushval)
-        else:  # name of one of the inputs (token not in [".", "*"])
-            stack.append(token)
 
     val = stack.pop()
     if isinstance(val, tuple):
         val = val[0]
-    return val, keys, shapes_var, keys_fromLeftSpl
+    return val, keys, keys_fromLeftSpl
 
 
 # dj: TODO: do I need keys?
@@ -636,17 +666,22 @@ def splits_groups(splitter_rpn, combiner=None, inner_inputs=None):
 
 
 def _single_op_splits(
-    op_single, inputs, inner_inputs, shapes_var, previous_states_ind, keys_fromLeftSpl
+    op_single,
+    inputs,
+    inner_inputs,
+    previous_states_ind,
+    keys_fromLeftSpl,
+    cont_dim=None,
 ):
     if op_single.startswith("_"):
         return (
             previous_states_ind[op_single][0],
             previous_states_ind[op_single][1],
-            None,
             keys_fromLeftSpl,
         )
-    shape = input_shape(inputs[op_single])
-    shapes_var[op_single] = shape
+    if cont_dim is None:
+        cont_dim = {}
+    shape = input_shape(inputs[op_single], cont_dim=cont_dim.get(op_single, 1))
     trmval = range(reduce(lambda x, y: x * y, shape))
     if op_single in inner_inputs:
         # TODO: have to be changed if differ length
@@ -659,11 +694,11 @@ def _single_op_splits(
         res = op["."](op_out, trmval)
         val = res
         keys = inner_inputs[op_single].keys_final + [op_single]
-        return val, keys, shapes_var, keys_fromLeftSpl
+        return val, keys, keys_fromLeftSpl
     else:
         val = op["*"](trmval)
         keys = [op_single]
-        return val, keys, shapes_var, keys_fromLeftSpl
+        return val, keys, keys_fromLeftSpl
 
 
 def _single_op_splits_groups(
@@ -727,10 +762,15 @@ def combine_final_groups(combiner, groups, groups_stack, keys):
     return keys_final, groups_final, groups_stack_final, combiner_all
 
 
-def map_splits(split_iter, inputs):
+def map_splits(split_iter, inputs, cont_dim=None):
     """Get a dictionary of prescribed splits."""
+    if cont_dim is None:
+        cont_dim = {}
     for split in split_iter:
-        yield {k: list(flatten(ensure_list(inputs[k])))[v] for k, v in split.items()}
+        yield {
+            k: list(flatten(ensure_list(inputs[k]), max_depth=cont_dim.get(k, None)))[v]
+            for k, v in split.items()
+        }
 
 
 # Functions for merging and completing splitters in states.

pydra/engine/state.py

@@ -346,7 +346,7 @@ def push_new_states(self):
                 stack = [gr + nr_gr_f for gr in stack]
                 self.groups_stack_final.append(stack)
 
-    def prepare_states(self, inputs):
+    def prepare_states(self, inputs, cont_dim=None):
         """
         Prepare a full list of state indices and state values.
 
@@ -357,6 +357,10 @@ def prepare_states(self, inputs):
             specific elements from inputs that can be used running interfaces
 
         """
+        if cont_dim is None:
+            self.cont_dim = {}
+        else:
+            self.cont_dim = cont_dim
         if isinstance(inputs, BaseSpec):
             self.inputs = hlpst.inputs_types_to_dict(self.name, inputs)
         else:
@@ -366,7 +370,7 @@ def prepare_states(self, inputs):
                 # I think now this if is never used
                 if not hasattr(st, "states_ind"):
                     # dj: should i provide different inputs?
-                    st.prepare_states(self.inputs)
+                    st.prepare_states(self.inputs, cont_dim=cont_dim)
                 self.inputs.update(st.inputs)
         self.prepare_states_ind()
         self.prepare_states_val()
@@ -395,8 +399,11 @@ def prepare_states_ind(self):
         partial_rpn = hlpst.remove_inp_from_splitter_rpn(
             deepcopy(self.splitter_rpn_compact), elements_to_remove
         )
-        values_out_pr, keys_out_pr, _, kL = hlpst.splits(
-            partial_rpn, self.inputs, inner_inputs=self.inner_inputs
+        values_out_pr, keys_out_pr, kL = hlpst.splits(
+            partial_rpn,
+            self.inputs,
+            inner_inputs=self.inner_inputs,
+            cont_dim=self.cont_dim,
         )
         values_pr = list(values_out_pr)
 
@@ -429,8 +436,11 @@ def prepare_states_combined_ind(self, elements_to_remove_comb):
         )
         # TODO: create a function for this!!
         if combined_rpn:
-            val_r, key_r, _, _ = hlpst.splits(
-                combined_rpn, self.inputs, inner_inputs=self.inner_inputs
+            val_r, key_r, _ = hlpst.splits(
+                combined_rpn,
+                self.inputs,
+                inner_inputs=self.inner_inputs,
+                cont_dim=self.cont_dim,
             )
             values = list(val_r)
         else:
@@ -464,7 +474,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(hlpst.map_splits(self.states_ind, self.inputs))
+        self.states_val = list(
+            hlpst.map_splits(self.states_ind, self.inputs, cont_dim=self.cont_dim)
+        )
         return self.states_val
 
     def prepare_inputs(self):
@@ -489,8 +501,11 @@ def prepare_inputs(self):
                 deepcopy(self.splitter_rpn_compact), elements_to_remove
             )
             if partial_rpn:
-                values_inp, keys_inp, _, _ = hlpst.splits(
-                    partial_rpn, self.inputs, inner_inputs=self.inner_inputs
+                values_inp, keys_inp, _ = hlpst.splits(
+                    partial_rpn,
+                    self.inputs,
+                    inner_inputs=self.inner_inputs,
+                    cont_dim=self.cont_dim,
                 )
                 inputs_ind = values_inp
             else: