Mypy fixes

.github/workflows/test_package.yml

@@ -15,7 +15,15 @@ on:
 jobs:
   test-code:
     runs-on: ubuntu-latest
-    container: ghcr.io/fenics/dolfinx/dolfinx:nightly
+    container: ghcr.io/fenics/dolfinx/dolfinx:${{ matrix.label }}
+    strategy:
+      fail-fast: false
+      matrix:
+        label: [
+          "stable",
+          "nightly"
+        ]
+
     env:
       OMPI_ALLOW_RUN_AS_ROOT: 1
       OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1

src/dolfinx_adjoint/blocks/function_assigner.py

@@ -1,4 +1,3 @@
-
 import dolfinx
 import numpy as np
 import numpy.typing as npt

src/dolfinx_adjoint/petsc_utils.py

@@ -99,7 +99,7 @@ def solve(
             except RuntimeError:
                 bcs0 = dolfinx.fem.bcs.bcs_by_block(dolfinx.fem.forms.extract_spaces(self._L), self.bcs)  # type: ignore
                 dolfinx.fem.petsc.set_bc(self._b, bcs0, alpha=0.0)
-
+        self._b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)  # type: ignore
         # 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

src/dolfinx_adjoint/solvers.py

@@ -9,12 +9,12 @@
 import ufl
 from dolfinx.fem.function import Function as _Function
 
-from .blocks.solvers import LinearProblemBlock, NonlinearProblemBlock, _ProblemBlockBase
+from .blocks.solvers import LinearProblemBlock, NonlinearProblemBlock, _ProblemBlockBase, collect_coefficients
 from .petsc_utils import HomogeneousBCLinearProblem
 from .types import Function
