Support join on a node with an itersource.

nipype/pipeline/tests/test_join.py

@@ -35,9 +35,9 @@ def _list_outputs(self):
  outputs['output1'] = self.inputs.input1 + self.inputs.inc
  return outputs
 
-_sum = 0
+_sums = []
 
-_sum_operands = None
+_sum_operands = []
 
 class SumInputSpec(nib.TraitedSpec):
  input1 = nib.traits.List(nib.traits.Int, mandatory=True, desc='input')
@@ -58,8 +58,10 @@ def _list_outputs(self):
  global _sum
  global _sum_operands
  outputs = self._outputs().get()
- _sum_operands = outputs['operands'] = self.inputs.input1
- _sum = outputs['output1'] = sum(self.inputs.input1)
+ outputs['operands'] = self.inputs.input1
+ _sum_operands.append(outputs['operands'])
+ outputs['output1'] = sum(self.inputs.input1)
+ _sums.append(outputs['output1'])
  return outputs
 
 
@@ -150,9 +152,11 @@ def test_join_expansion():
  # Nipype factors away the IdentityInterface.
  assert_equal(len(result.nodes()), 8, "The number of expanded nodes is incorrect.")
  # the join Sum result is (1 + 1 + 1) + (2 + 1 + 1)
- assert_equal(_sum, 7, "The join Sum output value is incorrect: %s." % _sum)
+ assert_equal(len(_sums), 1,
+ "The number of join outputs is incorrect")
+ assert_equal(_sums[0], 7, "The join Sum output value is incorrect: %s." % _sums[0])
  # the join input preserves the iterables input order
- assert_equal(_sum_operands, [3, 4], "The join Sum input is incorrect: %s." % _sum_operands)
+ assert_equal(_sum_operands[0], [3, 4], "The join Sum input is incorrect: %s." % _sum_operands[0])
  # there are two iterations of the post-join node in the iterable path
  assert_equal(len(_products), 2,
  "The number of iterated post-join outputs is incorrect")
@@ -209,12 +213,14 @@ def test_unique_join_node():
  wf.run()
 
  # the join length is the number of unique inputs
- assert_equal(_sum_operands, [4, 2, 3], "The unique join output value is incorrect: %s." % _sum_operands)
+ assert_equal(_sum_operands[0], [4, 2, 3], "The unique join output value is incorrect: %s." % _sum_operands[0])
 
  os.chdir(cwd)
  rmtree(wd)
 
 def test_identity_join_node():
+ global _sum_operands
+ _sum_operands = []
  cwd = os.getcwd()
  wd = mkdtemp()
  os.chdir(wd)
@@ -241,8 +247,9 @@ def test_identity_join_node():
  # node and 1 post-join node. Nipype factors away the iterable input
  # IdentityInterface but keeps the join IdentityInterface.
  assert_equal(len(result.nodes()), 5, "The number of expanded nodes is incorrect.")
- assert_equal(_sum_operands, [2, 3, 4],
- "The join Sum input is incorrect: %s." %_sum_operands)
+ assert_equal(_sum_operands[0], [2, 3, 4],
+ "The join Sum input is incorrect: %s." %_sum_operands[0])
+
  os.chdir(cwd)
  rmtree(wd)
 
@@ -282,9 +289,66 @@ def test_multifield_join_node():
  # the product inputs are [2, 4], [2, 5], [3, 4], [3, 5]
  assert_equal(_products, [8, 10, 12, 15],
  "The post-join products is incorrect: %s." % _products)
