Change interface of fem.petsc.create_vector to allow for general list of function spaces

python/dolfinx/fem/petsc.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2018-2025 Garth N. Wells and Jørgen S. Dokken
+# Copyright (C) 2018-2025 Garth N. Wells, Jørgen S. Dokken and Paul T. Kühner
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -23,7 +23,6 @@
 
 import contextlib
 import functools
-import itertools
 import typing
 from collections.abc import Iterable, Sequence
 
@@ -48,7 +47,7 @@
 from dolfinx.fem.assemble import apply_lifting as _apply_lifting
 from dolfinx.fem.bcs import DirichletBC
 from dolfinx.fem.bcs import bcs_by_block as _bcs_by_block
-from dolfinx.fem.forms import Form
+from dolfinx.fem.forms import Form, extract_function_spaces
 from dolfinx.fem.forms import extract_function_spaces as _extract_spaces
 from dolfinx.fem.forms import form as _create_form
 from dolfinx.fem.function import Function as _Function
@@ -106,17 +105,19 @@ def fn(form):
 
 
 def create_vector(
-    L: typing.Union[Form, Iterable[Form]], kind: typing.Optional[str] = None
+    container: typing.Union[Form, _FunctionSpace, Iterable[Form], Iterable[_FunctionSpace]],
+    kind: typing.Optional[str] = None,
 ) -> PETSc.Vec:
-    """Create a PETSc vector that is compatible with a linear form(s).
+    """Create a PETSc vector that is compatible with a linear form(s)
+    or functionspace(s).
 
     Three cases are supported:
 
-    1. For a single linear form ``L``, if ``kind`` is ``None`` or is
+    1. For a single linear form ``L`` or space ``V``, if ``kind`` is ``None`` or is
        ``PETSc.Vec.Type.MPI``, a ghosted PETSc vector which is
-       compatible with ``L`` is created.
+       compatible with ``L/V`` is created.
 
