def form_global_qbx_locals(self, src_weight_vecs):
local_exps = self.qbx_local_expansion_zeros()
events = []
geo_data = self.geo_data
if len(geo_data.global_qbx_centers()) == 0:
return (local_exps, SumpyTimingFuture(self.queue, events))
traversal = geo_data.traversal()
starts = traversal.neighbor_source_boxes_starts
lists = traversal.neighbor_source_boxes_lists
kwargs = self.extra_kwargs.copy()
kwargs.update(self.box_source_list_kwargs())
p2qbxl = self.code.p2qbxl(self.qbx_order)
evt, (result, ) = p2qbxl(
self.queue,
global_qbx_centers=geo_data.global_qbx_centers(),
qbx_center_to_target_box=geo_data.qbx_center_to_target_box(),
qbx_centers=geo_data.flat_centers(),
qbx_expansion_radii=geo_data.flat_expansion_radii(),
source_box_starts=starts,
source_box_lists=lists,
strengths=src_weight_vecs,
qbx_expansions=local_exps,
**kwargs)
events.append(evt)
assert local_exps is result
result.add_event(evt)
return (result, SumpyTimingFuture(self.queue, events))
def eval_qbx_expansions(self, qbx_expansions):
pot = self.full_output_zeros()
geo_data = self.geo_data
events = []
if len(geo_data.global_qbx_centers()) == 0:
return (pot, SumpyTimingFuture(self.queue, events))
ctt = geo_data.center_to_tree_targets()
qbxl2p = self.code.qbxl2p(self.qbx_order)
evt, pot_res = qbxl2p(self.queue,
qbx_centers=geo_data.flat_centers(),
qbx_expansion_radii=geo_data.flat_expansion_radii(),
global_qbx_centers=geo_data.global_qbx_centers(),
center_to_targets_starts=ctt.starts,
center_to_targets_lists=ctt.lists,
targets=self.tree.targets,
qbx_expansions=qbx_expansions,
result=pot,
**self.kernel_extra_kwargs.copy())
for pot_i, pot_res_i in zip(pot, pot_res):
assert pot_i is pot_res_i
return (pot, SumpyTimingFuture(self.queue, events))
def translate_box_multipoles_to_qbx_local(self, multipole_exps):
queue = multipole_exps.queue
qbx_expansions = self.qbx_local_expansion_zeros(multipole_exps)
events = []
geo_data = self.geo_data
if geo_data.ncenters == 0:
return (qbx_expansions, SumpyTimingFuture(queue, events))
traversal = geo_data.traversal()
wait_for = multipole_exps.events
for isrc_level, ssn in enumerate(traversal.from_sep_smaller_by_level):
m2qbxl = self.tree_indep.m2qbxl(
self.level_orders[isrc_level],
self.qbx_order)
source_level_start_ibox, source_mpoles_view = \
self.multipole_expansions_view(multipole_exps, isrc_level)
evt, (qbx_expansions_res,) = m2qbxl(queue,
qbx_center_to_target_box_source_level=(
geo_data.qbx_center_to_target_box_source_level(isrc_level)
),
centers=self.tree.box_centers,
qbx_centers=geo_data.flat_centers(),
qbx_expansion_radii=geo_data.flat_expansion_radii(),
src_expansions=source_mpoles_view,
src_base_ibox=source_level_start_ibox,
qbx_expansions=qbx_expansions,
src_box_starts=ssn.starts,
src_box_lists=ssn.lists,
src_rscale=level_to_rscale(self.tree, isrc_level),
wait_for=wait_for,
**self.kernel_extra_kwargs)
events.append(evt)
wait_for = [evt]
assert qbx_expansions_res is qbx_expansions
qbx_expansions.add_event(evt)
return (qbx_expansions, SumpyTimingFuture(queue, events))
def translate_box_local_to_qbx_local(self, local_exps):
queue = local_exps.queue
qbx_expansions = self.qbx_local_expansion_zeros(local_exps)
geo_data = self.geo_data
events = []
if geo_data.ncenters == 0:
return (qbx_expansions, SumpyTimingFuture(queue, events))
trav = geo_data.traversal()
wait_for = local_exps.events
for isrc_level in range(geo_data.tree().nlevels):
l2qbxl = self.tree_indep.l2qbxl(
self.level_orders[isrc_level],
self.qbx_order)
target_level_start_ibox, target_locals_view = \
self.local_expansions_view(local_exps, isrc_level)
evt, (qbx_expansions_res,) = l2qbxl(
queue,
qbx_center_to_target_box=geo_data.qbx_center_to_target_box(),
target_boxes=trav.target_boxes,
target_base_ibox=target_level_start_ibox,
centers=self.tree.box_centers,
qbx_centers=geo_data.flat_centers(),
qbx_expansion_radii=geo_data.flat_expansion_radii(),
expansions=target_locals_view,
qbx_expansions=qbx_expansions,
src_rscale=level_to_rscale(self.tree, isrc_level),
wait_for=wait_for,
**self.kernel_extra_kwargs)
events.append(evt)
wait_for = [evt]
assert qbx_expansions_res is qbx_expansions
qbx_expansions.add_event(evt)
return (qbx_expansions, SumpyTimingFuture(queue, events))
def eval_direct(
self,
target_boxes,
neighbor_source_boxes_starts,
neighbor_source_boxes_lists,
mode_coefs,
):
pot = self.output_zeros()
events = []
for i in range(len(self.code.out_kernels)):
# print("processing near-field of out_kernel", i)
pot[i], evt = self.eval_direct_single_out_kernel(
pot[i],
self.code.out_kernels[i],
target_boxes,
neighbor_source_boxes_starts,
neighbor_source_boxes_lists,
mode_coefs,
)
events.append(evt)
for out_pot in pot:
out_pot.finish()
return (pot, SumpyTimingFuture(self.queue, events))
def eval_direct(
self,
target_boxes,
neighbor_source_boxes_starts,
neighbor_source_boxes_lists,
mode_coefs,
):
pot = self.output_zeros()
if pot.dtype != np.object:
pot = make_obj_array([pot, ])
events = []
for i in range(len(self.code.out_kernels)):
# print("processing near-field of out_kernel", i)
pot[i], evt = self.eval_direct_single_out_kernel(
pot[i],
self.code.out_kernels[i],
target_boxes,
neighbor_source_boxes_starts,
neighbor_source_boxes_lists,
mode_coefs,
)
events.append(evt)
for out_pot in pot:
if isinstance(out_pot, cl.array.Array):
out_pot.finish()
# boxtree.pyfmmlib_integration handles things diffferently
# when out_kernels has only one element
if len(pot) == 1:
pot = pot[0]
return (pot, SumpyTimingFuture(self.queue, events))
def eval_target_specific_qbx_locals(self, src_weight_vecs):
return (self.full_output_zeros(), SumpyTimingFuture(self.queue, events=()))
def eval_target_specific_qbx_locals(self, src_weight_vecs):
template_ary = src_weight_vecs[0]
return (self.full_output_zeros(template_ary),
SumpyTimingFuture(template_ary.queue, events=()))