Splitting a Form into parts which do and don't contain a given Argument

[jmbigman]

I have a Form which I would like to separate into two Forms: the terms which contain a given Argument and the terms which do not. I will then modify these Forms separately and add them back together.

As an example, consider the modified Poisson equation
$$\Delta u + \nabla\cdot(\mathbf{b}z)=f$$
$$z=\mathbf{b}\cdot\nabla u$$
and Firedrake pseudocode

VCG = FunctionSpace(msh, "CG", degree)
VDG = FunctionSpace(msh, "DG", degree-1)
M=VCG*VDG

uz = Function(M)
u, z = split(uz)
eta, phi = TestFunctions(M)

# b is a constant vector, f is the forcing

F = inner(grad(eta), grad(u))*dx + inner(grad(eta), b*z)*dx + eta*f*dx \
      + phi*(z - dot(b, grad(u)))*dx

arg_ = phi

# Split into terms which contain arg_: phi*(z - dot(b, grad(u)))*dx
F_w_arg = ??
# and terms which do not: inner(grad(eta), grad(u))*dx + inner(grad(eta), b*z)*dx + eta*f*dx
F_wo_arg = ??

# Modify these two Forms

# Recombine
F = F_w_arg + F_wo_arg

I realize this is unnecessarily complicated for this example, but I am looking for a general method.

ufl.algorithms.formsplitter.extract_blocks nearly achieves what I'm looking for except the resulting terms appear to be associated with individual FunctionSpaces instead of the original MixedFunctionSpace. This caused errors when I tried to use the Form which resulted from adding the pieces together.

I appreciate any help! Please let me know if additional information is needed.

[ksagiyam]

I think you can use UFL replace with as_vector:

from firedrake import *


mesh = UnitIntervalMesh(1)
V0 = FunctionSpace(mesh, "DG", 0)
V1 = FunctionSpace(mesh, "DG", 1)
V = V0 * V1
v = TestFunction(V)
v0, v1 = split(v)
F = Constant(10.) * v0 * dx + Constant(11.) * v1 * dx
F0 = replace(F, {v: as_vector([v0, zero()])})
print(F0)  # { v_0[0] * [10.] } * dx(<Mesh #0>[everywhere], {})
F1 = replace(F, {v: as_vector([zero(), v1])})
print(F1)  # { v_0[1] * [11.] } * dx(<Mesh #0>[everywhere], {})
val = assemble(F0 + F1 - F)
print(val.dat.data_ro)  # (array([0.]), array([0., 0.]))

[jmbigman]

This seems to accomplish what I want. Thank you!