From ebe2cce22a9434f52e457e7de0fcca47ce96b7b1 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 21 Aug 2026 13:31:00 +0000 Subject: [PATCH 1/3] Fix unsafe id()-keyed matrix cache and stale recompute output in InterpolationBlock The interpolation matrix cache was a module-level dict keyed by (id(space_from), id(space_to)) without holding strong references, so a garbage-collected FunctionSpace's id could be reused by an unrelated object and silently return the wrong matrix. Cache the matrix as a per-instance attribute on InterpolationBlock instead, tied to the block's own space references. recompute_component also wrapped a separately cached Function instead of updating the tape's real output object, permanently disconnecting it from the Python object returned by interpolate(). It now mutates block_variable.saved_output in place, matching FunctionAssignBlock's convention. Note this only holds until the output is used as a dependency elsewhere, at which point pyadjoint freezes its own checkpoint copy (an existing, codebase-wide characteristic, not specific to interpolation) - documented inline. Also fixes the mypy failures in this file: the missing `dolfinx.fem.petsc` import that `petsc_mat=True` silently relied on via import order, and the monkey-patched `_row_vec`/`_col_vec` attributes on MatrixCSR. Co-Authored-By: Claude Sonnet 5 --- src/dolfinx_adjoint/blocks/interpolation.py | 109 +++++++++++--------- 1 file changed, 61 insertions(+), 48 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py index 8cf8b58..e26a5f2 100644 --- a/src/dolfinx_adjoint/blocks/interpolation.py +++ b/src/dolfinx_adjoint/blocks/interpolation.py @@ -4,24 +4,24 @@ from typing import Callable import dolfinx +import dolfinx.fem.petsc from pyadjoint import Block -from pyadjoint.overloaded_type import create_overloaded_object 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 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 + # mat._row_vec/_col_vec are monkey-patched on; cast to Any so mypy doesn't + # flag the dynamic attributes. + m = typing.cast(typing.Any, mat) + if not hasattr(m, "_row_vec"): + m._row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype) + if not hasattr(m, "_col_vec"): + m._col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype) + m._row_vec.array[:] = 0.0 + m._col_vec.array[:] = 0.0 def get_mult( @@ -36,22 +36,23 @@ 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) # Ensure working arrays are attached + m = typing.cast(typing.Any, mat) if transpose: # Prevent double-counting in parallel by zeroing ghosts of input vector # Calculate the exact number of local degrees of freedom - mat._row_vec.array[:in_size_local] = v_in.array[:in_size_local] - mat._row_vec.scatter_forward() # Ensure ghost values are updated before multiplication - mat._col_vec.array[:out_size_local] = 0.0 - mat.mult(mat._row_vec, mat._col_vec, transpose=True) - v_out.array[:out_size_local] = mat._col_vec.array[:out_size_local] + m._row_vec.array[:in_size_local] = v_in.array[:in_size_local] + m._row_vec.scatter_forward() # Ensure ghost values are updated before multiplication + m._col_vec.array[:out_size_local] = 0.0 + m.mult(m._row_vec, m._col_vec, transpose=True) + v_out.array[:out_size_local] = m._col_vec.array[:out_size_local] else: # Prevent double-counting in parallel by zeroing ghosts of input vector # Calculate the exact number of local degrees of freedom - mat._row_vec.array[:out_size_local] = 0 - mat._col_vec.array[:in_size_local] = v_in.array[:in_size_local] - mat._col_vec.scatter_forward() # Ensure ghost values are updated before multiplication - mat.mult(mat._col_vec, mat._row_vec) - v_out.array[:out_size_local] = mat._row_vec.array[:out_size_local] + m._row_vec.array[:out_size_local] = 0 + m._col_vec.array[:in_size_local] = v_in.array[:in_size_local] + m._col_vec.scatter_forward() # Ensure ghost values are updated before multiplication + m.mult(m._col_vec, m._row_vec) + v_out.array[:out_size_local] = m._row_vec.array[:out_size_local] v_out.scatter_forward() return mult @@ -74,25 +75,20 @@ def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector): raise TypeError("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)) - - 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() - # The built in interpolation matrix requires two working arrays - attach_working_array(mat) - - _INTERPOLATION_MATRIX_CACHE[key] = 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 - return _INTERPOLATION_MATRIX_CACHE[key] + mat = dolfinx.fem.interpolation_matrix(space_from, space_to) + mat.scatter_reverse() + # The built in interpolation matrix requires two working arrays + attach_working_array(mat) + return mat class InterpolationBlock(Block): @@ -116,15 +112,27 @@ def __init__( 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 + + # Interpolation matrix is fixed for the lifetime of this block (tied to + # self.space_from/self.space_to), so cache it per-instance rather than + # in a module-level dict keyed by id() (which can collide once a + # FunctionSpace is garbage collected). + self._interpolation_matrix: dolfinx.la.MatrixCSR | "PETSc.Mat" | None = None def __str__(self): return "interpolate_function" + def _get_interpolation_matrix(self) -> dolfinx.la.MatrixCSR | "PETSc.Mat": + if self._interpolation_matrix is None: + self._interpolation_matrix = _build_interpolation_matrix( + self.space_from, self.space_to, use_petsc=self._use_petsc + ) + return self._interpolation_matrix + # --- 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] @@ -143,7 +151,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar # --- 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] @@ -165,7 +173,7 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar # --- 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 @@ -191,12 +199,17 @@ 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() - - # Overload the object to ensure PyAdjoint tracks it properly - output = 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 From 1de33910556ece18e6335973c8665e122918e0d4 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 21 Aug 2026 13:45:15 +0000 Subject: [PATCH 2/3] Restore cross-block interpolation matrix cache, made safe against id() reuse The previous commit removed the shared matrix cache entirely (making it per-instance on InterpolationBlock) to close the id()-collision hole, but that gives up the redundant-assembly/MPI-communicator-exhaustion protection the cache was originally for, since every new block (e.g. each optimization iteration that creates fresh spaces) reassembles its own matrix from scratch. Bring back a shared, module-level cache, but keyed safely this time: each FunctionSpace that contributes to a cache key gets a weakref.finalize callback 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 - no stale matrix can be served, and entries don't accumulate forever for short-lived spaces. Add regression tests: same-space-pair reuse, cache entries getting purged once their space is garbage collected, and no unbounded growth across many transient spaces. Verified these tests fail against a naive id()-keyed cache without the weakref purge (i.e. they would have caught the original bug). Co-Authored-By: Claude Sonnet 5 --- src/dolfinx_adjoint/blocks/interpolation.py | 52 ++++++++++++---- tests/test_interpolate.py | 67 ++++++++++++++++++++- 2 files changed, 107 insertions(+), 12 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py index e26a5f2..1c0cf2d 100644 --- a/src/dolfinx_adjoint/blocks/interpolation.py +++ b/src/dolfinx_adjoint/blocks/interpolation.py @@ -1,6 +1,7 @@ from __future__ import annotations import typing +import weakref from typing import Callable import dolfinx @@ -10,6 +11,33 @@ if typing.TYPE_CHECKING: from petsc4py import PETSc +# 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], "dolfinx.la.MatrixCSR | 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 attach_working_array(mat: dolfinx.la.MatrixCSR): """Attach working arrays to a dolfinx.la.MatrixCSR for efficient matrix-vector multiplication.""" @@ -91,6 +119,18 @@ def _build_interpolation_matrix( return mat +def _get_interpolation_matrix( + space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False +) -> dolfinx.la.MatrixCSR | "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] + + class InterpolationBlock(Block): """Block for interpolating a dolfinx.fem.Function into another space.""" @@ -113,21 +153,11 @@ def __init__( self._tlm_output: dolfinx.fem.Function | None = None self._hessian_output: dolfinx.fem.Function | None = None - # Interpolation matrix is fixed for the lifetime of this block (tied to - # self.space_from/self.space_to), so cache it per-instance rather than - # in a module-level dict keyed by id() (which can collide once a - # FunctionSpace is garbage collected). - self._interpolation_matrix: dolfinx.la.MatrixCSR | "PETSc.Mat" | None = None - def __str__(self): return "interpolate_function" def _get_interpolation_matrix(self) -> dolfinx.la.MatrixCSR | "PETSc.Mat": - if self._interpolation_matrix is None: - self._interpolation_matrix = _build_interpolation_matrix( - self.space_from, self.space_to, use_petsc=self._use_petsc - ) - return self._interpolation_matrix + return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc) # --- Adjoint --- diff --git a/tests/test_interpolate.py b/tests/test_interpolate.py index 1328412..1106f8c 100644 --- a/tests/test_interpolate.py +++ b/tests/test_interpolate.py @@ -1,3 +1,5 @@ +import gc + from mpi4py import MPI import dolfinx @@ -7,7 +9,12 @@ import ufl from dolfinx_adjoint import Function, assemble_scalar, interpolate -from dolfinx_adjoint.blocks.interpolation import InterpolationBlock +from dolfinx_adjoint.blocks.interpolation import ( + _CACHE_KEYS_BY_SPACE_ID, + _INTERPOLATION_MATRIX_CACHE, + InterpolationBlock, + _get_interpolation_matrix, +) # Dynamically determine available matrix backends petsc_options = [False] @@ -148,3 +155,61 @@ def u_ex(mod, x_coords): assert np.isclose(min_rate, 3.0, rtol=1e-3, atol=1e-3) pyadjoint.get_working_tape().clear_tape() + + +# ============================================================================== +# Test 3: 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 From f2120e8a2e054a4a5ba54bc94775fd66932bb15b Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 21 Aug 2026 13:51:42 +0000 Subject: [PATCH 3/3] Replace monkey-patched matrix attributes with a typed workspace class attach_working_array() dynamically attached _row_vec/_col_vec onto dolfinx.la.MatrixCSR instances, which mypy can't see, so every access needed typing.cast(Any, ...). Replace it with _MatrixCSRWorkspace, a small class built once alongside the matrix that holds the working vectors as real, statically-typed fields, threaded through get_mult(), _build_interpolation_matrix(), and the cache in place of the bare matrix. No behavior change (the full working arrays are still zeroed before each multiply); no more casts or dynamic attributes. Co-Authored-By: Claude Sonnet 5 --- src/dolfinx_adjoint/blocks/interpolation.py | 73 +++++++++++---------- 1 file changed, 37 insertions(+), 36 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py index 1c0cf2d..dcd0c23 100644 --- a/src/dolfinx_adjoint/blocks/interpolation.py +++ b/src/dolfinx_adjoint/blocks/interpolation.py @@ -23,7 +23,7 @@ # 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], "dolfinx.la.MatrixCSR | PETSc.Mat"] = {} +_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int, bool], "_MatrixCSRWorkspace | PETSc.Mat"] = {} _CACHE_KEYS_BY_SPACE_ID: dict[int, set[tuple[int, int, bool]]] = {} @@ -39,48 +39,51 @@ def _register_cache_key(space: dolfinx.fem.FunctionSpace, key: tuple[int, int, b _CACHE_KEYS_BY_SPACE_ID.setdefault(space_id, set()).add(key) -def attach_working_array(mat: dolfinx.la.MatrixCSR): - """Attach working arrays to a dolfinx.la.MatrixCSR for efficient matrix-vector multiplication.""" - # mat._row_vec/_col_vec are monkey-patched on; cast to Any so mypy doesn't - # flag the dynamic attributes. - m = typing.cast(typing.Any, mat) - if not hasattr(m, "_row_vec"): - m._row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype) - if not hasattr(m, "_col_vec"): - m._col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype) - m._row_vec.array[:] = 0.0 - m._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, - transpose: bool = False, # type: ignore + mat: "PETSc.Mat" | _MatrixCSRWorkspace, + transpose: bool = False, ) -> Callable[[dolfinx.la.Vector, dolfinx.la.Vector], None]: """Return a function that performs matrix-vector multiplication with the given matrix.""" - if isinstance(mat, dolfinx.la.MatrixCSR): + if isinstance(mat, _MatrixCSRWorkspace): + workspace = mat def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector): - # Need to use vectors from 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) # Ensure working arrays are attached - m = typing.cast(typing.Any, 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: - # Prevent double-counting in parallel by zeroing ghosts of input vector # Calculate the exact number of local degrees of freedom - m._row_vec.array[:in_size_local] = v_in.array[:in_size_local] - m._row_vec.scatter_forward() # Ensure ghost values are updated before multiplication - m._col_vec.array[:out_size_local] = 0.0 - m.mult(m._row_vec, m._col_vec, transpose=True) - v_out.array[:out_size_local] = m._col_vec.array[:out_size_local] + 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) + v_out.array[:out_size_local] = workspace.col_vec.array[:out_size_local] else: - # Prevent double-counting in parallel by zeroing ghosts of input vector # Calculate the exact number of local degrees of freedom - m._row_vec.array[:out_size_local] = 0 - m._col_vec.array[:in_size_local] = v_in.array[:in_size_local] - m._col_vec.scatter_forward() # Ensure ghost values are updated before multiplication - m.mult(m._col_vec, m._row_vec) - v_out.array[:out_size_local] = m._row_vec.array[:out_size_local] + workspace.row_vec.array[:out_size_local] = 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) + v_out.array[:out_size_local] = workspace.row_vec.array[:out_size_local] v_out.scatter_forward() return mult @@ -105,7 +108,7 @@ def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector): def _build_interpolation_matrix( space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False -) -> dolfinx.la.MatrixCSR | "PETSc.Mat": +) -> _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) @@ -114,14 +117,12 @@ def _build_interpolation_matrix( mat = dolfinx.fem.interpolation_matrix(space_from, space_to) mat.scatter_reverse() - # The built in interpolation matrix requires two working arrays - attach_working_array(mat) - return mat + return _MatrixCSRWorkspace(mat) def _get_interpolation_matrix( space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False -) -> dolfinx.la.MatrixCSR | "PETSc.Mat": +) -> _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: @@ -156,7 +157,7 @@ def __init__( def __str__(self): return "interpolate_function" - def _get_interpolation_matrix(self) -> dolfinx.la.MatrixCSR | "PETSc.Mat": + def _get_interpolation_matrix(self) -> _MatrixCSRWorkspace | "PETSc.Mat": return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc) # --- Adjoint ---