def op_group_features(self, expr):
from sumpy.kernel import AxisTargetDerivativeRemover
result = (
expr.source, expr.density,
AxisTargetDerivativeRemover()(expr.kernel),
)
return result
def map_int_g(self, expr, name_hint=None):
try:
return self.expr_to_var[expr]
except KeyError:
# make sure operator assignments stand alone and don't get muddled
# up in vector arithmetic
density_var = self.assign_to_new_var(self.rec(expr.density))
group = self.group_to_operators[self.op_group_features(expr)]
names = [self.get_var_name() for op in group]
kernels = sorted({op.kernel for op in group}, key=repr)
kernel_to_index = {kernel: i for i, kernel in enumerate(kernels)}
from pytools import single_valued
from sumpy.kernel import AxisTargetDerivativeRemover
atdr = AxisTargetDerivativeRemover()
base_kernel = single_valued(atdr(kernel) for kernel in kernels)
for op in group:
assert op.qbx_forced_limit in [-2, -1, None, 1, 2]
kernel_arguments = {
arg_name: self.rec(arg_val)
for arg_name, arg_val in expr.kernel_arguments.items()
}
outputs = [
PotentialOutput(
name=name,
kernel_index=kernel_to_index[op.kernel],
target_name=op.target,
qbx_forced_limit=op.qbx_forced_limit,
) for name, op in zip(names, group)
]
self.code.append(
ComputePotentialInstruction(
outputs=outputs,
kernels=tuple(kernels),
kernel_arguments=kernel_arguments,
base_kernel=base_kernel,
density=density_var,
source=expr.source,
priority=max(getattr(op, "priority", 0) for op in group),
dep_mapper_factory=self.dep_mapper_factory))
from pymbolic.primitives import Variable
for name, group_expr in zip(names, group):
self.expr_to_var[group_expr] = Variable(name)
return self.expr_to_var[expr]
def get_fmm_kernel(self, kernels):
fmm_kernel = None
from sumpy.kernel import AxisTargetDerivativeRemover
for knl in kernels:
candidate_fmm_kernel = AxisTargetDerivativeRemover()(knl)
if fmm_kernel is None:
fmm_kernel = candidate_fmm_kernel
else:
assert fmm_kernel == candidate_fmm_kernel
return fmm_kernel
