Simplify assemble routines (#4241)

* Remove additional callable

* Typing

* Simplify and typing of pack_constants

* Simplify packing logic

* Remove deprecated set_bc

* Various

Author: schnellerhase

python/dolfinx/fem/__init__.py

@@ -21,7 +21,6 @@
     create_vector,
     pack_coefficients,
     pack_constants,
-    set_bc,
 )
 from dolfinx.fem.bcs import (
     DirichletBC,

python/dolfinx/fem/assemble.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2018-2022 Garth N. Wells, Jack S. Hale
+# Copyright (C) 2018-2022 Garth N. Wells, Jack S. Hale and Paul T. Kühner
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -9,7 +9,6 @@
 
 import functools
 import typing
-import warnings
 from collections.abc import Sequence
 
 import numpy as np
@@ -26,9 +25,15 @@
 from dolfinx.fem.function import FunctionSpace
 
 
-def pack_constants(
-    form: Form | Sequence[Form],
-) -> npt.NDArray | Sequence[npt.NDArray]:
+@typing.overload
+def pack_constants(form: Form) -> npt.NDArray: ...
+
+
+@typing.overload
+def pack_constants(form: Sequence[Form]) -> list[npt.NDArray]: ...
+
+
+def pack_constants(form):
     """Pack form constants for use in assembly.
 
     Pack the 'constants' that appear in forms. The packed constants can
@@ -46,23 +51,25 @@ def pack_constants(
     Returns:
         A ``constant`` array for each form.
     """
+    if form is None:
+        return None
+    elif isinstance(form, Sequence):
+        return list(map(pack_constants, form))
+    else:
+        return _pack_constants(form._cpp_object)
 
-    def _pack(form):
-        if form is None:
-            return None
-        elif isinstance(form, Sequence):
-            return list(map(lambda sub_form: _pack(sub_form), form))
-        else:
-            return _pack_constants(form._cpp_object)
 
-    return _pack(form)
+@typing.overload
+def pack_coefficients(form: None | Form) -> dict[tuple[IntegralType, int], npt.NDArray]: ...
 
 
+@typing.overload
 def pack_coefficients(
-    form: Form | Sequence[Form],
-) -> (
-    dict[tuple[IntegralType, int], npt.NDArray] | list[dict[tuple[IntegralType, int], npt.NDArray]]
-):
+    form: Sequence[Form],
+) -> list[dict[tuple[IntegralType, int], npt.NDArray]]: ...
+
+
+def pack_coefficients(form):
     """Pack form coefficients for use in assembly.
 
     Pack the ``coefficients`` that appear in forms. The packed
@@ -80,16 +87,12 @@ def pack_coefficients(
     Returns:
         Coefficients for each form.
     """
-
-    def _pack(form):
-        if form is None:
-            return {}
-        elif isinstance(form, Sequence):
-            return list(map(lambda sub_form: _pack(sub_form), form))
-        else:
-            return _pack_coefficients(form._cpp_object)
-
-    return _pack(form)
+    if form is None:
+        return {}
+    elif isinstance(form, Sequence):
+        return list(map(pack_coefficients, form))
+    else:
+        return _pack_coefficients(form._cpp_object)
 
 
 # -- Vector and matrix instantiation --------------------------------------
@@ -107,7 +110,7 @@ def create_vector(V: FunctionSpace, dtype: npt.DTypeLike = default_scalar_type)
     """
     # Can just take the first dofmap here, since all dof maps have the same
     # index map in mixed-topology meshes
-    dofmap = V.dofmaps[0]  # type: ignore[attr-defined]
+    dofmap = V.dofmaps[0]
     return la.vector(dofmap.index_map, dofmap.index_map_bs, dtype=dtype)
 
 
@@ -160,8 +163,12 @@ def assemble_scalar(
         To compute the functional value on the whole domain, the output
         of this function is typically summed across all MPI ranks.
     """
-    constants = pack_constants(M) if constants is None else constants  # type: ignore[assignment]
-    coeffs = pack_coefficients(M) if coeffs is None else coeffs  # type: ignore[assignment]
+    if constants is None:
+        constants = pack_constants(M)
+
+    if coeffs is None:
+        coeffs = pack_coefficients(M)
+
     return _cpp.fem.assemble_scalar(M._cpp_object, constants, coeffs)
 
 
