def to_numpy(self, actx: PyOpenCLArrayContext) -> "BlockProxyPoints":
from arraycontext import to_numpy
from dataclasses import replace
return replace(self,
points=to_numpy(self.points, actx),
centers=to_numpy(self.centers, actx),
radii=to_numpy(self.radii, actx))
def __init__(self,
dcoll: DiscretizationCollection,
array_container: ArrayOrContainerT,
remote_rank, tag=None):
actx = get_container_context_recursively(array_container)
btag = BTAG_PARTITION(remote_rank)
local_bdry_data = project(dcoll, "vol", btag, array_container)
comm = dcoll.mpi_communicator
self.dcoll = dcoll
self.array_context = actx
self.remote_btag = btag
self.bdry_discr = dcoll.discr_from_dd(btag)
self.local_bdry_data = local_bdry_data
self.local_bdry_data_np = \
to_numpy(flatten(self.local_bdry_data, actx), actx)
self.tag = self.base_tag
if tag is not None:
self.tag += tag
# Here, we initialize both send and recieve operations through
# mpi4py `Request` (MPI_Request) instances for comm.Isend (MPI_Isend)
# and comm.Irecv (MPI_Irecv) respectively. These initiate non-blocking
# point-to-point communication requests and require explicit management
# via the use of wait (MPI_Wait, MPI_Waitall, MPI_Waitany, MPI_Waitsome),
# test (MPI_Test, MPI_Testall, MPI_Testany, MPI_Testsome), and cancel
# (MPI_Cancel). The rank-local data `self.local_bdry_data_np` will have its
# associated memory buffer sent across connected ranks and must not be
# modified at the Python level during this process. Completion of the
# requests is handled in :meth:`finish`.
#
# For more details on the mpi4py semantics, see:
# https://mpi4py.readthedocs.io/en/stable/overview.html#nonblocking-communications
#
# NOTE: mpi4py currently (2021-11-03) holds a reference to the send
# memory buffer for (i.e. `self.local_bdry_data_np`) until the send
# requests is complete, however it is not clear that this is documented
# behavior. We hold on to the buffer (via the instance attribute)
# as well, just in case.
self.send_req = comm.Isend(self.local_bdry_data_np,
remote_rank,
tag=self.tag)
self.remote_data_host_numpy = np.empty_like(self.local_bdry_data_np)
self.recv_req = comm.Irecv(self.remote_data_host_numpy,
remote_rank,
tag=self.tag)
def gather_block_neighbor_points(
actx: PyOpenCLArrayContext,
discr: Discretization,
pxy: BlockProxyPoints,
max_particles_in_box: Optional[int] = None) -> BlockIndexRanges:
"""Generate a set of neighboring points for each range of points in
*discr*. Neighboring points of a range :math:`i` are defined
as all the points inside the proxy ball :math:`i` that do not also
belong to the range itself.
"""
if max_particles_in_box is None:
# FIXME: this is a fairly arbitrary value
max_particles_in_box = 32
# {{{ get only sources in indices
@memoize_in(actx,
(gather_block_neighbor_points, discr.ambient_dim, "picker_knl")
)
def prg():
knl = lp.make_kernel(
"{[idim, i]: 0 <= idim < ndim and 0 <= i < npoints}",
"""
result[idim, i] = ary[idim, srcindices[i]]
""",
[
lp.GlobalArg("ary",
None,
shape=(discr.ambient_dim, "ndofs"),
dim_tags="sep,C"),
lp.ValueArg("ndofs", np.int64), ...
],
name="picker_knl",
assumptions="ndim>=1 and npoints>=1",
fixed_parameters=dict(ndim=discr.ambient_dim),
lang_version=MOST_RECENT_LANGUAGE_VERSION,
)
knl = lp.tag_inames(knl, "idim*:unr")
knl = lp.split_iname(knl, "i", 64, outer_tag="g.0")
return knl
_, (sources, ) = prg()(actx.queue,
ary=flatten(discr.nodes(),
actx,
leaf_class=DOFArray),
srcindices=pxy.srcindices.indices)
# }}}
# {{{ perform area query
from boxtree import TreeBuilder
builder = TreeBuilder(actx.context)
tree, _ = builder(actx.queue,
sources,
max_particles_in_box=max_particles_in_box)
from boxtree.area_query import AreaQueryBuilder
builder = AreaQueryBuilder(actx.context)
query, _ = builder(actx.queue, tree, pxy.centers, pxy.radii)
# find nodes inside each proxy ball
tree = tree.get(actx.queue)
query = query.get(actx.queue)
# }}}
# {{{ retrieve results
from arraycontext import to_numpy
pxycenters = to_numpy(pxy.centers, actx)
pxyradii = to_numpy(pxy.radii, actx)
indices = pxy.srcindices
nbrindices = np.empty(indices.nblocks, dtype=object)
for iblock in range(indices.nblocks):
# get list of boxes intersecting the current ball
istart = query.leaves_near_ball_starts[iblock]
iend = query.leaves_near_ball_starts[iblock + 1]
iboxes = query.leaves_near_ball_lists[istart:iend]
if (iend - istart) <= 0:
nbrindices[iblock] = np.empty(0, dtype=np.int64)
continue
# get nodes inside the boxes
istart = tree.box_source_starts[iboxes]
iend = istart + tree.box_source_counts_cumul[iboxes]
isources = np.hstack([np.arange(s, e) for s, e in zip(istart, iend)])
nodes = np.vstack([s[isources] for s in tree.sources])
isources = tree.user_source_ids[isources]
# get nodes inside the ball but outside the current range
# FIXME: this assumes that only the points in `pxy.srcindices` should
# count as neighbors, not all the nodes in the discretization.
# FIXME: it also assumes that all the indices are sorted?
center = pxycenters[:, iblock].reshape(-1, 1)
radius = pxyradii[iblock]
mask = ((la.norm(nodes - center, axis=0) < radius)
& ((isources < indices.ranges[iblock])
| (indices.ranges[iblock + 1] <= isources)))
nbrindices[iblock] = indices.indices[isources[mask]]
# }}}
from pytential.linalg import make_block_index_from_array
return make_block_index_from_array(indices=nbrindices)
def __call__(self, actx: PyOpenCLArrayContext, source_dd,
indices: BlockIndexRanges, **kwargs) -> BlockProxyPoints:
"""Generate proxy points for each block in *indices* with nodes in
the discretization *source_dd*.
:arg source_dd: a :class:`~pytential.symbolic.primitives.DOFDescriptor`
for the discretization on which the proxy points are to be
generated.
"""
from pytential import sym
source_dd = sym.as_dofdesc(source_dd)
discr = self.places.get_discretization(source_dd.geometry,
source_dd.discr_stage)
# {{{ get proxy centers and radii
sources = flatten(discr.nodes(), actx, leaf_class=DOFArray)
knl = self.get_centers_knl(actx)
_, (centers_dev, ) = knl(actx.queue,
sources=sources,
srcindices=indices.indices,
srcranges=indices.ranges)
knl = self.get_radii_knl(actx)
_, (radii_dev, ) = knl(actx.queue,
sources=sources,
srcindices=indices.indices,
srcranges=indices.ranges,
radius_factor=self.radius_factor,
proxy_centers=centers_dev,
**kwargs)
# }}}
# {{{ build proxy points for each block
from arraycontext import to_numpy
centers = np.vstack(to_numpy(centers_dev, actx))
radii = to_numpy(radii_dev, actx)
nproxy = self.nproxy * indices.nblocks
proxies = np.empty((self.ambient_dim, nproxy), dtype=centers.dtype)
pxy_nr_base = 0
for i in range(indices.nblocks):
points = radii[i] * self.ref_points + centers[:, i].reshape(-1, 1)
proxies[:, pxy_nr_base:pxy_nr_base + self.nproxy] = points
pxy_nr_base += self.nproxy
# }}}
pxyindices = np.arange(0, nproxy, dtype=indices.indices.dtype)
pxyranges = np.arange(0, nproxy + 1, self.nproxy)
from arraycontext import freeze, from_numpy
from pytential.linalg import make_block_index_from_array
return BlockProxyPoints(
lpot_source=self.places.get_geometry(source_dd.geometry),
srcindices=indices,
indices=make_block_index_from_array(pxyindices, pxyranges),
points=freeze(from_numpy(proxies, actx), actx),
centers=freeze(centers_dev, actx),
radii=freeze(radii_dev, actx),
)
FACE_RESTR_INTERIOR, "all_faces",
flux(dcoll, interior_tpair)))
duh_by_dt = op.inverse_mass(dcoll,
np.dot([2 * np.pi], Su) - lift)
# forward euler time step
uh = uh + dt * duh_by_dt
t += dt
# ENDEXAMPLE
# Plot the solution:
def u_exact(x, t):
return actx.np.sin(x[0] - 2 * np.pi * t)
assert op.norm(dcoll,
uh - u_exact(x_vol, t_final),
p=2) <= 0.1
import matplotlib.pyplot as plt
from arraycontext import to_numpy
plt.plot(to_numpy(actx.np.ravel(x_vol[0][0]), actx),
to_numpy(actx.np.ravel(uh[0]), actx), label="Numerical")
plt.plot(to_numpy(actx.np.ravel(x_vol[0][0]), actx),
to_numpy(actx.np.ravel(u_exact(x_vol, t_final)[0]), actx), label="Exact")
plt.xlabel("$x$")
plt.ylabel("$u$")
plt.legend()
plt.show()
def to_numpy(self, ary):
from arraycontext import to_numpy
return to_numpy(ary, self.geo_data._setup_actx)