from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING
import dolfinx
import dolfinx.fem
if TYPE_CHECKING:
# Gives the `else:` branch below (petsc4py absent) the same `PETSc.Vec`
# annotations as the `if` branch, so mypy sees identical conditional
# function signatures. Postponed evaluation (`from __future__ import
# annotations`, above) means these annotations are never evaluated at
# runtime, so this import being TYPE_CHECKING-only (never executed) does
# not affect the `else:` branch's runtime behaviour.
from petsc4py import PETSc
if dolfinx.has_petsc4py:
from petsc4py import PETSc
# dolfinx.la.petsc's vector helpers are private API, wrapped here so a
# rename upstream only needs fixing in one place.
import dolfinx.la.petsc
def zero_petsc_vector(b: PETSc.Vec) -> None:
"""Zero a PETSc vector, including ghosts."""
dolfinx.la.petsc._zero_vector(b)
def ghost_update(x: PETSc.Vec, insert_mode, scatter_mode) -> None:
"""Ghost-update a PETSc vector."""
dolfinx.la.petsc._ghost_update(x, insert_mode, scatter_mode)
import dolfinx.fem.petsc
def set_bc(b: PETSc.Vec, bcs, x0=None, alpha: float = 1.0) -> None:
"""Set boundary condition values in a vector, across the nested-vs-single-form
`set_bc` API."""
try:
dolfinx.fem.petsc.set_bc(b, bcs, x0=x0, alpha=alpha)
except AttributeError:
for _bcs in bcs:
dolfinx.fem.petsc.set_bc(b, _bcs, x0=x0, alpha=alpha)
except TypeError:
# bcs shouldn't be a nested list here (e.g. for a single form).
assert len(bcs) == 1, "bcs should be a single DirichletBC or a list of DirichletBCs."
dolfinx.fem.petsc.set_bc(b, bcs[0], x0=x0, alpha=alpha)
def apply_lifting_and_set_bc(
b: PETSc.Vec,
a,
bcs,
x: PETSc.Vec | None = None,
alpha: float = 1.0,
) -> None:
r"""Apply lifting to a vector and set boundary conditions.
Convenience function to apply lifting and set boundary conditions
for single, blocked, or nested forms. Modifies `b` such that its free
rows become `b - alpha * A (g - x)` and its constrained rows
`alpha * (g - x)`, with `g` the boundary values and `A` the operator
of `a`.
Args:
b: The vector to apply lifting to.
a: Form or nested sequence of forms to apply lifting from.
bcs: The boundary conditions to apply.
x: Vector to subtract from the boundary conditions (e.g. in a Newton iteration).
alpha: Scaling factor for the boundary conditions.
"""
# Ruling F6 (revised): use this module's own bcs_by_block (Task 21)
# instead of scifem's `dolfinx.fem.bcs_by_block`, which on failure
# hands every bc to every space unfiltered -- silently wrong for a
# blocked/nested system. That much of scifem's try/except was a real
# bug and stays deleted.
#
# But `dolfinx.fem.extract_function_spaces` also refuses an explicit
# index for a bare single form (`ValueError: index must be None for
# a single form`, confirmed), and for that case `bcs0 = bcs1 = bcs`
# is not a workaround, it is the correct answer: a single form has
# exactly one block, so every bc in `bcs` belongs to it. The
# docstring above promises "single, blocked, or nested forms", so
# single-form support is part of this function's contract and must
# stay -- but selected by an explicit, positive check on `a`'s shape
# (never by catching whatever exception extract_function_spaces
# happens to raise, which could silently swallow an unrelated bug).
if isinstance(a, (list, tuple)):
# Blocked/nested: `a` is a 2D array of forms, one row per block.
bcs0 = bcs_by_block(dolfinx.fem.extract_function_spaces(a, 0), bcs)
bcs1 = bcs_by_block(dolfinx.fem.extract_function_spaces(a, 1), bcs)
a_lifting = a
x_lifting: PETSc.Vec | list[PETSc.Vec] | None = x
else:
# Single form: one block, so every bc applies to it. `set_bc`
# (below) accepts `bcs` in its bare/flat shape directly, but
# `dolfinx.fem.petsc.apply_lifting` always wants `a` and `bcs` one
# level more nested than that for a plain (non-blocked, non-nest)
# vector -- confirmed: a bare `a` raises `TypeError: 'Form'
# object is not iterable`, and a flat `bcs` raises `TypeError:
# 'DirichletBC' object is not iterable` inside apply_lifting.
# The same goes for `x`, which apply_lifting iterates over: a
# bare Vec raises `AttributeError` there, on every version.
bcs0 = bcs
bcs1 = [bcs]
a_lifting = [a]
x_lifting = None if x is None else [x]
dolfinx.fem.petsc.apply_lifting(b, a_lifting, bcs=bcs1, x0=x_lifting, alpha=alpha) # type: ignore[arg-type]
ghost_update(b, PETSc.InsertMode.ADD_VALUES, PETSc.ScatterMode.REVERSE)
set_bc(b, bcs0, x0=x, alpha=alpha)
ghost_update(b, PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
else:
[docs]
def zero_petsc_vector(b: PETSc.Vec) -> None:
raise RuntimeError("petsc4py is not available. Cannot zero vector.")
[docs]
def ghost_update(x: PETSc.Vec, insert_mode, scatter_mode) -> None:
raise RuntimeError("petsc4py is not available. Cannot ghost update vector.")
[docs]
def set_bc(b: PETSc.Vec, bcs, x0=None, alpha: float = 1.0) -> None:
raise RuntimeError("petsc4py is not available. Cannot set boundary conditions.")
[docs]
def apply_lifting_and_set_bc(
b: PETSc.Vec, a, bcs, x: PETSc.Vec | None = None, alpha: float = 1.0
) -> None:
raise RuntimeError(
"petsc4py is not available. Cannot apply lifting and set boundary conditions."
)
[docs]
def bcs_by_block(
spaces: Sequence[dolfinx.fem.FunctionSpace | None],
bcs: Sequence[dolfinx.fem.DirichletBC],
) -> list[list[dolfinx.fem.DirichletBC]]:
"""Arrange boundary conditions by the space they constrain, across dolfinx 0.11/0.12.
`dolfinx.fem.bcs_by_block` only works when the spaces and
`bc.function_space` are the same flavour: both cpp objects on 0.10
and 0.11, both Python wrappers from dolfinx PRs #4342 and #4350. That
holds for spaces from `dolfinx.fem.extract_function_spaces`, but not
for spaces the caller builds a list of. Normalizing both sides to cpp
objects and calling the cpp `contains` directly accepts either
flavour on every version.
"""
def _cpp(space):
return getattr(space, "_cpp_object", space)
return [
[bc for bc in bcs if _cpp(V).contains(_cpp(bc.function_space))] if V is not None else []
for V in spaces
]