diff --git a/pyproject.toml b/pyproject.toml index 0c7abe4..cf59b1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,8 @@ dependencies = [ [project.optional-dependencies] -test = ["pytest"] +scifem = ["scifem"] +test = ["pytest", "dolfinx-adjoint[scifem]"] dev = ["pdbpp", "ipython", "mypy", "ruff"] docs = [ "jupyter-book<2.0", @@ -29,7 +30,7 @@ docs = [ "networkx", "pygraphviz" ] -all = ["dolfinx-adjoint[test]", "dolfinx-adjoint[dev]", "dolfinx-adjoint[docs]"] +all = ["dolfinx-adjoint[test]", "dolfinx-adjoint[dev]", "dolfinx-adjoint[docs]", "dolfinx-adjoint[scifem]"] [tool.pytest.ini_options] addopts = ["--import-mode=importlib"] diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py index c81d81c..6a62daf 100644 --- a/src/dolfinx_adjoint/blocks/interpolation.py +++ b/src/dolfinx_adjoint/blocks/interpolation.py @@ -1,23 +1,59 @@ from __future__ import annotations import typing +import weakref from typing import Callable import dolfinx +import dolfinx.fem.petsc +import ufl from pyadjoint import Block, OverloadedType -from pyadjoint.overloaded_type import create_overloaded_object from pyadjoint.tape import stop_annotating - -import ufl -import scifem from ufl.algorithms.analysis import traverse_unique_terminals + from ..compat import get_interpolation_points if typing.TYPE_CHECKING: from petsc4py import PETSc -# Global cache to prevent redundant matrix assembly across multiple blocks/iterations -_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int], dolfinx.la.MatrixCSR] = {} + +def _import_scifem(): + """Import scifem lazily: only ExprInterpolationBlock needs it, so importing + dolfinx_adjoint (or using InterpolationBlock) must not require it to be installed. + """ + try: + import scifem + except ImportError as e: + raise ImportError("scifem is required to interpolate a UFL expression: pip install scifem") from e + return scifem + + +# Cache of assembled interpolation matrices, keyed by (id(space_from), id(space_to), +# use_petsc), shared across all InterpolationBlocks to avoid redundant matrix assembly +# (and, for PETSc matrices, MPI communicator exhaustion) when the same pair of spaces is +# interpolated between many times (e.g. across optimization iterations). +# +# Keying on id() would normally risk a stale/wrong entry once a FunctionSpace is garbage +# collected and its id is reused by an unrelated object (ids are only unique among +# simultaneously-alive objects). We close that hole with weakref.finalize: each space +# that contributes to a cache key gets a finalizer that purges every entry mentioning +# its id. Finalizers run at the point the object is actually deallocated, which is +# necessarily before CPython can hand that id out again, so a cache hit always +# corresponds to spaces that are still alive. +_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int, bool], "_MatrixCSRWorkspace | PETSc.Mat"] = {} +_CACHE_KEYS_BY_SPACE_ID: dict[int, set[tuple[int, int, bool]]] = {} + + +def _purge_cache_entries_for(space_id: int) -> None: + for key in _CACHE_KEYS_BY_SPACE_ID.pop(space_id, ()): + _INTERPOLATION_MATRIX_CACHE.pop(key, None) + + +def _register_cache_key(space: dolfinx.fem.FunctionSpace, key: tuple[int, int, bool]) -> None: + space_id = id(space) + if space_id not in _CACHE_KEYS_BY_SPACE_ID: + weakref.finalize(space, _purge_cache_entries_for, space_id) + _CACHE_KEYS_BY_SPACE_ID.setdefault(space_id, set()).add(key) def get_dependency_index(dependencies, dep_or_bv) -> int: @@ -44,48 +80,56 @@ def get_dependency_index(dependencies, dep_or_bv) -> int: raise ValueError(f"Could not locate dependency index for {dep_or_bv}") -def attach_working_array(mat: dolfinx.la.MatrixCSR): - """Attach working arrays to a dolfinx.la.MatrixCSR for efficient matrix-vector multiplication.""" - if not hasattr(mat, "_row_vec"): - mat._row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype) - if not hasattr(mat, "_col_vec"): - mat._col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype) - mat._row_vec.array[:] = 0.0 - mat._col_vec.array[:] = 0.0 +class _MatrixCSRWorkspace: + """A dolfinx.la.MatrixCSR paired with pre-allocated working vectors. + + dolfinx.la.MatrixCSR.mult requires vectors built from its own index maps + as scratch space. Wrapping them here (built once, alongside the matrix) + means they can be reused across repeated matrix-vector multiplications + without allocating on every adjoint/TLM/Hessian evaluation, and without + monkey-patching dynamic attributes onto the matrix object itself. + """ + + def __init__(self, mat: dolfinx.la.MatrixCSR): + self.mat = mat + self.row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype) + self.col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype) def get_mult( - mat: "PETSc.Mat" | dolfinx.la.MatrixCSR, + mat: "PETSc.Mat" | _MatrixCSRWorkspace, transpose: bool = False, accumulate: bool = False, ) -> Callable[[dolfinx.la.Vector, dolfinx.la.Vector], None]: """Return a function that performs matrix-vector multiplication, optionally accumulating results.""" - if isinstance(mat, dolfinx.la.MatrixCSR): + if isinstance(mat, _MatrixCSRWorkspace): + workspace = mat def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector): in_size_local = v_in.index_map.size_local * v_in.block_size out_size_local = v_out.index_map.size_local * v_out.block_size - attach_working_array(mat) - + # Zero the full working arrays (including ghosts) before each use + # to prevent double-counting in parallel. + workspace.row_vec.array[:] = 0.0 + workspace.col_vec.array[:] = 0.0 if transpose: - mat._row_vec.array[:in_size_local] = v_in.array[:in_size_local] - mat._row_vec.scatter_forward() - mat._col_vec.array[:out_size_local] = 0.0 - mat.mult(mat._row_vec, mat._col_vec, transpose=True) + workspace.row_vec.array[:in_size_local] = v_in.array[:in_size_local] + workspace.row_vec.scatter_forward() # Ensure ghost values are updated before multiplication + workspace.col_vec.array[:out_size_local] = 0.0 + workspace.mat.mult(workspace.row_vec, workspace.col_vec, transpose=True) if accumulate: - v_out.array[:out_size_local] += mat._col_vec.array[:out_size_local] + v_out.array[:out_size_local] += workspace.col_vec.array[:out_size_local] else: - v_out.array[:out_size_local] = mat._col_vec.array[:out_size_local] + v_out.array[:out_size_local] = workspace.col_vec.array[:out_size_local] else: - mat._row_vec.array[:out_size_local] = 0.0 - mat._col_vec.array[:in_size_local] = v_in.array[:in_size_local] - mat._col_vec.scatter_forward() - mat.mult(mat._col_vec, mat._row_vec) + workspace.row_vec.array[:out_size_local] = 0.0 + workspace.col_vec.array[:in_size_local] = v_in.array[:in_size_local] + workspace.col_vec.scatter_forward() # Ensure ghost values are updated before multiplication + workspace.mat.mult(workspace.col_vec, workspace.row_vec) if accumulate: - v_out.array[:out_size_local] += mat._row_vec.array[:out_size_local] + v_out.array[:out_size_local] += workspace.row_vec.array[:out_size_local] else: - v_out.array[:out_size_local] = mat._row_vec.array[:out_size_local] - + v_out.array[:out_size_local] = workspace.row_vec.array[:out_size_local] v_out.scatter_forward() return mult @@ -111,26 +155,32 @@ def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector): else: raise TypeError("Expected a PETSc.Mat when PETSc is available.") else: - raise TypeError("Matrix type not supported.") + raise TypeError(f"Matrix type not supported. Expected dolfinx.la.MatrixCSR or PETSc.Mat, got {type(mat)=}.") -def _get_interpolation_matrix( +def _build_interpolation_matrix( space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False -) -> dolfinx.la.MatrixCSR | "PETSc.Mat": - """Retrieve or compute the interpolation matrix for a pair of spaces.""" - key = (id(space_from), id(space_to)) +) -> _MatrixCSRWorkspace | "PETSc.Mat": + """Assemble the interpolation matrix for a pair of spaces.""" + if use_petsc: + petsc_mat = dolfinx.fem.petsc.interpolation_matrix(space_from, space_to) + petsc_mat.assemble() + return petsc_mat - if key not in _INTERPOLATION_MATRIX_CACHE: - if use_petsc: - mat = dolfinx.fem.petsc.interpolation_matrix(space_from, space_to) - mat.assemble() - else: - mat = dolfinx.fem.interpolation_matrix(space_from, space_to) - mat.scatter_reverse() - attach_working_array(mat) + mat = dolfinx.fem.interpolation_matrix(space_from, space_to) + mat.scatter_reverse() + return _MatrixCSRWorkspace(mat) - _INTERPOLATION_MATRIX_CACHE[key] = mat +def _get_interpolation_matrix( + space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False +) -> _MatrixCSRWorkspace | "PETSc.Mat": + """Retrieve or assemble the interpolation matrix for a pair of spaces, cached for reuse.""" + key = (id(space_from), id(space_to), use_petsc) + if key not in _INTERPOLATION_MATRIX_CACHE: + _INTERPOLATION_MATRIX_CACHE[key] = _build_interpolation_matrix(space_from, space_to, use_petsc=use_petsc) + _register_cache_key(space_from, key) + _register_cache_key(space_to, key) return _INTERPOLATION_MATRIX_CACHE[key] @@ -151,18 +201,21 @@ def __init__( self.add_dependency(func_from) + # Initialize internal caches for outputs to avoid MPI communicator exhaustion self._adj_output: dolfinx.fem.Function | None = None self._tlm_output: dolfinx.fem.Function | None = None self._hessian_output: dolfinx.fem.Function | None = None - self._recompute_output: dolfinx.fem.Function | None = None def __str__(self): return "interpolate_function" + def _get_interpolation_matrix(self) -> _MatrixCSRWorkspace | "PETSc.Mat": + return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc) + # --- Adjoint --- def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): - return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc) + return self._get_interpolation_matrix() def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): adj_input = adj_inputs[0] @@ -171,15 +224,17 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar if self._adj_output is None: self._adj_output = dolfinx.fem.Function(self.space_from) - self._adj_output.x.array[:] = 0.0 - mult = get_mult(mat, transpose=True, accumulate=False) + # Action of the adjoint: A^T * adj_input + self._adj_output.x.array[:] = 0.0 # Reset the output vector before accumulation + adj_input.x.scatter_forward() + mult = get_mult(mat, transpose=True) mult(adj_input.x, self._adj_output.x) return self._adj_output # --- Tangent Linear Model (TLM) --- def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs): - return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc) + return self._get_interpolation_matrix() def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None): tlm_input = tlm_inputs[0] @@ -191,15 +246,17 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar if self._tlm_output is None: self._tlm_output = dolfinx.fem.Function(self.space_to) - self._tlm_output.x.array[:] = 0.0 - mult = get_mult(mat, transpose=False, accumulate=False) + # Forward Jacobian action: A * tlm_input + tlm_input.x.scatter_forward() + self._tlm_output.x.array[:] = 0.0 # Reset the output vector before accumulation + mult = get_mult(mat, transpose=False) mult(tlm_input.x, self._tlm_output.x) return self._tlm_output # --- Hessian --- def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies): - return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc) + return self._get_interpolation_matrix() def evaluate_hessian_component( self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None @@ -210,8 +267,10 @@ def evaluate_hessian_component( if self._hessian_output is None: self._hessian_output = dolfinx.fem.Function(self.space_from) - self._hessian_output.x.array[:] = 0.0 - mult = get_mult(mat, transpose=True, accumulate=False) + # Action of the adjoint on the incoming Hessian sensitivity + hessian_input.x.scatter_forward() + self._hessian_output.x.array[:] = 0.0 # Reset the output vector before accumulation + mult = get_mult(mat, transpose=True) mult(hessian_input.x, self._hessian_output.x) return self._hessian_output @@ -223,13 +282,20 @@ def prepare_recompute_component(self, inputs, relevant_outputs): def recompute_component(self, inputs, block_variable, idx, prepared): func_from = inputs[0] - if self._recompute_output is None: - self._recompute_output = dolfinx.fem.Function(self.space_to) - - self._recompute_output.interpolate(func_from) - self._recompute_output.x.scatter_forward() - - return create_overloaded_object(self._recompute_output) + # Update the tape's actual output object in-place (rather than a + # separately cached Function) so that any Python reference the user + # is holding to the interpolated Function stays in sync after the + # tape is replayed, matching FunctionAssignBlock's convention. Note + # this only holds until the output is first used as a dependency of + # another block: pyadjoint then freezes a private checkpoint copy + # (see OverloadedType._ad_will_add_as_dependency), so the live + # Python object can still go stale after that point. This is an + # existing pyadjoint/dolfinx_adjoint-wide characteristic (identical + # behavior in FunctionAssignBlock), not specific to interpolation. + output = block_variable.saved_output + output.interpolate(func_from) + output.x.scatter_forward() + return output class ExprInterpolationBlock(Block): @@ -247,6 +313,11 @@ def __init__( self.space_to = func_to.function_space self._use_petsc = petsc_mat + # self._deps and self.get_dependencies() are always populated in lockstep by this + # loop, so position i means the same dependency in both: get_dependency_index is + # only needed for identity lookups from outside this alignment (e.g. tests locating + # a specific control), not for indices pyadjoint already hands back via idx or + # relevant_dependencies in the evaluate_*/prepare_evaluate_* methods below. self._deps = [] for op in traverse_unique_terminals(self.expr): if isinstance(op, OverloadedType): @@ -256,7 +327,6 @@ def __init__( self._adj_output: dict[int, dolfinx.fem.Function] = {} self._tlm_output: dolfinx.fem.Function | None = None self._hessian_output: dict[int, dolfinx.fem.Function] = {} - self._recompute_output: dolfinx.fem.Function | None = None def __str__(self): return f"interpolate_expression_{str(self.expr)}_to_{str(self.space_to)}" @@ -276,11 +346,11 @@ def _assemble_jacobian(self, idx: int, inputs: list | None = None): du = ufl.TrialFunction(V_in) dE = ufl.derivative(current_expr, target_dep, du) dE = ufl.algorithms.apply_derivatives.apply_derivatives(dE) + scifem = _import_scifem() if self._use_petsc: mat = scifem.petsc_interpolation_matrix(dE, self.space_to) else: - mat = scifem.interpolation_matrix(dE, self.space_to) - attach_working_array(mat) + mat = _MatrixCSRWorkspace(scifem.interpolation_matrix(dE, self.space_to)) return mat @@ -288,19 +358,17 @@ def _assemble_jacobian(self, idx: int, inputs: list | None = None): def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): matrices = {} - for dep in relevant_dependencies: - idx = get_dependency_index(self._deps, dep) + for idx, _dep in relevant_dependencies: matrices[idx] = self._assemble_jacobian(idx, inputs) return matrices def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): adj_input = adj_inputs[0] - dep_idx = get_dependency_index(self._deps, block_variable) if block_variable is not None else idx - mat = prepared[dep_idx] + mat = prepared[idx] - if dep_idx not in self._adj_output: - self._adj_output[dep_idx] = dolfinx.fem.Function(self._deps[dep_idx].function_space) - out_func = self._adj_output[dep_idx] + if idx not in self._adj_output: + self._adj_output[idx] = dolfinx.fem.Function(self._deps[idx].function_space) + out_func = self._adj_output[idx] out_func.x.array[:] = 0.0 mult = get_mult(mat, transpose=True, accumulate=False) @@ -312,8 +380,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs): matrices = {} - for dep in self.get_dependencies(): - idx = get_dependency_index(self._deps, dep) + for idx in range(len(self._deps)): matrices[idx] = self._assemble_jacobian(idx, inputs) return matrices @@ -325,7 +392,9 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar # Reset output vector to prepare for accumulation out_func.x.array[:] = 0.0 - # The TLM is the sum of the Jacobians applied to each perturbation + # The TLM is the sum of the Jacobians applied to each perturbation. tlm_inputs is + # aligned with self.get_dependencies(), which is aligned 1:1 with self._deps (see + # the comment in __init__), so its position doubles as the index into `prepared`. for dep_idx, tlm_input in enumerate(tlm_inputs): if tlm_input is None: continue @@ -354,19 +423,11 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ deps_bv = self.get_dependencies() for i, dep_bv in enumerate(deps_bv): - # EXHAUSTIVE SEARCH: Find the forward perturbation wherever PyAdjoint hid it - tlm_val = getattr(dep_bv, "saved_tlm_output", None) - if tlm_val is None: - tlm_val = getattr(dep_bv, "tlm_value", None) + tlm_val = dep_bv.tlm_value - # Fallbacks for controls or unrecorded raw inputs + # Fallback for controls or unrecorded raw inputs if tlm_val is None and hasattr(inputs[i], "block_variable") and inputs[i].block_variable is not None: - tlm_val = getattr(inputs[i].block_variable, "saved_tlm_output", None) - if tlm_val is None: - tlm_val = getattr(inputs[i].block_variable, "tlm_value", None) - - if tlm_val is None: - tlm_val = getattr(inputs[i], "saved_tlm_output", None) + tlm_val = inputs[i].block_variable.tlm_value if tlm_val is not None: target_dep = inputs[i] @@ -374,9 +435,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ dE_total = term if dE_total is None else dE_total + term # 3. Assemble both the Jacobian and the Hessian matrix for each dependency - for dep in relevant_dependencies: - idx = get_dependency_index(self._deps, dep) - + for idx, _dep in relevant_dependencies: # Matrix 1: The standard Jacobian (J) J_mat = self._assemble_jacobian(idx, inputs) @@ -397,11 +456,11 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_ if not isinstance(d2E, (int, float)): args = ufl.algorithms.extract_arguments(d2E) if len(args) > 0: + scifem = _import_scifem() if self._use_petsc: H_mat = scifem.petsc_interpolation_matrix(d2E, self.space_to) else: - H_mat = scifem.interpolation_matrix(d2E, self.space_to) - attach_working_array(H_mat) + H_mat = _MatrixCSRWorkspace(scifem.interpolation_matrix(d2E, self.space_to)) matrices[idx] = (J_mat, H_mat) @@ -413,12 +472,11 @@ def evaluate_hessian_component( hessian_input = hessian_inputs[0] adj_input = adj_inputs[0] # The incoming adjoint sensitivity (y*) - dep_idx = get_dependency_index(self._deps, block_variable) if block_variable is not None else idx - J_mat, H_mat = prepared[dep_idx] + J_mat, H_mat = prepared[idx] - if dep_idx not in self._hessian_output: - self._hessian_output[dep_idx] = dolfinx.fem.Function(self._deps[dep_idx].function_space) - out_func = self._hessian_output[dep_idx] + if idx not in self._hessian_output: + self._hessian_output[idx] = dolfinx.fem.Function(self._deps[idx].function_space) + out_func = self._hessian_output[idx] # Reset output vector to 0.0 before accumulation out_func.x.array[:] = 0.0 @@ -440,15 +498,15 @@ def prepare_recompute_component(self, inputs, relevant_outputs): return None def recompute_component(self, inputs, block_variable, idx, prepared): - if self._recompute_output is None: - self._recompute_output = dolfinx.fem.Function(self.space_to) - replace_map = {self._deps[i]: inputs[i] for i in range(len(self._deps))} updated_expr = ufl.replace(self.expr, replace_map) + # Update the tape's actual output object in-place, matching + # InterpolationBlock.recompute_component and FunctionAssignBlock's convention. + output = block_variable.saved_output with stop_annotating(): compiled_expr = dolfinx.fem.Expression(updated_expr, get_interpolation_points(self.space_to)) - self._recompute_output.interpolate(compiled_expr) - self._recompute_output.x.scatter_forward() + output.interpolate(compiled_expr) + output.x.scatter_forward() - return create_overloaded_object(self._recompute_output) + return output diff --git a/src/dolfinx_adjoint/interpolation.py b/src/dolfinx_adjoint/interpolation.py index a6ccb80..b7d61ac 100644 --- a/src/dolfinx_adjoint/interpolation.py +++ b/src/dolfinx_adjoint/interpolation.py @@ -1,8 +1,9 @@ import dolfinx +import ufl from pyadjoint.overloaded_type import create_overloaded_object from pyadjoint.tape import annotate_tape, get_working_tape, stop_annotating -import ufl -from .blocks.interpolation import InterpolationBlock, ExprInterpolationBlock + +from .blocks.interpolation import ExprInterpolationBlock, InterpolationBlock from .compat import get_interpolation_points diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index 165c36c..70a5114 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -10,7 +10,6 @@ from pyadjoint.overloaded_type import ( FloatingType, create_overloaded_object, - get_overloaded_class, register_overloaded_type, ) from pyadjoint.tape import annotate_tape, get_working_tape, no_annotations, stop_annotating @@ -69,7 +68,12 @@ def _ad_init_object(cls, obj): @no_annotations def _ad_create_checkpoint(self): - checkpoint = create_overloaded_object(self.copy()) + # Note: self.copy() (dolfinx.fem.Function.copy) always returns a plain + # dolfinx.fem.Function regardless of self's concrete type, so wrapping it with + # create_overloaded_object would silently downcast a Constant checkpoint to a + # plain Function. Use _ad_new_like() instead to preserve the concrete subclass. + checkpoint = self._ad_new_like() + checkpoint.x.array[:] = self.x.array[:] checkpoint.name = self.name + "_checkpoint" return checkpoint @@ -109,16 +113,29 @@ def _ad_dot(self, other: typing.Self, options: typing.Optional[dict] = None): else: raise NotImplementedError("Unknown Riesz representation %s" % riesz_representation) + def _ad_new_like(self) -> typing.Self: + """Create a new, zero-valued instance sharing this object's exact overloaded type and + function space. + + Constructing via ``type(self)(...)`` directly does not work here because subclasses + such as ``Constant`` take a different constructor signature (mesh and value, not a + function space). Going through ``__new__`` and ``Function.__init__`` bypasses that + constructor while still producing an instance of the correct concrete subclass. + """ + r = type(self).__new__(type(self), self.function_space) # type: ignore[call-arg] + Function.__init__(r, self.function_space) + return r + @no_annotations def _ad_mul(self, other: typing.Union[int, float]) -> typing.Self: """Multiplication of self with integer or floating value.""" - r = get_overloaded_class(dolfinx.fem.Function)(self.function_space) + r = self._ad_new_like() r.x.array[:] = self.x.array * other return r @no_annotations def _ad_add(self, other: typing.Self) -> typing.Self: - r = get_overloaded_class(dolfinx.fem.Function)(self.function_space) + r = self._ad_new_like() r.x.array[:] = self.x.array[:] + other.x.array[:] return r @@ -188,7 +205,7 @@ def _ad_to_list(m): def _ad_copy(self): """Create a (deep) copy of the function.""" - r = get_overloaded_class(dolfinx.fem.Function)(self.function_space) + r = self._ad_new_like() assign(self, r) return r diff --git a/tests/test_interpolate.py b/tests/test_interpolate.py index 1b797c7..640f8e7 100644 --- a/tests/test_interpolate.py +++ b/tests/test_interpolate.py @@ -1,3 +1,5 @@ +import gc + from mpi4py import MPI import dolfinx @@ -6,8 +8,14 @@ import pytest import ufl -from dolfinx_adjoint import Function, assemble_scalar, interpolate, Constant -from dolfinx_adjoint.blocks.interpolation import InterpolationBlock, ExprInterpolationBlock +from dolfinx_adjoint import Constant, Function, assemble_scalar, interpolate +from dolfinx_adjoint.blocks.interpolation import ( + _CACHE_KEYS_BY_SPACE_ID, + _INTERPOLATION_MATRIX_CACHE, + ExprInterpolationBlock, + InterpolationBlock, + _get_interpolation_matrix, +) # Dynamically determine available matrix backends petsc_options = [False] @@ -151,7 +159,7 @@ def u_ex(mod, x_coords): # ============================================================================== -# Test 1: Expression Interpolation Adjoint Property ( == ) +# Test 3: Expression Interpolation Adjoint Property ( == ) # ============================================================================== @@ -178,7 +186,7 @@ def test_expr_interpolation_adjoint_property(mesh_2D, use_petsc): v = Function(V_to, name="v_output") v.x.array[:] = rng.random(len(v.x.array)) - from dolfinx_adjoint.blocks.interpolation import ExprInterpolationBlock, get_dependency_index + from dolfinx_adjoint.blocks.interpolation import get_dependency_index block = ExprInterpolationBlock(expr, v, petsc_mat=use_petsc) @@ -195,8 +203,10 @@ def test_expr_interpolation_adjoint_property(mesh_2D, use_petsc): aligned_tlm_inputs[u_idx] = u # --- Test Tangent Linear Model and Adjoint --- + # relevant_dependencies mirrors pyadjoint's real contract (Block.evaluate_adj): a list of + # (idx, block_variable) tuples, where idx is the position in block.get_dependencies(). mat_tlm = block.prepare_evaluate_tlm(aligned_inputs, aligned_tlm_inputs, None) - mat_adj = block.prepare_evaluate_adj(aligned_inputs, [v], [u_bv]) + mat_adj = block.prepare_evaluate_adj(aligned_inputs, [v], [(u_idx, u_bv)]) tlm_output = block.evaluate_tlm_component( inputs=aligned_inputs, tlm_inputs=aligned_tlm_inputs, block_variable=None, idx=0, prepared=mat_tlm @@ -224,7 +234,7 @@ def test_expr_interpolation_adjoint_property(mesh_2D, use_petsc): # ============================================================================== -# Test 2: Taylor Test with Function and Constant Controls +# Test 4: Taylor Test with Function and Constant Controls # ============================================================================== @@ -274,6 +284,7 @@ def test_expr_interpolation_taylor_test_function(mesh_2D, use_petsc): min_rate = pyadjoint.taylor_test(Jh_u, u, du) assert np.isclose(min_rate, 2.0, rtol=1e-2, atol=1e-2) + Jh_u(u) dJdm_u = Jh_u.derivative()._ad_dot(du) hessian_u = Jh_u.hessian(du) dHddu_u = hessian_u._ad_dot(du) @@ -319,6 +330,7 @@ def test_expr_interpolation_taylor_test_constant(mesh_2D, use_petsc): min_rate = pyadjoint.taylor_test(Jh_c, c, dc) assert np.isclose(min_rate, 2.0, rtol=1e-2, atol=1e-2) + Jh_c(c) dJdm_c = Jh_c.derivative()._ad_dot(dc) hessian_c = Jh_c.hessian(dc) dHddu_c = hessian_c._ad_dot(dc) @@ -327,3 +339,61 @@ def test_expr_interpolation_taylor_test_constant(mesh_2D, use_petsc): assert np.isclose(min_rate, 3.0, rtol=1e-3, atol=1e-3) pyadjoint.get_working_tape().clear_tape() + + +# ============================================================================== +# Test 5: Interpolation Matrix Cache +# ============================================================================== + + +def test_interpolation_matrix_cache_reused_for_same_spaces(mesh_1D): + """The same (space_from, space_to) pair should reuse one assembled matrix.""" + V_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1)) + V_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2)) + + mat1 = _get_interpolation_matrix(V_from, V_to) + mat2 = _get_interpolation_matrix(V_from, V_to) + assert mat1 is mat2 + + # A different space pair must not share the cached matrix. + V_other = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 3)) + mat3 = _get_interpolation_matrix(V_from, V_other) + assert mat3 is not mat1 + + +def test_interpolation_matrix_cache_purged_when_space_is_garbage_collected(mesh_1D): + """Regression test for the id()-collision bug: a cache entry keyed on + id(space) must be dropped once that space is garbage collected, otherwise + a later, unrelated FunctionSpace object that happens to reuse the same id + could silently be served a stale matrix built for a completely different + pair of spaces. + """ + V_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2)) + + def build_and_return_id(): + # V_from only lives inside this function, so it is eligible for + # collection as soon as it returns. + v_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1)) + _get_interpolation_matrix(v_from, V_to) + return id(v_from) + + space_id = build_and_return_id() + gc.collect() + + assert space_id not in _CACHE_KEYS_BY_SPACE_ID + assert all(key[0] != space_id for key in _INTERPOLATION_MATRIX_CACHE) + + +def test_interpolation_matrix_cache_does_not_leak_across_transient_spaces(mesh_1D): + """The cache must not grow without bound as short-lived FunctionSpaces + (e.g. created once per optimization iteration) are used and discarded. + """ + baseline = len(_INTERPOLATION_MATRIX_CACHE) + + for _ in range(20): + v_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1)) + v_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2)) + _get_interpolation_matrix(v_from, v_to) + + gc.collect() + assert len(_INTERPOLATION_MATRIX_CACHE) <= baseline + 1