Skip to content

minor fix + test (test_tree_of_boxes) #118

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
inducer merged 3 commits into from
Jul 7, 2026
Merged

minor fix + test (test_tree_of_boxes) #118

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4599b6d
minor-fix + test
ShawnL00 Jun 6, 2026
66b34cc
Minor TOB typing
inducer Jul 7, 2026
d03237a
Update baseline
inducer Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 98 additions & 10 deletions .basedpyright/baseline.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32786,34 +32786,34 @@
}
},
{
"code": "reportUnknownVariableType",
"code": "reportOperatorIssue",
"range": {
"startColumn": 4,
"endColumn": 24,
"startColumn": 14,
"endColumn": 53,
"lineCount": 1
}
},
{
"code": "reportUnknownArgumentType",
"range": {
"startColumn": 32,
"endColumn": 52,
"startColumn": 14,
"endColumn": 53,
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 4,
"endColumn": 25,
"endColumn": 24,
"lineCount": 1
}
},
{
"code": "reportAny",
"code": "reportUnknownArgumentType",
"range": {
"startColumn": 35,
"endColumn": 70,
"startColumn": 32,
"endColumn": 52,
"lineCount": 1
}
},
Expand Down Expand Up @@ -32905,6 +32905,70 @@
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 4,
"endColumn": 14,
"lineCount": 1
}
},
{
"code": "reportUnknownArgumentType",
"range": {
"startColumn": 31,
"endColumn": 39,
"lineCount": 1
}
},
{
"code": "reportUnknownArgumentType",
"range": {
"startColumn": 41,
"endColumn": 49,
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 4,
"endColumn": 17,
"lineCount": 1
}
},
{
"code": "reportOperatorIssue",
"range": {
"startColumn": 20,
"endColumn": 39,
"lineCount": 1
}
},
{
"code": "reportIndexIssue",
"range": {
"startColumn": 4,
"endColumn": 18,
"lineCount": 1
}
},
{
"code": "reportUnnecessaryComparison",
"range": {
"startColumn": 19,
"endColumn": 37,
"lineCount": 1
}
},
{
"code": "reportIndexIssue",
"range": {
"startColumn": 4,
"endColumn": 17,
"lineCount": 1
}
},
{
"code": "reportArgumentType",
"range": {
Expand All @@ -32925,7 +32989,7 @@
"code": "reportUnnecessaryComparison",
"range": {
"startColumn": 27,
"endColumn": 51,
"endColumn": 79,
"lineCount": 1
}
},
Expand Down Expand Up @@ -34802,6 +34866,30 @@
"endColumn": 25,
"lineCount": 1
}
},
{
"code": "reportUnnecessaryComparison",
"range": {
"startColumn": 15,
"endColumn": 28,
"lineCount": 1
}
},
{
"code": "reportUnnecessaryComparison",
"range": {
"startColumn": 20,
"endColumn": 39,
"lineCount": 1
}
},
{
"code": "reportUnnecessaryComparison",
"range": {
"startColumn": 15,
"endColumn": 28,
"lineCount": 1
}
}
]
}
Expand Down
39 changes: 25 additions & 14 deletions boxtree/tree_of_boxes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -174,10 +174,8 @@ def _apply_refine_flags_without_sorting(
tob.nboxes
+ nchildren * np.arange(len(refine_parents)))

refine_parents_per_child = np.empty(
(nchildren, len(refine_parents)), dtype=np.intp)
refine_parents_per_child[:] = refine_parents.reshape(-1)
refine_parents_per_child = refine_parents_per_child.reshape(-1)
# Assign the parent box number for each new child box.
refine_parents_per_child = np.repeat(refine_parents, nchildren)
Comment on lines -177 to +178

@alexfikl alexfikl Jun 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this correct? The before code looks more like an np.tile.

@alexfikl alexfikl Jun 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, I see, this was a bug fix! Fair enough.


