def __init__(self,
ctx_getter=cl.create_some_context,
enable_extents=False):
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 = 9000000
ntargets = 9000000
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
self.tree = host_tree
self.trav = host_trav
self.input = [host_tree, weights, weights_sum, host_trav]
self.pot = None
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_pyfmmlib_fmm(ctx_getter):
logging.basicConfig(level=logging.INFO)
from pytest import importorskip
importorskip("pyfmmlib")
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 3000
ntargets = 1000
dims = 2
dtype = np.float64
helmholtz_k = 2
sources = p_normal(queue, nsources, dims, dtype, seed=15)
targets = (
p_normal(queue, ntargets, dims, dtype, seed=18)
+ np.array([2, 0]))
sources_host = particle_array_to_host(sources)
targets_host = particle_array_to_host(targets)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources, targets=targets,
max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
trav = trav.get(queue=queue)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=20)
weights = rng.uniform(queue, nsources, dtype=np.float64).get()
#weights = np.ones(nsources)
logger.info("computing direct (reference) result")
from pyfmmlib import hpotgrad2dall_vec
ref_pot, _, _ = hpotgrad2dall_vec(ifgrad=False, ifhess=False,
sources=sources_host.T, charge=weights,
targets=targets_host.T, zk=helmholtz_k)
from boxtree.pyfmmlib_integration import Helmholtz2DExpansionWrangler
wrangler = Helmholtz2DExpansionWrangler(trav.tree, helmholtz_k, nterms=10)
from boxtree.fmm import drive_fmm
pot = drive_fmm(trav, wrangler, weights)
rel_err = la.norm(pot - ref_pot) / la.norm(ref_pot)
logger.info("relative l2 error: %g" % rel_err)
assert rel_err < 1e-5
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 test_sumpy_fmm_timing_data_collection(ctx_getter):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(
ctx,
properties=cl.command_queue_properties.PROFILING_ENABLE)
nsources = 500
dtype = np.float64
from boxtree.tools import (
make_normal_particle_array as p_normal)
knl = LaplaceKernel(2)
local_expn_class = VolumeTaylorLocalExpansion
mpole_expn_class = VolumeTaylorMultipoleExpansion
order = 1
sources = p_normal(queue, nsources, knl.dim, dtype, seed=15)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources,
max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(ctx)
weights = rng.uniform(queue, nsources, dtype=np.float64)
out_kernels = [knl]
from functools import partial
from sumpy.fmm import SumpyExpansionWranglerCodeContainer
wcc = SumpyExpansionWranglerCodeContainer(
ctx,
partial(mpole_expn_class, knl),
partial(local_expn_class, knl),
out_kernels)
wrangler = wcc.get_wrangler(queue, tree, dtype,
fmm_level_to_order=lambda kernel, kernel_args, tree, lev: order)
from boxtree.fmm import drive_fmm
timing_data = {}
pot, = drive_fmm(trav, wrangler, weights, timing_data=timing_data)
print(timing_data)
assert timing_data
def test_sumpy_fmm_timing_data_collection(ctx_factory):
logging.basicConfig(level=logging.INFO)
ctx = ctx_factory()
queue = cl.CommandQueue(
ctx,
properties=cl.command_queue_properties.PROFILING_ENABLE)
nsources = 500
dtype = np.float64
from boxtree.tools import (
make_normal_particle_array as p_normal)
knl = LaplaceKernel(2)
local_expn_class = VolumeTaylorLocalExpansion
mpole_expn_class = VolumeTaylorMultipoleExpansion
order = 1
sources = p_normal(queue, nsources, knl.dim, dtype, seed=15)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources,
max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(ctx)
weights = rng.uniform(queue, nsources, dtype=np.float64)
out_kernels = [knl]
from functools import partial
from sumpy.fmm import SumpyExpansionWranglerCodeContainer
wcc = SumpyExpansionWranglerCodeContainer(
ctx,
partial(mpole_expn_class, knl),
partial(local_expn_class, knl),
out_kernels)
wrangler = wcc.get_wrangler(queue, tree, dtype,
fmm_level_to_order=lambda kernel, kernel_args, tree, lev: order)
from boxtree.fmm import drive_fmm
timing_data = {}
pot, = drive_fmm(trav, wrangler, (weights,), timing_data=timing_data)
print(timing_data)
assert timing_data
def test_level_to_order_lookup(ctx_getter, lookup_func, extra_args):
pytest.importorskip("pyfmmlib")
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 500
dtype = np.float64
epsilon = 1e-9
from boxtree.tools import (make_normal_particle_array as p_normal)
sources = p_normal(queue, nsources, 2, dtype, seed=15)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources, max_particles_in_box=30, debug=True)
levels = lookup_func(tree, *(extra_args + (epsilon, )))
assert len(levels) == tree.nlevels
# Should be monotonically nonincreasing.
assert all(levels[i] >= levels[i + 1] for i in range(tree.nlevels - 1))
def test_level_to_order_lookup(ctx_getter, lookup_func, extra_args):
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 500
dtype = np.float64
epsilon = 1e-9
from boxtree.tools import (
make_normal_particle_array as p_normal)
sources = p_normal(queue, nsources, 2, dtype, seed=15)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources, max_particles_in_box=30, debug=True)
levels = lookup_func(tree, *(extra_args + (epsilon,)))
assert len(levels) == tree.nlevels
# Should be monotonically nonincreasing.
assert all(levels[i] >= levels[i+1] for i in range(tree.nlevels - 1))
def test_sumpy_fmm_exclude_self(ctx_getter):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 500
dtype = np.float64
from boxtree.tools import (make_normal_particle_array as p_normal)
knl = LaplaceKernel(2)
local_expn_class = VolumeTaylorLocalExpansion
mpole_expn_class = VolumeTaylorMultipoleExpansion
order = 10
sources = p_normal(queue, nsources, knl.dim, dtype, seed=15)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources, max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(ctx)
weights = rng.uniform(queue, nsources, dtype=np.float64)
target_to_source = np.arange(tree.ntargets, dtype=np.int32)
self_extra_kwargs = {"target_to_source": target_to_source}
out_kernels = [knl]
from functools import partial
from sumpy.fmm import SumpyExpansionWranglerCodeContainer
wcc = SumpyExpansionWranglerCodeContainer(ctx,
partial(mpole_expn_class, knl),
partial(local_expn_class, knl),
out_kernels,
exclude_self=True)
wrangler = wcc.get_wrangler(
queue,
tree,
dtype,
fmm_level_to_order=lambda kernel, kernel_args, tree, lev: order,
self_extra_kwargs=self_extra_kwargs)
from boxtree.fmm import drive_fmm
pot, = drive_fmm(trav, wrangler, weights)
from sumpy import P2P
p2p = P2P(ctx, out_kernels, exclude_self=True)
evt, (ref_pot, ) = p2p(queue, sources, sources, (weights, ),
**self_extra_kwargs)
pot = pot.get()
ref_pot = ref_pot.get()
rel_err = la.norm(pot - ref_pot) / la.norm(ref_pot)
logger.info("order %d -> relative l2 error: %g" % (order, rel_err))
assert np.isclose(rel_err, 0, atol=1e-7)
def test_pyfmmlib_fmm(ctx_getter, dims, use_dipoles, helmholtz_k):
logging.basicConfig(level=logging.INFO)
from pytest import importorskip
importorskip("pyfmmlib")
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 3000
ntargets = 1000
dtype = np.float64
sources = p_normal(queue, nsources, dims, dtype, seed=15)
targets = (p_normal(queue, ntargets, dims, dtype, seed=18) +
np.array([2, 0, 0])[:dims])
sources_host = particle_array_to_host(sources)
targets_host = particle_array_to_host(targets)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue,
sources,
targets=targets,
max_particles_in_box=30,
debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
trav = trav.get(queue=queue)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=20)
weights = rng.uniform(queue, nsources, dtype=np.float64).get()
#weights = np.ones(nsources)
if use_dipoles:
np.random.seed(13)
dipole_vec = np.random.randn(dims, nsources)
else:
dipole_vec = None
if dims == 2 and helmholtz_k == 0:
base_nterms = 20
else:
base_nterms = 10
def fmm_level_to_nterms(tree, lev):
result = base_nterms
if lev < 3 and helmholtz_k:
# exercise order-varies-by-level capability
result += 5
if use_dipoles:
result += 1
return result
from boxtree.pyfmmlib_integration import FMMLibExpansionWrangler
wrangler = FMMLibExpansionWrangler(trav.tree,
helmholtz_k,
fmm_level_to_nterms=fmm_level_to_nterms,
dipole_vec=dipole_vec)
from boxtree.fmm import drive_fmm
timing_data = {}
pot = drive_fmm(trav, wrangler, weights, timing_data=timing_data)
print(timing_data)
assert timing_data
# {{{ ref fmmlib computation
logger.info("computing direct (reference) result")
import pyfmmlib
fmmlib_routine = getattr(
pyfmmlib, "%spot%s%ddall%s_vec" %
(wrangler.eqn_letter, "fld" if dims == 3 else "grad", dims,
"_dp" if use_dipoles else ""))
kwargs = {}
if dims == 3:
kwargs["iffld"] = False
else:
kwargs["ifgrad"] = False
kwargs["ifhess"] = False
if use_dipoles:
if helmholtz_k == 0 and dims == 2:
kwargs["dipstr"] = -weights * (dipole_vec[0] + 1j * dipole_vec[1])
else:
kwargs["dipstr"] = weights
kwargs["dipvec"] = dipole_vec
else:
kwargs["charge"] = weights
if helmholtz_k:
kwargs["zk"] = helmholtz_k
ref_pot = wrangler.finalize_potentials(
fmmlib_routine(sources=sources_host.T,
targets=targets_host.T,
**kwargs)[0])
rel_err = la.norm(pot - ref_pot, np.inf) / la.norm(ref_pot, np.inf)
logger.info("relative l2 error vs fmmlib direct: %g" % rel_err)
assert rel_err < 1e-5, rel_err
# }}}
# {{{ check against sumpy
try:
import sumpy # noqa
except ImportError:
have_sumpy = False
from warnings import warn
warn("sumpy unavailable: cannot compute independent reference "
"values for pyfmmlib")
else:
have_sumpy = True
if have_sumpy:
from sumpy.kernel import (LaplaceKernel, HelmholtzKernel,
DirectionalSourceDerivative)
from sumpy.p2p import P2P
sumpy_extra_kwargs = {}
if helmholtz_k:
knl = HelmholtzKernel(dims)
sumpy_extra_kwargs["k"] = helmholtz_k
else:
knl = LaplaceKernel(dims)
if use_dipoles:
knl = DirectionalSourceDerivative(knl)
sumpy_extra_kwargs["src_derivative_dir"] = dipole_vec
p2p = P2P(ctx, [knl], exclude_self=False)
evt, (sumpy_ref_pot, ) = p2p(queue,
targets,
sources, [weights],
out_host=True,
**sumpy_extra_kwargs)
sumpy_rel_err = (la.norm(pot - sumpy_ref_pot, np.inf) /
la.norm(sumpy_ref_pot, np.inf))
logger.info("relative l2 error vs sumpy direct: %g" % sumpy_rel_err)
assert sumpy_rel_err < 1e-5, sumpy_rel_err
def test_fmm_with_optimized_3d_m2l(ctx_factory, nsrcntgts, helmholtz_k,
well_sep_is_n_away):
logging.basicConfig(level=logging.INFO)
from pytest import importorskip
importorskip("pyfmmlib")
dims = 3
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nsources = ntargets = nsrcntgts // 2
dtype = np.float64
sources = p_normal(queue, nsources, dims, dtype, seed=15)
targets = (p_normal(queue, ntargets, dims, dtype, seed=18) +
np.array([2, 0, 0])[:dims])
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue,
sources,
targets=targets,
max_particles_in_box=30,
debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
trav = trav.get(queue=queue)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=20)
weights = rng.uniform(queue, nsources, dtype=np.float64).get()
base_nterms = 10
def fmm_level_to_nterms(tree, lev):
result = base_nterms
if lev < 3 and helmholtz_k:
# exercise order-varies-by-level capability
result += 5
return result
from boxtree.pyfmmlib_integration import (FMMLibExpansionWrangler,
FMMLibRotationData)
baseline_wrangler = FMMLibExpansionWrangler(
trav.tree, helmholtz_k, fmm_level_to_nterms=fmm_level_to_nterms)
optimized_wrangler = FMMLibExpansionWrangler(
trav.tree,
helmholtz_k,
fmm_level_to_nterms=fmm_level_to_nterms,
rotation_data=FMMLibRotationData(queue, trav))
from boxtree.fmm import drive_fmm
baseline_timing_data = {}
baseline_pot = drive_fmm(trav,
baseline_wrangler, (weights, ),
timing_data=baseline_timing_data)
optimized_timing_data = {}
optimized_pot = drive_fmm(trav,
optimized_wrangler, (weights, ),
timing_data=optimized_timing_data)
baseline_time = baseline_timing_data["multipole_to_local"][
"process_elapsed"]
if baseline_time is not None:
print("Baseline M2L time : %#.4g s" % baseline_time)
opt_time = optimized_timing_data["multipole_to_local"]["process_elapsed"]
if opt_time is not None:
print("Optimized M2L time: %#.4g s" % opt_time)
assert np.allclose(baseline_pot, optimized_pot, atol=1e-13, rtol=1e-13)
def test_compare_cl_and_py_cost_model(ctx_factory, nsources, ntargets, dims, dtype):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
# {{{ Generate sources, targets and target_radii
from boxtree.tools import make_normal_particle_array as p_normal
sources = p_normal(queue, nsources, dims, dtype, seed=15)
targets = p_normal(queue, ntargets, dims, dtype, seed=18)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=22)
target_radii = rng.uniform(
queue, ntargets, a=0, b=0.05, dtype=dtype
).get()
# }}}
# {{{ Generate tree and traversal
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)
# }}}
# {{{ Construct cost models
cl_cost_model = FMMCostModel(None)
python_cost_model = _PythonFMMCostModel(None)
constant_one_params = cl_cost_model.get_unit_calibration_params().copy()
for ilevel in range(trav.tree.nlevels):
constant_one_params["p_fmm_lev%d" % ilevel] = 10
xlat_cost = make_pde_aware_translation_cost_model(dims, trav.tree.nlevels)
# }}}
# {{{ Test process_form_multipoles
nlevels = trav.tree.nlevels
p2m_cost = np.zeros(nlevels, dtype=np.float64)
for ilevel in range(nlevels):
p2m_cost[ilevel] = evaluate(
xlat_cost.p2m(ilevel),
context=constant_one_params
)
p2m_cost_dev = cl.array.to_device(queue, p2m_cost)
queue.finish()
start_time = time.time()
cl_form_multipoles = cl_cost_model.process_form_multipoles(
queue, trav_dev, p2m_cost_dev
)
queue.finish()
logger.info("OpenCL time for process_form_multipoles: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_form_multipoles = python_cost_model.process_form_multipoles(
queue, trav, p2m_cost
)
logger.info("Python time for process_form_multipoles: {0}".format(
str(time.time() - start_time)
))
assert np.array_equal(cl_form_multipoles.get(), python_form_multipoles)
# }}}
# {{{ Test process_coarsen_multipoles
m2m_cost = np.zeros(nlevels - 1, dtype=np.float64)
for target_level in range(nlevels - 1):
m2m_cost[target_level] = evaluate(
xlat_cost.m2m(target_level + 1, target_level),
context=constant_one_params
)
m2m_cost_dev = cl.array.to_device(queue, m2m_cost)
queue.finish()
start_time = time.time()
cl_coarsen_multipoles = cl_cost_model.process_coarsen_multipoles(
queue, trav_dev, m2m_cost_dev
)
queue.finish()
logger.info("OpenCL time for coarsen_multipoles: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_coarsen_multipoles = python_cost_model.process_coarsen_multipoles(
queue, trav, m2m_cost
)
logger.info("Python time for coarsen_multipoles: {0}".format(
str(time.time() - start_time)
))
assert cl_coarsen_multipoles == python_coarsen_multipoles
# }}}
# {{{ Test process_direct
queue.finish()
start_time = time.time()
cl_ndirect_sources_per_target_box = \
cl_cost_model.get_ndirect_sources_per_target_box(queue, trav_dev)
cl_direct = cl_cost_model.process_direct(
queue, trav_dev, cl_ndirect_sources_per_target_box, 5.0
)
queue.finish()
logger.info("OpenCL time for process_direct: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_ndirect_sources_per_target_box = \
python_cost_model.get_ndirect_sources_per_target_box(queue, trav)
python_direct = python_cost_model.process_direct(
queue, trav, python_ndirect_sources_per_target_box, 5.0
)
logger.info("Python time for process_direct: {0}".format(
str(time.time() - start_time)
))
assert np.array_equal(cl_direct.get(), python_direct)
# }}}
# {{{ Test aggregate_over_boxes
start_time = time.time()
cl_direct_aggregate = cl_cost_model.aggregate_over_boxes(cl_direct)
queue.finish()
logger.info("OpenCL time for aggregate_over_boxes: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_direct_aggregate = python_cost_model.aggregate_over_boxes(python_direct)
logger.info("Python time for aggregate_over_boxes: {0}".format(
str(time.time() - start_time)
))
assert cl_direct_aggregate == python_direct_aggregate
# }}}
# {{{ Test process_list2
nlevels = trav.tree.nlevels
m2l_cost = np.zeros(nlevels, dtype=np.float64)
for ilevel in range(nlevels):
m2l_cost[ilevel] = evaluate(
xlat_cost.m2l(ilevel, ilevel),
context=constant_one_params
)
m2l_cost_dev = cl.array.to_device(queue, m2l_cost)
queue.finish()
start_time = time.time()
cl_m2l_cost = cl_cost_model.process_list2(queue, trav_dev, m2l_cost_dev)
queue.finish()
logger.info("OpenCL time for process_list2: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_m2l_cost = python_cost_model.process_list2(queue, trav, m2l_cost)
logger.info("Python time for process_list2: {0}".format(
str(time.time() - start_time)
))
assert np.array_equal(cl_m2l_cost.get(), python_m2l_cost)
# }}}
# {{{ Test process_list 3
m2p_cost = np.zeros(nlevels, dtype=np.float64)
for ilevel in range(nlevels):
m2p_cost[ilevel] = evaluate(
xlat_cost.m2p(ilevel),
context=constant_one_params
)
m2p_cost_dev = cl.array.to_device(queue, m2p_cost)
queue.finish()
start_time = time.time()
cl_m2p_cost = cl_cost_model.process_list3(queue, trav_dev, m2p_cost_dev)
queue.finish()
logger.info("OpenCL time for process_list3: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_m2p_cost = python_cost_model.process_list3(queue, trav, m2p_cost)
logger.info("Python time for process_list3: {0}".format(
str(time.time() - start_time)
))
assert np.array_equal(cl_m2p_cost.get(), python_m2p_cost)
# }}}
# {{{ Test process_list4
p2l_cost = np.zeros(nlevels, dtype=np.float64)
for ilevel in range(nlevels):
p2l_cost[ilevel] = evaluate(
xlat_cost.p2l(ilevel),
context=constant_one_params
)
p2l_cost_dev = cl.array.to_device(queue, p2l_cost)
queue.finish()
start_time = time.time()
cl_p2l_cost = cl_cost_model.process_list4(queue, trav_dev, p2l_cost_dev)
queue.finish()
logger.info("OpenCL time for process_list4: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_p2l_cost = python_cost_model.process_list4(queue, trav, p2l_cost)
logger.info("Python time for process_list4: {0}".format(
str(time.time() - start_time)
))
assert np.array_equal(cl_p2l_cost.get(), python_p2l_cost)
# }}}
# {{{ Test process_refine_locals
l2l_cost = np.zeros(nlevels - 1, dtype=np.float64)
for ilevel in range(nlevels - 1):
l2l_cost[ilevel] = evaluate(
xlat_cost.l2l(ilevel, ilevel + 1),
context=constant_one_params
)
l2l_cost_dev = cl.array.to_device(queue, l2l_cost)
queue.finish()
start_time = time.time()
cl_refine_locals_cost = cl_cost_model.process_refine_locals(
queue, trav_dev, l2l_cost_dev
)
queue.finish()
logger.info("OpenCL time for refine_locals: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_refine_locals_cost = python_cost_model.process_refine_locals(
queue, trav, l2l_cost
)
logger.info("Python time for refine_locals: {0}".format(
str(time.time() - start_time)
))
assert cl_refine_locals_cost == python_refine_locals_cost
# }}}
# {{{ Test process_eval_locals
l2p_cost = np.zeros(nlevels, dtype=np.float64)
for ilevel in range(nlevels):
l2p_cost[ilevel] = evaluate(
xlat_cost.l2p(ilevel),
context=constant_one_params
)
l2p_cost_dev = cl.array.to_device(queue, l2p_cost)
queue.finish()
start_time = time.time()
cl_l2p_cost = cl_cost_model.process_eval_locals(queue, trav_dev, l2p_cost_dev)
queue.finish()
logger.info("OpenCL time for process_eval_locals: {0}".format(
str(time.time() - start_time)
))
start_time = time.time()
python_l2p_cost = python_cost_model.process_eval_locals(queue, trav, l2p_cost)
logger.info("Python time for process_eval_locals: {0}".format(
str(time.time() - start_time)
))
assert np.array_equal(cl_l2p_cost.get(), python_l2p_cost)
def demo_cost_model():
if not SUPPORTS_PROCESS_TIME:
raise NotImplementedError(
"Currently this script uses process time which only works on Python>=3.3"
)
from boxtree.pyfmmlib_integration import FMMLibExpansionWrangler
nsources_list = [1000, 2000, 3000, 4000, 5000]
ntargets_list = [1000, 2000, 3000, 4000, 5000]
dims = 3
dtype = np.float64
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
traversals = []
traversals_dev = []
level_to_orders = []
timing_results = []
def fmm_level_to_nterms(tree, ilevel):
return 10
for nsources, ntargets in zip(nsources_list, ntargets_list):
# {{{ Generate sources, targets and target_radii
from boxtree.tools import make_normal_particle_array as p_normal
sources = p_normal(queue, nsources, dims, dtype, seed=15)
targets = p_normal(queue, ntargets, dims, dtype, seed=18)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=22)
target_radii = rng.uniform(queue, ntargets, a=0, b=0.05,
dtype=dtype).get()
# }}}
# {{{ Generate tree and traversal
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)
traversals.append(trav)
traversals_dev.append(trav_dev)
# }}}
wrangler = FMMLibExpansionWrangler(trav.tree, 0, fmm_level_to_nterms)
level_to_orders.append(wrangler.level_nterms)
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)
timing_results.append(timing_data)
time_field_name = "process_elapsed"
from boxtree.cost import FMMCostModel
from boxtree.cost import make_pde_aware_translation_cost_model
cost_model = FMMCostModel(make_pde_aware_translation_cost_model)
model_results = []
for icase in range(len(traversals) - 1):
traversal = traversals_dev[icase]
model_results.append(
cost_model.cost_per_stage(
queue,
traversal,
level_to_orders[icase],
FMMCostModel.get_unit_calibration_params(),
))
queue.finish()
params = cost_model.estimate_calibration_params(
model_results, timing_results[:-1], time_field_name=time_field_name)
predicted_time = cost_model.cost_per_stage(
queue,
traversals_dev[-1],
level_to_orders[-1],
params,
)
queue.finish()
for field in [
"form_multipoles", "eval_direct", "multipole_to_local",
"eval_multipoles", "form_locals", "eval_locals",
"coarsen_multipoles", "refine_locals"
]:
measured = timing_results[-1][field]["process_elapsed"]
pred_err = ((measured - predicted_time[field]) / measured)
logger.info("actual/predicted time for %s: %.3g/%.3g -> %g %% error",
field, measured, predicted_time[field],
abs(100 * pred_err))
def test_pyfmmlib_fmm(ctx_getter):
logging.basicConfig(level=logging.INFO)
from pytest import importorskip
importorskip("pyfmmlib")
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 3000
ntargets = 1000
dims = 2
dtype = np.float64
helmholtz_k = 2
sources = p_normal(queue, nsources, dims, dtype, seed=15)
targets = (p_normal(queue, ntargets, dims, dtype, seed=18) +
np.array([2, 0]))
sources_host = particle_array_to_host(sources)
targets_host = particle_array_to_host(targets)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue,
sources,
targets=targets,
max_particles_in_box=30,
debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
trav = trav.get(queue=queue)
from pyopencl.clrandom import RanluxGenerator
rng = RanluxGenerator(queue, seed=20)
weights = rng.uniform(queue, nsources, dtype=np.float64).get()
#weights = np.ones(nsources)
logger.info("computing direct (reference) result")
from pyfmmlib import hpotgrad2dall_vec
ref_pot, _, _ = hpotgrad2dall_vec(ifgrad=False,
ifhess=False,
sources=sources_host.T,
charge=weights,
targets=targets_host.T,
zk=helmholtz_k)
from boxtree.pyfmmlib_integration import Helmholtz2DExpansionWrangler
wrangler = Helmholtz2DExpansionWrangler(trav.tree, helmholtz_k, nterms=10)
from boxtree.fmm import drive_fmm
pot = drive_fmm(trav, wrangler, weights)
rel_err = la.norm(pot - ref_pot) / la.norm(ref_pot)
logger.info("relative l2 error: %g" % rel_err)
assert rel_err < 1e-5
def test_sumpy_fmm(ctx_getter, knl, local_expn_class, mpole_expn_class):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 1000
ntargets = 300
dtype = np.float64
from boxtree.tools import (make_normal_particle_array as p_normal)
sources = p_normal(queue, nsources, knl.dim, dtype, seed=15)
if 1:
offset = np.zeros(knl.dim)
offset[0] = 0.1
targets = (p_normal(queue, ntargets, knl.dim, dtype, seed=18) + offset)
del offset
else:
from sumpy.visualization import FieldPlotter
fp = FieldPlotter(np.array([0.5, 0]), extent=3, npoints=200)
from pytools.obj_array import make_obj_array
targets = make_obj_array([fp.points[i] for i in range(knl.dim)])
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue,
sources,
targets=targets,
max_particles_in_box=30,
debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
# {{{ plot tree
if 0:
host_tree = tree.get()
host_trav = trav.get()
if 1:
print("src_box", host_tree.find_box_nr_for_source(403))
print("tgt_box", host_tree.find_box_nr_for_target(28))
print(list(host_trav.target_or_target_parent_boxes).index(37))
print(host_trav.get_box_list("sep_bigger", 22))
from boxtree.visualization import TreePlotter
plotter = TreePlotter(host_tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
plotter.draw_box_numbers()
import matplotlib.pyplot as pt
pt.show()
# }}}
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(ctx, seed=44)
weights = rng.uniform(queue, nsources, dtype=np.float64)
logger.info("computing direct (reference) result")
from pytools.convergence import PConvergenceVerifier
pconv_verifier = PConvergenceVerifier()
extra_kwargs = {}
dtype = np.float64
order_values = [1, 2, 3]
if isinstance(knl, HelmholtzKernel):
extra_kwargs["k"] = 0.05
dtype = np.complex128
if knl.dim == 3:
order_values = [1, 2]
elif knl.dim == 2 and issubclass(local_expn_class, H2DLocalExpansion):
order_values = [10, 12]
elif isinstance(knl, YukawaKernel):
extra_kwargs["lam"] = 2
dtype = np.complex128
if knl.dim == 3:
order_values = [1, 2]
elif knl.dim == 2 and issubclass(local_expn_class, Y2DLocalExpansion):
order_values = [10, 12]
from functools import partial
for order in order_values:
out_kernels = [knl]
from sumpy.fmm import SumpyExpansionWranglerCodeContainer
wcc = SumpyExpansionWranglerCodeContainer(
ctx, partial(mpole_expn_class, knl),
partial(local_expn_class, knl), out_kernels)
wrangler = wcc.get_wrangler(
queue,
tree,
dtype,
fmm_level_to_order=lambda kernel, kernel_args, tree, lev: order,
kernel_extra_kwargs=extra_kwargs)
from boxtree.fmm import drive_fmm
pot, = drive_fmm(trav, wrangler, weights)
from sumpy import P2P
p2p = P2P(ctx, out_kernels, exclude_self=False)
evt, (ref_pot, ) = p2p(queue, targets, sources, (weights, ),
**extra_kwargs)
pot = pot.get()
ref_pot = ref_pot.get()
rel_err = la.norm(pot - ref_pot, np.inf) / la.norm(ref_pot, np.inf)
logger.info("order %d -> relative l2 error: %g" % (order, rel_err))
pconv_verifier.add_data_point(order, rel_err)
print(pconv_verifier)
pconv_verifier()
def test_sumpy_fmm(ctx_getter, knl, local_expn_class, mpole_expn_class):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 1000
ntargets = 300
dtype = np.float64
from boxtree.tools import (
make_normal_particle_array as p_normal)
sources = p_normal(queue, nsources, knl.dim, dtype, seed=15)
if 1:
offset = np.zeros(knl.dim)
offset[0] = 0.1
targets = (
p_normal(queue, ntargets, knl.dim, dtype, seed=18)
+ offset)
del offset
else:
from sumpy.visualization import FieldPlotter
fp = FieldPlotter(np.array([0.5, 0]), extent=3, npoints=200)
from pytools.obj_array import make_obj_array
targets = make_obj_array(
[fp.points[i] for i in range(knl.dim)])
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources, targets=targets,
max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
# {{{ plot tree
if 0:
host_tree = tree.get()
host_trav = trav.get()
if 1:
print("src_box", host_tree.find_box_nr_for_source(403))
print("tgt_box", host_tree.find_box_nr_for_target(28))
print(list(host_trav.target_or_target_parent_boxes).index(37))
print(host_trav.get_box_list("sep_bigger", 22))
from boxtree.visualization import TreePlotter
plotter = TreePlotter(host_tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
plotter.draw_box_numbers()
import matplotlib.pyplot as pt
pt.show()
# }}}
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(ctx, seed=44)
weights = rng.uniform(queue, nsources, dtype=np.float64)
logger.info("computing direct (reference) result")
from pytools.convergence import PConvergenceVerifier
pconv_verifier = PConvergenceVerifier()
extra_kwargs = {}
dtype = np.float64
order_values = [1, 2, 3]
if isinstance(knl, HelmholtzKernel):
extra_kwargs["k"] = 0.05
dtype = np.complex128
if knl.dim == 3:
order_values = [1, 2]
elif knl.dim == 2 and issubclass(local_expn_class, H2DLocalExpansion):
order_values = [10, 12]
elif isinstance(knl, YukawaKernel):
extra_kwargs["lam"] = 2
dtype = np.complex128
if knl.dim == 3:
order_values = [1, 2]
elif knl.dim == 2 and issubclass(local_expn_class, Y2DLocalExpansion):
order_values = [10, 12]
from functools import partial
for order in order_values:
out_kernels = [knl]
from sumpy.fmm import SumpyExpansionWranglerCodeContainer
wcc = SumpyExpansionWranglerCodeContainer(
ctx,
partial(mpole_expn_class, knl),
partial(local_expn_class, knl),
out_kernels)
wrangler = wcc.get_wrangler(queue, tree, dtype,
fmm_level_to_order=lambda kernel, kernel_args, tree, lev: order,
kernel_extra_kwargs=extra_kwargs)
from boxtree.fmm import drive_fmm
pot, = drive_fmm(trav, wrangler, weights)
from sumpy import P2P
p2p = P2P(ctx, out_kernels, exclude_self=False)
evt, (ref_pot,) = p2p(queue, targets, sources, (weights,),
**extra_kwargs)
pot = pot.get()
ref_pot = ref_pot.get()
rel_err = la.norm(pot - ref_pot, np.inf) / la.norm(ref_pot, np.inf)
logger.info("order %d -> relative l2 error: %g" % (order, rel_err))
pconv_verifier.add_data_point(order, rel_err)
print(pconv_verifier)
pconv_verifier()
def test_estimate_calibration_params(ctx_factory):
from boxtree.pyfmmlib_integration import FMMLibExpansionWrangler
nsources_list = [1000, 2000, 3000, 4000]
ntargets_list = [1000, 2000, 3000, 4000]
dims = 3
dtype = np.float64
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
traversals = []
traversals_dev = []
level_to_orders = []
timing_results = []
def fmm_level_to_nterms(tree, ilevel):
return 10
for nsources, ntargets in zip(nsources_list, ntargets_list):
# {{{ Generate sources, targets and target_radii
from boxtree.tools import make_normal_particle_array as p_normal
sources = p_normal(queue, nsources, dims, dtype, seed=15)
targets = p_normal(queue, ntargets, dims, dtype, seed=18)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=22)
target_radii = rng.uniform(
queue, ntargets, a=0, b=0.05, dtype=dtype
).get()
# }}}
# {{{ Generate tree and traversal
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)
traversals.append(trav)
traversals_dev.append(trav_dev)
# }}}
wrangler = FMMLibExpansionWrangler(trav.tree, 0, fmm_level_to_nterms)
level_to_orders.append(wrangler.level_nterms)
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)
timing_results.append(timing_data)
if SUPPORTS_PROCESS_TIME:
time_field_name = "process_elapsed"
else:
time_field_name = "wall_elapsed"
def test_params_sanity(test_params):
param_names = ["c_p2m", "c_m2m", "c_p2p", "c_m2l", "c_m2p", "c_p2l", "c_l2l",
"c_l2p"]
for name in param_names:
assert isinstance(test_params[name], np.float64)
def test_params_equal(test_params1, test_params2):
param_names = ["c_p2m", "c_m2m", "c_p2p", "c_m2l", "c_m2p", "c_p2l", "c_l2l",
"c_l2p"]
for name in param_names:
assert test_params1[name] == test_params2[name]
python_cost_model = _PythonFMMCostModel(make_pde_aware_translation_cost_model)
python_model_results = []
for icase in range(len(traversals)-1):
traversal = traversals[icase]
level_to_order = level_to_orders[icase]
python_model_results.append(python_cost_model.cost_per_stage(
queue, traversal, level_to_order,
_PythonFMMCostModel.get_unit_calibration_params(),
))
python_params = python_cost_model.estimate_calibration_params(
python_model_results, timing_results[:-1], time_field_name=time_field_name
)
test_params_sanity(python_params)
cl_cost_model = FMMCostModel(make_pde_aware_translation_cost_model)
cl_model_results = []
for icase in range(len(traversals_dev)-1):
traversal = traversals_dev[icase]
level_to_order = level_to_orders[icase]
cl_model_results.append(cl_cost_model.cost_per_stage(
queue, traversal, level_to_order,
FMMCostModel.get_unit_calibration_params(),
))
cl_params = cl_cost_model.estimate_calibration_params(
cl_model_results, timing_results[:-1], time_field_name=time_field_name
)
test_params_sanity(cl_params)
if SUPPORTS_PROCESS_TIME:
test_params_equal(cl_params, python_params)
def test_sumpy_fmm_exclude_self(ctx_getter):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 500
dtype = np.float64
from boxtree.tools import (
make_normal_particle_array as p_normal)
knl = LaplaceKernel(2)
local_expn_class = VolumeTaylorLocalExpansion
mpole_expn_class = VolumeTaylorMultipoleExpansion
order = 10
sources = p_normal(queue, nsources, knl.dim, dtype, seed=15)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources,
max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(ctx)
weights = rng.uniform(queue, nsources, dtype=np.float64)
target_to_source = np.arange(tree.ntargets, dtype=np.int32)
self_extra_kwargs = {"target_to_source": target_to_source}
out_kernels = [knl]
from functools import partial
from sumpy.fmm import SumpyExpansionWranglerCodeContainer
wcc = SumpyExpansionWranglerCodeContainer(
ctx,
partial(mpole_expn_class, knl),
partial(local_expn_class, knl),
out_kernels,
exclude_self=True)
wrangler = wcc.get_wrangler(queue, tree, dtype,
fmm_level_to_order=lambda kernel, kernel_args, tree, lev: order,
self_extra_kwargs=self_extra_kwargs)
from boxtree.fmm import drive_fmm
pot, = drive_fmm(trav, wrangler, weights)
from sumpy import P2P
p2p = P2P(ctx, out_kernels, exclude_self=True)
evt, (ref_pot,) = p2p(queue, sources, sources, (weights,),
**self_extra_kwargs)
pot = pot.get()
ref_pot = ref_pot.get()
rel_err = la.norm(pot - ref_pot) / la.norm(ref_pot)
logger.info("order %d -> relative l2 error: %g" % (order, rel_err))
assert np.isclose(rel_err, 0, atol=1e-7)
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"]
)
