some small edits to the state class, adding comments

Author: djarecka

pydra/engine/helpers_state.py

@@ -19,11 +19,11 @@ def splitter2rpn(splitter, other_states=None, state_fields=True):
     Parameters
     ----------
     splitter :
-        TODO
+        splitter (standard form)
     other_states :
-        TODO
+        other states that are connected to the state
     state_fields : :obj:`bool`
-        TODO
+        if False the splitter from the previous states are unwrapped
 
     """
     if not splitter:
@@ -174,7 +174,6 @@ def _iterate_list(element, sign, other_states, output_splitter, state_fields=Tru
         )
 
 
-# functions used in State to know which element should be used for a specific axis
 def converter_groups_to_input(group_for_inputs):
     """
     Return fields for each axis and number of all groups.
@@ -184,7 +183,7 @@ def converter_groups_to_input(group_for_inputs):
     Parameters
     ----------
     group_for_inputs :
-        TODO
+        specified axes (groups) for each input
 
     """
     input_for_axis = {}
@@ -199,7 +198,6 @@ def converter_groups_to_input(group_for_inputs):
     return input_for_axis, ngr
 
 
-# function used in State if combiner
 def remove_inp_from_splitter_rpn(splitter_rpn, inputs_to_remove):
     """
     Remove inputs due to combining.
@@ -211,7 +209,7 @@ def remove_inp_from_splitter_rpn(splitter_rpn, inputs_to_remove):
     splitter_rpn :
         The splitter in reverse polish notation
     inputs_to_remove :
-        TODO
+        input names that should be removed from the splitter
 
     """
     splitter_rpn_copy = splitter_rpn.copy()
@@ -258,12 +256,12 @@ def rpn2splitter(splitter_rpn):
     Parameters
     ----------
     splitter_rpn :
-        TODO
+        splitter in reverse polish notation
 
     Returns
     -------
     splitter :
-        TODO
+        splitter in the standard/original form
 
     """
     if splitter_rpn == []:
@@ -298,9 +296,8 @@ def rpn2splitter(splitter_rpn):
     return rpn2splitter(splitter_modified)
 
 
-# used in the Node to change names in a splitter and combiner
 def add_name_combiner(combiner, name):
-    """Add a combiner."""
+    """ adding a node's name to each field from the combiner"""
     combiner_changed = []
     for comb in combiner:
         if "." not in comb:
@@ -311,7 +308,7 @@ def add_name_combiner(combiner, name):
 
 
 def add_name_splitter(splitter, name):
-    """Change names of splitter: adding names of the node."""
+    """ adding a node's name to each field from the splitter"""
     if isinstance(splitter, str):
         return _add_name([splitter], name)[0]
     elif isinstance(splitter, list):
@@ -322,6 +319,7 @@ def add_name_splitter(splitter, name):
 
 
 def _add_name(mlist, name):
+    """ adding anem to each element from the list"""
     for i, elem in enumerate(mlist):
         if isinstance(elem, str):
             if "." in elem or elem.startswith("_"):
@@ -421,22 +419,15 @@ def splits(splitter_rpn, inputs, inner_inputs=None, cont_dim=None):
             for _, v in inner_inputs.items()
         }
         inner_inputs = {k: v for k, v in inner_inputs.items() if k in splitter_rpn}
-        keys_fromLeftSpl = ["_{}".format(st.name) for _, st in inner_inputs.items()]
     else:
         previous_states_ind = {}
         inner_inputs = {}
-        keys_fromLeftSpl = []
 
     # when splitter is a single element (no operators)
     if len(splitter_rpn) == 1:
         op_single = splitter_rpn[0]
         return _single_op_splits(
-            op_single,
-            inputs,
-            inner_inputs,
-            previous_states_ind,
-            keys_fromLeftSpl,
-            cont_dim=cont_dim,
+            op_single, inputs, inner_inputs, previous_states_ind, cont_dim=cont_dim
         )
 
     terms = {}
@@ -530,12 +521,41 @@ def splits(splitter_rpn, inputs, inner_inputs=None, cont_dim=None):
     val = stack.pop()
     if isinstance(val, tuple):
         val = val[0]
