From d7b08cc6a80679471164ade309d6afa1214d90be Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Mon, 8 Jan 2024 23:44:04 +0000 Subject: [PATCH] mat + sparsity: use dset.layout_vec.local_size for sizes --- pyop2/sparsity.pyx | 8 ++++---- pyop2/types/dataset.py | 25 ++----------------------- pyop2/types/mat.py | 27 ++++++--------------------- test/unit/test_matrices.py | 4 ++-- 4 files changed, 14 insertions(+), 50 deletions(-) diff --git a/pyop2/sparsity.pyx b/pyop2/sparsity.pyx index 23635fda1..d88bab456 100644 --- a/pyop2/sparsity.pyx +++ b/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)), diff --git a/pyop2/types/dataset.py b/pyop2/types/dataset.py index 4e114032a..e554bbcef 100644 --- a/pyop2/types/dataset.py +++ b/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. diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index aefd77de1..a5ad65f71 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -66,8 +66,6 @@ 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: @@ -75,9 +73,6 @@ def __init__(self, dsets, maps, *, iteration_regions=None, name=None, nest=None, 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): diff --git a/test/unit/test_matrices.py b/test/unit/test_matrices.py index c407df41e..6ae7fb284 100644 --- a/test/unit/test_matrices.py +++ b/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()