# Source code for dolfinx.fem.petsc

    # Copyright (C) 2018-2025 Garth N. Wells, Nathan Sime and Jørgen S. Dokken
    #
    # This file is part of DOLFINx (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later
    """High-level solver classes and functions for assembling PETSc objects.

    Functions in this module generally apply functions in :mod:`dolfinx.fem`
    to PETSc linear algebra objects and handle any PETSc-specific
    preparation.

    Note:
        The following does not apply to the high-level classes
        :class:`dolfinx.fem.petsc.LinearProblem`
        :class:`dolfinx.fem.petsc.NonlinearProblem`.

        Due to subtle issues in the interaction between petsc4py memory
        management and the Python garbage collector, it is recommended that
        the PETSc method ``destroy()`` is called on returned PETSc objects
        once the object is no longer required. Note that ``destroy()`` is
        collective over the object's MPI communicator.
    """

    from __future__ import annotations

    import contextlib
    import ctypes as _ctypes
    import functools
    import os
    import pathlib
    from collections.abc import Sequence

    from petsc4py import PETSc

    # ruff: noqa: E402
    import dolfinx

    assert dolfinx.has_petsc4py

    import warnings
    from functools import partial

    import numpy as np
    from numpy import typing as npt

    import dolfinx.cpp as _cpp
    import dolfinx.la.petsc
    import ufl
    from dolfinx.cpp.fem.petsc import discrete_curl as _discrete_curl
    from dolfinx.cpp.fem.petsc import discrete_gradient as _discrete_gradient
    from dolfinx.cpp.fem.petsc import interpolation_matrix as _interpolation_matrix
    from dolfinx.fem import IntegralType, pack_coefficients, pack_constants
    from dolfinx.fem.assemble import _assemble_vector_array
    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 extract_function_spaces as _extract_function_spaces
    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.mesh import EntityMap as _EntityMap

    __all__ = [
        "LinearProblem",
        "NewtonSolverNonlinearProblem",
        "NonlinearProblem",
        "apply_lifting",
        "assemble_jacobian",
        "assemble_matrix",
        "assemble_residual",
        "assemble_vector",
        "assign",
        "cffi_utils",
        "create_matrix",
        "create_vector",
        "ctypes_utils",
        "discrete_curl",
        "discrete_gradient",
        "interpolation_matrix",
        "numba_utils",
        "set_bc",
    ]

    # -- Vector instantiation -------------------------------------------------

    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 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 ``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 ``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 ``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]]

        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:
            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 ``V``. The
            vector is not initialised to zero.
        """
        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.")

        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 -------------------------------------------------

    def create_matrix(
        a: Form | Sequence[Sequence[Form]],
        kind: str | Sequence[Sequence[str]] | None = None,
    ) -> PETSc.Mat:  # type: ignore[name-defined]
        """Create a PETSc matrix that is compatible with the (sequence) of
        bilinear form(s).

        Three cases are supported:

        1. For a single bilinear form, it creates a compatible PETSc matrix
           of type ``kind``.
        2. For a rectangular array of bilinear forms, if ``kind`` is
           ``PETSc.Mat.Type.NEST`` or ``kind`` is an array of PETSc ``Mat``
           types (with the same shape as ``a``), a matrix of type
           ``PETSc.Mat.Type.NEST`` is created. The matrix is compatible
           with the forms ``a``.
        3. For a rectangular array of bilinear forms, it create a single
           (non-nested) matrix of type ``kind`` that is compatible with the
           array of for forms ``a``. If ``kind`` is ``None``, then the
           matrix is the default type.

           In this case, the matrix is arranged::

                 A = [a_00 ... a_0n]
                     [a_10 ... a_1n]
                     [     ...     ]
                     [a_m0 ..  a_mn]

        Args:
            a: A bilinear form or a nested sequence of bilinear forms.
            kind: The PETSc matrix type (``MatType``).

        Returns:
            A PETSc matrix.
        """
        if isinstance(a, Sequence):
            _a = [[None if form is None else form._cpp_object for form in arow] for arow in a]
            if kind == PETSc.Mat.Type.NEST:  # type: ignore[attr-defined]
                # Create nest matrix with default types
                return _cpp.fem.petsc.create_matrix_nest(_a, None)
            else:
                if kind is None or isinstance(kind, str):  # Single 'kind' type
                    return _cpp.fem.petsc.create_matrix_block(_a, kind)
                else:  # Array of 'kind' types
                    return _cpp.fem.petsc.create_matrix_nest(_a, kind)
        else:  # Single form
            return _cpp.fem.petsc.create_matrix(a._cpp_object, kind)

    # -- Vector assembly ------------------------------------------------------

    @functools.singledispatch
    def assemble_vector(
        L: Form | Sequence[Form],
        constants: npt.NDArray | Sequence[npt.NDArray] | None = None,
        coeffs: (
            dict[tuple[IntegralType, int], npt.NDArray]
            | Sequence[dict[tuple[IntegralType, int], npt.NDArray]]
            | None
        ) = None,
        kind: str | None = None,
    ) -> PETSc.Vec:  # type: ignore[name-defined]
        """Assemble linear form(s) into a new PETSc vector.

        Three cases are supported:

        1. If ``L`` is a single linear form, the form is assembled into a
           ghosted PETSc vector.

        2. If ``L`` is a sequence of linear forms and ``kind`` is ``None``
           or is ``PETSc.Vec.Type.MPI``, the forms are assembled into a
           vector ``b`` such that ``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]``. See :func:`create_vector`
           for a description of the offset blocks.

        3. If ``L`` is a sequence of linear forms and ``kind`` is
           ``PETSc.Vec.Type.NEST``, the forms are assembled into a PETSc
           nested vector ``b`` (a nest of ghosted PETSc vectors) such that
           ``L[i]`` is assembled into into the ith nested matrix in ``b``.

        Constant and coefficient data that appear in the forms(s) can be
        packed outside of this function to avoid re-packing by this
        function. The functions :func:`dolfinx.fem.pack_constants` and
        :func:`dolfinx.fem.pack_coefficients` can be used to 'pre-pack' the
        data.

        Note:
            The returned vector is not finalised, i.e. ghost values are not
            accumulated on the owning processes.

        Args:
            L: A linear form or sequence of linear forms.
            constants: Constants appearing in the form. For a single form,
                ``constants.ndim==1``. For multiple forms, the constants for
                form ``L[i]`` are  ``constants[i]``.
            coeffs: Coefficients appearing in the form. For a single form,
                ``coeffs.shape=(num_cells, n)``. For multiple forms, the
                coefficients for form ``L[i]`` are  ``coeffs[i]``.
            kind: PETSc vector type.

        Returns:
            An assembled vector.
        """
        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]

    @assemble_vector.register
    def _(
        b: PETSc.Vec,  # type: ignore[name-defined]
        L: Form | Sequence[Form],
        constants: npt.NDArray | Sequence[npt.NDArray] | None = None,
        coeffs: (
            dict[tuple[IntegralType, int], npt.NDArray]
            | Sequence[dict[tuple[IntegralType, int], npt.NDArray]]
            | None
        ) = None,
    ) -> PETSc.Vec:  # type: ignore[name-defined]
        """Assemble linear form(s) into a PETSc vector.

        The vector ``b`` must have been initialized with a size/layout that
        is consistent with the linear form. The vector ``b`` is normally
        created by :func:`create_vector`.

        Constants and coefficients that appear in the forms(s) can be passed
        to avoid re-computation of constants and coefficients. The functions
        :func:`dolfinx.fem.assemble.pack_constants` and
        :func:`dolfinx.fem.assemble.pack_coefficients` can be called.

        Note:
            The vector is not zeroed before assembly and it is not
            finalised, i.e. ghost values are not accumulated on the owning
            processes.

        Args:
            b: Vector to assemble the contribution of the linear form into.
            L: A linear form or sequence of linear forms to assemble into
                ``b``.
            constants: Constants appearing in the form. For a single form,
                ``constants.ndim==1``. For multiple forms, the constants for
                form ``L[i]`` are  ``constants[i]``.
            coeffs: Coefficients appearing in the form. For a single form,
                ``coeffs.shape=(num_cells, n)``. For multiple forms, the
                coefficients for form ``L[i]`` are  ``coeffs[i]``.

        Returns:
            Assembled vector.
        """
        if b.getType() == PETSc.Vec.Type.NEST:  # type: ignore[attr-defined]
            if not isinstance(L, Sequence):
                raise ValueError("Must provide a sequence of forms when assembling a nest vector")
            constants = [None] * len(L) if constants is None else constants  # type: ignore[list-item]
            coeffs = [None] * len(L) if coeffs is None else coeffs  # type: ignore[list-item]
            for b_sub, L_sub, const, coeff in zip(b.getNestSubVecs(), L, constants, coeffs):
                with b_sub.localForm() as b_local:
                    _assemble_vector_array(b_local.array_w, L_sub, const, coeff)  # type: ignore[arg-type]
        elif isinstance(L, Sequence):
            constants = pack_constants(L) if constants is None else constants
            coeffs = pack_coefficients(L) if coeffs is None else coeffs
            offset0, offset1 = b.getAttr("_blocks")
            with b.localForm() as b_l:
                for L_, const, coeff, off0, off1, offg0, offg1 in zip(
                    L, constants, coeffs, offset0, offset0[1:], offset1, offset1[1:]
                ):
                    bx_ = np.zeros((off1 - off0) + (offg1 - offg0), dtype=PETSc.ScalarType)  # type: ignore[attr-defined]
                    _assemble_vector_array(bx_, L_, const, coeff)  # type: ignore[arg-type]
                    size = off1 - off0
                    b_l.array_w[off0:off1] += bx_[:size]
                    b_l.array_w[offg0:offg1] += bx_[size:]
        else:
            with b.localForm() as b_local:
                _assemble_vector_array(b_local.array_w, L, constants, coeffs)  # type: ignore[arg-type]

        return b

    # -- Matrix assembly ------------------------------------------------------

    @functools.singledispatch
    def assemble_matrix(
        a: Form | Sequence[Sequence[Form]],
        bcs: Sequence[DirichletBC] | None = None,
        diag: float = 1,
        constants: Sequence[npt.NDArray] | Sequence[Sequence[npt.NDArray]] | None = None,
        coeffs: (
            dict[tuple[IntegralType, int], npt.NDArray]
            | Sequence[dict[tuple[IntegralType, int], npt.NDArray]]
            | None
        ) = None,
        kind=None,
    ):
        """Assemble a bilinear form into a matrix.

        The following cases are supported:

        1. If ``a`` is a single bilinear form, the form is assembled
           into PETSc matrix of type ``kind``.
        #. If ``a`` is a :math:`m \\times n` rectangular array of forms the
           forms in ``a`` are assembled into a matrix such that::

                A = [A_00 ... A_0n]
                    [A_10 ... A_1n]
                    [     ...     ]
                    [A_m0 ..  A_mn]

           where ``A_ij`` is the matrix associated with the form
           ``a[i][j]``.

           a. If ``kind`` is a ``PETSc.Mat.Type`` (other than
              ``PETSc.Mat.Type.NEST``) or is ``None``, the matrix type is
              ``kind`` of the default type (if ``kind`` is ``None``).
           #. If ``kind`` is ``PETSc.Mat.Type.NEST`` or a rectangular array
              of PETSc matrix types, the returned matrix has type
              ``PETSc.Mat.Type.NEST``.

        Rows/columns that are constrained by a Dirichlet boundary condition
        are zeroed, with the diagonal to set to ``diag``.

        Constant and coefficient data that appear in the forms(s) can be
        packed outside of this function to avoid re-packing by this
        function. The functions :func:`dolfinx.fem.pack_constants` and
        :func:`dolfinx.fem.pack_coefficients` can be used to 'pre-pack' the
        data.

        Note:
            The returned matrix is not 'assembled', i.e. ghost contributions
            are not accumulated.

        Args:
            a: Bilinear form(s) to assembled into a matrix.
            bc: Dirichlet boundary conditions applied to the system.
            diag: Value to set on the matrix diagonal for Dirichlet
                boundary condition constrained degrees-of-freedom belonging
                to the same trial and test space.
            constants: Constants appearing the in the form.
            coeffs: Coefficients appearing the in the form.

        Returns:
            Matrix representing the bilinear form.
        """
        A = create_matrix(a, kind)
        assemble_matrix(A, a, bcs, diag, constants, coeffs)  # type: ignore[arg-type]
        return A

    @assemble_matrix.register
    def _(
        A: PETSc.Mat,  # type: ignore[name-defined]
        a: Form | Sequence[Sequence[Form]],
        bcs: Sequence[DirichletBC] | None = None,
        diag: float = 1,
        constants: npt.NDArray | Sequence[Sequence[npt.NDArray]] | None = None,
        coeffs: (
            dict[tuple[IntegralType, int], npt.NDArray]
            | Sequence[Sequence[dict[tuple[IntegralType, int], npt.NDArray]]]
            | None
        ) = None,
    ) -> PETSc.Mat:  # type: ignore[name-defined]
        """Assemble bilinear form into a matrix.

        The matrix vector ``A`` must have been initialized with a
        size/layout that is consistent with the bilinear form(s). The PETSc
        matrix ``A`` is normally created by :func:`create_matrix`.

        The returned matrix is not finalised, i.e. ghost values are not
        accumulated.
        """
        if A.getType() == PETSc.Mat.Type.NEST:  # type: ignore[attr-defined]
            if not isinstance(a, Sequence):
                raise ValueError("Must provide a sequence of forms when assembling a nest matrix")
            constants = [pack_constants(forms) for forms in a] if constants is None else constants  # type: ignore[misc]
            coeffs = [pack_coefficients(forms) for forms in a] if coeffs is None else coeffs  # type: ignore[misc]
            for i, (a_row, const_row, coeff_row) in enumerate(zip(a, constants, coeffs)):
                for j, (a_block, const, coeff) in enumerate(zip(a_row, const_row, coeff_row)):
                    if a_block is not None:
                        Asub = A.getNestSubMatrix(i, j)
                        assemble_matrix(Asub, a_block, bcs, diag, const, coeff)  # type: ignore[arg-type]
                    elif i == j:
                        for bc in bcs:
                            row_forms = [row_form for row_form in a_row if row_form is not None]
                            assert len(row_forms) > 0
                            if row_forms[0].function_spaces[0].contains(bc.function_space):
                                raise RuntimeError(
                                    f"Diagonal sub-block ({i}, {j}) cannot be 'None'"
                                    " and have DirichletBC applied."
                                    " Consider assembling a zero block."
                                )
        elif isinstance(a, Sequence):  # Blocked
            consts = [pack_constants(forms) for forms in a] if constants is None else constants  # type: ignore[misc]
            coeffs = [pack_coefficients(forms) for forms in a] if coeffs is None else coeffs  # type: ignore[misc]
            V = (_extract_function_spaces(a, 0), _extract_function_spaces(a, 1))
            for index in range(2):
                # the check below is to ensure that a .dofmaps attribute is
                # available when creating is0 and is1 below
                Vi = V[index]
                assert isinstance(Vi, list)
                if all(Vsub is None for Vsub in Vi):
                    raise ValueError(
                        "Cannot have a entire {'row' if index == 0 else 'column'} of a full of None"
                    )
            is0 = _cpp.la.petsc.create_index_sets(
                [(Vsub.dofmaps(0).index_map, Vsub.dofmaps(0).index_map_bs) for Vsub in V[0]]  # type: ignore[union-attr]
            )
            is1 = _cpp.la.petsc.create_index_sets(
                [(Vsub.dofmaps(0).index_map, Vsub.dofmaps(0).index_map_bs) for Vsub in V[1]]  # type: ignore[union-attr]
            )

            _bcs = [bc._cpp_object for bc in bcs] if bcs is not None else []
            for i, a_row in enumerate(a):
                for j, a_sub in enumerate(a_row):
                    if a_sub is not None:
                        Asub = A.getLocalSubMatrix(is0[i], is1[j])
                        _cpp.fem.petsc.assemble_matrix(
                            Asub,
                            a_sub._cpp_object,
                            consts[i][j],
                            coeffs[i][j],  # type: ignore[index]
                            _bcs,
                            True,
                        )
                        A.restoreLocalSubMatrix(is0[i], is1[j], Asub)
                    elif i == j:
                        for bc in _bcs:
                            row_forms = [row_form for row_form in a_row if row_form is not None]
                            assert len(row_forms) > 0
                            if row_forms[0].function_spaces[0].contains(bc.function_space):
                                raise RuntimeError(
                                    f"Diagonal sub-block ({i}, {j}) cannot be 'None' "
                                    " and have DirichletBC applied."
                                    " Consider assembling a zero block."
                                )

            # Flush to enable switch from add to set in the matrix
            A.assemble(PETSc.Mat.AssemblyType.FLUSH)  # type: ignore[attr-defined]

            # Set diagonal
            for i, a_row in enumerate(a):
                for j, a_sub in enumerate(a_row):
                    if a_sub is not None:
                        Asub = A.getLocalSubMatrix(is0[i], is1[j])
                        if a_sub.function_spaces[0] is a_sub.function_spaces[1]:
                            _cpp.fem.petsc.insert_diagonal(Asub, a_sub.function_spaces[0], _bcs, diag)
                        A.restoreLocalSubMatrix(is0[i], is1[j], Asub)
        else:  # Non-blocked
            constants = pack_constants(a) if constants is None else constants  # type: ignore[assignment]
            coeffs = pack_coefficients(a) if coeffs is None else coeffs  # type: ignore[assignment]
            _bcs = [bc._cpp_object for bc in bcs] if bcs is not None else []
            _cpp.fem.petsc.assemble_matrix(A, a._cpp_object, constants, coeffs, _bcs)
            if a.function_spaces[0] is a.function_spaces[1]:
                A.assemblyBegin(PETSc.Mat.AssemblyType.FLUSH)  # type: ignore[attr-defined]
                A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH)  # type: ignore[attr-defined]
                _cpp.fem.petsc.insert_diagonal(A, a.function_spaces[0], _bcs, diag)

        return A

    # -- Modifiers for Dirichlet conditions -----------------------------------

    def apply_lifting(
        b: PETSc.Vec,  # type: ignore[name-defined]
        a: Sequence[Form] | Sequence[Sequence[Form]],
        bcs: Sequence[DirichletBC] | Sequence[Sequence[DirichletBC]] | None,
        x0: Sequence[PETSc.Vec] | None = None,  # type: ignore[name-defined]
        alpha: float = 1,
        constants: Sequence[npt.NDArray] | Sequence[Sequence[npt.NDArray]] | None = None,
        coeffs: (
            dict[tuple[IntegralType, int], npt.NDArray]
            | Sequence[Sequence[dict[tuple[IntegralType, int], npt.NDArray]]]
            | None
        ) = None,
    ) -> None:
        r"""Modify the right-hand side PETSc vector ``b`` to account for
        constraints (Dirichlet boundary conitions).

        See :func:`dolfinx.fem.apply_lifting` for a mathematical
        descriptions of the lifting operation.

        Args:
            b: Vector to modify in-place.
            a: List of bilinear forms. If ``b`` is not blocked or a nest,
                then ``a`` is a 1D sequence. If ``b`` is blocked or a nest,
                then ``a`` is  a 2D array of forms, with the ``a[i]`` forms
                used to modify the block/nest vector ``b[i]``.
            bcs: Boundary conditions to apply, which form a 2D array.
                If ``b`` is nested or blocked then ``bcs[i]`` are the
                boundary conditions to apply to block/nest ``i``.
                The function :func:`dolfinx.fem.bcs_by_block` can be
                used to prepare the 2D array of ``DirichletBC`` objects
                from the 2D sequence ``a``::

                    bcs1 = fem.bcs_by_block(
                        fem.extract_function_spaces(a, 1),
                        bcs
                    )

                If ``b`` is not blocked or nest, then ``len(bcs)`` must be
                equal to 1. The function :func:`dolfinx.fem.bcs_by_block`
                can be used to prepare the 2D array of ``DirichletBC``
                from the 1D sequence ``a``::

                    bcs1 = fem.bcs_by_block(
                        fem.extract_function_spaces([a], 1),
                        bcs
                    )

            x0: Vector to use in modify ``b`` (see
                :func:`dolfinx.fem.apply_lifting`). Treated as zero if
                ``None``.
            alpha: Scalar parameter in lifting (see
                :func:`dolfinx.fem.apply_lifting`).
            constants: Packed constant data appearing in the forms ``a``. If
                ``None``, the constant data will be packed by the function.
            coeffs: Packed coefficient data appearing in the forms ``a``. If
                ``None``, the coefficient data will be packed by the
                function.

        Note:
            Ghost contributions are not accumulated (not sent to owner).
            Caller is responsible for reverse-scatter to update the ghosts.

        Note:
            Boundary condition values are *not* set in ``b`` by this
            function. Use :func:`dolfinx.fem.DirichletBC.set` to set values
            in ``b``.
        """
        if b.getType() == PETSc.Vec.Type.NEST:  # type: ignore[attr-defined]
            x0 = [] if x0 is None else x0.getNestSubVecs()  # type: ignore[attr-defined]
            constants = [pack_constants(forms) for forms in a] if constants is None else constants  # type: ignore[assignment]
            coeffs = [pack_coefficients(forms) for forms in a] if coeffs is None else coeffs  # type: ignore[misc]
            for b_sub, a_sub, const, coeff in zip(b.getNestSubVecs(), a, constants, coeffs):  # type: ignore[arg-type]
                const_ = list(
                    map(lambda x: np.array([], dtype=PETSc.ScalarType) if x is None else x, const)  # type: ignore[attr-defined, call-overload]
                )
                apply_lifting(b_sub, a_sub, bcs, x0, alpha, const_, coeff)  # type: ignore[arg-type]
        else:
            with contextlib.ExitStack() as stack:
                if b.getAttr("_blocks") is not None:
                    if x0 is not None:
                        offset0, offset1 = x0.getAttr("_blocks")  # type: ignore[attr-defined]
                        xl = stack.enter_context(x0.localForm())  # type: ignore[attr-defined]
                        xlocal = [
                            np.concatenate((xl[off0:off1], xl[offg0:offg1]))
                            for (off0, off1, offg0, offg1) in zip(
                                offset0, offset0[1:], offset1, offset1[1:]
                            )
                        ]
                    else:
                        xlocal = None

                    offset0, offset1 = b.getAttr("_blocks")
                    with b.localForm() as b_l:
                        for i, (a_, off0, off1, offg0, offg1) in enumerate(
                            zip(a, offset0, offset0[1:], offset1, offset1[1:])
                        ):
                            const = pack_constants(a_) if constants is None else constants[i]  # type: ignore[arg-type]
                            coeff = pack_coefficients(a_) if coeffs is None else coeffs[i]  # type: ignore[arg-type, assignment, index]
                            const_ = [
                                np.empty(0, dtype=PETSc.ScalarType) if val is None else val  # type: ignore[attr-defined]
                                for val in const
                            ]
                            bx_ = np.concatenate((b_l[off0:off1], b_l[offg0:offg1]))
                            _apply_lifting(bx_, a_, bcs, xlocal, float(alpha), const_, coeff)  # type: ignore[arg-type]
                            size = off1 - off0
                            b_l.array_w[off0:off1] = bx_[:size]
                            b_l.array_w[offg0:offg1] = bx_[size:]
                else:
                    x0 = [] if x0 is None else x0
                    x0 = [stack.enter_context(x.localForm()) for x in x0]
                    x0_r = [x.array_r for x in x0]
                    b_local = stack.enter_context(b.localForm())
                    _apply_lifting(b_local.array_w, a, bcs, x0_r, alpha, constants, coeffs)  # type: ignore[arg-type]

        return b

    def set_bc(
        b: PETSc.Vec,  # type: ignore[name-defined]
        bcs: Sequence[DirichletBC] | Sequence[Sequence[DirichletBC]],
        x0: PETSc.Vec | None = None,  # type: ignore[name-defined]
        alpha: float = 1,
    ) -> None:
        r"""Set constraint (Dirchlet boundary condition) values in an vector.

        For degrees-of-freedoms that are constrained by a Dirichlet boundary
        condition, this function sets that degrees-of-freedom to ``alpha *
        (g - x0)``, where ``g`` is the boundary condition value.

        Only owned entries in ``b`` (owned by the MPI process) are modified
        by this function.

        Args:
            b: Vector to modify by setting  boundary condition values.
            bcs: Boundary conditions to apply. If ``b`` is nested or
                blocked, ``bcs`` is a 2D array and ``bcs[i]`` are the
                boundary conditions to apply to block/nest ``i``. Otherwise
                ``bcs`` should be a sequence of ``DirichletBC``\s. For
                block/nest problems, :func:`dolfinx.fem.bcs_by_block` can be
                used to prepare the 2D array of ``DirichletBC`` objects.
            x0: Vector used in the value that constrained entries are set
                to. If ``None``, ``x0`` is treated as zero.
            alpha: Scalar value used in the value that constrained entries
                are set to.
        """
        if len(bcs) == 0:
            return

        if not isinstance(bcs[0], Sequence):
            x0 = x0.array_r if x0 is not None else None
            for bc in bcs:
                bc.set(b.array_w, x0, alpha)  # type: ignore[union-attr]
        elif b.getType() == PETSc.Vec.Type.NEST:  # type: ignore[attr-defined]
            _b = b.getNestSubVecs()
            x0 = len(_b) * [None] if x0 is None else x0.getNestSubVecs()
            for b_sub, bc, x_sub in zip(_b, bcs, x0):  # type: ignore[assignment, arg-type]
                set_bc(b_sub, bc, x_sub, alpha)  # type: ignore[arg-type]
        else:  # block vector
            offset0, _ = b.getAttr("_blocks")
            b_array = b.getArray(readonly=False)
            x_array = x0.getArray(readonly=True) if x0 is not None else None
            for bcs, off0, off1 in zip(bcs, offset0, offset0[1:]):  # type: ignore[assignment]
                x0_sub = x_array[off0:off1] if x0 is not None else None  # type: ignore[index]
                for bc in bcs:
                    bc.set(b_array[off0:off1], x0_sub, alpha)  # type: ignore[arg-type, union-attr]

    # -- High-level interface for KSP ---------------------------------------

    class LinearProblem:
        """High-level class for solving a linear variational problem using
        a PETSc KSP.

        Solves problems of the form
        :math:`a_{ij}(u, v) = f_i(v), i,j=0,\\ldots,N\\
        \\forall v \\in V` where
        :math:`u=(u_0,\\ldots,u_N), v=(v_0,\\ldots,v_N)`
        using PETSc KSP as the linear solver.

        Note:
            This high-level class automatically handles PETSc memory
            management. The user does not need to manually call
            ``.destroy()`` on returned PETSc objects.
        """

        def __init__(
            self,
            a: ufl.Form | Sequence[Sequence[ufl.Form]],
            L: ufl.Form | Sequence[ufl.Form],
            *,
            petsc_options_prefix: str,
            bcs: Sequence[DirichletBC] | None = None,
            u: _Function | Sequence[_Function] | None = None,
            P: ufl.Form | Sequence[Sequence[ufl.Form]] | None = None,
            kind: str | Sequence[Sequence[str]] | None = None,
            petsc_options: dict | None = None,
            form_compiler_options: dict | None = None,
            jit_options: dict | None = None,
            entity_maps: Sequence[_EntityMap] | None = None,
        ) -> None:
            """Initialize solver for a linear variational problem.

            By default, the underlying KSP solver uses PETSc's default
            options, usually GMRES + ILU preconditioning. To use the robust
            combination of LU via MUMPS

            Example::

                problem = LinearProblem(a, L, bcs=[bc0, bc1],
                    petsc_options_prefix="basic_linear_problem",
                    petsc_options= {
                      "ksp_type": "preonly",
                      "pc_type": "lu",
                      "pc_factor_mat_solver_type": "mumps"
                })

            This class also supports nested block-structured problems.

            Example::

                problem = LinearProblem([[a00, a01], [None, a11]], [L0, L1],
                    bcs=[bc0, bc1], u=[uh0, uh1],
                    kind="nest",
                    petsc_options_prefix="nest_linear_problem")

            Every PETSc object created will have a unique options prefix set.
            We recommend discovering these prefixes dynamically via the
            petsc4py API rather than hard-coding each prefix value
            into the programme.

            Example::

                ksp_options_prefix = problem.solver.getOptionsPrefix()
                A_options_prefix = problem.A.getOptionsPrefix()

            Args:
                a: Bilinear UFL form or a nested sequence of bilinear
                    forms, the left-hand side of the variational problem.
                L: Linear UFL form or a sequence of linear forms, the
                    right-hand side of the variational problem.
                bcs: Sequence of Dirichlet boundary conditions to apply to
                     the variational problem and the preconditioner matrix.
                u: Solution function. It is created if not provided.
                P: Bilinear UFL form or a sequence of sequence of bilinear
                    forms, used as a preconditioner.
                kind: The PETSc matrix and vector kind. Common choices
                    are ``mpi`` and ``nest``. See
                    :func:`dolfinx.fem.petsc.create_matrix` and
                    :func:`dolfinx.fem.petsc.create_vector` for more
                    information.
                petsc_options_prefix: Mandatory named argument. Options prefix
                    used as root prefix on all internally created PETSc
                    objects. Typically ends with ``_``. Must be the same on
                    all ranks, and is usually unique within the programme.
                petsc_options: Options set on the underlying PETSc KSP only.
                    The options must be the same on all ranks. For available
                    choices for the ``petsc_options`` kwarg, see the `PETSc KSP
                    documentation
                    <https://petsc4py.readthedocs.io/en/stable/manual/ksp/>`_.
                    Options on other objects (matrices, vectors) should be set
                    explicitly by the user.
                form_compiler_options: Options used in FFCx compilation of
                    all forms. Run ``ffcx --help`` at the commandline to see
                    all available options.
                jit_options: Options used in CFFI JIT compilation of C
                    code generated by FFCx. See ``python/dolfinx/jit.py`` for
                    all available options. Takes priority over all other
                    option values.
                entity_maps: If any trial functions, test functions, or
                    coefficients in the form are not defined over the same mesh
                    as the integration domain, a corresponding :class:
                    `EntityMap<dolfinx.mesh.EntityMap>` must be provided.
            """
            self._a = _create_form(
                a,
                dtype=PETSc.ScalarType,  # type: ignore[attr-defined]
                form_compiler_options=form_compiler_options,
                jit_options=jit_options,
                entity_maps=entity_maps,
            )
            self._L = _create_form(
                L,
                dtype=PETSc.ScalarType,  # type: ignore[attr-defined]
                form_compiler_options=form_compiler_options,
                jit_options=jit_options,
                entity_maps=entity_maps,
            )
            self._A = create_matrix(self._a, kind=kind)
            self._preconditioner = _create_form(
                P,  # type: ignore[arg-type]
                dtype=PETSc.ScalarType,  # type: ignore[attr-defined]
                form_compiler_options=form_compiler_options,
                jit_options=jit_options,
                entity_maps=entity_maps,
            )
            self._P_mat = (
                create_matrix(self._preconditioner, kind=kind)
                if self._preconditioner is not None
                else None
            )

            # 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(_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:
                # Extract function space for unknown from the right hand
                # side of the equation.
                if isinstance(L, Sequence):
                    self._u = [_Function(Li.arguments()[0].ufl_function_space()) for Li in L]
                else:
                    self._u = _Function(L.arguments()[0].ufl_function_space())
            else:
                self._u = u

            self.bcs = [] if bcs is None else bcs

            self._solver = PETSc.KSP().create(self.A.comm)  # type: ignore[attr-defined]
            self.solver.setOperators(self.A, self.P_mat)

            if petsc_options_prefix == "":
                raise ValueError("PETSc options prefix cannot be empty.")

            self._petsc_options_prefix = petsc_options_prefix
            self.solver.setOptionsPrefix(petsc_options_prefix)
            self.A.setOptionsPrefix(f"{petsc_options_prefix}A_")
            self.b.setOptionsPrefix(f"{petsc_options_prefix}b_")
            self.x.setOptionsPrefix(f"{petsc_options_prefix}x_")
            if self.P_mat is not None:
                self.P_mat.setOptionsPrefix(f"{petsc_options_prefix}P_mat_")

            # Set options on KSP only
            if petsc_options is not None:
                opts = PETSc.Options()  # type: ignore[attr-defined]
                opts.prefixPush(self.solver.getOptionsPrefix())

                for k, v in petsc_options.items():
                    opts[k] = v

                self.solver.setFromOptions()

                # Tidy up global options
                for k in petsc_options.keys():
                    del opts[k]

                opts.prefixPop()

            if kind == "nest":
                # Transfer nest IS on self.A to PC of main KSP. This allows
                # fieldsplit preconditioning to be applied, if desired.
                nest_IS = self.A.getNestISs()
                fieldsplit_IS = tuple(
                    [
                        (f"{u.name + '_' if u.name != 'f' else ''}{i}", IS)
                        for i, (u, IS) in enumerate(zip(self.u, nest_IS[0]))
                    ]
                )
                self.solver.getPC().setFieldSplitIS(*fieldsplit_IS)

        def __del__(self):
            for obj in filter(
                lambda obj: obj is not None, (self._solver, self._A, self._b, self._x, self._P_mat)
            ):
                obj.destroy()

        def solve(self) -> _Function | Sequence[_Function]:
            """Solve the problem.

            This method updates the solution ``u`` function(s) stored in the
            problem instance.

            Note:
                The user is responsible for asserting convergence of the KSP
                solver e.g. ``problem.solver.getConvergedReason() > 0``.
                Alternatively, pass ``"ksp_error_if_not_converged" : True`` in
                ``petsc_options`` to raise a ``PETScError`` on failure.

            Returns:
                The solution function(s).
            """

            # Assemble lhs
            self.A.zeroEntries()
            assemble_matrix(self.A, self.a, bcs=self.bcs)  # type: ignore[arg-type, misc]
            self.A.assemble()

            # Assemble preconditioner
            if self.P_mat is not None:
                self.P_mat.zeroEntries()
                assemble_matrix(self.P_mat, self.preconditioner, bcs=self.bcs)  # type: ignore[arg-type, misc]
                self.P_mat.assemble()

            # Assemble rhs
            dolfinx.la.petsc._zero_vector(self.b)
            assemble_vector(self.b, self.L)  # type: ignore[arg-type]

            # Apply boundary conditions to the rhs
            if self.bcs is not None:
                if isinstance(self.u, Sequence):  # block or nest
                    bcs1 = _bcs_by_block(_extract_function_spaces(self.a, 1), self.bcs)  # type: ignore[arg-type]
                    apply_lifting(self.b, self.a, bcs=bcs1)  # type: ignore[arg-type]
                    dolfinx.la.petsc._ghost_update(
                        self.b,
                        PETSc.InsertMode.ADD,  # type: ignore[attr-defined]
                        PETSc.ScatterMode.REVERSE,  # type: ignore[attr-defined]
                    )
                    bcs0 = _bcs_by_block(_extract_function_spaces(self.L), self.bcs)  # type: ignore[arg-type]
                    dolfinx.fem.petsc.set_bc(self.b, bcs0)
                else:  # single form
                    apply_lifting(self.b, [self.a], bcs=[self.bcs])  # type: ignore[arg-type]
                    dolfinx.la.petsc._ghost_update(
                        self.b,
                        PETSc.InsertMode.ADD,  # type: ignore[attr-defined]
                        PETSc.ScatterMode.REVERSE,  # type: ignore[attr-defined]
                    )
                    for bc in self.bcs:
                        bc.set(self.b.array_w)
            else:
                dolfinx.la.petsc._ghost_update(self.b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore[attr-defined]

            # Solve linear system and update ghost values in the solution
            self.solver.solve(self.b, self.x)
            dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)  # type: ignore[attr-defined]
            dolfinx.fem.petsc.assign(self.x, self.u)
            return self.u

        @property
        def L(self) -> Form | Sequence[Form]:
            """The compiled linear form representing the left-hand side."""
            return self._L

        @property
        def a(self) -> Form | Sequence[Form]:
            """The compiled bilinear form representing the right-hand side."""
            return self._a

        @property
        def preconditioner(self) -> Form | Sequence[Form]:
            """The compiled bilinear form representing the preconditioner."""
            return self._preconditioner

        @property
        def A(self) -> PETSc.Mat:  # type: ignore[name-defined]
            """Left-hand side matrix."""
            return self._A

        @property
        def P_mat(self) -> PETSc.Mat:  # type: ignore[name-defined]
            """Preconditioner matrix."""
            return self._P_mat

        @property
        def b(self) -> PETSc.Vec:  # type: ignore[name-defined]
            """Right-hand side vector."""
            return self._b

        @property
        def x(self) -> PETSc.Vec:  # type: ignore[name-defined]
            """Solution vector.

            Note:
                The vector does not share memory with the solution
                function(s) ``u``.
            """
            return self._x

        @property
        def solver(self) -> PETSc.KSP:  # type: ignore[name-defined]
            """The PETSc KSP solver."""
            return self._solver

        @property
        def u(self) -> _Function | Sequence[_Function]:
            """Solution function(s).

            Note:
                The function(s) do not share memory with the solution
                vector ``x``.
            """
            return self._u

    # -- High-level interface for SNES ---------------------------------------

    def _assign_block_data(forms: Sequence[Form], vec: PETSc.Vec):  # type: ignore[name-defined]
        """Assign block data to a PETSc vector.

        Args:
            forms: List of forms to extract block data from.
            vec: PETSc vector to assign block data to.
        """

        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
        )

        return dolfinx.la.petsc._assign_block_data(maps, vec)

    def assemble_residual(
        u: _Function | Sequence[_Function],
        residual: Form | Sequence[Form],
        jacobian: Form | Sequence[Sequence[Form]],
        bcs: Sequence[DirichletBC],
        _snes: PETSc.SNES,  # type: ignore[name-defined]
        x: PETSc.Vec,  # type: ignore[name-defined]
        b: PETSc.Vec,  # type: ignore[name-defined]
    ):
        """Assemble the residual at ``x`` into the vector ``b``.

        A function conforming to the interface expected by ``SNES.setFunction``
        can be created by fixing the first four arguments, e.g.:

        Example::

            snes = PETSc.SNES().create(mesh.comm)
            assemble_residual = functools.partial(
                dolfinx.fem.petsc.assemble_residual,
                u, residual, jacobian, bcs)
            snes.setFunction(assemble_residual, x, b)

        Args:
            u: Function(s) tied to the solution vector within the residual and
               Jacobian.
            residual: Form of the residual. It can be a sequence of forms.
            jacobian: Form of the Jacobian. It can be a nested sequence of
                forms.
            bcs: List of Dirichlet boundary conditions to lift the residual.
            _snes: The solver instance.
            x: The vector containing the point to evaluate the residual at.
            b: Vector to assemble the residual into.
        """
        # Update input vector before assigning
        dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)  # type: ignore[attr-defined]

        # Assign the input vector to the unknowns
        assign(x, u)

        # Assign block data if block assembly is requested
        if isinstance(residual, Sequence) and b.getType() != PETSc.Vec.Type.NEST:  # type: ignore[attr-defined]
            _assign_block_data(residual, b)
            _assign_block_data(residual, x)

        # Assemble the residual
        dolfinx.la.petsc._zero_vector(b)
        assemble_vector(b, residual)  # type: ignore[arg-type]

        # Lift vector
        if isinstance(jacobian, Sequence):
            # Nest and blocked lifting
            bcs1 = _bcs_by_block(_extract_function_spaces(jacobian, 1), bcs)  # type: ignore[arg-type]
            apply_lifting(b, jacobian, bcs=bcs1, x0=x, alpha=-1.0)
            dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore[attr-defined]
            bcs0 = _bcs_by_block(_extract_function_spaces(residual), bcs)  # type: ignore[arg-type]
            set_bc(b, bcs0, x0=x, alpha=-1.0)
        else:
            # Single form lifting
            apply_lifting(b, [jacobian], bcs=[bcs], x0=[x], alpha=-1.0)
            dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore[attr-defined]
            set_bc(b, bcs, x0=x, alpha=-1.0)
        dolfinx.la.petsc._ghost_update(b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)  # type: ignore[attr-defined]

    def assemble_jacobian(
        u: Sequence[_Function] | _Function,
        jacobian: Form | Sequence[Sequence[Form]],
        preconditioner: Form | Sequence[Sequence[Form]] | None,
        bcs: Sequence[DirichletBC],
        _snes: PETSc.SNES,  # type: ignore[name-defined]
        x: PETSc.Vec,  # type: ignore[name-defined]
        J: PETSc.Mat,  # type: ignore[name-defined]
        P_mat: PETSc.Mat,  # type: ignore[name-defined]
    ):
        """Assemble the Jacobian and preconditioner matrices at ``x``
        into ``J`` and ``P_mat``.

        A function conforming to the interface expected by ``SNES.setJacobian``
        can be created by fixing the first four arguments e.g.:

        Example::

            snes = PETSc.SNES().create(mesh.comm)
            assemble_jacobian = functools.partial(
                dolfinx.fem.petsc.assemble_jacobian,
                u, jacobian, preconditioner, bcs)
            snes.setJacobian(assemble_jacobian, A, P_mat)

        Args:
            u: Function tied to the solution vector within the residual and
                jacobian.
            jacobian: Compiled form of the Jacobian.
            preconditioner: Compiled form of the preconditioner.
            bcs: List of Dirichlet boundary conditions to apply to the Jacobian
                 and preconditioner matrices.
            _snes: The solver instance.
            x: The vector containing the point to evaluate at.
            J: Matrix to assemble the Jacobian into.
            P_mat: Matrix to assemble the preconditioner into.
        """
        # Copy existing soultion into the function used in the residual and
        # Jacobian
        dolfinx.la.petsc._ghost_update(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)  # type: ignore[attr-defined]
        assign(x, u)

        # Assemble Jacobian
        J.zeroEntries()
        assemble_matrix(J, jacobian, bcs, diag=1.0)  # type: ignore[arg-type, misc]
        J.assemble()
        if preconditioner is not None:
            P_mat.zeroEntries()
            assemble_matrix(P_mat, preconditioner, bcs, diag=1.0)  # type: ignore[arg-type, misc]
            P_mat.assemble()

    class NonlinearProblem:
        """High-level class for solving nonlinear variational problems
        with PETSc SNES.

        Solves problems of the form
        :math:`F_i(u, v) = 0, i=0,\\ldots,N\\ \\forall v \\in V` where
        :math:`u=(u_0,\\ldots,u_N), v=(v_0,\\ldots,v_N)` using PETSc
        SNES as the non-linear solver.

        Note:
            The deprecated version of this class for use with
            :class:`dolfinx.nls.petsc.NewtonSolver` has been renamed
            :class:`dolfinx.fem.petsc.NewtonSolverNonlinearProblem`.

        Note:
            This high-level class automatically handles PETSc memory
            management. The user does not need to manually call
            ``.destroy()`` on returned PETSc objects.
        """

        def __init__(
            self,
            F: ufl.form.Form | Sequence[ufl.form.Form],
            u: _Function | Sequence[_Function],
            *,
            petsc_options_prefix: str,
            bcs: Sequence[DirichletBC] | None = None,
            J: ufl.form.Form | Sequence[Sequence[ufl.form.Form]] | None = None,
            P: ufl.form.Form | Sequence[Sequence[ufl.form.Form]] | None = None,
            kind: str | Sequence[Sequence[str]] | None = None,
            petsc_options: dict | None = None,
            form_compiler_options: dict | None = None,
            jit_options: dict | None = None,
            entity_maps: Sequence[_EntityMap] | None = None,
        ):
            """
            Initialize solver for a nonlinear variational problem.

            By default, the underlying SNES solver uses PETSc's default
            options. To use the robust combination of LU via MUMPS with
            a backtracking linesearch, pass:

            Example::

                petsc_options = {"ksp_type": "preonly",
                                 "pc_type": "lu",
                                 "pc_factor_mat_solver_type": "mumps",
                                 "snes_linesearch_type": "bt",
                }

            Every PETSc object will have a unique options prefix set. We
            recommend discovering these prefixes dynamically via the
            petsc4py API rather than hard-coding each prefix value into
            the programme.

            Example::

                snes_options_prefix = problem.solver.getOptionsPrefix()
                jacobian_options_prefix = problem.A.getOptionsPrefix()

            Args:
                F: UFL form(s) representing the residual :math:`F_i`.
                u: Function(s) used to define the residual and Jacobian.
                bcs: Dirichlet boundary conditions.
                J: UFL form(s) representing the Jacobian
                    :math:`J_{ij} = dF_i/du_j`. If not passed, derived
                    automatically.
                P: UFL form(s) representing the preconditioner.
                kind: The PETSc matrix and vector kind. Common choices
                    are ``mpi`` and ``nest``. See
                    :func:`dolfinx.fem.petsc.create_matrix` and
                    :func:`dolfinx.fem.petsc.create_vector` for more
                    information.
                petsc_options_prefix: Mandatory named argument.
                    Options prefix used as root prefix on all
                    internally created PETSc objects. Typically ends with `_`.
                    Must be the same on all ranks, and is usually unique within
                    the programme.
                petsc_options: Options set on the underlying PETSc SNES only.
                    The options must be the same on all ranks. For available
                    choices for ``petsc_options``, see the
                    `PETSc SNES documentation
                    <https://petsc4py.readthedocs.io/en/stable/manual/snes/>`_.
                    Options on other objects (matrices, vectors) should be set
                    explicitly by the user.
                form_compiler_options: Options used in FFCx compilation of all
                    forms. Run ``ffcx --help`` at the command line to see all
                    available options.
                jit_options: Options used in CFFI JIT compilation of C code
                    generated by FFCx. See ``python/dolfinx/jit.py`` for all
                    available options. Takes priority over all other option
                    values.
                entity_maps: If any trial functions, test functions, or
                    coefficients in the form are not defined over the same mesh
                    as the integration domain, a corresponding :class:
                    `EntityMap<dolfinx.mesh.EntityMap>` must be provided.
            """
            # Compile residual and Jacobian forms
            self._F = _create_form(
                F,
                form_compiler_options=form_compiler_options,
                jit_options=jit_options,
                entity_maps=entity_maps,
            )

            if J is None:
                J = derivative_block(F, u)

            self._J = _create_form(
                J,
                form_compiler_options=form_compiler_options,
                jit_options=jit_options,
                entity_maps=entity_maps,
            )

            if P is not None:
                self._preconditioner = _create_form(
                    P,
                    form_compiler_options=form_compiler_options,
                    jit_options=jit_options,
                    entity_maps=entity_maps,
                )
            else:
                self._preconditioner = None

            self._u = u
            # Set default values if not supplied
            bcs = [] if bcs is None else bcs

            # Create PETSc structures for the residual, Jacobian and solution
            # vector
            self._A = create_matrix(self.J, kind=kind)
            # Create PETSc structure for preconditioner if provided
            if self._preconditioner is not None:
                self._P_mat = create_matrix(self._preconditioner, kind=kind)
            else:
                self._P_mat = None

            # 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(_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
            self._snes = PETSc.SNES().create(self.A.comm)  # type: ignore[attr-defined]
            self.solver.setJacobian(
                partial(assemble_jacobian, u, self.J, self.preconditioner, bcs), self.A, self.P_mat
            )
            self.solver.setFunction(partial(assemble_residual, u, self.F, self.J, bcs), self.b)

            if petsc_options_prefix == "":
                raise ValueError("PETSc options prefix cannot be empty.")

            self.solver.setOptionsPrefix(petsc_options_prefix)
            self.A.setOptionsPrefix(f"{petsc_options_prefix}A_")
            if self.P_mat is not None:
                self.P_mat.setOptionsPrefix(f"{petsc_options_prefix}P_mat_")
            self.b.setOptionsPrefix(f"{petsc_options_prefix}b_")
            self.x.setOptionsPrefix(f"{petsc_options_prefix}x_")

            # Set options for SNES only
            if petsc_options is not None:
                opts = PETSc.Options()  # type: ignore[attr-defined]
                opts.prefixPush(self.solver.getOptionsPrefix())

                for k, v in petsc_options.items():
                    opts[k] = v

                self.solver.setFromOptions()

                # Tidy up global options
                for k in petsc_options.keys():
                    del opts[k]

                opts.prefixPop()

            if self.P_mat is not None and kind == "nest":
                # Transfer nest IS on self.P_mat to PC of main KSP. This allows
                # fieldsplit preconditioning to be applied, if desired.
                nest_IS = self.P_mat.getNestISs()
                fieldsplit_IS = tuple(
                    [
                        (f"{u.name + '_' if u.name != 'f' else ''}{i}", IS)
                        for i, (u, IS) in enumerate(zip(self.u, nest_IS[0]))
                    ]
                )
                self.solver.getKSP().getPC().setFieldSplitIS(*fieldsplit_IS)

        def solve(self) -> _Function | Sequence[_Function]:
            """Solve the problem.

            This method updates the solution ``u`` function(s) stored in the
            problem instance.

            Note:
                The user is responsible for asserting convergence of the SNES
                solver e.g. ``assert problem.solver.getConvergedReason() > 0``.
                Alternatively, pass ``"snes_error_if_not_converged": True`` and
                ``"ksp_error_if_not_converged" : True`` in ``petsc_options`` to
                raise a ``PETScError`` on failure.

            Returns:
                The solution function(s).
            """
            # Copy current iterate into the work array.
            assign(self.u, self.x)

            # Solve problem
            self.solver.solve(None, self.x)
            dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)  # type: ignore[attr-defined]

            # Copy solution back to function
            assign(self.x, self.u)

            return self.u

        def __del__(self):
            for obj in filter(
                lambda obj: obj is not None, (self._snes, self._A, self._b, self._x, self._P_mat)
            ):
                obj.destroy()

        @property
        def F(self) -> Form | Sequence[Form]:
            """The compiled residual."""
            return self._F

        @property
        def J(self) -> Form | Sequence[Sequence[Form]]:
            """The compiled Jacobian."""
            return self._J

        @property
        def preconditioner(self) -> Form | Sequence[Sequence[Form]] | None:
            """The compiled preconditioner."""
            return self._preconditioner

        @property
        def A(self) -> PETSc.Mat:  # type: ignore[name-defined]
            """Jacobian matrix."""
            return self._A

        @property
        def P_mat(self) -> PETSc.Mat | None:  # type: ignore[name-defined]
            """Preconditioner matrix."""
            return self._P_mat

        @property
        def b(self) -> PETSc.Vec:  # type: ignore[name-defined]
            """Residual vector."""
            return self._b

        @property
        def x(self) -> PETSc.Vec:  # type: ignore[name-defined]
            """Solution vector.

            Note:
                The vector does not share memory with the
                solution function(s) ``u``.
            """
            return self._x

        @property
        def solver(self) -> PETSc.SNES:  # type: ignore[name-defined]
            """The SNES solver."""
            return self._snes

        @property
        def u(self) -> _Function | Sequence[_Function]:
            """Solution function(s).

            Note:
                The function(s) do not share memory with the solution
                vector ``x``.
            """
            return self._u

    # -- Deprecated non-linear problem class for NewtonSolver -----------------

    class NewtonSolverNonlinearProblem:
        """(Deprecated) Nonlinear problem class for solving nonlinear
        problems using :class:`dolfinx.nls.petsc.NewtonSolver`.

        Solves problems of the form :math:`F(u, v) = 0 \\ \\forall v \\in V`
        using PETSc as the linear algebra backend.

        Note:
            This class is deprecated in favour of
            :class:`dolfinx.fem.petsc.NonlinearProblem`, a high level
            interface to SNES.

        Note:
            This class was previously called
            ``dolfinx.fem.petsc.NonlinearProblem``.
        """

        def __init__(
            self,
            F: ufl.form.Form,
            u: _Function,
            bcs: Sequence[DirichletBC] | None = None,
            J: ufl.form.Form = None,
            form_compiler_options: dict | None = None,
            jit_options: dict | None = None,
        ):
            """Initialize solver for solving a non-linear problem using
            Newton's method`.

            Args:
                F: The PDE residual F(u, v).
                u: The unknown.
                bcs: List of Dirichlet boundary conditions.
                J: UFL representation of the Jacobian (optional)
                form_compiler_options: Options used in FFCx
                    compilation of this form. Run ``ffcx --help`` at the
                    command line to see all available options.
                jit_options: Options used in CFFI JIT compilation of C
                    code generated by FFCx. See ``python/dolfinx/jit.py``
                    for all available options. Takes priority over all other
                    option values.

            Example::

                problem = NonlinearProblem(F, u, [bc0, bc1])
            """
            warnings.warn(
                (
                    "dolfinx.nls.petsc.NewtonSolver is deprecated. "
                    + "Use dolfinx.fem.petsc.NonlinearProblem, "
                    + "a high level interface to PETSc SNES, instead."
                ),
                DeprecationWarning,
            )

            self._L = _create_form(
                F, form_compiler_options=form_compiler_options, jit_options=jit_options
            )

            # Create the Jacobian matrix, dF/du
            if J is None:
                V = u.function_space
                du = ufl.TrialFunction(V)
                J = ufl.derivative(F, u, du)

            self._a = _create_form(
                J, form_compiler_options=form_compiler_options, jit_options=jit_options
            )
            self.bcs = bcs

        @property
        def L(self) -> Form:
            """The compiled linear form (the residual form)."""
            return self._L

        @property
        def a(self) -> Form:
            """The compiled bilinear form (the Jacobian form)."""
            return self._a

        def form(self, x: PETSc.Vec) -> None:  # type: ignore[name-defined]
            """This function is called before the residual or Jacobian is
            computed. This is usually used to update ghost values.

            Args:
               x: The vector containing the latest solution
            """
            x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)  # type: ignore[attr-defined]

        def F(self, x: PETSc.Vec, b: PETSc.Vec) -> None:  # type: ignore[name-defined]
            """Assemble the residual F into the vector b.

            Args:
                x: The vector containing the latest solution
                b: Vector to assemble the residual into
            """
            # Reset the residual vector
            dolfinx.la.petsc._zero_vector(b)
            assemble_vector(b, self._L)

            # Apply boundary condition
            if self.bcs is not None:
                apply_lifting(b, [self._a], bcs=[self.bcs], x0=[x], alpha=-1.0)
                b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)  # type: ignore[attr-defined]
                set_bc(b, self.bcs, x, -1.0)
            else:
                b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)  # type: ignore[attr-defined]

        def J(self, x: PETSc.Vec, A: PETSc.Mat) -> None:  # type: ignore[name-defined]
            """Assemble the Jacobian matrix.

            Args:
                x: The vector containing the latest solution
            """
            A.zeroEntries()
            assemble_matrix(A, self._a, self.bcs)  # type: ignore[arg-type]
            A.assemble()

    # -- Additional free helper functions (interpolations, assignments etc.) --

    def discrete_curl(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat:  # type: ignore[name-defined]
        """Assemble a discrete curl operator.

        Args:
            space0: H1 space to interpolate the gradient from.
            space1: H(curl) space to interpolate into.

        Returns:
            Discrete curl operator.
        """
        return _discrete_curl(space0._cpp_object, space1._cpp_object)

    def discrete_gradient(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat:  # type: ignore[name-defined]
        """Assemble a discrete gradient operator.

        The discrete gradient operator interpolates the gradient of a H1
        finite element function into a H(curl) space. It is assumed that the
        H1 space uses an identity map and the H(curl) space uses a covariant
        Piola map.

        Args:
            space0: H1 space to interpolate the gradient from.
            space1: H(curl) space to interpolate into.

        Returns:
            Discrete gradient operator.
        """
        return _discrete_gradient(space0._cpp_object, space1._cpp_object)

    def interpolation_matrix(V0: _FunctionSpace, V1: _FunctionSpace) -> PETSc.Mat:  # type: ignore[name-defined]
        r"""Assemble an interpolation operator matrix for discreye
        interpolation between finite element spaces.

        Consider is the vector of degrees-of-freedom  :math:`u_{i}`
        associated with a function in :math:`V_{i}`. This function returns
        the matrix :math:`\Pi` sucht that

        .. math::

            u_{1} = \Pi u_{0}.

        Args:
            V0: Space to interpolate from.
            V1: Space to interpolate into.

        Returns:
            The interpolation matrix :math:`\Pi`.

        Note:
            The returned matrix is not finalised, i.e. ghost values are not
            accumulated.
        """
        return _interpolation_matrix(V0._cpp_object, V1._cpp_object)

    @functools.singledispatch
    def assign(u: _Function | Sequence[_Function], x: PETSc.Vec):  # type: ignore[name-defined]
        """Assign :class:`Function` degrees-of-freedom to a vector.

        Assigns degree-of-freedom values in ``u``, which is possibly a
        sequence of ``Function``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:  # type: ignore[attr-defined]
            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
    def _(x: PETSc.Vec, u: _Function | Sequence[_Function]):  # type: ignore[name-defined]
        """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:  # type: ignore[attr-defined]
            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)

    def get_petsc_lib() -> pathlib.Path:
        """Find the full path of the PETSc shared library.

        Returns:
            Full path to the PETSc shared library.

        Raises:
            RuntimeError: If PETSc library cannot be found for if more than
                one library is found.
        """
        import petsc4py as _petsc4py

        petsc_dir = _petsc4py.get_config()["PETSC_DIR"]
        petsc_arch = _petsc4py.lib.getPathArchPETSc()[1]
        try:
            if np.issubdtype(_petsc4py.PETSc.ScalarType, np.complexfloating):
                scalar_type = "complex"
            else:
                scalar_type = "real"
        except AttributeError:
            # if petsc4py.PETSc is not available, read type from petsc_dir
            scalar_type = "complex" if "complex" in petsc_dir else "real"

        candidate_paths = [
            os.path.join(petsc_dir, petsc_arch, "lib", f"libpetsc_{scalar_type}.so"),
            os.path.join(petsc_dir, petsc_arch, "lib", "libpetsc.so"),
            os.path.join(petsc_dir, petsc_arch, "lib", "libpetsc.dylib"),
        ]
        exists_paths = []
        for candidate_path in candidate_paths:
            if os.path.exists(candidate_path):
                exists_paths.append(candidate_path)

        if len(exists_paths) == 0:
            raise RuntimeError(
                f"Could not find a PETSc shared library. Candidate paths: {candidate_paths}"
            )
        elif len(exists_paths) > 1:
            raise RuntimeError(f"More than one PETSc shared library found. Paths: {exists_paths}")

        return pathlib.Path(exists_paths[0])

    class numba_utils:
        """Utility attributes for working with Numba and PETSc.

        These attributes are convenience functions for calling PETSc C
        functions from within Numba functions.

        Note:
            `Numba <https://numba.pydata.org/>`_ must be available
            to use these utilities.

        Examples:
            A typical use of these utility functions is::

                import numpy as np
                import numpy.typing as npt
                def set_vals(A: int,
                             m: int, rows: npt.NDArray[PETSc.IntType],
                             n: int, cols: npt.NDArray[PETSc.IntType],
                             data: npt.NDArray[PETSc.ScalarTYpe], mode: int):
                    MatSetValuesLocal(A, m, rows.ctypes, n, cols.ctypes,
                                      data.ctypes, mode)
        """

        try:
            import petsc4py.PETSc as _PETSc

            import llvmlite as _llvmlite
            import numba as _numba

            _llvmlite.binding.load_library_permanently(str(get_petsc_lib()))

            _int = _numba.from_dtype(_PETSc.IntType)  # type: ignore
            _scalar = _numba.from_dtype(_PETSc.ScalarType)  # type: ignore
            _real = _numba.from_dtype(_PETSc.RealType)  # type: ignore
            _int_ptr = _numba.core.types.CPointer(_int)
            _scalar_ptr = _numba.core.types.CPointer(_scalar)
            _MatSetValues_sig = _numba.core.typing.signature(
                _numba.core.types.intc,
                _numba.core.types.uintp,
                _int,
                _int_ptr,
                _int,
                _int_ptr,
                _scalar_ptr,
                _numba.core.types.intc,
            )
            MatSetValuesLocal = _numba.core.types.ExternalFunction(
                "MatSetValuesLocal", _MatSetValues_sig
            )
            """See PETSc `MatSetValuesLocal
            <https://petsc.org/release/manualpages/Mat/MatSetValuesLocal>`_
            documentation."""

            MatSetValuesBlockedLocal = _numba.core.types.ExternalFunction(
                "MatSetValuesBlockedLocal", _MatSetValues_sig
            )
            """See PETSc `MatSetValuesBlockedLocal
            <https://petsc.org/release/manualpages/Mat/MatSetValuesBlockedLocal>`_
            documentation."""
        except ImportError:
            pass

    class ctypes_utils:
        """Utility attributes for working with ctypes and PETSc.

        These attributes are convenience functions for calling PETSc C
        functions, typically from within Numba functions.

        Examples:
            A typical use of these utility functions is::

                import numpy as np
                import numpy.typing as npt
                def set_vals(A: int,
                             m: int, rows: npt.NDArray[PETSc.IntType],
                             n: int, cols: npt.NDArray[PETSc.IntType],
                             data: npt.NDArray[PETSc.ScalarTYpe], mode: int):
                    MatSetValuesLocal(A, m, rows.ctypes, n, cols.ctypes,
                                      data.ctypes, mode)
        """

        try:
            import petsc4py.PETSc as _PETSc

            _lib_ctypes = _ctypes.cdll.LoadLibrary(str(get_petsc_lib()))

            # Note: ctypes does not have complex types, hence we use void* for
            # scalar data
            _int = np.ctypeslib.as_ctypes_type(_PETSc.IntType)  # type: ignore

            MatSetValuesLocal = _lib_ctypes.MatSetValuesLocal
            """See PETSc `MatSetValuesLocal
            <https://petsc.org/release/manualpages/Mat/MatSetValuesLocal>`_
            documentation."""
            MatSetValuesLocal.argtypes = [
                _ctypes.c_void_p,
                _int,
                _ctypes.POINTER(_int),
                _int,
                _ctypes.POINTER(_int),
                _ctypes.c_void_p,
                _ctypes.c_int,
            ]

            MatSetValuesBlockedLocal = _lib_ctypes.MatSetValuesBlockedLocal
            """See PETSc `MatSetValuesBlockedLocal
            <https://petsc.org/release/manualpages/Mat/MatSetValuesBlockedLocal>`_
            documentation."""
            MatSetValuesBlockedLocal.argtypes = [
                _ctypes.c_void_p,
                _int,
                _ctypes.POINTER(_int),
                _int,
                _ctypes.POINTER(_int),
                _ctypes.c_void_p,
                _ctypes.c_int,
            ]
        except ImportError:
            pass

    class cffi_utils:
        """Utility attributes for working with CFFI (ABI mode) and Numba.

        Registers Numba's complex types with CFFI.

        If PETSc is available, CFFI convenience functions for calling PETSc C
        functions are also created. These are typically called from within
        Numba functions.

        Note:
            `CFFI <https://cffi.readthedocs.io/>`_ and  `Numba
            <https://numba.pydata.org/>`_ must be available to use these
            utilities.

        Examples:
            A typical use of these utility functions is::

                import numpy as np
                import numpy.typing as npt
                def set_vals(A: int,
                             m: int, rows: npt.NDArray[PETSc.IntType],
                             n: int, cols: npt.NDArray[PETSc.IntType],
                             data: npt.NDArray[PETSc.ScalarType], mode: int):
                    MatSetValuesLocal(A, m, ffi.from_buffer(rows), n,
                                      ffi.from_buffer(cols),
                                      ffi.from_buffer(rows(data), mode)
        """

        import cffi as _cffi

        _ffi = _cffi.FFI()

        try:
            import numba as _numba
            import numba.core.typing.cffi_utils as _cffi_support

            # Register complex types
            _cffi_support.register_type(_ffi.typeof("float _Complex"), _numba.types.complex64)
            _cffi_support.register_type(_ffi.typeof("double _Complex"), _numba.types.complex128)

        except KeyError:
            pass
        except ImportError:
            from dolfinx.log import LogLevel, log

            log(
                LogLevel.DEBUG,
                "Could not import numba, so cffi/numba complex types were not registered.",
            )

        try:
            from petsc4py import PETSc as _PETSc

            _lib_cffi = _ffi.dlopen(str(get_petsc_lib()))

            _CTYPES = {
                np.int32: "int32_t",
                np.int64: "int64_t",
                np.float32: "float",
                np.float64: "double",
                np.complex64: "float _Complex",
                np.complex128: "double _Complex",
                np.longlong: "long long",
            }

            _c_int_t = _CTYPES[_PETSc.IntType]  # type: ignore
            _c_scalar_t = _CTYPES[_PETSc.ScalarType]  # type: ignore
            _ffi.cdef(
                f"""
                    int MatSetValuesLocal(void* mat, {_c_int_t} nrow, const {_c_int_t}* irow,
                                        {_c_int_t} ncol, const {_c_int_t}* icol,
                                        const {_c_scalar_t}* y, int addv);
                    int MatSetValuesBlockedLocal(void* mat, {_c_int_t} nrow, const {_c_int_t}* irow,
                                        {_c_int_t} ncol, const {_c_int_t}* icol,
                                        const {_c_scalar_t}* y, int addv);
                                        """
            )

            MatSetValuesLocal = _lib_cffi.MatSetValuesLocal  # type: ignore[attr-defined]
            """See PETSc `MatSetValuesLocal
            <https://petsc.org/release/manualpages/Mat/MatSetValuesLocal>`_
            documentation."""

            MatSetValuesBlockedLocal = _lib_cffi.MatSetValuesBlockedLocal  # type: ignore[attr-defined]
            """See PETSc `MatSetValuesBlockedLocal
            <https://petsc.org/release/manualpages/Mat/MatSetValuesBlockedLocal>`_
            documentation."""
        except KeyError:
            pass
        except ImportError:
            from dolfinx.log import LogLevel, log

            log(
                LogLevel.DEBUG,
                "Could not import petsc4py, so cffi/PETSc ABI mode interface was not created.",
            )