PETSc FEM python wrapper: attach field information

Author: stefanozampini

cpp/dolfinx/la/petsc.cpp

@@ -147,6 +147,52 @@ std::vector<IS> la::petsc::create_index_sets(
   return is;
 }
 //-----------------------------------------------------------------------------
+std::vector<IS> la::petsc::create_global_index_sets(
+    const std::vector<
+        std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps)
+{
+  std::vector<IS> is;
+
+  std::int64_t offset = 0;
+  std::int64_t merged_local_size = 0;
+  MPI_Comm comm = MPI_COMM_NULL;
+
+  for (auto& map : maps)
+  {
+    if (comm == MPI_COMM_NULL)
+    {
+      comm = map.first.get().comm();
+    }
+    int result;
+    MPI_Comm_compare(comm, map.first.get().comm(), &result);
+    if (result != MPI_IDENT && result != MPI_CONGRUENT)
+    {
+      throw std::runtime_error("Not supported on the different communicators.");
+    }
+    int bs = map.second;
+    std::int32_t size = map.first.get().size_local();
+    merged_local_size += size * bs;
+  }
+  if (comm == MPI_COMM_NULL)
+    return is;
+
+  int ierr = MPI_Exscan(&merged_local_size, &offset, 1, MPI_INT64_T,
+                        MPI_SUM, comm);
+  dolfinx::MPI::check_error(comm, ierr);
+
+  for (auto& map : maps)
+  {
+    int bs = map.second;
+    std::int32_t size = map.first.get().size_local();
+    IS _is;
+    ISCreateStride(map.first.get().comm(), bs * size, offset, 1, &_is);
+    is.push_back(_is);
+    offset += bs * size;
+  }
+
+  return is;
+}
+//-----------------------------------------------------------------------------
 std::vector<std::vector<PetscScalar>> la::petsc::get_local_vectors(
     const Vec x,
     const std::vector<

cpp/dolfinx/la/petsc.h

@@ -100,11 +100,25 @@ Vec create_vector_wrap(const la::Vector<V>& x)
 /// @note The caller is responsible for destruction of each IS.
 ///
 /// @param[in] maps Vector of IndexMaps and corresponding block sizes
-/// @return Vector of PETSc Index Sets, created on` PETSC_COMM_SELF`
+/// @return Vector of PETSc Index Sets, created on `PETSC_COMM_SELF`
 std::vector<IS> create_index_sets(
     const std::vector<
         std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
 
+/// @brief Compute PETSc IndexSets (IS) for a stack of index maps.
+///
+/// This function stacks the owned part of the maps and returns
+/// indices in the global space. The maps must have the same communicator.
+///
+/// @note Collective
+/// @note The caller is responsible for destruction of each IS.
+///
+/// @param[in] maps Vector of IndexMaps and corresponding block sizes
+/// @return Vector of PETSc Index Sets, created on the index map communicators
+std::vector<IS> create_global_index_sets(
+    const std::vector<
+        std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
+
 /// Copy blocks from Vec into local arrays
 std::vector<std::vector<PetscScalar>> get_local_vectors(
     const Vec x,

python/dolfinx/fem/petsc.py

@@ -875,6 +875,11 @@ def __init__(
         self._solver = PETSc.KSP().create(self.A.comm)  # type: ignore[attr-defined]
         self.solver.setOperators(self.A, self.P_mat)
 
+        # Attach problem information
+        dm = self.solver.getDM()
+        dm.setCreateFieldDecomposition(partial(create_field_decomposition, u, self.L))
+        self.solver.getPC().setDM(dm)
+
         if petsc_options_prefix == "":
             raise ValueError("PETSc options prefix cannot be empty.")
 
@@ -1346,6 +1351,11 @@ def __init__(
         )
         self.solver.setFunction(partial(assemble_residual, u, self.F, self.J, bcs), self.b)
 
+        # Attach problem information
+        dm = self.solver.getDM()
+        dm.setCreateFieldDecomposition(partial(create_field_decomposition, u, self.F))
+        self.solver.getKSP().setDM(dm)
+
         if petsc_options_prefix == "":
             raise ValueError("PETSc options prefix cannot be empty.")
 
@@ -1710,3 +1720,35 @@ def _(x: PETSc.Vec, u: typing.Union[_Function, Sequence[_Function]]):  # type: i
             dolfinx.la.petsc.assign(x, data0 + data1)
         else:
             dolfinx.la.petsc.assign(x, u.x.array)
+
+
+def create_field_decomposition(
+    u: typing.Union[Sequence[_Function], _Function],
+    form: typing.Union[Form, Sequence[Form]],
+    _dm: PETSc.DM,  # type: ignore[name-defined]
+):
+    """Return index sets of the fields and their associated names.
+
+    Args:
+        u: Function tied to the solution vector.
+        form: Form of the residual or of the right-hand side.
+            It can be a sequence of forms.
+        _dm: The DM instance.
+
+    Returns:
+        names: field names.
+        ises: list of index sets in global numbering.
+        dms: list of subDMs. This function returns `None`.
+    """
+
+    if not isinstance(form, Sequence):
+        form = [form]
+    spaces = _extract_function_spaces(form)
+    ises = _cpp.la.petsc.create_global_index_sets(
+        [(V.dofmaps(0).index_map, V.dofmaps(0).index_map_bs) for V in spaces]  # type: ignore[union-attr]
+    )
+    if isinstance(u, Sequence):
+        names = [f"{v.name + '_' if v.name != 'f' else ''}{i}" for i, v in enumerate(u)]
+    else:
+        names = [f"dolfinx_field_{i}" for i in range(len(form))]
+    return names, ises, None

python/dolfinx/wrappers/petsc.cpp

@@ -210,6 +210,31 @@ void petsc_la_module(nb::module_& m)
       },
       nb::arg("maps"));
 
+  m.def(
+      "create_global_index_sets",
+      [](const std::vector<std::pair<const dolfinx::common::IndexMap*, int>>&
+             maps)
+      {
+        using X = std::vector<std::pair<
+            std::reference_wrapper<const dolfinx::common::IndexMap>, int>>;
+        X _maps;
+        std::ranges::transform(maps, std::back_inserter(_maps),
+                               [](auto m) -> typename X::value_type
+                               { return {*m.first, m.second}; });
+        std::vector<IS> index_sets
+            = dolfinx::la::petsc::create_global_index_sets(_maps);
+
+        std::vector<nb::object> py_index_sets;
+        for (auto is : index_sets)
+        {
+          PyObject* obj = PyPetscIS_New(is);
+          PetscObjectDereference((PetscObject)is);
+          py_index_sets.push_back(nb::steal(obj));
+        }
+        return py_index_sets;
+      },
+      nb::arg("maps"));
+
   m.def(
       "scatter_local_vectors",
       [](Vec x,