Remove component tensors

test/test_algorithms.py

@@ -12,6 +12,7 @@
     FacetNormal,
     FunctionSpace,
     Mesh,
+    SpatialCoordinate,
     TestFunction,
     TrialFunction,
     adjoint,
@@ -21,9 +22,11 @@
     dx,
     grad,
     inner,
+    sin,
     triangle,
 )
 from ufl.algorithms import (
+    compute_form_data,
     expand_derivatives,
     expand_indices,
     extract_arguments,
@@ -182,3 +185,32 @@ def test_adjoint(domain):
     d = adjoint(b)
     d_arg_degrees = [arg.ufl_element().embedded_superdegree for arg in extract_arguments(d)]
     assert d_arg_degrees == [2, 1]
+
+
+def test_remove_component_tensors(domain):
+    x = SpatialCoordinate(domain)
+    u = sin(x[0])
+
+    f = div(grad(div(grad(u))))
+    form = f * dx
+
+    fd = compute_form_data(form)
+
+    assert "ComponentTensor" not in repr(fd.preprocessed_form)
+
+
+def test_grad_cellwise_constant(domain):
+    element = FiniteElement("Lagrange", triangle, 3, (), identity_pullback, H1)
+    space = FunctionSpace(domain, element)
+    u = Coefficient(space)
+
+    # Applying four derivatives to a cubic should simplify to zero
+    f = div(grad(div(grad(u))))
+    form = f * dx
+
+    fd = compute_form_data(
+        form,
+        do_apply_function_pullbacks=True,
+    )
+    assert fd.preprocessed_form.empty()
+    assert fd.num_coefficients == 0

test/test_apply_function_pullbacks.py

@@ -151,8 +151,8 @@ def test_apply_single_function_pullbacks_triangle3d():
         vc: as_vector(Jinv[j, i] * rvc[j], i),
         t: rt,
         s: as_tensor([[rs[0], rs[1], rs[2]], [rs[1], rs[3], rs[4]], [rs[2], rs[4], rs[5]]]),
-        cov2t: as_tensor(Jinv[k, i] * rcov2t[k, l] * Jinv[l, j], (i, j)),
-        contra2t: as_tensor((1.0 / detJ) ** 2 * J[i, k] * rcontra2t[k, l] * J[j, l], (i, j)),
+        cov2t: as_tensor(Jinv[k, i] * (rcov2t[k, l] * Jinv[l, j]), (i, j)),
+        contra2t: (1.0 / detJ) ** 2 * as_tensor(J[i, k] * (rcontra2t[k, l] * J[j, l]), (i, j)),
         # Mixed elements become a bit more complicated
         uml2: as_vector([ruml2[0] / detJ, ruml2[1] / detJ]),
         um: rum,

test/test_derivative.py

@@ -665,7 +665,7 @@ def test_vector_coefficient_scalar_derivatives(self):
     integrand = inner(f, g)
 
     i0, i1, i2, i3, i4 = [Index(count=c) for c in range(5)]
-    expected = as_tensor(df[i1] * dv, (i1,))[i0] * g[i0]
+    expected = as_tensor(df[i1], (i1,))[i0] * dv * g[i0]
 
     F = integrand * dx
     J = derivative(F, u, dv, cd)
@@ -693,7 +693,7 @@ def test_vector_coefficient_derivatives(self):
     integrand = inner(f, g)
 
     i0, i1, i2, i3, i4 = [Index(count=c) for c in range(5)]
-    expected = as_tensor(df[i2, i1] * dv[i1], (i2,))[i0] * g[i0]
+    expected = as_tensor(df[i2, i1], (i2,))[i0] * dv[i1] * g[i0]
 
     F = integrand * dx
     J = derivative(F, u, dv, cd)

ufl/algebra.py

@@ -13,6 +13,7 @@
 from ufl.core.operator import Operator
 from ufl.core.ufl_type import ufl_type
 from ufl.index_combination_utils import merge_unique_indices
+from ufl.indexed import Indexed
 from ufl.precedence import parstr
 from ufl.sorting import sorted_expr
 
@@ -89,6 +90,11 @@ def __init__(self, a, b):
         """Initialise."""
         Operator.__init__(self)
 
+    def _simplify_indexed(self, multiindex):
+        """Return a simplified Expr used in the constructor of Indexed(self, multiindex)."""
+        a, b = self.ufl_operands
+        return Sum(Indexed(a, multiindex), Indexed(b, multiindex))
+
     def evaluate(self, x, mapping, component, index_values):
         """Evaluate."""
         return sum(o.evaluate(x, mapping, component, index_values) for o in self.ufl_operands)

ufl/algorithms/apply_derivatives.py

@@ -12,6 +12,7 @@
 
 from ufl.action import Action
 from ufl.algorithms.analysis import extract_arguments
+from ufl.algorithms.estimate_degrees import SumDegreeEstimator
 from ufl.algorithms.map_integrands import map_integrand_dags
 from ufl.algorithms.replace_derivative_nodes import replace_derivative_nodes
 from ufl.argument import BaseArgument
@@ -562,6 +563,14 @@ def __init__(self, geometric_dimension):
         """Initialise."""
         GenericDerivativeRuleset.__init__(self, var_shape=(geometric_dimension,))
         self._Id = Identity(geometric_dimension)
+        self.degree_estimator = SumDegreeEstimator(1, {})
+
+    def is_cellwise_constant(self, o):
+        """More precise checks for cellwise constants."""
+        if is_cellwise_constant(o):
+            return True
+        degree = map_expr_dag(self.degree_estimator, o)
+        return degree == 0
 
     # --- Specialized rules for geometric quantities
 
@@ -572,7 +581,7 @@ def geometric_quantity(self, o):
         otherwise transform derivatives to reference derivatives.
         Override for specific types if other behaviour is needed.
         """
-        if is_cellwise_constant(o):
+        if self.is_cellwise_constant(o):
             return self.independent_terminal(o)
         else:
             domain = extract_unique_domain(o)
@@ -583,7 +592,7 @@ def geometric_quantity(self, o):
     def jacobian_inverse(self, o):
         """Differentiate a jacobian_inverse."""
         # grad(K) == K_ji rgrad(K)_rj
-        if is_cellwise_constant(o):
+        if self.is_cellwise_constant(o):
             return self.independent_terminal(o)
         if not o._ufl_is_terminal_:
             raise ValueError("ReferenceValue can only wrap a terminal")
@@ -653,9 +662,10 @@ def reference_value(self, o):
 
     def reference_grad(self, o):
         """Differentiate a reference_grad."""
+        if self.is_cellwise_constant(o):
+            return self.independent_terminal(o)
         # grad(o) == grad(rgrad(rv(f))) -> K_ji*rgrad(rgrad(rv(f)))_rj
         f = o.ufl_operands[0]
-
         valid_operand = f._ufl_is_in_reference_frame_ or isinstance(
             f, (JacobianInverse, SpatialCoordinate, Jacobian, JacobianDeterminant)
         )
@@ -676,7 +686,6 @@ def grad(self, o):
         # Check that o is a "differential terminal"
         if not isinstance(o.ufl_operands[0], (Grad, Terminal)):
             raise ValueError("Expecting only grads applied to a terminal.")
-
         return Grad(o)
 
     def _grad(self, o):

ufl/algorithms/compute_form_data.py

@@ -33,6 +33,7 @@
 from ufl.algorithms.formdata import FormData
 from ufl.algorithms.formtransformations import compute_form_arities
 from ufl.algorithms.remove_complex_nodes import remove_complex_nodes
+from ufl.algorithms.remove_component_tensors import remove_component_tensors
 from ufl.classes import Coefficient, Form, FunctionSpace, GeometricFacetQuantity
 from ufl.corealg.traversal import traverse_unique_terminals
 from ufl.domain import extract_unique_domain
@@ -328,6 +329,9 @@ def compute_form_data(
 
     form = apply_coordinate_derivatives(form)
 
+    # Remove component tensors
+    form = remove_component_tensors(form)
+
     # Propagate restrictions to terminals
     if do_apply_restrictions:
         form = apply_restrictions(form, apply_default=do_apply_default_restrictions)

ufl/algorithms/remove_component_tensors.py

@@ -0,0 +1,117 @@
+"""Remove component tensors.
+
+This module contains classes and functions to remove component tensors.
+"""
+# Copyright (C) 2008-2016 Martin Sandve Alnæs
+#
+# This file is part of UFL (https://www.fenicsproject.org)
+#
+# SPDX-License-Identifier:    LGPL-3.0-or-later
+
+from ufl.classes import ComponentTensor, Form, Index, MultiIndex, Zero
+from ufl.corealg.map_dag import map_expr_dag
+from ufl.corealg.multifunction import MultiFunction, memoized_handler
+
+
+class IndexReplacer(MultiFunction):
+    """Replace Indices."""
+
+    def __init__(self, fimap: dict):
+        """Initialise.
+
+        Args:
+           fimap: map for index replacements.
+
+        """
+        MultiFunction.__init__(self)
+        self.fimap = fimap
+        self._object_cache = {}
+
+    expr = MultiFunction.reuse_if_untouched
+
+    @memoized_handler
+    def zero(self, o):
+        """Handle Zero."""
+        free_indices = []
+        index_dimensions = []
+        for i, d in zip(o.ufl_free_indices, o.ufl_index_dimensions):
+            k = Index(i)
+            j = self.fimap.get(k, k)
+            if isinstance(j, Index):
+                free_indices.append(j.count())
+                index_dimensions.append(d)
+        return Zero(
+            shape=o.ufl_shape,
+            free_indices=tuple(free_indices),
+            index_dimensions=tuple(index_dimensions),
+        )
+
+    @memoized_handler
+    def multi_index(self, o):
+        """Handle MultiIndex."""
+        return MultiIndex(tuple(self.fimap.get(i, i) for i in o.indices()))
+
+
+class IndexRemover(MultiFunction):
+    """Remove Indexed."""
+
+    def __init__(self):
+        """Initialise."""
+        MultiFunction.__init__(self)
+        self._object_cache = {}
+
+    expr = MultiFunction.reuse_if_untouched
+
+    @memoized_handler
+    def _unary_operator(self, o):
+        """Simplify UnaryOperator(Zero)."""
+        (operand,) = o.ufl_operands
+        f = map_expr_dag(self, operand)
+        if isinstance(f, Zero):
+            return Zero(
+                shape=o.ufl_shape,
+                free_indices=o.ufl_free_indices,
+                index_dimensions=o.ufl_index_dimensions,
+            )
+        if f is operand:
+            # Reuse if untouched
+            return o
+        return o._ufl_expr_reconstruct_(f)
+
+    @memoized_handler
+    def indexed(self, o):
+        """Simplify Indexed."""
+        o1, i1 = o.ufl_operands
+        if isinstance(o1, ComponentTensor):
+            # Simplify Indexed ComponentTensor
+            o2, i2 = o1.ufl_operands
+            # Replace inner indices first
+            v = map_expr_dag(self, o2)
+            # Replace outer indices
+            assert len(i2) == len(i1)
+            fimap = dict(zip(i2, i1))
+            rule = IndexReplacer(fimap)
+            return map_expr_dag(rule, v)
+
+        expr = map_expr_dag(self, o1)
+        if expr is o1:
+            # Reuse if untouched
+            return o
+        return o._ufl_expr_reconstruct_(expr, i1)
+
+    reference_grad = _unary_operator
+    reference_value = _unary_operator
+
+
+def remove_component_tensors(o):
+    """Remove component tensors."""
+    if isinstance(o, Form):
+        integrals = []
+        for integral in o.integrals():
+            integrand = remove_component_tensors(integral.integrand())
+            if not isinstance(integrand, Zero):
+                integrals.append(integral.reconstruct(integrand=integrand))
+        return o._ufl_expr_reconstruct_(integrals)
+    else:
+        rule = IndexRemover()
+        return map_expr_dag(rule, o)

ufl/exproperators.py

@@ -406,20 +406,17 @@ def _getitem(self, component):
     mi = MultiIndex(all_indices)
     a = Indexed(self, mi)
 
-    # TODO: I think applying as_tensor after index sums results in
-    # cleaner expression graphs.
-
-    # If the Ellipsis or any slices were found, wrap as tensor valued
-    # with the slice indices created at the top here
-    if slice_indices:
-        a = as_tensor(a, slice_indices)
-
     # If any repeated indices were found, apply implicit summation
     # over those
     for i in repeated_indices:
         mi = MultiIndex((i,))
         a = IndexSum(a, mi)
 
+    # If the Ellipsis or any slices were found, wrap as tensor valued
+    # with the slice indices created at the top here
+    if slice_indices:
+        a = as_tensor(a, slice_indices)
+
     # Check for zero (last so we can get indices etc from a, could
     # possibly be done faster by checking early instead)
     if isinstance(self, Zero):