Propagate sub_mat_type, fix MatNest + MatIS

Author: pbrubeck

firedrake/assemble.py

@@ -353,7 +353,9 @@ def allocate(self):
             else:
                 test, trial = self._form.arguments()
                 sparsity = ExplicitMatrixAssembler._make_sparsity(test, trial, self._mat_type, self._sub_mat_type, self.maps_and_regions)
-                return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, options_prefix=self._options_prefix)
+                return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType,
+                                     sub_mat_type=self._sub_mat_type,
+                                     options_prefix=self._options_prefix)
         else:
             raise NotImplementedError("Only implemented for rank = 2 and diagonal = False")
 
@@ -1344,6 +1346,7 @@ def allocate(self):
                                                           self._sub_mat_type,
                                                           self._make_maps_and_regions())
         return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType,
+                             sub_mat_type=self._sub_mat_type,
                              options_prefix=self._options_prefix,
                              fc_params=self._form_compiler_params)
 
@@ -1366,10 +1369,21 @@ def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions):
         except SparsityFormatError:
             raise ValueError("Monolithic matrix assembly not supported for systems "
                              "with R-space blocks")
+
+        # TODO reconstruct dof_dset with the unghosted lgmap
         if mat_type == "is":
-            rlgmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space())
-            clgmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space())
-            sparsity._lgmaps = (rlgmap, clgmap)
+            rmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space())
+            cmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space())
+            sparsity._lgmaps = (rmap, cmap)
+
+        elif mat_type == "nest" and sub_mat_type == "is":
+            Vrow = test.function_space()
+            Vcol = trial.function_space()
+            for i, j in numpy.ndindex((len(Vrow), len(Vcol))):
+                block = sparsity[i, j]
+                rmap = unghosted_lgmap(block._dsets[0].lgmap, Vrow[i])
+                cmap = unghosted_lgmap(block._dsets[1].lgmap, Vcol[j])
+                block._lgmaps = (rmap, cmap)
         return sparsity
 
     def _make_maps_and_regions(self):
@@ -1465,6 +1479,15 @@ def _apply_bc(self, tensor, bc, u=None):
             # block is on the matrix diagonal and its index matches the
             # index of the function space the bc is defined on.
             op2tensor[index, index].set_local_diagonal_entries(bc.nodes, idx=component, diag_val=self.weight)
+            if self._mat_type == "nest" and self._sub_mat_type == "is":
+                for i in range(len(tensor.arguments()[1].function_space())):
+                    if i == index:
+                        continue
+                    op2tensor[index, i].zero_rows(bc.nodes, idx=component, diag_val=0)
+                for i in range(len(tensor.arguments()[0].function_space())):
+                    if i == index:
+                        continue
+                    op2tensor[i, index].zero_columns(bc.nodes, idx=component, diag_val=0)
 
             # Handle off-diagonal block involving real function space.
             # "lgmaps" is correctly constructed in _matrix_arg, but

pyop2/types/mat.py

@@ -561,6 +561,7 @@ class Mat(AbstractMat):
 
     def __init__(self, *args, **kwargs):
         self.mat_type = kwargs.pop("mat_type", None)
+        self.sub_mat_type = kwargs.pop("sub_mat_type", None)
         super().__init__(*args, **kwargs)
         self._init()
         self.assembly_state = Mat.ASSEMBLED
@@ -620,7 +621,7 @@ def _init_monolithic(self):
         rlgmap = rset.unblocked_lgmap
         clgmap = cset.unblocked_lgmap
         if self.mat_type == "is":
-            # TODO monolithic lgmaps
+            # FIXME monolithic lgmaps
             rlgmap, clgmap = self.sparsity._lgmaps
             create = mat.createIS
         else:
