def balance_subroutine(tree, depth):
"""
Perform balancing to enforce the 2:1 constraint between neighbouring
nodes in linear octree.
Strategy: Traverse the octree level by level, bottom-up, and check if
a given node's parent lies in the tree, add it and their respective
siblings. This enforces the 2:1 constraint.
"""
balanced = set(tree)
for level in range(depth, 0, -1):
work_items = [x for x in balanced if morton.find_level(x) == level]
for work_item in work_items:
neighbours = morton.find_neighbours(work_item)
for neighbour in neighbours:
parent = morton.find_parent(neighbour)
parent_level = level - 1
if ~(neighbour in balanced) and ~(parent in balanced):
balanced.add(parent)
if parent_level > 0:
siblings = morton.find_siblings(parent)
balanced.update(siblings)
return numba.typed.List(balanced)
def find_dense_v_list(key, depth):
"""
Find the V list of a key if it were in a non-adaptive tree.
"""
parent = morton.find_parent(key)
parent_neighbours = morton.find_neighbours(parent)
parent_neigbhours_children = morton.find_children_vec(
parent_neighbours).ravel()
are_adj = morton.are_adjacent_vec(key, parent_neigbhours_children, depth)
return parent_neigbhours_children[are_adj == 0]
def test_find_interaction_lists(balanced):
"""
Currently only tests interaction lists for nodes at leaf level, mainly
checking that the constraints on their level, and adjacency are
satisfied.
"""
depth = tree.find_depth(balanced)
complete = tree.complete_tree(balanced)
u, x, v, w = tree.find_interaction_lists(balanced, complete, depth)
for i in range(len(complete)):
key = complete[i]
if key in balanced:
# Check all u list members are adjacent
u_i = u[i][u[i] != -1]
u_adj_idxs = morton.are_adjacent_vec(key, u_i, depth)
assert np.all(u_adj_idxs == 1)
# Check all x list members are at the right level, and not adjacent
x_i = x[i][x[i] != -1]
x_nadj_idxs = morton.are_adjacent_vec(key, x_i, depth)
assert np.all(x_nadj_idxs == 0)
x_levels = morton.find_level(x_i)
assert np.all(
x_levels == morton.find_level(morton.find_parent(key)))
# Check all w list members are at right level, and not adjacent
w_i = w[i][w[i] != -1]
w_nadj_idxs = morton.are_adjacent_vec(key, w_i, depth)
assert np.all(w_nadj_idxs == 0)
w_levels = morton.find_level(w_i)
assert np.all(w_levels == (morton.find_level(key) - 1))
# Check all v list members are at right level, and not adjacent
v_i = v[i][v[i] != -1]
v_nadj_idxs = morton.are_adjacent_vec(key, v_i, depth)
assert np.all(v_nadj_idxs == 0)
v_levels = morton.find_level(v_i)
assert np.all(v_levels == morton.find_level(key))
def find_interaction_list(i, leaves, leaves_set, complete_set, depth, u, x,
v, w):
"""
Internal method to find interaction list for a given key.
Parameters:
-----------
i : np.int64
Index of key in the complete tree.
leaves : np.array(dtype=np.int64)
Linear octree, represented by its leaves.
leaves_set : set(np.int64)
Set containing the linear octree.
complete_set : set(np.int64)
Set containing the completed octree.
depth : np.int64
Depth of the octree.
u : np.array(shape=(n_complete, 90))
U list container, nearest neighbours.
x : np.array(shape=(n_complete, 20))
X list container, colleagues of a node's parent which are
non-adjacent to the node - conjugate to the W list.
v : np.array(shape=(n_complete, 189))
V list container, children of a node's parent's colleagues which
are not adjacent to the node.
w : np.array(shape=(n_complete, 208))
W list container, children of colleagues, which are not adjacent to
the node.
"""
def build_parent_level_leaf(key, leaves_set, colleagues_parents,
parent_colleagues, depth):
"""
Build the portion of the interaction list that is expected at the
parent level of a node. Contributes to the U and X lists for
leaf keys.
Parameters:
----------
key : np.int64
Key for the current node being considered.
leaves_set : set(np.int64)
colleagues_parents : np.array(np.int64)
The (unique) parents of the node's colleagues.
parent_colleagues : np.array(np.int64)
The colleagues of the node's parents.
depth : np.int64
Returns:
--------
(np.int64, np.array(dtype=np.int64))
Four-tuple (u_ptr, u, x_ptr, x), where 'u_ptr' is the length
of the U list contributions at this level.
"""
# U List (P2P)
cp_in_tree = np.zeros_like(colleagues_parents)
i = 0
for cp in colleagues_parents:
if cp in leaves_set:
cp_in_tree[i] = cp
i += 1
cp_in_tree = cp_in_tree[:i]
adj_idxs = morton.are_adjacent_vec(key, cp_in_tree, depth)
adjacent = cp_in_tree[adj_idxs == 1]
# X List (P2L)
pc_in_tree = np.zeros_like(parent_colleagues)
i = 0
for pc in parent_colleagues:
if pc in leaves_set:
pc_in_tree[i] = pc
i += 1
pc_in_tree = pc_in_tree[:i]
not_adj_idxs = morton.are_adjacent_vec(key, pc_in_tree, depth)
not_adjacent = pc_in_tree[not_adj_idxs == 0]
return len(adjacent), adjacent, len(not_adjacent), not_adjacent
def build_current_level_leaf(key, leaves_set, complete_set, colleagues,
parent_colleagues_children, depth):
"""
Build the portion of the interaction list that is expected at the
level of a node. Contributes to the U and V lists for leaf keys.
Parameters:
-----------
key : np.int64
Key for the current node being considered.
leaves_set : set(np.int64)
colleagues : np.array(np.int64)
The node's colleagues.
parent_colleagues_children : np.array(np.int64)
The children of a node's parent's colleagues.
depth : np.int64
Returns:
--------
(np.int64, np.array(dtype=np.int64))
Four-tuple (u_ptr, u, v_ptr, v), where 'u_ptr' is the length
of the U list contributions at this level.
"""
# U List (P2P)
c_in_tree = np.zeros_like(colleagues)
i = 0
for c in colleagues:
if c in leaves_set:
c_in_tree[i] = c
i += 1
c_in_tree = c_in_tree[:i]
adj_idxs = morton.are_adjacent_vec(key, c_in_tree, depth)
adjacent = c_in_tree[adj_idxs == 1]
# V List (M2L)
pcc_in_tree = np.zeros_like(parent_colleagues_children)
i = 0
for pcc in parent_colleagues_children:
if pcc in complete_set:
pcc_in_tree[i] = pcc
i += 1
pcc_in_tree = pcc_in_tree[:i]
adj_idxs = morton.are_adjacent_vec(key, pcc_in_tree, depth)
not_adjacent = pcc_in_tree[adj_idxs == 0]
return len(adjacent), adjacent, len(not_adjacent), not_adjacent
def build_current_level_non_leaf(key, complete_set,
parent_colleagues_children, depth):
"""
Build the portion of the interaction list that is expected at the
level of a node. Contributes to the V lists for non-leaf keys.
Parameters:
-----------
key : np.int64
Key for the current node being considered.
complete_set : set(np.int64)
parent_colleagues_children : np.array(np.int64)
The children of a node's parent's colleagues.
depth : np.int64
Returns:
--------
(np.int64, np.array(dtype=np.int64))
Tuple (v_ptr, v), where 'v_ptr' is the length of the V list
contributions at this level.
"""
# V List (M2L)
pcc_in_tree = np.zeros_like(parent_colleagues_children)
i = 0
for pcc in parent_colleagues_children:
if pcc in complete_set:
pcc_in_tree[i] = pcc
i += 1
pcc_in_tree = pcc_in_tree[:i]
adj_idxs = morton.are_adjacent_vec(key, pcc_in_tree, depth)
not_adjacent = pcc_in_tree[adj_idxs == 0]
return len(not_adjacent), not_adjacent
def build_child_level_leaf(key, leaves_set, colleagues_children,
depth):
"""
Build the portion of the interaction list that is expected at the
level of a node's children. Contributes to the U and W lists for
leaf keys.
Parameters:
-----------
key : np.int64
Key for the current node being considered.
leaves_set : set(np.int64)
colleagues_children : np.array(np.int64)
The children of a node's colleagues.
depth : np.int64
Returns:
--------
(np.int64, np.array(dtype=np.int64))
Four-tuple (u_ptr, u, w_ptr, w), where 'u_ptr' is the length
of the U list contributions at this level.
"""
# U List (P2P)
i = 0
cc_in_tree = np.zeros_like(colleagues_children)
for cc in colleagues_children:
if cc in leaves_set:
cc_in_tree[i] = cc
i += 1
cc_in_tree = cc_in_tree[:i]
adj_idxs = morton.are_adjacent_vec(key, cc_in_tree, depth)
adjacent = cc_in_tree[adj_idxs == 1]
# W List (M2P)
not_adjacent = cc_in_tree[adj_idxs == 0]
return len(adjacent), adjacent, len(not_adjacent), not_adjacent
u_ptr = 0
x_ptr = 0
v_ptr = 0
w_ptr = 0
key = complete[i]
if key in leaves_set:
parent = morton.find_parent(key)
colleagues = morton.find_neighbours(key)
colleagues_children = morton.find_children_vec(colleagues).ravel()
colleagues_parents = np.unique(morton.find_parent(colleagues))
parent_colleagues = morton.find_neighbours(parent)
parent_colleagues_children = morton.find_children_vec(
parent_colleagues).ravel()
pu_ptr, padj, px_ptr, pnadj = build_parent_level_leaf(
key, leaves_set, colleagues_parents, parent_colleagues, depth)
u[i][u_ptr:pu_ptr] = padj
u_ptr = pu_ptr
x[i][x_ptr:px_ptr] = pnadj
cu_ptr, cadj, pcc_ptr, pcc_nadj = build_current_level_leaf(
key, leaves_set, complete_set, colleagues,
parent_colleagues_children, depth)
u[i][u_ptr:u_ptr + cu_ptr] = cadj
u_ptr = pu_ptr + cu_ptr
v[i][v_ptr:pcc_ptr] = pcc_nadj
ccu_ptr, ccadj, ccw_ptr, ccnadj = build_child_level_leaf(
key, leaves_set, colleagues_children, depth)
u[i][u_ptr:u_ptr + ccu_ptr] = ccadj
w[i][w_ptr:ccw_ptr] = ccnadj
else:
parent = morton.find_parent(key)
parent_colleagues = morton.find_neighbours(parent)
parent_colleagues_children = morton.find_children_vec(
parent_colleagues).ravel()
pcc_ptr, pcc_nadj = build_current_level_non_leaf(
key, complete_set, parent_colleagues_children, depth)
v[i][v_ptr:pcc_ptr] = pcc_nadj
def test_find_parent(key, expected):
result = morton.find_parent(key)
assert result == expected