Source code for dolfinx.fem.petsc

# Copyright (C) 2018-2022 Garth N. Wells 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."""


from __future__ import annotations

import typing

if typing.TYPE_CHECKING:
    from dolfinx.fem.forms import FormMetaClass
    from dolfinx.fem.bcs import DirichletBCMetaClass

import contextlib
import functools

import ufl
from dolfinx import cpp as _cpp
from dolfinx import la
from dolfinx.cpp.fem import pack_coefficients as _pack_coefficients
from dolfinx.cpp.fem import pack_constants as _pack_constants
from dolfinx.fem import assemble
from dolfinx.fem.bcs import bcs_by_block as _bcs_by_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 petsc4py import PETSc


def _extract_function_spaces(a: typing.List[typing.List[FormMetaClass]]):
    """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 = map(partial(map, fn), a)

    # Compute spaces for each row/column block
    rows = [set() for i in range(len(a))]
    cols = [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: FormMetaClass) -> PETSc.Vec: """Create a PETSc vector that is compaible with a linear form. Args: L: A linear form. Returns: A PETSc vector with a layout that is compatible with `L`. """ dofmap = L.function_spaces[0].dofmap return la.create_petsc_vector(dofmap.index_map, dofmap.index_map_bs)
[docs]def create_vector_block(L: typing.List[FormMetaClass]) -> PETSc.Vec: """Create a PETSc vector (blocked) that is compaible with a list of linear forms. Args: L: List of linear forms. Returns: A PETSc vector with a layout that is compatible with `L`. """ maps = [(form.function_spaces[0].dofmap.index_map, form.function_spaces[0].dofmap.index_map_bs) for form in L] return _cpp.fem.petsc.create_vector_block(maps)
[docs]def create_vector_nest(L: typing.List[FormMetaClass]) -> PETSc.Vec: """Create a PETSc netsted vector (``VecNest``) that is compaible with a list of linear forms. Args: L: List of linear forms. Returns: A PETSc nested vector (``VecNest``) with a layout that is compatible with `L`. """ maps = [(form.function_spaces[0].dofmap.index_map, form.function_spaces[0].dofmap.index_map_bs) for form in L] return _cpp.fem.petsc.create_vector_nest(maps)
# -- Matrix instantiation ----------------------------------------------------
[docs]def create_matrix(a: FormMetaClass, mat_type=None) -> PETSc.Mat: """Create a PETSc matrix that is compaible with a bilinear form. Args: a: A bilinear form. mat_type: The PETSc matrix type (``MatType``). Returns: A PETSc matrix with a layout that is compatible with `a`. """ if mat_type is None: return _cpp.fem.petsc.create_matrix(a) else: return _cpp.fem.petsc.create_matrix(a, mat_type)
[docs]def create_matrix_block(a: typing.List[typing.List[FormMetaClass]]) -> PETSc.Mat: """Create a PETSc matrix that is compaible with a rectangular array of bilinear forms. Args: a: A rectangular array of bilinear forms. Returns: A PETSc matrix with a blocked layout that is compatible with `a`. """ return _cpp.fem.petsc.create_matrix_block(a)
[docs]def create_matrix_nest(a: typing.List[typing.List[FormMetaClass]]) -> PETSc.Mat: """Create a PETSc matrix (``MatNest``) that is compaible with a rectangular array of bilinear forms. Args: a: A rectangular array of bilinear forms. Returns: A PETSc matrix ('MatNest``) that is compatible with `a`. """ return _cpp.fem.petsc.create_matrix_nest(a)
# -- Vector assembly ---------------------------------------------------------
[docs]@functools.singledispatch def assemble_vector(L: FormMetaClass, constants=None, coeffs=None) -> PETSc.Vec: """Assemble linear form into a new PETSc vector. Note: The returned vector is not finalised, i.e. ghost values are not accumulated on the owning processes. Args: L: A linear form. Returns: An assembled vector. """ b = la.create_petsc_vector(L.function_spaces[0].dofmap.index_map, L.function_spaces[0].dofmap.index_map_bs) with b.localForm() as b_local: assemble.assemble_vector(b_local.array_w, L, constants, coeffs) return b
@assemble_vector.register(PETSc.Vec) def _(b: PETSc.Vec, L: FormMetaClass, constants=None, coeffs=None) -> PETSc.Vec: """Assemble linear form into an existing PETSc vector. 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 to assemble into `b`. Returns: An assembled vector. """ with b.localForm() as b_local: assemble.assemble_vector(b_local.array_w, L, constants, coeffs) return b
[docs]@functools.singledispatch def assemble_vector_nest(L: typing.List[FormMetaClass], constants=None, coeffs=None) -> PETSc.Vec: """Assemble linear forms into a new nested PETSc (VecNest) vector. The returned vector is not finalised, i.e. ghost values are not accumulated on the owning processes. """ maps = [(form.function_spaces[0].dofmap.index_map, form.function_spaces[0].dofmap.index_map_bs) for form in L] b = _cpp.fem.petsc.create_vector_nest(maps) for b_sub in b.getNestSubVecs(): with b_sub.localForm() as b_local: b_local.set(0.0) return assemble_vector_nest(b, L, constants, coeffs)
@assemble_vector_nest.register(PETSc.Vec) def _(b: PETSc.Vec, L: typing.List[FormMetaClass], constants=None, coeffs=None) -> PETSc.Vec: """Assemble linear forms into a nested PETSc (VecNest) vector. The vector is not zeroed before assembly and it is not finalised, i.e. ghost values are not accumulated on the owning processes. """ 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.assemble_vector(b_local.array_w, L_sub, const, coeff) return b # FIXME: Revise this interface
[docs]@functools.singledispatch def assemble_vector_block(L: typing.List[FormMetaClass], a: typing.List[typing.List[FormMetaClass]], bcs: typing.List[DirichletBCMetaClass] = [], x0: typing.Optional[PETSc.Vec] = None, scale: float = 1.0, constants_L=None, coeffs_L=None, constants_a=None, coeffs_a=None) -> PETSc.Vec: """Assemble linear forms into a monolithic vector. The vector is not finalised, i.e. ghost values are not accumulated. """ maps = [(form.function_spaces[0].dofmap.index_map, form.function_spaces[0].dofmap.index_map_bs) for form in L] b = _cpp.fem.petsc.create_vector_block(maps) with b.localForm() as b_local: b_local.set(0.0) return assemble_vector_block(b, L, a, bcs, x0, scale, constants_L, coeffs_L, constants_a, coeffs_a)
@assemble_vector_block.register(PETSc.Vec) def _(b: PETSc.Vec, L: typing.List[FormMetaClass], a, bcs: typing.List[DirichletBCMetaClass] = [], x0: typing.Optional[PETSc.Vec] = None, scale: float = 1.0, constants_L=None, coeffs_L=None, constants_a=None, coeffs_a=None) -> PETSc.Vec: """Assemble linear forms into a monolithic vector. The vector is not zeroed and it is not finalised, i.e. ghost values are not accumulated. """ maps = [(form.function_spaces[0].dofmap.index_map, form.function_spaces[0].dofmap.index_map_bs) for form in L] if x0 is not None: x0_local = _cpp.la.petsc.get_local_vectors(x0, maps) x0_sub = x0_local else: x0_local = [] x0_sub = [None] * len(maps) constants_L = [form and _pack_constants(form) for form in L] if constants_L is None else constants_L coeffs_L = [{} if form is None else _pack_coefficients( form) for form in L] if coeffs_L is None else coeffs_L constants_a = [[form and _pack_constants(form) for form in forms] for forms in a] if constants_a is None else constants_a coeffs_a = [[{} if form is None else _pack_coefficients( form) for form in forms] for forms in a] if coeffs_a is None else coeffs_a bcs1 = _bcs_by_block(_extract_spaces(a, 1), bcs) b_local = _cpp.la.petsc.get_local_vectors(b, maps) for b_sub, L_sub, a_sub, const_L, coeff_L, const_a, coeff_a in zip(b_local, L, a, constants_L, coeffs_L, constants_a, coeffs_a): _cpp.fem.assemble_vector(b_sub, L_sub, const_L, coeff_L) _cpp.fem.apply_lifting(b_sub, a_sub, const_a, coeff_a, bcs1, x0_local, scale) _cpp.la.petsc.scatter_local_vectors(b, b_local, maps) b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) bcs0 = _bcs_by_block(_extract_spaces(L), bcs) offset = 0 b_array = b.getArray(readonly=False) for submap, bc, _x0 in zip(maps, bcs0, x0_sub): size = submap[0].size_local * submap[1] _cpp.fem.set_bc(b_array[offset: offset + size], bc, _x0, scale) offset += size return b # -- Matrix assembly ---------------------------------------------------------
[docs]@functools.singledispatch def assemble_matrix(a: FormMetaClass, bcs: typing.List[DirichletBCMetaClass] = [], diagonal: float = 1.0, constants=None, coeffs=None) -> PETSc.Mat: """Assemble bilinear form into a matrix. The returned matrix is not finalised, i.e. ghost values are not accumulated. """ A = _cpp.fem.petsc.create_matrix(a) assemble_matrix(A, a, bcs, diagonal, constants, coeffs) return A
@assemble_matrix.register(PETSc.Mat) def _(A: PETSc.Mat, a: FormMetaClass, bcs: typing.List[DirichletBCMetaClass] = [], diagonal: float = 1.0, constants=None, coeffs=None) -> PETSc.Mat: """Assemble bilinear form into a matrix. The returned matrix is not finalised, i.e. ghost values are not accumulated. """ constants = _pack_constants(a) if constants is None else constants coeffs = _pack_coefficients(a) if coeffs is None else coeffs _cpp.fem.petsc.assemble_matrix(A, a, 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, diagonal) return A # FIXME: Revise this interface
[docs]@functools.singledispatch def assemble_matrix_nest(a: typing.List[typing.List[FormMetaClass]], bcs: typing.List[DirichletBCMetaClass] = [], mat_types=[], diagonal: float = 1.0, constants=None, coeffs=None) -> PETSc.Mat: """Assemble bilinear forms into matrix""" A = _cpp.fem.petsc.create_matrix_nest(a, mat_types) assemble_matrix_nest(A, a, bcs, diagonal, constants, coeffs) return A
@assemble_matrix_nest.register(PETSc.Mat) def _(A: PETSc.Mat, a: typing.List[typing.List[FormMetaClass]], bcs: typing.List[DirichletBCMetaClass] = [], diagonal: float = 1.0, constants=None, coeffs=None) -> PETSc.Mat: """Assemble bilinear forms into matrix""" constants = [[form and _pack_constants(form) for form in forms] for forms in a] if constants is None else constants coeffs = [[{} if form is None else _pack_coefficients( form) for form in 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, diagonal, const, coeff) elif i == j: for bc in bcs: if a_row[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.") return A # FIXME: Revise this interface
[docs]@functools.singledispatch def assemble_matrix_block(a: typing.List[typing.List[FormMetaClass]], bcs: typing.List[DirichletBCMetaClass] = [], diagonal: float = 1.0, constants=None, coeffs=None) -> PETSc.Mat: """Assemble bilinear forms into matrix""" A = _cpp.fem.petsc.create_matrix_block(a) return assemble_matrix_block(A, a, bcs, diagonal, constants, coeffs)
@assemble_matrix_block.register(PETSc.Mat) def _(A: PETSc.Mat, a: typing.List[typing.List[FormMetaClass]], bcs: typing.List[DirichletBCMetaClass] = [], diagonal: float = 1.0, constants=None, coeffs=None) -> PETSc.Mat: """Assemble bilinear forms into matrix""" constants = [[form and _pack_constants(form) for form in forms] for forms in a] if constants is None else constants coeffs = [[{} if form is None else _pack_coefficients( form) for form in forms] for forms in a] if coeffs is None else coeffs V = _extract_function_spaces(a) is_rows = _cpp.la.petsc.create_index_sets([(Vsub.dofmap.index_map, Vsub.dofmap.index_map_bs) for Vsub in V[0]]) is_cols = _cpp.la.petsc.create_index_sets([(Vsub.dofmap.index_map, Vsub.dofmap.index_map_bs) for Vsub in V[1]]) # Assemble form for i, a_row in enumerate(a): for j, a_sub in enumerate(a_row): if a_sub is not None: Asub = A.getLocalSubMatrix(is_rows[i], is_cols[j]) _cpp.fem.petsc.assemble_matrix(Asub, a_sub, constants[i][j], coeffs[i][j], bcs, True) A.restoreLocalSubMatrix(is_rows[i], is_cols[j], Asub) elif i == j: for bc in bcs: if a_row[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(is_rows[i], is_cols[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, diagonal) A.restoreLocalSubMatrix(is_rows[i], is_cols[j], Asub) return A # -- Modifiers for Dirichlet conditions ---------------------------------------
[docs]def apply_lifting(b: PETSc.Vec, a: typing.List[FormMetaClass], bcs: typing.List[typing.List[DirichletBCMetaClass]], x0: typing.Optional[typing.List[PETSc.Vec]] = [], scale: float = 1.0, constants=None, coeffs=None) -> None: """Apply the function :func:`dolfinx.fem.apply_lifting` to a PETSc Vector.""" with contextlib.ExitStack() as stack: 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()) assemble.apply_lifting(b_local.array_w, a, bcs, x0_r, scale, constants, coeffs)
[docs]def apply_lifting_nest(b: PETSc.Vec, a: typing.List[typing.List[FormMetaClass]], bcs: typing.List[DirichletBCMetaClass], x0: typing.Optional[PETSc.Vec] = None, scale: float = 1.0, constants=None, coeffs=None) -> PETSc.Vec: """Apply the function :func:`dolfinx.fem.apply_lifting` to each sub-vector in a nested PETSc Vector.""" x0 = [] if x0 is None else x0.getNestSubVecs() bcs1 = _bcs_by_block(_extract_spaces(a, 1), bcs) constants = [[form and _pack_constants(form) for form in forms] for forms in a] if constants is None else constants coeffs = [[{} if form is None else _pack_coefficients( form) for form in 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): apply_lifting(b_sub, a_sub, bcs1, x0, scale, const, coeff) return b
[docs]def set_bc(b: PETSc.Vec, bcs: typing.List[DirichletBCMetaClass], x0: typing.Optional[PETSc.Vec] = None, scale: float = 1.0) -> None: """Apply the function :func:`dolfinx.fem.set_bc` to a PETSc Vector.""" if x0 is not None: x0 = x0.array_r assemble.set_bc(b.array_w, bcs, x0, scale)
[docs]def set_bc_nest(b: PETSc.Vec, bcs: typing.List[typing.List[DirichletBCMetaClass]], x0: typing.Optional[PETSc.Vec] = None, scale: float = 1.0) -> None: """Apply the function :func:`dolfinx.fem.set_bc` to each sub-vector of a nested PETSc Vector.""" _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, scale)
[docs]class LinearProblem(): """Class for solving a linear variational problem of the form :math:`a(u, v) = L(v) \\, \\forall v \\in V` using PETSc as a linear algebra backend. """ def __init__(self, a: ufl.Form, L: ufl.Form, bcs: typing.List[DirichletBCMetaClass] = [], u: _Function = None, petsc_options={}, form_compiler_params={}, jit_params={}): """Initialize solver for a linear variational problem. Args: a: A bilinear UFL form, the left hand side of the variational problem. L: A linear UFL form, the right hand side of the variational problem. bcs: A list of Dirichlet boundary conditions. u: The solution function. It will be created if not provided. petsc_options: Parameters that is passed to the linear algebra backend PETSc. For available choices for the 'petsc_options' kwarg, see the `PETSc documentation <https://petsc4py.readthedocs.io/en/stable/manual/ksp/>`_. form_compiler_params: Parameters used in FFCx compilation of this form. Run ``ffcx --help`` at the commandline to see all available options. jit_params: Parameters used in CFFI JIT compilation of C code generated by FFCx. See `python/dolfinx/jit.py` for all available parameters. Takes priority over all other parameter values. Example:: problem = LinearProblem(a, L, [bc0, bc1], petsc_options={"ksp_type": "preonly", "pc_type": "lu"}) """ self._a = _create_form(a, form_compiler_params=form_compiler_params, jit_params=jit_params) self._A = create_matrix(self._a) self._L = _create_form(L, form_compiler_params=form_compiler_params, jit_params=jit_params) self._b = create_vector(self._L) if u is None: # Extract function space from TrialFunction (which is at the # end of the argument list as it is numbered as 1, while the # Test function is numbered as 0) self.u = _Function(a.arguments()[-1].ufl_function_space()) else: self.u = u self._x = _cpp.la.petsc.create_vector_wrap(self.u.x) self.bcs = bcs self._solver = PETSc.KSP().create(self.u.function_space.mesh.comm) self._solver.setOperators(self._A) # Give PETSc solver options a unique prefix problem_prefix = f"dolfinx_solve_{id(self)}" self._solver.setOptionsPrefix(problem_prefix) # Set PETSc options opts = PETSc.Options() opts.prefixPush(problem_prefix) for k, v in petsc_options.items(): opts[k] = v opts.prefixPop() self._solver.setFromOptions() # Set matrix and vector PETSc options self._A.setOptionsPrefix(problem_prefix) self._A.setFromOptions() self._b.setOptionsPrefix(problem_prefix) self._b.setFromOptions()
[docs] def solve(self) -> _Function: """Solve the problem.""" # Assemble lhs self._A.zeroEntries() assemble_matrix(self._A, self._a, bcs=self.bcs) self._A.assemble() # Assemble rhs with self._b.localForm() as b_loc: b_loc.set(0) assemble_vector(self._b, self._L) # Apply boundary conditions to the rhs apply_lifting(self._b, [self._a], bcs=[self.bcs]) self._b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) set_bc(self._b, self.bcs) # Solve linear system and update ghost values in the solution self._solver.solve(self._b, self._x) self.u.x.scatter_forward() return self.u
@property def L(self) -> FormMetaClass: """The compiled linear form""" return self._L @property def a(self) -> FormMetaClass: """The compiled bilinear form""" return self._a @property def A(self) -> PETSc.Mat: """Matrix operator""" return self._A @property def b(self) -> PETSc.Vec: """Right-hand side vector""" return self._b @property def solver(self) -> PETSc.KSP: """Linear solver object""" return self._solver
[docs]class NonlinearProblem: """Nonlinear problem class for solving the non-linear problem :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.List[DirichletBCMetaClass] = [], J: ufl.form.Form = None, form_compiler_params={}, jit_params={}): """Initialize solver for solving a non-linear problem using Newton's method, :math:`dF/du(u) du = -F(u)`. 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_params: Parameters used in FFCx compilation of this form. Run ``ffcx --help`` at the commandline to see all available options. jit_params: Parameters used in CFFI JIT compilation of C code generated by FFCx. See ``python/dolfinx/jit.py`` for all available parameters. Takes priority over all other parameter values. Example:: problem = LinearProblem(F, u, [bc0, bc1]) """ self._L = _create_form(F, form_compiler_params=form_compiler_params, jit_params=jit_params) # 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_params=form_compiler_params, jit_params=jit_params) self.bcs = bcs @property def L(self) -> FormMetaClass: """Compiled linear form (the residual form)""" return self._L @property def a(self) -> FormMetaClass: """Compiled bilinear form (the Jacobian form)""" return self._a
[docs] def form(self, x: PETSc.Vec): """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)
[docs] def F(self, x: PETSc.Vec, b: PETSc.Vec): """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 apply_lifting(b, [self._a], bcs=[self.bcs], x0=[x], scale=-1.0) b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) set_bc(b, self.bcs, x, -1.0)
[docs] def J(self, x: PETSc.Vec, A: PETSc.Mat): """Assemble the Jacobian matrix. Args: x: The vector containing the latest solution """ A.zeroEntries() assemble_matrix(A, self._a, self.bcs) A.assemble()