From e8af1289a265dd6d50958719b849d8b912e15378 Mon Sep 17 00:00:00 2001
From: Stefano Zampini <stefano.zampini@gmail.com>
Date: Mon, 10 Apr 2023 16:47:35 +0300
Subject: [PATCH] PetscWrappers::MatrixBase add column sizes accessors

---
 include/deal.II/lac/petsc_matrix_base.h | 18 ++++++++++++++
 source/lac/petsc_matrix_base.cc         | 33 +++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/include/deal.II/lac/petsc_matrix_base.h b/include/deal.II/lac/petsc_matrix_base.h
index 215f8a0c8c..148cc48944 100644
--- a/include/deal.II/lac/petsc_matrix_base.h
+++ b/include/deal.II/lac/petsc_matrix_base.h
@@ -678,6 +678,24 @@ namespace PETScWrappers
     bool
     in_local_range(const size_type index) const;
 
+    /**
+     * Return the local number of columns stored on the present MPI process.
+     *
+     * To figure out which elements exactly are stored locally, use
+     * local_domain().
+     */
+    size_type
+    local_domain_size() const;
+
+    /**
+     * Return a pair of indices indicating which columns of this matrix are
+     * stored locally. The first number is the index of the first column stored,
+     * the second the index of the one past the last one that is stored
+     * locally.
+     */
+    std::pair<size_type, size_type>
+    local_domain() const;
+
     /**
      * Return a reference to the MPI communicator object in use with this
      * matrix.
diff --git a/source/lac/petsc_matrix_base.cc b/source/lac/petsc_matrix_base.cc
index 43133ce0c3..35a5da4207 100644
--- a/source/lac/petsc_matrix_base.cc
+++ b/source/lac/petsc_matrix_base.cc
@@ -307,9 +307,9 @@ namespace PETScWrappers
   MatrixBase::size_type
   MatrixBase::local_size() const
   {
-    PetscInt n_rows, n_cols;
+    PetscInt n_rows;
 
-    const PetscErrorCode ierr = MatGetLocalSize(matrix, &n_rows, &n_cols);
+    const PetscErrorCode ierr = MatGetLocalSize(matrix, &n_rows, nullptr);
     AssertThrow(ierr == 0, ExcPETScError(ierr));
 
     return n_rows;
@@ -331,6 +331,35 @@ namespace PETScWrappers
 
 
 
+  MatrixBase::size_type
+  MatrixBase::local_domain_size() const
+  {
+    PetscInt n_cols;
+
+    const PetscErrorCode ierr = MatGetLocalSize(matrix, nullptr, &n_cols);
+    AssertThrow(ierr == 0, ExcPETScError(ierr));
+
+    return n_cols;
+  }
+
+
+
+  std::pair<MatrixBase::size_type, MatrixBase::size_type>
+  MatrixBase::local_domain() const
+  {
+    PetscInt begin, end;
+
+    const PetscErrorCode ierr =
+      MatGetOwnershipRangeColumn(static_cast<const Mat &>(matrix),
+                                 &begin,
+                                 &end);
+    AssertThrow(ierr == 0, ExcPETScError(ierr));
+
+    return std::make_pair(begin, end);
+  }
+
+
+
   std::uint64_t
   MatrixBase::n_nonzero_elements() const
   {
-- 
2.39.5