Extend adjoint action to work with mixed spaces (#352)

* Update version number.

* Fix

* Correct tag ref

* Use BaseArgument.__eq__ in Argument (#147)

(cherry picked from commit e683148)

* Oops! Bump version number.

* Fix.

* Bump version.

* Updated to .md README (#275)

* Kebab case in build-wheels.yml (#276)

* Fixes for pypa packaging.

* Correct documentation links (#277)

* Remove unecessary pip pinning (#278)

* Bump version.

* Update version to 2024.2.0

* Extending ufl.lhs and ufl.rhs to mixed domains variational forms

* Add arity optional argument to extract_blocks to be used in compute_form_rhs/lhs

* Fix ruff

* Fix ruff

* Fix lhs / rhs : No need to extract blocks if parts=()

* Extend `MixedFunctionSpace with `ufl.rhs`, `ufl.lhs` and thus `ufl.system`.

* Add `ufl.action` for MixedFunctionSpace

* Extend adjoint action to work with mixed spaces

---------

Co-authored-by: Jack S. Hale <mail@jackhale.co.uk>
Co-authored-by: Chris Richardson <chris@bpi.cam.ac.uk>
Co-authored-by: Michal Habera <michal.habera@gmail.com>
Co-authored-by: Garth N. Wells <gnw20@cam.ac.uk>
Co-authored-by: Cécile <cecile@simula.no>

Author: 6 people

test/test_mixed_function_space.py

@@ -11,6 +11,7 @@
     TestFunctions,
     TrialFunctions,
     action,
+    conj,
     dx,
     grad,
     inner,
@@ -23,6 +24,7 @@
 )
 from ufl.algorithms import expand_derivatives, renumbering
 from ufl.algorithms.formsplitter import extract_blocks
+from ufl.algorithms.formtransformations import compute_form_adjoint
 from ufl.finiteelement import FiniteElement
 from ufl.pullback import identity_pullback
 from ufl.sobolevspace import H1
@@ -134,7 +136,6 @@ def source2(f, v):
     f = Constant(domain)
     g = Constant(domain)
     F = mass(u, v) + mixed(u, q) + stiffness(p, q) + source1(f, v) + source2(g, q)
-
     a = lhs(F)
     a_blocked = extract_blocks(a)
     L = rhs(F)
@@ -180,7 +181,6 @@ def source2(f, v):
     F_reduced_renumbered = renumbering.renumber_indices(expand_derivatives(F_reduced))
     assert len(F_reduced_renumbered.coefficients()) == 1
     inserted_coeff = F_reduced_renumbered.coefficients()[0]
-
     # Create reference solution
     Fh = mass(inserted_coeff, v) + stiffness(inserted_coeff, q) + source1(f, v) + source2(g, q)
 
@@ -197,3 +197,32 @@ def source2(f, v):
     J_ref = replace(J_exp, {u: inserted_coeff, v: coefficients[0], q: coefficients[1]})
     # Verify
     assert J_renumbered == J_ref
+
+
+def test_adjoint():
+    V = FiniteElement("Lagrange", triangle, 1, (), identity_pullback, H1)
+    domain = Mesh(FiniteElement("Lagrange", triangle, 1, (2,), identity_pullback, H1))
+    space0 = FunctionSpace(domain, V)
+    space1 = FunctionSpace(domain, V)
+    mixed_space = MixedFunctionSpace(space0, space1)
+
+    u, p = TrialFunctions(mixed_space)
+    du, dp = TestFunctions(mixed_space)
+    c = Coefficient(space0)
+    Jh = (
+        inner(grad(u), grad(du)) * dx
+        + inner(c * dp.dx(0), u) * dx
+        - inner(du.dx(0), p.dx(1)) * dx
+        + inner(p, dp) * dx
+    )
+    Jh_adj = compute_form_adjoint(Jh)
+    blocked_adj = extract_blocks(Jh_adj)
+
+    ref_adj = [
+        [conj(inner(grad(du), grad(u))) * dx, conj(-inner(p.dx(0), du.dx(1))) * dx],
+        [conj(inner(c * u.dx(0), dp)) * dx, conj(inner(dp, p)) * dx],
+    ]
+
+    for i in range(2):
+        for j in range(2):
+            assert ref_adj[i][j] == blocked_adj[i][j]

ufl/algorithms/formtransformations.py

@@ -496,7 +496,39 @@ def compute_form_adjoint(form, reordered_arguments=None):
 
     parts = [arg.part() for arg in arguments]
     if set(parts) - {None}:
-        raise ValueError("compute_form_adjoint cannot handle parts.")
+        J = extract_blocks(form, arity=2)
+        num_blocks = len(J)
+        J_adj = 0
+        for i in range(num_blocks):
+            for j in range(num_blocks):
+                if J[i][j] is None:
+                    continue
+                v, u = J[i][j].arguments()
+                if reordered_arguments is None:
+                    reordered_u = Argument(u.ufl_function_space(), number=v.number(), part=v.part())
+                    reordered_v = Argument(v.ufl_function_space(), number=u.number(), part=u.part())
+                else:
+                    reordered_u, reordered_v = reordered_arguments[i]
+
+                if reordered_u.number() >= reordered_v.number():
+                    raise ValueError("Ordering of new arguments is the same as the old arguments!")
+
+                if reordered_u.part() != v.part():
+                    raise ValueError("Ordering of new arguments is the same as the old arguments!")
+                if reordered_v.part() != u.part():
+                    raise ValueError("Ordering of new arguments is the same as the old arguments!")
+
+                if reordered_u.ufl_function_space() != u.ufl_function_space():
+                    raise ValueError(
+                        "Element mismatch between new and old arguments (trial functions)."
+                    )
+                if reordered_v.ufl_function_space() != v.ufl_function_space():
+                    raise ValueError(
+                        "Element mismatch between new and old arguments (test functions)."
+                    )
+
+                J_adj += map_integrands(Conj, replace(J[i][j], {v: reordered_v, u: reordered_u}))
+        return J_adj
 
     if len(arguments) != 2:
         raise ValueError("Expecting bilinear form.")