Remove component tensors

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/compute_form_data.py

@@ -13,6 +13,7 @@
 
 from ufl.algorithms.analysis import extract_coefficients, extract_sub_elements, unique_tuple
 from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering
+from ufl.algorithms.remove_component_tensors import remove_component_tensors
 from ufl.algorithms.apply_derivatives import apply_coordinate_derivatives, apply_derivatives
 
 # These are the main symbolic processing steps:
@@ -328,6 +329,8 @@ def compute_form_data(
 
     form = apply_coordinate_derivatives(form)
 
+    form = remove_component_tensors(form)
+
     # Propagate restrictions to terminals
     if do_apply_restrictions:
         form = apply_restrictions(form, apply_default=do_apply_default_restrictions)

ufl/indexsum.py

@@ -11,6 +11,7 @@
 from ufl.core.multiindex import MultiIndex
 from ufl.core.operator import Operator
 from ufl.core.ufl_type import ufl_type
+from ufl.indexed import Indexed
 from ufl.precedence import parstr
 
 # --- Sum over an index ---
@@ -69,6 +70,11 @@ def ufl_shape(self):
         """Get UFL shape."""
         return self.ufl_operands[0].ufl_shape
 
+    def _simplify_indexed(self, multiindex):
+        """Return a simplified Expr used in the constructor of Indexed(self, multiindex)."""
+        A, i = self.ufl_operands
+        return IndexSum(Indexed(A, multiindex), i)
+
     def evaluate(self, x, mapping, component, index_values):
         """Evaluate."""
         (i,) = self.ufl_operands[1]