def test_groups(self):
"""
Testing connection between (aiida-)groups and (aiida-)nodes, which are treated
as if they both were (graph-)nodes.
"""
node1 = orm.Data().store()
node2 = orm.Data().store()
node3 = orm.Data().store()
node4 = orm.Data().store()
group1 = orm.Group(label='group-01').store()
group1.add_nodes(node1)
group2 = orm.Group(label='group-02').store()
group2.add_nodes(node2)
group2.add_nodes(node3)
group3 = orm.Group(label='group-03').store()
group3.add_nodes(node4)
group4 = orm.Group(label='group-04').store()
group4.add_nodes(node4)
# Rule that only gets nodes connected by the same group
queryb = orm.QueryBuilder()
queryb.append(orm.Node, tag='nodes_in_set')
queryb.append(orm.Group, with_node='nodes_in_set', tag='groups_considered', filters={'type_string': 'user'})
queryb.append(orm.Data, with_group='groups_considered')
initial_node = [node2.id]
basket_inp = Basket(nodes=initial_node)
tested_rule = UpdateRule(queryb, max_iterations=np.inf)
basket_out = tested_rule.run(basket_inp.copy())
obtained = basket_out['nodes'].keyset
expected = set([node2.id, node3.id])
self.assertEqual(obtained, expected)
obtained = basket_out['groups'].keyset
expected = set()
self.assertEqual(obtained, expected)
# But two rules chained should get both nodes and groups...
queryb = orm.QueryBuilder()
queryb.append(orm.Node, tag='nodes_in_set')
queryb.append(orm.Group, with_node='nodes_in_set', filters={'type_string': 'user'})
rule1 = UpdateRule(queryb)
queryb = orm.QueryBuilder()
queryb.append(orm.Group, tag='groups_in_set')
queryb.append(orm.Node, with_group='groups_in_set')
rule2 = UpdateRule(queryb)
ruleseq = RuleSequence((rule1, rule2), max_iterations=np.inf)
# ...both starting with a node
initial_node = [node2.id]
basket_inp = Basket(nodes=initial_node)
basket_out = ruleseq.run(basket_inp.copy())
obtained = basket_out['nodes'].keyset
expected = set([node2.id, node3.id])
self.assertEqual(obtained, expected)
obtained = basket_out['groups'].keyset
expected = set([group2.id])
self.assertEqual(obtained, expected)
# ...and starting with a group
initial_group = [group3.id]
basket_inp = Basket(groups=initial_group)
basket_out = ruleseq.run(basket_inp.copy())
obtained = basket_out['nodes'].keyset
expected = set([node4.id])
self.assertEqual(obtained, expected)
obtained = basket_out['groups'].keyset
expected = set([group3.id, group4.id])
self.assertEqual(obtained, expected)
# Testing a "group chain"
total_groups = 10
groups = []
for idx in range(total_groups):
new_group = orm.Group(label='group-{}'.format(idx)).store()
groups.append(new_group)
nodes = []
edges = set()
for idx in range(1, total_groups):
new_node = orm.Data().store()
groups[idx].add_nodes(new_node)
groups[idx - 1].add_nodes(new_node)
nodes.append(new_node)
edges.add(GroupNodeEdge(node_id=new_node.id, group_id=groups[idx].id))
edges.add(GroupNodeEdge(node_id=new_node.id, group_id=groups[idx - 1].id))
qb1 = orm.QueryBuilder()
qb1.append(orm.Node, tag='nodes_in_set')
qb1.append(orm.Group, with_node='nodes_in_set', filters={'type_string': 'user'})
rule1 = UpdateRule(qb1, track_edges=True)
qb2 = orm.QueryBuilder()
qb2.append(orm.Group, tag='groups_in_set')
qb2.append(orm.Node, with_group='groups_in_set')
rule2 = UpdateRule(qb2, track_edges=True)
ruleseq = RuleSequence((rule1, rule2), max_iterations=np.inf)
initial_node = [nodes[-1].id]
basket_inp = Basket(nodes=initial_node)
basket_out = ruleseq.run(basket_inp.copy())
obtained = basket_out['nodes'].keyset
expected = set(n.id for n in nodes)
self.assertEqual(obtained, expected)
obtained = basket_out['groups'].keyset
expected = set(g.id for g in groups)
self.assertEqual(obtained, expected)
# testing the edges between groups and nodes:
result = basket_out['groups_nodes'].keyset
self.assertEqual(result, edges)
def test_returns_calls(self):
"""Tests return calls (?)"""
create_reversed = False
return_reversed = False
rules = []
# linking all processes to input data:
queryb = orm.QueryBuilder()
queryb.append(orm.Data, tag='predecessor')
queryb.append(
orm.ProcessNode,
with_incoming='predecessor',
edge_filters={'type': {
'in': [LinkType.INPUT_CALC.value, LinkType.INPUT_WORK.value]
}}
)
rules.append(UpdateRule(queryb))
# CREATE/RETURN(ProcessNode, Data) - Forward
queryb = orm.QueryBuilder()
queryb.append(orm.ProcessNode, tag='predecessor')
queryb.append(
orm.Data,
with_incoming='predecessor',
edge_filters={'type': {
'in': [LinkType.CREATE.value, LinkType.RETURN.value]
}}
)
rules.append(UpdateRule(queryb))
# CALL(ProcessNode, ProcessNode) - Forward
queryb = orm.QueryBuilder()
queryb.append(orm.ProcessNode, tag='predecessor')
queryb.append(
orm.ProcessNode,
with_incoming='predecessor',
edge_filters={'type': {
'in': [LinkType.CALL_CALC.value, LinkType.CALL_WORK.value]
}}
)
rules.append(UpdateRule(queryb))
# CREATE(ProcessNode, Data) - Reversed
if create_reversed:
queryb = orm.QueryBuilder()
queryb.append(orm.ProcessNode, tag='predecessor', project=['id'])
queryb.append(orm.Data, with_incoming='predecessor', edge_filters={'type': {'in': [LinkType.CREATE.value]}})
rules.append(UpdateRule(queryb))
# Case 3:
# RETURN(ProcessNode, Data) - Reversed
if return_reversed:
queryb = orm.QueryBuilder()
queryb.append(orm.ProcessNode, tag='predecessor')
queryb.append(orm.Data, output_of='predecessor', edge_filters={'type': {'in': [LinkType.RETURN.value]}})
rules.append(UpdateRule(queryb))
# Test was doing the calculation but not checking the results. Will have to think
# how to do that now. Temporal replacement:
new_node = orm.Data().store()
basket = Basket(nodes=(new_node.id,))
ruleseq = RuleSequence(rules, max_iterations=np.inf)
resulting_set = ruleseq.run(basket.copy())
expecting_set = resulting_set
self.assertEqual(expecting_set, resulting_set)
def traverse_graph(starting_pks, max_iterations=None, get_links=False, links_forward=(), links_backward=()):
"""
This function will return the set of all nodes that can be connected
to a list of initial nodes through any sequence of specified links.
Optionally, it may also return the links that connect these nodes.
:type starting_pks: list or tuple or set
:param starting_pks: Contains the (valid) pks of the starting nodes.
:type max_iterations: int or None
:param max_iterations:
The number of iterations to apply the set of rules (a value of 'None' will
iterate until no new nodes are added).
:param bool get_links:
Pass True to also return the links between all nodes (found + initial).
:type links_forward: aiida.common.links.LinkType
:param links_forward:
List with all the links that should be traversed in the forward direction.
:type links_backward: aiida.common.links.LinkType
:param links_backward:
List with all the links that should be traversed in the backward direction.
"""
# pylint: disable=too-many-locals,too-many-statements,too-many-branches
from aiida import orm
from aiida.tools.graph.age_entities import Basket
from aiida.tools.graph.age_rules import UpdateRule, RuleSequence, RuleSaveWalkers, RuleSetWalkers
from aiida.common import exceptions
if max_iterations is None:
max_iterations = inf
elif not (isinstance(max_iterations, int) or max_iterations is inf):
raise TypeError('Max_iterations has to be an integer or infinity')
linktype_list = []
for linktype in links_forward:
if not isinstance(linktype, LinkType):
raise TypeError('links_forward should contain links, but one of them is: {}'.format(type(linktype)))
linktype_list.append(linktype.value)
filters_forwards = {'type': {'in': linktype_list}}
linktype_list = []
for linktype in links_backward:
if not isinstance(linktype, LinkType):
raise TypeError('links_backward should contain links, but one of them is: {}'.format(type(linktype)))
linktype_list.append(linktype.value)
filters_backwards = {'type': {'in': linktype_list}}
if not isinstance(starting_pks, (list, set, tuple)):
raise TypeError('starting_pks must be of type list, set or tuple\ninstead, it is {}'.format(type(starting_pks)))
if not starting_pks:
if get_links:
output = {'nodes': set(), 'links': set()}
else:
output = {'nodes': set(), 'links': None}
return output
if any([not isinstance(pk, int) for pk in starting_pks]):
raise TypeError('one of the starting_pks is not of type int:\n {}'.format(starting_pks))
operational_set = set(starting_pks)
query_nodes = orm.QueryBuilder()
query_nodes.append(orm.Node, project=['id'], filters={'id': {'in': operational_set}})
existing_pks = set(query_nodes.all(flat=True))
missing_pks = operational_set.difference(existing_pks)
if missing_pks:
raise exceptions.NotExistent(
'The following pks are not in the database and must be pruned before this call: {}'.format(missing_pks)
)
rules = []
basket = Basket(nodes=operational_set)
# When max_iterations is finite, the order of traversal may affect the result
# (its not the same to first go backwards and then forwards than vice-versa)
# In order to make it order-independent, the result of the first operation needs
# to be stashed and the second operation must be performed only on the nodes
# that were already in the set at the begining of the iteration: this way, both
# rules are applied on the same set of nodes and the order doesn't matter.
# The way to do this is saving and seting the walkers at the right moments only
# when both forwards and backwards rules are present.
if links_forward and links_backward:
stash = basket.get_template()
rules += [RuleSaveWalkers(stash)]
if links_forward:
query_outgoing = orm.QueryBuilder()
query_outgoing.append(orm.Node, tag='sources')
query_outgoing.append(orm.Node, edge_filters=filters_forwards, with_incoming='sources')
rule_outgoing = UpdateRule(query_outgoing, max_iterations=1, track_edges=get_links)
rules += [rule_outgoing]
if links_forward and links_backward:
rules += [RuleSetWalkers(stash)]
if links_backward:
query_incoming = orm.QueryBuilder()
query_incoming.append(orm.Node, tag='sources')
query_incoming.append(orm.Node, edge_filters=filters_backwards, with_outgoing='sources')
rule_incoming = UpdateRule(query_incoming, max_iterations=1, track_edges=get_links)
rules += [rule_incoming]
rulesequence = RuleSequence(rules, max_iterations=max_iterations)
results = rulesequence.run(basket)
output = {}
output['nodes'] = results.nodes.keyset
output['links'] = None
if get_links:
output['links'] = results['nodes_nodes'].keyset
return output
def test_stash(self):
"""Testing sequencies and 'stashing'
Testing the dependency on the order of the operations in RuleSequence and the
'stash' functionality. This will be performed in a graph that has a calculation
(calc_ca) with two input data nodes (data_i1 and data_i2) and two output data
nodes (data_o1 and data_o2), and another calculation (calc_cb) which takes both
one of the inputs and one of the outputs of the first one (data_i2 and data_o2)
as inputs to produce a final output (data_o3).
"""
nodes = self._create_branchy_graph()
basket = Basket(nodes=[nodes['data_1'].id])
queryb_inp = orm.QueryBuilder()
queryb_inp.append(orm.Node, tag='nodes_in_set')
queryb_inp.append(orm.Node, with_outgoing='nodes_in_set')
uprule_inp = UpdateRule(queryb_inp)
queryb_out = orm.QueryBuilder()
queryb_out.append(orm.Node, tag='nodes_in_set')
queryb_out.append(orm.Node, with_incoming='nodes_in_set')
uprule_out = UpdateRule(queryb_out)
expect_base = set([nodes['calc_1'].id, nodes['data_1'].id, nodes['calc_2'].id])
# First get outputs, then inputs.
rule_seq = RuleSequence((uprule_out, uprule_inp))
obtained = rule_seq.run(basket.copy())['nodes'].keyset
expected = expect_base.union(set([nodes['data_i'].id]))
self.assertEqual(obtained, expected)
# First get inputs, then outputs.
rule_seq = RuleSequence((uprule_inp, uprule_out))
obtained = rule_seq.run(basket.copy())['nodes'].keyset
expected = expect_base.union(set([nodes['data_o'].id]))
self.assertEqual(obtained, expected)
# Now using the stash option in either order.
stash = basket.get_template()
rule_save = RuleSaveWalkers(stash)
rule_load = RuleSetWalkers(stash)
# Checking whether Rule does the right thing
# (i.e. If I stash the result, the operational sets should be the original,
# set, whereas the stash contains the same data as the starting point)
obtained = rule_save.run(basket.copy())
expected = basket.copy()
self.assertEqual(obtained, expected)
self.assertEqual(stash, basket)
stash = basket.get_template()
rule_save = RuleSaveWalkers(stash)
rule_load = RuleSetWalkers(stash)
serule_io = RuleSequence((rule_save, uprule_inp, rule_load, uprule_out))
result_io = serule_io.run(basket.copy())['nodes'].keyset
self.assertEqual(result_io, expect_base)
stash = basket.get_template()
rule_save = RuleSaveWalkers(stash)
rule_load = RuleSetWalkers(stash)
serule_oi = RuleSequence((rule_save, uprule_out, rule_load, uprule_inp))
result_oi = serule_oi.run(basket.copy())['nodes'].keyset
self.assertEqual(result_oi, expect_base)
def traverse_graph(
starting_pks: Iterable[int],
max_iterations: Optional[int] = None,
get_links: bool = False,
links_forward: Iterable[LinkType] = (),
links_backward: Iterable[LinkType] = (),
missing_callback: Optional[Callable[[Iterable[int]], None]] = None
) -> TraverseGraphOutput:
"""
This function will return the set of all nodes that can be connected
to a list of initial nodes through any sequence of specified links.
Optionally, it may also return the links that connect these nodes.
:param starting_pks: Contains the (valid) pks of the starting nodes.
:param max_iterations:
The number of iterations to apply the set of rules (a value of 'None' will
iterate until no new nodes are added).
:param get_links: Pass True to also return the links between all nodes (found + initial).
:param links_forward: List with all the links that should be traversed in the forward direction.
:param links_backward: List with all the links that should be traversed in the backward direction.
:param missing_callback: A callback to handle missing starting_pks or if None raise NotExistent
"""
# pylint: disable=too-many-locals,too-many-statements,too-many-branches
if max_iterations is None:
max_iterations = cast(int, inf)
elif not (isinstance(max_iterations, int) or max_iterations is inf):
raise TypeError('Max_iterations has to be an integer or infinity')
linktype_list = []
for linktype in links_forward:
if not isinstance(linktype, LinkType):
raise TypeError(
f'links_forward should contain links, but one of them is: {type(linktype)}'
)
linktype_list.append(linktype.value)
filters_forwards = {'type': {'in': linktype_list}}
linktype_list = []
for linktype in links_backward:
if not isinstance(linktype, LinkType):
raise TypeError(
f'links_backward should contain links, but one of them is: {type(linktype)}'
)
linktype_list.append(linktype.value)
filters_backwards = {'type': {'in': linktype_list}}
if not isinstance(starting_pks, Iterable): # pylint: disable=isinstance-second-argument-not-valid-type
raise TypeError(
f'starting_pks must be an iterable\ninstead, it is {type(starting_pks)}'
)
if any([not isinstance(pk, int) for pk in starting_pks]):
raise TypeError(
f'one of the starting_pks is not of type int:\n {starting_pks}')
operational_set = set(starting_pks)
if not operational_set:
if get_links:
return {'nodes': set(), 'links': set()}
return {'nodes': set(), 'links': None}
query_nodes = orm.QueryBuilder()
query_nodes.append(orm.Node,
project=['id'],
filters={'id': {
'in': operational_set
}})
existing_pks = set(query_nodes.all(flat=True))
missing_pks = operational_set.difference(existing_pks)
if missing_pks and missing_callback is None:
raise exceptions.NotExistent(
f'The following pks are not in the database and must be pruned before this call: {missing_pks}'
)
elif missing_pks and missing_callback is not None:
missing_callback(missing_pks)
rules = []
basket = Basket(nodes=existing_pks)
# When max_iterations is finite, the order of traversal may affect the result
# (its not the same to first go backwards and then forwards than vice-versa)
# In order to make it order-independent, the result of the first operation needs
# to be stashed and the second operation must be performed only on the nodes
# that were already in the set at the begining of the iteration: this way, both
# rules are applied on the same set of nodes and the order doesn't matter.
# The way to do this is saving and seting the walkers at the right moments only
# when both forwards and backwards rules are present.
if links_forward and links_backward:
stash = basket.get_template()
rules += [RuleSaveWalkers(stash)]
if links_forward:
query_outgoing = orm.QueryBuilder()
query_outgoing.append(orm.Node, tag='sources')
query_outgoing.append(orm.Node,
edge_filters=filters_forwards,
with_incoming='sources')
rule_outgoing = UpdateRule(query_outgoing,
max_iterations=1,
track_edges=get_links)
rules += [rule_outgoing]
if links_forward and links_backward:
rules += [RuleSetWalkers(stash)]
if links_backward:
query_incoming = orm.QueryBuilder()
query_incoming.append(orm.Node, tag='sources')
query_incoming.append(orm.Node,
edge_filters=filters_backwards,
with_outgoing='sources')
rule_incoming = UpdateRule(query_incoming,
max_iterations=1,
track_edges=get_links)
rules += [rule_incoming]
rulesequence = RuleSequence(rules, max_iterations=max_iterations)
results = rulesequence.run(basket)
output = {}
output['nodes'] = results.nodes.keyset
output['links'] = None
if get_links:
output['links'] = results['nodes_nodes'].keyset
return cast(TraverseGraphOutput, output)