Unify PETSc matrix and vector factory functions (#3667)

* Create singelton constructors for PETSc matrix and vectors

* Modify docs

* Fix type hint

* Simplification

* Simplification

* Small fix

* Remove default arg in low level function

* Small update

---------

Co-authored-by: Garth N. Wells <[email protected]>

Author: jorgensd

cpp/dolfinx/fem/petsc.h

@@ -16,6 +16,7 @@
 #include <functional>
 #include <map>
 #include <memory>
+#include <optional>
 #include <petscmat.h>
 #include <petscvec.h>
 #include <span>
@@ -43,7 +44,7 @@ namespace petsc
 /// object.
 template <std::floating_point T>
 Mat create_matrix(const Form<PetscScalar, T>& a,
-                  std::string type = std::string())
+                  std::optional<std::string> type = std::nullopt)
 {
   la::SparsityPattern pattern = fem::create_sparsity_pattern(a);
   pattern.finalize();
@@ -52,6 +53,7 @@ Mat create_matrix(const Form<PetscScalar, T>& a,
 
 /// @brief Initialise a monolithic matrix for an array of bilinear
 /// forms.
+///
 /// @param[in] a Rectangular array of bilinear forms. The `a(i, j)` form
 /// will correspond to the `(i, j)` block in the returned matrix
 /// @param[in] type The type of PETSc Mat. If empty the PETSc default is
@@ -62,7 +64,7 @@ Mat create_matrix(const Form<PetscScalar, T>& a,
 template <std::floating_point T>
 Mat create_matrix_block(
     const std::vector<std::vector<const Form<PetscScalar, T>*>>& a,
-    std::string type = std::string())
+    std::optional<std::string> type = std::nullopt)
 {
   // Extract and check row/column ranges
   std::array<std::vector<std::shared_ptr<const FunctionSpace<T>>>, 2> V
@@ -191,16 +193,11 @@ Mat create_matrix_block(
 template <std::floating_point T>
 Mat create_matrix_nest(
     const std::vector<std::vector<const Form<PetscScalar, T>*>>& a,
-    const std::vector<std::vector<std::string>>& types)
+    std::optional<std::vector<std::vector<std::optional<std::string>>>> types)
 {
   // Extract and check row/column ranges
   auto V = fem::common_function_spaces(extract_function_spaces(a));
 
-  std::vector<std::vector<std::string>> _types(
-      a.size(), std::vector<std::string>(a.front().size()));
-  if (!types.empty())
-    _types = types;
-
   // Loop over each form and create matrix
   int rows = a.size();
   int cols = a.front().size();
@@ -212,7 +209,10 @@ Mat create_matrix_nest(
     {
       if (const Form<PetscScalar, T>* form = a[i][j]; form)
       {
-        mats[i * cols + j] = create_matrix(*form, _types[i][j]);
+        if (types)
+          mats[i * cols + j] = create_matrix(*form, types->at(i).at(j));
+        else
+          mats[i * cols + j] = create_matrix(*form, std::nullopt);
         mesh = form->mesh();
       }
     }
@@ -238,14 +238,14 @@ Mat create_matrix_nest(
   return A;
 }
 
-/// Initialise monolithic vector. Vector is not zeroed.
+/// @brief Initialise monolithic vector. Vector is not zeroed.
 ///
 /// The caller is responsible for destroying the Mat object
 Vec create_vector_block(
     const std::vector<
         std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
 
-/// Create nested (VecNest) vector. Vector is not zeroed.
+/// @brief Create nested (VecNest) vector. Vector is not zeroed.
 Vec create_vector_nest(
     const std::vector<
         std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
@@ -283,15 +283,16 @@ void assemble_vector(
   VecGhostRestoreLocalForm(b, &b_local);
 }
 
-/// Assemble linear form into an already allocated PETSc vector. Ghost
-/// contributions are not accumulated (not sent to owner). Caller is
-/// responsible for calling VecGhostUpdateBegin/End.
+/// @brief Assemble linear form into an already allocated PETSc vector.
 ///
-/// @param[in,out] b The PETsc vector to assemble the form into. The
-/// vector must already be initialised with the correct size. The
-/// process-local contribution of the form is assembled into this
-/// vector. It is not zeroed before assembly.
-/// @param[in] L The linear form to assemble
+/// Ghost contributions are not accumulated (not sent to owner). Caller
+/// is responsible for calling `VecGhostUpdateBegin`/`End`.
+///
+/// @param[in,out] b Vector to assemble the form into. The vector must
+/// already be initialised with the correct size. The process-local
+/// contribution of the form is assembled into this vector. It is not
+/// zeroed before assembly.
+/// @param[in] L Linear form to assemble.
 template <std::floating_point T>
 void assemble_vector(Vec b, const Form<PetscScalar, T>& L)
 {
@@ -462,7 +463,8 @@ void apply_lifting(
 template <std::floating_point T>
 void set_bc(
     Vec b,
-    const std::vector<std::reference_wrapper<const DirichletBC<PetscScalar, T>>> bcs,
+    const std::vector<std::reference_wrapper<const DirichletBC<PetscScalar, T>>>
+        bcs,
     const Vec x0, PetscScalar alpha = 1)
 {
   PetscInt n = 0;

cpp/dolfinx/la/petsc.cpp

@@ -231,7 +231,7 @@ void la::petsc::scatter_local_vectors(
 }
 //-----------------------------------------------------------------------------
 Mat la::petsc::create_matrix(MPI_Comm comm, const SparsityPattern& sp,
-                             std::string type)
+                             std::optional<std::string> type)
 {
   PetscErrorCode ierr;
   Mat A;
@@ -243,8 +243,8 @@ Mat la::petsc::create_matrix(MPI_Comm comm, const SparsityPattern& sp,
   std::array maps = {sp.index_map(0), sp.index_map(1)};
   const std::array bs = {sp.block_size(0), sp.block_size(1)};
 
-  if (!type.empty())
-    MatSetType(A, type.c_str());
+  if (type)
+    MatSetType(A, type->c_str());
 
   // Get global and local dimensions
   const std::int64_t M = bs[0] * maps[0]->size_global();
@@ -561,7 +561,7 @@ Mat petsc::Operator::mat() const { return _matA; }
 //-----------------------------------------------------------------------------
 //-----------------------------------------------------------------------------
 petsc::Matrix::Matrix(MPI_Comm comm, const SparsityPattern& sp,
-                      std::string type)
+                      std::optional<std::string> type)
     : Operator(petsc::create_matrix(comm, sp, type), false)
 {
   // Do nothing

cpp/dolfinx/la/petsc.h

@@ -13,6 +13,7 @@
 #include "utils.h"
 #include <boost/lexical_cast.hpp>
 #include <functional>
+#include <optional>
 #include <petscksp.h>
 #include <petscmat.h>
 #include <petscoptions.h>
@@ -118,7 +119,7 @@ void scatter_local_vectors(
 /// Create a PETSc Mat. Caller is responsible for destroying the
 /// returned object.
 Mat create_matrix(MPI_Comm comm, const SparsityPattern& sp,
-                  std::string type = std::string());
+                  std::optional<std::string> type = std::nullopt);
 
 /// Create PETSc MatNullSpace. Caller is responsible for destruction
 /// returned object.
@@ -390,7 +391,7 @@ class Matrix : public Operator
 
   /// Create holder for a PETSc Mat object from a sparsity pattern
   Matrix(MPI_Comm comm, const SparsityPattern& sp,
-         std::string type = std::string());
+         std::optional<std::string> type = std::nullopt);
 
   /// Create holder of a PETSc Mat object/pointer. The Mat A object
   /// should already be created. If inc_ref_count is true, the reference

python/demo/demo_stokes.py

@@ -249,7 +249,7 @@ def nested_iterative_solver():
     # a vector that spans the nullspace to the solver, and any component
     # of the solution in this direction will be eliminated during the
     # solution process.
-    null_vec = fem.petsc.create_vector_nest(L)
+    null_vec = fem.petsc.create_vector(L, "nest")
 
     # Set velocity part to zero and the pressure part to a non-zero
     # constant