-    2. If ``L`` is a sequence of linear forms and ``kind`` is ``None``
+    2. If ``L/V`` is a sequence of linear forms/functionspaces and ``kind`` is ``None``
        or is ``PETSc.Vec.Type.MPI``, a ghosted PETSc vector which is
        compatible with ``L`` is created. The created vector ``b`` is
        initialized such that on each MPI process ``b = [b_0, b_1, ...,
@@ -126,7 +127,7 @@ def create_vector(
 
        For this case, the returned vector has an attribute ``_blocks``
        that holds the local offsets into ``b`` for the (i) owned and
-       (ii) ghost entries for each ``L[i]``. It can be accessed by
+       (ii) ghost entries for each ``L_i/V_i``. It can be accessed by
        ``b.getAttr("_blocks")``. The offsets can be used to get views
        into ``b`` for blocks, e.g.::
 
@@ -140,50 +141,27 @@ def create_vector(
            >>> b1_owned = b.array[offsets0[1]:offsets0[2]]
            >>> b1_ghost = b.array[offsets1[1]:offsets1[2]]
 
-    3. If ``L`` is a sequence of linear forms and ``kind`` is
+    3. If ``L/V`` is a sequence of linear forms/functionspaces and ``kind`` is
        ``PETSc.Vec.Type.NEST``, a PETSc nested vector (a 'nest' of
-       ghosted PETSc vectors) which is compatible with ``L`` is created.
+       ghosted PETSc vectors) which is compatible with ``L/V`` is created.
 
     Args:
-        L: Linear form or a sequence of linear forms.
+        : Linear form/function space or a sequence of such.
         kind: PETSc vector type (``VecType``) to create.
 
     Returns:
         A PETSc vector with a layout that is compatible with ``L``. The
         vector is not initialised to zero.
     """
-    try:
-        dofmap = L.function_spaces[0].dofmaps(0)  # Single form case
-        return dolfinx.la.petsc.create_vector(dofmap.index_map, dofmap.index_map_bs)
-    except AttributeError:
-        maps = [
-            (
-                form.function_spaces[0].dofmaps(0).index_map,
-                form.function_spaces[0].dofmaps(0).index_map_bs,
-            )
-            for form in L
-        ]
-        if kind == PETSc.Vec.Type.NEST:
-            return _cpp.fem.petsc.create_vector_nest(maps)
-        elif kind == PETSc.Vec.Type.MPI:
-            off_owned = tuple(
-                itertools.accumulate(maps, lambda off, m: off + m[0].size_local * m[1], initial=0)
-            )
-            off_ghost = tuple(
-                itertools.accumulate(
-                    maps, lambda off, m: off + m[0].num_ghosts * m[1], initial=off_owned[-1]
-                )
-            )
+    if isinstance(container, (Form, _FunctionSpace)):
+        container = [container]
 
-            b = _cpp.fem.petsc.create_vector_block(maps)
-            b.setAttr("_blocks", (off_owned, off_ghost))
-            return b
-        else:
-            raise NotImplementedError(
-                "Vector type must be specified for blocked/nested assembly."
-                f"Vector type '{kind}' not supported."
-                "Did you mean 'nest' or 'mpi'?"
-            )
+    if len(container) == 0:
+        raise RuntimeError("No container provided.")
+
+    V = extract_function_spaces(container) if isinstance(container[0], Form) else container
+    maps = [(_V.dofmap.index_map, _V.dofmap.index_map_bs) for _V in V]
+    return dolfinx.la.petsc.create_vector(maps, kind=kind)
 
 
 # -- Matrix instantiation ----------------------------------------------------
@@ -299,6 +277,7 @@ def assemble_vector(
         An assembled vector.
     """
     b = create_vector(L, kind=kind)
+
     if kind == PETSc.Vec.Type.NEST:
         for b_sub in b.getNestSubVecs():
             with b_sub.localForm() as b_local:

python/dolfinx/la/petsc.py

@@ -14,6 +14,7 @@
 """
 
 import functools
+import itertools
 import typing
 
 from petsc4py import PETSc
@@ -22,29 +23,14 @@
 import numpy.typing as npt
 
 import dolfinx
-from dolfinx.la import IndexMap, Vector
+from dolfinx.common import IndexMap
+from dolfinx.la import Vector
 
 assert dolfinx.has_petsc4py
 
 __all__ = ["assign", "create_vector", "create_vector_wrap"]
 
 
-def create_vector(index_map: IndexMap, bs: int) -> PETSc.Vec:  # type: ignore[name-defined]
-    """Create a distributed PETSc vector.
-
-    Args:
-        index_map: Index map that describes the size and parallel layout of
-            the vector to create.
-        bs: Block size of the vector.
-
-    Returns:
-        PETSc Vec object.
-    """
-    ghosts = index_map.ghosts.astype(PETSc.IntType)  # type: ignore[attr-defined]
-    size = (index_map.size_local * bs, index_map.size_global * bs)
-    return PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=index_map.comm)  # type: ignore
-
-
 def create_vector_wrap(x: Vector) -> PETSc.Vec:  # type: ignore[name-defined]
     """Wrap a distributed DOLFINx vector as a PETSc vector.
 
@@ -63,6 +49,79 @@ def create_vector_wrap(x: Vector) -> PETSc.Vec:  # type: ignore[name-defined]
     )
 
 
+def create_vector(
+    maps: typing.Sequence[tuple[IndexMap, int]], kind: typing.Optional[str] = None
+) -> PETSc.Vec:
+    """Create a PETSc vector from a sequence of maps and blocksizes.
+
+    Three cases are supported:
+
+    1. If ``maps=[(im_0, bs_0), ..., (im_n, bs_n)]`` is a sequence of indexmaps and blocksizes and
+       ``kind`` is ``None``or is ``PETSc.Vec.Type.MPI``, a ghosted PETSc vector whith block
+       structure described by ``(im_i, bs_i)`` is created.
+       The created vector ``b`` is initialized such that on each MPI process ``b = [b_0, b_1, ...,
+       b_n, b_0g, b_1g, ..., b_ng]``, where ``b_i`` are the entries associated with the 'owned'
+       degrees-of-freedom for ``(im_i, bs_i)`` and ``b_ig`` are the 'unowned' (ghost) entries.
+
+       If more than one tuple is supplied, the returned vector has an attribute ``_blocks``
+       that holds the local offsets into ``b`` for the (i) owned and
+       (ii) ghost entries for each ``V[i]``. It can be accessed by
+       ``b.getAttr("_blocks")``. The offsets can be used to get views
+       into ``b`` for blocks, e.g.::
+
+           >>> offsets0, offsets1, = b.getAttr("_blocks")
+           >>> offsets0
+           (0, 12, 28)
+           >>> offsets1
+           (28, 32, 35)
+           >>> b0_owned = b.array[offsets0[0]:offsets0[1]]
+           >>> b0_ghost = b.array[offsets1[0]:offsets1[1]]
+           >>> b1_owned = b.array[offsets0[1]:offsets0[2]]
+           >>> b1_ghost = b.array[offsets1[1]:offsets1[2]]
+
+    2. If ``V=[(im_0, bs_0), ..., (im_n, bs_n)]`` is a sequence of function space and ``kind`` is
+       ``PETSc.Vec.Type.NEST``, a PETSc nested vector (a 'nest' of ghosted PETSc vectors) is
+       created.
+
+    Args:
+        map: Sequence of tuples of ``IndexMap`` and the associated block size.
+        kind: PETSc vector type (``VecType``) to create.
+
+    Returns:
+        A PETSc vector with the prescribed layout. The vector is not initialised to zero.
+    """
+    if len(maps) == 1:
+        # Single space case
+        index_map, bs = maps[0]
+        ghosts = index_map.ghosts.astype(PETSc.IntType)  # type: ignore[attr-defined]
+        size = (index_map.size_local * bs, index_map.size_global * bs)
+        return PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=index_map.comm)  # type: ignore
+
+    if kind is None or kind == PETSc.Vec.Type.MPI:
+        off_owned = tuple(
+            itertools.accumulate(maps, lambda off, m: off + m[0].size_local * m[1], initial=0)
+        )
+        off_ghost = tuple(
+            itertools.accumulate(
+                maps, lambda off, m: off + m[0].num_ghosts * m[1], initial=off_owned[-1]
+            )
+        )
+
+        b = dolfinx.cpp.fem.petsc.create_vector_block(maps)
+        b.setAttr("_blocks", (off_owned, off_ghost))
+        return b
+
+    elif kind == PETSc.Vec.Type.NEST:
+        return dolfinx.cpp.fem.petsc.create_vector_nest(maps)
+
+    else:
+        raise NotImplementedError(
+            "Vector type must be specified for blocked/nested assembly."
+            f"Vector type '{kind}' not supported."
+            "Did you mean 'nest' or 'mpi'?"
+        )
+
+
 @functools.singledispatch
 def assign(x0: typing.Union[npt.NDArray[np.inexact], list[npt.NDArray[np.inexact]]], x1: PETSc.Vec):  # type: ignore
     """Assign ``x0`` values to a PETSc vector ``x1``.

python/test/unit/nls/test_newton.py

@@ -200,8 +200,7 @@ def test_nonlinear_pde_snes(self):
         """Test Newton solver for a simple nonlinear PDE"""
         from petsc4py import PETSc
 
-        from dolfinx.fem.petsc import create_matrix
-        from dolfinx.la.petsc import create_vector
+        from dolfinx.fem.petsc import create_matrix, create_vector
 
         mesh = create_unit_square(MPI.COMM_WORLD, 12, 15)
         V = functionspace(mesh, ("Lagrange", 1))
@@ -220,7 +219,7 @@ def test_nonlinear_pde_snes(self):
         problem = NonlinearPDE_SNESProblem(F, u, bc)
 
         u.x.array[:] = 0.9
-        b = create_vector(V.dofmap.index_map, V.dofmap.index_map_bs)
+        b = create_vector(V)
         J = create_matrix(problem.a)
 
         # Create Newton solver and solve