-    return val, keys, keys_fromLeftSpl
+    return val, keys
+
+
+def _single_op_splits(
+    op_single, inputs, inner_inputs, previous_states_ind, cont_dim=None
+):
+    """ splits function if splitter is a singleton"""
+    if op_single.startswith("_"):
+        return (previous_states_ind[op_single][0], previous_states_ind[op_single][1])
+    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
+        inner_len = [shape[-1]] * reduce(lambda x, y: x * y, shape[:-1])
+        # this come from the previous node
+        outer_ind = inner_inputs[op_single].ind_l
+        op_out = itertools.chain.from_iterable(
+            itertools.repeat(x, n) for x, n in zip(outer_ind, inner_len)
+        )
+        res = op["."](op_out, trmval)
+        val = res
+        keys = inner_inputs[op_single].keys_final + [op_single]
+        return val, keys
+    else:
+        val = op["*"](trmval)
+        keys = [op_single]
+        return val, keys
 
 
-# dj: TODO: do I need keys?
 def splits_groups(splitter_rpn, combiner=None, inner_inputs=None):
-    """Process splitter rpn from left to right."""
+    """ splits inputs to groups (axes) and creates stacks for these groups
+        This is used to specify which input can be combined.
+    """
     if not splitter_rpn:
         return [], {}, [], []
     stack = []
@@ -556,9 +576,7 @@ def splits_groups(splitter_rpn, combiner=None, inner_inputs=None):
     # when splitter is a single element (no operators)
     if len(splitter_rpn) == 1:
         op_single = splitter_rpn[0]
-        return _single_op_splits_groups(
-            op_single, combiner, inner_inputs, previous_states_ind, groups
-        )
+        return _single_op_splits_groups(op_single, combiner, inner_inputs, groups)
 
     # len(splitter_rpn) > 1
     # iterating splitter_rpn
@@ -668,45 +686,8 @@ def splits_groups(splitter_rpn, combiner=None, inner_inputs=None):
         return keys, groups, groups_stack, []
 
 
-def _single_op_splits(
-    op_single,
-    inputs,
-    inner_inputs,
-    previous_states_ind,
-    keys_fromLeftSpl,  # TODO NOW do I need it?
-    cont_dim=None,
-):
-    if op_single.startswith("_"):
-        return (
-            previous_states_ind[op_single][0],
-            previous_states_ind[op_single][1],
-            keys_fromLeftSpl,
-        )
-    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
-        inner_len = [shape[-1]] * reduce(lambda x, y: x * y, shape[:-1])
-        # this come from the previous node
-        outer_ind = inner_inputs[op_single].ind_l
-        op_out = itertools.chain.from_iterable(
-            itertools.repeat(x, n) for x, n in zip(outer_ind, inner_len)
-        )
-        res = op["."](op_out, trmval)
-        val = res
-        keys = inner_inputs[op_single].keys_final + [op_single]
-        return val, keys, keys_fromLeftSpl
-    else:
-        val = op["*"](trmval)
-        keys = [op_single]
-        return val, keys, keys_fromLeftSpl
-
-
-def _single_op_splits_groups(
-    op_single, combiner, inner_inputs, previous_states_ind, groups
-):
+def _single_op_splits_groups(op_single, combiner, inner_inputs, groups):
+    """ splits_groups function if splitter is a singleton"""
     if op_single in inner_inputs:
         # TODO: have to be changed if differ length
         # TODO: i think I don't want to add here from left part
@@ -768,7 +749,7 @@ def combine_final_groups(combiner, groups, groups_stack, keys):
 
 
 def map_splits(split_iter, inputs, cont_dim=None):
-    """Get a dictionary of prescribed splits."""
+    """generate a dictionary of inputs prescribed by the splitter."""
     if cont_dim is None:
         cont_dim = {}
     for split in split_iter:

pydra/engine/state.py

@@ -132,6 +132,7 @@ def splitter_rpn_final(self):
 
     @property
     def splitter_final(self):
+        """ final splitter, after removing the combined fields"""
         return hlpst.rpn2splitter(self.splitter_rpn_final)
 
     @property
@@ -153,6 +154,7 @@ def splitter_rpn_compact(self):
 
     @property
     def right_splitter(self):
+        """ current state splitter (i.e. the Right part)"""
         lr_flag = self._left_right_check(self.splitter)
         if lr_flag == "Left":
             return None
@@ -161,8 +163,19 @@ def right_splitter(self):
         elif lr_flag == "[Left, Right]":
             return self.splitter[1]
 
+    @property
+    def right_splitter_rpn(self):
+        if self.right_splitter:
+            right_splitter_rpn = hlpst.splitter2rpn(
+                self.right_splitter, other_states=self.other_states
+            )
+            return right_splitter_rpn
+        else:
+            return []
+
     @property
     def left_splitter(self):
+        """ splitters from the previous stated (i.e. the Light part)"""
         if hasattr(self, "_left_splitter"):
             return self._left_splitter
         else:
@@ -261,15 +274,13 @@ def inner_inputs(self):
             return {}
 
     def update_connections(self, new_other_states=None, new_combiner=None):
