Extend derivative block to derive residuals from functionals

python/dolfinx/fem/forms.py

@@ -1,5 +1,5 @@
 # Copyright (C) 2017-2026 Chris N. Richardson, Garth N. Wells,
-# Michal Habera, Jørgen S. Dokken and Jack S. Hale
+# Michal Habera, Jørgen S. Dokken, Jack S. Hale and Jose Fernandez
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -681,35 +681,107 @@ def create_form(
     return Form(f, form.ufcx_form, form.code)
 
 
+def _derive_univariate_residual(
+    F: ufl.Form,
+    u: Function,
+    du: ufl.Argument | None = None,
+) -> ufl.Form:
+    if du is None:
+        du = ufl.TestFunction(u.function_space)
+    return ufl.derivative(F, u, du)
+
+
+def _derive_block_residual(
+    F: ufl.Form,
+    u: Sequence[ufl.Form],
+    du: Sequence[ufl.Argument] | None = None,
+) -> Sequence[ufl.Form]:
+    if du is None:
+        du = ufl.TestFunctions(ufl.MixedFunctionSpace(*(u_i.function_space for u_i in u)))
+    return ufl.extract_blocks(ufl.derivative(F, u, du))
+
+
+def _derive_univariate_jacobian(
+    F: ufl.Form,
+    u: Function,
+    du: ufl.Argument | None = None,
+) -> ufl.Form:
+    if du is None:
+        du = ufl.TrialFunction(u.function_space)
+    return ufl.derivative(F, u, du)
+
+
+def _derive_block_jacobian(
+    F: Sequence[ufl.Form],
+    u: Sequence[ufl.Form],
+    du: Sequence[ufl.Argument] | None = None,
+) -> Sequence[Sequence[ufl.Form]]:
+    if not isinstance(u, Sequence):
+        raise ValueError("When F is a sequence, u must be a sequence")
+    if du is None:
+        du = [ufl.TrialFunction(u_i.function_space) for u_i in u]
+    elif (not isinstance(du, Sequence) or not len(u) == len(du)):
+        raise ValueError(
+            "When F is a list of N forms, du must be a sequence "
+            "containing N functions"
+        )
+    return [[ufl.derivative(F_i, u_j, du_j) for u_j, du_j in zip(u, du)] for F_i in F]
+
+
 def derivative_block(
     F: ufl.Form | Sequence[ufl.Form],
     u: Function | Sequence[Function],
     du: ufl.Argument | Sequence[ufl.Argument] | None = None,
-) -> ufl.Form | Sequence[Sequence[ufl.Form]]:
-    """Return the UFL derivative of a (list of) UFL rank one form(s).
+) -> ufl.Form | Sequence[ufl.Form] | Sequence[Sequence[ufl.Form]]:
+    """Return the UFL derivative of a UFL rank zero form, or the UFL derivative
+    of a (list of) rank one form(s).
+
+    This is commonly used to derive a block residual from a functional, or to
+    derive a block Jacobian from a block residual.
+
+    Four cases are supported:
+
+    1. ``F`` is a rank-zero ``ufl.Form``, and ``u`` is a ``ufl.Function``.
+        Returns a ``ufl.Form`` representing the residual :math:`R =
+        \\frac{\\partial F}{\\partial u}[\\delta u]`. This is equivalent to
+        calling {py:func}`ufl.derivative` directly.
 
-    This is commonly used to derive a block Jacobian from a block
-    residual.
+    2. ``F`` is a rank-zero `ufl.Form``, and ``u`` is a list of ``ufl.Function``.
+        Returns a list of ``ufl.Form`` representing the block residual :math:`R`,
+        with :math:`R_i = \\frac{\\partial F}{\\partial u_i}[\\delta u_i]`, where
+        :math:`\\delta u_i` is a test subfunction of the mixed space defined by
+        ``u``. This is equivalent to calling {py:func}`ufl.extract_blocks` on the
+        result from {py:func}`ufl.derivative`.
 
-    If ``F_i`` is a list of forms, the Jacobian is a list of lists with
-    :math:`J_{ij} = \\frac{\\partial F_i}{u_j}[\\delta u_j]` using
-    ``ufl.derivative`` called component-wise.
+    3. ``F`` is a rank-one `ufl.Form``, and ``u`` is a ``ufl.Function``.
+       Returns a ``ufl.Form`` representing the Jacobian :math:`J =
+       \\frac{\\partial F}{\\partial u}[\\delta u]`. This is equivalent to
+       calling {py:func}`ufl.derivative` directly.
 
-    If ``F`` is a form, the Jacobian is computed as :math:`J =
-    \\frac{\\partial F}{\\partial u}[\\delta u]`. This is identical to
-    calling ``ufl.derivative`` directly.
+    4. ``F`` is a list of rank-one `ufl.Form``, and ``u`` is a list of
+       ``ufl.Function``. Returns a list of lists representing the block Jacobian
+       :math:`J`, with :math:`J_{ij} = \\frac{\\partial F_i}{u_j}[\\delta u_j]`
+       using {py:func}`ufl.derivative` called component-wise.
+
+    Args:
+        F: UFL form(s) to be derived.
+        u: Function(s) with respect to the derivative is computed.
+        du: UFL argument(s) representing the direction of the derivative.
     """  # noqa: D301
-    if isinstance(F, ufl.Form):
-        if not isinstance(u, Function):
-            raise ValueError("Must provide a single function when F is a UFL form")
-        if du is None:
-            du = ufl.TrialFunction(u.function_space)
-        return ufl.derivative(F, u, du)
-    else:
-        assert all([isinstance(Fi, ufl.Form) for Fi in F]), "F must be a sequence of UFL forms"
-        assert len(F) == len(u), "Number of forms and functions must be equal"
-        if du is not None:
-            assert len(F) == len(du), "Number of forms and du must be equal"
+
+    if isinstance(F, ufl.Form) and not F.arguments():
+        if isinstance(u, Function):
+            return _derive_univariate_residual(F, u, du)
+        elif isinstance(u, Sequence):
+            return _derive_block_residual(F, u, du)
         else:
-            du = [ufl.TrialFunction(u_i.function_space) for u_i in u]
-        return [[ufl.derivative(Fi, u_j, du_j) for u_j, du_j in zip(u, du)] for Fi in F]
+            raise ValueError("u must be either a ufl.Function or a sequence of ufl.Function")
+    elif isinstance(F, ufl.Form) and len(F.arguments()) == 1:
+        return _derive_univariate_jacobian(F, u, du)
+    elif isinstance(F, Sequence):
+        return _derive_block_jacobian(F, u, du)
+    else:
+        raise ValueError(
+            "F must be either a UFL form (with rank zero or one), or a sequence of "
+            "rank-one UFL forms."
+        )

python/test/unit/fem/test_forms.py

@@ -1,6 +1,6 @@
 """Tests for DOLFINx integration of various form operations."""
 
-# Copyright (C) 2021 Garth N. Wells
+# Copyright (C) 2021-2026 Garth N. Wells and Jose Fernandez
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -15,9 +15,9 @@
 import basix.ufl
 import dolfinx
 from dolfinx.fem import IntegralType, extract_function_spaces, form, functionspace
-from dolfinx.fem.forms import form_cpp_class
+from dolfinx.fem.forms import form_cpp_class, derivative_block
 from dolfinx.mesh import create_unit_square
-from ufl import Measure, SpatialCoordinate, TestFunction, TrialFunction, dx, inner
+from ufl import Measure, SpatialCoordinate, TestFunction, TrialFunction, dx, inner, Form as ufl_form, TrialFunctions, TestFunctions, MixedFunctionSpace
 
 
 def test_extract_forms():
@@ -132,3 +132,54 @@ def test_multiple_measures_one_subdomain_data():
     J_local = dolfinx.fem.assemble_scalar(J)
     J_global = comm.allreduce(J_local, op=MPI.SUM)
     assert np.isclose(J_global, 1 / 3 + 1 / 2)
+
+
+def test_derivative_block():
+    """Test the function derivative_block"""
+    mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 10)
+    V0 = functionspace(mesh, ("Lagrange", 1))
+    V1 = functionspace(mesh, ("Lagrange", 2))
+    V = MixedFunctionSpace(V0, V1)
+
+    f0, f1 = dolfinx.fem.Function(V0), dolfinx.fem.Function(V1)
+    v0, v1 = TestFunctions(V)
+    u0, u1 = TrialFunctions(V)
+
+
+    M = f0**2 * dx  # univariate functional
+
+    F = derivative_block(M, f0)
+    assert isinstance(F, ufl_form) and len(F.arguments()) == 1
+
+    F = derivative_block(M, f0, v0)
+    assert isinstance(F, ufl_form) and len(F.arguments()) == 1
+
+    J = derivative_block(F, f0)
+    assert isinstance(J, ufl_form) and len(J.arguments()) == 2
+
+    J = derivative_block(F, f0, u0)
+    assert isinstance(J, ufl_form) and len(J.arguments()) == 2
+
+
+    M_block = f0**2 * f1 * dx  # multivariate functional
+
+    F_block = derivative_block(M_block, [f0, f1])
+    assert all([isinstance(F_i, ufl_form) and len(F_i.arguments()) == 1 for F_i in F_block])
+
+    F_block = derivative_block(M_block, [f0, f1], [v0, v1])
+    assert all([isinstance(F_i, ufl_form) and len(F_i.arguments()) == 1 for F_i in F_block])
+
+    with pytest.raises(ValueError):
+        derivative_block(F_block, f0)  # second argument not a sequence
+
+    with pytest.raises(ValueError):
+        derivative_block(F_block, [f0, f1], [u0])  # third argument has wrong length
+
+    with pytest.raises(ValueError):
+        derivative_block(F_block, [f0, f1], u0) # third argument not a sequence
+
+    J_block = derivative_block(F_block, [f0, f1])
+    assert all([isinstance(J_ij, ufl_form) and len(J_ij.arguments()) == 2 for J_i in J_block for J_ij in J_i])
+
+    J_block = derivative_block(F_block, [f0, f1], [u0, u1])
+    assert all([isinstance(J_ij, ufl_form) and len(J_ij.arguments()) == 2 for J_i in J_block for J_ij in J_i])