Add arity property to form

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)

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.