+
+ os.chdir(cwd)
+ rmtree(wd)
+
+def test_itersource_join_source_node():
+ cwd = os.getcwd()
+ wd = mkdtemp()
+ os.chdir(wd)
+
+ # Make the workflow.
+ wf = pe.Workflow(name='test')
+ # the iterated input node
+ inputspec = pe.Node(IdentityInterface(fields=['n']), name='inputspec')
+ inputspec.iterables = [('n', [1, 2])]
+ # an intermediate node in the first iteration path
+ pre_join1 = pe.Node(IncrementInterface(), name='pre_join1')
+ wf.connect(inputspec, 'n', pre_join1, 'input1')
+ # an iterable pre-join node with an itersource
+ pre_join2 = pe.Node(ProductInterface(), name='pre_join2')
+ pre_join2.itersource = 'inputspec'
+ pre_join2.iterables = ('input1', {1: [3, 4], 2: [5, 6]})
+ wf.connect(pre_join1, 'output1', pre_join2, 'input2')
+ # an intermediate node in the second iteration path
+ pre_join3 = pe.Node(IncrementInterface(), name='pre_join3')
+ wf.connect(pre_join2, 'output1', pre_join3, 'input1')
+ # the join node
+ join = pe.JoinNode(IdentityInterface(fields=['vector']), joinsource='pre_join2',
+ joinfield='vector', name='join')
+ wf.connect(pre_join3, 'output1', join, 'vector')
+ # a join successor node
+ post_join1 = pe.Node(SumInterface(), name='post_join1')
+ wf.connect(join, 'vector', post_join1, 'input1')
+
+ result = wf.run()
+
+ # the expanded graph contains
+ # 1 pre_join1 replicate for each inputspec iteration,
+ # 2 pre_join2 replicates for each inputspec iteration,
+ # 1 pre_join3 for each pre_join2 iteration,
+ # 1 join replicate for each inputspec iteration and
+ # 1 post_join1 replicate for each join replicate =
+ # 2 + (2 * 2) + 4 + 2 + 2 = 14 expansion graph nodes.
+ # Nipype factors away the iterable input
+ # IdentityInterface but keeps the join IdentityInterface.
+ assert_equal(len(result.nodes()), 14, "The number of expanded nodes is incorrect.")
+ # The first join inputs are:
+ # 1 + (3 * 2) and 1 + (4 * 2)
+ # The second join inputs are:
+ # 1 + (5 * 3) and 1 + (6 * 3)
+ # the post-join nodes execution order is indeterminate;
+ # therefore, compare the lists item-wise.
+ assert_true([16, 19] in _sum_operands,
+ "The join Sum input is incorrect: %s." % _sum_operands)
+ assert_true([7, 9] in _sum_operands,
+ "The join Sum input is incorrect: %s." % _sum_operands)
+
  os.chdir(cwd)
  rmtree(wd)
 
+
 if __name__ == "__main__":
  import nose
 

nipype/pipeline/utils.py

@@ -591,13 +591,34 @@ def generate_expanded_graph(graph_in):
 
  # the iterable nodes
  inodes = _iterable_nodes(graph_in)
+ logger.debug("Detected iterable nodes %s" % inodes)
  # record the iterable fields, since expansion removes them
  iter_fld_dict = {inode.name: inode.iterables.keys()
  for inode in inodes}
  # while there is an iterable node, expand the iterable node's
  # subgraphs
  while inodes:
  inode = inodes[0]
+ logger.debug("Expanding the iterable node %s..." % inode)
+
+ # the join successor nodes of the current iterable node
+ jnodes = [node for node in graph_in.nodes_iter()
+ if hasattr(node, 'joinsource')
+ and inode.name == node.joinsource
+ and nx.has_path(graph_in, inode, node)]
+
+ # excise the join in-edges. save the excised edges in a
+ # {jnode: {source name: (destination name, edge data)}}
+ # dictionary
+ jedge_dict = {}
+ for jnode in jnodes:
+ in_edges = jedge_dict[jnode] = {}
+ for src, dest, data in graph_in.in_edges_iter(jnode, True):
+ in_edges[src._id] = data
+ graph_in.remove_edge(src, dest)
+ logger.debug("Excised the %s -> %s join node in-edge."
+ % (src, dest))
+
  if inode.itersource:
  # find the unique iterable source node in the graph
  iter_src = None
