diff --git a/boxtree/traversal.py b/boxtree/traversal.py index a35c60d08cfcc9850b49f50fd7fd4b561c1019d4..69f04b09401405cca8c075ebd12353a33269ffa9 100644 --- a/boxtree/traversal.py +++ b/boxtree/traversal.py @@ -810,9 +810,21 @@ void generate(LIST_ARG_DECL USER_ARG_DECL box_id_t target_box_number) // box, but our stick-out regions might still have a // non-empty intersection. - if (meets_sep_crit - && box_source_counts_cumul[walk_box_id] - >= from_sep_smaller_min_nsources_cumul) + // If the number of particles in this box is below the + // source count threshold, it can be moved to a "close" list. + // This is a performance optimization. + + bool close_lists_exist = ${"true" \ + if sources_have_extent or targets_have_extent \ + else "false"}; + + bool force_close_list_for_low_interaction_count = + close_lists_exist && + (box_source_counts_cumul[walk_box_id] + < from_sep_smaller_min_nsources_cumul); + + if (meets_sep_crit && + !force_close_list_for_low_interaction_count) { if (from_sep_smaller_source_level == walk_level) APPEND_from_sep_smaller(walk_box_id); diff --git a/test/test_fmm.py b/test/test_fmm.py index cb5e502818e184e330eb4dcd0b3e8ea731ed7e3e..082b43176d84385098d8298fd1f50e56c7383810 100644 --- a/test/test_fmm.py +++ b/test/test_fmm.py @@ -656,6 +656,64 @@ def test_pyfmmlib_fmm(ctx_getter, dims, use_dipoles, helmholtz_k): # }}} +# {{{ test particle count thresholding in traversal generation + +@pytest.mark.parametrize("enable_extents", [True, False]) +def test_interaction_list_particle_count_thresholding(ctx_getter, enable_extents): + ctx = ctx_getter() + queue = cl.CommandQueue(ctx) + + logging.basicConfig(level=logging.INFO) + + dims = 2 + nsources = 1000 + ntargets = 1000 + dtype = np.float + + max_particles_in_box = 30 + # Ensure that we have underfilled boxes. + from_sep_smaller_min_nsources_cumul = 1 + max_particles_in_box + + from boxtree.fmm import drive_fmm + sources = p_normal(queue, nsources, dims, dtype, seed=15) + targets = p_normal(queue, ntargets, dims, dtype, seed=15) + + from pyopencl.clrandom import PhiloxGenerator + rng = PhiloxGenerator(queue.context, seed=12) + + if enable_extents: + target_radii = 2**rng.uniform(queue, ntargets, dtype=dtype, a=-10, b=0) + else: + target_radii = None + + from boxtree import TreeBuilder + tb = TreeBuilder(ctx) + + tree, _ = tb(queue, sources, targets=targets, + max_particles_in_box=max_particles_in_box, + target_radii=target_radii, + debug=True, stick_out_factor=0.25) + + from boxtree.traversal import FMMTraversalBuilder + tbuild = FMMTraversalBuilder(ctx) + trav, _ = tbuild(queue, tree, debug=True, + _from_sep_smaller_min_nsources_cumul=from_sep_smaller_min_nsources_cumul) + + weights = np.ones(nsources) + weights_sum = np.sum(weights) + + host_trav = trav.get(queue=queue) + host_tree = host_trav.tree + + wrangler = ConstantOneExpansionWrangler(host_tree) + + pot = drive_fmm(host_trav, wrangler, weights) + + assert (pot == weights_sum).all() + +# }}} + + # You can test individual routines by typing # $ python test_fmm.py 'test_routine(cl.create_some_context)'