test_forms.py

"""Tests for DOLFINx integration of various form operations."""

# Copyright (C) 2021-2026 Garth N. Wells and Jose Fernandez
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later

from mpi4py import MPI

import numpy as np
import pytest

import basix
import basix.ufl
import dolfinx
from dolfinx.fem import IntegralType, extract_function_spaces, form, functionspace
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, Form as ufl_form, TrialFunctions, TestFunctions, MixedFunctionSpace


def test_extract_forms():
    """Test extraction on unique function spaces for rows and columns of
    a block system.
    """
    mesh = create_unit_square(MPI.COMM_WORLD, 32, 31)
    V0 = functionspace(mesh, ("Lagrange", 1))
    V1 = functionspace(mesh, ("Lagrange", 2))
    V2 = V0.clone()
    V3 = V1.clone()

    v0, u0 = TestFunction(V0), TrialFunction(V0)
    v1, u1 = TestFunction(V1), TrialFunction(V1)
    v2, u2 = TestFunction(V2), TrialFunction(V2)
    v3, u3 = TestFunction(V3), TrialFunction(V3)

    a = form([[inner(u0, v0) * dx, inner(u1, v1) * dx], [inner(u2, v2) * dx, inner(u3, v3) * dx]])
    with pytest.raises(AssertionError):
        extract_function_spaces(a, 0)
    with pytest.raises(AssertionError):
        extract_function_spaces(a, 1)

    a = form([[inner(u0, v0) * dx, inner(u2, v1) * dx], [inner(u0, v2) * dx, inner(u2, v2) * dx]])
    with pytest.raises(AssertionError):
        extract_function_spaces(a, 0)
    Vc = extract_function_spaces(a, 1)
    assert Vc[0] is V0._cpp_object
    assert Vc[1] is V2._cpp_object

    a = form([[inner(u0, v0) * dx, inner(u1, v0) * dx], [inner(u2, v1) * dx, inner(u3, v1) * dx]])
    Vr = extract_function_spaces(a, 0)
    assert Vr[0] is V0._cpp_object
    assert Vr[1] is V1._cpp_object
    with pytest.raises(AssertionError):
        extract_function_spaces(a, 1)


def test_incorrect_element():
    """Test that an error is raised if an incorrect element is used."""
    mesh = create_unit_square(MPI.COMM_WORLD, 32, 31)
    element = basix.ufl.element(
        "Lagrange",
        "triangle",
        4,
        lagrange_variant=basix.LagrangeVariant.gll_warped,
        dtype=dolfinx.default_real_type,
    )
    incorrect_element = basix.ufl.element(
        "Lagrange",
        "triangle",
        4,
        lagrange_variant=basix.LagrangeVariant.equispaced,
        dtype=dolfinx.default_real_type,
    )

    space = functionspace(mesh, element)
    incorrect_space = functionspace(mesh, incorrect_element)

    u = TrialFunction(space)
    v = TestFunction(space)

    a = inner(u, v) * dx

    dtype = dolfinx.default_scalar_type
    ftype = form_cpp_class(dtype)

    ufcx_form, module, code = dolfinx.jit.ffcx_jit(
        mesh.comm, a, form_compiler_options={"scalar_type": dtype}
    )

    f = ftype(
        [module.ffi.cast("uintptr_t", module.ffi.addressof(ufcx_form))],
        [space._cpp_object, space._cpp_object],
        [],
        [],
        {IntegralType.cell: []},
        [],
        mesh._cpp_object,
    )
    dolfinx.fem.Form(f, ufcx_form, code)

    with pytest.raises(RuntimeError):
        f = ftype(
            [module.ffi.cast("uintptr_t", module.ffi.addressof(ufcx_form))],
            [incorrect_space._cpp_object, incorrect_space._cpp_object],
            [],
            [],
            {IntegralType.cell: []},
            [],
            mesh._cpp_object,
        )
        dolfinx.fem.Form(f, ufcx_form, code)


def test_multiple_measures_one_subdomain_data():
    comm = MPI.COMM_WORLD
    mesh = dolfinx.mesh.create_unit_interval(comm, 10)
    x = SpatialCoordinate(mesh)
    num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local
    ct = dolfinx.mesh.meshtags(
        mesh,
        mesh.topology.dim,
        np.arange(num_cells_local, dtype=np.int32),
        np.arange(num_cells_local, dtype=np.int32),
    )

    dx = Measure("dx", domain=mesh, subdomain_data=ct)
    dx_stand = Measure("dx", domain=mesh)

    J = dolfinx.fem.form(x[0] ** 2 * dx + x[0] * dx_stand)
    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])