How to use UFL PartExtracter

Question

I (think) I want to use PartExtracter to split a form into parts involving different arguments. The result I would like in a ufl.Form a_u1 is the first part of the ufl.Form a.

Any ideas? I'm guessing one of the UFL masters knows how to do this.

from dolfin import *

mesh = UnitSquareMesh(10, 10)
V = FunctionSpace(mesh, "CG", 1)
U = MixedFunctionSpace([V, V])

u = TrialFunction(U)
v = TestFunction(U)
u1, u2 = split(u)
v1, v2 = split(v)

a = inner(u1, v1)*dx + inner(u2, v2)*dx
print dir(a)

import ufl.algorithms.formtransformations
extractor = ufl.algorithms.formtransformations.PartExtracter(u1)
a_u1 = extractor.visit(a)

The above code fails with the error:

Traceback (most recent call last):
  File "extractor.py", line 17, in <module>
    a_u1 = extractor.visit(a)
  File "/home/fenics/build/lib/python2.7/site-packages/ufl/algorithms/transformer.py", line 94, in visit
    h, visit_children_first = self._handlers[o._ufl_typecode_]
AttributeError: 'Form' object has no attribute '_ufl_typecode_'

Comments

Well, I've learned a lot about UFL, and it seems that this is not an easy task! I believe it is tricky because when you have an Expression defined on an Argument created with a MixedElement the next Operator up is an Index object which extracts the component of the Argument. Therefore one cannot just ask PartExtracter for u1 as that never actually exists in the Expression tree.

Interestingly, there is some manipulation in Firedrake that can separate out UFL Expressions by Argument into new trees, but it doesn't work with DOLFIN as-is.

Another thing that is not possible in UFL is to replace a non-Terminal operator.