Refactor LA into native and PETSc part and add assigners (#3659)

* Take directly from other branch.

* Make implementations private.

* Prototype dispatcher.

* Remove trailing comma

* Looser typing

* Polishing.

* Sketch function for assigning arrays of arrays to PETSc vectors

Untested.

* Use update in tests

* Parallel fix

* Remove ghost update

* Remove ghost update that is called.

* Add reverse assign direction

* Consistent use of assign

* More simplifications

* Add assign from Function to Vec

* More simplification

* Work on docs

* Doc fixes

* Update python/dolfinx/fem/petsc.py

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Update python/dolfinx/fem/petsc.py

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Update python/dolfinx/fem/petsc.py

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Update python/dolfinx/fem/petsc.py

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Fix.

* Tidy up typing.

A Collection has no ordering - Sequence is correct.
Use np.inexact instead of np.floating

* Make la into a directory-based module - no API changes.

* Fix some tests.

* Fix.

* Fix.

* Remove duplicate code

* Add back workflow dispatch for Windows.

* Fix.

* Update __init__.py

* Don't understand docstring issue.

* Update petsc.py

* Fix docstring ala Garth.

* Partial renaming.

* Refactor.

* Fix up.

* ruff

* Improve typing

* Fix typo

* Guard typechecked definitions, remove noqa, ruff

* Tidy up typing.

* Fix bug picked up by pyright - looks like a fluke that it works?

* Update python/dolfinx/fem/petsc.py

Co-authored-by: Michal Habera <michal.habera@gmail.com>

* Remove now obvious assign comments

---------

Co-authored-by: Garth N. Wells <gnw20@cam.ac.uk>
Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>
Co-authored-by: Michal Habera <michal.habera@gmail.com>

Author: 4 people

.github/workflows/windows.yml

@@ -12,7 +12,7 @@ on:
   # merge_group:
   #   branches:
   #     - main
-  # workflow_dispatch:
+  workflow_dispatch:
 
 jobs:
 

cpp/dolfinx/la/petsc.cpp

@@ -215,12 +215,10 @@ void la::petsc::scatter_local_vectors(
   int offset_ghost = offset_owned; // Ghost DoFs start after owned
   for (std::size_t i = 0; i < maps.size(); ++i)
   {
-    const std::int32_t size_owned
-        = maps[i].first.get().size_local() * maps[i].second;
-    const std::int32_t size_ghost
-        = maps[i].first.get().num_ghosts() * maps[i].second;
-
+    std::int32_t size_owned = maps[i].first.get().size_local() * maps[i].second;
     std::copy_n(x_b[i].begin(), size_owned, std::next(_x.begin(), offset));
+
+    std::int32_t size_ghost = maps[i].first.get().num_ghosts() * maps[i].second;
     std::copy_n(std::next(x_b[i].begin(), size_owned), size_ghost,
                 std::next(_x.begin(), offset_ghost));
 

python/demo/demo_mixed-poisson.py

@@ -112,8 +112,9 @@
 import dolfinx.fem.petsc
 import ufl
 from basix.ufl import element
-from dolfinx import fem, la, mesh
+from dolfinx import fem, mesh
 from dolfinx.fem.petsc import discrete_gradient, interpolation_matrix
+from dolfinx.la.petsc import create_vector_wrap
 from dolfinx.mesh import CellType, create_unit_square
 
 # Solution scalar (e.g., float32, complex128) and geometry (float32/64)
@@ -369,7 +370,7 @@
 # `u` solution (degree-of-freedom) vectors and solve.
 
 # +
-x = PETSc.Vec().createNest([la.create_petsc_vector_wrap(sigma.x), la.create_petsc_vector_wrap(u.x)])
+x = PETSc.Vec().createNest([create_vector_wrap(sigma.x), create_vector_wrap(u.x)])
 ksp.solve(b, x)
 reason = ksp.getConvergedReason()
 assert reason > 0, f"Krylov solver has not converged {reason}."

python/demo/demo_stokes.py

@@ -115,6 +115,7 @@
 )
 from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block
 from dolfinx.io import XDMFFile
+from dolfinx.la.petsc import create_vector_wrap
 from dolfinx.mesh import CellType, create_rectangle, locate_entities_boundary
 
 # We create a {py:class}`Mesh <dolfinx.mesh.Mesh>`, define functions for
@@ -291,7 +292,7 @@ def nested_iterative_solver():
     # space `Q`). The vectors for `u` and `p` are combined to form a
     # nested vector and the system is solved.
     u, p = Function(V), Function(Q)
-    x = PETSc.Vec().createNest([la.create_petsc_vector_wrap(u.x), la.create_petsc_vector_wrap(p.x)])
+    x = PETSc.Vec().createNest([create_vector_wrap(u.x), create_vector_wrap(p.x)])
     ksp.solve(b, x)
 
     # Save solution to file in XDMF format for visualization, e.g. with

python/dolfinx/fem/petsc.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2018-2022 Garth N. Wells and Jørgen S. Dokken
+# Copyright (C) 2018-2025 Garth N. Wells and Jørgen S. Dokken
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -16,6 +16,7 @@
 import contextlib
 import functools
 import typing
+from collections.abc import Sequence
 
 from petsc4py import PETSc
 
@@ -27,8 +28,8 @@
 import numpy as np
 
 import dolfinx.cpp as _cpp
+import dolfinx.la.petsc
 import ufl
-from dolfinx import la
 from dolfinx.cpp.fem import pack_coefficients as _pack_coefficients
 from dolfinx.cpp.fem import pack_constants as _pack_constants
 from dolfinx.cpp.fem.petsc import discrete_curl as _discrete_curl
@@ -42,7 +43,6 @@
 from dolfinx.fem.forms import form as _create_form
 from dolfinx.fem.function import Function as _Function
 from dolfinx.fem.function import FunctionSpace as _FunctionSpace
-from dolfinx.la import create_petsc_vector
 
 __all__ = [
     "LinearProblem",
@@ -55,6 +55,7 @@
     "assemble_vector",
     "assemble_vector_block",
     "assemble_vector_nest",
+    "assign",
     "create_matrix",
     "create_matrix_block",
     "create_matrix_nest",
@@ -118,7 +119,7 @@ def create_vector(L: Form) -> PETSc.Vec:
         A PETSc vector with a layout that is compatible with ``L``.
     """
     dofmap = L.function_spaces[0].dofmaps(0)
-    return create_petsc_vector(dofmap.index_map, dofmap.index_map_bs)
+    return dolfinx.la.petsc.create_vector(dofmap.index_map, dofmap.index_map_bs)
 
 
 def create_vector_block(L: list[Form]) -> PETSc.Vec:
@@ -261,7 +262,7 @@ def assemble_vector(L: typing.Any, constants=None, coeffs=None) -> PETSc.Vec:
     Returns:
         An assembled vector.
     """
-    b = create_petsc_vector(
+    b = dolfinx.la.petsc.create_vector(
         L.function_spaces[0].dofmaps(0).index_map, L.function_spaces[0].dofmaps(0).index_map_bs
     )
     with b.localForm() as b_local:
@@ -881,7 +882,7 @@ def __init__(
         else:
             self.u = u
 
-        self._x = la.create_petsc_vector_wrap(self.u.x)
+        self._x = dolfinx.la.petsc.create_vector_wrap(self.u.x)
         self.bcs = bcs
 
         self._solver = PETSc.KSP().create(self.u.function_space.mesh.comm)
@@ -939,12 +940,12 @@ def solve(self) -> _Function:
 
     @property
     def L(self) -> Form:
-        """The compiled linear form"""
+        """The compiled linear form."""
         return self._L
 
     @property
     def a(self) -> Form:
-        """The compiled bilinear form"""
+        """The compiled bilinear form."""
         return self._a
 
     @property
@@ -1015,12 +1016,12 @@ def __init__(
 
     @property
     def L(self) -> Form:
-        """Compiled linear form (the residual form)"""
+        """The compiled linear form (the residual form)."""
         return self._L
 
     @property
     def a(self) -> Form:
-        """Compiled bilinear form (the Jacobian form)"""
+        """The compiled bilinear form (the Jacobian form)."""
         return self._a
 
     def form(self, x: PETSc.Vec) -> None:
@@ -1118,3 +1119,61 @@ def interpolation_matrix(space0: _FunctionSpace, space1: _FunctionSpace) -> PETS
         Interpolation matrix.
     """
     return _interpolation_matrix(space0._cpp_object, space1._cpp_object)
+
+
+@functools.singledispatch
+def assign(u: typing.Union[_Function, Sequence[_Function]], x: PETSc.Vec):
+    """Assign :class:`Function` degrees-of-freedom to a vector.
+
+    Assigns degree-of-freedom values in values of ``u``, which is possibly a
+    Sequence of ``Functions``s, to ``x``. When ``u`` is a Sequence of
+    ``Function``s, degrees-of-freedom for the ``Function``s in ``u`` are
+    'stacked' and assigned to ``x``. See :func:`assign` for documentation on
+    how stacked assignment is handled.
+
+    Args:
+        u: ``Function`` (s) to assign degree-of-freedom value from.
+        x: Vector to assign degree-of-freedom values in ``u`` to.
+    """
+    if x.getType() == PETSc.Vec.Type().NEST:
+        dolfinx.la.petsc.assign([v.x.array for v in u], x)
+    else:
+        if isinstance(u, Sequence):
+            data0, data1 = [], []
+            for v in u:
+                bs = v.function_space.dofmap.bs
+                n = v.function_space.dofmap.index_map.size_local
+                data0.append(v.x.array[: bs * n])
+                data1.append(v.x.array[bs * n :])
+            dolfinx.la.petsc.assign(data0 + data1, x)
+        else:
+            dolfinx.la.petsc.assign(u.x.array, x)
+
+
+@assign.register(PETSc.Vec)
+def _(x: PETSc.Vec, u: typing.Union[_Function, Sequence[_Function]]):
+    """Assign vector entries to :class:`Function` degrees-of-freedom.
+
+    Assigns values in ``x`` to the degrees-of-freedom of ``u``, which is
+    possibly a Sequence of ``Function``s. When ``u`` is a Sequence of
+    ``Function``s, values in ``x`` are assigned block-wise to the
+    ``Function``s. See :func:`assign` for documentation on how blocked
+    assignment is handled.
+
+    Args:
+        x: Vector with values to assign values from.
+        u: ``Function`` (s) to assign degree-of-freedom values to.
+    """
+    if x.getType() == PETSc.Vec.Type().NEST:
+        dolfinx.la.petsc.assign(x, [v.x.array for v in u])
+    else:
+        if isinstance(u, Sequence):
+            data0, data1 = [], []
+            for v in u:
+                bs = v.function_space.dofmap.bs
+                n = v.function_space.dofmap.index_map.size_local
+                data0.append(v.x.array[: bs * n])
+                data1.append(v.x.array[bs * n :])
+            dolfinx.la.petsc.assign(x, data0 + data1)
+        else:
+            dolfinx.la.petsc.assign(x, u.x.array)

python/dolfinx/la.py python/dolfinx/la/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2017-2021 Garth N. Wells
+# Copyright (C) 2017-2025 Garth N. Wells, Jack S. Hale
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -10,6 +10,7 @@
 import numpy as np
 import numpy.typing as npt
 
+import dolfinx
 from dolfinx import cpp as _cpp
 from dolfinx.cpp.common import IndexMap
 from dolfinx.cpp.la import BlockMode, InsertMode, Norm
@@ -19,9 +20,9 @@
     "MatrixCSR",
     "Norm",
     "Vector",
-    "create_petsc_vector",
     "is_orthonormal",
     "matrix_csr",
+    "norm",
     "orthonormalize",
     "vector",
 ]
@@ -93,8 +94,12 @@ def petsc_vec(self):
           When the object is destroyed it will destroy the underlying petsc4py
           vector automatically.
         """
+        assert dolfinx.has_petsc4py
+
+        from dolfinx.la.petsc import create_vector_wrap
+
         if self._petsc_x is None:
-            self._petsc_x = create_petsc_vector_wrap(self)
+            self._petsc_x = create_vector_wrap(self)
         return self._petsc_x
 
     def scatter_forward(self) -> None:
@@ -326,56 +331,8 @@ def vector(map, bs=1, dtype: npt.DTypeLike = np.float64) -> Vector:
     return Vector(vtype(map, bs))
 
 
-def create_petsc_vector_wrap(x: Vector):
-    """Wrap a distributed DOLFINx vector as a PETSc vector.
-
-    Note:
-        Due to subtle issues in the interaction between petsc4py memory management
-        and the Python garbage collector, it is recommended that the method ``PETSc.Vec.destroy()``
-        is called on the returned object once the object is no longer required. Note that
-        ``PETSc.Vec.destroy()`` is collective over the object's MPI communicator.
-
-    Args:
-        x: The vector to wrap as a PETSc vector.
-
-    Returns:
-        A PETSc vector that shares data with ``x``.
-    """
-    from petsc4py import PETSc
-
-    map = x.index_map
-    ghosts = map.ghosts.astype(PETSc.IntType)  # type: ignore
-    bs = x.block_size
-    size = (map.size_local * bs, map.size_global * bs)
-    return PETSc.Vec().createGhostWithArray(ghosts, x.array, size=size, bsize=bs, comm=map.comm)  # type: ignore
-
-
-def create_petsc_vector(map, bs: int):
-    """Create a distributed PETSc vector.
-
-    Note:
-        Due to subtle issues in the interaction between petsc4py memory management
-        and the Python garbage collector, it is recommended that the method ``PETSc.Vec.destroy()``
-        is called on the returned object once the object is no longer required. Note that
-        ``PETSc.Vec.destroy()`` is collective over the object's MPI communicator.
-
-    Args:
-        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.
-    """
-    from petsc4py import PETSc
-
-    ghosts = map.ghosts.astype(PETSc.IntType)  # type: ignore
-    size = (map.size_local * bs, map.size_global * bs)
-    return PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=map.comm)  # type: ignore
-
-
 def orthonormalize(basis: list[Vector]):
-    """Orthogonalise set of PETSc vectors in-place."""
+    """Orthogonalise set of vectors in-place."""
     _cpp.la.orthonormalize([x._cpp_object for x in basis])