From 0651e63e354ba55add2b3d29f9ed9368a51a4559 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 10 Jul 2026 00:40:16 +0100 Subject: [PATCH 1/2] Cancel contractions of the Jacobian with its inverse 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 --- ufl/algorithms/cancel_jacobian_products.py | 355 +++++++++++++++++++++ ufl/algorithms/compute_form_data.py | 36 ++- 2 files changed, 388 insertions(+), 3 deletions(-) create mode 100644 ufl/algorithms/cancel_jacobian_products.py diff --git a/ufl/algorithms/cancel_jacobian_products.py b/ufl/algorithms/cancel_jacobian_products.py new file mode 100644 index 000000000..b3bb9c429 --- /dev/null +++ b/ufl/algorithms/cancel_jacobian_products.py @@ -0,0 +1,355 @@ +"""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 ufl.algorithms.map_integrands import map_integrand_dags +from ufl.algorithms.remove_component_tensors import IndexReplacer +from ufl.classes import ( + Division, + 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.map_dag import map_expr_dag +from ufl.corealg.multifunction import MultiFunction +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(MultiFunction): + """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): + """Initialise.""" + MultiFunction.__init__(self) + self._rules = {} + + expr = MultiFunction.reuse_if_untouched + + 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,))) + + def index_sum(self, o, summand, multiindex): + """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 + + def indexed(self, o, A, ii): + """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() + return self.expr(o, 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(MultiFunction): + """Cancel reciprocal factors within products.""" + + expr = MultiFunction.reuse_if_untouched + + def product(self, o, a, b): + """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_integrand_dags(JacobianCanceller(), o) + o = map_integrand_dags(IdentityEliminator(), o) + o = map_integrand_dags(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. From ced0cd4bde68d938bd5052530c6d234bfb98f2c8 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 11 Jul 2026 12:35:48 +0100 Subject: [PATCH 2/2] Address review: convert to DAGTraverser, add tests 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. --- test/test_cancel_jacobian_products.py | 105 +++++++++++++++++++++ ufl/algorithms/cancel_jacobian_products.py | 70 ++++++++++---- 2 files changed, 159 insertions(+), 16 deletions(-) create mode 100644 test/test_cancel_jacobian_products.py 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 index b3bb9c429..d0e81f5c8 100644 --- a/ufl/algorithms/cancel_jacobian_products.py +++ b/ufl/algorithms/cancel_jacobian_products.py @@ -40,11 +40,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from collections import defaultdict +from functools import singledispatchmethod -from ufl.algorithms.map_integrands import map_integrand_dags +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, @@ -55,8 +57,8 @@ ) 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.corealg.multifunction import MultiFunction from ufl.domain import extract_unique_domain @@ -78,7 +80,7 @@ def _make_product(factors): return result -class IndexSumSimplifier(MultiFunction): +class IndexSumSimplifier(DAGTraverser): """Base class for simplifying contractions in IndexSum nodes. Subclasses implement the ``match`` method, which rewrites the @@ -87,12 +89,25 @@ class IndexSumSimplifier(MultiFunction): through nested IndexSums by interchanging the order of summation. """ - def __init__(self): + def __init__( + self, + compress: bool | None = True, + visited_cache: dict[tuple, Expr] | None = None, + result_cache: dict[Expr, Expr] | None = None, + ) -> None: """Initialise.""" - MultiFunction.__init__(self) - self._rules = {} + super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache) + self._rules: dict[tuple, IndexReplacer] = {} - expr = MultiFunction.reuse_if_untouched + @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.""" @@ -154,7 +169,9 @@ def _index_sum(self, summand, k): return result return IndexSum(summand, MultiIndex((k,))) - def index_sum(self, o, summand, multiindex): + @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) @@ -238,13 +255,24 @@ def match(self, with_k, rest, k): return self._substitute(_make_product(others), k, a) return None - def indexed(self, o, A, ii): + # 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() - return self.expr(o, A, ii) + if o.ufl_operands == (A, ii): + return o + return Indexed(A, ii) def _as_base_exponent(f): @@ -286,12 +314,22 @@ def _make_power(base, exponent): return Power(base, as_ufl(exponent)) -class ReciprocalCanceller(MultiFunction): +class ReciprocalCanceller(DAGTraverser): """Cancel reciprocal factors within products.""" - expr = MultiFunction.reuse_if_untouched + @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 product(self, o, a, b): + @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, [])) @@ -349,7 +387,7 @@ def cancel_jacobian_products(o): Args: o: An Expr or Form. """ - o = map_integrand_dags(JacobianCanceller(), o) - o = map_integrand_dags(IdentityEliminator(), o) - o = map_integrand_dags(ReciprocalCanceller(), o) + o = map_integrands(JacobianCanceller(), o) + o = map_integrands(IdentityEliminator(), o) + o = map_integrands(ReciprocalCanceller(), o) return o