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

python/dolfinx/fem/petsc.py

@@ -23,7 +23,6 @@
 
 import contextlib
 import functools
-import itertools
 import typing
 from collections.abc import Iterable, Sequence
 
@@ -51,7 +50,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, derivative_block
+from dolfinx.fem.forms import Form, derivative_block, 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
@@ -112,27 +111,30 @@ 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
-       ``PETSc.Vec.Type.MPI``, a ghosted PETSc vector which is
-       compatible with ``L`` is created.
+    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/V`` is created.
 
-    2. If ``L`` is a sequence of linear forms 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, ...,
+    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, ...,
        b_n, b_0g, b_1g, ..., b_ng]``, where ``b_i`` are the entries
        associated with the 'owned' degrees-of-freedom for ``L[i]`` and
        ``b_ig`` are the 'unowned' (ghost) entries for ``L[i]``.
 
        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.::
 
@@ -146,50 +148,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
-       ``PETSc.Vec.Type.NEST``, a PETSc nested vector (a 'nest' of
-       ghosted PETSc vectors) which is compatible with ``L`` is created.
+    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/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("Empty sequence of functionspaces/forms 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 -------------------------------------------------
@@ -1056,26 +1035,16 @@ def _assign_block_data(forms: typing.Iterable[dolfinx.fem.Form], vec: PETSc.Vec)
         forms: List of forms to extract block data from.
         vec: PETSc vector to assign block data to.
     """
-    # Early exit if the vector already has block data or is a nest vector
-    if vec.getAttr("_blocks") is not None or vec.getType() == "nest":
-        return
 
-    maps = [
+    maps = (
         (
             form.function_spaces[0].dofmaps(0).index_map,
             form.function_spaces[0].dofmaps(0).index_map_bs,
         )
-        for form in forms  # type: ignore
-    ]
-    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]
-        )
+        for form in forms
     )
-    vec.setAttr("_blocks", (off_owned, off_ghost))
+
+    return dolfinx.la.petsc._assign_block_data(maps, vec)
 
 
 def assemble_residual(

python/dolfinx/la/petsc.py

@@ -14,15 +14,18 @@
 """
 
 import functools
+import itertools
 import typing
+from collections.abc import Iterable
 
 from petsc4py import PETSc
 
 import numpy as np
 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
 
@@ -51,22 +54,6 @@ def _zero_vector(x: PETSc.Vec):  # type: ignore
             x_local.set(0.0)
 
 
-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.
 
@@ -85,6 +72,75 @@ 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:
+        b = dolfinx.cpp.fem.petsc.create_vector_block(maps)
+        _assign_block_data(maps, b)
+        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``.
@@ -157,3 +213,27 @@ def _(x0: PETSc.Vec, x1: typing.Union[npt.NDArray[np.inexact], list[npt.NDArray[
                     start = end
             except IndexError:
                 x1[:] = _x0.array_r[:]
+
+
+def _assign_block_data(maps: Iterable[tuple[IndexMap, int]], vec: PETSc.Vec):
+    """Assign block data to a PETSc vector.
+
+    Args:
+        maps: Iterable of tuples each containing an ``IndexMap`` and the
+        associated block size ``bs``.
+        vec: PETSc vector to assign block data to.
+    """
+    # Early exit if the vector already has block data or is a nest vector
+    if vec.getAttr("_blocks") is not None or vec.getType() == "nest":
+        return
+
+    maps = [(index_map, bs) for index_map, bs in maps]
+    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]
+        )
+    )
+    vec.setAttr("_blocks", (off_owned, off_ghost))

python/test/unit/fem/test_assembler.py

@@ -1221,6 +1221,7 @@ def test_lifting_coefficients(self, kind):
         # Apply lifting with input coefficient
         coeffs = pack_coefficients(J)
         b = create_vector(L, kind=kind)
+        assert b.equal(create_vector([V, Q]))
         apply_lifting(b, J, bcs=bcs1, coeffs=coeffs)
         b.assemble()
 

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