adding comments

Author: djarecka

pydra/engine/helpers_state.py

@@ -363,7 +363,7 @@ def iter_splits(iterable, keys):
 
 
 def input_shape(inp, cont_dim=1):
-    """Get input shape."""
+    """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
     shape = [len(inp)]
@@ -396,6 +396,9 @@ def splits(splitter_rpn, inputs, inner_inputs=None, cont_dim=None):
         input variables
     inner_inputs: dict, optional
         inner input specification
+    cont_dim: dict, optional
+        container dimension for input variable, specifies how nested is the intput,
+        if not specified 1 will be used for all inputs (so will not be flatten)
 
 
     Returns

pydra/engine/state.py

@@ -357,6 +357,7 @@ def prepare_states(self, inputs, cont_dim=None):
             specific elements from inputs that can be used running interfaces
 
         """
+        # container dimension for each input, specifies how nested the input is
         if cont_dim is None:
             self.cont_dim = {}
         else:

pydra/engine/tests/test_helpers_state.py

@@ -267,7 +267,7 @@ def test_splits_groups_comb(
         ),
         (
             (["a", "v"], "x"),
-            {"x": 2},
+            {"x": 2},  # input x is treated as 2d container, so x will be flatten
             [((0, 0), 0), ((0, 1), 1), ((1, 0), 2), ((1, 1), 3)],
             ["a", "v", "x"],
             [
@@ -306,7 +306,7 @@ def test_splits_1b(splitter, cont_dim, values, keys, splits):
         ((["a", "v"], "c"), None, {"a": [1, 2], "v": ["a", "b"], "c": [3, 4]}, True),
         (
             (["a", "v"], "c"),
-            {"c": 2},
+            {"c": 2},  # c is treated as 2d container
             {"a": [1, 2], "v": ["a", "b"], "c": [[3, 4], [5, 6]]},
             False,
         ),
@@ -442,7 +442,7 @@ def test_splits_2(splitter_rpn, inner_inputs, values, keys, splits):
         #     [["d211", "d212"], ["d221", "d222"]],
         # ],
     }
-    cont_dim = {"NB.b": 2}
+    cont_dim = {"NB.b": 2}  # will be treated as 2d container
     values_out, keys_out, _ = hlpst.splits(
         splitter_rpn, inputs, inner_inputs=inner_inputs, cont_dim=cont_dim
     )

pydra/engine/tests/test_state.py

@@ -566,7 +566,7 @@ def test_state_connect_innerspl_1():
 
     st2.prepare_states(
         inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]]},
-        cont_dim={"NB.b": 2},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
 
     assert st2.states_ind == [
@@ -613,7 +613,7 @@ def test_state_connect_innerspl_1a():
 
     st2.prepare_states(
         inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]]},
-        cont_dim={"NB.b": 2},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
 
     assert st2.states_ind == [
@@ -667,7 +667,7 @@ def test_state_connect_innerspl_2():
 
     st2.prepare_states(
         inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]},
-        cont_dim={"NB.b": 2},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
     assert st2.states_ind == [
         {"NB.c": 0, "NA.a": 0, "NB.b": 0},
@@ -732,7 +732,7 @@ def test_state_connect_innerspl_2a():
 
     st2.prepare_states(
         inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]},
-        cont_dim={"NB.b": 2},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
 
     assert st2.states_ind == [
@@ -803,7 +803,7 @@ def test_state_connect_innerspl_3():
             "NB.c": [13, 17],
             "NC.d": [33, 77],
         },
-        cont_dim={"NB.b": 2},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
 
     assert st2.states_ind == [
@@ -943,7 +943,7 @@ def test_state_connect_innerspl_4():
             "NC.f": [[23, 27], [33, 37]],
             "NC.d": [1, 2],
         },
-        cont_dim={"NC.f": 2},
+        cont_dim={"NC.f": 2},  # will be treated as 2d container
     )
     assert st3.states_ind == [
         {"NA.a": 0, "NB.b": 0, "NB.c": 0, "NC.d": 0},
@@ -1209,7 +1209,7 @@ def test_state_connect_innerspl_combine_1():
 
     st2.prepare_states(
         inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]},
-        cont_dim={"NB.b": 2},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
     # NOW TODO: checking st2.states_ind_final!!!
     assert st2.states_ind == [
@@ -1283,7 +1283,7 @@ def test_state_connect_innerspl_combine_2():
 
     st2.prepare_states(
         inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]},
-        cont_dim={"NB.b": 2},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
     assert st2.states_ind == [
         {"NB.c": 0, "NA.a": 0, "NB.b": 0},