def _build_outputs(self) -> List[DataPlaceholder]:
port = len(self._nodes)
outputs = []
for i in range(self._n_outputs):
name = make_name(make_name(self._name, port, sep=":"), i)
outputs.append(DataPlaceholder(self, port, name))
return outputs
def build_output_edges(self, model, outer_port, container):
root_name = make_name(model.name, outer_port, sep=":")
keys = self.get_innermost_outputs_keys(model, outer_port)
for (outer_port,
inner_output), output in safezip2(keys, model._internal_outputs):
src = self.node_names[outer_port, inner_output.node]
dst = make_name(root_name, output.name)
label = output.name
container.add_node(dummy_dot_node(dst))
container.add_edge(dot_edge(src, dst, label, "black"))
def _generate_unique_name(self):
name = self.__class__.__name__
n_instances = self._names.get(name, 0)
unique_name = make_name(name, n_instances, sep="_")
n_instances += 1
self._names[name] = n_instances
return unique_name
def test_make_name(sep, expected):
assert make_name("x", "y", sep=sep)
def name(self):
return make_name(self.step.name, self.port, sep=":")
def _build_outputs(self) -> List[DataPlaceholder]:
outputs = []
for i in range(self._n_outputs):
name = make_name(self._name, i)
outputs.append(DataPlaceholder(self, name))
return outputs
def transform(self, model, outer_port=0, container=None, level=None):
"""Transform model graph to a dot graph. It will transform nested sub-models
recursively, in which case it returns the dot nodes of the sub-model where the
enclosing model should connect the edges of the steps that precede and follow
the sub-model.
"""
container = (pydot.Dot(graph_type="digraph", **self.dot_kwargs)
if container is None else container)
level = 0 if level is None else level
root_name = make_name(model.name, outer_port, sep=":")
# Add nodes
for node in model.graph:
if _is_input(node):
name = make_name(root_name, node.step.name)
label = node.step.name
container.add_node(dot_input_node(name, label))
elif _is_model(node) and self.expand_nested:
name = make_name(root_name, node.name)
label = node.name
cluster = dot_cluster(name, label)
container.add_subgraph(cluster)
self.inner_dot_nodes[outer_port, node] = self.transform(
node.step, node.port, cluster, level + 1)
else:
name = make_name(root_name, node.name)
label = node.name
container.add_node(dot_node(name, label))
self.node_names[outer_port, node] = name
# Add edges
for parent_node, node, dataplaceholders in model.graph.edges:
for d in dataplaceholders:
color = "orange" if d in node.targets else "black"
if (_is_model(parent_node)
or _is_model(node)) and self.expand_nested:
if _is_model(parent_node):
output_srcs, *_ = self.inner_dot_nodes[outer_port,
parent_node]
src = output_srcs[parent_node.outputs.index(d)]
else:
src = self.node_names[outer_port, parent_node]
if _is_model(node):
if d in node.targets:
*_, target_dsts = self.inner_dot_nodes[outer_port,
node]
dst = target_dsts[node.targets.index(d)]
else:
_, input_dsts, _ = self.inner_dot_nodes[outer_port,
node]
dst = input_dsts[node.inputs.index(d)]
else:
dst = self.node_names[outer_port, node]
else:
# Not expanded case, or step -> step case
color = "orange" if d in node.targets else "black"
src = self.node_names[outer_port, parent_node]
dst = self.node_names[outer_port, node]
label = d.name
container.add_edge(dot_edge(src, dst, label, color))
if self.expand_nested and level > 0:
return self.get_internal_dot_nodes(model, outer_port)
else:
self.build_output_edges(model, outer_port, container)
return container