Skip to content
test_tree.py 37 KiB
Newer Older
    if visualize:
Matt Wala's avatar
Matt Wala committed
        import matplotlib.pyplot as pt
        np_particles = actx.to_numpy(particles)
        pt.plot(np_particles[0], np_particles[1], "x")
Matt Wala's avatar
Matt Wala committed

    from boxtree import TreeBuilder
    tb = TreeBuilder(actx.context)
    actx.queue.finish()
    tree, _ = tb(actx.queue, particles, max_particles_in_box=30, debug=True)
Matt Wala's avatar
Matt Wala committed

    nballs = 10**4
    ball_centers = make_normal_particle_array(actx.queue, nballs, dims, dtype)
    ball_radii = 0.1 + actx.zeros(nballs, dtype)
Matt Wala's avatar
Matt Wala committed

    from boxtree.area_query import (
        LeavesToBallsLookupBuilder, SpaceInvaderQueryBuilder)

    siqb = SpaceInvaderQueryBuilder(actx.context)
Matt Wala's avatar
Matt Wala committed
    # We can use leaves-to-balls lookup to get the set of overlapping balls for
    # each box, and from there to compute the outer space invader distance.
    lblb = LeavesToBallsLookupBuilder(actx.context)
    siq, _ = siqb(actx.queue, tree, ball_centers, ball_radii)
    lbl, _ = lblb(actx.queue, tree, ball_centers, ball_radii)
Matt Wala's avatar
Matt Wala committed

    # get data to host for test
    tree = tree.get(queue=actx.queue)
    siq = siq.get(queue=actx.queue)
    lbl = lbl.get(queue=actx.queue)
    ball_centers = np.array([actx.to_numpy(x) for x in ball_centers])
    ball_radii = actx.to_numpy(ball_radii)
Matt Wala's avatar
Matt Wala committed

    # Find leaf boxes.
    from boxtree import box_flags_enum

    outer_space_invader_dist = np.zeros(tree.nboxes)

    for ibox in range(tree.nboxes):
        # We only want leaves here.
        if tree.box_flags[ibox] & box_flags_enum.HAS_CHILDREN:
            continue

        start, end = lbl.balls_near_box_starts[ibox:ibox + 2]
        space_invaders = lbl.balls_near_box_lists[start:end]
        if len(space_invaders) > 0:
            outer_space_invader_dist[ibox] = np.max(np.abs(
                    tree.box_centers[:, ibox].reshape((-1, 1))
                    - ball_centers[:, space_invaders]))

    assert np.allclose(siq, outer_space_invader_dist)

# }}}


# {{{ test_same_tree_with_zero_weight_particles

@pytest.mark.parametrize("dims", [2, 3])
def test_same_tree_with_zero_weight_particles(actx_factory, dims):
    actx = actx_factory()

    ntargets_values = [300, 400, 500]
    stick_out_factors = [0, 0.1, 0.3, 1]
    nsources = 20

    from boxtree import TreeBuilder
    tb = TreeBuilder(actx.context)
    rng = np.random.default_rng(10)
    for stick_out_factor in stick_out_factors:
        for ntargets in [40]:
            sources = rng.random((dims, nsources))**2
            sources[:, 0] = -0.1
            sources[:, 1] = 1.1

            targets = rng.random((dims, max(ntargets_values)))[:, :ntargets].copy()
            target_radii = rng.random(max(ntargets_values))[:ntargets]
            sources = actx.from_numpy(sources)
            targets = actx.from_numpy(targets)
            refine_weights = actx.empty(nsources + ntargets, np.int32)
            refine_weights[:nsources] = 1
            refine_weights[nsources:] = 0

            tree, _ = tb(actx.queue, sources, targets=targets,
                    target_radii=target_radii,
                    stick_out_factor=stick_out_factor,
                    max_leaf_refine_weight=10,
                    refine_weights=refine_weights,
                    debug=True)
            tree = tree.get(queue=actx.queue)
            trees.append(tree)

            print("TREE:", tree.nboxes)

    if 0:
        import matplotlib.pyplot as plt
        for tree in trees:
            plt.figure()
            tree.plot()

        plt.show()

# {{{ test_max_levels_error

def test_max_levels_error(actx_factory):
    actx = actx_factory()

    from boxtree import TreeBuilder
    tb = TreeBuilder(actx.context)
    sources = [actx.zeros(11, np.float64) for i in range(2)]
    from boxtree.tree_build import MaxLevelsExceeded
    with pytest.raises(MaxLevelsExceeded):
        tree, _ = tb(actx.queue, sources, max_particles_in_box=10, debug=True)
Andreas Klöckner's avatar
Andreas Klöckner committed
# You can test individual routines by typing
Andreas Klöckner's avatar
Andreas Klöckner committed
# $ python test_tree.py 'test_routine(cl.create_some_context)'
Andreas Klöckner's avatar
Andreas Klöckner committed

if __name__ == "__main__":
    if len(sys.argv) > 1:
        exec(sys.argv[1])
    else:
        from pytest import main
Andreas Klöckner's avatar
Andreas Klöckner committed
        main([__file__])

# vim: fdm=marker