Interpolated Test Functions in Variational Forms: Handling post_jacobian_callback Modifications with Fieldsplit Solvers and R Spaces

[lrtfm]

I am using Firedrake to solve a variational problem involving an interpolated test function, for example
dot(grad(u), grad(I_h(g*v))) * dx,
where $u$ and $v$ are the trial and test functions, $g$ is a known function, and $I_h$ denotes an interpolation operator.

I can handle this formulation by using post_jacobian_callback when defining the solver. However, when the problem involves an $\mathbb{R}$ (Real) space, this approach no longer works. In this case, the solver uses field splitting, and the subproblems obtained after splitting cannot apply the same post_jacobian_callback modification.

I tried patching the _SNESContext.split function to add callback functions to the subproblems, modifying both the AssembledPC to correct the matrix and the ImplicitMatrixContext to correct the operator action. However, the results still seem incorrect. Am I missing something?

Do you have any suggestions on how to handle this situation?

[lrtfm]

Inspired by @colinjcotter, I have figured out how to implement the term involving grad for $P1$ elements. However, I still do not know how to handle the term without a gradient.

---

Assume that $u$ and $v$ are the trial and test functions in the $P1$ finite element space (CG1).
We would like to implement the bilinear forms

$$ \int_\Omega \nabla u \cdot \nabla I_h(g v)\ dx $$

and

$$ \int_\Omega u I_h(g v)\ dx $$

in Firedrake, where $I_h$ is the interpolation operator into the $P1$ element space, and $g$ is a given function.

Implementing the bilinear form with a gradient

Element-wise formulation of the gradient term

For the gradient term, on an element $K$, we have

$$ \begin{aligned} \int_K \nabla u \cdot \nabla I_h(g v), dx &= \sum_{1 \le i \le n,\ 1 \le j \le n} \int_K u_i \nabla \phi_i(x) \cdot g_j v_j \nabla \phi_j(x)\ dx \\ &= \sum_{1 \le i \le n,\ 1 \le j \le n} u_i v_j \int_K \nabla \phi_i(x) \cdot g_j \nabla \phi_j(x)\ dx . \end{aligned} $$

Here, $n = \mathrm{dim} + 1$ is the number of vertices of the element $K$, and $\phi_i$ is the basis function associated with vertex $i$.
Using the fact that $\nabla \phi_i$ is constant on the element $K$, the element matrix associated with the gradient term can be computed as

$$ \begin{aligned} \int_K \nabla \phi_i(x) \cdot \nabla \phi_j(x)\ g_j\ dx &= \nabla \phi_i\big|_{K} \cdot \nabla \phi_j\big|_{K}\ g_j\ |K| \\ &= \sum_{k=1}^n \nabla \phi_i\big|_{x_k} \cdot \nabla \phi_j\big|_{x_k} g_j \phi_j(x_k)\ |K| \\ &= \sum_{k=1}^n \nabla \phi_i\big|_{x_k} \cdot \nabla \phi_j\big|_{x_k} g(x_k) \phi_j(x_k)\ |K| \\ &= n \sum_{k=1}^n \frac{1}{n} \nabla \phi_i\big|_{x_k} \cdot \nabla \phi_j\big|_{x_k} g(x_k) \phi_j(x_k)\ |K| \\ &= n \left( \nabla \phi_i(x) \cdot f(x), g(x) \phi_j(x) \right)_h . \end{aligned} $$

Here, $f$ is a vector-valued function such that $f|_{K}(x_k) = \nabla \phi_k$ for vertex $k$ of the element $K$.
It belongs to

VectorFunctionSpace(mesh, "DG", 1)

Firedrake implementation

We compute $f$, denoted as grad_v, using the following code:

V_DG1 = FunctionSpace(mesh, "DG", 1, variant="equispaced")
v_dg1 = TestFunction(V_DG1)
area = CellVolume(mesh)

V_v_DG1 = VectorFunctionSpace(mesh, "DG", 1, variant="equispaced")
grad_v = Function(V_v_DG1)
grad_v_dg1 = assemble(grad(v_dg1)[0]/area*dx)
dim = mesh.geometric_dimension()

for i in range(dim):
    assemble(grad(v_dg1)[i]/area*dx, tensor=grad_v_dg1)
    grad_v.dat.data[:, i] = grad_v_dg1.dat.data[:]

With this construction, the gradient bilinear form can be implemented as

(dim + 1) * inner(grad(u), g * v * grad_v) * dx(scheme=vertex_rule)

where dim = mesh.topology_dimension(), and vertex_rule specifies a vertex-based quadrature scheme.
The vertex_rule can be obtained by

def get_vertex_qrule(topological_dim):
    import FIAT, finat
    dim2cell = {
        1: FIAT.reference_element.UFCInterval(),
        2: FIAT.reference_element.UFCTriangle(),
        3: FIAT.reference_element.UFCTetrahedron()
    }
    ref_cell = dim2cell.get(topological_dim, None)
    if ref_cell is None:
        raise NotImplementedError
    point_set = finat.quadrature.PointSet(ref_cell.vertices)
    weights = [1 / math.factorial(len(ref_cell.vertices))] * len(ref_cell.vertices)

    return finat.quadrature.QuadratureRule(point_set, weights)