diff --git a/python/dolfinx/fem/petsc.py b/python/dolfinx/fem/petsc.py index 3b6fc8b33ca..801e73d49fc 100644 --- a/python/dolfinx/fem/petsc.py +++ b/python/dolfinx/fem/petsc.py @@ -27,7 +27,6 @@ import contextlib import functools -import itertools import typing from collections.abc import Iterable, Sequence @@ -55,7 +54,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 @@ -117,27 +116,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.:: @@ -151,50 +153,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 ------------------------------------------------- @@ -1064,26 +1043,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( diff --git a/python/dolfinx/la/petsc.py b/python/dolfinx/la/petsc.py index 27c37102d9d..0e358ab1fe4 100644 --- a/python/dolfinx/la/petsc.py +++ b/python/dolfinx/la/petsc.py @@ -14,7 +14,9 @@ """ import functools +import itertools import typing +from collections.abc import Iterable from petsc4py import PETSc @@ -22,7 +24,8 @@ 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)) diff --git a/python/test/unit/fem/test_assembler.py b/python/test/unit/fem/test_assembler.py index 787f6935a87..556a2e34858 100644 --- a/python/test/unit/fem/test_assembler.py +++ b/python/test/unit/fem/test_assembler.py @@ -1217,6 +1217,7 @@ def test_lifting_coefficients(self, kind): # Apply lifting with input coefficient coeffs = pack_coefficients(J) b = petsc_create_vector(L, kind=kind) + assert b.equal(petsc_create_vector([V, Q])) petsc_apply_lifting(b, J, bcs=bcs1, coeffs=coeffs) b.assemble() diff --git a/python/test/unit/nls/test_newton.py b/python/test/unit/nls/test_newton.py index 48f61f64c3e..12761b87c5d 100644 --- a/python/test/unit/nls/test_newton.py +++ b/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