Cancel contractions of the Jacobian with its inverse#68
Closed
pbrubeck wants to merge 2 commits into
Closed
Conversation
This was referenced Jul 9, 2026
Preserve Jacobian, JacobianInverse and JacobianDeterminant during geometry lowering, and cancel J-Jinv contractions into Kronecker deltas before the inverse is expanded into individual matrix entries. Eliminate the resulting Identity tensors by contraction against the remaining factors, and cancel reciprocal factors within products, such as detJ**2 * (1/detJ)**2. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
pbrubeck
force-pushed
the
pbrubeck/cancel-jacobian
branch
from
23cffd8 to
0651e63
Compare
pbrubeck
commented
Jul 10, 2026
| return result | ||
|
|
||
|
|
||
| class IndexSumSimplifier(MultiFunction): |
pbrubeck
commented
Jul 10, 2026
| return Power(base, as_ufl(exponent)) | ||
|
|
||
|
|
||
| class ReciprocalCanceller(MultiFunction): |
pbrubeck
commented
Jul 10, 2026
| return result | ||
|
|
||
|
|
||
| def cancel_jacobian_products(o): |
IndexSumSimplifier (and its JacobianCanceller/IdentityEliminator subclasses) and ReciprocalCanceller were plain MultiFunctions; convert them to DAGTraverser, matching the convention used by the rest of the newer algorithms modules (apply_derivatives.py, apply_coefficient_split.py). This switches the driver function from map_integrand_dags (the MultiFunction-compatible entry point) to map_integrands, since a DAGTraverser is called directly rather than through map_expr_dag. Add test/test_cancel_jacobian_products.py, covering JacobianCanceller and IdentityEliminator on both contraction orders (J.K and K.J), ReciprocalCanceller on both a fully- and a partially-cancelling reciprocal product, the full pass end-to-end on a div-div form over a contravariant Piola-mapped (Raviart-Thomas) element (verifying both the size reduction and that the opt-in flag defaults to off), and a no-op check on a form with no Jacobian contractions to cancel.
Author
|
Superseded by FEniCS#496, which targets origin (FEniCS/ufl) instead of this fork. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Piola-mapped elements generate expressions of the form
J * inv(J)that were not being simplified, becauseJacobianInversewas expanded into individual matrix entries too early in the form preprocessing.This PR delays that expansion: during geometry lowering,
Jacobian,JacobianInverse, andJacobianDeterminantare preserved as opaque terminals (which also letsis_cellwise_constantfold their derivatives to zero early on affine cells). A new algorithm modulecancel_jacobian_productsthen simplifies the integrand in three traversals:JacobianCanceller:IndexSum(J[a, k] * K[k, b] * factors, k) -> Identity[a, b] * factors(and the K·J order, which is also valid for pseudo-inverses on immersed manifolds). Contractions are chased through nestedIndexSums by interchanging the order of summation.IdentityEliminator:IndexSum(Identity[a, k] * factors, k) -> factors[k -> a], and folds fixed-index Identity entries into scalar constants. This also cleans up user-written Identity contractions.ReciprocalCanceller: cancels reciprocal factors within products, such asdetJ**2 * (1/detJ)**2 -> 1.The surviving Jacobian quantities are lowered as usual afterwards, so form compilers see no new terminal types. The behaviour is opt-in via a new
compute_form_datakwargdo_cancel_jacobian_products(default off, so ffcx is unaffected).For the Johnson-Mercier element, the div-div integrand shrinks from 9630 to 664 characters, becoming literally
(1/detJ**2) J : rdiv(sigma)on affine cells. Together with companion GEM changes this gives up to -42% flops and -51% compile time (see companion PRs in firedrakeproject/fiat and firedrakeproject/firedrake).🤖 Generated with Claude Code