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_empty_input(self): """Testing empty input.""" all_links = [ LinkType.INPUT_CALC, LinkType.CALL_CALC, LinkType.CREATE, LinkType.INPUT_WORK, LinkType.CALL_WORK, LinkType.RETURN ] obtained_results = traverse_graph([], links_forward=all_links, links_backward=all_links) self.assertEqual(obtained_results['nodes'], set()) self.assertEqual(obtained_results['links'], None) obtained_results = traverse_graph([], get_links=True, links_forward=all_links, links_backward=all_links) self.assertEqual(obtained_results['nodes'], set()) self.assertEqual(obtained_results['links'], set())
def test_traversal_errors(self): """This will test the errors of the traversers.""" from aiida.common.exceptions import NotExistent from aiida import orm test_node = orm.Data().store() false_node = -1 with self.assertRaises(NotExistent): _ = traverse_graph([false_node]) with self.assertRaises(TypeError): _ = traverse_graph(['not a node']) with self.assertRaises(TypeError): _ = traverse_graph('not a list') with self.assertRaises(TypeError): _ = traverse_graph([test_node], links_forward=1984) with self.assertRaises(TypeError): _ = traverse_graph([test_node], links_backward=['not a link'])
def _single_test(self, starting_nodes=(), expanded_nodes=(), links_forward=(), links_backward=()): """Auxiliary method to perform a single test run and assertion""" obtained_nodes = traverse_graph( starting_nodes, links_forward=links_forward, links_backward=links_backward, )['nodes'] expected_nodes = set(starting_nodes + expanded_nodes) self.assertEqual(obtained_nodes, expected_nodes)
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)
def test_traversal_cycle(self): """ This will test that cycles don't go into infinite loops by testing a graph with two data nodes data_take and data_drop and a work_select that takes both as input but returns only data_take """ from aiida import orm data_take = orm.Data().store() data_drop = orm.Data().store() work_select = orm.WorkflowNode() work_select.add_incoming(data_take, link_type=LinkType.INPUT_WORK, link_label='input_take') work_select.add_incoming(data_drop, link_type=LinkType.INPUT_WORK, link_label='input_drop') work_select.store() data_take.add_incoming(work_select, link_type=LinkType.RETURN, link_label='return_link') data_take = data_take.pk data_drop = data_drop.pk work_select = work_select.pk every_node = [data_take, data_drop, work_select] for single_node in every_node: expected_nodes = set([single_node]) obtained_nodes = traverse_graph([single_node])['nodes'] self.assertEqual(obtained_nodes, expected_nodes) links_forward = [LinkType.INPUT_WORK, LinkType.RETURN] links_backward = [] # Forward: data_drop to (input) work_select to (return) data_take obtained_nodes = traverse_graph([data_drop], links_forward=links_forward, links_backward=links_backward)['nodes'] expected_nodes = set(every_node) self.assertEqual(obtained_nodes, expected_nodes) # Forward: data_take to (input) work_select (data_drop is not returned) obtained_nodes = traverse_graph([data_take], links_forward=links_forward, links_backward=links_backward)['nodes'] expected_nodes = set([work_select, data_take]) self.assertEqual(obtained_nodes, expected_nodes) # Forward: work_select to (return) data_take (data_drop is not returned) obtained_nodes = traverse_graph([work_select], links_forward=links_forward, links_backward=links_backward)['nodes'] self.assertEqual(obtained_nodes, expected_nodes) links_forward = [] links_backward = [LinkType.INPUT_WORK, LinkType.RETURN] # Backward: data_drop is not returned so it has no backward link expected_nodes = set([data_drop]) obtained_nodes = traverse_graph([data_drop], links_forward=links_forward, links_backward=links_backward)['nodes'] self.assertEqual(obtained_nodes, expected_nodes) # Backward: data_take to (return) work_select to (input) data_drop expected_nodes = set(every_node) obtained_nodes = traverse_graph([data_take], links_forward=links_forward, links_backward=links_backward)['nodes'] self.assertEqual(obtained_nodes, expected_nodes) # Backward: work_select to (input) data_take and data_drop obtained_nodes = traverse_graph([work_select], links_forward=links_forward, links_backward=links_backward)['nodes'] self.assertEqual(obtained_nodes, expected_nodes)