def test_graph_add_outgoing(self):
""" test adding a node and all its outgoing nodes to a graph"""
nodes = self.create_provenance()
graph = graph_mod.Graph()
graph.add_outgoing(nodes.calcf1)
self.assertEqual(graph.nodes,
set([nodes.calcf1.pk, nodes.pd3.pk, nodes.fd1.pk]))
self.assertEqual(
graph.edges,
set([(nodes.calcf1.pk, nodes.pd3.pk,
LinkPair(LinkType.CREATE, 'output1')),
(nodes.calcf1.pk, nodes.fd1.pk,
LinkPair(LinkType.CREATE, 'output2'))]))
def add_outgoing(self, node, link_types=(), annotate_links=None, return_pks=True):
"""add nodes and edges for outgoing links to a node
:param node: node or node pk
:type node: aiida.orm.nodes.node.Node or int
:param link_types: filter by link types (Default value = ())
:type link_types: str or tuple[str] or aiida.common.links.LinkType or tuple[aiida.common.links.LinkType]
:param annotate_links: label edges with the link 'label', 'type' or 'both' (Default value = None)
:type annotate_links: bool or str
:param return_pks: whether to return a list of nodes, or list of node pks (Default value = True)
:type return_pks: bool
:returns: list of nodes or node pks
"""
if annotate_links not in [None, False, 'label', 'type', 'both']:
raise AssertionError('annotate_links must be one of False, "label", "type" or "both"')
node = self.add_node(node)
nodes = []
for link_triple in node.get_outgoing(link_type=self._convert_link_types(link_types)).link_triples:
self.add_node(link_triple.node)
link_pair = LinkPair(link_triple.link_type, link_triple.link_label)
style = self._link_styles(
link_pair, add_label=annotate_links in ['label', 'both'], add_type=annotate_links in ['type', 'both']
)
self.add_edge(node, link_triple.node, link_pair, style=style)
nodes.append(link_triple.node.pk if return_pks else link_triple.node)
return nodes
def test_graph_recurse_ancestors_filter_links(self):
""" test adding nodes and all its (recursed) incoming nodes to a graph, but filter link types"""
nodes = self.create_provenance()
graph = graph_mod.Graph()
graph.recurse_ancestors(nodes.rd1, link_types=['create', 'input_calc'])
self.assertEqual(
graph.nodes,
set([nodes.rd1.pk, nodes.calc1.pk, nodes.pd0.pk, nodes.pd1.pk]))
self.assertEqual(
graph.edges,
set([(nodes.calc1.pk, nodes.rd1.pk,
LinkPair(LinkType.CREATE, 'output')),
(nodes.pd0.pk, nodes.calc1.pk,
LinkPair(LinkType.INPUT_CALC, 'input1')),
(nodes.pd1.pk, nodes.calc1.pk,
LinkPair(LinkType.INPUT_CALC, 'input2'))]))
def test_graph_add_incoming(self):
""" test adding a node and all its incoming nodes to a graph"""
nodes = self.create_provenance()
graph = graph_mod.Graph()
graph.add_incoming(nodes.calc1)
self.assertEqual(
graph.nodes,
set([nodes.calc1.pk, nodes.pd0.pk, nodes.pd1.pk, nodes.wc1.pk]))
self.assertEqual(
graph.edges,
set([(nodes.pd0.pk, nodes.calc1.pk,
LinkPair(LinkType.INPUT_CALC, 'input1')),
(nodes.pd1.pk, nodes.calc1.pk,
LinkPair(LinkType.INPUT_CALC, 'input2')),
(nodes.wc1.pk, nodes.calc1.pk,
LinkPair(LinkType.CALL_CALC, 'call1'))]))
def add_outgoing(self, node, link_types=(), annotate_links=None, return_pks=True):
"""add nodes and edges for outgoing links to a node
:param node: node or node pk
:type node: aiida.orm.nodes.node.Node or int
:param link_types: filter by link types (Default value = ())
:type link_types: str or tuple[str] or aiida.common.links.LinkType or tuple[aiida.common.links.LinkType]
:param annotate_links: label edges with the link 'label', 'type' or 'both' (Default value = None)
:type annotate_links: bool or str
:param return_pks: whether to return a list of nodes, or list of node pks (Default value = True)
:type return_pks: bool
:returns: list of nodes or node pks
"""
if annotate_links not in [None, False, 'label', 'type', 'both']:
raise ValueError(
'annotate_links must be one of False, "label", "type" or "both"\ninstead, it is: {}'.
format(annotate_links)
)
# outgoing nodes are found traversing forwards
node_pk = self._load_node(node).pk
valid_link_types = self._default_link_types(link_types)
valid_link_types = self._convert_link_types(valid_link_types)
traversed_graph = traverse_graph(
(node_pk,),
max_iterations=1,
get_links=True,
links_forward=valid_link_types,
)
traversed_nodes = orm.QueryBuilder().append(
orm.Node,
filters={'id': {
'in': traversed_graph['nodes']
}},
project=['id', '*'],
tag='node',
)
traversed_nodes = {query_result[0]: query_result[1] for query_result in traversed_nodes.all()}
for _, traversed_node in traversed_nodes.items():
self.add_node(traversed_node, style_override=None)
for link in traversed_graph['links']:
source_node = traversed_nodes[link.source_id]
target_node = traversed_nodes[link.target_id]
link_pair = LinkPair(self._convert_link_types(link.link_type)[0], link.link_label)
link_style = self._link_styles(
link_pair, add_label=annotate_links in ['label', 'both'], add_type=annotate_links in ['type', 'both']
)
self.add_edge(source_node, target_node, link_pair, style=link_style)
if return_pks:
return list(traversed_nodes.keys())
# else:
return list(traversed_nodes.values())
def test_graph_recurse_descendants(self):
""" test adding nodes and all its (recursed) incoming nodes to a graph"""
nodes = self.create_provenance()
graph = graph_mod.Graph()
graph.recurse_descendants(nodes.rd1)
self.assertEqual(
graph.nodes,
set([nodes.rd1.pk, nodes.calcf1.pk, nodes.pd3.pk, nodes.fd1.pk]))
self.assertEqual(
graph.edges,
set([
(nodes.rd1.pk, nodes.calcf1.pk,
LinkPair(LinkType.INPUT_CALC, 'input1')),
(nodes.calcf1.pk, nodes.pd3.pk,
LinkPair(LinkType.CREATE, 'output1')),
(nodes.calcf1.pk, nodes.fd1.pk,
LinkPair(LinkType.CREATE, 'output2')),
]))
def recurse_ancestors(self,
origin,
depth=None,
link_types=(),
annotate_links=False,
origin_style=None,
include_process_outputs=False,
print_func=None):
"""add nodes and edges from an origin recursively,
following incoming links
:param origin: node or node pk/uuid
:type origin: aiida.orm.nodes.node.Node or int
:param depth: if not None, stop after travelling a certain depth into the graph (Default value = None)
:type depth: None or int
:param link_types: filter by subset of link types (Default value = ())
:type link_types: tuple or str
:param annotate_links: label edges with the link 'label', 'type' or 'both' (Default value = False)
:type annotate_links: bool
:param origin_style: node style map for origin node (Default value = None)
:type origin_style: None or dict
:param include_process_outputs: include outgoing links for all processes (Default value = False)
:type include_process_outputs: bool
:param print_func: a function to stream information to, i.e. print_func(str)
.. deprecated:: 1.1.0
`print_func` will be removed in `v2.0.0`
"""
# pylint: disable=too-many-arguments,too-many-locals
import warnings
from aiida.common.warnings import AiidaDeprecationWarning
if print_func:
warnings.warn( # pylint: disable=no-member
'`print_func` is deprecated because graph traversal has been refactored',
AiidaDeprecationWarning)
# Get graph traversal rules where the given link types and direction are all set to True,
# and all others are set to False
origin_pk = origin if isinstance(origin, int) else origin.pk
valid_link_types = self._default_link_types(link_types)
valid_link_types = self._convert_link_types(valid_link_types)
traversed_graph = traverse_graph(
(origin_pk, ),
max_iterations=depth,
get_links=True,
links_backward=valid_link_types,
)
# Traverse forward along input_work and input_calc links from all nodes traversed in the previous step
# and join the result with the original traversed graph. This includes calculation outputs in the Graph
if include_process_outputs:
traversed_outputs = traverse_graph(
traversed_graph['nodes'],
max_iterations=1,
get_links=True,
links_forward=[LinkType.CREATE, LinkType.RETURN])
traversed_graph['nodes'] = traversed_graph['nodes'].union(
traversed_outputs['nodes'])
traversed_graph['links'] = traversed_graph['links'].union(
traversed_outputs['links'])
# Do one central query for all nodes in the Graph and generate a {id: Node} dictionary
traversed_nodes = orm.QueryBuilder().append(
orm.Node,
filters={'id': {
'in': traversed_graph['nodes']
}},
project=['id', '*'],
tag='node',
)
traversed_nodes = {
query_result[0]: query_result[1]
for query_result in traversed_nodes.all()
}
# Pop the origin node and add it to the graph, applying custom styling
origin_node = traversed_nodes.pop(origin_pk)
self.add_node(origin_node, style_override=origin_style)
# Add all traversed nodes to the graph with default styling
for _, traversed_node in traversed_nodes.items():
self.add_node(traversed_node, style_override=None)
# Add the origin node back into traversed nodes so it can be found for adding edges
traversed_nodes[origin_pk] = origin_node
# Add all links to the Graph, using the {id: Node} dictionary for queryless Node retrieval, applying
# appropriate styling
for link in traversed_graph['links']:
source_node = traversed_nodes[link.source_id]
target_node = traversed_nodes[link.target_id]
link_pair = LinkPair(
self._convert_link_types(link.link_type)[0], link.link_label)
link_style = self._link_styles(link_pair,
add_label=annotate_links
in ['label', 'both'],
add_type=annotate_links
in ['type', 'both'])
self.add_edge(source_node,
target_node,
link_pair,
style=link_style)
