Refactor iterable standardization to account for the itersource alternate format.

Author: FredLoney

nipype/pipeline/utils.py

@@ -584,23 +584,15 @@ def generate_expanded_graph(graph_in):
  """
  logger.debug("PE: expanding iterables")
  graph_in = _remove_nonjoin_identity_nodes(graph_in, keep_iterables=True)
- # convert list of tuples to dict fields
+ # standardize the iterables as {(field, function)} dictionaries
  for node in graph_in.nodes_iter():
- if isinstance(node.iterables, tuple):
- node.iterables = [node.iterables]
- for node in graph_in.nodes_iter():
- if isinstance(node.iterables, list):
- node.iterables = dict(map(lambda(x): (x[0],
- lambda: x[1]),
- node.iterables))
+ if node.iterables:
+ _standardize_iterables(node)
  allprefixes = list('abcdefghijklmnopqrstuvwxyz')
 
  # 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:
@@ -626,22 +618,25 @@ def generate_expanded_graph(graph_in):
  % (src, dest))
 
  if inode.itersource:
+ # the itersource is a (node name, fields) tuple
+ src_name, src_fields = inode.itersource
+ # convert a single field to a list
+ if isinstance(src_fields, str):
+ src_fields = [src_fields]
  # find the unique iterable source node in the graph
- iter_src = None
- for node in graph_in.nodes_iter():
- if (node.name == inode.itersource
- and nx.has_path(graph_in, node, inode)):
- iter_src = node
- break
- if not iter_src or not iter_fld_dict.has_key(inode.itersource):
+ try:
+ iter_src = next((node for node in graph_in.nodes_iter()
+ if node.name == src_name
+ and nx.has_path(graph_in, node, inode)))
+ except StopIteration:
  raise ValueError("The node %s itersource %s was not found"
- " among the iterable nodes %s"
- % (inode, inode.itersource, iter_fld_dict.keys()))
+ " among the iterable predecessor nodes"
+ % (inode, src_name))
+ logger.debug("The node %s has iterable source node %s"
+ % (inode, iter_src))
  # look up the iterables for this particular itersource descendant
  # using the iterable source ancestor values as a key
  iterables = {}
- # the source node iterables fields
- src_fields = iter_fld_dict[inode.itersource]
  # the source node iterables values
  src_values = [getattr(iter_src.inputs, field) for field in src_fields]
  # if there is one source field, then the key is the the source value,
@@ -714,9 +709,12 @@ def generate_expanded_graph(graph_in):
  for src_id, edge_data in old_edge_dict.iteritems():
  if node._id.startswith(src_id):
  expansions[src_id].append(node)
+ for in_id, in_nodes in expansions.iteritems():
+ logger.debug("The join node %s input %s was expanded"
+ " to %d nodes." %(jnode, in_id, len(in_nodes)))
  # preserve the node iteration order by sorting on the node id
- for src_nodes in expansions.itervalues():
- src_nodes.sort(key=lambda node: node._id)
+ for in_nodes in expansions.itervalues():
+ in_nodes.sort(key=lambda node: node._id)
 
  # the number of iterations.
  iter_cnt = count_iterables(iterables, inode.synchronize)
@@ -731,28 +729,28 @@ def generate_expanded_graph(graph_in):
  # field 'in' are qualified as ('out_file', 'in1') and
  # ('out_file', 'in2'), resp. This preserves connection port
  # integrity.
- for old_id, src_nodes in expansions.iteritems():
+ for old_id, in_nodes in expansions.iteritems():
  # reconnect each replication of the current join in-edge
  # source
- for si, src in enumerate(src_nodes):
+ for in_idx, in_node in enumerate(in_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]
- slots = slot_dicts[si]
- for ci, connect in enumerate(connects):
+ slots = slot_dicts[in_idx]
+ for con_idx, connect in enumerate(connects):
  src_field, dest_field = connect
  # qualify a join destination field name
  if dest_field in slots:
  slot_field = slots[dest_field]
- connects[ci] = (src_field, slot_field)
+ connects[con_idx] = (src_field, slot_field)
  logger.debug("Qualified the %s -> %s join field"
  " %s as %s." %
- (src, jnode, dest_field, slot_field))
- graph_in.add_edge(src, jnode, newdata)
+ (in_node, jnode, dest_field, slot_field))
+ graph_in.add_edge(in_node, jnode, newdata)
  logger.debug("Connected the join node %s subgraph to the"
- " expanded join point %s" % (jnode, src))
+ " expanded join point %s" % (jnode, in_node))
 
  #nx.write_dot(graph_in, '%s_post.dot' % node)
  # the remaining iterable nodes
@@ -792,6 +790,65 @@ def _iterable_nodes(graph_in):
  inodes_src = [node for node in inodes if node.itersource]
  inodes_no_src.reverse()
  return inodes_no_src + inodes_src
+
+def _standardize_iterables(node):
+ """Converts the given iterables to a {field: function} dictionary,
+ if necessary, where the function returns a list."""
+ # trivial case
+ if not node.iterables:
+ return
+ iterables = node.iterables
+ # The candidate iterable fields
+ fields = set(node.inputs.copyable_trait_names())
+ 
+ # Convert a tuple to a list
+ if isinstance(iterables, tuple):
+ iterables = [iterables]
+ # Convert a list to a dictionary
+ if isinstance(iterables, list):
+ # Synchronize iterables can be in [fields, value tuples] format
+ # rather than [(field, value list), (field, value list), ...]
+ if node.synchronize and len(iterables) == 2:
+ first, last = iterables
+ if all((isinstance(item, str) and item in fields
+ for item in first)):
+ iterables = _transpose_iterables(first, last)
+ # Validate the format
+ for item in iterables:
+ try:
+ if len(item) != 2:
+ raise ValueError("The %s iterables do not consist of"
+ " (field, values) pairs" % node.name)
+ except TypeError, e:
+ raise TypeError("The %s iterables is not iterable: %s"
+ % (node.name, e))
+ # Convert the values to functions. This is a legacy Nipype
+ # requirement with unknown rationale.
+ iter_items = map(lambda(field, value): (field, lambda: value),
+ iterables)
+ # Make the iterables dictionary
+ iterables = dict(iter_items)
+ elif not isinstance(iterables, dict):
+ raise ValueError("The %s iterables type is not a list or a dictionary:"
+ " %s" % (node.name, iterables.__class__))
+ 
+ # Validate the iterable fields
+ for field in iterables.iterkeys():
+ if field not in fields:
+ raise ValueError("The %s iterables field is unrecognized: %s"
+ % (node.name, field))
+ 
+ # Assign to the standard form
+ node.iterables = iterables
+
+def _transpose_iterables(fields, values):
+ """
+ Converts the given fields and tuple values into a list of
+ iterable (field: value list) pairs, suitable for setting
+ a node iterables property.
+ """
+ return zip(fields, [filter(lambda(v): v != None, transpose)
+ for transpose in zip(*values)])
 
 def export_graph(graph_in, base_dir=None, show=False, use_execgraph=False,
  show_connectinfo=False, dotfilename='graph.dot', format='png',