Fix ZeroBaseForm holes in derivative expansion and action

test/test_duals.py

@@ -28,9 +28,11 @@
     triangle,
 )
 from ufl.algorithms.ad import expand_derivatives
+from ufl.algorithms.analysis import extract_type
 from ufl.constantvalue import Zero
+from ufl.differentiation import CoefficientDerivative
 from ufl.duals import is_dual, is_primal
-from ufl.form import ZeroBaseForm
+from ufl.form import Form, ZeroBaseForm
 
 
 def test_mixed_functionspace(self):
@@ -268,6 +270,23 @@ def test_action():
     res = action(a + a2, u)
     assert res == Action(a, u) + Action(a2, u)
 
+    # `action` must expand derivatives even when the top-level object is a
+    # FormSum (not just a bare Form), so that a CoefficientDerivative left
+    # over from an affine residual (here `F - b`, with `b` a Cofunction
+    # independent of `w`) does not survive unexpanded inside the resulting
+    # Action (regression test for compute_form_action choking on a
+    # CoefficientDerivative reachable only through a FormSum).
+    w = Coefficient(U)
+    dw = Argument(U, 1)
+    F = inner(w**3, u_a) * dx
+    residual = F - ustar
+    J = derivative(residual, w, dw)
+    assert isinstance(J, FormSum)
+    dw2 = Coefficient(U)
+    Jw = action(J, dw2)
+    assert isinstance(Jw, Form)
+    assert not extract_type(Jw, CoefficientDerivative)
+
 
 def test_differentiation():
     domain_2d = Mesh(LagrangeElement(triangle, 1, (2,)))
@@ -315,6 +334,18 @@ def test_differentiation():
     assert dMdu.arguments() == M.arguments() + (Argument(u.ufl_function_space(), 2),)
     # Check compatibility with int/float
     assert dMdu == 0
+    assert dMdu.empty()
+
+    # Second derivative of a Matrix: differentiating an already-zero
+    # ZeroBaseForm must not raise (regression test for a missing
+    # GateauxDerivativeRuleset rule for ZeroBaseForm).
+    d2Mdu2 = expand_derivatives(derivative(derivative(M, u), u))
+    assert isinstance(d2Mdu2, ZeroBaseForm)
+    # One new direction Argument is picked up per differentiation, on top
+    # of the two Arguments of M.
+    assert d2Mdu2.arguments()[:2] == M.arguments()
+    assert len(d2Mdu2.arguments()) == 4
+    assert d2Mdu2.empty()
 
     # -- Action -- #
     Ac = Action(w, u)
@@ -336,6 +367,19 @@ def test_differentiation():
     # Distribute differentiation over FormSum components
     assert dFsdu == inner(what * uhat, v) * dx
 
+    # A FormSum whose components all vanish under differentiation must
+    # collapse to a properly shaped ZeroBaseForm rather than a bare `0`
+    # that silently drops the arguments (regression test for
+    # map_integrands losing arguments via `sum([])`).
+    p = Coefficient(U)
+    b1 = Cofunction(U.dual())
+    b2 = Cofunction(U.dual())
+    residual = b1 - b2
+    dresidualdp = expand_derivatives(derivative(residual, p, uhat))
+    assert isinstance(dresidualdp, ZeroBaseForm)
+    assert dresidualdp.arguments() == b1.arguments() + (uhat,)
+    assert dresidualdp.empty()
+
 
 def test_zero_base_form_mult():
     domain_2d = Mesh(LagrangeElement(triangle, 1, (2,)))

ufl/algorithms/apply_derivatives.py

@@ -1650,6 +1650,13 @@ def _(self, M: Matrix) -> BaseForm:
         # the appropriate space
         return ZeroBaseForm(M.arguments() + self._v)
 
+    @process.register(ZeroBaseForm)  # type: ignore
+    def _(self, o: BaseForm) -> BaseForm:
+        """Differentiate a zero_base_form."""
+        # ZeroBaseForm is idempotent under differentiation: it stays zero,
+        # gaining the new derivative direction as an extra argument.
+        return ZeroBaseForm(o.arguments() + self._v)
+
 
 class BaseFormOperatorDerivativeRuleset(GateauxDerivativeRuleset):
     """Apply AFD (Automatic Functional Differentiation) to BaseFormOperator.
@@ -2013,8 +2020,9 @@ def apply_derivatives(expression):
         and isinstance(dexpression_dvar, int)
         and dexpression_dvar == 0
     ):
-        # The arguments got lost, just keep an empty Form
-        dexpression_dvar = Form([])
+        # Algebraic cancellation collapsed everything to a bare `0`: rebuild
+        # a properly shaped `ZeroBaseForm` rather than an argument-less Form.
+        dexpression_dvar = ZeroBaseForm(expression.arguments())
 
     # Get the recorded delayed operations
     pending_operations = dag_traverser.pending_operations

ufl/algorithms/map_integrands.py

@@ -51,6 +51,12 @@ def map_integrands(function, form, only_integral_type=None):
             if component != 0
         ]
 
+        if not nonzero_components:
+            # Every component vanished: preserve the argument shape (which may
+            # include a new derivative direction picked up while mapping)
+            # instead of collapsing to a bare `0` and losing it.
+            return ZeroBaseForm(mapped_components[0].arguments())
+
         # Simplify case with one nonzero component and the corresponding weight is 1
         if len(nonzero_components) == 1 and nonzero_components[0][1] == 1:
             return nonzero_components[0][0]

ufl/form.py

@@ -912,6 +912,10 @@ def ufl_domains(self):
             self._analyze_domains()
         return self._domains
 
+    def empty(self):
+        """Returns whether the ZeroBaseForm has no components, which is always true."""
+        return True
+
     def __ne__(self, other):
         """Overwrite BaseForm.__neq__ which relies on `equals`."""
         return not self == other

ufl/formoperators.py

@@ -197,10 +197,11 @@ def action(form, coefficient=None, derivatives_expanded=None):
     is_coefficient_valid = not isinstance(coefficient, BaseForm) or (
         isinstance(coefficient, BaseFormOperator) and len(coefficient.arguments()) == 1
     )
-    # Can't expand derivatives on objects that are not Form or Expr (e.g. Matrix)
-    if isinstance(form, Form | BaseFormOperator) and is_coefficient_valid:
+    if isinstance(form, BaseForm) and is_coefficient_valid:
         if not derivatives_expanded:
-            # For external operators differentiation may turn a Form into a FormSum
+            # For external operators differentiation may turn a Form into a FormSum,
+            # and a FormSum's components may themselves simplify back into a single
+            # Form (e.g. a term cancelling against a ZeroBaseForm).
             form = expand_derivatives(form)
         if isinstance(form, Form):
             return compute_form_action(form, coefficient)