def __init__(self, queue, geo_data, use_target_specific_qbx):
from pytential.qbx.utils import ToHostTransferredGeoDataWrapper
geo_data = ToHostTransferredGeoDataWrapper(queue, geo_data)
self.geo_data = geo_data
self.trav = geo_data.traversal()
self.using_tsqbx = (
use_target_specific_qbx
# None means use by default if possible
or use_target_specific_qbx is None)
ConstantOneExpansionWrangler.__init__(self, geo_data.tree())
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()
def test_fmm_float32(ctx_getter=cl.create_some_context, enable_extents=True):
from time import time
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
from pyopencl.characterize import has_struct_arg_count_bug
if has_struct_arg_count_bug(queue.device):
pytest.xfail("won't work on devices with the struct arg count issue")
logging.basicConfig(level=logging.INFO)
dims = 2
nsources = 3000000
ntargets = 3000000
dtype = np.float32
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=30,
target_radii=target_radii,stick_out_factor=0.25,
debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
weights = np.ones(nsources)
weights_sum = np.sum(weights)
host_trav = trav.get(queue=queue)
host_tree = host_trav.tree
wrangler = ConstantOneExpansionWrangler(host_tree)
ti = time()
pot = drive_fmm(host_trav, wrangler, weights)
print(time() - ti)
assert (pot == weights_sum).all()
def __init__(self, tree, filtered_targets):
ConstantOneExpansionWrangler.__init__(self, tree)
self.filtered_targets = filtered_targets
def test_fmm_completeness(ctx_getter, dims, nsources_req, ntargets_req,
who_has_extent, source_gen, target_gen, filter_kind,
well_sep_is_n_away, extent_norm,
from_sep_smaller_crit):
"""Tests whether the built FMM traversal structures and driver completely
capture all interactions.
"""
sources_have_extent = "s" in who_has_extent
targets_have_extent = "t" in who_has_extent
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
dtype = np.float64
try:
sources = source_gen(queue, nsources_req, dims, dtype, seed=15)
nsources = len(sources[0])
if ntargets_req is None:
# This says "same as sources" to the tree builder.
targets = None
ntargets = ntargets_req
else:
targets = target_gen(queue, ntargets_req, dims, dtype, seed=16)
ntargets = len(targets[0])
except ImportError:
pytest.skip("loo.py not available, but needed for particle array "
"generation")
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=12)
if sources_have_extent:
source_radii = 2**rng.uniform(queue, nsources, dtype=dtype, a=-10, b=0)
else:
source_radii = None
if targets_have_extent:
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=30,
source_radii=source_radii,
target_radii=target_radii,
debug=True,
stick_out_factor=0.25,
extent_norm=extent_norm)
if 0:
tree.get().plot()
import matplotlib.pyplot as pt
pt.show()
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx,
well_sep_is_n_away=well_sep_is_n_away,
from_sep_smaller_crit=from_sep_smaller_crit)
trav, _ = tbuild(queue, tree, debug=True)
if who_has_extent:
pre_merge_trav = trav
trav = trav.merge_close_lists(queue)
#weights = np.random.randn(nsources)
weights = np.ones(nsources)
weights_sum = np.sum(weights)
host_trav = trav.get(queue=queue)
host_tree = host_trav.tree
if who_has_extent:
pre_merge_host_trav = pre_merge_trav.get(queue=queue)
from boxtree.tree import ParticleListFilter
plfilt = ParticleListFilter(ctx)
if filter_kind:
flags = rng.uniform(queue, ntargets or nsources, np.int32, a=0, b=2) \
.astype(np.int8)
if filter_kind == "user":
filtered_targets = plfilt.filter_target_lists_in_user_order(
queue, tree, flags)
wrangler = ConstantOneExpansionWranglerWithFilteredTargetsInUserOrder(
host_tree, filtered_targets.get(queue=queue))
elif filter_kind == "tree":
filtered_targets = plfilt.filter_target_lists_in_tree_order(
queue, tree, flags)
wrangler = ConstantOneExpansionWranglerWithFilteredTargetsInTreeOrder(
host_tree, filtered_targets.get(queue=queue))
else:
raise ValueError("unsupported value of 'filter_kind'")
else:
wrangler = ConstantOneExpansionWrangler(host_tree)
flags = cl.array.empty(queue, ntargets or nsources, dtype=np.int8)
flags.fill(1)
if ntargets is None and not filter_kind:
# This check only works for targets == sources.
assert (wrangler.reorder_potentials(
wrangler.reorder_sources(weights)) == weights).all()
from boxtree.fmm import drive_fmm
pot = drive_fmm(host_trav, wrangler, weights)
if filter_kind:
pot = pot[flags.get() > 0]
rel_err = la.norm((pot - weights_sum) / nsources)
good = rel_err < 1e-8
# {{{ build, evaluate matrix (and identify incorrect interactions)
if 0 and not good:
mat = np.zeros((ntargets, nsources), dtype)
from pytools import ProgressBar
logging.getLogger().setLevel(logging.WARNING)
pb = ProgressBar("matrix", nsources)
for i in range(nsources):
unit_vec = np.zeros(nsources, dtype=dtype)
unit_vec[i] = 1
mat[:, i] = drive_fmm(host_trav, wrangler, unit_vec)
pb.progress()
pb.finished()
logging.getLogger().setLevel(logging.INFO)
import matplotlib.pyplot as pt
if 0:
pt.imshow(mat)
pt.colorbar()
pt.show()
incorrect_tgts, incorrect_srcs = np.where(mat != 1)
if 1 and len(incorrect_tgts):
from boxtree.visualization import TreePlotter
plotter = TreePlotter(host_tree)
plotter.draw_tree(fill=False, edgecolor="black")
plotter.draw_box_numbers()
plotter.set_bounding_box()
tree_order_incorrect_tgts = \
host_tree.indices_to_tree_target_order(incorrect_tgts)
tree_order_incorrect_srcs = \
host_tree.indices_to_tree_source_order(incorrect_srcs)
src_boxes = [
host_tree.find_box_nr_for_source(i)
for i in tree_order_incorrect_srcs
]
tgt_boxes = [
host_tree.find_box_nr_for_target(i)
for i in tree_order_incorrect_tgts
]
print(src_boxes)
print(tgt_boxes)
# plot all sources/targets
if 0:
pt.plot(host_tree.targets[0],
host_tree.targets[1],
"v",
alpha=0.9)
pt.plot(host_tree.sources[0],
host_tree.sources[1],
"gx",
alpha=0.9)
# plot offending sources/targets
if 0:
pt.plot(host_tree.targets[0][tree_order_incorrect_tgts],
host_tree.targets[1][tree_order_incorrect_tgts], "rv")
pt.plot(host_tree.sources[0][tree_order_incorrect_srcs],
host_tree.sources[1][tree_order_incorrect_srcs], "go")
pt.gca().set_aspect("equal")
from boxtree.visualization import draw_box_lists
draw_box_lists(
plotter, pre_merge_host_trav if who_has_extent else host_trav,
22)
# from boxtree.visualization import draw_same_level_non_well_sep_boxes
# draw_same_level_non_well_sep_boxes(plotter, host_trav, 2)
pt.show()
# }}}
if 0 and not good:
import matplotlib.pyplot as pt
pt.plot(pot - weights_sum)
pt.show()
if 0 and not good:
import matplotlib.pyplot as pt
filt_targets = [
host_tree.targets[0][flags.get() > 0],
host_tree.targets[1][flags.get() > 0],
]
host_tree.plot()
bad = np.abs(pot - weights_sum) >= 1e-3
bad_targets = [
filt_targets[0][bad],
filt_targets[1][bad],
]
print(bad_targets[0].shape)
pt.plot(filt_targets[0], filt_targets[1], "x")
pt.plot(bad_targets[0], bad_targets[1], "v")
pt.show()
assert good
def test_cost_model_op_counts_agree_with_constantone_wrangler(
ctx_factory, nsources, ntargets, dims, dtype):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
from boxtree.tools import make_normal_particle_array as p_normal
sources = p_normal(queue, nsources, dims, dtype, seed=16)
targets = p_normal(queue, ntargets, dims, dtype, seed=19)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=20)
target_radii = rng.uniform(queue, ntargets, a=0, b=0.04, dtype=dtype).get()
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(
queue, sources, targets=targets, target_radii=target_radii,
stick_out_factor=0.15, max_particles_in_box=30, debug=True
)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx, well_sep_is_n_away=2)
trav_dev, _ = tg(queue, tree, debug=True)
trav = trav_dev.get(queue=queue)
from boxtree.tools import ConstantOneExpansionWrangler
wrangler = ConstantOneExpansionWrangler(trav.tree)
timing_data = {}
from boxtree.fmm import drive_fmm
src_weights = np.random.rand(tree.nsources).astype(tree.coord_dtype)
drive_fmm(trav, wrangler, (src_weights,), timing_data=timing_data)
cost_model = FMMCostModel(
translation_cost_model_factory=OpCountingTranslationCostModel
)
level_to_order = np.array([1 for _ in range(tree.nlevels)])
modeled_time = cost_model.cost_per_stage(
queue, trav_dev, level_to_order,
FMMCostModel.get_unit_calibration_params(),
)
mismatches = []
for stage in timing_data:
if timing_data[stage]["ops_elapsed"] != modeled_time[stage]:
mismatches.append(
(stage, timing_data[stage]["ops_elapsed"], modeled_time[stage]))
assert not mismatches, "\n".join(str(s) for s in mismatches)
# {{{ Test per-box cost
total_cost = 0.0
for stage in timing_data:
total_cost += timing_data[stage]["ops_elapsed"]
per_box_cost = cost_model.cost_per_box(
queue, trav_dev, level_to_order,
FMMCostModel.get_unit_calibration_params(),
)
total_aggregate_cost = cost_model.aggregate_over_boxes(per_box_cost)
assert total_cost == (
total_aggregate_cost
+ modeled_time["coarsen_multipoles"]
+ modeled_time["refine_locals"]
)