+from .typing_utils import NestedSequence
 from .ufl_utils import (
     assign_mixed_parts,
-    collect_coefficients,
     compute_adjoint,
     get_sorted_arguments,
     recursive_replace,
@@ -107,13 +107,13 @@ class HessianTemplates(typing.NamedTuple):
             per-dependency (not per-row) shape as ``fixed``.
     """
 
-    soa_self: dolfinx.fem.Form | list[dolfinx.fem.Form]
+    soa_self: NestedSequence[dolfinx.fem.Form]
     soa_cross: dict
     fixed: dict
     cross: dict
 
 
-def _pad_blocks_by_part(form: ufl.form.BaseForm, test_funcs: typing.Sequence[ufl.Argument]) -> list[ufl.form.BaseForm]:
+def _pad_blocks_by_part(form: ufl.Form, test_funcs: typing.Sequence[ufl.Argument]) -> list[ufl.Form | ufl.ZeroBaseForm]:
     """Split a blocked one-form into one entry per ``test_funcs`` part, in part order.
 
     {py:func}`ufl.extract_blocks` only returns an entry for a part that actually appears in
@@ -128,9 +128,10 @@ def _pad_blocks_by_part(form: ufl.form.BaseForm, test_funcs: typing.Sequence[ufl
     no parts for {py:func}`ufl.extract_blocks` to find, so every row is padded to zero directly
     instead.
     """
-    padded: list[ufl.form.BaseForm] = [ufl.ZeroBaseForm((test,)) for test in test_funcs]
+    padded: list[ufl.Form | ufl.ZeroBaseForm] = [ufl.ZeroBaseForm((test,)) for test in test_funcs]
     if form.empty():
         return padded
+    assert isinstance(form, ufl.Form)
     for block in ufl.extract_blocks(form):
         args = block.arguments()
         assert len(args) == 1, "Expected a single test function in the block."
@@ -147,7 +148,7 @@ def _build_soa_self_template(
     jit_options: dict | None,
     form_compiler_options: dict | None,
     entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None,
-) -> dolfinx.fem.Form:
+) -> NestedSequence[dolfinx.fem.Form]:
     """Build the SOA self-term ``adjoint(d2F/du2) . adjoint_solution``.
 
     The same computation for both {py:class}`~dolfinx_adjoint.LinearProblem` and
@@ -164,7 +165,7 @@ def _build_soa_self_template(
         soa_self_form = ufl.ZeroBaseForm((dFdu_template.arguments()[0],))
     else:
         soa_self_form = ufl.action(ufl.adjoint(d2Fdu2), adjoint_solution_placeholder)
-    return dolfinx.fem.form(
+    return dolfinx.fem.form(  # type: ignore[call-overload]
         soa_self_form,
         jit_options=jit_options,
         form_compiler_options=form_compiler_options,
@@ -264,7 +265,7 @@ def _init_adjoint_state(self) -> None:
         self._tlm_solver: HomogeneousBCLinearProblem | None = None
         self._residual_template: ufl.Form | None = None
         self._residual_state_placeholder: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] | None = None
-        self._dFdu_template: ufl.Form | typing.Sequence | None = None
+        self._dFdu_template: ufl.Form | None = None
         self._dFdu_adj_template: ufl.Form | typing.Sequence | None = None
         self._tlm_rhs_templates: dict | None = None
         self._tlm_seed_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = {}
@@ -296,7 +297,7 @@ def _get_or_build_residual_template(
             a blocked problem).
         """
 
-    def _get_or_build_dFdu_template(self) -> ufl.Form | typing.Sequence:
+    def _get_or_build_dFdu_template(self) -> ufl.Form:
         """Build (once) and return dF/du, evaluated at the residual template's state placeholder."""
         # Shared by both classes: derived from _get_or_build_residual_template by
         # symbolic differentiation (free at compile time) rather than
@@ -417,7 +418,7 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates:
             assert isinstance(dFdu_adj_template, ufl.Form)
 
             blocked = isinstance(self._u, list)
-            soa_self: dolfinx.fem.Form | list[dolfinx.fem.Form]
+            soa_self: NestedSequence[dolfinx.fem.Form]
             if blocked:
                 # One placeholder Function per output block, mirroring how
                 # _get_or_build_hessian_templates's scalar branch below uses a
@@ -451,7 +452,7 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates:
                 # happened to eliminate entirely.
                 soa_self = [
                     dolfinx.fem.form(
-                        form_i,
+                        form_i,  # type: ignore[arg-type]
                         jit_options=self._jit_options,
                         form_compiler_options=self._form_compiler_options,
                         entity_maps=self._entity_maps,
@@ -507,7 +508,7 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates:
                     if blocked:
                         soa_cross_templates[c] = [
                             dolfinx.fem.form(
-                                form_i,
+                                form_i,  # type: ignore[arg-type]
                                 jit_options=self._jit_options,
                                 form_compiler_options=self._form_compiler_options,
                                 entity_maps=self._entity_maps,
@@ -790,9 +791,10 @@ def __init__(
             coefficients |= collect_coefficients(P)
             if set(u_list).issubset(coefficients):
                 raise ValueError("The unknown `u` should not be part of the coefficients of a linear problem.")
-
+        # Has to be sorted when creating placeholders, as function creation is a collective operation
+        sorted_coefficients = sorted(coefficients, key=lambda c: c.ufl_id())
         self._value_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = {
-            c: dolfinx.fem.Function(c.function_space) for c in coefficients
+            c: dolfinx.fem.Function(c.function_space) for c in sorted_coefficients
         }
         a_R, L_R, P_R = recursive_replace((a, L, P), self._value_placeholders)  # type: ignore[misc]
         super().__init__(

src/dolfinx_adjoint/types/dirichletbc.py

@@ -50,7 +50,7 @@ def __init__(self, g: Function, dofs: npt.NDArray[np.int32], **kwargs):
         bc_kwargs: dict[str, Any] = {}
         # If dolfinx-version is 0.12 we need to pass the following
         # due to https://github.com/FEniCS/dolfinx/pull/4342/
-        if Version(dolfinx.__version__).minor >= 11:
+        if Version(dolfinx.__version__).minor > 11:
             bc_kwargs["V"] = g.function_space
             bc_kwargs["g"] = g
 

src/dolfinx_adjoint/types/function.py

@@ -13,6 +13,7 @@
     register_overloaded_type,
 )
 from pyadjoint.tape import no_annotations
+from ufl.core.ufl_id import attach_ufl_id
 
 from ..blocks._vector import _SpecialVector, _vector
 from ..blocks.assembly import assemble_compiled_form
@@ -35,6 +36,7 @@ def _create_function(
     return Function(V, x=x, annotate=False)
 
 
+@attach_ufl_id
 class Function(dolfinx.fem.Function, FloatingType):
     """A class overloading `dolfinx.fem.Function` to support it being used as a control variable
     in the adjoint framework.
@@ -56,6 +58,8 @@ def __init__(
         dtype: npt.DTypeLike = dolfinx.default_scalar_type,
         **kwargs,
     ):
+        ufl_id = kwargs.pop("ufl_id", None)
+        self._ufl_id = self._init_ufl_id(ufl_id)
         super(Function, self).__init__(
             V,
             x,
@@ -249,6 +253,7 @@ def x(self) -> dolfinx.la.Vector:
         return self._x
 
 
+@attach_ufl_id
 class Constant(Function):
     """A class overloading {py:class}`dolfinx.fem.Constant`
     to support it being used as a control variable in
@@ -274,7 +279,10 @@ def __init__(
         self,
         domain: dolfinx.mesh.Mesh,
         c: float | numpy.floating | complex | numpy.complexfloating | typing.Sequence | numpy.ndarray,
+        ufl_id: int | None = None,
     ):
+        self._ufl_id = self._init_ufl_id(ufl_id)
+
         value_shape = numpy.shape(c)
         try:
             el = basix.ufl.real_element(domain.basix_cell(), value_shape=numpy.shape(c))

src/dolfinx_adjoint/ufl_utils.py

@@ -9,7 +9,7 @@
 
 
 def recursive_space_discovery(
-    obj: NestedSequence[ufl.Form], indices: tuple[int, ...], spaces: dict[int, ufl.FunctionSpace]
+    obj: NestedSequence[ufl.Form | None], indices: tuple[int, ...], spaces: dict[int, ufl.FunctionSpace]
 ) -> None:
     """Recursively discover, for each row/column index, the function space of the
     (as yet unassigned) argument occupying that position.
@@ -34,15 +34,17 @@ def recursive_space_discovery(
         for i, item in enumerate(obj):
             if item is not None:
                 recursive_space_discovery(item, indices + (i,), spaces)
+    elif obj is None:
+        return
     else:
         raise TypeError(f"Expected ufl.Form or iterable, got {type(obj)}")
 
 
 def build_argument_replacement_map(
-    obj: NestedSequence[ufl.Form],
+    obj: NestedSequence[ufl.Form | None],
     indices: tuple[int, ...],
-    test_functions: typing.Sequence[ufl.TestFunction],
-    trial_functions: typing.Sequence[ufl.TrialFunction],
+    test_functions: typing.Sequence[ufl.Argument],
+    trial_functions: typing.Sequence[ufl.Argument],
     replace_map: dict[ufl.Argument, ufl.Argument],
 ) -> None:
     """
@@ -66,17 +68,21 @@ def build_argument_replacement_map(
         for i, item in enumerate(obj):
             if item is not None:
                 build_argument_replacement_map(item, indices + (i,), test_functions, trial_functions, replace_map)
+    elif obj is None:
+        return
+    else:
+        raise TypeError(f"Expected ufl.Form or iterable, got {type(obj)}")
 
 
 @typing.overload
-def assign_mixed_parts[T: NestedSequence[ufl.Form]](form1: T, /) -> T: ...
+def assign_mixed_parts[T: NestedSequence[ufl.Form | None]](form1: T, /) -> T: ...
 @typing.overload
-def assign_mixed_parts[T: NestedSequence[ufl.Form], S: NestedSequence[ufl.Form]](
+def assign_mixed_parts[T: NestedSequence[ufl.Form | None], S: NestedSequence[ufl.Form | None]](
     form1: T, form2: S, /
 ) -> tuple[T, S]: ...
 def assign_mixed_parts(
-    *form_structs: NestedSequence[ufl.Form],
-) -> NestedSequence[ufl.Form] | tuple[NestedSequence[ufl.Form], ...]:
+    *form_structs: NestedSequence[ufl.Form | None],
+) -> NestedSequence[ufl.Form | None] | tuple[NestedSequence[ufl.Form | None], ...]:
     """
     Recursively assigns mixed-space `part` indices to {py:class}`ufl.Argument`
     (test and trial functions), within nested iterables of forms.
@@ -106,7 +112,7 @@ def assign_mixed_parts(
         {py:class}`ufl.MixedFunctionSpace` built from the row/column function spaces
         discovered while walking the structure.
     """
-    spaces: dict[int, ufl.functionspace.AbstractFunctionSpace] = {}
+    spaces: dict[int, ufl.FunctionSpace] = {}
     for struct in form_structs:
         recursive_space_discovery(struct, (), spaces)
 
@@ -133,26 +139,6 @@ def get_sorted_arguments(arguments: typing.Iterable[ufl.Argument], number: int)
     return sorted(filter(lambda x: x.number() == number, arguments), key=lambda a: a.part())
 
 
-def collect_coefficients(form: ufl.Form | typing.Sequence | None) -> set[ufl.Coefficient]:
-    """Return the set of UFL coefficients appearing anywhere in ``form``.
-
-    ``form`` may be a single form or an arbitrarily nested sequence of forms
-    (entries may be ``None``, e.g. a zero block in a blocked system). Plain set
-    union rather than ``sum_form``: unlike summing, this never requires the
-    sub-forms' arguments to be mutually compatible (e.g. carry matching
-    ``part()`` tags), which a blocked ``NonlinearProblem``'s forms are not
-    required to be before ``assign_mixed_parts`` runs.
-    """
-    if form is None:
-        return set()
-    if isinstance(form, ufl.Form):
-        return set(form.coefficients())
-    coefficients: set = set()
-    for f in form:
-        coefficients |= collect_coefficients(f)
-    return coefficients
-
-
 def sum_form(form: NestedSequence[ufl.Form | None]) -> ufl.Form | None:
     """Sum a blocked form into a single form."""
     # Handle top-level None
@@ -195,7 +181,7 @@ def compute_adjoint(form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl.Form]
     return ufl.extract_blocks(compute_form_adjoint(form))
 
 
-def recursive_replace(form: ufl.Form | typing.Sequence | None, placeholders: dict) -> ufl.Form | typing.Sequence | None:
+def recursive_replace(form: NestedSequence[ufl.Form | None], placeholders: dict) -> NestedSequence[ufl.Form | None]:
     """Recursively apply {py:func}`ufl.replace` to a (possibly nested) form structure.
 
     Args: