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

python/demo/demo_stokes.py

@@ -238,7 +238,7 @@ def nested_iterative_solver_high_level():
     # 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 = create_vector(L, "nest")
+    null_vec = create_vector(extract_function_spaces(L), "nest")
 
     # Set velocity part to zero and the pressure part to a non-zero
     # constant
@@ -322,7 +322,7 @@ def nested_iterative_solver_low_level():
 
     # Set the nullspace for pressure (since pressure is determined only
     # up to a constant)
-    null_vec = create_vector(L, "nest")
+    null_vec = create_vector(extract_function_spaces(L), "nest")
     null_vecs = null_vec.getNestSubVecs()
     null_vecs[0].set(0.0), null_vecs[1].set(1.0)
     null_vec.normalize()

python/dolfinx/fem/assemble.py

@@ -17,12 +17,13 @@
 
 import dolfinx
 from dolfinx import cpp as _cpp
-from dolfinx import la
+from dolfinx import default_scalar_type, la
 from dolfinx.cpp.fem import pack_coefficients as _pack_coefficients
 from dolfinx.cpp.fem import pack_constants as _pack_constants
 from dolfinx.fem import IntegralType
 from dolfinx.fem.bcs import DirichletBC
 from dolfinx.fem.forms import Form
+from dolfinx.fem.function import FunctionSpace
 
 
 def pack_constants(
@@ -94,19 +95,19 @@ def _pack(form):
 # -- Vector and matrix instantiation --------------------------------------
 
 
-def create_vector(L: Form) -> la.Vector:
-    """Create a Vector that is compatible with a given linear form.
+def create_vector(V: FunctionSpace, dtype: npt.DTypeLike = default_scalar_type) -> la.Vector:
+    """Create a Vector that is compatible with the given function space.
 
     Args:
-        L: A linear form.
+        V: A function space.
 
     Returns:
-        A vector that the form can be assembled into.
+        A vector compatible with the function space.
     """
     # Can just take the first dofmap here, since all dof maps have the same
     # index map in mixed-topology meshes
-    dofmap = L.function_spaces[0].dofmaps(0)  # type: ignore[attr-defined]
-    return la.vector(dofmap.index_map, dofmap.index_map_bs, dtype=L.dtype)
+    dofmap = V.dofmaps(0)  # type: ignore[attr-defined]
+    return la.vector(dofmap.index_map, dofmap.index_map_bs, dtype=dtype)
 
 
 def create_matrix(a: Form, block_mode: la.BlockMode | None = None) -> la.MatrixCSR:
@@ -204,7 +205,7 @@ def _assemble_vector_form(
         :func:`dolfinx.la.Vector.scatter_reverse` on the return vector
         can accumulate ghost contributions.
     """
-    b = create_vector(L)
+    b = create_vector(L.function_spaces[0], L.dtype)
     b.array[:] = 0
     constants = pack_constants(L) if constants is None else constants  # type: ignore[assignment]
     coeffs = pack_coefficients(L) if coeffs is None else coeffs  # type: ignore[assignment]

python/dolfinx/fem/forms.py

@@ -443,7 +443,7 @@ def _create_form(form):
 def extract_function_spaces(
     forms: Form | Sequence[Form] | Sequence[Sequence[Form]],
     index: int = 0,
-) -> list[None | FunctionSpace]:
+) -> FunctionSpace | list[None | FunctionSpace]:
     """Extract common function spaces from an array of forms.
 
     If ``forms`` is a list of linear forms, this function returns of list
@@ -463,7 +463,8 @@ def extract_function_spaces(
     """
     _forms = np.array(forms)
     if _forms.ndim == 0:
-        raise RuntimeError("Expected an array for forms, not a single form")
+        form: Form = _forms.tolist()
+        return form.function_spaces[0] if form is not None else None
     elif _forms.ndim == 1:
         assert index == 0, "Expected index=0 for 1D array of forms"
         for form in _forms:

python/dolfinx/fem/function.py

@@ -19,7 +19,7 @@
 import ufl
 from dolfinx import cpp as _cpp
 from dolfinx import default_scalar_type, jit, la
-from dolfinx.fem import dofmap
+from dolfinx.fem.dofmap import DofMap
 from dolfinx.fem.element import FiniteElement, finiteelement
 from dolfinx.geometry import PointOwnershipData
 
@@ -747,9 +747,12 @@ def element(self) -> FiniteElement:
         return FiniteElement(self._cpp_object.element)
 
     @property
-    def dofmap(self) -> dofmap.DofMap:
+    def dofmap(self) -> DofMap:
         """Degree-of-freedom map associated with the function space."""
-        return dofmap.DofMap(self._cpp_object.dofmap)  # type: ignore
+        return DofMap(self._cpp_object.dofmap)
+
+    def dofmaps(self, idx: int) -> DofMap:
+        return DofMap(self._cpp_object.dofmaps(idx))
 
     @property
     def mesh(self) -> Mesh:

python/dolfinx/fem/petsc.py

@@ -25,7 +25,6 @@
 
 import contextlib
 import functools
