Backport bug fixes (#459)

* Fix split for symmetric tensors (#448)

* BUG: prevent ReferenceGradient(Zero) on CoefficientSplitter (#455)

* BUG: prevent ReferenceGradient(Zero) on CoefficientSplitter

* test

* fix for mixed dimension cells

* cleanup

* Pullback: remove FunctionSpace (#452)

* Pullback: remove FunctionSpace

* cast to int

* import Zero

* CoefficientSplitter: fix again ReferenceGrad(Zero) (#460)

Author: pbrubeck

test/test_apply_coefficent_split.py

@@ -7,6 +7,7 @@
     MeshSequence,
     triangle,
 )
+from ufl.algorithms import extract_coefficients
 from ufl.algorithms.apply_coefficient_split import apply_coefficient_split
 from ufl.classes import (
     ComponentTensor,
@@ -64,3 +65,25 @@ def test_apply_coefficient_split(self):
     op1_, (idx1_,) = op1.ufl_operands
     assert op1_ == PositiveRestricted(ReferenceGrad(ReferenceValue(f1)))
     assert idx1_ == idx1
+
+
+def test_derivative_zero_simplication():
+    cell = triangle
+    mesh = Mesh(LagrangeElement(cell, 1, (2,)))
+    elem0 = LagrangeElement(cell, 0)
+    elem1 = LagrangeElement(cell, 3)
+    V0 = FunctionSpace(mesh, elem0)
+    V1 = FunctionSpace(mesh, elem1)
+    f0 = Coefficient(V0)
+    f1 = Coefficient(V1)
+    elem = MixedElement([elem0, elem1])
+    V = FunctionSpace(mesh, elem)
+    f = Coefficient(V)
+
+    coefficient_split = {f: (f0, f1)}
+    expr = ReferenceGrad(ReferenceGrad(ReferenceValue(f)))
+    expr_split = apply_coefficient_split(expr, coefficient_split)
+
+    coefficients = extract_coefficients(expr_split)
+    assert f1 in coefficients
+    assert f0 not in coefficients

test/test_derivative.py

@@ -3,7 +3,7 @@
 
 from itertools import chain
 
-from utils import FiniteElement, LagrangeElement, MixedElement
+from utils import FiniteElement, LagrangeElement, MixedElement, SymmetricElement
 
 from ufl import (
     CellDiameter,
@@ -574,6 +574,37 @@ def test_multiple_indexed_coefficient_derivative(self):
     assertEqualBySampling(actual, expected)
 
 
+def test_split_tensor_coefficient_derivative(self):
+    cell = triangle
+    dim = 2
+    domain = Mesh(LagrangeElement(cell, 1, (dim,)))
+    symmetry = {(0, 0): 0, (0, 1): 1, (1, 0): 1, (1, 1): 2}
+    ncomp = (dim * (dim + 1)) // 2
+
+    S = SymmetricElement(symmetry, [LagrangeElement(cell, 1)] * ncomp)
+    V = LagrangeElement(cell, 1, (dim,))
+
+    element = MixedElement([S, V])
+    Z = FunctionSpace(domain, element)
+    z = Coefficient(Z)
+    test = TestFunction(Z)
+    s, u = split(z)
+    t, v = split(test)
+
+    J = 0.5 * inner(s, s) + inner(s, nabla_grad(u))
+
+    dJ_du = derivative(J, u, v)
+    dJ_ds = derivative(J, s, t)
+
+    actual = dJ_ds + dJ_du
+
+    expected = (
+        0.5 * inner(t, s) + 0.5 * inner(s, t) + inner(t, nabla_grad(u)) + inner(s, nabla_grad(v))
+    )
+
+    assertEqualBySampling(actual, expected)
+
+
 def test_segregated_derivative_of_convection(self):
     cell = tetrahedron
     V = LagrangeElement(cell, 1)

test/test_split.py

@@ -43,10 +43,10 @@ def test_split(self):
 
     # Shapes of subelements are reproduced:
     g = Coefficient(m_space)
-    (s,) = g.ufl_shape
+    (size,) = g.ufl_shape
     for g2 in split(g):
-        s -= product(g2.ufl_shape)
-    assert s == 0
+        size -= product(g2.ufl_shape)
+    assert size == 0
 
     # Mixed elements of non-scalar subelements are flattened
     v2 = MixedElement([v, v])
@@ -78,3 +78,12 @@ def test_split(self):
     assert split(split(split(t)[0])[1]) == (t[2],)
     assert split(split(split(t)[1])[0]) == (t[3],)
     assert split(split(split(t)[1])[1]) == (t[4], t[5])
+
+    # Split twice on nested mixed elements with symmetry
+    vs = MixedElement([v, s])
+    vs_space = FunctionSpace(domain, vs)
+    vs_test = TestFunction(vs_space)
+
+    v_test, s_test = split(vs_test)
+    assert split(v_test) == (vs_test[0], vs_test[1])
+    assert split(s_test) == (vs_test[2], vs_test[3], vs_test[4], vs_test[5])

ufl/algorithms/apply_coefficient_split.py

@@ -20,6 +20,7 @@
     ReferenceValue,
     Restricted,
     Terminal,
+    Zero,
 )
 from ufl.core.multiindex import indices
 from ufl.corealg.dag_traverser import DAGTraverser
@@ -209,8 +210,16 @@ def _handle_terminal(
         c = o
         if reference_value:
             c = ReferenceValue(c)
-        for _ in range(reference_grad):
+        for k in range(reference_grad):
             c = ReferenceGrad(c)
+            if isinstance(c, Zero):
+                gdim = c.ufl_shape[-1]
+                c = Zero(
+                    c.ufl_shape + (gdim,) * (reference_grad - k - 1),
+                    c.ufl_free_indices,
+                    c.ufl_index_dimensions,
+                )
+                break
         if restricted == "+":
             c = PositiveRestricted(c)
         elif restricted == "-":

ufl/formoperators.py

@@ -9,6 +9,8 @@
 # Modified by Massimiliano Leoni, 2016
 # Modified by Cecile Daversin-Catty, 2018
 
+import numpy as np
+
 from ufl.action import Action
 from ufl.adjoint import Adjoint
 from ufl.algorithms import (
@@ -270,7 +272,9 @@ def set_list_item(li, i, v):
 def _handle_derivative_arguments(form, coefficient, argument):
     """Handle derivative arguments."""
     # Wrap single coefficient in tuple for uniform treatment below
-    if isinstance(coefficient, list | tuple | ListTensor):
+    if isinstance(coefficient, ListTensor):
+        coefficients = tuple(coefficient[i] for i in np.ndindex(coefficient.ufl_shape))
+    elif isinstance(coefficient, list | tuple):
         coefficients = tuple(coefficient)
     else:
         coefficients = (coefficient,)

ufl/pullback.py

@@ -15,7 +15,6 @@
 from ufl.core.expr import Expr
 from ufl.core.multiindex import indices
 from ufl.domain import AbstractDomain, MeshSequence, extract_unique_domain
-from ufl.functionspace import FunctionSpace
 from ufl.tensors import as_tensor
 
 if TYPE_CHECKING:
@@ -428,15 +427,14 @@ def apply(self, expr, domain=None):
             subdomain = domain[i] if isinstance(domain, MeshSequence) else None
             rmapped = subelem.pullback.apply(rsub, domain=subdomain)
             # Flatten into the pulled back expression for the whole thing
-            g_components.extend([rmapped[idx] for idx in np.ndindex(rmapped.ufl_shape)])
+            g_components.extend(rmapped[idx] for idx in np.ndindex(rmapped.ufl_shape))
             offset += subelem.reference_value_size
         # And reshape appropriately
-        space = FunctionSpace(domain, self._element)
-        f = as_tensor(np.asarray(g_components).reshape(space.value_shape))
-        if f.ufl_shape != space.value_shape:
+        value_shape = self.physical_value_shape(self._element, domain)
+        f = as_tensor(np.asarray(g_components).reshape(value_shape))
+        if f.ufl_shape != value_shape:
             raise ValueError(
-                "Expecting pulled back expression with shape "
-                f"'{space.value_shape}', got '{f.ufl_shape}'"
+                f"Expecting pulled back expression with shape '{value_shape}', got '{f.ufl_shape}'"
             )
         return f
 
@@ -453,7 +451,8 @@ def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         assert element == self._element
         domains = domain.iterable_like(element)
         dim = sum(
-            FunctionSpace(d, e).value_size for d, e in zip(domains, self._element.sub_elements)
+            int(np.prod(e.pullback.physical_value_shape(e, d), dtype=int))
+            for d, e in zip(domains, self._element.sub_elements)
         )
         return (dim,)
 
@@ -496,8 +495,6 @@ def apply(self, expr, domain=None):
 
         Returns: The function pulled back to the reference cell
         """
-        domain = extract_unique_domain(expr)
-        space = FunctionSpace(domain, self._element)
         rflat = [expr[idx] for idx in np.ndindex(expr.ufl_shape)]
         g_components = []
         offsets = [0]
@@ -520,13 +517,13 @@ def apply(self, expr, domain=None):
             subdomain = domain[i] if isinstance(domain, MeshSequence) else None
             rmapped = subelem.pullback.apply(rsub, domain=subdomain)
             # Flatten into the pulled back expression for the whole thing
-            g_components.extend([rmapped[idx] for idx in np.ndindex(rmapped.ufl_shape)])
+            g_components.extend(rmapped[idx] for idx in np.ndindex(rmapped.ufl_shape))
         # And reshape appropriately
-        f = as_tensor(np.asarray(g_components).reshape(space.value_shape))
-        if f.ufl_shape != space.value_shape:
+        value_shape = self.physical_value_shape(self._element, domain)
+        f = as_tensor(np.asarray(g_components).reshape(value_shape))
+        if f.ufl_shape != value_shape:
             raise ValueError(
-                f"Expecting pulled back expression with shape "
-                f"'{space.value_shape}', got '{f.ufl_shape}'"
+                f"Expecting pulled back expression with shape '{value_shape}', got '{f.ufl_shape}'"
             )
         return f
 
@@ -543,7 +540,7 @@ def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         assert isinstance(element, type(self._element))
         subelem = element.sub_elements[0]
         pvs = subelem.pullback.physical_value_shape(subelem, domain)
-        return tuple(i + 1 for i in max(self._symmetry.keys())) + pvs
+        return tuple(int(i) + 1 for i in max(self._symmetry.keys())) + pvs
 
 
 class PhysicalPullback(AbstractPullback):

ufl/split_functions.py

@@ -6,12 +6,16 @@
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 #
 # Modified by Anders Logg, 2008
+# Modified by Pablo Brubeck, 2025
+
+import numpy as np
 
 from ufl.domain import extract_unique_domain
 from ufl.functionspace import FunctionSpace
 from ufl.indexed import Indexed
 from ufl.permutation import compute_indices
-from ufl.tensors import ListTensor, as_matrix, as_vector
+from ufl.pullback import SymmetricPullback
+from ufl.tensors import ListTensor, as_tensor
 from ufl.utils.indexflattening import flatten_multiindex, shape_to_strides
 from ufl.utils.sequences import product
 
@@ -33,7 +37,8 @@ def split(v):
 
     elif isinstance(v, ListTensor):
         # Special case: split previous output of split again
-        ops = v.ufl_operands
+        ops = tuple(v[i] for i in np.ndindex(v.ufl_shape))
+
         if all(isinstance(comp, Indexed) for comp in ops):
             args = [comp.ufl_operands[0] for comp in ops]
             if all(args[0] == args[i] for i in range(1, len(args))):
@@ -85,8 +90,18 @@ def split(v):
     # Build expressions representing the subfunction of v for each subelement
     offset = begin
     sub_functions = []
-    domains = domain.iterable_like(element)
-    for i, (d, e) in enumerate(zip(domains, element.sub_elements)):
+
+    # Deal with symmetry
+    if isinstance(element.pullback, SymmetricPullback):
+        symmetry = element.pullback._symmetry
+        indices = (comp for idx, comp in sorted(symmetry.items()))
+    else:
+        indices = range(len(element.sub_elements))
+
+    domains = tuple(domain.iterable_like(element))
+    for i in indices:
+        d = domains[i]
+        e = element.sub_elements[i]
         # Get shape, size, indices, and v components
         # corresponding to subelement value
         shape = FunctionSpace(d, e).value_shape
@@ -99,16 +114,8 @@ def split(v):
         # Shape components into same shape as subelement
         if rank == 0:
             (subv,) = components
-        elif rank <= 1:
-            subv = as_vector(components)
-        elif rank == 2:
-            subv = as_matrix(
-                [components[i * shape[1] : (i + 1) * shape[1]] for i in range(shape[0])]
-            )
         else:
-            raise ValueError(
-                f"Don't know how to split functions with sub functions of rank {rank}."
-            )
+            subv = as_tensor(np.reshape(components, shape))
 
         offset += sub_size
         sub_functions.append(subv)