def target_info(self):
"""Return a :class:`TargetInfo`. |cached|"""
from pytential.utils import flatten_if_needed
code_getter = self.code_getter
queue = self.array_context.queue
ntargets = self.ncenters
target_discr_starts = []
for target_discr, qbx_side in self.target_discrs_and_qbx_sides:
target_discr_starts.append(ntargets)
ntargets += target_discr.ndofs
target_discr_starts.append(ntargets)
targets = cl.array.empty(self.cl_context, (self.ambient_dim, ntargets),
self.coord_dtype)
code_getter.copy_targets_kernel()(queue,
targets=targets[:, :self.ncenters],
points=self.flat_centers())
for start, (target_discr, _) in zip(target_discr_starts,
self.target_discrs_and_qbx_sides):
code_getter.copy_targets_kernel()(
queue,
targets=targets[:, start:start + target_discr.ndofs],
points=flatten_if_needed(self.array_context,
target_discr.nodes()))
return TargetInfo(targets=targets,
target_discr_starts=target_discr_starts,
ntargets=ntargets).with_queue(None)
def check_element_prop_threshold(self,
element_property,
threshold,
refine_flags,
debug,
wait_for=None):
knl = self.code_container.element_prop_threshold_checker()
if debug:
npanels_to_refine_prev = cl.array.sum(refine_flags).get()
from pytential.utils import flatten_if_needed
element_property = flatten_if_needed(self.array_context,
element_property)
evt, out = knl(self.queue,
element_property=element_property,
refine_flags=refine_flags,
refine_flags_updated=np.array(0),
threshold=np.array(threshold),
wait_for=wait_for)
cl.wait_for_events([evt])
if debug:
npanels_to_refine = cl.array.sum(refine_flags).get()
if npanels_to_refine > npanels_to_refine_prev:
logger.debug("refiner: found {} panel(s) to refine".format(
npanels_to_refine - npanels_to_refine_prev))
return (out["refine_flags_updated"] == 1).all()
def target_info(self):
code_getter = self.code_getter
lpot_src = self.lpot_source
target_discrs = self.target_discrs
ntargets = 0
target_discr_starts = []
for target_discr in target_discrs:
target_discr_starts.append(ntargets)
ntargets += target_discr.ndofs
target_discr_starts.append(ntargets)
targets = self.array_context.empty(
(lpot_src.ambient_dim, ntargets),
self.coord_dtype)
from pytential.utils import flatten_if_needed
for start, target_discr in zip(target_discr_starts, target_discrs):
code_getter.copy_targets_kernel()(
self.array_context.queue,
targets=targets[:, start:start+target_discr.ndofs],
points=flatten_if_needed(
self.array_context, target_discr.nodes()))
return _TargetInfo(
targets=targets,
target_discr_starts=target_discr_starts,
ntargets=ntargets).with_queue(None)
def get_fmm_expansion_wrangler_extra_kwargs(self, actx, out_kernels,
tree_user_source_ids,
arguments, evaluator):
# This contains things like the Helmholtz parameter k or
# the normal directions for double layers.
queue = actx.queue
def reorder_sources(source_array):
if isinstance(source_array, cl.array.Array):
return (source_array.with_queue(queue)
[tree_user_source_ids].with_queue(None))
else:
return source_array
kernel_extra_kwargs = {}
source_extra_kwargs = {}
from sumpy.tools import gather_arguments, gather_source_arguments
from pytools.obj_array import obj_array_vectorize
from pytential.utils import flatten_if_needed
for func, var_dict in [
(gather_arguments, kernel_extra_kwargs),
(gather_source_arguments, source_extra_kwargs),
]:
for arg in func(out_kernels):
var_dict[arg.name] = obj_array_vectorize(
reorder_sources,
flatten_if_needed(actx, evaluator(arguments[arg.name])))
return kernel_extra_kwargs, source_extra_kwargs
def flatten(self, ary):
# Return a flat version of *ary*. The returned value is suitable for
# use with solvers whose API expects a one-dimensional array.
if not self._operator_uses_obj_array:
ary = [ary]
result = self.array_context.empty(self.total_dofs, self.dtype)
from pytential.utils import flatten_if_needed
for res_i, (start, end) in zip(ary, self.starts_and_ends):
result[start:end] = flatten_if_needed(self.array_context, res_i)
return result
def exec_compute_potential_insn_direct(self, actx: PyOpenCLArrayContext,
insn, bound_expr, evaluate):
kernel_args = {}
from pytential.utils import flatten_if_needed
from meshmode.dof_array import flatten, thaw, unflatten
for arg_name, arg_expr in insn.kernel_arguments.items():
kernel_args[arg_name] = flatten_if_needed(actx, evaluate(arg_expr))
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) for strength in strengths]
results = []
p2p = None
for o in insn.outputs:
target_discr = bound_expr.places.get_discretization(
o.target_name.geometry, o.target_name.discr_stage)
if p2p is None:
p2p = self.get_p2p(actx, source_kernels=insn.source_kernels,
target_kernels=insn.target_kernels)
evt, output_for_each_kernel = p2p(actx.queue,
flatten_if_needed(actx, target_discr.nodes()),
flatten(thaw(actx, self.density_discr.nodes())),
flat_strengths, **kernel_args)
from meshmode.discretization import Discretization
result = output_for_each_kernel[o.target_kernel_index]
if isinstance(target_discr, Discretization):
result = unflatten(actx, target_discr, result)
results.append((o.name, result))
timing_data = {}
return results, timing_data
def _get_layer_potential_args(mapper, expr, include_args=None):
"""
:arg mapper: a :class:`~pytential.symbolic.matrix.MatrixBuilderBase`.
:arg expr: symbolic layer potential expression.
:return: a mapping of kernel arguments evaluated by the *mapper*.
"""
kernel_args = {}
for arg_name, arg_expr in six.iteritems(expr.kernel_arguments):
if (include_args is not None and arg_name not in include_args):
continue
kernel_args[arg_name] = flatten_if_needed(mapper.array_context,
mapper.rec(arg_expr))
return kernel_args
def exec_compute_potential_insn(self, actx, insn, bound_expr, evaluate,
return_timing_data):
if return_timing_data:
from warnings import warn
warn("Timing data collection not supported.",
category=UnableToCollectTimingData)
p2p = None
kernel_args = {}
for arg_name, arg_expr in insn.kernel_arguments.items():
kernel_args[arg_name] = evaluate(arg_expr)
strengths = evaluate(insn.density)
# FIXME: Do this all at once
results = []
for o in insn.outputs:
target_discr = bound_expr.places.get_discretization(
o.target_name.geometry, o.target_name.discr_stage)
# no on-disk kernel caching
if p2p is None:
p2p = self.get_p2p(actx, insn.kernels)
from pytential.utils import flatten_if_needed
evt, output_for_each_kernel = p2p(
actx.queue, flatten_if_needed(actx, target_discr.nodes()),
self._nodes, [strengths], **kernel_args)
from meshmode.discretization import Discretization
result = output_for_each_kernel[o.kernel_index]
if isinstance(target_discr, Discretization):
from meshmode.dof_array import unflatten
result = unflatten(actx, target_discr, result)
results.append((o.name, result))
timing_data = {}
return results, timing_data
def exec_compute_potential_insn_direct(self, actx, insn, bound_expr,
evaluate, return_timing_data):
from pytential import bind, sym
if return_timing_data:
from pytential.source import UnableToCollectTimingData
from warnings import warn
warn("Timing data collection not supported.",
category=UnableToCollectTimingData)
lpot_applier = self.get_lpot_applier(insn.target_kernels,
insn.source_kernels)
p2p = None
lpot_applier_on_tgt_subset = None
from pytential.utils import flatten_if_needed
kernel_args = {}
for arg_name, arg_expr in insn.kernel_arguments.items():
kernel_args[arg_name] = flatten_if_needed(actx, evaluate(arg_expr))
waa = bind(
bound_expr.places,
sym.weights_and_area_elements(self.ambient_dim,
dofdesc=insn.source))(actx)
strength_vecs = [waa * evaluate(density) for density in insn.densities]
from meshmode.discretization import Discretization
flat_strengths = [flatten(strength) for strength in strength_vecs]
source_discr = bound_expr.places.get_discretization(
insn.source.geometry, insn.source.discr_stage)
# FIXME: Do this all at once
results = []
for o in insn.outputs:
source_dd = insn.source.copy(discr_stage=o.target_name.discr_stage)
target_discr = bound_expr.places.get_discretization(
o.target_name.geometry, o.target_name.discr_stage)
density_discr = bound_expr.places.get_discretization(
source_dd.geometry, source_dd.discr_stage)
is_self = density_discr is target_discr
if is_self:
# QBXPreprocessor is supposed to have taken care of this
assert o.qbx_forced_limit is not None
assert abs(o.qbx_forced_limit) > 0
expansion_radii = bind(
bound_expr.places,
sym.expansion_radii(self.ambient_dim,
dofdesc=o.target_name))(actx)
centers = bind(
bound_expr.places,
sym.expansion_centers(self.ambient_dim,
o.qbx_forced_limit,
dofdesc=o.target_name))(actx)
evt, output_for_each_kernel = lpot_applier(
actx.queue,
flatten(thaw(actx, target_discr.nodes())),
flatten(thaw(actx, source_discr.nodes())),
flatten(centers),
flat_strengths,
expansion_radii=flatten(expansion_radii),
**kernel_args)
result = output_for_each_kernel[o.target_kernel_index]
if isinstance(target_discr, Discretization):
result = unflatten(actx, target_discr, result)
results.append((o.name, result))
else:
# no on-disk kernel caching
if p2p is None:
p2p = self.get_p2p(actx, insn.target_kernels,
insn.source_kernels)
if lpot_applier_on_tgt_subset is None:
lpot_applier_on_tgt_subset = self.get_lpot_applier_on_tgt_subset(
insn.target_kernels, insn.source_kernels)
queue = actx.queue
flat_targets = flatten_if_needed(actx, target_discr.nodes())
flat_sources = flatten(thaw(actx, source_discr.nodes()))
evt, output_for_each_kernel = p2p(queue, flat_targets,
flat_sources, flat_strengths,
**kernel_args)
qbx_forced_limit = o.qbx_forced_limit
if qbx_forced_limit is None:
qbx_forced_limit = 0
target_discrs_and_qbx_sides = ((target_discr,
qbx_forced_limit), )
geo_data = self.qbx_fmm_geometry_data(
bound_expr.places,
insn.source.geometry,
target_discrs_and_qbx_sides=target_discrs_and_qbx_sides)
# center-related info is independent of targets
# First ncenters targets are the centers
tgt_to_qbx_center = (
geo_data.user_target_to_center()[geo_data.ncenters:].copy(
queue=queue).with_queue(queue))
qbx_tgt_numberer = self.get_qbx_target_numberer(
tgt_to_qbx_center.dtype)
qbx_tgt_count = cl.array.empty(queue, (), np.int32)
qbx_tgt_numbers = cl.array.empty_like(tgt_to_qbx_center)
qbx_tgt_numberer(tgt_to_qbx_center,
qbx_tgt_numbers,
qbx_tgt_count,
queue=queue)
qbx_tgt_count = int(qbx_tgt_count.get())
if (o.qbx_forced_limit is not None
and abs(o.qbx_forced_limit) == 1
and qbx_tgt_count < target_discr.ndofs):
raise RuntimeError("Did not find a matching QBX center "
"for some targets")
qbx_tgt_numbers = qbx_tgt_numbers[:qbx_tgt_count]
qbx_center_numbers = tgt_to_qbx_center[qbx_tgt_numbers]
qbx_center_numbers.finish()
tgt_subset_kwargs = kernel_args.copy()
for i, res_i in enumerate(output_for_each_kernel):
tgt_subset_kwargs[f"result_{i}"] = res_i
if qbx_tgt_count:
lpot_applier_on_tgt_subset(
queue,
targets=flat_targets,
sources=flat_sources,
centers=geo_data.flat_centers(),
expansion_radii=geo_data.flat_expansion_radii(),
strengths=flat_strengths,
qbx_tgt_numbers=qbx_tgt_numbers,
qbx_center_numbers=qbx_center_numbers,
**tgt_subset_kwargs)
result = output_for_each_kernel[o.target_kernel_index]
if isinstance(target_discr, Discretization):
result = unflatten(actx, target_discr, result)
results.append((o.name, result))
timing_data = {}
return results, timing_data
def build_tree_with_qbx_metadata(actx: PyOpenCLArrayContext,
places,
tree_builder,
particle_list_filter,
sources_list=(),
targets_list=(),
use_stage2_discr=False):
"""Return a :class:`TreeWithQBXMetadata` built from the given layer
potential source. This contains particles of four different types:
* source particles either from
:class:`~pytential.symbolic.primitives.QBX_SOURCE_STAGE1` or
:class:`~pytential.symbolic.primitives.QBX_SOURCE_QUAD_STAGE2`.
* centers from
:class:`~pytential.symbolic.primitives.QBX_SOURCE_STAGE1`.
* targets from ``targets_list``.
:arg actx: A :class:`PyOpenCLArrayContext`
:arg places: An instance of
:class:`~pytential.symbolic.execution.GeometryCollection`.
:arg targets_list: A list of :class:`pytential.target.TargetBase`
:arg use_stage2_discr: If *True*, builds a tree with stage 2 sources.
If *False*, the tree is built with stage 1 sources.
"""
# The ordering of particles is as follows:
# - sources go first
# - then centers
# - then targets
from pytential import bind, sym
stage1_density_discrs = []
density_discrs = []
for source_name in sources_list:
dd = sym.as_dofdesc(source_name)
discr = places.get_discretization(dd.geometry)
stage1_density_discrs.append(discr)
if use_stage2_discr:
discr = places.get_discretization(dd.geometry,
sym.QBX_SOURCE_QUAD_STAGE2)
density_discrs.append(discr)
# TODO: update code to work for multiple source discretizations
if len(sources_list) != 1:
raise RuntimeError("can only build a tree for a single source")
def _make_centers(discr):
return bind(discr,
sym.interleaved_expansion_centers(discr.ambient_dim))(actx)
stage1_density_discr = stage1_density_discrs[0]
density_discr = density_discrs[0]
from meshmode.dof_array import flatten, thaw
from pytential.utils import flatten_if_needed
sources = flatten(thaw(actx, density_discr.nodes()))
centers = flatten(_make_centers(stage1_density_discr))
targets = [flatten_if_needed(actx, tgt.nodes()) for tgt in targets_list]
queue = actx.queue
particles = tuple(
cl.array.concatenate(dim_coords, queue=queue)
for dim_coords in zip(sources, centers, *targets))
# Counts
nparticles = len(particles[0])
npanels = density_discr.mesh.nelements
nsources = len(sources[0])
ncenters = len(centers[0])
# Each source gets an interior / exterior center.
assert 2 * nsources == ncenters or use_stage2_discr
ntargets = sum(tgt.ndofs for tgt in targets_list)
# Slices
qbx_user_source_slice = slice(0, nsources)
center_slice_start = nsources
qbx_user_center_slice = slice(center_slice_start,
center_slice_start + ncenters)
panel_slice_start = center_slice_start + ncenters
target_slice_start = panel_slice_start
qbx_user_target_slice = slice(target_slice_start,
target_slice_start + ntargets)
# Build tree with sources and centers. Split boxes
# only because of sources.
refine_weights = cl.array.zeros(queue, nparticles, np.int32)
refine_weights[:nsources].fill(1)
refine_weights.finish()
tree, evt = tree_builder(queue,
particles,
max_leaf_refine_weight=MAX_REFINE_WEIGHT,
refine_weights=refine_weights)
# Compute box => particle class relations
flags = refine_weights
del refine_weights
particle_classes = {}
for class_name, particle_slice, fixup in (("box_to_qbx_source",
qbx_user_source_slice, 0),
("box_to_qbx_target",
qbx_user_target_slice,
-target_slice_start),
("box_to_qbx_center",
qbx_user_center_slice,
-center_slice_start)):
flags.fill(0)
flags[particle_slice].fill(1)
flags.finish()
box_to_class = (particle_list_filter.filter_target_lists_in_user_order(
queue, tree, flags).with_queue(actx.queue))
if fixup:
box_to_class.target_lists += fixup
particle_classes[class_name + "_starts"] = box_to_class.target_starts
particle_classes[class_name + "_lists"] = box_to_class.target_lists
del flags
del box_to_class
# Compute panel => source relation
qbx_panel_to_source_starts = cl.array.empty(queue,
npanels + 1,
dtype=tree.particle_id_dtype)
el_offset = 0
node_nr_base = 0
for group in density_discr.groups:
qbx_panel_to_source_starts[el_offset:el_offset + group.nelements] = \
cl.array.arange(queue, node_nr_base,
node_nr_base + group.ndofs,
group.nunit_dofs,
dtype=tree.particle_id_dtype)
node_nr_base += group.ndofs
el_offset += group.nelements
qbx_panel_to_source_starts[-1] = nsources
# Compute panel => center relation
qbx_panel_to_center_starts = (2 * qbx_panel_to_source_starts
if not use_stage2_discr else None)
# Transfer all tree attributes.
tree_attrs = {}
for attr_name in tree.__class__.fields:
try:
tree_attrs[attr_name] = getattr(tree, attr_name)
except AttributeError:
pass
tree_attrs.update(particle_classes)
return TreeWithQBXMetadata(
qbx_panel_to_source_starts=qbx_panel_to_source_starts,
qbx_panel_to_center_starts=qbx_panel_to_center_starts,
qbx_user_source_slice=qbx_user_source_slice,
qbx_user_center_slice=qbx_user_center_slice,
qbx_user_target_slice=qbx_user_target_slice,
nqbxpanels=npanels,
nqbxsources=nsources,
nqbxcenters=ncenters,
nqbxtargets=ntargets,
**tree_attrs).with_queue(None)
