Do not split off-diagonal blocks if we only want the diagonal

firedrakeproject · Jan 2, 2025 · af53302 · af53302
1 parent d82039d
commit af53302
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py
@@ -196,14 +196,15 @@ def split_form(form, diagonal=False):
     args = form.arguments()
     shape = tuple(len(a.function_space()) for a in args)
     forms = []
+    arity = len(shape)
     if diagonal:
-        assert len(shape) == 2
+        assert arity == 2
+        arity = 1
     for idx in numpy.ndindex(shape):
-        f = splitter.split(form, idx)
         if diagonal:
             i, j = idx
             if i != j:
                 continue
-            idx = (i, )
-        forms.append(SplitForm(indices=idx, form=f))
+        f = splitter.split(form, idx)
+        forms.append(SplitForm(indices=idx[:arity], form=f))
     return tuple(forms)