def _update_event_counts(self):
root = self.query_lifted
self.taxonomy.n_samples += 1
for n in root.H:
t = findall(self.taxonomy, filter_=lambda t: t.name == n)
t[0].n_gains += 1
for n in root.L:
t = findall(self.taxonomy, filter_=lambda t: t.name == n)
t[0].n_losses += 1
def NavigateState(self, graph_root: AnyNode, node: AnyNode):
"""
This function sets the state of a node depending on its (position in the) corresponding tree (-> DAG review as Tree)
:param graph_root:AnyNode: tree root
:param node:AnyNode: node from tree you want to update
"""
try:
if isNotNone(node.label) and isNone(node.content):
label = node.label
regex = str('\\b' + label + '\\b')
desired = []
tmp_desired = findall(graph_root,
lambda node: node.label in label)
for i in tmp_desired:
match = re.findall(regex, i.label)
if len(match) > 0:
desired.append(i)
if (len(desired) < 1): print(node.state)
elif (len(desired) == 1): self.NormalState(node)
else:
node.followerNodes = desired[0].followerNodes
node.hasFollowerNodes = desired[0].hasFollowerNodes
node.hasInputNode = desired[0].hasInputNode
node.state = 'navigator'
node.name = 'navigator_' + str(node.id)
else:
self.NormalState(node)
except Exception as ex:
template = "An exception of type {0} occurred in [TParser.NavigateState]. Arguments:\n{1!r}"
message = template.format(type(ex).__name__, ex.args)
print(message)
def st_make_halo(stree):
"""
Add NodeHalos to a ScheduleTree. A NodeHalo captures the halo exchanges
that should take place before executing the sub-tree; these are described
by means of a HaloScheme.
"""
# Build a HaloScheme for each expression bundle
halo_schemes = {}
for n in findall(stree, lambda i: i.is_Exprs):
try:
halo_schemes[n] = HaloScheme(n.exprs, n.ispace)
except HaloSchemeException as e:
if configuration['mpi']:
raise RuntimeError(str(e))
# Insert the HaloScheme at a suitable level in the ScheduleTree
mapper = {}
for k, hs in halo_schemes.items():
for f, v in hs.fmapper.items():
spot = k
ancestors = [n for n in k.ancestors if n.is_Iteration]
for n in ancestors:
test0 = any(n.dim is i.dim for i in v.halos)
test1 = n.dim not in [i.root for i in v.loc_indices]
if test0 or test1:
spot = n
break
mapper.setdefault(spot, []).append((f, v))
for spot, entries in mapper.items():
insert(NodeHalo(HaloScheme(fmapper=dict(entries))), spot.parent, [spot])
return stree
def _hover(self, event):
if event.inaxes == self._subplot:
pos = [event.x, event.y]
if self.has_cached_hovered_wedge() and self.latest_hovered_wedge[
'use_cached'] and self.latest_hovered_wedge[
'wedge'].contains_point(pos):
self._update_wedge_annotation(pos)
else:
for series_index, wedges in enumerate(self._wedge_series):
for wedge_index, w in enumerate(wedges):
corresponding_data = self._chart_data[-series_index -
1][wedge_index]
if not corresponding_data.is_filler and w.contains_point(
pos):
node = findall(self.current_plotted_root,
filter_=lambda n: n.id ==
corresponding_data.node_id)
self.latest_hovered_wedge['wedge'] = w
self.latest_hovered_wedge[
'data'] = corresponding_data
self.latest_hovered_wedge['use_cached'] = len(
node[0].parent.children) > 1
self._update_wedge_annotation(pos)
return
else:
self._wedge_annotation.set_visible(False)
self._blit()
def st_make_halo(stree):
"""
Add :class:`NodeHalo`s to a :class:`ScheduleTree`. A HaloNode captures
the halo exchanges that should take place before executing the sub-tree;
these are described by means of a :class:`HaloScheme`.
"""
# Build a HaloScheme for each expression bundle
halo_schemes = {}
for n in findall(stree, lambda i: i.is_Exprs):
try:
halo_schemes[n] = HaloScheme(n.exprs, n.ispace, n.dspace)
except HaloSchemeException as e:
if configuration['mpi']:
raise RuntimeError(str(e))
# Insert the HaloScheme at a suitable level in the ScheduleTree
mapper = {}
for k, hs in halo_schemes.items():
for f, v in hs.fmapper.items():
spot = k
ancestors = [n for n in k.ancestors if n.is_Iteration]
for n in ancestors:
test0 = any(n.dim is i.dim for i in v.halos)
test1 = n.dim not in [i.root for i in v.loc_indices]
if test0 or test1:
spot = n
break
mapper.setdefault(spot, []).append((f, v))
for spot, entries in mapper.items():
insert(NodeHalo(HaloScheme(fmapper=dict(entries))), spot.parent,
[spot])
return stree
def set_evaluator_value(self):
current_state, current_player = self.get_tree_height(), self.first_player
while current_state >= 0:
for node in findall(self.tree[0], filter_=lambda n: n.depth == current_state):
node.evaluator_value = self.calculate_evaluator_value(node, current_player, current_state)
current_player = not current_player
current_state -= 1
def on_file_open_clicked(self):
file_path, _ = QFileDialog.getOpenFileName(self)
if os.path.isfile(file_path):
# Sets windows title based on open file.
self.setWindowTitle(file_path + " - " + self.window_title)
# Initializes game session from file.
game_session = GameSession.from_file(file_path)
# Builds the tree based on the game session object hierarchy.
tree = proto_tree.construct_protocol_tree(game_session)
# Creates network event list model and sets to view.
network_event_list_model = NetworkEventListModel()
network_event_list_model.set_root(tree)
self.network_event_list.setModel(network_event_list_model)
message_list_model = MessageListModel()
message_list = anytree.findall(tree, filter_=lambda node: node.class_name == 'Message')
message_list_model.set_message_list(message_list)
self.message_list.setModel(message_list_model)
# Creates object tree model and sets to view.
object_tree_model = ObjectTreeModel()
object_tree_model.set_root(tree)
self.object_tree.setModel(object_tree_model)
# Selects first element of network event list.
event_list_first_item = self.network_event_list.model().index(0, 0, QModelIndex())
self.network_event_list.setCurrentIndex(event_list_first_item)
self.on_network_event_list_activated(event_list_first_item)
def search_models(
self,
desired_ids: Optional[List[Union[ModelGroupsIdsEnum,
ModelTypesIdsEnum]]] = None,
desired_metainfo: Optional[ModelMetaInfoTemplate] = None
) -> Tuple[List[ModelTypesIdsEnum], List[ModelMetaInfo]]:
desired_ids = [
ModelGroupsIdsEnum.all
] if desired_ids is None or not desired_ids else desired_ids
results = findall(self._tree,
filter_=lambda node: isinstance(node, ModelType) and
self._is_in_path(node, desired_ids))
if desired_metainfo is not None:
results = [
result for result in results if isinstance(result, ModelType)
and desired_metainfo.is_suits_for_template(result.meta_info)
]
models_ids = [
result.name for result in results
if (result.name in self.model_types)
]
models_metainfo = [
result.meta_info for result in results
if (result.name in self.model_types)
]
return models_ids, models_metainfo
def list_downstream_populations(self, population: str) -> list or None:
"""For a given population find all dependencies
Parameters
----------
population : str
population name
Returns
-------
list or None
List of populations dependent on given population
Raises
------
AssertionError
If Population does not exist
"""
assert population in self.tree.keys(), f'population {population} does not exist; ' \
f'valid population names include: {self.tree.keys()}'
root = self.tree['root']
node = self.tree[population]
dependencies = [
x.name
for x in anytree.findall(root, filter_=lambda n: node in n.path)
]
return [p for p in dependencies if p != population]
def __call__(self, roots):
#aaa = [[child.num_visited for child in findall(root, lambda node: node.depth == self.numActionPlaned)] for root in roots]
grandchildren_visit = np.sum([[
child.num_visited for child in findall(
root, lambda node: node.depth == self.numActionPlaned)
] for root in roots],
axis=0)
#__import__('ipdb').set_trace()
#for root in roots:
# for comboIndex in range(self.numActionIndexCombos):
# currParent = root
# for actionOrder in range(self.numActionPlaned):
# child = currParent.children[self.actionIndexCombosForActionPlaned[comboIndex][actionOrder]
# currParent = child
# numVisits[comboIndex] + = child.num_visited
maxIndex = np.argwhere(
grandchildren_visit == np.max(grandchildren_visit)).flatten()
#print(maxIndex,grandchildren_visit)
selectedActionIndexCombos = np.random.choice(maxIndex)
action = [
self.actionSpace[actionIndex] for actionIndex in
self.actionIndexCombosForActionPlaned[selectedActionIndexCombos]
]
return action
def get_group_ip(group):
""" given a group, get the ip of the group switch """
root = get_node_tree()
group_node = findall(
root, lambda node: next((True for x in node.properties['handles']
if x == group), False))[0]
return group_node.id if group_node else None
def find_candidate_minima(node: Node):
"""find everyone above this node that we've already visited.
Previously visited nodes have a cumulative cost."""
touched = anytree.findall(node.root,
filter_=lambda x: x.depth < node.depth and not x.
is_root and not visited(x))
# and visited(x)) # and hasattr(x, 'cumulative_cost'))
return touched
def ComputeItemPath(item1, item2, maxlength):
if item1 != item2:
str1path = findall(root, filter_=lambda node: node.name in (item1))
str2path = findall(root, filter_=lambda node: node.name in (item2))
str1path = str(str1path)
str2path = str(str2path)
str1lst = str1path.split('/')
str2lst = str2path.split('/')
for l in range(min(len(str1lst), len(str2lst))):
if str1lst[l] != str2lst[l]:
cmpindex = l
break
itempath = (len(str1lst) - cmpindex) + (len(str2lst) - cmpindex)
cost = round(itempath / maxlength, 3)
else:
cost = 0
return cost
def ComputeItemPathCost(matrix, i, j, item1, item2, maxlength):
if ((matrix[i - 1][j] + 1) > matrix[i - 1][j - 1]) and (
(matrix[i][j - 1] + 1) > matrix[i - 1][j - 1]):
str1char = findall(root, filter_=lambda node: node.name in (item1))
str2char = findall(root, filter_=lambda node: node.name in (item2))
str1char = str(str1char)
str2char = str(str2char)
str1lst = str1char.split('/')
str2lst = str2char.split('/')
for l in range(min(len(str1lst), len(str2lst))):
if str1lst[l] != str2lst[l]:
cmpindex = l
break
itempath = (len(str1lst) - cmpindex) + (len(str2lst) - cmpindex)
cost = round(itempath / maxlength, 3)
else:
cost = 1
return cost
def stree_section(stree):
"""
Add NodeSections to a ScheduleTree. A NodeSection, or simply "section",
defines a sub-tree with the following properties:
* The root is a node of type NodeSection;
* The immediate children of the root are nodes of type NodeIteration;
* The Dimensions of the immediate children are either:
* identical, OR
* different, but all of type SubDimension;
* The Dimension of the immediate children cannot be a TimeDimension.
"""
class Section(object):
def __init__(self, node):
self.parent = node.parent
try:
self.dim = node.dim
except AttributeError:
self.dim = None
self.nodes = [node]
def is_compatible(self, node):
return self.parent == node.parent and self.dim.root == node.dim.root
# Search candidate sections
sections = []
for i in range(stree.height):
# Find all sections at depth `i`
section = None
for n in findall(stree, filter_=lambda n: n.depth == i):
if any(p in flatten(s.nodes for s in sections)
for p in n.ancestors):
# Already within a section
continue
elif n.is_Sync:
# SyncNodes are self-contained
sections.append(Section(n))
section = None
elif n.is_Iteration:
if n.dim.is_Time and SEQUENTIAL in n.properties:
# If n.dim.is_Time, we end up here in 99.9% of the cases.
# Sometimes, however, time is a PARALLEL Dimension (e.g.,
# think of `norm` Operators)
section = None
elif section is None or not section.is_compatible(n):
section = Section(n)
sections.append(section)
else:
section.nodes.append(n)
else:
section = None
# Transform the schedule tree by adding in sections
for i in sections:
insert(NodeSection(), i.parent, i.nodes)
return stree
def section(stree):
"""
Create sections in a :class:`ScheduleTree`. A section is a sub-tree with
the following properties: ::
* The root is a node of type :class:`NodeSection`;
* The immediate children of the root are nodes of type :class:`NodeIteration`
and have same parent.
* The :class:`Dimension` of the immediate children are either: ::
* identical, OR
* different, but all of type :class:`SubDimension`;
* The :class:`Dimension` of the immediate children cannot be a
:class:`TimeDimension`.
"""
class Section(object):
def __init__(self, node):
self.parent = node.parent
self.dim = node.dim
self.nodes = [node]
def is_compatible(self, node):
return (self.parent == node.parent
and (self.dim == node.dim or node.dim.is_Sub))
# Search candidate sections
sections = []
for i in range(stree.height):
# Find all sections at depth `i`
section = None
for n in findall(stree, filter_=lambda n: n.depth == i):
if any(p in flatten(s.nodes for s in sections)
for p in n.ancestors):
# Already within a section
continue
elif not n.is_Iteration or n.dim.is_Time:
section = None
elif section is None or not section.is_compatible(n):
section = Section(n)
sections.append(section)
else:
section.nodes.append(n)
# Transform the schedule tree by adding in sections
for i in sections:
node = NodeSection()
processed = []
for n in list(i.parent.children):
if n in i.nodes:
n.parent = node
if node not in processed:
processed.append(node)
else:
processed.append(n)
i.parent.children = processed
return stree
def get_ip_by_handle(handle):
""" given a handle, get the ip of the handle switch """
if "." in handle: # check if handle is IP address
return handle
root = get_node_tree()
handle_node = findall(
root, lambda node: next(
(True for x in node.properties['handles'] if x == handle), False))
return handle_node[0].id if handle_node else None
def walk(self, node: Node):
self.iter_count += 1
print(f'current node: {node.name}, depth: {node.depth}')
compute_cumulative_cost(node)
visit(node)
if not node.is_leaf:
print('\tThis is not a leaf; continuing down')
self.walk(find_min(node.children))
else:
print('\tThis is a leaf; checking incumbency.')
if self.incumbent is None or self.incumbent.cumulative_cost > node.cumulative_cost:
print(
f'\t\tThere is no incumbent yet, or this leaf has a lower cost. Promoting {node.name} to incumbent.'
)
self.incumbent = node
to_prune = anytree.findall(
node.root,
filter_=lambda x: not x.is_root and x.depth < node.depth
and has_cumulative_cost(
x) and x.cumulative_cost > node.cumulative_cost)
for x in to_prune:
x.parent = None
pruned = len(to_prune)
print(f'\t\tPruned {pruned} nodes')
self.pruned_count += pruned
self.history.append(node)
next_node = find_min(
[x for x in node.root.children if not visited(x)])
print(
f'\t\tIdentified new minimum node to visit: {next_node.name}, depth {next_node.depth}'
)
self.walk(next_node)
else:
candidate_minima = [
x for x in node.root.children
if not visited(x) and x not in node.ancestors
]
if candidate_minima:
next_node = find_min(candidate_minima)
print(
f"\t\tLeaf {node.name} has a larger cumulatuve cost than the current incumbent."
)
print(f'\t\tReturning to a higher node {next_node.name}')
self.walk(next_node)
else:
print(
f'\t\t\tTree is depleted. This leaf is not the new incumbent but no higher nodes have a lower cumulative cost'
)
print(
f'Pruning complete. {self.pruned_count} nodes were removed in {self.iter_count} passes'
)
retval = node.root
return retval
def get_contacts(root_node):
""" Get all the contacts from our tree """
log.debug(render_tree(root_node))
clean_tree(root_node)
roll_up(root_node)
log.debug(render_tree(root_node))
all_nodes = findall(
root_node,
filter_=lambda x: x.contact and x.contact.name and x.contact.position)
[n.contact.get_names() for n in all_nodes]
return set([n.contact for n in all_nodes])
def find_node_by_name(self, name: str) -> Optional[XsdNode]:
if name == self.expand_element.name:
return self.expand_element
result = findall(
self.root,
filter_=lambda node: node.name == name,
)
if len(result) > 0:
return cast(XsdNode, result[0])
return None
def verifymap(orbitmap):
"""Verifies an orbitmap, returns the checksum
>>> verifymap(readmap("COM)B\\nB)C\\nC)D\\nD)E\\nE)F\\nB)G\\nG)H\\nD)I\\nE)J\\nJ)K\\nK)L"))
42
"""
count = 0
orbiters = findall(orbitmap, filter_=lambda node: node.name != "COM")
for orbiter in orbiters:
count += orbiter.depth
return count
def ComputeDiagonalCost(matrix, i, j, str1, str2, root):
maxlength = SearchLongestPath(root)
if ((matrix[i - 1][j] + 1) > matrix[i - 1][j - 1]) and (
(matrix[i][j - 1] + 1) > matrix[i - 1][j - 1]):
str1char = findall(root,
filter_=lambda node: node.name in (str1[i - 1]))
str2char = findall(root,
filter_=lambda node: node.name in (str2[j - 1]))
str1char = str(str1char)
str2char = str(str2char)
str1lst = str1char.split('/')
str2lst = str2char.split('/')
for l in range(min(len(str1lst), len(str2lst))):
if str1lst[l] != str2lst[l]:
cmpindex = l
break
itempath = (len(str1lst) - cmpindex) + (len(str2lst) - cmpindex)
cost = round(itempath / maxlength, 3)
else:
cost = 1
return cost
def _create_tree(target, t_paths):
# This is to create a tree with the information of the dominant species
root = Node(target, order=0)
for idx, ds in enumerate(t_paths):
for pa, v in ds.items():
sps = np.concatenate(v)
for sp in sps:
p = findall(root,
filter_=lambda n: n.name == pa and n.order == idx)
for m in p:
Node(sp, parent=m, order=idx + 1)
return root
def find_name(self,
root,
key,
script=False,
directory=False,
startpath=None,
parentfromtree=False,
fmt='native'):
'''
find files based on their names
@script: output script
@directory: only search for directories
@startpath: node to start with
@parentfromtree: get path from parent instead of stored relpath
@fmt: output format
'''
self._debug('searching for \"{}\"'.format(key))
start = root
if startpath:
start = self.get_node(root, startpath)
self.term = key
found = anytree.findall(start, filter_=self._find_name)
paths = []
for f in found:
if f.type == self.TYPE_STORAGE:
# ignore storage nodes
continue
if directory and f.type != self.TYPE_DIR:
# ignore non directory
continue
# print the node
if fmt == 'native':
self._print_node(f,
withpath=True,
withdepth=True,
withstorage=True,
recalcparent=parentfromtree)
elif fmt == 'csv':
self._node_to_csv(f)
if parentfromtree:
paths.append(self._get_parents(f))
else:
paths.append(f.relpath)
if script:
tmp = ['${source}/' + x for x in paths]
cmd = 'op=file; source=/media/mnt; $op {}'.format(' '.join(tmp))
Logger.info(cmd)
return found
def extract_tetrads_to_csv(tree: Node, outfile: str):
tetrads = {}
leaves = anytree.findall(tree, filter_=lambda x: x.is_leaf)
for leaf in leaves:
tetrad = leaf.path[1:] # omit root
tetrads[tetrad] = (gain(tetrad), cumulative_cost(leaf))
df = pd.DataFrame({make_tetrad_str(k): v
for k, v in tetrads.items()}).transpose()
df.columns = ['gain', 'cumulative_cost']
df = df.sort_values('cumulative_cost')
df.to_csv(outfile)
print(f"wrote {os.path.abspath(outfile)}")
def main():
with open('input.txt', 'r') as f:
lines = f.read().splitlines()
root=find_root(lines)
tree=Node(root)
tree=add_children(lines,root,tree)
print_tree(tree)
count=0
for o in (findall(tree)):
count += o.depth
print(count)
def _make_matches_from_leaf_nodes(self, last_node_before_leaf, prefix):
nodes = anytree.findall(last_node_before_leaf,
lambda node: node.name.startswith(prefix),
maxlevel=2)
matches = []
for node in nodes:
if node != last_node_before_leaf:
this_match = ":".join(
[i.name for i in node.path if i.name != ''])
if len(node.children) != 0:
this_match = this_match + ':'
matches = matches + [this_match]
return matches
def __call__(self, roots):
grandchildrenVisit = np.sum([[
child.numVisited for child in findall(
root, lambda node: node.depth == self.numActionPlaned)
] for root in roots],
axis=0)
maxIndex = np.argwhere(
grandchildrenVisit == np.max(grandchildrenVisit)).flatten()
selectedActionIndexCombos = np.random.choice(maxIndex)
action = [
self.actionSpace[actionIndex] for actionIndex in
self.actionIndexCombosForActionPlaned[selectedActionIndexCombos]
]
return action
def test_enum():
class Animals(IntEnum):
Mammal = 1
Cat = 2
Dog = 3
root = Node("ANIMAL")
mammal = Node(Animals.Mammal, parent=root)
cat = Node(Animals.Cat, parent=mammal)
dog = Node(Animals.Dog, parent=mammal)
eq_(findall(root), (root, mammal, cat, dog))
eq_(findall_by_attr(root, Animals.Cat), (cat, ))
def add_gate(self, gate):
"""
Add a gate to the gating strategy, see `gates` module. The gate ID must be unique in the gating strategy.
:param gate: instance from a sub-class of the Gate class
:return: None
"""
if not isinstance(gate, Gate):
raise ValueError("gate must be a sub-class of the Gate class")
parent_id = gate.parent
if parent_id is None:
parent_id = 'root'
# TODO: check for uniqueness of gate ID + gate path combo
matched_nodes = anytree.findall(
self._gate_tree,
filter_=lambda g_node: g_node.name == gate.id and g_node.parent.
name == parent_id)
if len(matched_nodes) != 0:
raise KeyError("Gate ID '%s' is already defined" % gate.id)
parent_id = gate.parent
if parent_id is None:
parent_node = self._gate_tree
else:
matching_nodes = anytree.search.findall_by_attr(
self._gate_tree, parent_id)
if len(matching_nodes) == 0:
# TODO: could be in a quadrant gate
raise ValueError(
"Parent gate %s does not exist in the gating strategy" %
parent_id)
elif len(matching_nodes) > 1:
raise ValueError(
"Multiple gates exist matching parent ID %s, specify full gate path"
% parent_id)
parent_node = matching_nodes[0]
node = anytree.Node(gate.id, parent=parent_node, gate=gate)
# Quadrant gates need special handling to add their individual quadrants as children.
# Other gates cannot have the main quadrant gate as a parent, they can only reference
# the individual quadrants as parents.
if isinstance(gate, fk_gates.QuadrantGate):
for q_id, q in gate.quadrants.items():
anytree.Node(q_id, parent=node, gate=q)
def main():
# part 1
print("# part1:")
f = open("input7.txt", "r")
lines = f.readlines()
f.close()
root = Node("root")
# generate first level
for n_line, line in enumerate(lines):
splited_line = line.replace("\n","").split(" contain ")
Node(splited_line[0].strip().replace("bags", "bag").replace(".", ""), parent=root, bag_value=0)
# generate n-th levels
for pre, fill, node in RenderTree(root):
for n_line, line in enumerate(lines):
splited_line = line.replace("\n","").split(" contain ")
if(node.name in splited_line[0]):
desc = splited_line[1].split(",")
for d in desc:
d_value = ''.join([i for i in d if i.isdigit()]).strip()
#if(len(d_value) == 0): d_value = 0
d_name = ''.join([i for i in d if not i.isdigit()]).strip()
d_name = d_name.replace("bags", "bag").replace(".", "")
if(d_name not in "no other bag"):
Node(d_name, parent=node, bag_value=d_value)
#print the tree
# for pre, fill, node in RenderTree(root):
# print("%s%s" % (pre, node.name))
results = []
for pre, fill, node in RenderTree(root, maxlevel=2):
finds = findall(node, filter_=lambda n: "shiny gold" in n.name)
if(len(finds) > 0): results.append(finds)
print("total part1 :", len(results) - 2)
#part 2
r = Resolver('name')
shiny_gold_bag_node = r.get(root, "shiny gold bag")
shiny_gold_bag_node.bag_value = 1
# print the tree
# for pre, fill, node in RenderTree(shiny_gold_bag_node):
# print("%s%s(%s)" % (pre, node.name, node.bag_value))
print("total part2 :", count_bags(shiny_gold_bag_node))
def st_section(stree):
"""
Add NodeSections to a ScheduleTree. A NodeSection, or simply "section",
defines a sub-tree with the following properties:
* The root is a node of type NodeSection;
* The immediate children of the root are nodes of type NodeIteration;
* The Dimensions of the immediate children are either:
* identical, OR
* different, but all of type SubDimension;
* The Dimension of the immediate children cannot be a TimeDimension.
"""
class Section(object):
def __init__(self, node):
self.parent = node.parent
self.dim = node.dim
self.nodes = [node]
def is_compatible(self, node):
return self.parent == node.parent and self.dim.root == node.dim.root
# Search candidate sections
sections = []
for i in range(stree.height):
# Find all sections at depth `i`
section = None
for n in findall(stree, filter_=lambda n: n.depth == i):
if any(p in flatten(s.nodes for s in sections) for p in n.ancestors):
# Already within a section
continue
elif not n.is_Iteration or n.dim.is_Time:
section = None
elif section is None or not section.is_compatible(n):
section = Section(n)
sections.append(section)
else:
section.nodes.append(n)
# Transform the schedule tree by adding in sections
for i in sections:
insert(NodeSection(), i.parent, i.nodes)
return stree
