diff --git a/test/test_cancel_jacobian_products.py b/test/test_cancel_jacobian_products.py new file mode 100644 index 000000000..a827dab75 --- /dev/null +++ b/test/test_cancel_jacobian_products.py @@ -0,0 +1,105 @@ +from utils import FiniteElement, LagrangeElement + +from ufl import ( + Coefficient, + FunctionSpace, + Mesh, + TestFunction, + TrialFunction, + div, + dx, + inner, + triangle, +) +from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering +from ufl.algorithms.cancel_jacobian_products import ( + IdentityEliminator, + JacobianCanceller, + ReciprocalCanceller, + cancel_jacobian_products, +) +from ufl.algorithms.compute_form_data import compute_form_data +from ufl.classes import Identity, Jacobian, JacobianDeterminant, JacobianInverse +from ufl.constantvalue import as_ufl +from ufl.operators import dot +from ufl.pullback import contravariant_piola +from ufl.sobolevspace import HDiv + + +def test_jacobian_canceller_and_identity_eliminator(): + domain = Mesh(LagrangeElement(triangle, 1, (2,))) + J = Jacobian(domain) + K = JacobianInverse(domain) + + # cancel_jacobian_products assumes component tensors have already + # been removed, so contractions appear as IndexSum over Indexed + # terminals rather than as a Dot operator. + expr = apply_algebra_lowering(dot(J, K)) + assert JacobianCanceller()(expr) == Identity(2) + # Idempotent: nothing left to cancel the second time round. + assert IdentityEliminator()(JacobianCanceller()(expr)) == Identity(2) + + # The other contraction order (K . J) is also a valid Kronecker delta, + # since K is a left inverse of J even for immersed manifolds. + expr2 = apply_algebra_lowering(dot(K, J)) + assert JacobianCanceller()(expr2) == Identity(2) + + +def test_reciprocal_canceller(): + domain = Mesh(LagrangeElement(triangle, 1, (2,))) + detJ = JacobianDeterminant(domain) + + expr = detJ**2 * (as_ufl(1) / detJ) ** 2 + assert ReciprocalCanceller()(expr) == as_ufl(1) + + # A reciprocal factor that does not fully cancel is simplified to a + # single net power, not left alone or over-simplified to nothing. + expr2 = detJ**3 * (as_ufl(1) / detJ) ** 2 + assert ReciprocalCanceller()(expr2) == detJ + + +def test_cancel_jacobian_products_on_form(): + domain = Mesh(LagrangeElement(triangle, 1, (2,))) + RT = FiniteElement("Raviart-Thomas", triangle, 1, (2,), contravariant_piola, HDiv) + V = FunctionSpace(domain, RT) + u = Coefficient(V) + v = TestFunction(V) + + form = inner(div(u), div(v)) * dx + fd_kwargs = dict( + do_apply_function_pullbacks=True, + do_apply_geometry_lowering=True, + do_apply_default_restrictions=True, + do_apply_restrictions=True, + ) + + fd_plain = compute_form_data(form, **fd_kwargs) + fd_cancelled = compute_form_data(form, do_cancel_jacobian_products=True, **fd_kwargs) + + integrand_plain = fd_plain.preprocessed_form.integrals()[0].integrand() + integrand_cancelled = fd_cancelled.preprocessed_form.integrals()[0].integrand() + + # The whole point of the pass: the div-div integrand of a + # contravariant Piola-mapped element is dominated by J-Jinv + # contractions that cancel, so the cancelled integrand is much + # smaller than the one form compilers would otherwise have to deal + # with, and does not introduce any new terminal type of its own. + assert len(str(integrand_cancelled)) < len(str(integrand_plain)) / 4 + + # do_cancel_jacobian_products defaults to False, so existing form + # compiler callers (e.g. ffcx) that do not pass it are unaffected. + fd_default = compute_form_data(form, **fd_kwargs) + assert fd_default.preprocessed_form.signature() == fd_plain.preprocessed_form.signature() + + +def test_cancel_jacobian_products_is_a_no_op_on_affine_elements(): + domain = Mesh(LagrangeElement(triangle, 1, (2,))) + V = FunctionSpace(domain, LagrangeElement(triangle, 1)) + u = TrialFunction(V) + v = TestFunction(V) + + # No Jacobian-inverse contractions or reciprocal Jacobian + # determinants appear here to begin with, so cancellation must + # reproduce the input form exactly. + form = inner(u, v) * dx + assert cancel_jacobian_products(form).signature() == form.signature() diff --git a/ufl/algorithms/cancel_jacobian_products.py b/ufl/algorithms/cancel_jacobian_products.py new file mode 100644 index 000000000..d0e81f5c8 --- /dev/null +++ b/ufl/algorithms/cancel_jacobian_products.py @@ -0,0 +1,393 @@ +"""Cancel contractions of the Jacobian with its inverse. + +This module simplifies index contractions of the Jacobian with its +inverse, before the Jacobian inverse is expanded into individual +matrix entries. This is done in two expression traversals. + +The first traversal replaces contractions of the Jacobian with its +inverse by Kronecker deltas, represented as an indexed Identity: + + IndexSum(Jacobian[a, k] * JacobianInverse[k, b] * factors, k) + -> Identity[a, b] * factors, + IndexSum(JacobianInverse[a, k] * Jacobian[k, b] * factors, k) + -> Identity[a, b] * factors. + +The second traversal eliminates the Identity tensors by contraction +against the remaining factors: + + IndexSum(Identity[a, k] * factors, k) -> factors[k -> a], + +and folds Identity entries at fixed indices into scalar constants. + +A third traversal cancels reciprocal factors within products, such as +those introduced by the Piola maps and their inverses: + + JacobianDeterminant**2 * (1 / JacobianDeterminant)**2 -> 1. + +These patterns arise from Piola-mapped elements, where the pullback +inserts Jacobian factors and spatial derivatives insert Jacobian +inverse factors that cancel out, e.g. in the divergence of a +contravariant Piola-mapped function. + +This pass assumes that derivatives have been expanded and that +component tensors have been removed, so that the contractions appear +as IndexSum nodes over products of indexed terminals. +""" +# Copyright (C) 2026 Pablo Brubeck +# +# This file is part of UFL (https://www.fenicsproject.org) +# +# SPDX-License-Identifier: LGPL-3.0-or-later + +from collections import defaultdict +from functools import singledispatchmethod + +from ufl.algorithms.map_integrands import map_integrands +from ufl.algorithms.remove_component_tensors import IndexReplacer +from ufl.classes import ( + Division, + Expr, + Identity, + Indexed, + IndexSum, + Jacobian, + JacobianInverse, + Power, + Product, +) +from ufl.constantvalue import ScalarValue, Zero, as_ufl +from ufl.core.multiindex import FixedIndex, MultiIndex +from ufl.corealg.dag_traverser import DAGTraverser +from ufl.corealg.map_dag import map_expr_dag +from ufl.domain import extract_unique_domain + + +def _flatten_product(expr, factors): + """Flatten nested Products into a list of scalar factors.""" + if isinstance(expr, Product): + for op in expr.ufl_operands: + _flatten_product(op, factors) + else: + factors.append(expr) + return factors + + +def _make_product(factors): + """Rebuild a Product from a list of scalar factors.""" + result = factors[0] + for f in factors[1:]: + result = Product(result, f) + return result + + +class IndexSumSimplifier(DAGTraverser): + """Base class for simplifying contractions in IndexSum nodes. + + Subclasses implement the ``match`` method, which rewrites the + product of the factors of a sum over a given index, or returns + None when no simplification applies. Contractions are chased + through nested IndexSums by interchanging the order of summation. + """ + + def __init__( + self, + compress: bool | None = True, + visited_cache: dict[tuple, Expr] | None = None, + result_cache: dict[Expr, Expr] | None = None, + ) -> None: + """Initialise.""" + super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache) + self._rules: dict[tuple, IndexReplacer] = {} + + @singledispatchmethod + def process(self, o: Expr) -> Expr: + """Process ``o``.""" + return super().process(o) + + @process.register(Expr) + def _(self, o: Expr) -> Expr: + """Reuse if untouched.""" + return self.reuse_if_untouched(o) + + def _substitute(self, expr, k, a): + """Replace the index k with a in expr.""" + rule = self._rules.get((k, a)) + if rule is None: + rule = IndexReplacer({k: a}) + self._rules[(k, a)] = rule + return map_expr_dag(rule, expr) + + def match(self, with_k, rest, k): + """Simplify IndexSum(product(with_k + rest), k), or return None. + + Args: + with_k: factors that have k as a free index. + rest: factors that do not depend on k. + k: the summation index. + """ + raise NotImplementedError(f"match() not implemented by {type(self).__name__}.") + + def _cancel(self, factors, k): + """Simplify IndexSum(product(factors), k), or return None. + + Only rewrites when a contraction over k cancels out; the + expression is otherwise left alone. + """ + with_k = [] + rest = [] + for f in factors: + if k.count() in f.ufl_free_indices: + with_k.append(f) + else: + rest.append(f) + + result = self.match(with_k, rest, k) + if result is not None: + return result + + if len(with_k) == 1 and isinstance(with_k[0], IndexSum): + # Interchange sums: sum_k sum_j f(j, k) = sum_j sum_k f(j, k) + summand, (j,) = with_k[0].ufl_operands + inner = self._cancel(_flatten_product(summand, []), k) + if inner is not None: + return _make_product(rest + [self._index_sum(inner, j)]) + + if len(with_k) == 2: + # Try to push an Indexed factor into an inner IndexSum + for f1, f2 in ((with_k[0], with_k[1]), (with_k[1], with_k[0])): + if isinstance(f1, Indexed) and isinstance(f2, IndexSum): + summand, (j,) = f2.ufl_operands + inner = self._cancel(_flatten_product(summand, [f1]), k) + if inner is not None: + return _make_product(rest + [self._index_sum(inner, j)]) + return None + + def _index_sum(self, summand, k): + """Construct IndexSum(summand, k), applying cancellations.""" + result = self._cancel(_flatten_product(summand, []), k) + if result is not None: + return result + return IndexSum(summand, MultiIndex((k,))) + + @process.register(IndexSum) + @DAGTraverser.postorder + def _(self, o: Expr, summand: Expr, multiindex: MultiIndex) -> Expr: + """Simplify IndexSum.""" + (k,) = multiindex + result = self._cancel(_flatten_product(summand, []), k) + if result is not None: + return result + if o.ufl_operands[0] is summand: + return o + return IndexSum(summand, multiindex) + + +def _delta_cancellation(f1, f2, k): + """Match a J-Jinv contraction over the index k. + + Given two scalar factors of a sum over k, return the equivalent + Kronecker delta if they contract to one, otherwise return None. + """ + for fa, fb in ((f1, f2), (f2, f1)): + if not (isinstance(fa, Indexed) and isinstance(fb, Indexed)): + continue + A, ia = fa.ufl_operands + B, ib = fb.ufl_operands + if not (isinstance(A, JacobianInverse) and isinstance(B, Jacobian)): + continue + domain = extract_unique_domain(A) + if domain != extract_unique_domain(B): + continue + ia, ib = tuple(ia), tuple(ib) + if ia[1] == k and ib[0] == k and ia[0] != k and ib[1] != k: + # sum_k K[a, k] * J[k, b] = Identity[a, b] + # This holds for pseudo-inverses on immersed manifolds, as + # K is a left inverse of J. + tdim = domain.topological_dimension + return Indexed(Identity(tdim), MultiIndex((ia[0], ib[1]))) + if ib[1] == k and ia[0] == k and ib[0] != k and ia[1] != k: + # sum_k J[a, k] * K[k, b] = Identity[a, b] + # Only valid when the Jacobian is square and invertible. + gdim = domain.geometric_dimension + tdim = domain.topological_dimension + if gdim == tdim: + return Indexed(Identity(gdim), MultiIndex((ib[0], ia[1]))) + return None + + +class JacobianCanceller(IndexSumSimplifier): + """Cancel Jacobian-inverse contractions into Kronecker deltas.""" + + def match(self, with_k, rest, k): + """Replace a J-Jinv contraction over k by a Kronecker delta.""" + if len(with_k) == 2: + delta = _delta_cancellation(with_k[0], with_k[1], k) + if delta is not None: + # The contraction over k is consumed by the delta + return _make_product(rest + [delta]) + return None + + +def _identity_index(f, k): + """If f is Identity[a, k] or Identity[k, a] with a != k, return a.""" + if isinstance(f, Indexed): + A, ii = f.ufl_operands + if isinstance(A, Identity): + a, b = tuple(ii) + if a == k and b != k: + return b + if b == k and a != k: + return a + return None + + +class IdentityEliminator(IndexSumSimplifier): + """Eliminate Kronecker deltas represented by indexed Identity tensors.""" + + def match(self, with_k, rest, k): + """Contract an Identity factor over k against the remaining factors.""" + for i, f in enumerate(with_k): + a = _identity_index(f, k) + if a is not None: + others = with_k[:i] + with_k[i + 1 :] + rest + if not others: + return None + return self._substitute(_make_product(others), k, a) + return None + + # Work around singledispatchmethod inheritance issue; + # see https://bugs.python.org/issue36457. + @singledispatchmethod + def process(self, o: Expr) -> Expr: + """Process ``o``.""" + return super().process(o) + + @process.register(Indexed) + @DAGTraverser.postorder + def _(self, o: Expr, A: Expr, ii: MultiIndex) -> Expr: + """Fold Identity entries at fixed indices into scalar constants.""" + if isinstance(A, Identity): + a, b = tuple(ii) + if isinstance(a, FixedIndex) and isinstance(b, FixedIndex): + return as_ufl(1) if int(a) == int(b) else Zero() + if o.ufl_operands == (A, ii): + return o + return Indexed(A, ii) + + +def _as_base_exponent(f): + """Destructure a factor into a (base, exponent) pair. + + Returns None when the factor does not have a constant real exponent. + """ + if isinstance(f, Power): + base, exponent = f.ufl_operands + if isinstance(exponent, ScalarValue) and not isinstance(exponent._value, complex): + pair = _as_base_exponent(base) + if pair is not None: + base, inner = pair + return base, inner * exponent._value + return None + elif isinstance(f, Division): + numerator, denominator = f.ufl_operands + if isinstance(numerator, ScalarValue) and numerator._value == 1: + pair = _as_base_exponent(denominator) + if pair is not None: + base, inner = pair + return base, -inner + return None + else: + return f, 1 + + +def _make_power(base, exponent): + """Construct base ** exponent for a constant real exponent.""" + if exponent == int(exponent): + exponent = int(exponent) + if exponent == 1: + return base + elif exponent == -1: + return Division(as_ufl(1), base) + elif exponent < 0: + return Division(as_ufl(1), Power(base, as_ufl(-exponent))) + else: + return Power(base, as_ufl(exponent)) + + +class ReciprocalCanceller(DAGTraverser): + """Cancel reciprocal factors within products.""" + + @singledispatchmethod + def process(self, o: Expr) -> Expr: + """Process ``o``.""" + return super().process(o) + + @process.register(Expr) + def _(self, o: Expr) -> Expr: + """Reuse if untouched.""" + return self.reuse_if_untouched(o) + + @process.register(Product) + @DAGTraverser.postorder + def _(self, o: Expr, a: Expr, b: Expr) -> Expr: + """Cancel bases that appear with exponents of opposite signs.""" + factors = _flatten_product(b, _flatten_product(a, [])) + + # Collect the exponents of each base + exponents = defaultdict(list) + for f in factors: + pair = _as_base_exponent(f) + if pair is not None: + base, exponent = pair + exponents[base].append(exponent) + + # Only rebuild bases that appear with exponents of opposite signs + mixed = { + base + for base, es in exponents.items() + if any(e > 0 for e in es) and any(e < 0 for e in es) + } + if not mixed: + if o.ufl_operands == (a, b): + return o + return Product(a, b) + + parts = [] + emitted = set() + for f in factors: + pair = _as_base_exponent(f) + base = pair[0] if pair is not None else None + if base in mixed: + if base not in emitted: + emitted.add(base) + net = sum(exponents[base]) + if net != 0: + parts.append(_make_power(base, net)) + else: + parts.append(f) + + if not parts: + result = as_ufl(1) + else: + result = parts[0] + for f in parts[1:]: + result = Product(result, f) + + # Do not rebuild if cancellation would drop free indices + if result.ufl_free_indices != o.ufl_free_indices: + if o.ufl_operands == (a, b): + return o + return Product(a, b) + return result + + +def cancel_jacobian_products(o): + """Cancel contractions of the Jacobian with its inverse. + + Args: + o: An Expr or Form. + """ + o = map_integrands(JacobianCanceller(), o) + o = map_integrands(IdentityEliminator(), o) + o = map_integrands(ReciprocalCanceller(), o) + return o diff --git a/ufl/algorithms/compute_form_data.py b/ufl/algorithms/compute_form_data.py index d5164de63..85c8a6014 100644 --- a/ufl/algorithms/compute_form_data.py +++ b/ufl/algorithms/compute_form_data.py @@ -16,6 +16,7 @@ from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering from ufl.algorithms.apply_integral_scaling import apply_integral_scaling +from ufl.algorithms.cancel_jacobian_products import cancel_jacobian_products from ufl.algorithms.comparison_checker import do_comparison_check # See TODOs at the call sites of these below: @@ -27,7 +28,7 @@ from ufl.algorithms.formdata import FormData from ufl.algorithms.remove_complex_nodes import remove_complex_nodes from ufl.algorithms.remove_component_tensors import remove_component_tensors -from ufl.classes import Form +from ufl.classes import Form, Jacobian, JacobianDeterminant, JacobianInverse def attach_estimated_degrees(form): @@ -88,6 +89,7 @@ def compute_form_data( do_apply_integral_scaling=False, do_apply_geometry_lowering=False, preserve_geometry_types=(), + do_cancel_jacobian_products=False, do_apply_default_restrictions=True, do_apply_restrictions=True, do_estimate_degrees=True, @@ -110,6 +112,12 @@ def compute_form_data( quantities to a smaller subset of quantities preserve_geometry_types: Set of quantities not to lower, and keep at its present stage for the form-compiler. + do_cancel_jacobian_products: Delay the expansion of the Jacobian + inverse into individual matrix entries, and cancel out index + contractions of the Jacobian with its inverse. This + simplifies the expressions that Piola-mapped elements + generate, before lowering the surviving Jacobian quantities + as usual. do_apply_default_restrictions: Apply default restrictions, defined in {py:mod}`ufl.algorithms.apply_restrictions` to integrals if no restriction has been set. @@ -164,12 +172,25 @@ def compute_form_data( if do_apply_integral_scaling: form = apply_integral_scaling(form) + # Keep the Jacobian, its inverse, and its determinant as opaque + # terminals while derivatives are expanded, so that contractions + # of the Jacobian with its inverse can be cancelled out before the + # inverse is expanded into individual matrix entries. + if do_cancel_jacobian_products: + lowering_preserve_types = set(preserve_geometry_types) | { + Jacobian, + JacobianInverse, + JacobianDeterminant, + } + else: + lowering_preserve_types = preserve_geometry_types + # Lower abstractions for geometric quantities into a smaller set # of quantities, allowing the form compiler to deal with a smaller # set of types and treating geometric quantities like any other # expressions w.r.t. loop-invariant code motion etc. if do_apply_geometry_lowering: - form = apply_geometry_lowering(form, preserve_geometry_types) + form = apply_geometry_lowering(form, lowering_preserve_types) # Apply differentiation again, because the algorithms above can # generate new derivatives or rewrite expressions inside @@ -180,10 +201,19 @@ def compute_form_data( # Neverending story: apply_derivatives introduces new Jinvs, # which needs more geometry lowering if do_apply_geometry_lowering: - form = apply_geometry_lowering(form, preserve_geometry_types) + form = apply_geometry_lowering(form, lowering_preserve_types) # Lower derivatives that may have appeared form = apply_derivatives(form) + if do_cancel_jacobian_products: + # Cancel contractions of the Jacobian with its inverse, + # which requires component tensors to be removed first + form = remove_component_tensors(form) + form = cancel_jacobian_products(form) + # Lower the Jacobian quantities that were preserved above + form = apply_geometry_lowering(form, preserve_geometry_types) + form = apply_derivatives(form) + form = apply_coordinate_derivatives(form) # If in real mode, remove any complex nodes introduced during form processing.