Cancel contractions of the Jacobian with its inverse#496
Open
pbrubeck wants to merge 3 commits into
Open
Conversation
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>
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.
This was referenced Jul 11, 2026
pbrubeck
commented
Jul 11, 2026
The div-div test only checked that the cancelled integrand's string representation was shorter than the uncancelled one -- a check that would pass even for a wrong simplification, as long as it happened to produce something smaller. Replace it with two meaningful checks against the intermediate result cancel_jacobian_products actually produces (before further geometry lowering, which erases Jacobian/JacobianInverse regardless of whether cancellation ran): - extract_type confirms Jacobian/JacobianInverse are fully eliminated, and that JacobianDeterminant (which cannot cancel further on its own here) correctly survives. - The result is compared, via renumber_indices (independently built expressions carry differently-numbered dummy indices), against a hand-coded expected result: the classical Piola divergence identity div(w) = (1/detJ) * ref_div(w), derived independently of cancel_jacobian_products itself.
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