# 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
"""Assembly functions into PETSc objects for variational forms.
Functions in this module generally apply functions in :mod:`dolfinx.fem`
to PETSc linear algebra objects and handle any PETSc-specific
preparation.
Note:
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.
"""
# mypy: ignore-errors
from __future__ import annotations
import contextlib
import functools
import itertools
import typing
from collections.abc import Iterable, 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_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 Mesh as _Mesh
__all__ = [
"LinearProblem",
"NonlinearProblem",
"apply_lifting",
"assemble_jacobian",
"assemble_matrix",
"assemble_residual",
"assemble_vector",
"assign",
"create_matrix",
"create_vector",
"discrete_curl",
"discrete_gradient",
"interpolation_matrix",
"set_bc",
]
def _extract_function_spaces(
a: Iterable[Iterable[Form]],
) -> tuple[Sequence[_FunctionSpace], Sequence[_FunctionSpace]]:
"""From a rectangular array of bilinear forms, extract the function
spaces for each block row and block column.
"""
assert len({len(cols) for cols in a}) == 1, "Array of function spaces is not rectangular"
# Extract (V0, V1) pair for each block in 'a'
def fn(form):
return form.function_spaces if form is not None else None
from functools import partial
Vblock: typing.Iterable = map(partial(map, fn), a)
# Compute spaces for each row/column block
rows: Sequence[set] = [set() for i in range(len(a))]
cols: Sequence[set] = [set() for i in range(len(a[0]))]
for i, Vrow in enumerate(Vblock):
for j, V in enumerate(Vrow):
if V is not None:
rows[i].add(V[0])
cols[j].add(V[1])
rows = [e for row in rows for e in row]
cols = [e for col in cols for e in col]
assert len(rows) == len(a)
assert len(cols) == len(a[0])
return rows, cols
# -- Vector instantiation -------------------------------------------------
[docs]
def create_vector(
L: typing.Union[Form, Iterable[Form]], kind: typing.Optional[str] = None
) -> PETSc.Vec:
"""Create a PETSc vector that is compatible with a linear form(s).
Three cases are supported:
1. For a single linear form ``L``, if ``kind`` is ``None`` or is
``PETSc.Vec.Type.MPI``, a ghosted PETSc vector which is
compatible with ``L`` is created.
2. If ``L`` is a sequence of linear forms and ``kind`` is ``None``
or is ``PETSc.Vec.Type.MPI``, a ghosted PETSc vector which is
compatible with ``L`` is created. The created vector ``b`` is
initialized such that on each MPI process ``b = [b_0, b_1, ...,
b_n, b_0g, b_1g, ..., b_ng]``, where ``b_i`` are the entries
associated with the 'owned' degrees-of-freedom for ``L[i]`` and
``b_ig`` are the 'unowned' (ghost) entries for ``L[i]``.
For this case, the returned vector has an attribute ``_blocks``
that holds the local offsets into ``b`` for the (i) owned and
(ii) ghost entries for each ``L[i]``. It can be accessed by
``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`` is a sequence of linear forms and ``kind`` is
``PETSc.Vec.Type.NEST``, a PETSc nested vector (a 'nest' of
ghosted PETSc vectors) which is compatible with ``L`` is created.
Args:
L: Linear form or a sequence of linear forms.
kind: PETSc vector type (``VecType``) to create.
Returns:
A PETSc vector with a layout that is compatible with ``L``. The
vector is not initialised to zero.
"""
try:
dofmap = L.function_spaces[0].dofmaps(0) # Single form case
return dolfinx.la.petsc.create_vector(dofmap.index_map, dofmap.index_map_bs)
except AttributeError:
maps = [
(
form.function_spaces[0].dofmaps(0).index_map,
form.function_spaces[0].dofmaps(0).index_map_bs,
)
for form in L
]
if kind == PETSc.Vec.Type.NEST:
return _cpp.fem.petsc.create_vector_nest(maps)
elif kind == PETSc.Vec.Type.MPI:
off_owned = tuple(
itertools.accumulate(maps, lambda off, m: off + m[0].size_local * m[1], initial=0)
)
off_ghost = tuple(
itertools.accumulate(
maps, lambda off, m: off + m[0].num_ghosts * m[1], initial=off_owned[-1]
)
)
b = _cpp.fem.petsc.create_vector_block(maps)
b.setAttr("_blocks", (off_owned, off_ghost))
return b
else:
raise NotImplementedError(
"Vector type must be specified for blocked/nested assembly."
f"Vector type '{kind}' not supported."
"Did you mean 'nest' or 'mpi'?"
)
# -- Matrix instantiation -------------------------------------------------
[docs]
def create_matrix(
a: typing.Union[Form, Iterable[Iterable[Form]]],
kind: typing.Optional[typing.Union[str, Iterable[Iterable[str]]]] = None,
) -> PETSc.Mat:
"""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.
"""
try:
return _cpp.fem.petsc.create_matrix(a._cpp_object, kind) # Single form
except AttributeError: # ``a`` is a nested 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: # Create nest matrix with default types
return _cpp.fem.petsc.create_matrix_nest(_a, None)
else:
try:
return _cpp.fem.petsc.create_matrix_block(_a, kind) # Single 'kind' type
except TypeError:
return _cpp.fem.petsc.create_matrix_nest(_a, kind) # Array of 'kind' types
# -- Vector assembly ------------------------------------------------------
[docs]
@functools.singledispatch
def assemble_vector(
L: typing.Union[Form, Iterable[Form]],
constants: typing.Optional[npt.NDArray, Iterable[npt.NDArray]] = None,
coeffs: typing.Optional[
typing.Union[
dict[tuple[IntegralType, int], npt.NDArray],
Iterable[dict[tuple[IntegralType, int], npt.NDArray]],
]
] = None,
kind: typing.Optional[str] = None,
) -> PETSc.Vec:
"""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(L, kind=kind)
if kind == PETSc.Vec.Type.NEST:
for b_sub in b.getNestSubVecs():
with b_sub.localForm() as b_local:
b_local.set(0.0)
else:
with b.localForm() as b_local:
b_local.set(0)
return assemble_vector(b, L, constants, coeffs)
@assemble_vector.register(PETSc.Vec)
def _assemble_vector_vec(
b: PETSc.Vec,
L: typing.Union[Form, Iterable[Form]],
constants: typing.Optional[npt.NDArray, Iterable[npt.NDArray]] = None,
coeffs: typing.Optional[
typing.Union[
dict[tuple[IntegralType, int], npt.NDArray],
Iterable[dict[tuple[IntegralType, int], npt.NDArray]],
]
] = None,
) -> PETSc.Vec:
"""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:
constants = [None] * len(L) if constants is None else constants
coeffs = [None] * len(L) if coeffs is None else coeffs
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)
elif isinstance(L, Iterable):
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)
_assemble_vector_array(bx_, L_, const, coeff)
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)
return b
# -- Matrix assembly ------------------------------------------------------
[docs]
@functools.singledispatch
def assemble_matrix(
a: typing.Union[Form, Iterable[Iterable[Form]]],
bcs: typing.Optional[Iterable[DirichletBC]] = None,
diag: float = 1,
constants: typing.Optional[
typing.Union[Iterable[np.ndarray], Iterable[Iterable[np.ndarray]]]
] = None,
coeffs: typing.Optional[
typing.Union[
dict[tuple[IntegralType, int], npt.NDArray],
Iterable[dict[tuple[IntegralType, int], npt.NDArray]],
]
] = 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.
"""
try:
A = _cpp.fem.petsc.create_matrix(a._cpp_object, kind)
except AttributeError:
A = create_matrix(a, kind)
assemble_matrix(A, a, bcs, diag, constants, coeffs)
return A
def _assemble_matrix_block_mat(
A: PETSc.Mat,
a: Iterable[Iterable[Form]],
bcs: typing.Optional[Iterable[DirichletBC]],
diag: float,
constants: typing.Optional[Iterable[npt.NDArray]] = None,
coeffs: typing.Optional[Iterable[Iterable[dict[tuple[IntegralType, int], npt.NDArray]]]] = None,
) -> PETSc.Mat:
"""Assemble bilinear forms into a blocked matrix."""
consts = [pack_constants(forms) for forms in a] if constants is None else constants
coeffs = [pack_coefficients(forms) for forms in a] if coeffs is None else coeffs
V = _extract_function_spaces(a)
is0 = _cpp.la.petsc.create_index_sets(
[(Vsub.dofmaps(0).index_map, Vsub.dofmaps(0).index_map_bs) for Vsub in V[0]]
)
is1 = _cpp.la.petsc.create_index_sets(
[(Vsub.dofmaps(0).index_map, Vsub.dofmaps(0).index_map_bs) for Vsub in V[1]]
)
_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], _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)
# 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)
return A
@assemble_matrix.register
def assemble_matrix_mat(
A: PETSc.Mat,
a: typing.Union[Form, Iterable[Iterable[Form]]],
bcs: typing.Optional[Iterable[DirichletBC]] = None,
diag: float = 1,
constants: typing.Optional[typing.Union[np.ndarray, Iterable[Iterable[np.ndarray]]]] = None,
coeffs: typing.Optional[
typing.Union[
dict[tuple[IntegralType, int], npt.NDArray],
Iterable[Iterable[dict[tuple[IntegralType, int], npt.NDArray]]],
]
] = None,
) -> PETSc.Mat:
"""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:
constants = [pack_constants(forms) for forms in a] if constants is None else constants
coeffs = [pack_coefficients(forms) for forms in a] if coeffs is None else coeffs
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)
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, Iterable): # Blocked
_assemble_matrix_block_mat(A, a, bcs, diag, constants, coeffs)
else: # Non-blocked
constants = pack_constants(a) if constants is None else constants
coeffs = pack_coefficients(a) if coeffs is None else coeffs
_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)
A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH)
_cpp.fem.petsc.insert_diagonal(A, a.function_spaces[0], _bcs, diag)
return A
# -- Modifiers for Dirichlet conditions -----------------------------------
[docs]
def apply_lifting(
b: PETSc.Vec,
a: typing.Union[Iterable[Form], Iterable[Iterable[Form]]],
bcs: typing.Optional[typing.Union[Iterable[DirichletBC], Iterable[Iterable[DirichletBC]]]],
x0: typing.Optional[Iterable[PETSc.Vec]] = None,
alpha: float = 1,
constants: typing.Optional[
typing.Union[Iterable[np.ndarray], Iterable[Iterable[np.ndarray]]]
] = None,
coeffs=None,
) -> None:
"""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 used to modify ``b`` (see
:func:`dolfinx.fem.apply_lifting`). Two cases are supported:
1. The boundary conditions ``bcs`` are a
'sequence-of-sequences' such that ``bcs[j]`` are the
Dirichlet boundary conditionns associated with the forms in
the ``j`` th colulmn of ``a``. Helper functions exist to
create a sequence-of-sequences of `DirichletBC` from the 2D
``a`` and a flat Sequence of `DirichletBC` objects ``bcs``::
bcs1 = fem.bcs_by_block(
fem.extract_function_spaces(a, 1), bcs
)
2. ``bcs`` is a sequence of :class:`dolfinx.fem.DirichletBC`
objects. The function deduces which `DiricletBC` objects
apply to each column of ``a`` by matching the
:class:`dolfinx.fem.FunctionSpace`.
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:
try:
bcs = _bcs_by_block(_extract_spaces(a, 1), bcs)
except AttributeError:
pass
x0 = [] if x0 is None else x0.getNestSubVecs()
constants = [pack_constants(forms) for forms in a] if constants is None else constants
coeffs = [pack_coefficients(forms) for forms in a] if coeffs is None else coeffs
for b_sub, a_sub, const, coeff in zip(b.getNestSubVecs(), a, constants, coeffs):
const_ = list(
map(lambda x: np.array([], dtype=PETSc.ScalarType) if x is None else x, const)
)
apply_lifting(b_sub, a_sub, bcs, x0, alpha, const_, coeff)
else:
with contextlib.ExitStack() as stack:
if b.getAttr("_blocks") is not None:
if x0 is not None:
offset0, offset1 = x0.getAttr("_blocks")
xl = stack.enter_context(x0.localForm())
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
try:
bcs = _bcs_by_block(_extract_spaces(a, 1), bcs)
except AttributeError:
pass
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]
coeff = pack_coefficients(a_) if coeffs is None else coeffs[i]
const_ = [
np.empty(0, dtype=PETSc.ScalarType) if val is None else val
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)
size = off1 - off0
b_l.array_w[off0:off1] = bx_[:size]
b_l.array_w[offg0:offg1] = bx_[size:]
else:
try:
bcs = _bcs_by_block(_extract_spaces([a], 1), bcs)
except AttributeError:
pass
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)
return b
[docs]
def set_bc(
b: PETSc.Vec,
bcs: typing.Union[Iterable[DirichletBC], Iterable[Iterable[DirichletBC]]],
x0: typing.Optional[PETSc.Vec] = None,
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 b.getType() == PETSc.Vec.Type.NEST:
_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):
set_bc(b_sub, bc, x_sub, alpha)
else:
try:
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:]):
x0_sub = x_array[off0:off1] if x0 is not None else None
for bc in bcs:
bc.set(b_array[off0:off1], x0_sub, alpha)
except TypeError:
x0 = x0.array_r if x0 is not None else None
for bc in bcs:
bc.set(b.array_w, x0, alpha)
# -- High-level interface for KSP ---------------------------------------
[docs]
class LinearProblem:
"""Class for solving a linear variational problem.
Solves of the form :math:`a(u, v) = L(v) \\, \\forall v \\in V`
using PETSc as a linear algebra backend.
"""
def __init__(
self,
a: typing.Union[ufl.Form, Iterable[Iterable[ufl.Form]]],
L: typing.Union[ufl.Form, Iterable[ufl.Form]],
bcs: typing.Optional[Iterable[DirichletBC]] = None,
u: typing.Optional[typing.Union[_Function, Iterable[_Function]]] = None,
P: typing.Optional[typing.Union[ufl.Form, Iterable[Iterable[ufl.Form]]]] = None,
kind: typing.Optional[typing.Union[str, Iterable[Iterable[str]]]] = None,
petsc_options: typing.Optional[dict] = None,
form_compiler_options: typing.Optional[dict] = None,
jit_options: typing.Optional[dict] = None,
entity_maps: typing.Optional[dict[_Mesh, npt.NDArray[np.int32]]] = None,
) -> None:
"""Initialize solver for a linear variational problem.
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 type. See
:func:`create_matrix` for options.
petsc_options: Options set on the underlying PETSc KSP only.
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, `entity_maps` must be supplied.
For each key (a mesh, different to the integration domain
mesh) a map should be provided relating the entities in the
integration domain mesh to the entities in the key mesh
e.g. for a key-value pair (msh, emap) in `entity_maps`,
`emap[i]` is the entity in `msh` corresponding to entity
`i` in the integration domain mesh.
Example::
problem = LinearProblem(a, L, [bc0, bc1], petsc_options={
"ksp_type": "preonly",
"pc_type": "lu",
"pc_factor_mat_solver_type": "mumps"
})
problem = LinearProblem([[a00, a01], [None, a11]], [L0, L1],
bcs=[bc0, bc1], u=[uh0, uh1])
"""
self._a = _create_form(
a,
dtype=PETSc.ScalarType,
form_compiler_options=form_compiler_options,
jit_options=jit_options,
entity_maps=entity_maps,
)
self._L = _create_form(
L,
dtype=PETSc.ScalarType,
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,
dtype=PETSc.ScalarType,
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
self._b = create_vector(self.L, kind=kind)
self._x = create_vector(self.L, kind=kind)
if u is None:
try:
# Extract function space for unknown from the right hand
# side of the equation.
self._u = _Function(L.arguments()[0].ufl_function_space())
except AttributeError:
self._u = [_Function(Li.arguments()[0].ufl_function_space()) for Li in L]
else:
self._u = u
self.bcs = bcs
try:
comm = self.u.function_space.mesh.comm
except AttributeError:
comm = self.u[0].function_space.mesh.comm
self._solver = PETSc.KSP().create(comm)
self.solver.setOperators(self.A, self.P_mat)
# Give PETSc objects a unique prefix
problem_prefix = f"dolfinx_linearproblem_{id(self)}"
self.solver.setOptionsPrefix(problem_prefix)
self.A.setOptionsPrefix(f"{problem_prefix}_A")
self.b.setOptionsPrefix(f"{problem_prefix}_b")
self.x.setOptionsPrefix(f"{problem_prefix}_x")
if self.P_mat is not None:
self.P_mat.setOptionsPrefix(f"{problem_prefix}_P_mat")
# Set options on KSP only
if petsc_options is not None:
opts = PETSc.Options()
opts.prefixPush(problem_prefix)
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()
def __del__(self):
self._solver.destroy()
self._A.destroy()
self._b.destroy()
self._x.destroy()
[docs]
def solve(self) -> typing.Union[_Function, Iterable[_Function]]:
"""Solve the problem."""
# Assemble lhs
self.A.zeroEntries()
assemble_matrix(self.A, self.a, bcs=self.bcs)
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)
self.P_mat.assemble()
# Assemble rhs
if self.b.getType() == PETSc.Vec.Type.NEST:
for b_sub in self.b.getNestSubVecs():
with b_sub.localForm() as b_local:
b_local.set(0.0)
else:
with self.b.localForm() as b_loc:
b_loc.set(0)
assemble_vector(self.b, self.L)
# Apply boundary conditions to the rhs
if self.bcs is not None:
try:
apply_lifting(self.b, [self.a], bcs=[self.bcs])
dolfinx.la.petsc._ghost_update(
self.b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE
)
for bc in self.bcs:
bc.set(self.b.array_w)
except RuntimeError:
bcs1 = _bcs_by_block(_extract_spaces(self.a, 1), self.bcs) # type: ignore
apply_lifting(self.b, self.a, bcs=bcs1) # type: ignore
dolfinx.la.petsc._ghost_update(
self.b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE
)
bcs0 = _bcs_by_block(_extract_spaces(self.L), self.bcs) # type: ignore
dolfinx.fem.petsc.set_bc(self.b, bcs0)
else:
dolfinx.la.petsc._ghost_update(self.b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)
# 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)
dolfinx.fem.petsc.assign(self.x, self.u)
return self.u
@property
def L(self) -> typing.Union[Form, Iterable[Form]]:
"""The compiled linear form representing the left-hand side."""
return self._L
@property
def a(self) -> typing.Union[Form, Iterable[Form]]:
"""The compiled bilinear form representing the right-hand side."""
return self._a
@property
def preconditioner(self) -> typing.Union[Form, Iterable[Form]]:
"""The compiled bilinear form representing the preconditioner."""
return self._preconditioner
@property
def A(self) -> PETSc.Mat:
"""Left-hand side matrix.
Note:
The matrix has an options prefix set.
"""
return self._A
@property
def P_mat(self) -> PETSc.Mat:
"""Preconditioner matrix.
Note:
The matrix has an options prefix set.
"""
return self._P_mat
@property
def b(self) -> PETSc.Vec:
"""Right-hand side vector.
Note:
The vector has an options prefix set.
"""
return self._b
@property
def x(self) -> PETSc.Vec:
"""Solution vector.
Note:
This vector does not share memory with the solution
Function `u`.
Note:
The vector has an options prefix set.
"""
return self._x
@property
def solver(self) -> PETSc.KSP:
"""The PETSc KSP solver.
Note:
The KSP solver has an options prefix set.
"""
return self._solver
@property
def u(self) -> typing.Union[_Function, Iterable[_Function]]:
"""Solution function.
Note:
This vector does not share memory with the solution
vector `x`.
"""
return self._u
# -- High-level interface for SNES ---------------------------------------
def _ghostUpdate(x: PETSc.Vec, insert_mode: PETSc.InsertMode, scatter_mode: PETSc.ScatterMode): # type: ignore
"""Helper function for ghost updating PETSc vectors"""
try:
for x_sub in x.getNestSubVecs():
x_sub.ghostUpdate(addv=insert_mode, mode=scatter_mode)
except PETSc.Error: # type: ignore
x.ghostUpdate(addv=insert_mode, mode=scatter_mode)
def _zero_vector(x: PETSc.Vec): # type: ignore
"""Helper function for zeroing out PETSc vectors"""
try:
for x_sub in x.getNestSubVecs():
with x_sub.localForm() as x_sub_local:
x_sub_local.set(0.0)
except PETSc.Error: # type: ignore
with x.localForm() as x_local:
x_local.set(0.0)
def _assign_block_data(forms: typing.Iterable[dolfinx.fem.Form], vec: PETSc.Vec):
"""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.
"""
# Early exit if the vector already has block data or is a nest vector
if vec.getAttr("_blocks") is not None or vec.getType() == "nest":
return
maps = [
(
form.function_spaces[0].dofmaps(0).index_map,
form.function_spaces[0].dofmaps(0).index_map_bs,
)
for form in forms # type: ignore
]
off_owned = tuple(
itertools.accumulate(maps, lambda off, m: off + m[0].size_local * m[1], initial=0)
)
off_ghost = tuple(
itertools.accumulate(
maps, lambda off, m: off + m[0].num_ghosts * m[1], initial=off_owned[-1]
)
)
vec.setAttr("_blocks", (off_owned, off_ghost))
[docs]
def assemble_residual(
u: typing.Union[_Function, Sequence[_Function]],
residual: typing.Union[Form, typing.Iterable[Form]],
jacobian: typing.Union[Form, typing.Iterable[typing.Iterable[Form]]],
bcs: typing.Iterable[DirichletBC],
_snes: PETSc.SNES, # type: ignore
x: PETSc.Vec, # type: ignore
b: PETSc.Vec, # type: ignore
):
"""Assemble the residual into the vector `b`.
A function conforming to the interface expected by SNES.setResidual can
be created by fixing the first four arguments:
functools.partial(assemble_residual, u, residual, jacobian, bcs)
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
_ghostUpdate(x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
# Assign the input vector to the unknowns
assign(x, u)
# Assemble the residual
_zero_vector(b)
try:
# Single form and nest assembly
assemble_vector(b, residual)
except TypeError:
# Block assembly
_assign_block_data(residual, b) # type: ignore
assemble_vector(b, residual) # type: ignore
# Lift vector
try:
# Nest and blocked lifting
bcs1 = _bcs_by_block(_extract_spaces(jacobian, 1), bcs) # type: ignore
_assign_block_data(residual, x) # type: ignore
apply_lifting(b, jacobian, bcs=bcs1, x0=x, alpha=-1.0) # type: ignore
_ghostUpdate(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
bcs0 = _bcs_by_block(_extract_spaces(residual), bcs) # type: ignore
set_bc(b, bcs0, x0=x, alpha=-1.0)
except RuntimeError:
# Single form lifting
apply_lifting(b, [jacobian], bcs=[bcs], x0=[x], alpha=-1.0) # type: ignore
_ghostUpdate(b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # type: ignore
set_bc(b, bcs, x0=x, alpha=-1.0)
_ghostUpdate(b, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) # type: ignore
[docs]
def assemble_jacobian(
u: typing.Union[Sequence[_Function], _Function],
jacobian: typing.Union[Form, typing.Iterable[typing.Iterable[Form]]],
preconditioner: typing.Optional[typing.Union[Form, typing.Iterable[typing.Iterable[Form]]]],
bcs: typing.Iterable[DirichletBC],
_snes: PETSc.SNES, # type: ignore
x: PETSc.Vec, # type: ignore
J: PETSc.Mat, # type: ignore
P: PETSc.Mat, # type: ignore
):
"""Assemble the Jacobian and preconditioner matrices.
A function conforming to the interface expected by SNES.setJacobian can
be created by fixing the first four arguments:
functools.partial(assemble_jacobian, u, jacobian, preconditioner,
bcs)
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: Matrix to assemble the preconditioner into.
"""
# Copy existing soultion into the function used in the residual and
# Jacobian
try:
x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) # type: ignore
except PETSc.Error: # type: ignore
for x_sub in x.getNestSubVecs():
x_sub.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) # type: ignore
assign(x, u)
# Assemble Jacobian
J.zeroEntries()
assemble_matrix(J, jacobian, bcs, diag=1.0) # type: ignore
J.assemble()
if preconditioner is not None:
P.zeroEntries()
assemble_matrix(P, preconditioner, bcs, diag=1.0) # type: ignore
P.assemble()
# -- High-level interface for SNES ---------------------------------------
def _ghostUpdate(x: PETSc.Vec, insert_mode: PETSc.InsertMode, scatter_mode: PETSc.ScatterMode): # type: ignore
"""Helper function for ghost updating PETSc vectors"""
try:
for x_sub in x.getNestSubVecs():
x_sub.ghostUpdate(addv=insert_mode, mode=scatter_mode)
except PETSc.Error: # type: ignore
x.ghostUpdate(addv=insert_mode, mode=scatter_mode)
def _zero_vector(x: PETSc.Vec): # type: ignore
"""Helper function for zeroing out PETSc vectors"""
try:
for x_sub in x.getNestSubVecs():
with x_sub.localForm() as x_sub_local:
x_sub_local.set(0.0)
except PETSc.Error: # type: ignore
with x.localForm() as x_local:
x_local.set(0.0)
def _assign_block_data(forms: typing.Iterable[dolfinx.fem.Form], vec: PETSc.Vec):
"""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.
"""
# Early exit if the vector already has block data or is a nest vector
if vec.getAttr("_blocks") is not None or vec.getType() == "nest":
return
maps = [
(
form.function_spaces[0].dofmaps(0).index_map,
form.function_spaces[0].dofmaps(0).index_map_bs,
)
for form in forms # type: ignore
]
off_owned = tuple(
itertools.accumulate(maps, lambda off, m: off + m[0].size_local * m[1], initial=0)
)
off_ghost = tuple(
itertools.accumulate(
maps, lambda off, m: off + m[0].num_ghosts * m[1], initial=off_owned[-1]
)
)
vec.setAttr("_blocks", (off_owned, off_ghost))
[docs]
class NonlinearProblem:
def __init__(
self,
F: typing.Union[ufl.form.Form, Sequence[ufl.form.Form]],
u: typing.Union[_Function, Sequence[_Function]],
bcs: typing.Optional[Sequence[DirichletBC]] = None,
J: typing.Optional[typing.Union[ufl.form.Form, Sequence[Sequence[ufl.form.Form]]]] = None,
P: typing.Optional[typing.Union[ufl.form.Form, Sequence[Sequence[ufl.form.Form]]]] = None,
kind: typing.Optional[typing.Union[str, typing.Iterable[typing.Iterable[str]]]] = None,
form_compiler_options: typing.Optional[dict] = None,
jit_options: typing.Optional[dict] = None,
petsc_options: typing.Optional[dict] = None,
entity_maps: typing.Optional[dict[dolfinx.mesh.Mesh, npt.NDArray[np.int32]]] = None,
):
"""Class for solving nonlinear problems with SNES.
Solves problems of the form
:math:`F_i(u, v) = 0, i=0,...N\\ \\forall v \\in V` where
:math:`u=(u_0,...,u_N), v=(v_0,...,v_N)` using PETSc SNES as the
non-linear solver.
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",
}
Note:
The deprecated version of this class for use with
NewtonSolver has been renamed NewtonSolverNonlinearProblem.
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 type(s) for the Jacobian and
preconditioner (``MatType``).
See :func:`dolfinx.fem.petsc.create_matrix` for more
information.
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.
petsc_options: Options that are set on the underlying
PETSc SNES object only. For available choices for the
'petsc_options' kwarg, 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.
entity_maps: If any trial functions, test functions, or
coefficients in the form are not defined over the same mesh
as the integration domain, ``entity_maps`` must be
supplied. For each key (a mesh, different to the
integration domain mesh) a map should be provided relating
the entities in the integration domain mesh to the entities
in the key mesh e.g. for a key-value pair ``(msh, emap)``
in ``entity_maps``, ``emap[i]`` is the entity in ``msh``
corresponding to entity ``i`` in the integration domain
mesh.
"""
# 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
self._b = create_vector(self.F, kind=kind)
self._x = create_vector(self.F, kind=kind)
# Create the SNES solver and attach the corresponding Jacobian and
# residual computation functions
self._snes = PETSc.SNES().create(comm=self.A.comm) # type: ignore
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)
# Set PETSc options prefixes
problem_prefix = f"dolfinx_nonlinearproblem_{id(self)}"
self.solver.setOptionsPrefix(problem_prefix)
self.A.setOptionsPrefix(f"{problem_prefix}_A")
if self.P_mat is not None:
self.P_mat.setOptionsPrefix(f"{problem_prefix}_P_mat")
self.b.setOptionsPrefix(f"{problem_prefix}_b")
self.x.setOptionsPrefix(f"{problem_prefix}_x")
# Set options for SNES only
if petsc_options is not None:
opts = PETSc.Options() # type: ignore
opts.prefixPush(problem_prefix)
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()
[docs]
def solve(self) -> tuple[PETSc.Vec, int, int]: # type: ignore
"""Solve the problem and update the solution in the problem
instance.
Note:
The user is responsible for asserting convergence of the SNES
solver e.g. `assert converged_reason > 0`. Alternatively, pass
`"snes_error_if_not_converged": True` and
`"ksp_error_if_not_converged" : True` in `petsc_options`.
Returns:
The solution, convergence reason and number of iterations.
"""
# Move current iterate into the work array.
assign(self._u, self.x)
# Solve problem
self.solver.solve(None, self.x)
# Move solution back to function
assign(self.x, self._u)
converged_reason = self.solver.getConvergedReason()
return self.x, converged_reason, self.solver.getIterationNumber()
def __del__(self):
self._snes.destroy()
self._x.destroy()
self._A.destroy()
self._b.destroy()
if self._P_mat is not None:
self._P_mat.destroy()
@property
def F(self) -> typing.Union[Form, Iterable[Form]]:
"""The compiled residual."""
return self._F
@property
def J(self) -> typing.Union[Form, Iterable[Iterable[Form]]]:
"""The compiled Jacobian."""
return self._J
@property
def preconditioner(self) -> typing.Optional[typing.Union[Form, Iterable[Iterable[Form]]]]:
"""The compiled preconditioner."""
return self._preconditioner
@property
def A(self) -> PETSc.Mat:
"""Jacobian matrix.
Note:
The matrix has an options prefix set.
"""
return self._A
@property
def b(self) -> PETSc.Vec:
"""Residual vector.
Note:
The vector has an options prefix set.
"""
return self._b
@property
def x(self) -> PETSc.Vec:
"""Solution vector.
Note:
The vector does not share memory with the
solution Function `u`.
Note:
The vector has an options prefix set.
"""
return self._x
@property
def P_mat(self) -> typing.Optional[PETSc.Vec]:
"""Preconditioner matrix.
Note:
The matrix has an options prefix set.
"""
return self._P_mat
@property
def solver(self) -> PETSc.SNES:
"""The SNES solver.
Note:
The SNES solver has an options prefix set.
"""
return self._snes
@property
def u(self) -> typing.Union[_Function, Iterable[_Function]]:
"""Solution function.
Note:
The Function does not share memory with the solution
vector `x`.
"""
return self._u
# -- Deprecated non-linear problem class for NewtonSolver -----------------
class NewtonSolverNonlinearProblem:
"""Nonlinear problem class for solving the non-linear problems using
NewtonSolver.
Note:
This class is deprecated in favour of NonlinearProblem, a high
level interface to SNES.
Note:
This class was previously called NonlinearProblem.
Solves problems of the form :math:`F(u, v) = 0 \\ \\forall v \\in V`
using PETSc as the linear algebra backend.
"""
def __init__(
self,
F: ufl.form.Form,
u: _Function,
bcs: typing.Optional[Iterable[DirichletBC]] = None,
J: ufl.form.Form = None,
form_compiler_options: typing.Optional[dict] = None,
jit_options: typing.Optional[dict] = 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:
"""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)
def F(self, x: PETSc.Vec, b: PETSc.Vec) -> None:
"""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
with b.localForm() as b_local:
b_local.set(0.0)
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)
set_bc(b, self.bcs, x, -1.0)
else:
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
def J(self, x: PETSc.Vec, A: PETSc.Mat) -> None:
"""Assemble the Jacobian matrix.
Args:
x: The vector containing the latest solution
"""
A.zeroEntries()
assemble_matrix(A, self._a, self.bcs)
A.assemble()
# -- Additional free helper functions (interpolations, assignments etc.) --
[docs]
def discrete_curl(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat:
"""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)
[docs]
def discrete_gradient(space0: _FunctionSpace, space1: _FunctionSpace) -> PETSc.Mat:
"""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)
[docs]
def interpolation_matrix(V0: _FunctionSpace, V1: _FunctionSpace) -> PETSc.Mat:
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`.
"""
return _interpolation_matrix(V0._cpp_object, V1._cpp_object)
[docs]
@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 ``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:
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)