def _query_kdtree(self, node: cKDTreeNode,
bb: Tuple[np.ndarray, np.ndarray]) -> List[np.ndarray]:
def substitute_dim(bound: np.ndarray, sub_dim: int, sub: float):
# replace bound[sub_dim] with sub
return np.array(
[bound[i] if i != sub_dim else sub for i in range(3)])
if node is None:
return []
if node.split_dim != -1:
# if the split location is above the roi, no need to split (cannot be above None)
if bb[1][node.split_dim] is not None and node.split > bb[1][
node.split_dim]:
return self._query_kdtree(node.lesser, bb)
elif (bb[0][node.split_dim] is not None
and node.split < bb[0][node.split_dim]):
return self._query_kdtree(node.greater, bb)
else:
return self._query_kdtree(
node.greater, bb) + self._query_kdtree(node.lesser, bb)
else:
# handle leaf node
bbox = Roi(
Coordinate(bb[0]),
Coordinate(
tuple(y - x if x is not None and y is not None else y
for x, y in zip(*bb))),
)
points = [
ind for ind, point in zip(node.indices, node.data_points)
if bbox.contains(np.round(point))
]
return points
def _query_kdtree(
self, node: cKDTreeNode, bb: Tuple[np.ndarray, np.ndarray]
) -> List[np.ndarray]:
def substitute_dim(bound: np.ndarray, sub_dim: int, sub: float):
# replace bound[sub_dim] with sub
return np.array([bound[i] if i != sub_dim else sub for i in range(3)])
if node is None:
return []
if node.split_dim != -1:
# recursive handling of child nodes
greater_roi = (substitute_dim(bb[0], node.split_dim, node.split), bb[1])
lesser_roi = (bb[0], substitute_dim(bb[1], node.split_dim, node.split))
return self._query_kdtree(node.greater, greater_roi) + self._query_kdtree(
node.lesser, lesser_roi
)
else:
# handle leaf node
# TODO: handle bounding box properly. bb[0], and bb[1] may not be integers.
bbox = Roi(Coordinate(bb[0]), Coordinate(bb[1] - bb[0]))
points = [point for point in node.data_points if bbox.contains(point)]
return points
def _grow_tree(self, root: np.ndarray, roi: Roi) -> nx.DiGraph:
tree = nx.DiGraph()
tree.add_node(0, pos=root)
leaves = deque([0])
next_id = 1
while len(leaves) > 0 and len(tree.nodes) < self.num_nodes[1]:
current_id = leaves.pop()
current_loc = tree.nodes[current_id]["pos"]
tail = [current_loc]
while len(tail) < self.ma and len(list(tree.predecessors(current_id))) == 1:
preds = list(tree.predecessors(current_id))
tail.append(tree.nodes[preds[0]]["pos"])
if len(tail) == 1:
tail.append(np.random.random(tail[0].shape)-0.5 + tail[0])
moving_direction = np.array(tail[0] - tail[-1])
if len(tree.nodes) >= self.num_nodes[0] and random.random() < self.p_die:
continue
else:
num_branches = np.random.choice(
range(1, self.max_split + 1), p=self.split_ps
)
next_points = self._gen_n(num_branches, current_loc, moving_direction)
for point in next_points:
if not roi.contains(point):
continue
tree.add_node(next_id, pos=point)
tree.add_edge(current_id, next_id)
leaves.appendleft(next_id)
next_id += 1
return tree