Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove component tensors #339

Open
wants to merge 13 commits into
base: main
Choose a base branch
from

Conversation

pbrubeck
Copy link
Contributor

@pbrubeck pbrubeck commented Jan 18, 2025

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.

@pbrubeck pbrubeck force-pushed the pbrubeck/remove-component-tensors branch from 2e3b737 to 7351d88 Compare January 18, 2025 13:12
@pbrubeck pbrubeck force-pushed the pbrubeck/remove-component-tensors branch from 7351d88 to d328fe5 Compare January 18, 2025 18:11
@pbrubeck pbrubeck force-pushed the pbrubeck/remove-component-tensors branch from d328fe5 to eae33f2 Compare January 19, 2025 10:10
@pbrubeck pbrubeck force-pushed the pbrubeck/remove-component-tensors branch from d83702d to d1b850e Compare January 22, 2025 14:24
@pbrubeck pbrubeck force-pushed the pbrubeck/remove-component-tensors branch 2 times, most recently from 909326f to 83609f8 Compare January 25, 2025 11:23
@pbrubeck pbrubeck force-pushed the pbrubeck/remove-component-tensors branch from 438c594 to 4713c06 Compare January 26, 2025 17:05
…akeproject/ufl into pbrubeck/remove-component-tensors
@jorgensd
Copy link
Member

jorgensd commented Feb 4, 2025

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

pbrubeck and others added 2 commits February 4, 2025 22:32
Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>
@pbrubeck pbrubeck force-pushed the pbrubeck/remove-component-tensors branch from 56ba52d to ec00649 Compare February 4, 2025 22:40
"""More precise checks for cellwise constants."""
if is_cellwise_constant(o):
return True
degree = map_expr_dag(self.degree_estimator, o)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

@@ -653,9 +662,10 @@ def reference_value(self, o):

def reference_grad(self, o):
"""Differentiate a reference_grad."""
if self.is_cellwise_constant(o):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this correct? Looks like one derivative is lost.

Copy link
Contributor Author

@pbrubeck pbrubeck Feb 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor

@michalhabera michalhabera Feb 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 pbrubeck force-pushed the pbrubeck/remove-component-tensors branch from ee35788 to 35472a4 Compare March 13, 2025 20:54
@pbrubeck
Copy link
Contributor Author

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

@pbrubeck pbrubeck force-pushed the pbrubeck/remove-component-tensors branch from fc76acd to 13e68c8 Compare March 17, 2025 10:06
@pbrubeck pbrubeck requested a review from michalhabera March 17, 2025 15:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants