diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 242c88ded1..db86d82813 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) # @@ -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." + ) diff --git a/python/test/unit/fem/test_forms.py b/python/test/unit/fem/test_forms.py index d08c07f295..35f6c730f6 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,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])