Merge branch 'main' into pbrubeck/merge-upstream

Author: pbrubeck

test/test_derivative.py

@@ -60,7 +60,7 @@
 from ufl.algorithms.apply_derivatives import apply_derivatives
 from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering
 from ufl.classes import Indexed, MultiIndex, ReferenceGrad
-from ufl.constantvalue import as_ufl
+from ufl.constantvalue import Zero, as_ufl
 from ufl.domain import extract_unique_domain
 from ufl.finiteelement import FiniteElement, MixedElement
 from ufl.pullback import identity_pullback
@@ -902,6 +902,39 @@ def test_index_simplification_reference_grad(self):
     assert expr.ufl_shape == ()
 
 
+def test_zero_shape(self):
+    cell = triangle
+    shape = (2, 3, 4)
+    P1 = FiniteElement("Lagrange", cell, 1, shape, identity_pullback, H1)
+    domain = Mesh(FiniteElement("Lagrange", cell, 1, (2,), identity_pullback, H1))
+    V = FunctionSpace(domain, P1)
+    v = TestFunction(V)
+    u = Coefficient(V)
+    w = Coefficient(V)
+
+    (i,) = indices(1)
+    z = zero(shape)
+    zi = z[:, i, :]
+    wi = w[:, i, :]
+    assert isinstance(zi, Zero)
+    assert wi.ufl_shape == (shape[0], shape[-1])
+    assert wi.ufl_shape == zi.ufl_shape
+
+    a = derivative(conditional(u[0, 0, 0] < 1, zi, wi), u, v)
+    assert not isinstance(a, Zero)
+
+    assert a.ufl_shape == zi.ufl_shape
+    assert a.ufl_free_indices == zi.ufl_free_indices
+    assert a.ufl_index_dimensions == zi.ufl_index_dimensions
+
+    expr = apply_derivatives(apply_geometry_lowering(apply_algebra_lowering(a)))
+
+    assert isinstance(expr, Zero)
+    assert expr.ufl_shape == a.ufl_shape
+    assert expr.ufl_free_indices == a.ufl_free_indices
+    assert expr.ufl_index_dimensions == a.ufl_index_dimensions
+
+
 # --- Scratch space
 
 

ufl/algorithms/apply_derivatives.py

@@ -179,6 +179,9 @@ def independent_operator(self, o):
     # Constants are independent of any differentiation
     constant = independent_terminal
 
+    # Zero may have free indices
+    zero = independent_operator
+
     # Rules for form arguments must be specified in specialized rule set
     form_argument = override
 

ufl/algorithms/remove_component_tensors.py

@@ -14,6 +14,7 @@
 from ufl.classes import ComponentTensor, Index, MultiIndex, Zero
 from ufl.corealg.map_dag import map_expr_dag
 from ufl.corealg.multifunction import MultiFunction
+from ufl.index_combination_utils import unique_sorted_indices
 
 
 class IndexReplacer(MultiFunction):
@@ -38,18 +39,19 @@ def zero(self, o):
             # Reuse if untouched
             return o
 
-        free_indices = []
-        index_dimensions = []
+        fi = []
         for i, d in zip(indices, o.ufl_index_dimensions):
             j = self.fimap.get(i, i)
             if isinstance(j, Index):
-                free_indices.append(j.count())
-                index_dimensions.append(d)
+                fi.append((j.count(), d))
+
+        fi = unique_sorted_indices(sorted(fi))
+        free_indices, index_dimensions = zip(*fi)
 
         return Zero(
             shape=o.ufl_shape,
-            free_indices=tuple(free_indices),
-            index_dimensions=tuple(index_dimensions),
+            free_indices=free_indices,
+            index_dimensions=index_dimensions,
         )
 
     def multi_index(self, o):

ufl/conditional.py

@@ -13,7 +13,6 @@
 from ufl.core.operator import Operator
 from ufl.core.ufl_type import ufl_type
 from ufl.exprequals import expr_equals
-from ufl.indexed import Indexed
 from ufl.precedence import parstr
 
 # --- Condition classes ---
@@ -292,7 +291,7 @@ def __init__(self, condition, true_value, false_value):
             raise ValueError("Shape mismatch between conditional branches.")
         tfi = true_value.ufl_free_indices
         ffi = false_value.ufl_free_indices
-        if tuple(sorted(tfi)) != tuple(sorted(ffi)):
+        if tfi != ffi:
             raise ValueError("Free index mismatch between conditional branches.")
         if isinstance(condition, (EQ, NE)):
             if not all(
@@ -307,10 +306,6 @@ def __init__(self, condition, true_value, false_value):
         Operator.__init__(self, (condition, true_value, false_value))
         self._initialised = True
 
-    def _simplify_indexed(self, multiindex):
-        (c, a, b) = self.ufl_operands
-        return Conditional(c, Indexed(a, multiindex), Indexed(b, multiindex))
-
     def evaluate(self, x, mapping, component, index_values):
         """Evaluate."""
         c = self.ufl_operands[0].evaluate(x, mapping, component, index_values)

ufl/exprequals.py

@@ -20,7 +20,7 @@ def expr_equals(self, other):
         return False
 
     # Large objects are costly to compare with themselves
-    if self is other or self.ufl_operands is other.ufl_operands:
+    if (self is other) or (self.ufl_operands is other.ufl_operands):
         return True
 
     # Modelled after pre_traversal to avoid recursion:
@@ -37,7 +37,7 @@ def expr_equals(self, other):
             # Delve into subtrees
             so = s.ufl_operands
             oo = o.ufl_operands
-            # Skip subtrees if operands are the same
+            # Skip subtree if operands are the same
             if so is oo:
                 continue
             if len(so) != len(oo):
@@ -48,7 +48,7 @@ def expr_equals(self, other):
                 if s._ufl_typecode_ != o._ufl_typecode_:
                     return False
                 # Skip subtree if objects are the same
-                if s is o or s.ufl_operands is o.ufl_operands:
+                if s is o:
                     continue
                 if (id(s), id(o)) in equal_pairs:
                     continue

ufl/sorting.py

@@ -143,7 +143,7 @@ def cmp_expr(a, b):
             # Delve into subtrees
             aops = a.ufl_operands
             bops = b.ufl_operands
-            # Skip subtrees if operands are the same
+            # Skip subtree if operands are the same
             if aops is bops:
                 continue