@@ -209,8 +216,13 @@ def _assemble_vector_form(
     """
     b = create_vector(L.function_spaces[0], L.dtype)
     b.array[:] = 0
-    constants = pack_constants(L) if constants is None else constants  # type: ignore[assignment]
-    coeffs = pack_coefficients(L) if coeffs is None else coeffs  # type: ignore[assignment]
+
+    if constants is None:
+        constants = pack_constants(L)
+
+    if coeffs is None:
+        coeffs = pack_coefficients(L)
+
     _assemble_vector_array(b.array, L, constants, coeffs)
     return b
 
@@ -244,8 +256,12 @@ def _assemble_vector_array(
         :func:`dolfinx.la.Vector.scatter_reverse` on the return vector
         can accumulate ghost contributions.
     """
-    constants = pack_constants(L) if constants is None else constants  # type: ignore[assignment]
-    coeffs = pack_coefficients(L) if coeffs is None else coeffs  # type: ignore[assignment]
+    if constants is None:
+        constants = pack_constants(L)
+
+    if coeffs is None:
+        coeffs = pack_coefficients(L)
+
     _cpp.fem.assemble_vector(b, L._cpp_object, constants, coeffs)
     return b
 
@@ -287,8 +303,10 @@ def assemble_matrix(
         The returned matrix is not finalised, i.e. ghost values are not
         accumulated.
     """
-    bcs = [] if bcs is None else bcs
-    A: la.MatrixCSR = create_matrix(a, block_mode)
+    if bcs is None:
+        bcs = []
+
+    A = create_matrix(a, block_mode)
     _assemble_matrix_csr(A, a, bcs, diag, constants, coeffs)
     return A
 
@@ -326,8 +344,13 @@ def _assemble_matrix_csr(
         accumulated.
     """
     bcs = [] if bcs is None else [bc._cpp_object for bc in bcs]
-    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]
+
+    if constants is None:
+        constants = pack_constants(a)
+
+    if coeffs is None:
+        coeffs = pack_coefficients(a)
+
     _cpp.fem.assemble_matrix(A._cpp_object, a._cpp_object, constants, coeffs, bcs)
 
     # If matrix is a 'diagonal'block, set diagonal entry for constrained
@@ -346,8 +369,8 @@ def apply_lifting(
     bcs: Sequence[Sequence[DirichletBC]],
     x0: Sequence[npt.NDArray] | None = None,
     alpha: float = 1,
-    constants: npt.NDArray | None = None,
-    coeffs: dict[tuple[IntegralType, int], npt.NDArray] | None = None,
+    constants: Sequence[npt.NDArray] | None = None,
+    coeffs: Sequence[dict[tuple[IntegralType, int], npt.NDArray]] | None = None,
 ) -> None:
     """Modify right-hand side for lifting of Dirichlet conditions.
 
@@ -442,40 +465,17 @@ def apply_lifting(
         function. Use :func:`dolfinx.fem.DirichletBC.set` to set values
         in ``b``.
     """  # noqa: D301
-    x0 = [] if x0 is None else x0
-    constants = (
-        [pack_constants(form) if form is not None else np.array([], dtype=b.dtype) for form in a]  # type: ignore[assignment]
-        if constants is None
-        else constants
-    )
-    coeffs = (
-        [{} if form is None else pack_coefficients(form) for form in a]  # type: ignore[assignment]
-        if coeffs is None
-        else coeffs
-    )
-    _a = [None if form is None else form._cpp_object for form in a]
-    _bcs = [[bc._cpp_object for bc in bcs0] for bcs0 in bcs]
-    _cpp.fem.apply_lifting(b, _a, constants, coeffs, _bcs, x0, alpha)
+    if x0 is None:
+        x0 = []
 
+    if constants is None:
+        constants = [
+            pack_constants(form) if form is not None else np.array([], dtype=b.dtype) for form in a
+        ]
 
-def set_bc(
-    b: npt.NDArray,
-    bcs: Sequence[DirichletBC],
-    x0: npt.NDArray | None = None,
-    alpha: float = 1,
-) -> None:
-    """Insert boundary condition values into vector.
-
-    Note:
-        This function is deprecated.
+    if coeffs is None:
+        coeffs = [pack_coefficients(form) for form in a]
 
-    Only local (owned) entries are set, hence communication after
-    calling this function is not required unless ghost entries need to
-    be updated to the boundary condition value.
-    """
-    warnings.warn(
-        "dolfinx.fem.assembler.set_bc is deprecated. Use dolfinx.fem.DirichletBC.set instead.",
-        DeprecationWarning,
-    )
-    for bc in bcs:
-        bc.set(b, x0, alpha)
+    _a = [None if form is None else form._cpp_object for form in a]
+    _bcs = [[bc._cpp_object for bc in bcs0] for bcs0 in bcs]
+    _cpp.fem.apply_lifting(b, _a, constants, coeffs, _bcs, x0, alpha)

python/dolfinx/fem/petsc.py

@@ -609,8 +609,8 @@ def apply_lifting(
                     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 = pack_constants(a_) if constants is None else constants[i]  # type: ignore[arg-type, call-overload]
+                        coeff = pack_coefficients(a_) if coeffs is None else coeffs[i]  # type: ignore[arg-type, assignment, index, call-overload]
                         const_ = [
                             np.empty(0, dtype=PETSc.ScalarType) if val is None else val  # type: ignore[attr-defined]
                             for val in const