WIP

Author: pbrubeck

ufl/algorithms/remove_component_tensors.py

@@ -9,6 +9,7 @@
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
 from ufl.algorithms.estimate_degrees import SumDegreeEstimator
+from ufl.checks import is_cellwise_constant
 from ufl.classes import (
     ComponentTensor,
     Form,
@@ -42,13 +43,10 @@ def zero(self, o):
         free_indices = []
         index_dimensions = []
         for i, d in zip(o.ufl_free_indices, o.ufl_index_dimensions):
-            if Index(i) in self.fimap:
-                ind_j = self.fimap[Index(i)]
-                if isinstance(ind_j, Index):
-                    free_indices.append(ind_j.count())
-                    index_dimensions.append(d)
-            else:
-                free_indices.append(i)
+            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,
@@ -71,24 +69,45 @@ def __init__(self):
         self._object_cache = {}
         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
+
     expr = MultiFunction.reuse_if_untouched
 
+    @memoized_handler
+    def sum(self, o):
+        """Simplify Sum, allow adding Zeros with different free indices."""
+        oa, ob = o.ufl_operands
+        a = map_expr_dag(self, oa)
+        b = map_expr_dag(self, ob)
+        if isinstance(b, Zero):
+            return a
+        if isinstance(a, Zero):
+            return b
+        if a is oa and b is ob:
+            # Reuse if untouched
+            return o
+        return o._ufl_expr_reconstruct_(a, b)
+
     @memoized_handler
     def reference_grad(self, o):
         """Simplify ReferenceGrad(Constant)."""
         (operand,) = o.ufl_operands
-        operand = map_expr_dag(self, operand)
-        degree = map_expr_dag(self.degree_estimator, operand)
-        if degree == 0:
+        f = map_expr_dag(self, operand)
+        if self.is_cellwise_constant(f):
             return Zero(
                 shape=o.ufl_shape,
                 free_indices=o.ufl_free_indices,
                 index_dimensions=o.ufl_index_dimensions,
             )
-        if operand is o.ufl_operands[0]:
+        if f is operand:
             # Reuse if untouched
             return o
-        return o._ufl_expr_reconstruct_(operand)
+        return o._ufl_expr_reconstruct_(f)
 
     @memoized_handler
     def indexed(self, o):