diff --git a/python/dune/perftool/ufl/transformations/blockpreconditioner.py b/python/dune/perftool/ufl/transformations/blockpreconditioner.py new file mode 100644 index 0000000000000000000000000000000000000000..c16c11520ccef27a5cf23d22ba387846b434b4a0 --- /dev/null +++ b/python/dune/perftool/ufl/transformations/blockpreconditioner.py @@ -0,0 +1,82 @@ +""" Derive block preconditioners from residual forms """ + +from dune.perftool.ufl.modified_terminals import Restriction + +from ufl.algorithms import MultiFunction +from ufl.algorithms.map_integrands import map_integrands + +import ufl.classes as uc +import itertools + + +class OffDiagonalBlockSwitcher(MultiFunction): + def __init__(self, restrictions): + self.restrictions = restrictions + self.res = Restriction.NONE + MultiFunction.__init__(self) + + def expr(self, o): + return self.reuse_if_untouched(o, *tuple(self(op) for op in o.ufl_operands)) + + def positive_restricted(self, o): + self.res = Restriction.POSITIVE + ret = self(o.ufl_operands[0]) + self.rest = Restriction.NONE + if isinstance(ret, uc.Zero): + return ret + else: + return o + + def negative_restricted(self, o): + self.res = Restriction.NEGATIVE + ret = self(o.ufl_operands[0]) + self.res = Restriction.NONE + if isinstance(ret, uc.Zero): + return ret + else: + return o + + def reference_value(self, o): + ret = self(o.ufl_operands[0]) + if isinstance(ret, uc.Zero): + return ret + else: + return o + + def argument(self, o): + if self.res == self.restrictions[o.number()]: + return o + else: + return uc.Zero(shape=o.ufl_shape, + free_indices=o.ufl_free_indices, + index_dimensions=o.ufl_index_dimensions) + + +def list_restriction_tuples(diagonal): + if diagonal: + yield (Restriction.NONE, Restriction.NONE) + + res = (Restriction.POSITIVE, Restriction.NEGATIVE) + amount = 1 if diagonal else 2 + + for rtup in itertools.product(res, res): + if len(set(rtup)) == amount: + yield rtup + + +def _block_jacobian(form, diagonal=True): + assert(len(form.arguments()) == 2) + + forms = [] + for rtup in list_restriction_tuples(diagonal): + forms.append(map_integrands(OffDiagonalBlockSwitcher(rtup), form)) + + return sum(forms) + + +def diagonal_block_jacobian(form): + return _block_jacobian(form) + + +def offdiagonal_block_jacobian(form): + return _block_jacobian(form, False)