@@ -671,7 +672,8 @@ def _init_nest(self):
             row = []
             for j in range(cols):
                 row.append(Mat(self.sparsity[i, j], self.dtype,
-                           '_'.join([self.name, str(i), str(j)])))
+                               '_'.join([self.name, str(i), str(j)]),
+                               mat_type=self.sub_mat_type))
             self._blocks.append(row)
         # PETSc Mat.createNest wants a flattened list of Mats
         mat.createNest([[m.handle for m in row_] for row_ in self._blocks],
@@ -794,15 +796,43 @@ def zero(self):
         self.handle.zeroEntries()
 
     @mpi.collective
-    def zero_rows(self, rows, diag_val=1.0):
+    def zero_rows(self, rows, diag_val=1.0, idx=None):
         """Zeroes the specified rows of the matrix, with the exception of the
         diagonal entry, which is set to diag_val. May be used for applying
         strong boundary conditions.
 
         :param rows: a :class:`Subset` or an iterable"""
+        rows = rows.indices if isinstance(rows, Subset) else rows
+        rows = np.asarray(rows, dtype=dtypes.IntType)
+        rbs, _ = self.dims[0][0]
+        if rbs > 1:
+            if idx is not None:
+                rows = rbs * rows + idx
+            else:
+                rows = np.dstack([rbs*rows + i for i in range(rbs)]).flatten()
         self.assemble()
+        self.handle.zeroRowsLocal(rows, diag_val)
+
+    @mpi.collective
+    def zero_columns(self, rows, diag_val=1.0, idx=None):
+        """Zeroes the specified columns of the matrix, with the exception of the
+        diagonal entry, which is set to diag_val. May be used for applying
+        strong boundary conditions.
+
+        :param rows: a :class:`Subset` or an iterable"""
         rows = rows.indices if isinstance(rows, Subset) else rows
+        rows = np.asarray(rows, dtype=dtypes.IntType)
+        rbs, _ = self.dims[0][0]
+        if rbs > 1:
+            if idx is not None:
+                rows = rbs * rows + idx
+            else:
+                rows = np.dstack([rbs*rows + i for i in range(rbs)]).flatten()
+        self.assemble()
+        # FIXME implement zeroColumnsLocal
+        self.handle.transpose()
         self.handle.zeroRowsLocal(rows, diag_val)
+        self.handle.transpose()
 
     def _flush_assembly(self):
         self.handle.assemble(assembly=PETSc.Mat.AssemblyType.FLUSH)

tests/firedrake/regression/test_assemble.py

@@ -152,27 +152,46 @@ def test_assemble_matis(mesh, shape, mat_type, dirichlet_bcs):
         V = FunctionSpace(mesh, "CG", 1)
         if shape == "mixed":
             V = V * V
-
     if V.value_size == 1:
         A = 1
     else:
         A = as_matrix([[2, -1], [-1, 2]])
 
     u = TrialFunction(V)
     v = TestFunction(V)
-    a = inner(A*grad(u), grad(v))*dx
+    a = inner(A * grad(u), grad(v))*dx
     if dirichlet_bcs:
         bcs = [DirichletBC(V.sub(i), 0, (i % 4+1, (i+2) % 4+1)) for i in range(V.value_size)]
     else:
         bcs = None
 
     ais = assemble(a, bcs=bcs, mat_type=mat_type, sub_mat_type="is").petscmat
-    aijnew = PETSc.Mat()
-    ais.convert("aij", aijnew)
 
-    aij = assemble(a, bcs=bcs, mat_type="aij").petscmat
-    aij.axpy(-1, aijnew)
-    assert np.allclose(aij[:, :], 0)
+    aij = PETSc.Mat()
+    if ais.type == "nest":
+        blocks = []
+        for i in range(len(V)):
+            row = []
+            for j in range(len(V)):
+                bis = ais.getNestSubMatrix(i, j)
+                assert bis.type == "is"
+                bij = PETSc.Mat()
+                bis.convert("aij", bij)
+                row.append(bij)
+            blocks.append(row)
+        anest = PETSc.Mat()
+        anest.createNest(blocks,
+                         isrows=V.dof_dset.field_ises,
+                         iscols=V.dof_dset.field_ises,
+                         comm=ais.comm)
+        anest.convert("aij", aij)
+    else:
+        assert ais.type == "is"
+        ais.convert("aij", aij)
+
+    aij_ref = assemble(a, bcs=bcs, mat_type="aij").petscmat
+    aij_ref.axpy(-1, aij)
+    assert np.allclose(aij_ref[:, :], 0)
 
 
 def test_assemble_diagonal(mesh):