Re: [firedrake] replacing indexed mixed coefficients

23 Nov 2018

...
On 23 Nov 2018, at 09:19, Cotter, Colin J <colin.cotter@imperial.ac.uk> wrote:
Dear Firedrake,
  In the Gusto equations rewrite project we are bumping into an issue with replacing components of mixed coefficients/test functions.
The setup is that we want to write down a UFL form describing a mixed problem, and extract the subproblem for one of the components. The code that describes this is:
from firedrake import *
mesh = UnitSquareMesh(10,10)
V = FunctionSpace(mesh,"CG",1)
W = V * V
w1, w2 = TestFunctions(W)
u = Function(W)
u1, u2 = split(u)
F = (u1*w1 + u2*w2)*dx
w = TestFunction(V)
u0 = Function(V)
F1 = replace(F, {w1:0})
F2 = replace(F1, {w2:w})
F3 = replace(F2, {u2,u0})
This does not work because w1 is not a terminal, it is an indexed mixed test function. I think what is happening is that the index is the terminal, but I don't really know what I am talking about.
What is the solution to this?
For extract subblocks of arguments:

from firedrake.formmanipulation import ExtractSubBlock

splitter = ExtractSubBlock()

You can now use the split method to extract subblocks of the form indexed by test and trial function position. The interface is a little internal but, if I understand correctly, what you want is to ignore the w1 test function and extract the piece tested against w2 on its own, followed by replacing u2 by u:

# Get the bit of the form whose test function has the index 1.
F2 = splitter.split(F, ((1, ), ))

(This function can handle extracting arbitrary subblocks, useful if you wanted to extract the block corresponding to the 0th and 2nd test functions, say).

Now you want to replace the coefficient, by ufl.replace doesn't like the u2 is indexed either.

So you need to write a tiny multifunction (we could probably put this in Firedrake or UFL if we think through the consequences properly).

from ufl.corealg.multifunction import MultiFunction
from ufl.algorithms.map_integrads import map_integrand_dags

class Replacer(MultiFunction):
    def __init__(self, mapping):
        super().__init__()
        self.mapping = mapping

    def expr(self, o, *args):
        if o in self.mapping:
            return self.mapping[o]
        else:
            return self.reuse_if_untouched(o, *args)

replacer = Replacer({u2: u0})

F3 = map_integrand_dags(replacer, F2)

I think in this case, one can use this replacer to kill all the pieces, but I'm not sure it works in the more general settings:

replace_all = Replacer({u2: u0, w1: zero(w1.ufl_shape), w2: w})

F3 = map_integrand_dags(replace_all, F)

Lawrence

Re: [firedrake] replacing indexed mixed coefficients

Lawrence Mitchell