+        """ updating states connections and input groups"""
         if new_other_states:
             self.other_states = new_other_states
         self._connect_splitters()
         if new_combiner:
             self.combiner = new_combiner
         self.set_input_groups()
-        self.states_val = []
-        self.inputs_ind = []
-        self.final_combined_ind_mapping = {}
 
     def _connect_splitters(self):
         """
@@ -308,7 +319,7 @@ def _connect_splitters(self):
             self.splitter = deepcopy(self._left_splitter)
 
     def _complete_left(self, left=None):
-        """Add all splitters from previous nodes (completing left part)."""
+        """Add all splitters from previous nodes (completing the Left part)."""
         if left:
             rpn_left = hlpst.splitter2rpn(
                 left, other_states=self.other_states, state_fields=False
@@ -470,6 +481,7 @@ def push_new_states(self):
                 self.groups_stack_final.append(stack)
 
     def splitter_validation(self):
+        """ validating if the splitter is correct (after all states are connected)"""
         for spl in self.splitter_rpn_compact:
             if not (
                 spl in [".", "*"]
@@ -483,6 +495,7 @@ def splitter_validation(self):
                 )
 
     def combiner_validation(self):
+        """ validating if the combiner is correct (after all states are connected)"""
         if self.combiner:
             if not self.splitter:
                 raise Exception("splitter has to be set before setting combiner")
@@ -516,7 +529,6 @@ def prepare_states(self, inputs, cont_dim=None):
             for nm, (st, _) in self.other_states.items():
                 # 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, cont_dim=cont_dim)
                 self.inputs.update(st.inputs)
         self.prepare_states_ind()
@@ -532,6 +544,7 @@ def prepare_states_ind(self):
         # removing elements that are connected to inner splitter
         # (they will be taken into account in hlpst.splits anyway)
         # _comb part will be used in prepare_states_combined_ind
+        # TODO: need tests in test_Workflow.py
         elements_to_remove = []
         elements_to_remove_comb = []
         for name, (st, inp) in self.other_states.items():
@@ -546,7 +559,7 @@ 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(
+        values_out_pr, keys_out_pr = hlpst.splits(
             partial_rpn,
             self.inputs,
             inner_inputs=self.inner_inputs,
@@ -581,9 +594,8 @@ def prepare_states_combined_ind(self, elements_to_remove_comb):
         combined_rpn = hlpst.remove_inp_from_splitter_rpn(
             deepcopy(partial_rpn), self.right_combiner_all + self.left_combiner_all
         )
-        # TODO: create a function for this!!
         if combined_rpn:
-            val_r, key_r, _ = hlpst.splits(
+            val_r, key_r = hlpst.splits(
                 combined_rpn,
                 self.inputs,
                 inner_inputs=self.inner_inputs,
@@ -596,7 +608,6 @@ def prepare_states_combined_ind(self, elements_to_remove_comb):
 
         keys_out = key_r
         if values:
-            # NOW TODO: move to init?
             self.ind_l_final = values
             self.keys_final = keys_out
             # groups after combiner
@@ -628,28 +639,19 @@ def prepare_states_val(self):
 
     def prepare_inputs(self):
         """
-        Get inputs ready.
+        Preparing inputs indices, merges input from previous states.
 
-        1. Remove elements that come from connected states.
-        2. Merge elements that come from outputs of previous nodes.
-        3. Remove elements connected to the inner splitter.
+        Includes indices for fields from inner splitters
+        (removes elements connected to the inner splitters fields).
 
         """
         if not self.other_states:
             self.inputs_ind = self.states_ind
         else:
-            # removing elements that come from connected states
-            elements_to_remove = [
-                spl
-                for spl in self.splitter_rpn_compact
-                if spl[1:] in self.other_states.keys()
-            ]
-            partial_rpn = hlpst.remove_inp_from_splitter_rpn(
-                deepcopy(self.splitter_rpn_compact), elements_to_remove
-            )
-            if partial_rpn:
-                values_inp, keys_inp, _ = hlpst.splits(
-                    partial_rpn,
+            # elements from the current node (the Right part)
+            if self.right_splitter_rpn:
+                values_inp, keys_inp = hlpst.splits(
+                    self.right_splitter_rpn,
                     self.inputs,
                     inner_inputs=self.inner_inputs,
                     cont_dim=self.cont_dim,
@@ -700,5 +702,6 @@ def prepare_inputs(self):
             # iter_splits using inputs from current state/node
             self.inputs_ind = list(hlpst.iter_splits(inputs_ind, keys_inp))
             # removing elements that are connected to inner splitter
+            # TODO - add tests to test_workflow.py (not sure if we want to remove it)
             for el in connected_to_inner:
                 [dict.pop(el) for dict in self.inputs_ind]