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)