-import itertools
 from collections.abc import Sequence
 
 from petsc4py import PETSc
@@ -81,26 +80,30 @@
 # -- Vector instantiation -------------------------------------------------
 
 
-def create_vector(L: Form | Sequence[Form], kind: str | None = None) -> PETSc.Vec:  # type: ignore[name-defined]
-    """Create a PETSc vector that is compatible with a linear form(s).
-
+def create_vector(
+    V: _FunctionSpace | Sequence[_FunctionSpace | None],
+    /,
+    kind: str | None = None,
+) -> PETSc.Vec:  # type: ignore[name-defined]
+    """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 space ``V``, if ``kind`` is ``None`` or is
+       ``PETSc.Vec.Type.MPI``, a ghosted PETSc vector which is compatible
+       with ``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 ``V`` is a sequence of functionspaces and ``kind`` is ``None`` or
+       is ``PETSc.Vec.Type.MPI``, a ghosted PETSc vector which is
+       compatible with ``V`` 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]``.
+       associated with the 'owned' degrees-of-freedom for ``V[i]`` and
+       ``b_ig`` are the 'unowned' (ghost) entries for ``V[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 ``V_i``. It can be accessed by
        ``b.getAttr("_blocks")``. The offsets can be used to get views
        into ``b`` for blocks, e.g.::
 
@@ -114,50 +117,27 @@ def create_vector(L: Form | Sequence[Form], kind: str | None = None) -> PETSc.Ve
            >>> 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.
+        V: 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
+        A PETSc vector with a layout that is compatible with ``V``. The
         vector is not initialised to zero.
     """
-    if isinstance(L, Sequence):
-        maps = [
-            (
-                form.function_spaces[0].dofmaps(0).index_map,  # type: ignore[attr-defined]
-                form.function_spaces[0].dofmaps(0).index_map_bs,  # type: ignore[attr-defined]
-            )
-            for form in L
-        ]
-        if kind == PETSc.Vec.Type.NEST:  # type: ignore[attr-defined]
-            return _cpp.fem.petsc.create_vector_nest(maps)
-        elif kind == PETSc.Vec.Type.MPI or kind is None:  # type: ignore[attr-defined]
-            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(
+        V, _FunctionSpace | _cpp.fem.FunctionSpace_float32 | _cpp.fem.FunctionSpace_float64
+    ):
+        V = [V]
+    elif any(_V is None for _V in V):
+        raise RuntimeError("Can not create vector for None block.")
 
-            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'?"
-            )
-    else:
-        dofmap = L.function_spaces[0].dofmaps(0)  # type: ignore[attr-defined]
-        return dolfinx.la.petsc.create_vector(dofmap.index_map, dofmap.index_map_bs)
+    maps = [(_V.dofmap.index_map, _V.dofmap.index_map_bs) for _V in V]  # type: ignore
+    return dolfinx.la.petsc.create_vector(maps, kind=kind)
 
 
 # -- Matrix instantiation -------------------------------------------------
@@ -273,7 +253,7 @@ def assemble_vector(
     Returns:
         An assembled vector.
     """
-    b = create_vector(L, kind=kind)
+    b = create_vector(_extract_function_spaces(L), kind=kind)  # type: ignore
     dolfinx.la.petsc._zero_vector(b)
     return assemble_vector(b, L, constants, coeffs)  # type: ignore[arg-type]
 
@@ -838,8 +818,8 @@ def __init__(
         # For nest matrices kind can be a nested list.
         kind = "nest" if self.A.getType() == PETSc.Mat.Type.NEST else kind  # type: ignore[attr-defined]
         assert kind is None or isinstance(kind, str)
-        self._b = create_vector(self.L, kind=kind)
-        self._x = create_vector(self.L, kind=kind)
+        self._b = create_vector(_extract_function_spaces(self.L), kind=kind)  # type: ignore
+        self._x = create_vector(_extract_function_spaces(self.L), kind=kind)  # type: ignore
 
         self._u: _Function | Sequence[_Function]
         if u is None:
@@ -1031,26 +1011,16 @@ def _assign_block_data(forms: Sequence[Form], vec: PETSc.Vec):  # type: ignore[n
         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,  # type: ignore[attr-defined]
             form.function_spaces[0].dofmaps(0).index_map_bs,  # type: ignore[attr-defined]
         )
         for form in forms
-    ]
-    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))
+
+    return dolfinx.la.petsc._assign_block_data(maps, vec)
 
 
 def assemble_residual(
@@ -1315,8 +1285,8 @@ def __init__(
         # Determine the vector kind based on the matrix type
         kind = "nest" if self._A.getType() == PETSc.Mat.Type.NEST else kind  # type: ignore[attr-defined]
         assert kind is None or isinstance(kind, str)
-        self._b = create_vector(self.F, kind=kind)
-        self._x = create_vector(self.F, kind=kind)
+        self._b = create_vector(_extract_function_spaces(self.F), kind=kind)  # type: ignore
+        self._x = create_vector(_extract_function_spaces(self.F), kind=kind)  # type: ignore
 
         # Create the SNES solver and attach the corresponding Jacobian and
         # residual computation functions