MatIS: support DirichletBC, add a test

Author: pbrubeck

pyop2/parloop.py

@@ -296,7 +296,8 @@ def replace_lgmaps(self):
                     olgmaps = []
                     for m, lgmaps in zip(pl_arg.data, pl_arg.lgmaps):
                         olgmaps.append(m.handle.getLGMap())
-                        m.handle.setLGMap(*lgmaps)
+                        if m.handle.type != "is":
+                            m.handle.setLGMap(*lgmaps)
                     orig_lgmaps.append(olgmaps)
         return tuple(orig_lgmaps)
 
@@ -309,7 +310,8 @@ def restore_lgmaps(self, orig_lgmaps):
         for arg, d in reversed(list(zip(self.global_kernel.arguments, self.arguments))):
             if isinstance(arg, (MatKernelArg, MixedMatKernelArg)) and d.lgmaps is not None:
                 for m, lgmaps in zip(d.data, orig_lgmaps.pop()):
-                    m.handle.setLGMap(*lgmaps)
+                    if m.handle.type != "is":
+                        m.handle.setLGMap(*lgmaps)
 
     @cached_property
     def _has_mats(self):

pyop2/types/mat.py

@@ -825,9 +825,13 @@ def set_local_diagonal_entries(self, rows, diag_val=1.0, idx=None):
         rows = rows.reshape(-1, 1)
         self.change_assembly_state(Mat.INSERT_VALUES)
         if len(rows) > 0:
-            values = np.full(rows.shape, diag_val, dtype=dtypes.ScalarType)
-            self.handle.setValuesLocalRCV(rows, rows, values,
-                                          addv=PETSc.InsertMode.INSERT_VALUES)
+            if self.handle.type == "is":
+                self.handle.assemble()
+                self.handle.zeroRowsColumnsLocal(rows, diag_val)
+            else:
+                values = np.full(rows.shape, diag_val, dtype=dtypes.ScalarType)
+                self.handle.setValuesLocalRCV(rows, rows, values,
+                                              addv=PETSc.InsertMode.INSERT_VALUES)
 
     @mpi.collective
     def assemble(self):

tests/firedrake/regression/test_assemble.py

@@ -142,6 +142,25 @@ def test_mat_nest_real_block_assembler_correctly_reuses_tensor(mesh):
 
     assert A2.M is A1.M
 
+@pytest.mark.parametrize("dirichlet_bcs", [False, True])
+def test_assemble_matis(mesh, dirichlet_bcs):
+    V = FunctionSpace(mesh, "CG", 1)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    a = inner(grad(u), grad(v))*dx
+    if dirichlet_bcs:
+        bcs = DirichletBC(V, 0, (1, 3))
+    else:
+        bcs = None
+
+    ais = assemble(a, bcs=bcs, mat_type="is").petscmat
+    aijnew = PETSc.Mat()
+    ais.convert("aij", aijnew)
+
+    aij = assemble(a, bcs=bcs, mat_type="aij").petscmat
+    aij.axpy(-1, aijnew)
+    aij.view()
+    assert np.allclose(aij[:, :], 0)
 
 def test_assemble_diagonal(mesh):
     V = FunctionSpace(mesh, "P", 3)