Support empty Jacobian blocks

firedrakeproject · Jan 1, 2025 · cf743a1 · cf743a1
1 parent 2286596
commit cf743a1
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py
@@ -26,9 +26,11 @@ def indexed(self, o, child, multiindex):
             indices = multiindex.indices()
             if isinstance(child, ListTensor) and all(isinstance(i, FixedIndex) for i in indices):
                 if len(indices) == 1:
-                    return child.ufl_operands[indices[0]._value]
+                    return child[indices[0]]
+                elif len(indices) == len(child.ufl_operands and all(k == int(i) for k, i in enumerate(indices))):
+                    return child
                 else:
-                    return ListTensor(*(child.ufl_operands[i._value] for i in multiindex.indices()))
+                    return ListTensor(*(child[i] for i in indices))
             return self.expr(o, child, multiindex)
 
     index_inliner = IndexInliner()
@@ -119,12 +121,11 @@ def argument(self, o):
                 c = indices.index(i)
                 a_ = a[c]
                 if len(a_.ufl_shape) == 0:
-                    args += [a_]
+                    args.append(a_)
                 else:
-                    args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)]
+                    args.extend(a_[j] for j in numpy.ndindex(a_.ufl_shape))
             else:
-                args += [Zero()
-                         for j in numpy.ndindex(V_is[i].value_shape)]
+                args.extend(Zero() for j in numpy.ndindex(V_is[i].value_shape))
         return self._arg_cache.setdefault(o, as_vector(args))
 
 
@@ -168,7 +169,7 @@ def split_form(form, diagonal=False):
         assert len(shape) == 2
     for idx in numpy.ndindex(shape):
         f = splitter.split(form, idx)
-        if len(f.integrals()) > 0:
+        if not f.empty():
             if diagonal:
                 i, j = idx
                 if i != j:

diff --git a/tests/firedrake/slate/test_assemble_tensors.py b/tests/firedrake/slate/test_assemble_tensors.py
@@ -249,9 +249,10 @@ def test_matrix_subblocks(mesh):
     refs = dict(split_form(A.form))
     _A = A.blocks
     for x, y in indices:
-        ref = assemble(refs[x, y]).M.values
         block = _A[x, y]
-        assert np.allclose(assemble(block).M.values, ref, rtol=1e-14)
+        if not block.form.empty():
+            ref = assemble(refs[x, y]).M.values
+            assert np.allclose(assemble(block).M.values, ref, rtol=1e-14)
 
     # Mixed blocks
     A0101 = _A[:2, :2]
@@ -267,17 +268,13 @@ def test_matrix_subblocks(mesh):
     A0101_10 = _A0101[1, 0]
     A1212_00 = _A1212[0, 0]
     A1212_11 = _A1212[1, 1]
-    A1212_01 = _A1212[0, 1]
-    A1212_10 = _A1212[1, 0]
 
     items = [(A0101_00, refs[(0, 0)]),
              (A0101_11, refs[(1, 1)]),
              (A0101_01, refs[(0, 1)]),
              (A0101_10, refs[(1, 0)]),
              (A1212_00, refs[(1, 1)]),
-             (A1212_11, refs[(2, 2)]),
-             (A1212_01, refs[(1, 2)]),
-             (A1212_10, refs[(2, 1)])]
+             (A1212_11, refs[(2, 2)])]
 
     # Test assembly of blocks of mixed blocks
     for tensor, form in items: