Skip to content

Matrix evaluation updates #89

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
inducer merged 5 commits into from
Jun 3, 2021
Merged

Matrix evaluation updates #89

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3d0d636
use geometry collection to evaluate matrix intg arguments
alexfikl Jun 3, 2021
ba97292
mark matrix p2p flatten with strict=False
alexfikl Jun 3, 2021
4b37a26
use base_kernel when evaluating p2p
alexfikl Jun 3, 2021
9bc58ec
allow weighting p2p matrices by area elements
alexfikl Jun 3, 2021
48e6d1e
allow non-self evaluation in block p2p
alexfikl Jun 3, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 64 additions & 48 deletions pytential/symbolic/matrix.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,21 +40,25 @@ def is_zero(x):
return isinstance(x, (int, float, complex, np.number)) and x == 0


def _get_layer_potential_args(mapper, expr, include_args=None):
def _get_layer_potential_args(actx, places, expr, context=None, 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*.
"""
from pytential import bind
if context is None:
context = {}

kernel_args = {}
for arg_name, arg_expr in expr.kernel_arguments.items():
if (include_args is not None
and arg_name not in include_args):
if include_args is not None and arg_name not in include_args:
continue

kernel_args[arg_name] = mapper.rec(arg_expr)
kernel_args[arg_name] = flatten(
bind(places, arg_expr)(actx, **context),
strict=False)

return kernel_args

Expand Down Expand Up @@ -255,20 +259,6 @@ def __init__(self, actx, dep_expr, other_dep_exprs,
places, context)
self.index_set = index_set

@property
@memoize_method
def _mat_mapper(self):
# mat_mapper is used to compute any kernel arguments that needs to
# be computed on the full discretization, ignoring our index_set,
# e.g the normal in a double layer potential

return MatrixBuilderBase(self.array_context,
self.dep_expr,
self.other_dep_exprs,
self.dep_source,
self.dep_discr,
self.places, self.context)

@property
@memoize_method
def _blk_mapper(self):
Expand Down Expand Up @@ -376,7 +366,8 @@ def map_int_g(self, expr):
raise NotImplementedError("layer potentials on non-variables")

actx = self.array_context
kernel_args = _get_layer_potential_args(self, expr)
kernel_args = _get_layer_potential_args(
actx, self.places, expr, context=self.context)
local_expn = lpot_source.get_expansion_for_qbx_direct_eval(
kernel.get_base_kernel(), (expr.target_kernel,))

Expand All @@ -396,7 +387,7 @@ def map_int_g(self, expr):
dofdesc=expr.target))(actx)

_, (mat,) = mat_gen(actx.queue,
targets=flatten(thaw(target_discr.nodes(), actx)),
targets=flatten(thaw(target_discr.nodes(), actx), strict=False),
sources=flatten(thaw(source_discr.nodes(), actx)),
centers=flatten(centers),
expansion_radii=flatten(radii),
Expand All @@ -406,10 +397,9 @@ def map_int_g(self, expr):
waa = bind(self.places, sym.weights_and_area_elements(
source_discr.ambient_dim,
dofdesc=expr.source))(actx)
mat[:, :] *= actx.to_numpy(flatten(waa))
mat = mat.dot(rec_density)
mat[:, :] *= flatten_to_numpy(actx, waa)

result += mat
result += mat @ rec_density

return result

Expand All @@ -420,12 +410,14 @@ def map_int_g(self, expr):

class P2PMatrixBuilder(MatrixBuilderBase):
def __init__(self, actx, dep_expr, other_dep_exprs,
dep_source, dep_discr, places, context, exclude_self=True):
dep_source, dep_discr, places, context,
weighted=False, exclude_self=True):
super().__init__(actx,
dep_expr, other_dep_exprs, dep_source, dep_discr,
places, context)

self.exclude_self = exclude_self
self.weighted = weighted

def map_int_g(self, expr):
source_discr = self.places.get_discretization(
Expand All @@ -444,29 +436,41 @@ def map_int_g(self, expr):
raise NotImplementedError("layer potentials on non-variables")

# NOTE: copied from pytential.symbolic.primitives.IntG
kernel_args = kernel.get_args() + kernel.get_source_args()
base_kernel = kernel.get_base_kernel()
kernel_args = base_kernel.get_args() + base_kernel.get_source_args()
kernel_args = {arg.loopy_arg.name for arg in kernel_args}

actx = self.array_context
kernel_args = _get_layer_potential_args(self,
expr, include_args=kernel_args)
kernel_args = _get_layer_potential_args(
actx, self.places, expr, context=self.context,
include_args=kernel_args)
if self.exclude_self:
kernel_args["target_to_source"] = actx.from_numpy(
np.arange(0, target_discr.ndofs, dtype=np.int64)
)

from sumpy.p2p import P2PMatrixGenerator
mat_gen = P2PMatrixGenerator(actx.context,
source_kernels=(kernel,),
target_kernels=(expr.target_kernel,),
source_kernels=(base_kernel,),
target_kernels=(expr.target_kernel.get_base_kernel(),),
exclude_self=self.exclude_self)

_, (mat,) = mat_gen(actx.queue,
targets=flatten(thaw(target_discr.nodes(), actx)),
sources=flatten(thaw(source_discr.nodes(), actx)),
targets=flatten(thaw(target_discr.nodes(), actx), strict=False),
sources=flatten(thaw(source_discr.nodes(), actx), strict=False),
**kernel_args)
mat = actx.to_numpy(mat)

result += actx.to_numpy(mat).dot(rec_density)
from meshmode.discretization import Discretization
if self.weighted and isinstance(source_discr, Discretization):
from pytential import bind, sym
waa = bind(self.places, sym.weights_and_area_elements(
source_discr.ambient_dim,
dofdesc=expr.source))(actx)

mat[:, :] *= flatten_to_numpy(actx, waa)

result += mat @ rec_density

return result

Expand Down Expand Up @@ -500,7 +504,6 @@ def map_int_g(self, expr):
raise NotImplementedError

result = 0

for kernel, density in zip(expr.source_kernels, expr.densities):
rec_density = self._blk_mapper.rec(density)
if is_zero(rec_density):
Expand All @@ -510,7 +513,8 @@ def map_int_g(self, expr):
raise NotImplementedError

actx = self.array_context
kernel_args = _get_layer_potential_args(self._mat_mapper, expr)
kernel_args = _get_layer_potential_args(
actx, self.places, expr, context=self.context)
local_expn = lpot_source.get_expansion_for_qbx_direct_eval(
kernel.get_base_kernel(), (expr.target_kernel,))

Expand All @@ -529,7 +533,7 @@ def map_int_g(self, expr):
dofdesc=expr.target))(actx)

_, (mat,) = mat_gen(actx.queue,
targets=flatten(thaw(target_discr.nodes(), actx)),
targets=flatten(thaw(target_discr.nodes(), actx), strict=False),
sources=flatten(thaw(source_discr.nodes(), actx)),
centers=flatten(centers),
expansion_radii=flatten(radii),
Expand All @@ -542,18 +546,21 @@ def map_int_g(self, expr):
waa = flatten(waa)

mat *= waa[self.index_set.linear_col_indices]
result += rec_density * actx.to_numpy(mat)
result += actx.to_numpy(mat) * rec_density

return result


class FarFieldBlockBuilder(MatrixBlockBuilderBase):
def __init__(self, queue, dep_expr, other_dep_exprs, dep_source, dep_discr,
places, index_set, context, exclude_self=False):
places, index_set, context,
weighted=False, exclude_self=False):
super().__init__(queue,
dep_expr, other_dep_exprs, dep_source, dep_discr,
places, index_set, context)

self.exclude_self = exclude_self
self.weighted = weighted

def get_dep_variable(self):
queue = self.array_context.queue
Expand All @@ -568,9 +575,6 @@ def map_int_g(self, expr):
target_discr = self.places.get_discretization(
expr.target.geometry, expr.target.discr_stage)

if source_discr is not target_discr:
raise NotImplementedError

result = 0
for kernel, density in zip(expr.source_kernels, expr.densities):
rec_density = self._blk_mapper.rec(density)
Expand All @@ -581,30 +585,42 @@ def map_int_g(self, expr):
raise NotImplementedError

# NOTE: copied from pytential.symbolic.primitives.IntG
kernel_args = kernel.get_args() + kernel.get_source_args()
base_kernel = kernel.get_base_kernel()
kernel_args = base_kernel.get_args() + base_kernel.get_source_args()
kernel_args = {arg.loopy_arg.name for arg in kernel_args}

actx = self.array_context
kernel_args = _get_layer_potential_args(self._mat_mapper,
expr, include_args=kernel_args)
kernel_args = _get_layer_potential_args(
actx, self.places, expr, context=self.context,
include_args=kernel_args)
if self.exclude_self:
kernel_args["target_to_source"] = actx.from_numpy(
np.arange(0, target_discr.ndofs, dtype=np.int64)
)

from sumpy.p2p import P2PMatrixBlockGenerator
mat_gen = P2PMatrixBlockGenerator(actx.context,
source_kernels=(kernel,),
target_kernels=(expr.target_kernel,),
source_kernels=(base_kernel,),
target_kernels=(expr.target_kernel.get_base_kernel(),),
exclude_self=self.exclude_self)

_, (mat,) = mat_gen(actx.queue,
targets=flatten(thaw(target_discr.nodes(), actx)),
sources=flatten(thaw(source_discr.nodes(), actx)),
targets=flatten(thaw(target_discr.nodes(), actx), strict=False),
sources=flatten(thaw(source_discr.nodes(), actx), strict=False),
index_set=self.index_set,
**kernel_args)

result += rec_density * actx.to_numpy(mat)
from meshmode.discretization import Discretization
if self.weighted and isinstance(source_discr, Discretization):
from pytential import bind, sym
waa = bind(self.places, sym.weights_and_area_elements(
source_discr.ambient_dim,
dofdesc=expr.source))(actx)
waa = flatten(waa)

mat *= waa[self.index_set.linear_col_indices]

result += actx.to_numpy(mat) * rec_density

return result

Expand Down