def __init__(self, code_container, queue, geo_data, dtype,
qbx_order, fmm_level_to_order,
source_extra_kwargs, kernel_extra_kwargs):
SumpyExpansionWrangler.__init__(self,
code_container, queue, geo_data.tree(),
dtype, fmm_level_to_order, source_extra_kwargs, kernel_extra_kwargs)
self.qbx_order = qbx_order
self.geo_data = geo_data
def __init__(self, code_container, queue, geo_data, dtype,
qbx_order, fmm_level_to_order,
source_extra_kwargs, kernel_extra_kwargs,
_use_target_specific_qbx=None):
if _use_target_specific_qbx:
raise ValueError("TSQBX is not implemented in sumpy")
SumpyExpansionWrangler.__init__(self,
code_container, queue, geo_data.tree(),
dtype, fmm_level_to_order, source_extra_kwargs, kernel_extra_kwargs)
self.qbx_order = qbx_order
self.geo_data = geo_data
self.using_tsqbx = False
def refine_locals(
self,
level_start_target_or_target_parent_box_nrs,
target_or_target_parent_boxes,
local_exps,
):
return SumpyExpansionWrangler.refine_locals(
self,
level_start_target_or_target_parent_box_nrs,
target_or_target_parent_boxes,
local_exps,
)
def form_locals(
self,
level_start_target_or_target_parent_box_nrs,
target_or_target_parent_boxes,
starts,
lists,
src_weights,
):
return SumpyExpansionWrangler.form_locals(
self,
level_start_target_or_target_parent_box_nrs,
target_or_target_parent_boxes,
starts,
lists,
src_weights,
)
def multipole_to_local(
self,
level_start_target_box_nrs,
target_boxes,
src_box_starts,
src_box_lists,
mpole_exps,
):
return SumpyExpansionWrangler.multipole_to_local(
self,
level_start_target_box_nrs,
target_boxes,
src_box_starts,
src_box_lists,
mpole_exps,
)
def full_potential_zeros(self):
# The superclass generates a full field of zeros, for all
# (not just non-QBX) targets.
return SumpyExpansionWrangler.potential_zeros(self)
def eval_direct_p2p(
self, target_boxes, source_box_starts, source_box_lists, src_weights
):
return SumpyExpansionWrangler.eval_direct(
self, target_boxes, source_box_starts, source_box_lists, src_weights
)
def eval_locals(self, level_start_target_box_nrs, target_boxes, local_exps):
return SumpyExpansionWrangler.eval_locals(
self, level_start_target_box_nrs, target_boxes, local_exps
)
def local_expansion_zeros(self):
return SumpyExpansionWrangler.local_expansion_zeros(self)
def coarsen_multipoles(
self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles
):
return SumpyExpansionWrangler.coarsen_multipoles(
self, level_start_source_parent_box_nrs, source_parent_boxes, mpoles
)
def form_multipoles(self, level_start_source_box_nrs, source_boxes, src_weights):
return SumpyExpansionWrangler.form_multipoles(
self, level_start_source_box_nrs, source_boxes, src_weights
)
def finalize_potentials(self, potentials):
# return potentials
return SumpyExpansionWrangler.finalize_potentials(self, potentials)
def reorder_potentials(self, potentials):
return SumpyExpansionWrangler.reorder_potentials(self, potentials)
def reorder_sources(self, source_array):
return SumpyExpansionWrangler.reorder_sources(self, source_array)
def output_zeros(self):
return SumpyExpansionWrangler.output_zeros(self)
def __init__(self, code_container, queue, geo_data,
dtype, source_extra_kwargs, kernel_extra_kwargs):
SumpyExpansionWrangler.__init__(self,
code_container, queue, geo_data.tree(),
dtype, source_extra_kwargs, kernel_extra_kwargs)
self.geo_data = geo_data
def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, insn,
bound_expr, evaluate):
# {{{ gather unique target discretizations used
target_name_to_index = {}
targets = []
for o in insn.outputs:
assert o.qbx_forced_limit not in (-1, 1)
if o.target_name in target_name_to_index:
continue
target_name_to_index[o.target_name] = len(targets)
targets.append(
bound_expr.places.get_geometry(o.target_name.geometry))
targets = tuple(targets)
# }}}
# {{{ get wrangler
geo_data = self.fmm_geometry_data(targets)
from pytential import bind, sym
waa = bind(
bound_expr.places,
sym.weights_and_area_elements(self.ambient_dim,
dofdesc=insn.source))(actx)
strengths = [waa * evaluate(density) for density in insn.densities]
flat_strengths = [flatten(strength, actx) for strength in strengths]
fmm_kernel = self.get_fmm_kernel(insn.target_kernels)
output_and_expansion_dtype = (self.get_fmm_output_and_expansion_dtype(
insn.target_kernels, strengths[0]))
kernel_extra_kwargs, source_extra_kwargs = (
self.get_fmm_expansion_wrangler_extra_kwargs(
actx, insn.target_kernels + insn.source_kernels,
geo_data.tree().user_source_ids, insn.kernel_arguments,
evaluate))
tree_indep = self._tree_indep_data_for_wrangler(
fmm_kernel,
target_kernels=insn.target_kernels,
source_kernels=insn.source_kernels)
from sumpy.fmm import SumpyExpansionWrangler
wrangler = SumpyExpansionWrangler(
tree_indep,
geo_data.traversal(),
output_and_expansion_dtype,
self.fmm_level_to_order,
source_extra_kwargs=source_extra_kwargs,
kernel_extra_kwargs=kernel_extra_kwargs)
# }}}
from boxtree.fmm import drive_fmm
all_potentials_on_every_tgt = drive_fmm(wrangler,
flat_strengths,
timing_data=None)
# {{{ postprocess fmm
results = []
for o in insn.outputs:
target_index = target_name_to_index[o.target_name]
target_slice = slice(*geo_data.target_info(
).target_discr_starts[target_index:target_index + 2])
target_discr = targets[target_index]
result = all_potentials_on_every_tgt[
o.target_kernel_index][target_slice]
from meshmode.discretization import Discretization
if isinstance(target_discr, Discretization):
template_ary = thaw(target_discr.nodes()[0], actx)
result = unflatten(template_ary, result, actx, strict=False)
results.append((o.name, result))
# }}}
timing_data = {}
return results, timing_data
def eval_multipoles(
self, target_boxes_by_source_level, source_boxes_by_level, mpole_exps
):
return SumpyExpansionWrangler.eval_multipoles(
self, target_boxes_by_source_level, source_boxes_by_level, mpole_exps
)
def full_output_zeros(self):
# The superclass generates a full field of zeros, for all
# (not just non-QBX) targets.
return SumpyExpansionWrangler.output_zeros(self)
def multipole_expansion_zeros(self):
return SumpyExpansionWrangler.multipole_expansion_zeros(self)
