diff --git a/test/test_form.py b/test/test_form.py index 0d680f82e..25ac96988 100755 --- a/test/test_form.py +++ b/test/test_form.py @@ -21,7 +21,9 @@ nabla_grad, triangle, ) +from ufl.argument import TestFunctions, TrialFunctions from ufl.form import BaseForm +from ufl.functionspace import MixedFunctionSpace @pytest.fixture @@ -37,6 +39,11 @@ def domain(): return Mesh(LagrangeElement(cell, 1, (2,))) +@pytest.fixture +def functional(domain): + return 1 * dx(domain) + + @pytest.fixture def mass(domain): cell = triangle @@ -104,6 +111,54 @@ def test_form_arguments(mass, stiffness, convection, load): assert ((f * v) * u * dx + (u * 3) * (v / 2) * dx(2)).arguments() == (v, u) +def test_form_arity(functional, mass, stiffness, convection, load) -> None: + assert functional.arity == 0 + assert mass.arity == 2 + assert stiffness.arity == 2 + assert convection.arity == 2 + assert load.arity == 1 + + assert (functional + load).arity is None + assert (functional + mass).arity is None + assert (functional + stiffness).arity is None + assert (load + mass).arity is None + assert (load + stiffness).arity is None + + +def test_form_arity_mixed(domain) -> None: + domain = Mesh(LagrangeElement(triangle, 1, (2,))) + V = FunctionSpace(domain, LagrangeElement(triangle, 1, (2,))) + W = FunctionSpace(domain, LagrangeElement(triangle, 1, (3,))) + + U = MixedFunctionSpace(V, W) + + v, sigma = TestFunctions(U) + u, tau = TrialFunctions(U) + + linear = v[0] * dx + sigma[0] * dx + bilinear = v[0] * u[0] * dx + sigma[0] * tau[0] * dx + + assert linear.arity == 1 + assert bilinear.arity == 2 + assert (linear + bilinear).arity is None + + # combined mixed form integrals (unsupported) + with pytest.raises( + RuntimeError, match=r"Arity does not support mixed arguments in an integral." + ): + assert ((v[0] + sigma[0]) * dx).arity == 1 + + with pytest.raises( + RuntimeError, match=r"Arity does not support mixed arguments in an integral." + ): + assert ((v[0] * u[0] + sigma[0] * tau[0]) * dx).arity == 2 + + with pytest.raises( + RuntimeError, match=r"Arity does not support mixed arguments in an integral." + ): + assert ((v[0] + sigma[0] + v[0] * u[0] + sigma[0] * tau[0]) * dx).arity is None + + def test_form_coefficients(element, domain): space = FunctionSpace(domain, element) v = TestFunction(space) diff --git a/ufl/form.py b/ufl/form.py index 9388d8f25..c155bf215 100644 --- a/ufl/form.py +++ b/ufl/form.py @@ -340,6 +340,37 @@ def empty(self): """Returns whether the form has no integrals.""" return len(self.integrals()) == 0 + @property + def arity(self) -> int | None: + """Arity of the form. + + Note: + Mixed function spaces are supported if the parts are not mixed in a single integral. + + Returns: + Number of arguments if all integrals share the argument count, otherwise None. + """ + if not self._integrals: + return 0 + + from ufl.algorithms.analysis import extract_arguments + + arity = None + for integral in self._integrals: + args = extract_arguments(integral.integrand()) + + if len(set(arg.part() for arg in args)) > 1: + raise RuntimeError("Arity does not support mixed arguments in an integral.") + + _arity = max((arg.number() + 1 for arg in args), default=0) + + if arity is None: + arity = _arity + elif arity != _arity: + return None + + return arity + def ufl_domains(self): """Return the geometric integration domains occuring in the form.