Split EquationBCSplit

Author: pbrubeck

firedrake/formmanipulation.py

@@ -160,7 +160,7 @@ def cofunction(self, o):
             return Cofunction(W, val=MixedDat(o.dat[i] for i in indices))
 
     def matrix(self, o):
-        from firedrake.bcs import DirichletBC
+        from firedrake.bcs import DirichletBC, EquationBCSplit
         ises = []
         args = []
         for a in o.arguments():
@@ -180,6 +180,7 @@ def matrix(self, o):
             args.append(asplit)
 
         submat = o.petscmat.createSubMatrix(*ises)
+
         bcs = []
         spaces = [a.function_space() for a in o.arguments()]
         for bc in o.bcs:
@@ -191,11 +192,14 @@ def matrix(self, o):
             V = args[number].function_space()
             field = self.blocks[number]
             if isinstance(bc, DirichletBC):
-                sub_bc = bc.reconstruct(field=field, V=V, g=bc.function_arg)
-            else:
-                raise NotImplementedError(f"Extracting matrix subblocks not supported with {type(bc).__name__}. Please get in touch if you need this.")
-            if sub_bc is not None:
-                bcs.append(sub_bc)
+                bc_temp = bc.reconstruct(field=field, V=V, g=bc.function_arg, use_split=True)
+            elif isinstance(bc, EquationBCSplit):
+                row_field = self.blocks[0]
+                col_field = self.blocks[1]
+                bc_temp = bc.reconstruct(field=field, V=V, row_field=row_field, col_field=col_field, use_split=True)
+            if bc_temp is not None:
+                bcs.append(bc_temp)
+
         return AssembledMatrix(tuple(args), tuple(bcs), submat)
 
     def zero_base_form(self, o):

firedrake/solving_utils.py

@@ -523,15 +523,6 @@ def form_jacobian(snes, X, J, P):
         ctx.set_nullspace(ctx._nullspace_T, ises, transpose=True, near=False)
         ctx.set_nullspace(ctx._near_nullspace, ises, transpose=False, near=True)
 
-        # Bump petsc matrix state of each split by assembling them.
-        # Ensures that if the matrix changed, the preconditioner is
-        # updated if necessary.
-        for fields, splits in ctx._splits.items():
-            for subctx in splits:
-                subctx._jac.assemble()
-                if subctx.Jp is not None:
-                    subctx._pjac.assemble()
-
     @staticmethod
     def compute_operators(ksp, J, P):
         r"""Form the Jacobian for this problem