@@ -607,8 +628,9 @@ def generate_expanded_graph(graph_in):
  iter_src = node
  break
  if not iter_src or not iter_fld_dict.has_key(inode.itersource):
- raise ValueError("Iterable node %s source node not found: %s"
- % (inode, inode.itersource))
+ raise ValueError("The node %s itersource %s was not found"
+ " among the iterable nodes %s"
+ % (inode, inode.itersource, iter_fld_dict.keys()))
  # look up the iterables for this particular itersource descendant
  # using the iterable source ancestor values as a key
  iterables = {}
@@ -674,16 +696,22 @@ def generate_expanded_graph(graph_in):
  graph_in = _merge_graphs(graph_in, subnodes,
  subgraph, inode._hierarchy + inode._id,
  iterables, iterable_prefix)
-
+ 
  # reconnect the join nodes
  for jnode in jnodes:
- # the {node name: replicated nodes} dictionary
- node_name_dict = defaultdict(list)
+ # the {node id: edge data} dictionary for edges connecting
+ # to the join node in the unexpanded graph
+ old_edge_dict = jedge_dict[jnode]
+ # the edge source node replicates
+ expansions = defaultdict(list)
  for node in graph_in.nodes_iter():
- node_name_dict[node.name].append(node)
+ for src_id, edge_data in old_edge_dict.iteritems():
+ if node._id.startswith(src_id):
+ expansions[src_id].append(node)
  # preserve the node iteration order by sorting on the node id
- for nodes in node_name_dict.values():
- nodes.sort(key=str)
+ for src_nodes in expansions.itervalues():
+ src_nodes.sort(key=lambda node: node._id)
+
  # the number of iterations. this magic formula is borrowed
  # from _merge_graphs.
  iter_cnt = len(list(walk(iterables.items())))
@@ -698,30 +726,12 @@ def generate_expanded_graph(graph_in):
  # field 'in' are qualified as ('out_file', 'in1') and
  # ('out_file', 'in2'), resp. This preserves connection port
  # integrity.
- # the {source name: (destination name, edge data)} dictionary
- in_edges = jedge_dict[jnode]
- # for each join in-edge source name, the replicated source
- # nodes are the expansion graph nodes which match on the
- # source name. reconnect each replicated source node to the
- # join node.
- for src_name, tgt in in_edges.iteritems():
- dest_name, edge_data = tgt
- # there is a single join destination, since the join node
- # is not expanded
- dests = node_name_dict[dest_name]
- if not dests:
- raise Exception("The execution graph does not contain"
- " the join node: %s" % dest_name)
- elif len(dests) > 1:
- raise Exception("The execution graph contains more than"
- " one join node named %s: %s"
- % (dest_name, dests))
- else:
- dest = dests[0]
+ for old_id, src_nodes in expansions.iteritems():
  # reconnect each replication of the current join in-edge
  # source
- for si, src in enumerate(node_name_dict[src_name]):
- newdata = deepcopy(edge_data)
+ for si, src in enumerate(src_nodes):
+ olddata = old_edge_dict[old_id]
+ newdata = deepcopy(olddata)
  connects = newdata['connect']
  join_fields = [field for _, field in connects
  if field in dest.joinfield]
@@ -734,10 +744,10 @@ def generate_expanded_graph(graph_in):
  connects[ci] = (src_field, slot_field)
  logger.debug("Qualified the %s -> %s join field"
  " %s as %s." %
- (src, dest, dest_field, slot_field))
- graph_in.add_edge(src, dest, newdata)
+ (src, jnode, dest_field, slot_field))
+ graph_in.add_edge(src, jnode, newdata)
  logger.debug("Connected the join node %s subgraph to the"
- " expanded join point %s" % (dest, src))
+ " expanded join point %s" % (jnode, src))
 
  #nx.write_dot(graph_in, '%s_post.dot' % node)
  # the remaining iterable nodes