Remove component tensors

[pbrubeck]

Consider the bilinear form
a(v, u) = inner(u, v)*dx + inner(div(u), div(v))*dx
where u is a (doubly) Piola mapped tensor in H(div, S) discretized with linear Johnson-Mercier macroelements.

Here we introduce two enhacements targeted to optimize code generation for a(v, u):

1. We simplify RefGrad(RefGrad(x)) -> 0 on affine simplex meshes. These zero terms would arise from the product rule on the doubly Piola mapped tensor.
2. We eagerly remove ComponentTensors, to make the UFL easier to traverse by the Form compiler.

Since a Form always maps to a scalar, it is always possible to remove ComponentTensor nodes by pushing Indexed inside expressions and simplifying Indexed(ComponentTensor) with the approapiate index replacements.

This PR makes more agressive simplifications within the Indexed constructor, so it's pushed inside sums, and adds a multifunction to handle index replacements. The removal of ComponentTensors is applied as a final preprocessing step in compute_form_data.

[jorgensd]

In general this looks sensible to me. Some minor comments/suggestions only.

[michalhabera]

You've introduced here that is_cellwise_constant logic is based on degree estimation, which is the only place in the code. How does it work, since the degree estimation for geometric quantities calls is_cellwise_constant itself?

ufl/ufl/algorithms/estimate_degrees.py

Lines 46 to 56 in 1ab30c2

	def geometric_quantity(self, v):
	"""Apply to geometric_quantity.
	Some geometric quantities are cellwise constant. Others are
	nonpolynomial and thus hard to estimate.
	"""
	if is_cellwise_constant(v):
	return 0
	else:
	# As a heuristic, just returning domain degree to bump up degree somewhat
	return extract_unique_domain(v).ufl_coordinate_element().embedded_superdegree

If there is no infinite loop, could this be used in the generic cell-wise constant check? Or is it that only GradRuleset does not cause infinite loop?

[pbrubeck]

I first tried to add this logic to the generic is_cellwise_constant in #334, but simplifying grad(cell_wise_constant) on instantiation turns out to be problematic for nested grad expression simply because grad(f) expects f to have a domain and Zero does not have a domain, so grad(grad(grad(linear))) would simplify to grad(Zero).

[michalhabera]

Is this correct? Looks like one derivative is lost.

[pbrubeck]

Here we are simplifying grad(o) where o = grad(f), if o is constant then we should return zero with the correct shape.

[michalhabera]

I see, thanks. I tried following simple UFL code

import ufl
import ufl.algorithms
import ufl.algorithms.apply_algebra_lowering
import ufl.algorithms.apply_function_pullbacks
import ufl.utils.formatting


mesh_el = ufl.finiteelement.FiniteElement("P", ufl.Cell("triangle"), 1, (2, ), ufl.pullback.IdentityPullback(), ufl.sobolevspace.H1)
mesh = ufl.Mesh(mesh_el)

V_el = ufl.finiteelement.FiniteElement("JM", ufl.Cell("triangle"), 1, (2, 2), ufl.pullback.DoubleCovariantPiola(), ufl.sobolevspace.L2)
V = ufl.FunctionSpace(mesh, V_el)
u = ufl.Coefficient(V)
x = ufl.SpatialCoordinate(mesh)
a = ufl.div(u) 

a = ufl.algorithms.apply_algebra_lowering.apply_algebra_lowering(a)
a = ufl.algorithms.apply_derivatives.apply_derivatives(a)
a = ufl.algorithms.apply_function_pullbacks.apply_function_pullbacks(a)
a = ufl.algorithms.apply_derivatives.apply_derivatives(a)

print(ufl.utils.formatting.tree_format(a))

which mimics the order in compute_form_data. Yet, I do not see the ReferenceGrad applied to JacobianInverse in the tree. Geometric quantities makes it correctly outside of the ReferenceGrad. For non-affine cell ReferenceGrad(JacobianInverse) is present as expected. Do you have UFL snippet that shows the problem?

But lets assume that such node will exist and remain unsimplified even if it shouldn't. Would ReferenceGradruleset.jacobian_inverse be more appropriate place for improved constness check?

[michalhabera]

What is always confusing is the order of preprocessing operations in compute_form_data. So my understanding is that for such expressions as div(u) where u is (some) Piola mapped you we essentially do

1. apply algebra lowering, div(u) -> sum(grad(u)[i]),
2. 1st derivate pass, grad(u) -> grad(u),
3. function pullbacks, grad(u) -> grad(piola_maps*reference_value(u)),
4. 2nd derivative pass, grad(piola_maps*u) -> piola_maps*reference_grad(reference_value(u)),
5. 3rd derivative pass, piola_maps*reference_grad(reference_value(u)) -> piola_maps*reference_grad(reference_value(u))

In the code above I am seeing GradRuleset.jacobian_inverse hit constness of JacobianInverse, that makes it in step 4.

[pbrubeck]

In Firedrake we are getting lots of grad(grad(x[j])). This could be because we do not include JacobianInverse in preserve_geometry_types https://github.com/firedrakeproject/firedrake/blob/96501362c73e17e28131294ccc7ee611b9b64e64/tsfc/ufl_utils.py#L38

[pbrubeck]

I have removed the degree estimation, and handled RefGrad(RefGrad(x)) -> 0 as a special case. @michalhabera could you review again?