diff --git a/doc/linalg.rst b/doc/linalg.rst index a492e4e8b..b75f7dc5c 100644 --- a/doc/linalg.rst +++ b/doc/linalg.rst @@ -10,6 +10,8 @@ scheme is used: proxy-based skeletonization of the direct solver algorithms. Clusters are represented by a :class:`~pytential.linalg.TargetAndSourceClusterList`. +.. _direct_solver: + Hierarchical Direct Solver -------------------------- @@ -21,4 +23,9 @@ Hierarchical Direct Solver .. automodule:: pytential.linalg.proxy .. automodule:: pytential.linalg.utils +Internal Functionality +---------------------- + +.. automodule:: pytential.linalg.direct_solver_symbolic + .. vim: sw=4:tw=75 diff --git a/pytential/linalg/direct_solver_symbolic.py b/pytential/linalg/direct_solver_symbolic.py new file mode 100644 index 000000000..7c7e6786e --- /dev/null +++ b/pytential/linalg/direct_solver_symbolic.py @@ -0,0 +1,189 @@ +__copyright__ = "Copyright (C) 2022 Alexandru Fikl" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +from pytential.symbolic.mappers import ( + IdentityMapper, OperatorCollector, LocationTagger) + +__doc__ = """ +.. autoclass:: KernelTransformationRemover +.. autoclass:: IntGTermCollector +.. autoclass:: DOFDescriptorReplacer +""" + + +# {{{ KernelTransformationRemover + +class KernelTransformationRemover(IdentityMapper): + r"""A mapper that removes the transformations from the kernel of all + :class:`~pytential.symbolic.primitives.IntG`\ s in the expression. + + This includes source and target derivatives and other such transformations. + Any unnecessary kernel arguments are also removed from + :attr:`~pytential.symbolic.primitives.IntG.kernel_arguments`. + + This mapper is meant to be used in the directs solver for proxy interaction, + where it is not possible to evaluate source or target directional derivatives. + """ + + def __init__(self): + from sumpy.kernel import ( + TargetTransformationRemover, + SourceTransformationRemover) + self.sxr = SourceTransformationRemover() + self.txr = TargetTransformationRemover() + + def map_int_g(self, expr): + target_kernel = self.txr(expr.target_kernel) + source_kernels = tuple([self.sxr(kernel) for kernel in expr.source_kernels]) + if (target_kernel == expr.target_kernel + and source_kernels == expr.source_kernels): + return expr + + # remove all args that come from the source transformations + source_args = { + arg.name for kernel in expr.source_kernels + for arg in kernel.get_source_args()} + kernel_arguments = { + name: self.rec(arg) for name, arg in expr.kernel_arguments.items() + if name not in source_args + } + + return expr.copy(target_kernel=target_kernel, + source_kernels=source_kernels, + densities=self.rec(expr.densities), + kernel_arguments=kernel_arguments) + +# }}} + + +# {{{ IntGTermCollector + +class IntGTermCollector(IdentityMapper): + r"""A mapper that removes all non-:class:`~pytential.symbolic.primitives.IntG` + terms from the expression and all their non-constant factors. + + In particular, an expression of the type + + .. math:: + + \sum_{i = 0}^N f_i(\mathbf{x}, \sigma) + + \sum_{i = 0}^M c_i g_i(\mathbf{x}) \mathrm{IntG}_i(\mathbf{x}) + + is reduced to + + .. math:: + + \sum_{i = 0}^M c_i \mathrm{IntG}_i(\mathbf{x}). + + The intended used of this transformation is to allow the evaluation of + the proxy interactions in the direct solver for a given expression + meant for self-evaluation. + """ + + def map_sum(self, expr): + collector = OperatorCollector() + + children = [] + for child in expr.children: + rec_child = self.rec(child) + if collector(rec_child): + children.append(rec_child) + + from pymbolic.primitives import flattened_sum + return flattened_sum(children) + + def map_product(self, expr): + + collector = OperatorCollector() + + from pymbolic.primitives import is_constant + children_const = [] + children_int_g = [] + for child in expr.children: + if is_constant(child): + children_const.append(child) + else: + rec_child = self.rec(child) + if collector(rec_child): + if children_int_g: + raise RuntimeError( + f"{type(self).__name__}.map_product does not " + "support products of IntGs") + + children_int_g.append(rec_child) + + from pymbolic.primitives import flattened_product + return flattened_product(children_const + children_int_g) + + def map_int_g(self, expr): + return expr + +# }}} + + +# {{{ DOFDescriptorReplacer + +class _LocationReplacer(LocationTagger): + """Unlike :class:`LocationTagger`, this mapper removes the heuristic for + target and source tagging and forcefully replaces existing + :class:`~pytential.symbolic.dof_desc.DOFDescriptor` in the expression. + """ + + def _default_dofdesc(self, dofdesc): + return self.default_target + + def map_int_g(self, expr): + return type(expr)( + expr.target_kernel, expr.source_kernels, + densities=self.operand_rec(expr.densities), + qbx_forced_limit=expr.qbx_forced_limit, + source=self.default_source, target=self.default_target, + kernel_arguments={ + name: self.operand_rec(arg_expr) + for name, arg_expr in expr.kernel_arguments.items() + } + ) + + +class DOFDescriptorReplacer(_LocationReplacer): + r"""Mapper that replaces all the + :class:`~pytential.symbolic.dof_desc.DOFDescriptor`\ s in the expression + with the given ones. + + This mapper is meant to allow for evaluation of proxy interactions in + the direct solver when the given expression is already partially + (or fully) tagged. + + .. automethod:: __init__ + """ + + def __init__(self, source, target): + """ + :param source: a descriptor for all expressions to be evaluated on + the source geometry. + :param target: a descriptor for all expressions to be evaluate on + the target geometry. + """ + super().__init__(target, default_source=source) + self.operand_rec = _LocationReplacer(source, default_source=source) + +# }}} diff --git a/pytential/symbolic/dof_desc.py b/pytential/symbolic/dof_desc.py index 7e1c36235..65cfded7e 100644 --- a/pytential/symbolic/dof_desc.py +++ b/pytential/symbolic/dof_desc.py @@ -222,7 +222,10 @@ def __str__(self) -> str: elif self.geometry == DEFAULT_TARGET: name.append("t") else: - name.append(str(self.geometry)) + name.append( + self.geometry.__name__ + if isinstance(self.geometry, type) + else str(self.geometry)) if self.discr_stage == QBX_SOURCE_STAGE2: name.append("stage2") diff --git a/pytential/symbolic/mappers.py b/pytential/symbolic/mappers.py index c5cd0d71b..d9cb0f78c 100644 --- a/pytential/symbolic/mappers.py +++ b/pytential/symbolic/mappers.py @@ -68,6 +68,8 @@ def rec_int_g_arguments(mapper, expr): return densities, kernel_arguments, changed +# {{{ IdentityMapper + class IdentityMapper(IdentityMapperBase): def map_node_sum(self, expr): operand = self.rec(expr.operand) @@ -137,6 +139,10 @@ def map_interpolation(self, expr): return type(expr)(expr.from_dd, expr.to_dd, operand) +# }}} + + +# {{{ CombineMapper class CombineMapper(CombineMapperBase): def map_node_sum(self, expr): @@ -169,6 +175,10 @@ def map_is_shape_class(self, expr): map_error_expression = map_is_shape_class +# }}} + + +# {{{ Collector class Collector(CollectorBase, CombineMapper): def map_ones(self, expr): @@ -187,6 +197,10 @@ def map_int_g(self, expr): class DependencyMapper(DependencyMapperBase, Collector): pass +# }}} + + +# {{{ EvaluationMapper class EvaluationMapper(EvaluationMapperBase): """Unlike :mod:`pymbolic.mapper.evaluation.EvaluationMapper`, this class @@ -250,8 +264,10 @@ def map_common_subexpression(self, expr): expr.prefix, expr.scope) +# }}} + -# {{{ dofdesc tagging +# {{{ LocationTagger class LocationTagger(CSECachingMapperMixin, IdentityMapper): """Used internally by :class:`ToTargetTagger`.""" @@ -384,6 +400,10 @@ def __init__(self, default_source, default_target): self.operand_rec = LocationTagger(default_source, default_source=default_source) +# }}} + + +# {{{ DiscretizationStageTagger class DiscretizationStageTagger(IdentityMapper): """Descends into an expression tree and changes the @@ -427,7 +447,7 @@ def map_num_reference_derivative(self, expr): # }}} -# {{{ derivative binder +# {{{ DerivativeBinder class DerivativeTaker(Mapper): def __init__(self, ambient_axis): @@ -498,7 +518,7 @@ def take_derivative(self, ambient_axis, expr): # }}} -# {{{ Unregularized preprocessor +# {{{ UnregularizedPreprocessor class UnregularizedPreprocessor(IdentityMapper): @@ -524,7 +544,7 @@ def map_int_g(self, expr): # }}} -# {{{ interpolation preprocessor +# {{{ InterpolationPreprocessor class InterpolationPreprocessor(IdentityMapper): """Handle expressions that require upsampling or downsampling by inserting @@ -596,7 +616,7 @@ def map_int_g(self, expr): # }}} -# {{{ QBX preprocessor +# {{{ QBXPreprocessor class QBXPreprocessor(IdentityMapper): def __init__(self, geometry, places): @@ -656,7 +676,7 @@ def map_int_g(self, expr): # }}} -# {{{ stringifier +# {{{ StringifyMapper def stringify_where(where): return str(prim.as_dofdesc(where)) @@ -769,13 +789,13 @@ def map_is_shape_class(self, expr, enclosing_prec): return "IsShape[{}]({})".format(stringify_where(expr.dofdesc), expr.shape.__name__) -# }}} - class PrettyStringifyMapper( CSESplittingStringifyMapperMixin, StringifyMapper): pass +# }}} + # {{{ graphviz diff --git a/pytential/symbolic/pde/scalar.py b/pytential/symbolic/pde/scalar.py index 0bc0373fb..9ce3cb8f4 100644 --- a/pytential/symbolic/pde/scalar.py +++ b/pytential/symbolic/pde/scalar.py @@ -400,8 +400,10 @@ def Sp(density): if isinstance(self.kernel, HelmholtzKernel): DpS0u = ( sym.Dp( - self.kernel - self.laplace_kernel, - laplace_s_inv_sqrt_w_u, + self.kernel, laplace_s_inv_sqrt_w_u, + qbx_forced_limit=+1, **kwargs) + - sym.Dp( + self.laplace_kernel, laplace_s_inv_sqrt_w_u, qbx_forced_limit=+1, **kwargs) + Dp0S0u) elif isinstance(self.kernel, LaplaceKernel): diff --git a/pytential/symbolic/primitives.py b/pytential/symbolic/primitives.py index 82f12786e..a4a652c71 100644 --- a/pytential/symbolic/primitives.py +++ b/pytential/symbolic/primitives.py @@ -1279,6 +1279,12 @@ class IntG(Expression): where :math:`\sigma_k` is the k-th *density*, :math:`G` is a Green's function, :math:`S_k` are source derivative operators and :math:`T` is a target derivative operator. + + .. attribute:: target_kernel + .. attribute:: source_kernels + .. attribute:: densities + .. attribute:: qbx_forced_limit + .. attribute:: kernel_arguments """ init_arg_names = ("target_kernel", "source_kernels", "densities", diff --git a/test/test_symbolic.py b/test/test_symbolic.py index 63b9edb62..6b1017ad5 100644 --- a/test/test_symbolic.py +++ b/test/test_symbolic.py @@ -123,7 +123,7 @@ def test_mean_curvature(actx_factory, discr_name, resolutions, from meshmode.dof_array import flat_norm h_error = flat_norm(mean_curvature - ref_mean_curvature, np.inf) eoc.add_data_point(h, actx.to_numpy(h_error)) - print(eoc) + logger.info("eoc:\n%s", eoc) order = min([g.order for g in discr.groups]) assert eoc.order_estimate() > order - 1.1 @@ -394,6 +394,107 @@ def test_derivative_binder_expr(): # }}} +# {{{ test_mapper_kernel_transformation_remover + +def _make_operator(ambient_dim: int, op_name: str, k: float, *, side: int = +1): + from sumpy.kernel import LaplaceKernel, HelmholtzKernel + if k == 0: + kernel = LaplaceKernel(ambient_dim) + kernel_arguments = {} + else: + kernel = HelmholtzKernel(ambient_dim) + kernel_arguments = {"k": sym.var("k")} + + import pytential.symbolic.pde.scalar as ops + if op_name == "dirichlet": + op = ops.DirichletOperator( + kernel, side, use_l2_weighting=True, + kernel_arguments=kernel_arguments) + elif op_name == "neumann": + op = ops.NeumannOperator( + kernel, side, use_l2_weighting=True, + kernel_arguments=kernel_arguments) + else: + raise ValueError(f"unknown operator: '{op_name}'") + + return op + + +@pytest.mark.parametrize("op_name", ["dirichlet", "neumann"]) +@pytest.mark.parametrize("k", [0, 5]) +def test_mapper_kernel_transformation_remover(op_name, k): + ambient_dim = 3 + op = _make_operator(ambient_dim, op_name, k) + expr = op.operator(op.get_density_var("sigma")) + + from pytential.linalg.direct_solver_symbolic import ( + OperatorCollector, KernelTransformationRemover) + expr_without_transformations = KernelTransformationRemover()(expr) + intgs = OperatorCollector()(expr_without_transformations) + + def is_base_kernel(knl): + return knl.get_base_kernel() == knl + + for intg in intgs: + assert is_base_kernel(intg.target_kernel) + assert all(is_base_kernel(kernel) for kernel in intg.source_kernels) + +# }}} + + +# {{{ test_mapper_int_g_term_collector + +@pytest.mark.parametrize("op_name", ["dirichlet", "neumann"]) +def test_mapper_int_g_term_collector(op_name, k=0): + ambient_dim = 3 + op = _make_operator(ambient_dim, op_name, k) + expr = op.operator(op.get_density_var("sigma")) + + from pytential.linalg.direct_solver_symbolic import IntGTermCollector + expr_only_intgs = IntGTermCollector()(expr) + + # FIXME: how to check this did something? + sigma = sym.cse(op.get_density_var("sigma") / op.get_sqrt_weight()) + if op_name == "dirichlet": + expected_expr = -1 * sym.D(op.kernel, sigma) + elif op_name == "neumann": + int_g = sym.S(op.kernel, sigma, qbx_forced_limit="avg") + expected_expr = sym.div([int_g] * ambient_dim) + else: + raise ValueError(f"unknown operator name: {op_name}") + + assert expr_only_intgs == expected_expr + +# }}} + + +# {{{ test_mapper_dof_descriptor_replacer + +@pytest.mark.parametrize("op_name", ["dirichlet", "neumann"]) +@pytest.mark.parametrize("k", [0, 5]) +def test_mapper_dof_descriptor_replacer(op_name, k): + ambient_dim = 3 + op = _make_operator(ambient_dim, op_name, k) + expr = op.operator(op.get_density_var("sigma")) + + from pytential.symbolic.mappers import ToTargetTagger + from pytential.linalg.direct_solver_symbolic import DOFDescriptorReplacer + source_dd = sym.as_dofdesc(sym.DEFAULT_SOURCE) + target_dd = sym.as_dofdesc(sym.DEFAULT_TARGET) + tagged_expr = ToTargetTagger(source_dd, target_dd)(expr) + + source_new_dd = sym.as_dofdesc("source") + target_new_dd = sym.as_dofdesc("target") + replaced_expr = DOFDescriptorReplacer(source_new_dd, target_new_dd)(tagged_expr) + + from testlib import DOFDescriptorCollector + collector = DOFDescriptorCollector() + assert collector(tagged_expr) == {source_dd, target_dd} + assert collector(replaced_expr) == {source_new_dd, target_new_dd} + +# }}} + + # You can test individual routines by typing # $ python test_symbolic.py 'test_routine()' diff --git a/test/testlib.py b/test/testlib.py new file mode 100644 index 000000000..7ad102023 --- /dev/null +++ b/test/testlib.py @@ -0,0 +1,42 @@ +from functools import reduce + +from pytential.symbolic.mappers import Collector + + +class DOFDescriptorCollector(Collector): + r"""Gathers all the :class:`~pytential.symbolic.dof_desc.DOFDescriptor`\ s + in an expression. + """ + + def map_ones(self, expr): + return {expr.dofdesc} + + map_is_shape_class = map_ones + map_q_weight = map_ones + map_node_coordinate_component = map_ones + + def map_num_reference_derivative(self, expr): + return {expr.dofdesc} | self.rec(expr.operand) + + def map_interpolation(self, expr): + return {expr.from_dd, expr.to_dd} | self.rec(expr.operand) + + def map_node_sum(self, expr): + return self.rec(expr.operand) + + map_node_max = map_node_sum + map_node_min = map_node_sum + + def map_elementwise_sum(self, expr): + return {expr.dofdesc} | self.rec(expr.operand) + + map_elementwise_max = map_elementwise_sum + map_elementwise_min = map_elementwise_sum + + def map_int_g(self, expr): + import operator + return ({expr.source, expr.target} + | reduce(operator.or_, (self.rec(d) for d in expr.densities)) + | reduce(operator.or_, + (self.rec(v) for v in expr.kernel_arguments.values()), set()) + )