diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 1ab16e02d5..77434579d4 100644 --- a/python/dolfinx/fem/forms.py +++ b/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) # @@ -685,35 +685,123 @@ 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[Function], + 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[Function], + 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 form or list of forms. + + 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 :func:`ufl.derivative` directly. + + 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], - This is commonly used to derive a block Jacobian from a block - residual. + where :math:`\\delta u_i` is a test subfunction of the mixed + space defined by ``u``. This is equivalent to calling + :func:`ufl.extract_blocks` on the result from + :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 - 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. + .. math:: + + J = \\frac{\\partial F}{\\partial u}[\\delta u]. + + This is equivalent to calling :func:`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}{\\partial u_j}[\\delta u_j] + + using :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) # type: ignore[arg-type] + 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." + ) diff --git a/python/test/unit/fem/test_forms.py b/python/test/unit/fem/test_forms.py index d08c07f295..9963c84108 100644 --- a/python/test/unit/fem/test_forms.py +++ b/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,20 @@ 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 derivative_block, form_cpp_class from dolfinx.mesh import create_unit_square -from ufl import Measure, SpatialCoordinate, TestFunction, TrialFunction, dx, inner +from ufl import Form as ufl_form +from ufl import ( + Measure, + MixedFunctionSpace, + SpatialCoordinate, + TestFunction, + TestFunctions, + TrialFunction, + TrialFunctions, + dx, + inner, +) def test_extract_forms(): @@ -132,3 +143,56 @@ 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 + )