def _visualize_refinement(actx: PyOpenCLArrayContext,
discr,
niter,
stage_nr,
stage_name,
flags,
visualize=False):
if not visualize:
return
if stage_nr not in (1, 2):
raise ValueError("unexpected stage number")
flags = actx.to_numpy(flags)
logger.info("for stage %s: splitting %d/%d stage-%d elements", stage_name,
np.sum(flags), discr.mesh.nelements, stage_nr)
from meshmode.discretization.visualization import make_visualizer
vis = make_visualizer(actx, discr, 3)
assert len(flags) == discr.mesh.nelements
flags = flags.astype(bool)
nodes_flags_template = discr.zeros(actx)
nodes_flags = []
for grp in discr.groups:
meg = grp.mesh_el_group
nodes_flags_grp = actx.to_numpy(nodes_flags_template[grp.index])
nodes_flags_grp[flags[meg.element_nr_base:meg.nelements +
meg.element_nr_base]] = 1
nodes_flags.append(actx.from_numpy(nodes_flags_grp))
nodes_flags = DOFArray(actx, tuple(nodes_flags))
vis_data = [
("refine_flags", nodes_flags),
]
if 0:
from pytential import sym, bind
bdry_normals = bind(discr, sym.normal(
discr.ambient_dim))(actx).as_vector(dtype=object)
vis_data.append(("bdry_normals", bdry_normals), )
vis.write_vtk_file(f"refinement-{stage_name}-{niter:03d}.vtu", vis_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 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 arraycontext import flatten
sources = flatten(density_discr.nodes(), actx, leaf_class=DOFArray)
centers = flatten(_make_centers(stage1_density_discr), actx, leaf_class=DOFArray)
targets = [
flatten(tgt.nodes(), actx, leaf_class=DOFArray)
for tgt in targets_list]
queue = actx.queue
particles = tuple(
actx.np.concatenate(dim_coords)
for dim_coords in zip(sources, centers, *targets))
# Counts
nparticles = len(particles[0])
nelements = 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)
element_slice_start = center_slice_start + ncenters
target_slice_start = element_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 = actx.zeros(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 element => source relation
qbx_element_to_source_starts = actx.empty(nelements + 1, tree.particle_id_dtype)
el_offset = 0
node_nr_base = 0
for group in density_discr.groups:
group_element_starts = np.arange(
node_nr_base, node_nr_base + group.ndofs, group.nunit_dofs,
dtype=tree.particle_id_dtype)
qbx_element_to_source_starts[el_offset:el_offset + group.nelements] = \
actx.from_numpy(group_element_starts)
node_nr_base += group.ndofs
el_offset += group.nelements
qbx_element_to_source_starts[-1] = nsources
# Compute element => center relation
qbx_element_to_center_starts = (
2 * qbx_element_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_element_to_source_starts=qbx_element_to_source_starts,
qbx_element_to_center_starts=qbx_element_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,
nqbxelements=nelements,
nqbxsources=nsources,
nqbxcenters=ncenters,
nqbxtargets=ntargets,
**tree_attrs).with_queue(None)