box_parents = _resized_array(tob.box_parent_ids, nboxes_new)
box_centers = _resized_array(tob.box_centers, nboxes_new)
Expand Down Expand Up @@ -230,27 +228,30 @@ def _apply_coarsen_flags(
raise ValueError("attempting to coarsen non-leaf")

coarsen_sources, = np.where(coarsen_flags)
if np.any(tob.box_parent_ids[coarsen_sources] < 0):
raise ValueError("cannot coarsen the root box")
if coarsen_sources.size == 0:
return tob

coarsen_parents = tob.box_parent_ids[coarsen_sources]
coarsen_peers = tob.box_child_ids[:, coarsen_parents].reshape(-1)
coarsen_peers = tob.box_child_ids[:, coarsen_parents]
coarsen_peer_is_leaf = box_is_leaf[coarsen_peers]
coarsen_exec_flags = np.all(coarsen_peer_is_leaf, axis=0)

# when a leaf box marked for coarsening has non-leaf peers
coarsen_flags_ignored = np.sum(coarsen_exec_flags != coarsen_flags)
if coarsen_flags_ignored:
msg = (f"{coarsen_flags_ignored} out of {np.sum(coarsen_flags)} coarsening "
"flags ignored to prevent removing non-leaf boxes")
coarsen_flags_ignored = ~coarsen_exec_flags
if np.any(coarsen_flags_ignored):
msg = (f"{np.sum(coarsen_flags_ignored)} out of "
f"{np.sum(coarsen_flags)} coarsening flags ignored "
"to prevent removing non-leaf boxes")
if error_on_ignored_flags:
raise RuntimeError(msg)
else:
import warnings
warnings.warn(msg, stacklevel=3)

# deleted boxes are marked as:
# level = inf
# level = sentinel (max int)
# parent = -1
coarsen_parents = coarsen_parents[coarsen_exec_flags]
coarsen_peers = coarsen_peers[:, coarsen_exec_flags]
Expand All @@ -259,7 +260,7 @@ def _apply_coarsen_flags(
box_children = tob.box_child_ids.copy()
box_children[:, coarsen_parents] = 0
box_levels = tob.box_levels.copy()
box_levels[coarsen_peers] = np.inf
box_levels[coarsen_peers] = np.iinfo(box_levels.dtype).max
Comment thread
inducer marked this conversation as resolved.

return TreeOfBoxes(
box_centers=tob.box_centers,
Expand All @@ -285,12 +286,22 @@ def _sort_boxes_by_level(tob: TreeOfBoxes) -> TreeOfBoxes:
if not np.any(np.diff(tob.box_levels) < 0):
return tob

# reorder boxes to into non-decreasing levels
# reorder boxes into non-decreasing levels
neworder = np.argsort(tob.box_levels)
old_to_new = np.empty_like(neworder)
old_to_new[neworder] = np.arange(len(neworder))

box_centers = tob.box_centers[:, neworder]
box_levels = tob.box_levels[neworder]
box_parent_ids = tob.box_parent_ids[neworder]
box_child_ids = tob.box_child_ids[:, neworder]
box_levels = tob.box_levels[neworder]

# box_parent_ids and box_child_ids need to be updated to
# reflect the new box numbering
parent_has_id = box_parent_ids >= 0
box_parent_ids[parent_has_id] = old_to_new[box_parent_ids[parent_has_id]]
child_has_id = box_child_ids != 0
box_child_ids[child_has_id] = old_to_new[box_child_ids[child_has_id]]
Comment on lines +301 to +304

@alexfikl alexfikl Jun 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure it's clear to me why the above box_parent_ids = tob.box_parent_ids[neworder] isn't enough. Maybe add a little comment to this block to explain what it's trying to do?


return TreeOfBoxes(
box_centers=box_centers,
Expand All @@ -314,7 +325,7 @@ def _sort_boxes_by_level(tob: TreeOfBoxes) -> TreeOfBoxes:

def _sort_and_prune_deleted_boxes(tob: TreeOfBoxes) -> TreeOfBoxes:
tob = _sort_boxes_by_level(tob)
n_stale_boxes = np.sum(tob.box_levels == np.inf)
n_stale_boxes = np.sum(tob.box_levels == np.iinfo(tob.box_levels.dtype).max)
newn = tob.nboxes - n_stale_boxes

return TreeOfBoxes(
Expand Down
91 changes: 87 additions & 4 deletions test/test_tree_of_boxes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,11 @@
import numpy as np
import pytest

from arraycontext import pytest_generate_tests_for_array_contexts
from arraycontext import (
ArrayContext,
ArrayContextFactory,
pytest_generate_tests_for_array_contexts,
)

# This means boxtree's tests have a hard dependency on meshmode. That's OK.
from meshmode import _acf # noqa: F401
Expand All @@ -40,8 +44,11 @@
from pytools import obj_array

from boxtree import (
TreeOfBoxes,
coarsen_tree_of_boxes,
make_meshmode_mesh_from_leaves,
make_tree_of_boxes_root,
refine_tree_of_boxes,
uniformly_refine_tree_of_boxes,
)

Expand Down Expand Up @@ -71,7 +78,7 @@ class FromMeshmodeContextFactory(PytestPyOpenCLArrayContextFactory):

# {{{ make_global_leaf_quadrature

def make_global_leaf_quadrature(actx, tob, order):
def make_global_leaf_quadrature(actx: ArrayContext, tob: TreeOfBoxes, order: int):
from meshmode.discretization.poly_element import (
GaussLegendreTensorProductGroupFactory,
)
Expand Down Expand Up @@ -143,7 +150,11 @@ def test_uniform_tree_of_boxes(actx_factory, dim, order, nlevels):

@pytest.mark.parametrize("dim", [1, 2, 3])
@pytest.mark.parametrize("order", [1, 2, 3])
def test_uniform_tree_of_boxes_convergence(actx_factory, dim, order):
def test_uniform_tree_of_boxes_convergence(
actx_factory: ArrayContextFactory,
dim: int,
order: int
):
actx = actx_factory()

radius = np.pi
Expand Down Expand Up @@ -226,7 +237,7 @@ def test_tree_plot():
# {{{ test_traversal_from_tob


def test_traversal_from_tob(actx_factory):
def test_traversal_from_tob(actx_factory: ArrayContextFactory):
actx = actx_factory()

radius = np.pi
Expand Down Expand Up @@ -261,6 +272,78 @@ def test_traversal_from_tob(actx_factory):
# }}}


# {{{ test_refine_parent_child_consistency

@pytest.mark.parametrize("dim", [1, 2, 3])
def test_refine_parent_child_consistency(dim: int):
tob = make_tree_of_boxes_root((np.zeros(dim), np.ones(dim)))

tob = uniformly_refine_tree_of_boxes(tob)
tob = uniformly_refine_tree_of_boxes(tob)

nboxes = tob.nboxes
for box_id in range(nboxes):
for child_id in tob.box_child_ids[:, box_id]:
if child_id == 0:
continue
assert 0 < child_id < nboxes
assert tob.box_parent_ids[child_id] == box_id
assert tob.box_levels[child_id] == tob.box_levels[box_id] + 1

# }}}


# {{{ test_coarsen_tree_of_boxes

@pytest.mark.parametrize("dim", [1, 2, 3])
def test_coarsen_tree_of_boxes(dim: int):
tob = make_tree_of_boxes_root((np.zeros(dim), np.ones(dim)))
tob = uniformly_refine_tree_of_boxes(tob)

coarsen_flags = np.zeros(tob.nboxes, dtype=bool)
coarsen_flags[1:] = True
tob = coarsen_tree_of_boxes(tob, coarsen_flags)

assert tob.nboxes == 1
assert tob.box_parent_ids[0] == -1
assert tob.box_levels[0] == 0
assert np.all(tob.box_child_ids[:, 0] == 0)

# }}}


# {{{ test_out_of_order_refine

@pytest.mark.parametrize("dim", [1, 2, 3])
def test_out_of_order_refine(dim: int):
tob = make_tree_of_boxes_root((np.zeros(dim), np.ones(dim)))
tob = uniformly_refine_tree_of_boxes(tob)

flags = np.zeros(tob.nboxes, dtype=bool)
flags[1] = True
tob = refine_tree_of_boxes(tob, flags)

deep = np.where(tob.box_levels == 2)[0][0]
flags = np.zeros(tob.nboxes, dtype=bool)
flags[deep] = True
tob = refine_tree_of_boxes(tob, flags)

flags = np.zeros(tob.nboxes, dtype=bool)
flags[2] = True
tob = refine_tree_of_boxes(tob, flags)

nboxes = tob.nboxes
for box_id in range(nboxes):
for child_id in tob.box_child_ids[:, box_id]:
if child_id == 0:
continue
assert 0 < child_id < nboxes
assert tob.box_parent_ids[child_id] == box_id
assert tob.box_levels[child_id] == tob.box_levels[box_id] + 1

# }}}


# You can test individual routines by typing
# $ python test_tree.py 'test_routine(cl.create_some_context)'

Expand Down
Loading