def test_plot_traversal(ctx_factory, well_sep_is_n_away=1, plot=False):
pytest.importorskip("matplotlib")
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
#for dims in [2, 3]:
for dims in [2]:
nparticles = 10**4
dtype = np.float64
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=15)
from pytools.obj_array import make_obj_array
particles = make_obj_array([
rng.normal(queue, nparticles, dtype=dtype)
for i in range(dims)])
# if do_plot:
# pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx, well_sep_is_n_away=well_sep_is_n_away)
trav, _ = tg(queue, tree)
tree = tree.get(queue=queue)
trav = trav.get(queue=queue)
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black")
#plotter.draw_box_numbers()
plotter.set_bounding_box()
from random import randrange, seed # noqa
seed(7)
from boxtree.visualization import draw_box_lists
#draw_box_lists(randrange(tree.nboxes))
if well_sep_is_n_away == 1:
draw_box_lists(plotter, trav, 380)
elif well_sep_is_n_away == 2:
draw_box_lists(plotter, trav, 320)
#plotter.draw_box_numbers()
if plot:
import matplotlib.pyplot as pt
pt.gca().set_xticks([])
pt.gca().set_yticks([])
pt.show()
def plot_traversal(ctx_getter, do_plot=False, well_sep_is_n_away=1):
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
#for dims in [2, 3]:
for dims in [2]:
nparticles = 10**4
dtype = np.float64
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=15)
from pytools.obj_array import make_obj_array
particles = make_obj_array([
rng.normal(queue, nparticles, dtype=dtype)
for i in range(dims)])
# if do_plot:
# pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, particles, max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx, well_sep_is_n_away=well_sep_is_n_away)
trav, _ = tg(queue, tree)
tree = tree.get(queue=queue)
trav = trav.get(queue=queue)
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black")
#plotter.draw_box_numbers()
plotter.set_bounding_box()
from random import randrange, seed # noqa
seed(7)
from boxtree.visualization import draw_box_lists
#draw_box_lists(randrange(tree.nboxes))
draw_box_lists(plotter, trav, 320)
#plotter.draw_box_numbers()
import matplotlib.pyplot as pt
pt.show()
def test_tree_connectivity(ctx_getter, dims, sources_are_targets):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
dtype = np.float64
sources = make_normal_particle_array(queue, 1 * 10**5, dims, dtype)
if sources_are_targets:
targets = None
else:
targets = make_normal_particle_array(queue, 2 * 10**5, dims, dtype)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue,
sources,
max_particles_in_box=30,
targets=targets,
debug=True)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav, _ = tg(queue, tree, debug=True)
tree = tree.get(queue=queue)
trav = trav.get(queue=queue)
levels = tree.box_levels
parents = tree.box_parent_ids.T
children = tree.box_child_ids.T
centers = tree.box_centers.T
# {{{ parent and child relations, levels match up
for ibox in range(1, tree.nboxes):
# /!\ Not testing box 0, has no parents
parent = parents[ibox]
assert levels[parent] + 1 == levels[ibox]
assert ibox in children[parent], ibox
# }}}
if 0:
import matplotlib.pyplot as pt
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black")
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.show()
# {{{ neighbor_source_boxes (list 1) consists of source boxes
for itgt_box, ibox in enumerate(trav.target_boxes):
start, end = trav.neighbor_source_boxes_starts[itgt_box:itgt_box + 2]
nbl = trav.neighbor_source_boxes_lists[start:end]
if sources_are_targets:
assert ibox in nbl
for jbox in nbl:
assert (0 == children[jbox]).all(), (ibox, jbox, children[jbox])
logger.info("list 1 consists of source boxes")
# }}}
# {{{ separated siblings (list 2) are actually separated
for itgt_box, tgt_ibox in enumerate(trav.target_or_target_parent_boxes):
start, end = trav.sep_siblings_starts[itgt_box:itgt_box + 2]
seps = trav.sep_siblings_lists[start:end]
assert (levels[seps] == levels[tgt_ibox]).all()
# three-ish box radii (half of size)
mindist = 2.5 * 0.5 * 2**-int(levels[tgt_ibox]) * tree.root_extent
icenter = centers[tgt_ibox]
for jbox in seps:
dist = la.norm(centers[jbox] - icenter)
assert dist > mindist, (dist, mindist)
logger.info("separated siblings (list 2) are actually separated")
# }}}
if sources_are_targets:
# {{{ sep_{smaller,bigger} are duals of each other
assert (trav.target_or_target_parent_boxes == np.arange(
tree.nboxes)).all()
# {{{ list 4 <= list 3
for itarget_box, ibox in enumerate(trav.target_boxes):
for ssn in trav.sep_smaller_by_level:
start, end = ssn.starts[itarget_box:itarget_box + 2]
for jbox in ssn.lists[start:end]:
rstart, rend = trav.sep_bigger_starts[jbox:jbox + 2]
assert ibox in trav.sep_bigger_lists[rstart:rend], (ibox,
jbox)
# }}}
# {{{ list 4 <= list 3
box_to_target_box_index = np.empty(tree.nboxes, tree.box_id_dtype)
box_to_target_box_index.fill(-1)
box_to_target_box_index[trav.target_boxes] = np.arange(
len(trav.target_boxes), dtype=tree.box_id_dtype)
assert (trav.source_boxes == trav.target_boxes).all()
assert (trav.target_or_target_parent_boxes == np.arange(
tree.nboxes, dtype=tree.box_id_dtype)).all()
for ibox in range(tree.nboxes):
start, end = trav.sep_bigger_starts[ibox:ibox + 2]
for jbox in trav.sep_bigger_lists[start:end]:
# In principle, entries of sep_bigger_lists are
# source boxes. In this special case, source and target boxes
# are the same thing (i.e. leaves--see assertion above), so we
# may treat them as targets anyhow.
jtgt_box = box_to_target_box_index[jbox]
assert jtgt_box != -1
good = False
for ssn in trav.sep_smaller_by_level:
rstart, rend = ssn.starts[jtgt_box:jtgt_box + 2]
good = good or ibox in ssn.lists[rstart:rend]
if not good:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
plotter.draw_box(ibox, facecolor='green', alpha=0.5)
plotter.draw_box(jbox, facecolor='red', alpha=0.5)
import matplotlib.pyplot as pt
pt.gca().set_aspect("equal")
pt.show()
# This assertion failing means that ibox's list 4 contains a box
# 'jbox' whose list 3 does not contain ibox.
assert good, (ibox, jbox)
# }}}
logger.info("list 3, 4 are duals")
# }}}
# {{{ sep_smaller satisfies relative level assumption
for itarget_box, ibox in enumerate(trav.target_boxes):
for ssn in trav.sep_smaller_by_level:
start, end = ssn.starts[itarget_box:itarget_box + 2]
for jbox in ssn.lists[start:end]:
assert levels[ibox] < levels[jbox]
logger.info("list 3 satisfies relative level assumption")
# }}}
# {{{ sep_bigger satisfies relative level assumption
for itgt_box, tgt_ibox in enumerate(trav.target_or_target_parent_boxes):
start, end = trav.sep_bigger_starts[itgt_box:itgt_box + 2]
for jbox in trav.sep_bigger_lists[start:end]:
assert levels[tgt_ibox] > levels[jbox]
logger.info("list 4 satisfies relative level assumption")
# }}}
# {{{ level_start_*_box_nrs lists make sense
for name, ref_array in [("level_start_source_box_nrs", trav.source_boxes),
("level_start_source_parent_box_nrs",
trav.source_parent_boxes),
("level_start_target_box_nrs", trav.target_boxes),
("level_start_target_or_target_parent_box_nrs",
trav.target_or_target_parent_boxes)]:
level_starts = getattr(trav, name)
for lev in range(tree.nlevels):
start, stop = level_starts[lev:lev + 2]
box_nrs = ref_array[start:stop]
assert (tree.box_levels[box_nrs] == lev).all(), name
def plot_traversal(ctx_getter, do_plot=False):
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
#for dims in [2, 3]:
for dims in [2]:
nparticles = 10**4
dtype = np.float64
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=15)
from pytools.obj_array import make_obj_array
particles = make_obj_array(
[rng.normal(queue, nparticles, dtype=dtype) for i in range(dims)])
# if do_plot:
# pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree = tb(queue, particles, max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav = tg(queue, tree).get()
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black")
#plotter.draw_box_numbers()
plotter.set_bounding_box()
from random import randrange, seed
seed(7)
# {{{ generic box drawing helper
def draw_some_box_lists(starts, lists, key_to_box=None, count=5):
actual_count = 0
while actual_count < count:
if key_to_box is not None:
key = randrange(len(key_to_box))
ibox = key_to_box[key]
else:
key = ibox = randrange(tree.nboxes)
start, end = starts[key:key + 2]
if start == end:
continue
#print ibox, start, end, lists[start:end]
for jbox in lists[start:end]:
plotter.draw_box(jbox, facecolor='yellow')
plotter.draw_box(ibox, facecolor='red')
actual_count += 1
# }}}
if 0:
# colleagues
draw_some_box_lists(trav.colleagues_starts, trav.colleagues_lists)
elif 0:
# near neighbors ("list 1")
draw_some_box_lists(trav.neighbor_leaves_starts,
trav.neighbor_leaves_lists,
key_to_box=trav.source_boxes)
elif 0:
# well-separated siblings (list 2)
draw_some_box_lists(trav.sep_siblings_starts,
trav.sep_siblings_lists)
elif 1:
# separated smaller (list 3)
draw_some_box_lists(trav.sep_smaller_starts,
trav.sep_smaller_lists,
key_to_box=trav.source_boxes)
elif 1:
# separated bigger (list 4)
draw_some_box_lists(trav.sep_bigger_starts, trav.sep_bigger_lists)
import matplotlib.pyplot as pt
pt.show()
def plot(self, **kwargs):
from boxtree.visualization import TreePlotter
plotter = TreePlotter(self)
plotter.draw_tree(**kwargs)
plotter.set_bounding_box()
def main():
print("*************************")
print("* Setting up...")
print("*************************")
dim = 3
# download precomputation results for the 3D Laplace kernel
download_table = True
table_filename = "nft_laplace3d.hdf5"
logger.info("Using table cache: " + table_filename)
q_order = 7 # quadrature order
n_levels = 5
use_multilevel_table = False
adaptive_mesh = False
n_refinement_loops = 100
refined_n_cells = 5e5
rratio_top = 0.2
rratio_bot = 0.5
dtype = np.float64
m_order = 10 # multipole order
force_direct_evaluation = False
logger.info("Multipole order = " + str(m_order))
logger.info("Quad order = " + str(q_order))
logger.info("N_levels = " + str(n_levels))
# a solution that is nearly zero at the boundary
# exp(-40) = 4.25e-18
alpha = 80
x = pmbl.var("x")
y = pmbl.var("y")
z = pmbl.var("z")
expp = pmbl.var("exp")
norm2 = x**2 + y**2 + z**2
source_expr = -(4 * alpha**2 * norm2 - 6 * alpha) * expp(-alpha * norm2)
solu_expr = expp(-alpha * norm2)
logger.info("Source expr: " + str(source_expr))
logger.info("Solu expr: " + str(solu_expr))
# bounding box
a = -0.5
b = 0.5
root_table_source_extent = 2
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
# logger.info("Summary of params: " + get_param_summary())
source_eval = Eval(dim, source_expr, [x, y, z])
# {{{ generate quad points
import volumential.meshgen as mg
# Show meshgen info
mg.greet()
mesh = mg.MeshGen3D(q_order, n_levels, a, b, queue=queue)
if not adaptive_mesh:
mesh.print_info()
q_points = mesh.get_q_points()
q_weights = mesh.get_q_weights()
else:
iloop = -1
while mesh.n_active_cells() < refined_n_cells:
iloop += 1
cell_centers = mesh.get_cell_centers()
cell_measures = mesh.get_cell_measures()
density_vals = source_eval(
queue,
np.array([[center[d] for center in cell_centers]
for d in range(dim)]))
crtr = np.abs(cell_measures * density_vals)
mesh.update_mesh(crtr, rratio_top, rratio_bot)
if iloop > n_refinement_loops:
print("Max number of refinement loops reached.")
break
mesh.print_info()
q_points = mesh.get_q_points()
q_weights = mesh.get_q_weights()
if 1:
try:
mesh.generate_gmsh("box_grid.msh")
except Exception as e:
print(e)
pass
legacy_msh_file = True
if legacy_msh_file:
import os
os.system("gmsh box_grid.msh convert_grid -")
assert len(q_points) == len(q_weights)
assert q_points.shape[1] == dim
q_points = np.ascontiguousarray(np.transpose(q_points))
from pytools.obj_array import make_obj_array
q_points = make_obj_array(
[cl.array.to_device(queue, q_points[i]) for i in range(dim)])
q_weights = cl.array.to_device(queue, q_weights)
# }}}
# {{{ discretize the source field
logger.info("discretizing source field")
source_vals = cl.array.to_device(
queue,
source_eval(queue, np.array([coords.get() for coords in q_points])))
# particle_weigt = source_val * q_weight
# }}} End discretize the source field
# {{{ build tree and traversals
from boxtree.tools import AXIS_NAMES
axis_names = AXIS_NAMES[:dim]
from pytools import single_valued
coord_dtype = single_valued(coord.dtype for coord in q_points)
from boxtree.bounding_box import make_bounding_box_dtype
bbox_type, _ = make_bounding_box_dtype(ctx.devices[0], dim, coord_dtype)
bbox = np.empty(1, bbox_type)
for ax in axis_names:
bbox["min_" + ax] = a
bbox["max_" + ax] = b
# tune max_particles_in_box to reconstruct the mesh
# TODO: use points from FieldPlotter are used as target points for better
# visuals
print("building tree")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(
queue,
particles=q_points,
targets=q_points,
bbox=bbox,
max_particles_in_box=q_order**3 * 8 - 1,
kind="adaptive-level-restricted",
)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav, _ = tg(queue, tree)
# }}} End build tree and traversals
# {{{ build near field potential table
from volumential.table_manager import NearFieldInteractionTableManager
import os
if download_table and (not os.path.isfile(table_filename)):
import json
with open("table_urls.json", 'r') as fp:
urls = json.load(fp)
print("Downloading table from %s" % urls['Laplace3D'])
import subprocess
subprocess.call(["wget", "-q", urls['Laplace3D'], table_filename])
tm = NearFieldInteractionTableManager(table_filename,
root_extent=root_table_source_extent,
queue=queue)
if use_multilevel_table:
logger.info("Using multilevel tables")
assert (abs(
int((b - a) / root_table_source_extent) *
root_table_source_extent - (b - a)) < 1e-15)
nftable = []
for lev in range(0, tree.nlevels + 1):
print("Getting table at level", lev)
tb, _ = tm.get_table(
dim,
"Laplace",
q_order,
source_box_level=lev,
compute_method="DrosteSum",
queue=queue,
n_brick_quad_points=120,
adaptive_level=False,
use_symmetry=True,
alpha=0,
n_levels=1,
)
nftable.append(tb)
print("Using table list of length", len(nftable))
else:
logger.info("Using single level table")
force_recompute = False
# 15 levels are sufficient (the inner most brick is 1e-15**3 in volume)
nftable, _ = tm.get_table(
dim,
"Laplace",
q_order,
force_recompute=force_recompute,
compute_method="DrosteSum",
queue=queue,
n_brick_quad_points=120,
adaptive_level=False,
use_symmetry=True,
alpha=0,
n_levels=1,
)
# }}} End build near field potential table
# {{{ sumpy expansion for laplace kernel
from sumpy.expansion import DefaultExpansionFactory
from sumpy.kernel import LaplaceKernel
knl = LaplaceKernel(dim)
out_kernels = [knl]
expn_factory = DefaultExpansionFactory()
local_expn_class = expn_factory.get_local_expansion_class(knl)
mpole_expn_class = expn_factory.get_multipole_expansion_class(knl)
exclude_self = True
from volumential.expansion_wrangler_fpnd import (
FPNDExpansionWrangler, FPNDExpansionWranglerCodeContainer)
wcc = FPNDExpansionWranglerCodeContainer(
ctx,
partial(mpole_expn_class, knl),
partial(local_expn_class, knl),
out_kernels,
exclude_self=exclude_self,
)
if exclude_self:
target_to_source = np.arange(tree.ntargets, dtype=np.int32)
self_extra_kwargs = {"target_to_source": target_to_source}
else:
self_extra_kwargs = {}
wrangler = FPNDExpansionWrangler(
code_container=wcc,
queue=queue,
tree=tree,
near_field_table=nftable,
dtype=dtype,
fmm_level_to_order=lambda kernel, kernel_args, tree, lev: m_order,
quad_order=q_order,
self_extra_kwargs=self_extra_kwargs,
)
# }}} End sumpy expansion for laplace kernel
print("*************************")
print("* Performing FMM ...")
print("*************************")
# {{{ conduct fmm computation
from volumential.volume_fmm import drive_volume_fmm
import time
queue.finish()
t0 = time.time()
pot, = drive_volume_fmm(trav,
wrangler,
source_vals * q_weights,
source_vals,
direct_evaluation=force_direct_evaluation,
list1_only=False)
t1 = time.time()
print("Finished in %.2f seconds." % (t1 - t0))
print("(%e points per second)" % (len(q_weights) / (t1 - t0)))
# }}} End conduct fmm computation
print("*************************")
print("* Postprocessing ...")
print("*************************")
# {{{ postprocess and plot
# print(pot)
solu_eval = Eval(dim, solu_expr, [x, y, z])
# x = q_points[0].get()
# y = q_points[1].get()
# z = q_points[2].get()
test_x = np.array([0.0])
test_y = np.array([0.0])
test_z = np.array([0.0])
test_nodes = make_obj_array(
# get() first for CL compatibility issues
[
cl.array.to_device(queue, test_x),
cl.array.to_device(queue, test_y),
cl.array.to_device(queue, test_z),
])
from volumential.volume_fmm import interpolate_volume_potential
ze = solu_eval(queue, np.array([test_x, test_y, test_z]))
zs = interpolate_volume_potential(test_nodes, trav, wrangler, pot).get()
print_error = True
if print_error:
err = np.max(np.abs(ze - zs))
print("Error =", err)
# Boxtree
if 0:
import matplotlib.pyplot as plt
if dim == 2:
plt.plot(q_points[0].get(), q_points[1].get(), ".")
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree.get(queue=queue))
plotter.draw_tree(fill=False, edgecolor="black")
# plotter.draw_box_numbers()
plotter.set_bounding_box()
plt.gca().set_aspect("equal")
plt.draw()
plt.show()
# plt.savefig("tree.png")
# Direct p2p
if 0:
print("Performing P2P")
pot_direct, = drive_volume_fmm(trav,
wrangler,
source_vals * q_weights,
source_vals,
direct_evaluation=True)
zds = pot_direct.get()
zs = pot.get()
print("P2P-FMM diff =", np.max(np.abs(zs - zds)))
print("P2P Error =", np.max(np.abs(ze - zds)))
# Write vtk
if 0:
from meshmode.mesh.io import read_gmsh
modemesh = read_gmsh("box_grid.msh", force_ambient_dim=None)
from meshmode.discretization.poly_element import (
LegendreGaussLobattoTensorProductGroupFactory, )
from meshmode.array_context import PyOpenCLArrayContext
from meshmode.discretization import Discretization
actx = PyOpenCLArrayContext(queue)
box_discr = Discretization(
actx, modemesh,
LegendreGaussLobattoTensorProductGroupFactory(q_order))
box_nodes_x = box_discr.nodes()[0].with_queue(queue).get()
box_nodes_y = box_discr.nodes()[1].with_queue(queue).get()
box_nodes_z = box_discr.nodes()[2].with_queue(queue).get()
box_nodes = make_obj_array(
# get() first for CL compatibility issues
[
cl.array.to_device(queue, box_nodes_x),
cl.array.to_device(queue, box_nodes_y),
cl.array.to_device(queue, box_nodes_z),
])
visual_order = 1
from meshmode.discretization.visualization import make_visualizer
vis = make_visualizer(queue, box_discr, visual_order)
from volumential.volume_fmm import interpolate_volume_potential
volume_potential = interpolate_volume_potential(
box_nodes, trav, wrangler, pot)
# qx = q_points[0].get()
# qy = q_points[1].get()
# qz = q_points[2].get()
exact_solution = cl.array.to_device(
queue,
solu_eval(queue, np.array([box_nodes_x, box_nodes_y,
box_nodes_z])))
# clean up the mess
def clean_file(filename):
import os
try:
os.remove(filename)
except OSError:
pass
vtu_filename = "laplace3d.vtu"
clean_file(vtu_filename)
vis.write_vtk_file(
vtu_filename,
[
("VolPot", volume_potential),
# ("SrcDensity", source_density),
("ExactSol", exact_solution),
("Error", volume_potential - exact_solution),
],
)
print("Written file " + vtu_filename)
def plot_traversal(ctx_getter, do_plot=False):
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
#for dims in [2, 3]:
for dims in [2]:
nparticles = 10**4
dtype = np.float64
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=15)
from pytools.obj_array import make_obj_array
particles = make_obj_array([
rng.normal(queue, nparticles, dtype=dtype)
for i in range(dims)])
# if do_plot:
# pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree = tb(queue, particles, max_particles_in_box=30, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav = tg(queue, tree).get()
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black")
#plotter.draw_box_numbers()
plotter.set_bounding_box()
from random import randrange, seed
seed(7)
# {{{ generic box drawing helper
def draw_some_box_lists(starts, lists, key_to_box=None,
count=5):
actual_count = 0
while actual_count < count:
if key_to_box is not None:
key = randrange(len(key_to_box))
ibox = key_to_box[key]
else:
key = ibox = randrange(tree.nboxes)
start, end = starts[key:key+2]
if start == end:
continue
#print ibox, start, end, lists[start:end]
for jbox in lists[start:end]:
plotter.draw_box(jbox, facecolor='yellow')
plotter.draw_box(ibox, facecolor='red')
actual_count += 1
# }}}
if 0:
# colleagues
draw_some_box_lists(
trav.colleagues_starts,
trav.colleagues_lists)
elif 0:
# near neighbors ("list 1")
draw_some_box_lists(
trav.neighbor_leaves_starts,
trav.neighbor_leaves_lists,
key_to_box=trav.source_boxes)
elif 0:
# well-separated siblings (list 2)
draw_some_box_lists(
trav.sep_siblings_starts,
trav.sep_siblings_lists)
elif 1:
# separated smaller (list 3)
draw_some_box_lists(
trav.sep_smaller_starts,
trav.sep_smaller_lists,
key_to_box=trav.source_boxes)
elif 1:
# separated bigger (list 4)
draw_some_box_lists(
trav.sep_bigger_starts,
trav.sep_bigger_lists)
import matplotlib.pyplot as pt
pt.show()
def plot(self, **kwargs):
from boxtree.visualization import TreePlotter
plotter = TreePlotter(self)
plotter.draw_tree(**kwargs)
plotter.set_bounding_box()
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_extent_tree(ctx_getter, dims, extent_norm, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 100000
ntargets = 200000
dtype = np.float64
npoint_sources_per_source = 16
sources = make_normal_particle_array(queue, nsources, dims, dtype,
seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype,
seed=19)
refine_weights = cl.array.zeros(queue, nsources+ntargets, np.int32)
refine_weights[:nsources] = 1
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=13)
source_radii = 2**rng.uniform(queue, nsources, dtype=dtype,
a=-10, b=0)
target_radii = 2**rng.uniform(queue, ntargets, dtype=dtype,
a=-10, b=0)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
dev_tree, _ = tb(queue, sources, targets=targets,
source_radii=source_radii,
target_radii=target_radii,
extent_norm=extent_norm,
refine_weights=refine_weights,
max_leaf_refine_weight=20,
#max_particles_in_box=10,
# Set artificially small, to exercise the reallocation code.
nboxes_guess=10,
debug=True,
stick_out_factor=0)
logger.info("transfer tree, check orderings")
tree = dev_tree.get(queue=queue)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(sources[0].get(), sources[1].get(), "rx")
pt.plot(targets[0].get(), targets[1].get(), "g+")
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.gca().set_aspect("equal", "datalim")
pt.show()
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
sorted_source_radii = tree.source_radii
sorted_target_radii = tree.target_radii
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
unsorted_source_radii = source_radii.get()
unsorted_target_radii = target_radii.get()
assert (sorted_sources
== unsorted_sources[:, tree.user_source_ids]).all()
assert (sorted_source_radii
== unsorted_source_radii[tree.user_source_ids]).all()
# {{{ test box structure, stick-out criterion
logger.info("test box structure, stick-out criterion")
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets, dtype=np.intp)
if ntargets:
assert (sorted_targets
== unsorted_targets[:, user_target_ids]).all()
assert (sorted_target_radii
== unsorted_target_radii[user_target_ids]).all()
all_good_so_far = True
# {{{ check sources, targets
assert np.sum(tree.box_source_counts_nonchild) == nsources
assert np.sum(tree.box_target_counts_nonchild) == ntargets
for ibox in range(tree.nboxes):
kid_sum = sum(
tree.box_target_counts_cumul[ichild_box]
for ichild_box in tree.box_child_ids[:, ibox]
if ichild_box != 0)
assert (
tree.box_target_counts_cumul[ibox]
== (
tree.box_target_counts_nonchild[ibox]
+ kid_sum)), ibox
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low
>= tree.bounding_box[0] - 1e-12*tree.root_extent).all(), ibox
assert (extent_high
<= tree.bounding_box[1] + 1e-12*tree.root_extent).all(), ibox
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox]
+ np.sum(tree.box_target_counts_cumul[existing_children])
== tree.box_target_counts_cumul[ibox])
del existing_children
del box_children
for ibox in range(tree.nboxes):
lev = int(tree.box_levels[ibox])
box_radius = 0.5 * tree.root_extent / (1 << lev)
box_center = tree.box_centers[:, ibox]
extent_low = box_center - box_radius
extent_high = box_center + box_radius
stick_out_dist = tree.stick_out_factor * box_radius
radius_with_stickout = (1 + tree.stick_out_factor) * box_radius
for what, starts, counts, points, radii in [
("source", tree.box_source_starts, tree.box_source_counts_cumul,
sorted_sources, sorted_source_radii),
("target", tree.box_target_starts, tree.box_target_counts_cumul,
sorted_targets, sorted_target_radii),
]:
bstart = starts[ibox]
bslice = slice(bstart, bstart+counts[ibox])
check_particles = points[:, bslice]
check_radii = radii[bslice]
if extent_norm == "linf":
good = (
(check_particles + check_radii
< extent_high[:, np.newaxis] + stick_out_dist)
& # noqa: W504
(extent_low[:, np.newaxis] - stick_out_dist
<= check_particles - check_radii)
).all(axis=0)
elif extent_norm == "l2":
center_dists = np.sqrt(
np.sum(
(check_particles - box_center.reshape(-1, 1))**2,
axis=0))
good = (
(center_dists + check_radii)**2
< dims * radius_with_stickout**2)
else:
raise ValueError("unexpected value of extent_norm")
all_good_here = good.all()
if not all_good_here:
print("BAD BOX %s %d level %d"
% (what, ibox, tree.box_levels[ibox]))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_here
# }}}
assert all_good_so_far
# }}}
# {{{ create, link point sources
logger.info("creating point sources")
np.random.seed(20)
from pytools.obj_array import make_obj_array
point_sources = make_obj_array([
cl.array.to_device(queue,
unsorted_sources[i][:, np.newaxis]
+ unsorted_source_radii[:, np.newaxis]
* np.random.uniform(
-1, 1, size=(nsources, npoint_sources_per_source))
)
for i in range(dims)])
point_source_starts = cl.array.arange(queue,
0, (nsources+1)*npoint_sources_per_source, npoint_sources_per_source,
dtype=tree.particle_id_dtype)
from boxtree.tree import link_point_sources
dev_tree = link_point_sources(queue, dev_tree,
point_source_starts, point_sources,
debug=True)
def test_source_target_tree(ctx_getter, dims, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
nsources = 2 * 10**5
ntargets = 3 * 10**5
dtype = np.float64
sources = make_normal_particle_array(queue, nsources, dims, dtype,
seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype,
seed=19)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(sources[0].get(), sources[1].get(), "rx")
pt.plot(targets[0].get(), targets[1].get(), "g+")
pt.show()
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue, sources, targets=targets,
max_particles_in_box=10, debug=True)
tree = tree.get(queue=queue)
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
assert (sorted_sources
== unsorted_sources[:, tree.user_source_ids]).all()
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets, dtype=np.intp)
assert (sorted_targets
== unsorted_targets[:, user_target_ids]).all()
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
tol = 1e-15
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low
>= tree.bounding_box[0] - 1e-12*tree.root_extent).all(), ibox
assert (extent_high
<= tree.bounding_box[1] + 1e-12*tree.root_extent).all(), ibox
src_start = tree.box_source_starts[ibox]
tgt_start = tree.box_target_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox]
+ np.sum(tree.box_target_counts_cumul[existing_children])
== tree.box_target_counts_cumul[ibox])
for what, particles in [
("sources", sorted_sources[:,
src_start:src_start+tree.box_source_counts_cumul[ibox]]),
("targets", sorted_targets[:,
tgt_start:tgt_start+tree.box_target_counts_cumul[ibox]]),
]:
good = (
(particles < extent_high[:, np.newaxis] + tol)
& (extent_low[:, np.newaxis] - tol <= particles)
).all(axis=0)
all_good_here = good.all()
if do_plot and not all_good_here:
pt.plot(
particles[0, np.where(~good)[0]],
particles[1, np.where(~good)[0]], "ro")
plotter.draw_box(ibox, edgecolor="red")
pt.show()
if not all_good_here:
print("BAD BOX %s %d" % (what, ibox))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_so_far
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
def run_build_test(builder,
queue,
dims,
dtype,
nparticles,
do_plot,
max_particles_in_box=None,
max_leaf_refine_weight=None,
refine_weights=None,
**kwargs):
dtype = np.dtype(dtype)
if dtype == np.float32:
tol = 1e-4
elif dtype == np.float64:
tol = 1e-12
else:
raise RuntimeError("unsupported dtype: %s" % dtype)
logger.info(75 * "-")
if max_particles_in_box is not None:
logger.info(
"%dD %s - %d particles - max %d per box - %s" %
(dims, dtype.type.__name__, nparticles, max_particles_in_box,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
else:
logger.info(
"%dD %s - %d particles - max leaf weight %d - %s" %
(dims, dtype.type.__name__, nparticles, max_leaf_refine_weight,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
logger.info(75 * "-")
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
queue.finish()
tree, _ = builder(queue,
particles,
max_particles_in_box=max_particles_in_box,
refine_weights=refine_weights,
max_leaf_refine_weight=max_leaf_refine_weight,
debug=True,
**kwargs)
tree = tree.get(queue=queue)
sorted_particles = np.array(list(tree.sources))
unsorted_particles = np.array([pi.get() for pi in particles])
assert (
sorted_particles == unsorted_particles[:, tree.user_source_ids]).all()
if refine_weights is not None:
refine_weights_reordered = refine_weights.get()[tree.user_source_ids]
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
from boxtree import box_flags_enum as bfe
scaled_tol = tol * tree.root_extent
for ibox in range(tree.nboxes):
# Empty boxes exist in non-pruned trees--which themselves are undocumented.
# These boxes will fail these tests.
if not (tree.box_flags[ibox] & bfe.HAS_OWN_SRCNTGTS):
continue
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >= tree.bounding_box[0] - scaled_tol).all(), (
ibox, extent_low, tree.bounding_box[0])
assert (extent_high <= tree.bounding_box[1] + scaled_tol).all(), (
ibox, extent_high, tree.bounding_box[1])
start = tree.box_source_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox] +
np.sum(tree.box_source_counts_cumul[existing_children]) ==
tree.box_source_counts_cumul[ibox])
box_particles = sorted_particles[:, start:start +
tree.box_source_counts_cumul[ibox]]
good = ((box_particles < extent_high[:, np.newaxis] + scaled_tol)
& (extent_low[:, np.newaxis] - scaled_tol <= box_particles))
all_good_here = good.all()
if do_plot and not all_good_here and all_good_so_far:
pt.plot(box_particles[0, np.where(~good)[1]],
box_particles[1, np.where(~good)[1]], "ro")
plotter.draw_box(ibox, edgecolor="red")
if not all_good_here:
print("BAD BOX", ibox)
if not (tree.box_flags[ibox] & bfe.HAS_CHILDREN):
# Check that leaf particle density is as promised.
nparticles_in_box = tree.box_source_counts_cumul[ibox]
if max_particles_in_box is not None:
if nparticles_in_box > max_particles_in_box:
print("too many particles ({0} > {1}); box {2}".format(
nparticles_in_box, max_particles_in_box, ibox))
all_good_here = False
else:
assert refine_weights is not None
box_weight = np.sum(
refine_weights_reordered[start:start + nparticles_in_box])
if box_weight > max_leaf_refine_weight:
print("refine weight exceeded ({0} > {1}); box {2}".format(
box_weight, max_leaf_refine_weight, ibox))
all_good_here = False
all_good_so_far = all_good_so_far and all_good_here
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
assert all_good_so_far
def test_extent_tree(ctx_factory, dims, extent_norm, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nsources = 100000
ntargets = 200000
dtype = np.float64
npoint_sources_per_source = 16
sources = make_normal_particle_array(queue, nsources, dims, dtype, seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype, seed=19)
refine_weights = cl.array.zeros(queue, nsources + ntargets, np.int32)
refine_weights[:nsources] = 1
from pyopencl.clrandom import PhiloxGenerator
rng = PhiloxGenerator(queue.context, seed=13)
source_radii = 2**rng.uniform(queue, nsources, dtype=dtype, a=-10, b=0)
target_radii = 2**rng.uniform(queue, ntargets, dtype=dtype, a=-10, b=0)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
dev_tree, _ = tb(
queue,
sources,
targets=targets,
source_radii=source_radii,
target_radii=target_radii,
extent_norm=extent_norm,
refine_weights=refine_weights,
max_leaf_refine_weight=20,
#max_particles_in_box=10,
# Set artificially small, to exercise the reallocation code.
nboxes_guess=10,
debug=True,
stick_out_factor=0)
logger.info("transfer tree, check orderings")
tree = dev_tree.get(queue=queue)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(sources[0].get(), sources[1].get(), "rx")
pt.plot(targets[0].get(), targets[1].get(), "g+")
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.gca().set_aspect("equal", "datalim")
pt.show()
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
sorted_source_radii = tree.source_radii
sorted_target_radii = tree.target_radii
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
unsorted_source_radii = source_radii.get()
unsorted_target_radii = target_radii.get()
assert (sorted_sources == unsorted_sources[:, tree.user_source_ids]).all()
assert (sorted_source_radii == unsorted_source_radii[tree.user_source_ids]
).all()
# {{{ test box structure, stick-out criterion
logger.info("test box structure, stick-out criterion")
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets,
dtype=np.intp)
if ntargets:
assert (sorted_targets == unsorted_targets[:, user_target_ids]).all()
assert (sorted_target_radii == unsorted_target_radii[user_target_ids]
).all()
all_good_so_far = True
# {{{ check sources, targets
assert np.sum(tree.box_source_counts_nonchild) == nsources
assert np.sum(tree.box_target_counts_nonchild) == ntargets
for ibox in range(tree.nboxes):
kid_sum = sum(tree.box_target_counts_cumul[ichild_box]
for ichild_box in tree.box_child_ids[:, ibox]
if ichild_box != 0)
assert (tree.box_target_counts_cumul[ibox] == (
tree.box_target_counts_nonchild[ibox] + kid_sum)), ibox
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >=
tree.bounding_box[0] - 1e-12 * tree.root_extent).all(), ibox
assert (extent_high <=
tree.bounding_box[1] + 1e-12 * tree.root_extent).all(), ibox
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox] +
np.sum(tree.box_source_counts_cumul[existing_children]) ==
tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox] +
np.sum(tree.box_target_counts_cumul[existing_children]) ==
tree.box_target_counts_cumul[ibox])
del existing_children
del box_children
for ibox in range(tree.nboxes):
lev = int(tree.box_levels[ibox])
box_radius = 0.5 * tree.root_extent / (1 << lev)
box_center = tree.box_centers[:, ibox]
extent_low = box_center - box_radius
extent_high = box_center + box_radius
stick_out_dist = tree.stick_out_factor * box_radius
radius_with_stickout = (1 + tree.stick_out_factor) * box_radius
for what, starts, counts, points, radii in [
("source", tree.box_source_starts, tree.box_source_counts_cumul,
sorted_sources, sorted_source_radii),
("target", tree.box_target_starts, tree.box_target_counts_cumul,
sorted_targets, sorted_target_radii),
]:
bstart = starts[ibox]
bslice = slice(bstart, bstart + counts[ibox])
check_particles = points[:, bslice]
check_radii = radii[bslice]
if extent_norm == "linf":
good = ((check_particles + check_radii <
extent_high[:, np.newaxis] + stick_out_dist)
& # noqa: W504
(extent_low[:, np.newaxis] - stick_out_dist <=
check_particles - check_radii)).all(axis=0)
elif extent_norm == "l2":
center_dists = np.sqrt(
np.sum((check_particles - box_center.reshape(-1, 1))**2,
axis=0))
good = ((center_dists + check_radii)**2 <
dims * radius_with_stickout**2)
else:
raise ValueError("unexpected value of extent_norm")
all_good_here = good.all()
if not all_good_here:
print("BAD BOX %s %d level %d" %
(what, ibox, tree.box_levels[ibox]))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_here
# }}}
assert all_good_so_far
# }}}
# {{{ create, link point sources
logger.info("creating point sources")
np.random.seed(20)
from pytools.obj_array import make_obj_array
point_sources = make_obj_array([
cl.array.to_device(
queue, unsorted_sources[i][:, np.newaxis] +
unsorted_source_radii[:, np.newaxis] * np.random.uniform(
-1, 1, size=(nsources, npoint_sources_per_source)))
for i in range(dims)
])
point_source_starts = cl.array.arange(queue,
0, (nsources + 1) *
npoint_sources_per_source,
npoint_sources_per_source,
dtype=tree.particle_id_dtype)
from boxtree.tree import link_point_sources
dev_tree = link_point_sources(queue,
dev_tree,
point_source_starts,
point_sources,
debug=True)
def test_source_target_tree(ctx_factory, dims, do_plot=False):
logging.basicConfig(level=logging.INFO)
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
nsources = 2 * 10**5
ntargets = 3 * 10**5
dtype = np.float64
sources = make_normal_particle_array(queue, nsources, dims, dtype, seed=12)
targets = make_normal_particle_array(queue, ntargets, dims, dtype, seed=19)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(sources[0].get(), sources[1].get(), "rx")
pt.plot(targets[0].get(), targets[1].get(), "g+")
pt.show()
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
queue.finish()
tree, _ = tb(queue,
sources,
targets=targets,
max_particles_in_box=10,
debug=True)
tree = tree.get(queue=queue)
sorted_sources = np.array(list(tree.sources))
sorted_targets = np.array(list(tree.targets))
unsorted_sources = np.array([pi.get() for pi in sources])
unsorted_targets = np.array([pi.get() for pi in targets])
assert (sorted_sources == unsorted_sources[:, tree.user_source_ids]).all()
user_target_ids = np.empty(tree.ntargets, dtype=np.intp)
user_target_ids[tree.sorted_target_ids] = np.arange(tree.ntargets,
dtype=np.intp)
assert (sorted_targets == unsorted_targets[:, user_target_ids]).all()
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
tol = 1e-15
for ibox in range(tree.nboxes):
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >=
tree.bounding_box[0] - 1e-12 * tree.root_extent).all(), ibox
assert (extent_high <=
tree.bounding_box[1] + 1e-12 * tree.root_extent).all(), ibox
src_start = tree.box_source_starts[ibox]
tgt_start = tree.box_target_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox] +
np.sum(tree.box_source_counts_cumul[existing_children]) ==
tree.box_source_counts_cumul[ibox])
assert (tree.box_target_counts_nonchild[ibox] +
np.sum(tree.box_target_counts_cumul[existing_children]) ==
tree.box_target_counts_cumul[ibox])
for what, particles in [
("sources", sorted_sources[:, src_start:src_start +
tree.box_source_counts_cumul[ibox]]),
("targets", sorted_targets[:, tgt_start:tgt_start +
tree.box_target_counts_cumul[ibox]]),
]:
good = ((particles < extent_high[:, np.newaxis] + tol)
&
(extent_low[:, np.newaxis] - tol <= particles)).all(axis=0)
all_good_here = good.all()
if do_plot and not all_good_here:
pt.plot(particles[0, np.where(~good)[0]],
particles[1, np.where(~good)[0]], "ro")
plotter.draw_box(ibox, edgecolor="red")
pt.show()
if not all_good_here:
print("BAD BOX %s %d" % (what, ibox))
all_good_so_far = all_good_so_far and all_good_here
assert all_good_so_far
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
def main():
print("*************************")
print("* Setting up...")
print("*************************")
dim = 2
# download precomputation results for the 2D Laplace kernel
download_table = True
table_filename = "nft_laplace2d.hdf5"
root_table_source_extent = 2
print("Using table cache:", table_filename)
q_order = 9 # quadrature order
n_levels = 6 # 2^(n_levels-1) subintervals in 1D
use_multilevel_table = False
adaptive_mesh = False
n_refinement_loops = 100
refined_n_cells = 2000
rratio_top = 0.2
rratio_bot = 0.5
dtype = np.float64
m_order = 20 # multipole order
force_direct_evaluation = False
print("Multipole order =", m_order)
alpha = 160
x = pmbl.var("x")
y = pmbl.var("y")
expp = pmbl.var("exp")
norm2 = x**2 + y**2
source_expr = -(4 * alpha**2 * norm2 - 4 * alpha) * expp(-alpha * norm2)
solu_expr = expp(-alpha * norm2)
logger.info("Source expr: " + str(source_expr))
logger.info("Solu expr: " + str(solu_expr))
# bounding box
a = -0.5
b = 0.5
root_table_source_extent = 2
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
source_eval = Eval(dim, source_expr, [x, y])
# {{{ generate quad points
import volumential.meshgen as mg
# Show meshgen info
mg.greet()
mesh = mg.MeshGen2D(q_order, n_levels, a, b, queue=queue)
if not adaptive_mesh:
mesh.print_info()
q_points = mesh.get_q_points()
q_weights = mesh.get_q_weights()
else:
iloop = -1
while mesh.n_active_cells() < refined_n_cells:
iloop += 1
crtr = np.abs(
source_eval(mesh.get_cell_centers) * mesh.get_cell_measures)
mesh.update_mesh(crtr, rratio_top, rratio_bot)
if iloop > n_refinement_loops:
print("Max number of refinement loops reached.")
break
mesh.print_info()
q_points = mesh.get_q_points()
q_weights = mesh.get_q_weights()
assert len(q_points) == len(q_weights)
assert q_points.shape[1] == dim
q_points = np.ascontiguousarray(np.transpose(q_points))
from pytools.obj_array import make_obj_array
q_points = make_obj_array(
[cl.array.to_device(queue, q_points[i]) for i in range(dim)])
q_weights = cl.array.to_device(queue, q_weights)
# q_radii = cl.array.to_device(queue, q_radii)
# }}}
# {{{ discretize the source field
source_vals = cl.array.to_device(
queue,
source_eval(queue, np.array([coords.get() for coords in q_points])))
# particle_weigt = source_val * q_weight
# }}} End discretize the source field
# {{{ build tree and traversals
from boxtree.tools import AXIS_NAMES
axis_names = AXIS_NAMES[:dim]
from pytools import single_valued
coord_dtype = single_valued(coord.dtype for coord in q_points)
from boxtree.bounding_box import make_bounding_box_dtype
bbox_type, _ = make_bounding_box_dtype(ctx.devices[0], dim, coord_dtype)
bbox = np.empty(1, bbox_type)
for ax in axis_names:
bbox["min_" + ax] = a
bbox["max_" + ax] = b
# tune max_particles_in_box to reconstruct the mesh
# TODO: use points from FieldPlotter are used as target points for better
# visuals
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(
queue,
particles=q_points,
targets=q_points,
bbox=bbox,
max_particles_in_box=q_order**2 * 4 - 1,
kind="adaptive-level-restricted",
)
bbox2 = np.array([[a, b], [a, b]])
tree2, _ = tb(
queue,
particles=q_points,
targets=q_points,
bbox=bbox2,
max_particles_in_box=q_order**2 * 4 - 1,
kind="adaptive-level-restricted",
)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav, _ = tg(queue, tree)
# }}} End build tree and traversals
# {{{ build near field potential table
from volumential.table_manager import NearFieldInteractionTableManager
import os
if download_table and (not os.path.isfile(table_filename)):
import json
with open("table_urls.json", 'r') as fp:
urls = json.load(fp)
print("Downloading table from %s" % urls['Laplace2D'])
import subprocess
subprocess.call(["wget", "-q", urls['Laplace2D'], table_filename])
tm = NearFieldInteractionTableManager(table_filename,
root_extent=root_table_source_extent,
queue=queue)
if use_multilevel_table:
assert (abs(
int((b - a) / root_table_source_extent) *
root_table_source_extent - (b - a)) < 1e-15)
nftable = []
for lev in range(0, tree.nlevels + 1):
print("Getting table at level", lev)
tb, _ = tm.get_table(
dim,
"Laplace",
q_order,
source_box_level=lev,
compute_method="DrosteSum",
queue=queue,
n_brick_quad_points=100,
adaptive_level=False,
use_symmetry=True,
alpha=0.1,
nlevels=15,
)
nftable.append(tb)
print("Using table list of length", len(nftable))
else:
nftable, _ = tm.get_table(
dim,
"Laplace",
q_order,
force_recompute=False,
compute_method="DrosteSum",
queue=queue,
n_brick_quad_points=100,
adaptive_level=False,
use_symmetry=True,
alpha=0.1,
nlevels=15,
)
# }}} End build near field potential table
# {{{ sumpy expansion for laplace kernel
from sumpy.expansion import DefaultExpansionFactory
from sumpy.kernel import LaplaceKernel
knl = LaplaceKernel(dim)
out_kernels = [knl]
expn_factory = DefaultExpansionFactory()
local_expn_class = expn_factory.get_local_expansion_class(knl)
mpole_expn_class = expn_factory.get_multipole_expansion_class(knl)
exclude_self = True
from volumential.expansion_wrangler_fpnd import (
FPNDExpansionWranglerCodeContainer, FPNDExpansionWrangler)
wcc = FPNDExpansionWranglerCodeContainer(
ctx,
partial(mpole_expn_class, knl),
partial(local_expn_class, knl),
out_kernels,
exclude_self=exclude_self,
)
if exclude_self:
target_to_source = np.arange(tree.ntargets, dtype=np.int32)
self_extra_kwargs = {"target_to_source": target_to_source}
else:
self_extra_kwargs = {}
wrangler = FPNDExpansionWrangler(
code_container=wcc,
queue=queue,
tree=tree,
near_field_table=nftable,
dtype=dtype,
fmm_level_to_order=lambda kernel, kernel_args, tree, lev: m_order,
quad_order=q_order,
self_extra_kwargs=self_extra_kwargs,
)
# }}} End sumpy expansion for laplace kernel
print("*************************")
print("* Performing FMM ...")
print("*************************")
# {{{ conduct fmm computation
from volumential.volume_fmm import drive_volume_fmm
import time
queue.finish()
t0 = time.time()
pot, = drive_volume_fmm(
trav,
wrangler,
source_vals * q_weights,
source_vals,
direct_evaluation=force_direct_evaluation,
)
queue.finish()
t1 = time.time()
print("Finished in %.2f seconds." % (t1 - t0))
print("(%e points per second)" % (len(q_weights) / (t1 - t0)))
# }}} End conduct fmm computation
print("*************************")
print("* Postprocessing ...")
print("*************************")
# {{{ postprocess and plot
# print(pot)
solu_eval = Eval(dim, solu_expr, [x, y])
x = q_points[0].get()
y = q_points[1].get()
ze = solu_eval(queue, np.array([x, y]))
zs = pot.get()
print_error = True
if print_error:
err = np.max(np.abs(ze - zs))
print("Error =", err)
# Interpolated surface
if 0:
h = 0.005
out_x = np.arange(a, b + h, h)
out_y = np.arange(a, b + h, h)
oxx, oyy = np.meshgrid(out_x, out_y)
out_targets = make_obj_array([
cl.array.to_device(queue, oxx.flatten()),
cl.array.to_device(queue, oyy.flatten()),
])
from volumential.volume_fmm import interpolate_volume_potential
# src = source_field([q.get() for q in q_points])
# src = cl.array.to_device(queue, src)
interp_pot = interpolate_volume_potential(out_targets, trav, wrangler,
pot)
opot = interp_pot.get()
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
plt3d = plt.figure()
ax = Axes3D(plt3d) # noqa
surf = ax.plot_surface(oxx, oyy, opot.reshape(oxx.shape)) # noqa
# ax.scatter(x, y, src.get())
# ax.set_zlim(-0.25, 0.25)
plt.draw()
plt.show()
# Boxtree
if 0:
import matplotlib.pyplot as plt
if dim == 2:
# plt.plot(q_points[0].get(), q_points[1].get(), ".")
pass
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree.get(queue=queue))
plotter.draw_tree(fill=False, edgecolor="black")
# plotter.draw_box_numbers()
plotter.set_bounding_box()
plt.gca().set_aspect("equal")
plt.draw()
# plt.show()
plt.savefig("tree.png")
# Direct p2p
if 0:
print("Performing P2P")
pot_direct, = drive_volume_fmm(trav,
wrangler,
source_vals * q_weights,
source_vals,
direct_evaluation=True)
zds = pot_direct.get()
zs = pot.get()
print("P2P-FMM diff =", np.max(np.abs(zs - zds)))
print("P2P Error =", np.max(np.abs(ze - zds)))
"""
import matplotlib.pyplot as plt
import matplotlib.cm as cm
x = q_points[0].get()
y = q_points[1].get()
plt.scatter(x, y, c=np.log(abs(zs-zds)) / np.log(10), cmap=cm.jet)
plt.colorbar()
plt.xlabel("Multipole order = " + str(m_order))
plt.draw()
plt.show()
"""
# Scatter plot
if 0:
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
x = q_points[0].get()
y = q_points[1].get()
ze = solu_eval(queue, np.array([x, y]))
zs = pot.get()
plt3d = plt.figure()
ax = Axes3D(plt3d)
ax.scatter(x, y, zs, s=1)
# ax.scatter(x, y, source_field([q.get() for q in q_points]), s=1)
# import matplotlib.cm as cm
# ax.scatter(x, y, zs, c=np.log(abs(zs-zds)), cmap=cm.jet)
# plt.gca().set_aspect("equal")
# ax.set_xlim3d([-1, 1])
# ax.set_ylim3d([-1, 1])
# ax.set_zlim3d([np.min(z), np.max(z)])
# ax.set_zlim3d([-0.002, 0.00])
plt.draw()
plt.show()
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 run_build_test(builder, queue, dims, dtype, nparticles, do_plot,
max_particles_in_box=None, max_leaf_refine_weight=None,
refine_weights=None, **kwargs):
dtype = np.dtype(dtype)
if dtype == np.float32:
tol = 1e-4
elif dtype == np.float64:
tol = 1e-12
else:
raise RuntimeError("unsupported dtype: %s" % dtype)
logger.info(75*"-")
if max_particles_in_box is not None:
logger.info("%dD %s - %d particles - max %d per box - %s" % (
dims, dtype.type.__name__, nparticles, max_particles_in_box,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
else:
logger.info("%dD %s - %d particles - max leaf weight %d - %s" % (
dims, dtype.type.__name__, nparticles, max_leaf_refine_weight,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
logger.info(75*"-")
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
queue.finish()
tree, _ = builder(queue, particles,
max_particles_in_box=max_particles_in_box,
refine_weights=refine_weights,
max_leaf_refine_weight=max_leaf_refine_weight,
debug=True, **kwargs)
tree = tree.get(queue=queue)
sorted_particles = np.array(list(tree.sources))
unsorted_particles = np.array([pi.get() for pi in particles])
assert (sorted_particles
== unsorted_particles[:, tree.user_source_ids]).all()
if refine_weights is not None:
refine_weights_reordered = refine_weights.get()[tree.user_source_ids]
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
from boxtree import box_flags_enum as bfe
scaled_tol = tol*tree.root_extent
for ibox in range(tree.nboxes):
# Empty boxes exist in non-pruned trees--which themselves are undocumented.
# These boxes will fail these tests.
if not (tree.box_flags[ibox] & bfe.HAS_OWN_SRCNTGTS):
continue
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >= tree.bounding_box[0] - scaled_tol).all(), (
ibox, extent_low, tree.bounding_box[0])
assert (extent_high <= tree.bounding_box[1] + scaled_tol).all(), (
ibox, extent_high, tree.bounding_box[1])
start = tree.box_source_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
box_particles = sorted_particles[:,
start:start+tree.box_source_counts_cumul[ibox]]
good = (
(box_particles < extent_high[:, np.newaxis] + scaled_tol)
& (extent_low[:, np.newaxis] - scaled_tol <= box_particles))
all_good_here = good.all()
if do_plot and not all_good_here and all_good_so_far:
pt.plot(
box_particles[0, np.where(~good)[1]],
box_particles[1, np.where(~good)[1]], "ro")
plotter.draw_box(ibox, edgecolor="red")
if not all_good_here:
print("BAD BOX", ibox)
if not (tree.box_flags[ibox] & bfe.HAS_CHILDREN):
# Check that leaf particle density is as promised.
nparticles_in_box = tree.box_source_counts_cumul[ibox]
if max_particles_in_box is not None:
if nparticles_in_box > max_particles_in_box:
print("too many particles ({0} > {1}); box {2}".format(
nparticles_in_box, max_particles_in_box, ibox))
all_good_here = False
else:
assert refine_weights is not None
box_weight = np.sum(
refine_weights_reordered[start:start+nparticles_in_box])
if box_weight > max_leaf_refine_weight:
print("refine weight exceeded ({0} > {1}); box {2}".format(
box_weight, max_leaf_refine_weight, ibox))
all_good_here = False
all_good_so_far = all_good_so_far and all_good_here
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
assert all_good_so_far
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_fmm_completeness(
ctx_getter, dims, nsources_req, ntargets_req, who_has_extent, source_gen, target_gen, filter_kind
):
"""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 RanluxGenerator
rng = RanluxGenerator(queue, seed=13)
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,
)
if 0:
tree.get().plot()
import matplotlib.pyplot as pt
pt.show()
from boxtree.traversal import FMMTraversalBuilder
tbuild = FMMTraversalBuilder(ctx)
trav, _ = tbuild(queue, tree, debug=True)
if trav.sep_close_smaller_starts is not None:
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 filter_kind:
flags = rng.uniform(queue, ntargets or nsources, np.int32, a=0, b=2).astype(np.int8)
if filter_kind == "user":
from boxtree.tree import filter_target_lists_in_user_order
filtered_targets = filter_target_lists_in_user_order(queue, tree, flags)
wrangler = ConstantOneExpansionWranglerWithFilteredTargetsInUserOrder(
host_tree, filtered_targets.get(queue=queue)
)
elif filter_kind == "tree":
from boxtree.tree import filter_target_lists_in_tree_order
filtered_targets = 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)
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)
# {{{ build, evaluate matrix (and identify missing interactions)
if 0:
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 1:
pt.spy(mat)
pt.show()
missing_tgts, missing_srcs = np.where(mat == 0)
if 1 and len(missing_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_missing_tgts = host_tree.indices_to_tree_target_order(missing_tgts)
tree_order_missing_srcs = host_tree.indices_to_tree_source_order(missing_srcs)
src_boxes = [host_tree.find_box_nr_for_source(i) for i in tree_order_missing_srcs]
tgt_boxes = [host_tree.find_box_nr_for_target(i) for i in tree_order_missing_tgts]
print(src_boxes)
print(tgt_boxes)
pt.plot(host_tree.targets[0][tree_order_missing_tgts], host_tree.targets[1][tree_order_missing_tgts], "rv")
pt.plot(host_tree.sources[0][tree_order_missing_srcs], host_tree.sources[1][tree_order_missing_srcs], "go")
pt.gca().set_aspect("equal")
pt.show()
# }}}
if filter_kind:
pot = pot[flags.get() > 0]
rel_err = la.norm((pot - weights_sum) / nsources)
good = rel_err < 1e-8
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
# ENDEXAMPLE
# -----------------------------------------------------------------------------
# plot the tree
# -----------------------------------------------------------------------------
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "+")
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree.get(queue=queue))
plotter.draw_tree(fill=False, edgecolor="black")
#plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.gca().set_aspect("equal")
pt.tight_layout()
pt.tick_params(
axis='x', # changes apply to the x-axis
which='both', # both major and minor ticks are affected
bottom='off', # ticks along the bottom edge are off
top='off', # ticks along the top edge are off
labelbottom='off')
pt.tick_params(
axis='y',
which='both',
left='off',
top='off',
labelleft='off')
pt.savefig("tree.pdf")
[rng.normal(queue, nparticles, dtype=np.float64) for i in range(dims)])
# -----------------------------------------------------------------------------
# build tree and traversals (lists)
# -----------------------------------------------------------------------------
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, particles, max_particles_in_box=30)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav, _ = tg(queue, tree)
# ENDEXAMPLE
# -----------------------------------------------------------------------------
# plot the tree
# -----------------------------------------------------------------------------
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree.get(queue=queue))
plotter.draw_tree(fill=False, edgecolor="black")
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.gca().set_aspect("equal")
pt.savefig("tree.png")
def test_tree_connectivity(ctx_getter, dims, sources_are_targets):
logging.basicConfig(level=logging.INFO)
ctx = ctx_getter()
queue = cl.CommandQueue(ctx)
dtype = np.float64
sources = make_normal_particle_array(queue, 1 * 10**5, dims, dtype)
if sources_are_targets:
targets = None
else:
targets = make_normal_particle_array(queue, 2 * 10**5, dims, dtype)
from boxtree import TreeBuilder
tb = TreeBuilder(ctx)
tree, _ = tb(queue, sources, max_particles_in_box=30,
targets=targets, debug=True)
from boxtree.traversal import FMMTraversalBuilder
tg = FMMTraversalBuilder(ctx)
trav, _ = tg(queue, tree, debug=True)
tree = tree.get(queue=queue)
trav = trav.get(queue=queue)
levels = tree.box_levels
parents = tree.box_parent_ids.T
children = tree.box_child_ids.T
centers = tree.box_centers.T
# {{{ parent and child relations, levels match up
for ibox in range(1, tree.nboxes):
# /!\ Not testing box 0, has no parents
parent = parents[ibox]
assert levels[parent] + 1 == levels[ibox]
assert ibox in children[parent], ibox
# }}}
if 0:
import matplotlib.pyplot as pt
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black")
plotter.draw_box_numbers()
plotter.set_bounding_box()
pt.show()
# {{{ neighbor_source_boxes (list 1) consists of source boxes
for itgt_box, ibox in enumerate(trav.target_boxes):
start, end = trav.neighbor_source_boxes_starts[itgt_box:itgt_box+2]
nbl = trav.neighbor_source_boxes_lists[start:end]
if sources_are_targets:
assert ibox in nbl
for jbox in nbl:
assert (0 == children[jbox]).all(), (ibox, jbox, children[jbox])
logger.info("list 1 consists of source boxes")
# }}}
# {{{ separated siblings (list 2) are actually separated
for itgt_box, tgt_ibox in enumerate(trav.target_or_target_parent_boxes):
start, end = trav.sep_siblings_starts[itgt_box:itgt_box+2]
seps = trav.sep_siblings_lists[start:end]
assert (levels[seps] == levels[tgt_ibox]).all()
# three-ish box radii (half of size)
mindist = 2.5 * 0.5 * 2**-int(levels[tgt_ibox]) * tree.root_extent
icenter = centers[tgt_ibox]
for jbox in seps:
dist = la.norm(centers[jbox]-icenter)
assert dist > mindist, (dist, mindist)
logger.info("separated siblings (list 2) are actually separated")
# }}}
if sources_are_targets:
# {{{ sep_{smaller,bigger} are duals of each other
assert (trav.target_or_target_parent_boxes == np.arange(tree.nboxes)).all()
# {{{ list 4 <= list 3
for itarget_box, ibox in enumerate(trav.target_boxes):
for ssn in trav.sep_smaller_by_level:
start, end = ssn.starts[itarget_box:itarget_box+2]
for jbox in ssn.lists[start:end]:
rstart, rend = trav.sep_bigger_starts[jbox:jbox+2]
assert ibox in trav.sep_bigger_lists[rstart:rend], (ibox, jbox)
# }}}
# {{{ list 4 <= list 3
box_to_target_box_index = np.empty(tree.nboxes, tree.box_id_dtype)
box_to_target_box_index.fill(-1)
box_to_target_box_index[trav.target_boxes] = np.arange(
len(trav.target_boxes), dtype=tree.box_id_dtype)
assert (trav.source_boxes == trav.target_boxes).all()
assert (trav.target_or_target_parent_boxes == np.arange(
tree.nboxes, dtype=tree.box_id_dtype)).all()
for ibox in range(tree.nboxes):
start, end = trav.sep_bigger_starts[ibox:ibox+2]
for jbox in trav.sep_bigger_lists[start:end]:
# In principle, entries of sep_bigger_lists are
# source boxes. In this special case, source and target boxes
# are the same thing (i.e. leaves--see assertion above), so we
# may treat them as targets anyhow.
jtgt_box = box_to_target_box_index[jbox]
assert jtgt_box != -1
good = False
for ssn in trav.sep_smaller_by_level:
rstart, rend = ssn.starts[jtgt_box:jtgt_box+2]
good = good or ibox in ssn.lists[rstart:rend]
if not good:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
plotter.draw_box(ibox, facecolor='green', alpha=0.5)
plotter.draw_box(jbox, facecolor='red', alpha=0.5)
import matplotlib.pyplot as pt
pt.gca().set_aspect("equal")
pt.show()
# This assertion failing means that ibox's list 4 contains a box
# 'jbox' whose list 3 does not contain ibox.
assert good, (ibox, jbox)
# }}}
logger.info("list 3, 4 are duals")
# }}}
# {{{ sep_smaller satisfies relative level assumption
for itarget_box, ibox in enumerate(trav.target_boxes):
for ssn in trav.sep_smaller_by_level:
start, end = ssn.starts[itarget_box:itarget_box+2]
for jbox in ssn.lists[start:end]:
assert levels[ibox] < levels[jbox]
logger.info("list 3 satisfies relative level assumption")
# }}}
# {{{ sep_bigger satisfies relative level assumption
for itgt_box, tgt_ibox in enumerate(trav.target_or_target_parent_boxes):
start, end = trav.sep_bigger_starts[itgt_box:itgt_box+2]
for jbox in trav.sep_bigger_lists[start:end]:
assert levels[tgt_ibox] > levels[jbox]
logger.info("list 4 satisfies relative level assumption")
# }}}
# {{{ level_start_*_box_nrs lists make sense
for name, ref_array in [
("level_start_source_box_nrs", trav.source_boxes),
("level_start_source_parent_box_nrs", trav.source_parent_boxes),
("level_start_target_box_nrs", trav.target_boxes),
("level_start_target_or_target_parent_box_nrs",
trav.target_or_target_parent_boxes)
]:
level_starts = getattr(trav, name)
for lev in range(tree.nlevels):
start, stop = level_starts[lev:lev+2]
box_nrs = ref_array[start:stop]
assert (tree.box_levels[box_nrs] == lev).all(), name
def run_build_test(builder, queue, dims, dtype, nparticles, do_plot,
max_particles_in_box=30, **kwargs):
dtype = np.dtype(dtype)
if dtype == np.float32:
tol = 1e-4
elif dtype == np.float64:
tol = 1e-12
else:
raise RuntimeError("unsupported dtype: %s" % dtype)
if (dtype == np.float32
and dims == 2
and queue.device.platform.name == "Portable Computing Language"):
pytest.xfail("2D float doesn't work on POCL")
logger.info(75*"-")
logger.info("%dD %s - %d particles - max %d per box - %s" % (
dims, dtype.type.__name__, nparticles, max_particles_in_box,
" - ".join("%s: %s" % (k, v) for k, v in six.iteritems(kwargs))))
logger.info(75*"-")
particles = make_normal_particle_array(queue, nparticles, dims, dtype)
if do_plot:
import matplotlib.pyplot as pt
pt.plot(particles[0].get(), particles[1].get(), "x")
queue.finish()
tree, _ = builder(queue, particles,
max_particles_in_box=max_particles_in_box, debug=True,
**kwargs)
tree = tree.get(queue=queue)
sorted_particles = np.array(list(tree.sources))
unsorted_particles = np.array([pi.get() for pi in particles])
assert (sorted_particles
== unsorted_particles[:, tree.user_source_ids]).all()
all_good_so_far = True
if do_plot:
from boxtree.visualization import TreePlotter
plotter = TreePlotter(tree)
plotter.draw_tree(fill=False, edgecolor="black", zorder=10)
plotter.set_bounding_box()
from boxtree import box_flags_enum as bfe
scaled_tol = tol*tree.root_extent
for ibox in range(tree.nboxes):
# Empty boxes exist in non-pruned trees--which themselves are undocumented.
# These boxes will fail these tests.
if not (tree.box_flags[ibox] & bfe.HAS_OWN_SRCNTGTS):
continue
extent_low, extent_high = tree.get_box_extent(ibox)
assert (extent_low >= tree.bounding_box[0] - scaled_tol).all(), (
ibox, extent_low, tree.bounding_box[0])
assert (extent_high <= tree.bounding_box[1] + scaled_tol).all(), (
ibox, extent_high, tree.bounding_box[1])
start = tree.box_source_starts[ibox]
box_children = tree.box_child_ids[:, ibox]
existing_children = box_children[box_children != 0]
assert (tree.box_source_counts_nonchild[ibox]
+ np.sum(tree.box_source_counts_cumul[existing_children])
== tree.box_source_counts_cumul[ibox])
box_particles = sorted_particles[:,
start:start+tree.box_source_counts_cumul[ibox]]
good = (
(box_particles < extent_high[:, np.newaxis] + scaled_tol)
&
(extent_low[:, np.newaxis] - scaled_tol <= box_particles)
)
all_good_here = good.all()
if do_plot and not all_good_here and all_good_so_far:
pt.plot(
box_particles[0, np.where(~good)[1]],
box_particles[1, np.where(~good)[1]], "ro")
plotter.draw_box(ibox, edgecolor="red")
if not all_good_here:
print("BAD BOX", ibox)
all_good_so_far = all_good_so_far and all_good_here
if do_plot:
pt.gca().set_aspect("equal", "datalim")
pt.show()
assert all_good_so_far