mat + sparsity: use dset.layout_vec.local_size for sizes (#714)

pyop2/sparsity.pyx

@@ -128,13 +128,13 @@ def build_sparsity(sparsity):
     preallocator.setType(PETSc.Mat.Type.PREALLOCATOR)
     if mixed:
         # Sparsity is the dof sparsity.
-        nrows = sum(s.size*s.cdim for s in rset)
-        ncols = sum(s.size*s.cdim for s in cset)
+        nrows = rset.layout_vec.local_size
+        ncols = cset.layout_vec.local_size
         preallocator.setLGMap(rmap=rset.unblocked_lgmap, cmap=cset.unblocked_lgmap)
     else:
         # Sparsity is the block sparsity
-        nrows = rset.size
-        ncols = cset.size
+        nrows = rset.layout_vec.local_size // rset.layout_vec.block_size
+        ncols = cset.layout_vec.local_size // cset.layout_vec.block_size
         preallocator.setLGMap(rmap=rset.scalar_lgmap, cmap=cset.scalar_lgmap)
 
     preallocator.setSizes(size=((nrows, None), (ncols, None)),

pyop2/types/dataset.py

@@ -156,11 +156,11 @@ def field_ises(self):
         ises = []
         nlocal_rows = 0
         for dset in self:
-            nlocal_rows += dset.size * dset.cdim
+            nlocal_rows += dset.layout_vec.local_size
         offset = self.comm.scan(nlocal_rows)
         offset -= nlocal_rows
         for dset in self:
-            nrows = dset.size * dset.cdim
+            nrows = dset.layout_vec.local_size
             iset = PETSc.IS().createStride(nrows, first=offset, step=1,
                                            comm=self.comm)
             iset.setBlockSize(dset.cdim)
@@ -284,27 +284,6 @@ def unblocked_lgmap(self):
                                          bsize=1, comm=self.lgmap.comm)
             return lgmap
 
-    @utils.cached_property
-    def field_ises(self):
-        """A list of PETSc ISes defining the global indices for each set in
-        the DataSet.
-
-        Used when extracting blocks from matrices for solvers."""
-        ises = []
-        nlocal_rows = 0
-        for dset in self:
-            nlocal_rows += dset.size * dset.cdim
-        offset = self.comm.scan(nlocal_rows)
-        offset -= nlocal_rows
-        for dset in self:
-            nrows = dset.size * dset.cdim
-            iset = PETSc.IS().createStride(nrows, first=offset, step=1,
-                                           comm=self.comm)
-            iset.setBlockSize(dset.cdim)
-            ises.append(iset)
-            offset += nrows
-        return tuple(ises)
-
     @utils.cached_property
     def local_ises(self):
         """A list of PETSc ISes defining the local indices for each set in the DataSet.

pyop2/types/mat.py

@@ -66,18 +66,13 @@ def __init__(self, dsets, maps, *, iteration_regions=None, name=None, nest=None,
             self._dims = (((1, 1),),)
             self._d_nnz = None
             self._o_nnz = None
-            self._nrows = None if isinstance(dsets[0], GlobalDataSet) else self._rmaps[0].toset.size
-            self._ncols = None if isinstance(dsets[1], GlobalDataSet) else self._cmaps[0].toset.size
             self.lcomm = mpi.internal_comm(dsets[0].comm if isinstance(dsets[0], GlobalDataSet) else self._rmaps[0].comm)
             self.rcomm = mpi.internal_comm(dsets[1].comm if isinstance(dsets[1], GlobalDataSet) else self._cmaps[0].comm)
         else:
             self.lcomm = mpi.internal_comm(self._rmaps[0].comm)
             self.rcomm = mpi.internal_comm(self._cmaps[0].comm)
 
             rset, cset = self.dsets
-            # All rmaps and cmaps have the same data set - just use the first.
-            self._nrows = rset.size
-            self._ncols = cset.size
 
             self._has_diagonal = (rset == cset)
 
@@ -277,16 +272,6 @@ def shape(self):
         return (len(self._dsets[0] or [1]),
                 len(self._dsets[1] or [1]))
 
-    @utils.cached_property
-    def nrows(self):
-        """The number of rows in the ``Sparsity``."""
-        return self._nrows
-
-    @utils.cached_property
-    def ncols(self):
-        """The number of columns in the ``Sparsity``."""
-        return self._ncols
-
     @utils.cached_property
     def nested(self):
         r"""Whether a sparsity is monolithic (even if it has a block structure).
@@ -376,8 +361,6 @@ def __init__(self, parent, i, j):
         self._dsets = (parent.dsets[0][i], parent.dsets[1][j])
         self._rmaps = tuple(m.split[i] for m in parent.rmaps)
         self._cmaps = tuple(m.split[j] for m in parent.cmaps)
-        self._nrows = self._dsets[0].size
-        self._ncols = self._dsets[1].size
         self._has_diagonal = i == j and parent._has_diagonal
         self._parent = parent
         self._dims = tuple([tuple([parent.dims[i][j]])])
@@ -520,7 +503,7 @@ def dims(self):
     @utils.cached_property
     def nrows(self):
         "The number of rows in the matrix (local to this process)"
-        return sum(d.size * d.cdim for d in self.sparsity.dsets[0])
+        return self.sparsity.dsets[0].layout_vec.local_size
 
     @utils.cached_property
     def nblock_rows(self):
@@ -530,7 +513,8 @@ def nblock_rows(self):
         by the dimension of the row :class:`DataSet`.
         """
         assert len(self.sparsity.dsets[0]) == 1, "Block rows don't make sense for mixed Mats"
-        return self.sparsity.dsets[0].size
+        layout_vec = self.sparsity.dsets[0].layout_vec
+        return layout_vec.local_size // layout_vec.block_size
 
     @utils.cached_property
     def nblock_cols(self):
@@ -540,12 +524,13 @@ def nblock_cols(self):
         divided by the dimension of the column :class:`DataSet`.
         """
         assert len(self.sparsity.dsets[1]) == 1, "Block cols don't make sense for mixed Mats"
-        return self.sparsity.dsets[1].size
+        layout_vec = self.sparsity.dsets[1].layout_vec
+        return layout_vec.local_size // layout_vec.block_size
 
     @utils.cached_property
     def ncols(self):
         "The number of columns in the matrix (local to this process)"
-        return sum(d.size * d.cdim for d in self.sparsity.dsets[1])
+        return self.sparsity.dsets[1].layout_vec.local_size
 
     @utils.cached_property
     def sparsity(self):

test/unit/test_matrices.py

@@ -582,7 +582,7 @@ def test_invalid_mode(self, elements, elem_node, mat, mode):
     def test_mat_set_diagonal(self, nodes, elem_node, n):
         "Set the diagonal of the entire matrix to 1.0"
         mat = op2.Mat(op2.Sparsity(nodes**n, elem_node), valuetype)
-        nrows = mat.sparsity.nrows
+        nrows = mat.nblock_rows
         mat.set_local_diagonal_entries(list(range(nrows)))
         mat.assemble()
         assert (mat.values == np.identity(nrows * n)).all()
@@ -591,7 +591,7 @@ def test_mat_set_diagonal(self, nodes, elem_node, n):
     def test_mat_repeated_set_diagonal(self, nodes, elem_node, n):
         "Set the diagonal of the entire matrix to 1.0"
         mat = op2.Mat(op2.Sparsity(nodes**n, elem_node), valuetype)
-        nrows = mat.sparsity.nrows
+        nrows = mat.nblock_rows
         mat.set_local_diagonal_entries(list(range(nrows)))
         mat.assemble()
         assert (mat.values == np.identity(nrows * n)).all()