From a75161e91d691ea1185a0c26657fb315cddaf21f Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 14 Aug 2026 07:38:44 +0000 Subject: [PATCH 01/11] Implement pointobservation --- _toc.yml | 1 + demos/point_observations.py | 233 ++++++++ src/dolfinx_adjoint/__init__.py | 3 + src/dolfinx_adjoint/blocks/__init__.py | 2 + src/dolfinx_adjoint/blocks/observation.py | 114 ++++ src/dolfinx_adjoint/observation.py | 485 ++++++++++++++++ tests/test_observation.py | 674 ++++++++++++++++++++++ 7 files changed, 1512 insertions(+) create mode 100644 demos/point_observations.py create mode 100644 src/dolfinx_adjoint/blocks/observation.py create mode 100644 src/dolfinx_adjoint/observation.py create mode 100644 tests/test_observation.py diff --git a/_toc.yml b/_toc.yml index 6c632a1..f8d221e 100644 --- a/_toc.yml +++ b/_toc.yml @@ -6,6 +6,7 @@ parts: chapters: - file: "demos/poisson_mother.py" - file: "demos/time_distributed_control.py" + - file: "demos/point_observations.py" - caption: Python API chapters: - file: "docs/api" diff --git a/demos/point_observations.py b/demos/point_observations.py new file mode 100644 index 0000000..b1ca324 --- /dev/null +++ b/demos/point_observations.py @@ -0,0 +1,233 @@ +# # Inverting for a source term from point measurements +# *Section author: Henrik Finsberg ([henriknf@simula.no](mailto:henriknf@simula.no))*. + +# The [Poisson mother problem](./poisson_mother.py) demo matches the computed state against a +# desired profile known *everywhere* in the domain. Real data is rarely like that: sensors sit +# at a handful of locations, wells are drilled at specific coordinates, and medical images +# sample the domain on their own grid rather than on the finite element mesh. + +# This demo solves the same mother problem, but against data that only exists at **points**. +# The tool for that is `dolfinx_adjoint.PointObservation`, which evaluates a finite element +# function at a given set of coordinates, +# +# $$ +# d = B u, \qquad B_{ij} = \phi_j(x_i), +# $$ +# +# so that row $i$ picks out the value of $u$ at the point $x_i$. Together with a PDE solve, +# $m \mapsto B u(m)$ is what is known as the *parameter-to-observable map*. + +# ## Problem definition +# We recover an unknown source $f$ from noisy measurements of the state at a set of sensor +# locations. The state solves +# +# $$ +# \begin{align} +# -\Delta u &= f && \text{in } \Omega, \\ +# u &= 0 && \text{on } \partial\Omega, +# \end{align} +# $$ +# +# and we minimize +# +# $$ +# \min_{f} J(f) = \frac{1}{2} \lVert B u(f) - d \rVert^2 +# + \frac{\alpha}{2} \int_\Omega f^2 ~\mathrm{d}x. +# $$ +# +# The first term is a sum over the sensors rather than an integral over $\Omega$. The data is +# never interpolated onto the mesh, so nothing is invented at locations where nothing was +# measured -- which matters, because interpolating sparse data onto a fine mesh manufactures +# exactly the information the inverse problem is supposed to extract. + +# ## Implementation + +# + +from mpi4py import MPI + +import dolfinx +import numpy as np +import pyadjoint +import ufl + +import dolfinx_adjoint + +# - + +# We work on the unit square with both the state and the source in $P_1$. + +# + +domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 32, 32) +V = dolfinx.fem.functionspace(domain, ("Lagrange", 1)) + +domain.topology.create_connectivity(domain.topology.dim - 1, domain.topology.dim) +boundary_facets = dolfinx.mesh.exterior_facet_indices(domain.topology) +boundary_dofs = dolfinx.fem.locate_dofs_topological(V, domain.topology.dim - 1, boundary_facets) + + +def solve_forward(f: dolfinx_adjoint.Function, prefix: str) -> dolfinx_adjoint.Function: + """Solve the Poisson problem with source ``f`` and homogeneous Dirichlet conditions.""" + u = dolfinx_adjoint.Function(V, name="state") + trial, test = ufl.TrialFunction(V), ufl.TestFunction(V) + a = ufl.inner(ufl.grad(trial), ufl.grad(test)) * ufl.dx + L = ufl.inner(f, test) * ufl.dx + + u_bc = dolfinx_adjoint.Function(V, name="u_bc") + bc = dolfinx_adjoint.dirichletbc(u_bc, boundary_dofs) + dolfinx_adjoint.LinearProblem(a, L, u=u, bcs=[bc], petsc_options_prefix=prefix).solve() + return u + + +# - + +# ## Generating synthetic data +# +# The source we are trying to recover is a pair of Gaussian bumps. + + +# + +def true_source(x): + return 5.0 * np.exp(-((x[0] - 0.35) ** 2 + (x[1] - 0.65) ** 2) / 0.05) + 3.0 * np.exp( + -((x[0] - 0.7) ** 2 + (x[1] - 0.3) ** 2) / 0.03 + ) + + +f_true = dolfinx_adjoint.Function(V, name="f_true") +f_true.interpolate(true_source) +# - + +# The sensors sit on a slightly jittered grid. Creating the observation operator is a matter of +# handing it the function space and the coordinates. + +# + +rng = np.random.default_rng(2024) +axis = np.linspace(0.08, 0.92, 12) +sensors = np.stack(np.meshgrid(axis, axis, indexing="ij"), axis=-1).reshape(-1, 2) +sensors += rng.uniform(-0.01, 0.01, size=sensors.shape) + +B = dolfinx_adjoint.PointObservation(V, sensors) +print(f"Placed {B.num_found} of {B.num_points} sensors inside the mesh") +# - + +# `B.evaluate(u)` gives the value of `u` at each sensor. We solve the forward problem with the +# true source and add measurement noise to get our synthetic data. + +# + +NOISE_STD = 2e-3 + +with pyadjoint.stop_annotating(): + clean = B.evaluate(solve_forward(f_true, "demo_truth_")) + +observations = clean + rng.normal(0.0, NOISE_STD, size=clean.shape) +print(f"Observations: range [{observations.min():.4f}, {observations.max():.4f}], noise std {NOISE_STD}") +# - + +# ## The inverse problem +# +# Now we forget the true source, start from zero, and try to recover it from the measurements. +# `point_observation_misfit` records the misfit on the tape, so it composes with the PDE solve +# exactly like `assemble_scalar` does and the two terms can simply be added. + +# + +ALPHA = 1e-6 # Tikhonov regularization parameter + +pyadjoint.get_working_tape().clear_tape() + +f = dolfinx_adjoint.Function(V, name="f") +f.x.array[:] = 0.0 + +u = solve_forward(f, "demo_inverse_") + +misfit = dolfinx_adjoint.point_observation_misfit(u, B, observations) +regularization = dolfinx_adjoint.assemble_scalar(0.5 * ALPHA * ufl.inner(f, f) * ufl.dx) +J = misfit + regularization + +Jhat = pyadjoint.ReducedFunctional(J, pyadjoint.Control(f)) +# - + +# Before optimizing, we check that the gradient of the combined functional is correct, by +# perturbing the source in some direction and confirming that the error left over after the +# first-order Taylor expansion shrinks quadratically. + +# + +h = dolfinx_adjoint.Function(V) +h.interpolate(lambda x: 10.0 * np.sin(4 * np.pi * x[0]) * np.cos(3 * np.pi * x[1])) + +rate = pyadjoint.taylor_test(Jhat, f, h) +print(f"Taylor convergence rate: {rate:.3f} (expected ~2)") +assert rate > 1.9 +# - + +# Now minimize. Recovering a source from measurements of the (much smoother) state is a +# classically ill-conditioned problem, so L-BFGS-B needs a fair number of iterations here. + +# + +Jhat(f) # re-evaluate at the initial guess, undoing the Taylor test's perturbations +f_opt = pyadjoint.minimize(Jhat, method="L-BFGS-B", options={"maxiter": 500}) +# - + +# ## Results +# +# Two things are worth measuring: how well the recovered source reproduces the measurements, +# and how close it is to the source we started from. + +# + +with pyadjoint.stop_annotating(): + predicted = B.evaluate(solve_forward(f_opt, "demo_check_")) + l2_error = dolfinx_adjoint.error_norm(f_true, f_opt, norm_type="L2") + l2_truth = np.sqrt(dolfinx_adjoint.assemble_scalar(ufl.inner(f_true, f_true) * ufl.dx)) + +rms = np.sqrt(np.mean((predicted - observations) ** 2)) +print(f"Sensor residual RMS: {rms:.5f} (injected noise std was {NOISE_STD})") +print(f"Relative L2 error in f: {l2_error / l2_truth:.4f}") +# - + +# The residual settles *at* the noise floor rather than below it. That is the outcome to hope +# for: driving it to zero would mean fitting the noise, and the regularization term is what +# stops that from happening. +# +# The recovered source is a smoothed version of the truth, with lower peaks than the real +# bumps. That is not a defect of the method but a property of the data: a few hundred +# measurements of a smoothing operator's output do not determine the fine structure of the +# source, and the regularization supplies what the data cannot. + +# + tags=["hide-input"] +try: + import pyvista + +except ImportError: + print("Install pyvista to visualize the result") + +else: + # pyvista.set_jupyter_backend("html") + cells, types, geometry = dolfinx.plot.vtk_mesh(V) + sensor_cloud = pyvista.PolyData(np.column_stack([sensors, np.zeros(len(sensors))])) + + plotter = pyvista.Plotter(shape=(1, 2), window_size=(900, 400)) + for column, (field, title) in enumerate([(f_true, "True source"), (f_opt, "Recovered source")]): + grid = pyvista.UnstructuredGrid(cells, types, geometry) + grid.point_data["f"] = field.x.array.real + grid.set_active_scalars("f") + plotter.subplot(0, column) + plotter.add_text(title, font_size=10) + plotter.add_mesh(grid, show_edges=False, clim=[0.0, 5.0]) + plotter.add_mesh(sensor_cloud, color="black", point_size=5, render_points_as_spheres=True) + plotter.view_xy() + if pyvista.OFF_SCREEN: + plotter.screenshot("point_observations.png") + else: + plotter.show() + +# - + +# ## Going further +# +# * **Sensors outside the domain.** `B.found` flags any points that fall outside the mesh, and +# those are left out of the misfit rather than being treated as measurements of zero. +# * **Unreliable sensors.** `point_observation_misfit(..., weights=...)` reweights individual +# measurements; a zero weight drops one entirely. +# * **Noisier data.** `point_observation_misfit(..., noise_variance=...)` divides the misfit by +# $\sigma^2$, which is what makes it a negative log-likelihood when the regularization term is +# a genuine prior. +# * **Time-dependent data.** Build `B` once and call `point_observation_misfit` once per +# observation time, summing the results -- the operator only depends on the sensor positions. diff --git a/src/dolfinx_adjoint/__init__.py b/src/dolfinx_adjoint/__init__.py index 7a0a688..aa04754 100644 --- a/src/dolfinx_adjoint/__init__.py +++ b/src/dolfinx_adjoint/__init__.py @@ -6,6 +6,7 @@ import pyadjoint as _pyad from .assembly import assemble_scalar, error_norm +from .observation import PointObservation, point_observation_misfit from .solvers import LinearProblem, NonlinearProblem from .types import Constant, Function, dirichletbc from .types.function import assign @@ -24,12 +25,14 @@ __all__ = [ "Constant", "Function", + "PointObservation", "dirichletbc", "LinearProblem", "NonlinearProblem", "assemble_scalar", "assign", "error_norm", + "point_observation_misfit", "__version__", "__author__", "__license__", diff --git a/src/dolfinx_adjoint/blocks/__init__.py b/src/dolfinx_adjoint/blocks/__init__.py index 4c689e5..cff2a4a 100644 --- a/src/dolfinx_adjoint/blocks/__init__.py +++ b/src/dolfinx_adjoint/blocks/__init__.py @@ -1,7 +1,9 @@ from .assembly import AssembleBlock from .function_assigner import FunctionAssignBlock +from .observation import PointObservationBlock __all__ = [ "AssembleBlock", "FunctionAssignBlock", + "PointObservationBlock", ] diff --git a/src/dolfinx_adjoint/blocks/observation.py b/src/dolfinx_adjoint/blocks/observation.py new file mode 100644 index 0000000..83bb379 --- /dev/null +++ b/src/dolfinx_adjoint/blocks/observation.py @@ -0,0 +1,114 @@ +"""Tape block for the pointwise observation misfit.""" + +from __future__ import annotations + +import typing + +from mpi4py import MPI + +import dolfinx +import numpy as np +import numpy.typing as npt +from pyadjoint import Block +from pyadjoint.overloaded_type import create_overloaded_object + +from ._vector import _SpecialVector, _vector + +if typing.TYPE_CHECKING: # pragma: no cover + from ..observation import PointObservation + +__all__ = ["PointObservationBlock"] + + +class PointObservationBlock(Block): + r"""Block for :math:`J(u) = \frac{1}{2\sigma^2}\,\lVert W(Bu - d)\rVert^2`. + + The functional is quadratic in the state, so the derivatives are available in closed + form: the adjoint is :math:`\sigma^{-2} B^T W^2 (Bu - d)` and the Hessian action is + :math:`\sigma^{-2} B^T W^2 B \hat{u}`. No linearization point needs to be stored. + + Args: + u: The observed state. + observation: The observation operator :math:`B`. + data: Measured values, already restricted to this rank's rows. + noise_variance: :math:`\sigma^2`. + weights: Optional per-row weights :math:`W`, restricted to this rank's rows. + ad_block_tag: Optional tag for the block on the tape. + """ + + def __init__( + self, + u: dolfinx.fem.Function, + observation: "PointObservation", + data: npt.NDArray[np.float64], + noise_variance: float, + weights: npt.NDArray[np.float64] | None = None, + ad_block_tag: str | None = None, + ) -> None: + super().__init__(ad_block_tag=ad_block_tag) + self.add_dependency(u) + self.observation = observation + self.data = data + self.noise_variance = noise_variance + self.weights = weights + + def __str__(self) -> str: + return f"point_observation_misfit({self.observation.num_found} points)" + + def _residual(self, u: dolfinx.fem.Function) -> npt.NDArray[np.float64]: + return self.observation.apply(u) - self.data + + def _apply_weights(self, values: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + if self.weights is None: + return values + return self.weights * values + + def _transpose_action(self, residual: npt.NDArray[np.float64], scale: float) -> _SpecialVector: + """:math:`\\mathrm{scale} \\cdot \\sigma^{-2} B^T W^2 r`, as a DOF vector.""" + V = self.observation.function_space + out = _vector(V.dofmap.index_map, V.dofmap.bs, V, dtype=V.mesh.geometry.x.dtype) + out.array[:] = 0.0 + # W is applied twice: once to the residual, once from differentiating ||W r||^2. + weighted = self._apply_weights(self._apply_weights(residual)) * (scale / self.noise_variance) + self.observation.apply_transpose(weighted, out=out) + return out + + def recompute_component(self, inputs, block_variable, idx, prepared=None): + from ..observation import misfit_value + + value = misfit_value(self.observation, inputs[0], self.data, self.noise_variance, self.weights) + return create_overloaded_object(value) + + def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): + adj_input = 1.0 if adj_inputs[0] is None else float(adj_inputs[0]) + return self._transpose_action(self._residual(inputs[0]), adj_input) + + def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None): + tlm_u = tlm_inputs[0] + if tlm_u is None: + return None + residual = self._apply_weights(self._residual(inputs[0])) + directional = self._apply_weights(self.observation.apply(tlm_u)) + local = float(np.dot(residual, directional)) / self.noise_variance + return self.observation.comm.allreduce(local, op=MPI.SUM) + + def evaluate_hessian_component( + self, + inputs, + hessian_inputs, + adj_inputs, + block_variable, + idx, + relevant_dependencies, + prepared=None, + ): + hessian_input = 0.0 if hessian_inputs[0] is None else float(hessian_inputs[0]) + adj_input = 1.0 if adj_inputs[0] is None else float(adj_inputs[0]) + + # Second-order seed, propagated through the first derivative ... + out = self._transpose_action(self._residual(inputs[0]), hessian_input) + # ... plus the curvature of J applied to the TLM direction. + tlm_u = block_variable.tlm_value + if tlm_u is not None: + out.array[:] += self._transpose_action(self.observation.apply(tlm_u), adj_input).array[:] + return out diff --git a/src/dolfinx_adjoint/observation.py b/src/dolfinx_adjoint/observation.py new file mode 100644 index 0000000..6aef371 --- /dev/null +++ b/src/dolfinx_adjoint/observation.py @@ -0,0 +1,485 @@ +"""Pointwise observation of a finite element function. + +Inverse problems are frequently posed against data that lives at *points* rather than on +the computational mesh: sensor readings, well measurements, or the voxel centres of an +image. This module provides the discrete observation operator + +.. math:: + + d = B u, \\qquad B \\in \\mathbb{R}^{n_d \\times N}, + \\qquad B_{ij} = \\phi_j(x_i), + +whose :math:`i`-th row evaluates a finite element function at the point :math:`x_i`, and a +differentiable least-squares misfit built on top of it, + +.. math:: + + J(u) = \\frac{1}{2\\sigma^2} \\lVert W (B u - d) \\rVert^2 . + +Composed with a PDE solve, this is the *parameter-to-observable map* of PDE-constrained +optimization and Bayesian inversion: :math:`m \\mapsto B u(m)`. + +Parallel semantics +------------------ +``points`` must be *replicated* on every MPI rank -- observation points typically come from +a file or an instrument layout that every rank can read. Each point is assigned to exactly +one owning rank: the lowest-numbered rank whose *owned* cells contain it. Points that no +rank can locate are reported in :attr:`PointObservation.found` and excluded from the +operator, rather than becoming zero rows -- a zero row in a misfit silently contributes +:math:`d_i^2` and biases the result. +""" + +from __future__ import annotations + +from mpi4py import MPI + +import dolfinx +import numpy as np +import numpy.typing as npt +import pyadjoint +from pyadjoint.overloaded_type import create_overloaded_object +from pyadjoint.tape import annotate_tape, get_working_tape, stop_annotating + +from .blocks.observation import PointObservationBlock + +__all__ = ["PointObservation", "point_observation_misfit"] + + +def _geometry_dofmap(mesh: dolfinx.mesh.Mesh): + """``mesh.geometry.dofmap`` was renamed to ``dofmaps[0]``.""" + try: + return mesh.geometry.dofmaps[0] + except AttributeError: + return mesh.geometry.dofmap + + +def _coordinate_map(mesh: dolfinx.mesh.Mesh): + """``mesh.geometry.cmap`` is deprecated in favour of ``cmaps[0]``.""" + try: + return mesh.geometry.cmaps[0] + except (AttributeError, IndexError): + return mesh.geometry.cmap + + +#: Cell types whose degree-1 coordinate map is affine. Tensor-product cells +#: (quadrilateral, hexahedron, prism, pyramid) are multilinear rather than affine even at +#: degree 1, so they take the general Newton pull-back. +AFFINE_CELL_TYPES = frozenset( + { + dolfinx.mesh.CellType.point, + dolfinx.mesh.CellType.interval, + dolfinx.mesh.CellType.triangle, + dolfinx.mesh.CellType.tetrahedron, + } +) + + +def _is_affine_simplex(mesh: dolfinx.mesh.Mesh) -> bool: + """True if every cell map is affine, i.e. a degree-1 simplex geometry.""" + return _coordinate_map(mesh).degree == 1 and mesh.topology.cell_type in AFFINE_CELL_TYPES + + +def _default_padding(mesh: dolfinx.mesh.Mesh) -> float: + """Padding on the scale of rounding error relative to the mesh extent. + + Reduced over the *global* bounding box rather than each process's own, so that the + operator locates the same points however the mesh happens to be partitioned. + """ + x = mesh.geometry.x + empty = x.shape[0] == 0 + local_lower = np.full(3, np.inf) if empty else x.min(axis=0) + local_upper = np.full(3, -np.inf) if empty else x.max(axis=0) + lower = np.empty(3) + upper = np.empty(3) + mesh.comm.Allreduce(np.ascontiguousarray(local_lower, dtype=np.float64), lower, op=MPI.MIN) + mesh.comm.Allreduce(np.ascontiguousarray(local_upper, dtype=np.float64), upper, op=MPI.MAX) + diagonal = float(np.linalg.norm(upper - lower)) if np.all(np.isfinite(lower)) else 0.0 + return 1e-10 * max(diagonal, 1.0) + + +def _pad_points(points: npt.ArrayLike) -> np.ndarray: + """Return points as a contiguous ``(num_points, 3)`` float64 array.""" + padded_input = np.atleast_2d(np.asarray(points, dtype=np.float64)) + if padded_input.ndim != 2: + raise ValueError(f"Points must be 2D with shape (num_points, dim), got shape {padded_input.shape}") + if padded_input.shape[1] > 3: + raise ValueError(f"Points can have at most 3 components, got {padded_input.shape[1]}") + padded = np.zeros((padded_input.shape[0], 3), dtype=np.float64) + padded[:, : padded_input.shape[1]] = padded_input + return padded + + +class PointObservation: + """The operator :math:`B` evaluating a finite element function at a set of points. + + Args: + V: Function space the observed state lives in. + points: Observation points, ``shape=(num_points, gdim)`` or ``(num_points, 3)``. + Must be identical on every MPI rank. + padding: Absolute padding of the mesh bounding boxes used when searching for the + cell containing each point. Defaults to a small multiple of the mesh extent, + which makes points sitting exactly on the boundary robustly detectable. Large + values are expensive -- keep this at cell scale at most. + + Attributes: + num_points: Total (global) number of input points. + found: Boolean array of length ``num_points``, ``True`` where the point was located + in the mesh on some rank. Identical on every rank. + owner: Rank owning each point, ``-1`` where it was not found. Identical on every + rank. + local_indices: Indices into the global point array of the points owned by this + rank. This is the row ordering used by :meth:`apply`. + + Note: + For a vector space with block size ``bs``, row ``i * bs + c`` observes component + ``c`` at point ``i``. + + Example: + .. code-block:: python + + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + B = PointObservation(V, np.array([[0.25, 0.25], [0.75, 0.5]])) + values = B.evaluate(u) # u(0.25, 0.25) and u(0.75, 0.5) + """ + + def __init__( + self, + V: dolfinx.fem.FunctionSpace, + points: npt.ArrayLike, + padding: float | None = None, + ) -> None: + mesh = V.mesh + comm = mesh.comm + self.function_space = V + self.comm = comm + + # Validate the element up front, on every process. Doing it lazily during assembly + # would raise only on the processes that end up owning points, and a rank-divergent + # exception deadlocks at the next collective call instead of surfacing cleanly. + if V.element.needs_dof_transformations: + raise NotImplementedError( + "PointObservation does not support elements requiring DOF transformations " + "(for instance high-order elements on oriented, non-simplex cells)." + ) + try: + self._basix_element = V.element.basix_element + except RuntimeError as exc: + raise NotImplementedError( + "PointObservation needs a function space backed by a Basix element, which mixed " + "elements are not. Collapse the sub-space you want to observe, for example " + "`V.sub(0).collapse()[0]`, and build the operator on that." + ) from exc + + padded_points = _pad_points(points) + num_points = padded_points.shape[0] + if comm.allreduce(num_points, op=MPI.MIN) != comm.allreduce(num_points, op=MPI.MAX): + raise ValueError("`points` must be replicated on all processes (got differing lengths).") + # A checksum costs one scalar reduction and catches the far nastier case of equally + # many but *different* points, which would otherwise silently corrupt the operator. + checksum = float(padded_points.sum()) + if comm.allreduce(checksum, op=MPI.MIN) != comm.allreduce(checksum, op=MPI.MAX): + raise ValueError("`points` must be replicated on all processes (got differing coordinates).") + self.num_points = num_points + self.points = padded_points + + tdim = mesh.topology.dim + gdim = mesh.geometry.dim + self.padding = _default_padding(mesh) if padding is None else float(padding) + + # Search owned cells only: a point interior to a cell is then found on exactly one + # rank, and only points on a shared facet or vertex remain ambiguous. + num_owned_cells = mesh.topology.index_map(tdim).size_local + first_cell = np.full(num_points, -1, dtype=np.int32) + if num_owned_cells > 0 and num_points > 0: + owned_cells = np.arange(num_owned_cells, dtype=np.int32) + tree = dolfinx.geometry.bb_tree(mesh, tdim, padding=self.padding, entities=owned_cells) + candidates = dolfinx.geometry.compute_collisions_points(tree, padded_points) + colliding = dolfinx.geometry.compute_colliding_cells(mesh, candidates, padded_points) + offsets = colliding.offsets + has_cell = offsets[1:] > offsets[:-1] + first_cell[has_cell] = colliding.array[offsets[:-1][has_cell]] + + # Break ties by lowest rank, so every rank agrees on a single owner per point. + candidate_owner = np.where(first_cell >= 0, comm.rank, comm.size).astype(np.int32) + owner = np.empty_like(candidate_owner) + comm.Allreduce(candidate_owner, owner, op=MPI.MIN) + + self.found = owner < comm.size + self.owner = np.where(self.found, owner, -1).astype(np.int32) + self.num_found = int(self.found.sum()) + self.local_indices = np.flatnonzero(owner == comm.rank).astype(np.int32) + + self._assemble(padded_points[self.local_indices], first_cell[self.local_indices], gdim) + + # -------------------------------------------------------------------- assembly --- + def _assemble(self, points: np.ndarray, cells: np.ndarray, gdim: int) -> None: + """Tabulate the basis functions of the containing cell for every owned point. + + The operator is stored as two dense ``(num_local_rows, num_dofs_per_cell)`` arrays + -- the basis values and the corresponding *ghosted* local DOF indices -- rather + than a sparse matrix. Every row has exactly the same number of entries, so this is + both smaller and faster than a general sparse format, and keeps the package free + of a sparse-matrix dependency. + """ + V = self.function_space + dofmap = V.dofmap + bs = dofmap.bs + self._num_local_dofs = (dofmap.index_map.size_local + dofmap.index_map.num_ghosts) * bs + + if len(points) == 0: + num_dofs_per_cell = dofmap.dof_layout.num_dofs + self._basis = np.zeros((0, num_dofs_per_cell)) + self._cols = np.zeros((0, num_dofs_per_cell), dtype=np.int32) + return + + reference_points = self._pull_back(points, cells, gdim) + + # The reference basis is shared by all cells, so a single tabulate call suffices. + basis = np.asarray(self._basix_element.tabulate(0, reference_points))[0, :, :, 0] + cell_dofs = np.asarray(dofmap.list)[cells] + + if bs == 1: + self._basis = basis + self._cols = cell_dofs.astype(np.int32) + else: + # Unroll to one row per (point, component): row i*bs + c reads the DOFs of + # component c, weighted by the same scalar basis values. + num_rows, num_dofs_per_cell = basis.shape + self._basis = np.repeat(basis, bs, axis=0) + components = np.arange(bs, dtype=np.int32) + cols = cell_dofs[:, None, :] * bs + components[None, :, None] + self._cols = cols.reshape(num_rows * bs, num_dofs_per_cell).astype(np.int32) + + def _pull_back(self, points: np.ndarray, cells: np.ndarray, gdim: int) -> np.ndarray: + """Map physical points to the reference coordinates of their containing cell.""" + mesh = self.function_space.mesh + geometry_dofmap = _geometry_dofmap(mesh) + geometry_x = mesh.geometry.x + tdim = mesh.topology.dim + dtype = geometry_x.dtype + + if _is_affine_simplex(mesh) and gdim == tdim: + # X = x0 + J xi, so xi = J^{-1} (X - x0) with a constant Jacobian per cell. + coords = geometry_x[np.asarray(geometry_dofmap)[cells]][:, : tdim + 1, :gdim] + jacobian = np.swapaxes(coords[:, 1:, :] - coords[:, :1, :], 1, 2) + offsets = (points[:, :gdim] - coords[:, 0, :])[..., None] + return np.linalg.solve(jacobian, offsets)[..., 0].astype(dtype) + + # Otherwise fall back to the coordinate element's Newton iteration, which handles + # one cell at a time; group the points by cell to call it as rarely as possible. + cmap = _coordinate_map(mesh) + reference_points = np.zeros((len(points), tdim), dtype=dtype) + order = np.argsort(cells, kind="stable") + boundaries = np.flatnonzero(np.diff(cells[order])) + 1 + for group in np.split(order, boundaries): + cell_coords = geometry_x[geometry_dofmap[cells[group[0]]]][:, :gdim] + physical = np.ascontiguousarray(points[group, :gdim], dtype=dtype) + reference_points[group] = cmap.pull_back(physical, cell_coords) + return reference_points + + # --------------------------------------------------------------------- actions --- + @property + def block_size(self) -> int: + """Block size of the observed function space.""" + return self.function_space.dofmap.bs + + @property + def num_local_rows(self) -> int: + """Number of rows owned by this rank, including the block-size unrolling.""" + return self._basis.shape[0] + + def apply(self, u: dolfinx.fem.Function) -> npt.NDArray[np.float64]: + """Evaluate :math:`Bu` for the rows owned by this rank. + + Args: + u: Function in :attr:`function_space`. + + Returns: + The point values of the rows owned by this rank, ordered as + :attr:`local_indices` (component-fastest for vector spaces). + """ + u.x.scatter_forward() + if self.num_local_rows == 0: + return np.zeros(0) + return np.einsum("ij,ij->i", self._basis, u.x.array[self._cols]) + + def evaluate(self, u: dolfinx.fem.Function, fill: float = np.nan) -> npt.NDArray[np.float64]: + """Evaluate ``u`` at every observation point. + + The convenience combination of :meth:`apply` and :meth:`gather`: the result has one + entry per point (``block_size`` entries per point for a vector space) and is the + same on every process, so it can be used directly as a data vector. + + Args: + u: Function in :attr:`function_space`. + fill: Value used for points that lie outside the mesh. + + Note: + This does not touch the tape. Use :func:`point_observation_misfit` to build a + differentiable functional out of the observations. + """ + return self.gather(self.apply(u), fill=fill) + + def apply_transpose(self, values: npt.ArrayLike, out: dolfinx.la.Vector | None = None) -> dolfinx.la.Vector: + """Accumulate :math:`B^T v` into a DOF vector. + + Args: + values: Row values in the layout produced by :meth:`apply`. + out: Optional vector to accumulate into; created when not supplied. + + Returns: + The vector holding :math:`B^T v`, with ghost contributions reduced onto their + owners and scattered back out. + """ + V = self.function_space + if out is None: + out = dolfinx.la.vector(V.dofmap.index_map, V.dofmap.bs, dtype=V.mesh.geometry.x.dtype) + out.array[:] = 0.0 + if self.num_local_rows > 0: + weights = self._basis * np.asarray(values, dtype=np.float64)[:, None] + contributions = np.bincount( + self._cols.reshape(-1), weights=weights.reshape(-1), minlength=self._num_local_dofs + ) + out.array[:] += contributions + out.scatter_reverse(dolfinx.la.InsertMode.add) + out.scatter_forward() + return out + + def restrict(self, data: npt.ArrayLike) -> npt.NDArray[np.float64]: + """Select the entries of a global, replicated array belonging to this rank's rows. + + Args: + data: Array of length ``num_points * block_size``, identical on every rank. + """ + values = np.asarray(data, dtype=np.float64) + bs = self.block_size + if values.shape[0] != self.num_points * bs: + raise ValueError(f"Expected data of length {self.num_points * bs}, got {values.shape[0]}") + if bs == 1: + return values[self.local_indices] + return values.reshape(self.num_points, bs)[self.local_indices].reshape(-1) + + def gather(self, values: npt.ArrayLike, fill: float = np.nan) -> npt.NDArray[np.float64]: + """Assemble local row values into a global array replicated on every rank. + + Args: + values: Row values in the layout produced by :meth:`apply`. + fill: Value used for points that were not found in the mesh. + """ + bs = self.block_size + buffer = np.zeros(self.num_points * bs, dtype=np.float64) + local = np.asarray(values, dtype=np.float64) + if bs == 1: + buffer[self.local_indices] = local + else: + buffer.reshape(self.num_points, bs)[self.local_indices] = local.reshape(-1, bs) + total = np.empty_like(buffer) + self.comm.Allreduce(buffer, total, op=MPI.SUM) + if not np.all(self.found): + missing = ~self.found + if bs == 1: + total[missing] = fill + else: + total.reshape(self.num_points, bs)[missing] = fill + return total + + def to_scipy(self): + """This rank's block of :math:`B` as a ``scipy.sparse`` CSR matrix. + + Columns index the ghosted local DOF array (``u.x.array``). Requires ``scipy``, + which is not a dependency of this package; :meth:`apply` and + :meth:`apply_transpose` do not need it. + """ + import scipy.sparse + + num_rows, num_dofs_per_cell = self._basis.shape + rows = np.repeat(np.arange(num_rows), num_dofs_per_cell) + return scipy.sparse.csr_matrix( + (self._basis.reshape(-1), (rows, self._cols.reshape(-1))), + shape=(num_rows, self._num_local_dofs), + ) + + +def point_observation_misfit( + u: dolfinx.fem.Function, + observation: PointObservation, + data: npt.ArrayLike, + noise_variance: float = 1.0, + weights: npt.ArrayLike | None = None, + ad_block_tag: str | None = None, + **kwargs, +) -> pyadjoint.AdjFloat: + """Least-squares misfit between a state observed at points and measured data. + + Computes :math:`\\frac{1}{2\\sigma^2}\\lVert W(Bu - d)\\rVert^2` and records it on the + tape, so that it can be differentiated with respect to anything ``u`` depends on. + + Args: + u: The state to observe. + observation: The observation operator :math:`B`. + data: The measured values :math:`d`. Either a global array of length + ``observation.num_points * observation.block_size``, replicated on every rank, + or an array already restricted to this rank's rows. + noise_variance: :math:`\\sigma^2`. Use ``1.0`` for a purely deterministic inverse + problem; for a Bayesian one this is the variance of the additive Gaussian + observation noise. + weights: Optional per-row weights :math:`W`, in the same layout as ``data``. A 0/1 + array masks individual observations out. + ad_block_tag: Optional tag for the tape block. + kwargs: ``annotate`` may be passed to control whether the tape records this. + + Returns: + The misfit, as a ``pyadjoint.AdjFloat``. + + Raises: + ZeroDivisionError: If ``noise_variance`` is zero. + """ + if noise_variance == 0: + raise ZeroDivisionError("noise_variance must not be 0.0; use 1.0 for a deterministic inverse problem.") + + local_data = _as_local(observation, data) + local_weights = None if weights is None else _as_local(observation, weights) + + annotate = annotate_tape(kwargs) + with stop_annotating(): + output = misfit_value(observation, u, local_data, noise_variance, local_weights) + + overloaded = create_overloaded_object(output) + + if annotate: + block = PointObservationBlock( + u, observation, local_data, noise_variance, local_weights, ad_block_tag=ad_block_tag + ) + get_working_tape().add_block(block) + block.add_output(overloaded.block_variable) + + return overloaded + + +def misfit_value( + observation: PointObservation, + u: dolfinx.fem.Function, + data: npt.NDArray[np.float64], + noise_variance: float, + weights: npt.NDArray[np.float64] | None, +) -> float: + """The misfit value, summed over ranks. Used by both the forward pass and the block.""" + residual = observation.apply(u) - data + if weights is not None: + residual = weights * residual + local = 0.5 * float(np.dot(residual, residual)) / noise_variance + return observation.comm.allreduce(local, op=MPI.SUM) + + +def _as_local(observation: PointObservation, values: npt.ArrayLike) -> npt.NDArray[np.float64]: + """Accept either a global (replicated) array or one already restricted to local rows.""" + array = np.asarray(values, dtype=np.float64).reshape(-1) + expected_global = observation.num_points * observation.block_size + if array.shape[0] == expected_global: + return observation.restrict(array) + if array.shape[0] == observation.num_local_rows: + return array + raise ValueError( + f"Expected an array of length {expected_global} (global) " + f"or {observation.num_local_rows} (local), got {array.shape[0]}" + ) diff --git a/tests/test_observation.py b/tests/test_observation.py new file mode 100644 index 0000000..753fcff --- /dev/null +++ b/tests/test_observation.py @@ -0,0 +1,674 @@ +"""Tests for the pointwise observation operator and its misfit.""" + +from mpi4py import MPI + +import basix.ufl +import dolfinx +import numpy as np +import pyadjoint +import pytest +import ufl + +import dolfinx_adjoint +from dolfinx_adjoint import observation + + +def unit_square(comm: MPI.Intracomm, n: int = 8) -> dolfinx.mesh.Mesh: + return dolfinx.mesh.create_unit_square(comm, n, n) + + +def sample_points(n: int = 37, seed: int = 0) -> np.ndarray: + """Points strictly inside the unit square, identical on every rank.""" + rng = np.random.default_rng(seed) + return 0.05 + 0.9 * rng.random((n, 2)) + + +def owned_dofs(V: dolfinx.fem.FunctionSpace) -> int: + return V.dofmap.index_map.size_local * V.dofmap.bs + + +# --------------------------------------------------------------------------- +# The observation operator +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("degree", [1, 2]) +def test_exact_on_polynomials(degree): + """B u reproduces the exact values of a function that lies in the FE space.""" + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", degree)) + + def expression(x): + if degree == 1: + return 1.0 + 2.0 * x[0] - 3.0 * x[1] + return 1.0 + 2.0 * x[0] - 3.0 * x[1] + 0.5 * x[0] * x[1] - x[0] ** 2 + 2 * x[1] ** 2 + + u = dolfinx.fem.Function(V) + u.interpolate(expression) + + points = sample_points() + B = dolfinx_adjoint.PointObservation(V, points) + + assert B.found.all() + assert np.allclose(B.gather(B.apply(u)), expression(points.T), atol=1e-12) + + +def test_evaluate_combines_apply_and_gather(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: 2.0 * x[0] - x[1]) + + points = sample_points(17, seed=101) + B = dolfinx_adjoint.PointObservation(V, points) + + assert np.allclose(B.evaluate(u), B.gather(B.apply(u)), equal_nan=True) + assert np.allclose(B.evaluate(u), 2.0 * points[:, 0] - points[:, 1], atol=1e-12) + + +def test_evaluate_marks_points_outside_the_mesh(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: x[0]) + + B = dolfinx_adjoint.PointObservation(V, np.array([[0.5, 0.5], [7.0, 7.0]])) + values = B.evaluate(u) + assert np.isclose(values[0], 0.5) + assert np.isnan(values[1]) + # The fill value is configurable, for callers that prefer a sentinel to NaN. + assert B.evaluate(u, fill=0.0)[1] == 0.0 + + +def test_partition_of_unity(): + """Each row of B sums to one, so a constant is observed as that constant.""" + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 2)) + u = dolfinx.fem.Function(V) + u.x.array[:] = 1.0 + + B = dolfinx_adjoint.PointObservation(V, sample_points()) + assert np.allclose(B.gather(B.apply(u)), 1.0) + + +def test_matches_dolfinx_point_evaluation(): + """B u agrees with dolfinx' own evaluation of a non-polynomial function.""" + comm = MPI.COMM_WORLD + mesh = unit_square(comm, 12) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 2)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: np.sin(3 * x[0]) * np.cos(2 * x[1])) + + points = sample_points(23, seed=3) + values = dolfinx_adjoint.PointObservation(V, points).gather(dolfinx_adjoint.PointObservation(V, points).apply(u)) + + points_3d = np.zeros((points.shape[0], 3)) + points_3d[:, :2] = points + tree = dolfinx.geometry.bb_tree(mesh, mesh.topology.dim) + candidates = dolfinx.geometry.compute_collisions_points(tree, points_3d) + colliding = dolfinx.geometry.compute_colliding_cells(mesh, candidates, points_3d) + for i in range(points.shape[0]): + cells = colliding.links(i) + if len(cells) == 0: + continue # this point is not on this rank + assert np.isclose(u.eval(points_3d[i], cells[0])[0], values[i], atol=1e-12) + + +def test_vector_space_rows_are_component_fastest(): + """Row i*bs + c observes component c at point i.""" + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 1, (2,))) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: np.vstack([x[0] + 2 * x[1], 3.0 - x[0]])) + + points = sample_points(11, seed=5) + B = dolfinx_adjoint.PointObservation(V, points) + assert B.block_size == 2 + + values = B.gather(B.apply(u)).reshape(-1, 2) + assert np.allclose(values[:, 0], points[:, 0] + 2 * points[:, 1], atol=1e-12) + assert np.allclose(values[:, 1], 3.0 - points[:, 0], atol=1e-12) + + +def test_points_outside_the_mesh_are_reported_and_dropped(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 1)) + points = np.array([[0.5, 0.5], [5.0, 5.0], [0.25, 0.75], [-1.0, 0.5]]) + + B = dolfinx_adjoint.PointObservation(V, points) + assert B.found.tolist() == [True, False, True, False] + assert B.num_found == 2 + assert B.owner[1] == -1 and B.owner[3] == -1 + + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: x[0] + x[1]) + values = B.gather(B.apply(u)) + assert np.isclose(values[0], 1.0) and np.isclose(values[2], 1.0) + assert np.isnan(values[1]) and np.isnan(values[3]) + + +def test_points_on_the_boundary_are_found(): + """The default padding makes points sitting exactly on the boundary detectable.""" + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 1)) + points = np.array([[0.0, 0.0], [1.0, 1.0], [0.0, 0.5], [0.5, 1.0]]) + B = dolfinx_adjoint.PointObservation(V, points) + assert B.found.all() + + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: x[0] + x[1]) + assert np.allclose(B.gather(B.apply(u)), [0.0, 2.0, 0.5, 1.5], atol=1e-12) + + +def test_ownership_is_unique_and_agreed_on_by_every_rank(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 10), ("Lagrange", 1)) + B = dolfinx_adjoint.PointObservation(V, sample_points(101, seed=7)) + + claimed = np.zeros(B.num_points, dtype=np.int64) + claimed[B.local_indices] = 1 + total = np.empty_like(claimed) + comm.Allreduce(claimed, total, op=MPI.SUM) + assert np.all(total[B.found] == 1) + assert np.all(total[~B.found] == 0) + + for other in comm.allgather(B.owner): + assert np.array_equal(other, B.owner) + + +def test_transpose_is_the_adjoint_of_apply(): + """ == globally, including ghost contributions.""" + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 9), ("Lagrange", 2)) + rng = np.random.default_rng(11) + + u = dolfinx.fem.Function(V) + u.x.array[:] = rng.random(u.x.array.shape) + u.x.scatter_forward() + + B = dolfinx_adjoint.PointObservation(V, sample_points(47, seed=13)) + v = B.restrict(rng.random(B.num_points)) + + lhs = comm.allreduce(float(np.dot(B.apply(u), v)), op=MPI.SUM) + Btv = B.apply_transpose(v) + n = owned_dofs(V) + rhs = comm.allreduce(float(np.dot(Btv.array[:n], u.x.array[:n])), op=MPI.SUM) + + assert np.isclose(lhs, rhs) + + +def test_replicated_points_are_required(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + if comm.size == 1: + pytest.skip("differing point counts require more than one rank") + points = sample_points(5) if comm.rank == 0 else sample_points(6) + with pytest.raises(ValueError, match="replicated"): + dolfinx_adjoint.PointObservation(V, points) + + +def test_mixed_element_is_rejected_on_every_process(): + """A mixed space must fail identically everywhere, not only where points landed. + + The element is validated up front rather than during assembly: a process that owns no + points skips assembly entirely, so a lazy check would raise on some processes and not + others and deadlock at the next collective call instead of surfacing the error. + """ + comm = MPI.COMM_WORLD + mesh = unit_square(comm, 6) + element = basix.ufl.mixed_element( + [ + basix.ufl.element("Lagrange", mesh.basix_cell(), 1), + basix.ufl.element("Lagrange", mesh.basix_cell(), 2), + ] + ) + W = dolfinx.fem.functionspace(mesh, element) + + # A single point, so at most one process would ever reach assembly. + with pytest.raises(NotImplementedError, match="Basix element"): + dolfinx_adjoint.PointObservation(W, np.array([[0.3, 0.4]])) + + +def test_points_must_match_across_processes(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + if comm.size == 1: + pytest.skip("mismatched points require more than one process") + + # Same number of points, different coordinates: caught by the checksum. + points = sample_points(5, seed=109) + if comm.rank == 1: + points = points + 0.01 + with pytest.raises(ValueError, match="differing coordinates"): + dolfinx_adjoint.PointObservation(V, points) + + +def test_padding_does_not_depend_on_the_partition(): + """The default padding comes from the global bounding box, not each process's own.""" + comm = MPI.COMM_WORLD + mesh = unit_square(comm, 8) + padding = observation._default_padding(mesh) + for other in comm.allgather(padding): + assert other == padding + # Unit square: the diagonal is sqrt(2) regardless of how the mesh is split up. + assert np.isclose(padding, 1e-10 * np.sqrt(2.0)) + + +def test_operator_survives_a_process_owning_no_points(): + """Clustered points leave some processes with nothing to do.""" + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 8), ("Lagrange", 1)) + points = np.array([[0.01, 0.01], [0.02, 0.02], [0.03, 0.01]]) + + B = dolfinx_adjoint.PointObservation(V, points) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: x[0] + x[1]) + + assert np.allclose(B.evaluate(u), points.sum(axis=1), atol=1e-12) + # An empty local block must still transpose, and contribute nothing. + result = B.apply_transpose(np.zeros(B.num_local_rows)) + assert np.allclose(result.array, 0.0) + + +def test_operator_with_no_points_at_all(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + B = dolfinx_adjoint.PointObservation(V, np.zeros((0, 2))) + + u = dolfinx.fem.Function(V) + u.x.array[:] = 1.0 + assert B.num_points == 0 + assert B.num_local_rows == 0 + assert B.evaluate(u).shape == (0,) + assert np.allclose(B.apply_transpose(np.zeros(0)).array, 0.0) + + +def test_too_many_components_raises(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + with pytest.raises(ValueError, match="at most 3 components"): + dolfinx_adjoint.PointObservation(V, np.zeros((3, 4))) + + +@pytest.mark.parametrize( + "cell_type", + [ + dolfinx.mesh.CellType.tetrahedron, + dolfinx.mesh.CellType.hexahedron, + dolfinx.mesh.CellType.prism, + ], +) +@pytest.mark.parametrize("degree", [1, 2]) +def test_three_dimensional_cell_types(cell_type, degree): + """Simplex and tensor-product cells alike, in 3D. + + Only the simplex has an affine coordinate map; the others go through the general + Newton pull-back, which must be just as exact. + """ + comm = MPI.COMM_WORLD + mesh = dolfinx.mesh.create_unit_cube(comm, 3, 3, 3, cell_type=cell_type) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", degree)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: 1.0 + 2.0 * x[0] - 3.0 * x[1] + 0.5 * x[2]) + + points = np.array([[0.21, 0.33, 0.47], [0.62, 0.11, 0.88], [0.5, 0.5, 0.5]]) + B = dolfinx_adjoint.PointObservation(V, points) + expected = 1.0 + 2.0 * points[:, 0] - 3.0 * points[:, 1] + 0.5 * points[:, 2] + + assert B.found.all() + assert np.allclose(B.evaluate(u), expected, atol=1e-12) + + +def test_distorted_hexahedra_are_handled(): + """A hexahedron's coordinate map is trilinear, so a distorted mesh is truly non-affine.""" + comm = MPI.COMM_WORLD + mesh = dolfinx.mesh.create_unit_cube(comm, 4, 4, 4, cell_type=dolfinx.mesh.CellType.hexahedron) + # Warp the geometry so that cell faces are no longer planar. + x = mesh.geometry.x + x[:, 0] += 0.08 * np.sin(3.0 * x[:, 1]) * x[:, 2] + x[:, 1] += 0.06 * np.cos(2.0 * x[:, 0]) * x[:, 2] + + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + # Q1 on a distorted hex reproduces affine fields exactly, unlike general quadratics. + u.interpolate(lambda x: 0.75 + x[0] - 2.0 * x[1] + 3.0 * x[2]) + + points = np.array([[0.5, 0.5, 0.5], [0.3, 0.62, 0.25], [0.7, 0.2, 0.8]]) + B = dolfinx_adjoint.PointObservation(V, points) + expected = 0.75 + B.points[:, 0] - 2.0 * B.points[:, 1] + 3.0 * B.points[:, 2] + + assert B.found.all() + assert np.allclose(B.evaluate(u), expected, atol=1e-11) + + +def test_affine_cell_types_classification(): + """Only degree-1 simplices take the fast affine pull-back.""" + comm = MPI.COMM_WORLD + triangles = unit_square(comm, 4) + assert observation._is_affine_simplex(triangles) + assert triangles.topology.cell_type in observation.AFFINE_CELL_TYPES + + quads = dolfinx.mesh.create_rectangle( + comm, + [np.array([0.0, 0.0]), np.array([1.0, 1.0])], + [3, 3], + cell_type=dolfinx.mesh.CellType.quadrilateral, + ) + assert not observation._is_affine_simplex(quads) + assert quads.topology.cell_type not in observation.AFFINE_CELL_TYPES + + +def test_quadrilateral_mesh_uses_the_generic_pull_back(): + """Non-simplex geometry falls back to the Newton pull-back and stays exact.""" + comm = MPI.COMM_WORLD + mesh = dolfinx.mesh.create_rectangle( + comm, + [np.array([0.0, 0.0]), np.array([1.0, 1.0])], + [5, 5], + cell_type=dolfinx.mesh.CellType.quadrilateral, + ) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: 2.0 * x[0] + 3.0 * x[1]) + + points = sample_points(13, seed=83) + B = dolfinx_adjoint.PointObservation(V, points) + assert np.allclose(B.gather(B.apply(u)), 2.0 * points[:, 0] + 3.0 * points[:, 1], atol=1e-12) + + +def test_to_scipy_matches_apply(): + comm = MPI.COMM_WORLD + scipy_sparse = pytest.importorskip("scipy.sparse") + assert scipy_sparse is not None + + V = dolfinx.fem.functionspace(unit_square(comm, 6), ("Lagrange", 1)) + B = dolfinx_adjoint.PointObservation(V, sample_points(15, seed=89)) + + rng = np.random.default_rng(97) + u = dolfinx.fem.Function(V) + u.x.array[:] = rng.random(u.x.array.shape) + u.x.scatter_forward() + + matrix = B.to_scipy() + assert matrix.shape[0] == B.num_local_rows + assert np.allclose(matrix @ u.x.array, B.apply(u)) + assert np.allclose(np.asarray(matrix.sum(axis=1)).reshape(-1), 1.0) + + +# --------------------------------------------------------------------------- +# The misfit +# --------------------------------------------------------------------------- + + +def test_misfit_matches_a_manual_computation(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: x[0] + x[1]) + + points = sample_points(15, seed=19) + B = dolfinx_adjoint.PointObservation(V, points) + data = np.random.default_rng(23).random(B.num_points) + noise_variance = 0.25 + + with pyadjoint.stop_annotating(): + J = dolfinx_adjoint.point_observation_misfit(u, B, data, noise_variance=noise_variance) + + expected = 0.5 * np.sum(((points[:, 0] + points[:, 1]) - data) ** 2) / noise_variance + assert np.isclose(float(J), expected) + + +def test_misfit_vanishes_at_the_exact_data(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: 2.0 * x[0] - x[1]) + + B = dolfinx_adjoint.PointObservation(V, sample_points(13, seed=29)) + with pyadjoint.stop_annotating(): + J = dolfinx_adjoint.point_observation_misfit(u, B, B.gather(B.apply(u))) + assert np.isclose(float(J), 0.0, atol=1e-20) + + +def test_misfit_weights_mask_rows_out(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm), ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: x[0]) + + points = sample_points(10, seed=31) + B = dolfinx_adjoint.PointObservation(V, points) + weights = np.zeros(B.num_points) + weights[:4] = 1.0 + + with pyadjoint.stop_annotating(): + J = dolfinx_adjoint.point_observation_misfit(u, B, np.zeros(B.num_points), weights=weights) + assert np.isclose(float(J), 0.5 * np.sum(points[:4, 0] ** 2)) + + +def test_zero_noise_variance_raises(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + B = dolfinx_adjoint.PointObservation(V, sample_points(4, seed=37)) + with pytest.raises(ZeroDivisionError): + dolfinx_adjoint.point_observation_misfit(u, B, np.zeros(B.num_points), noise_variance=0.0) + + +def test_mismatched_data_length_raises(): + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + u = dolfinx.fem.Function(V) + B = dolfinx_adjoint.PointObservation(V, sample_points(6, seed=41)) + with pytest.raises(ValueError, match="Expected an array of length"): + dolfinx_adjoint.point_observation_misfit(u, B, np.zeros(B.num_points + 3)) + + +# --------------------------------------------------------------------------- +# Differentiability +# --------------------------------------------------------------------------- + + +def test_gradient_taylor_test(): + comm = MPI.COMM_WORLD + pyadjoint.get_working_tape().clear_tape() + V = dolfinx.fem.functionspace(unit_square(comm, 6), ("Lagrange", 1)) + + u = dolfinx_adjoint.Function(V, name="u") + u.interpolate(lambda x: x[0] * x[1]) + + B = dolfinx_adjoint.PointObservation(V, sample_points(21, seed=37)) + rng = np.random.default_rng(41) + data = rng.random(B.num_points) + + control = pyadjoint.Control(u) + J = dolfinx_adjoint.point_observation_misfit(u, B, data, noise_variance=0.5) + Jhat = pyadjoint.ReducedFunctional(J, control) + + h = dolfinx_adjoint.Function(V) + h.x.array[:] = rng.random(h.x.array.shape) + h.x.scatter_forward() + + assert pyadjoint.taylor_test(Jhat, u, h) > 1.9 + pyadjoint.get_working_tape().clear_tape() + + +def test_gradient_matches_the_closed_form(): + """dJ/du == sigma^-2 B^T (Bu - d).""" + comm = MPI.COMM_WORLD + pyadjoint.get_working_tape().clear_tape() + V = dolfinx.fem.functionspace(unit_square(comm, 6), ("Lagrange", 1)) + + u = dolfinx_adjoint.Function(V, name="u") + u.interpolate(lambda x: x[0] + 0.5 * x[1]) + + B = dolfinx_adjoint.PointObservation(V, sample_points(19, seed=43)) + data = np.random.default_rng(47).random(B.num_points) + noise_variance = 0.75 + + J = dolfinx_adjoint.point_observation_misfit(u, B, data, noise_variance=noise_variance) + gradient = pyadjoint.ReducedFunctional(J, pyadjoint.Control(u)).derivative() + + expected = B.apply_transpose((B.apply(u) - B.restrict(data)) / noise_variance) + n = owned_dofs(V) + assert np.allclose(gradient.x.array[:n], expected.array[:n]) + pyadjoint.get_working_tape().clear_tape() + + +def test_weighted_derivatives_apply_the_weights_twice(): + """With weights, dJ/du is sigma^-2 B^T W^2 (Bu - d), not W applied once. + + Differentiating ||W r||^2 brings down one W from the norm and one from the residual, + so a non-binary weight vector is the case that would expose getting this wrong. + """ + pyadjoint.get_working_tape().clear_tape() + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 6), ("Lagrange", 1)) + + rng = np.random.default_rng(5) + B = dolfinx_adjoint.PointObservation(V, sample_points(17, seed=113)) + data = rng.random(B.num_points) + weights = rng.random(B.num_points) + noise_variance = 0.7 + + u = dolfinx_adjoint.Function(V, name="u") + u.interpolate(lambda x: x[0] * x[1]) + + J = dolfinx_adjoint.point_observation_misfit(u, B, data, noise_variance=noise_variance, weights=weights) + Jhat = pyadjoint.ReducedFunctional(J, pyadjoint.Control(u)) + + h = dolfinx_adjoint.Function(V) + h.interpolate(lambda x: np.sin(3.0 * x[0]) + x[1]) + assert pyadjoint.taylor_test(Jhat, u, h) > 1.9 + + # taylor_test leaves the control perturbed, so re-evaluate before differentiating. + Jhat(u) + gradient = Jhat.derivative() + + local_weights = B.restrict(weights) + residual = B.apply(u) - B.restrict(data) + expected = B.apply_transpose(local_weights**2 * residual / noise_variance) + + n = owned_dofs(V) + assert np.allclose(gradient.x.array[:n], expected.array[:n]) + + curvature = Jhat.hessian(h) + expected_curvature = B.apply_transpose(local_weights**2 * B.apply(h) / noise_variance) + assert np.allclose(curvature.x.array[:n], expected_curvature.array[:n]) + pyadjoint.get_working_tape().clear_tape() + + +def test_hessian_is_the_gauss_newton_operator(): + """The misfit is quadratic, so its Hessian action is sigma^-2 B^T B h.""" + comm = MPI.COMM_WORLD + pyadjoint.get_working_tape().clear_tape() + V = dolfinx.fem.functionspace(unit_square(comm, 6), ("Lagrange", 1)) + + u = dolfinx_adjoint.Function(V, name="u") + u.interpolate(lambda x: x[0] - x[1]) + + B = dolfinx_adjoint.PointObservation(V, sample_points(17, seed=53)) + rng = np.random.default_rng(59) + noise_variance = 2.0 + + J = dolfinx_adjoint.point_observation_misfit(u, B, rng.random(B.num_points), noise_variance=noise_variance) + Jhat = pyadjoint.ReducedFunctional(J, pyadjoint.Control(u)) + Jhat(u) + Jhat.derivative() + + h = dolfinx_adjoint.Function(V) + h.x.array[:] = rng.random(h.x.array.shape) + h.x.scatter_forward() + + expected = B.apply_transpose(B.apply(h) / noise_variance) + n = owned_dofs(V) + assert np.allclose(Jhat.hessian(h).x.array[:n], expected.array[:n]) + pyadjoint.get_working_tape().clear_tape() + + +def test_taylor_test_through_a_pde_solve(): + """The full parameter-to-observable map is differentiable: m -> u(m) -> B u(m).""" + comm = MPI.COMM_WORLD + pyadjoint.get_working_tape().clear_tape() + mesh = unit_square(comm, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + # -div(exp(m) grad u) = f, with m the log-diffusivity we differentiate against. + m = dolfinx_adjoint.Function(V, name="m") + m.x.array[:] = 0.3 + + u = dolfinx_adjoint.Function(V, name="u") + trial, test = ufl.TrialFunction(V), ufl.TestFunction(V) + x = ufl.SpatialCoordinate(mesh) + f = ufl.sin(np.pi * x[0]) * ufl.sin(np.pi * x[1]) + a = ufl.inner(ufl.exp(m) * ufl.grad(trial), ufl.grad(test)) * ufl.dx + L = ufl.inner(f, test) * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, facets) + u_bc = dolfinx_adjoint.Function(V, name="u_bc") + bc = dolfinx_adjoint.dirichletbc(u_bc, dofs) + + problem = dolfinx_adjoint.LinearProblem(a, L, u=u, bcs=[bc], petsc_options_prefix="test_observation_") + problem.solve() + + B = dolfinx_adjoint.PointObservation(V, sample_points(15, seed=61)) + rng = np.random.default_rng(67) + data = B.gather(B.apply(u)) + 0.01 * rng.standard_normal(B.num_points) + + J = dolfinx_adjoint.point_observation_misfit(u, B, data, noise_variance=0.01) + Jhat = pyadjoint.ReducedFunctional(J, pyadjoint.Control(m)) + + h = dolfinx_adjoint.Function(V) + h.x.array[:] = 0.1 * rng.standard_normal(h.x.array.shape) + h.x.scatter_forward() + + assert pyadjoint.taylor_test(Jhat, m, h) > 1.9 + pyadjoint.get_working_tape().clear_tape() + + +def test_recovers_a_known_source_from_point_data(): + """A small end-to-end inversion: point data is enough to pin down a scalar source.""" + comm = MPI.COMM_WORLD + pyadjoint.get_working_tape().clear_tape() + mesh = unit_square(comm, 8) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + true_scale = 3.0 + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, facets) + + def solve_for(scale_value, prefix): + """-div(grad u) = scale * sin(pi x) sin(pi y), with u = 0 on the boundary.""" + scale = dolfinx_adjoint.Constant(mesh, scale_value) + state = dolfinx_adjoint.Function(V, name="u") + trial, test = ufl.TrialFunction(V), ufl.TestFunction(V) + x = ufl.SpatialCoordinate(mesh) + a = ufl.inner(ufl.grad(trial), ufl.grad(test)) * ufl.dx + L = ufl.inner(scale * ufl.sin(np.pi * x[0]) * ufl.sin(np.pi * x[1]), test) * ufl.dx + u_bc = dolfinx_adjoint.Function(V, name="u_bc") + bc = dolfinx_adjoint.dirichletbc(u_bc, dofs) + dolfinx_adjoint.LinearProblem(a, L, u=state, bcs=[bc], petsc_options_prefix=prefix).solve() + return scale, state + + points = sample_points(12, seed=71) + with pyadjoint.stop_annotating(): + _, truth = solve_for(true_scale, "test_recover_truth_") + B_truth = dolfinx_adjoint.PointObservation(V, points) + observations = B_truth.gather(B_truth.apply(truth)) + + scale, state = solve_for(1.0, "test_recover_") + B = dolfinx_adjoint.PointObservation(V, points) + J = dolfinx_adjoint.point_observation_misfit(state, B, observations) + Jhat = pyadjoint.ReducedFunctional(J, pyadjoint.Control(scale)) + + # The observations depend linearly on the source scale, so the misfit is an exact + # quadratic and a single Newton step from the initial guess lands on the true value. + gradient = Jhat.derivative() + curvature = Jhat.hessian(dolfinx_adjoint.Constant(mesh, 1.0)) + recovered = 1.0 - float(gradient.x.array[0]) / float(curvature.x.array[0]) + + assert np.isclose(recovered, true_scale, rtol=1e-8) + pyadjoint.get_working_tape().clear_tape() From c8d789ba59ef29dd1c6be0c28594189ddc61f7ff Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 14 Aug 2026 13:00:50 +0000 Subject: [PATCH 02/11] Use interpolation from fenicsx_ii and point mesh --- pyproject.toml | 1 + src/dolfinx_adjoint/observation.py | 214 ++++++++++++----------------- tests/test_observation.py | 77 ++++++++--- 3 files changed, 146 insertions(+), 146 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0c7abe4..c532f58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ readme = "README.md" dependencies = [ "fenics-dolfinx>=0.10.0", "pyadjoint-ad>=2025.10.0", + "fenicsx-ii", "typing_extensions; python_version < '3.11'", "packaging>=24.2", ] diff --git a/src/dolfinx_adjoint/observation.py b/src/dolfinx_adjoint/observation.py index 6aef371..a138362 100644 --- a/src/dolfinx_adjoint/observation.py +++ b/src/dolfinx_adjoint/observation.py @@ -37,6 +37,8 @@ import numpy as np import numpy.typing as npt import pyadjoint +from fenicsx_ii import create_interpolation_matrix +from fenicsx_ii.quadrature import Quadrature from pyadjoint.overloaded_type import create_overloaded_object from pyadjoint.tape import annotate_tape, get_working_tape, stop_annotating @@ -45,38 +47,35 @@ __all__ = ["PointObservation", "point_observation_misfit"] -def _geometry_dofmap(mesh: dolfinx.mesh.Mesh): - """``mesh.geometry.dofmap`` was renamed to ``dofmaps[0]``.""" - try: - return mesh.geometry.dofmaps[0] - except AttributeError: - return mesh.geometry.dofmap - - -def _coordinate_map(mesh: dolfinx.mesh.Mesh): - """``mesh.geometry.cmap`` is deprecated in favour of ``cmaps[0]``.""" - try: - return mesh.geometry.cmaps[0] - except (AttributeError, IndexError): - return mesh.geometry.cmap +class _PointCloudTrace: + """Reduction operator handing :mod:`fenicsx_ii` the coordinates of a point cloud. + ``fenicsx_ii.PointwiseTrace`` is written for 1D line meshes and obtains the physical + coordinates by compiling a ``SpatialCoordinate`` expression, which FFCx cannot do on a + point cell. On a point mesh each cell *is* a single geometry node, so the coordinates + can simply be read off the geometry. + """ -#: Cell types whose degree-1 coordinate map is affine. Tensor-product cells -#: (quadrilateral, hexahedron, prism, pyramid) are multilinear rather than affine even at -#: degree 1, so they take the general Newton pull-back. -AFFINE_CELL_TYPES = frozenset( - { - dolfinx.mesh.CellType.point, - dolfinx.mesh.CellType.interval, - dolfinx.mesh.CellType.triangle, - dolfinx.mesh.CellType.tetrahedron, - } -) + def __init__(self, mesh: dolfinx.mesh.Mesh) -> None: + self._mesh = mesh + + def compute_quadrature(self, cells: npt.NDArray[np.int32], reference_points: npt.NDArray[np.floating]): + gdim = self._mesh.geometry.dim + nodes = np.asarray(self._mesh.geometry.dofmaps[0])[cells].reshape(-1) + points = self._mesh.geometry.x[nodes][:, :gdim] + return Quadrature( + name="PointCloud", + points=points, + weights=np.ones((points.shape[0], 1), dtype=points.dtype), + scales=np.ones(points.shape[0], dtype=points.dtype), + ) + @property + def num_points(self) -> int: + return 1 -def _is_affine_simplex(mesh: dolfinx.mesh.Mesh) -> bool: - """True if every cell map is affine, i.e. a degree-1 simplex geometry.""" - return _coordinate_map(mesh).degree == 1 and mesh.topology.cell_type in AFFINE_CELL_TYPES + def __str__(self) -> str: + return f"PointCloudTrace({self._mesh})" def _default_padding(mesh: dolfinx.mesh.Mesh) -> float: @@ -162,7 +161,7 @@ def __init__( "(for instance high-order elements on oriented, non-simplex cells)." ) try: - self._basix_element = V.element.basix_element + V.element.basix_element except RuntimeError as exc: raise NotImplementedError( "PointObservation needs a function space backed by a Basix element, which mixed " @@ -183,11 +182,17 @@ def __init__( self.points = padded_points tdim = mesh.topology.dim - gdim = mesh.geometry.dim self.padding = _default_padding(mesh) if padding is None else float(padding) - # Search owned cells only: a point interior to a cell is then found on exactly one - # rank, and only points on a shared facet or vertex remain ambiguous. + # Locate the points with an exact containment test, over owned cells only. + # + # This is the one piece that is *not* delegated. `fenicsx_ii` locates points with + # `dolfinx.geometry.determine_point_ownership`, and asserts that every point is + # found -- but that routine snaps points to nearby cells rather than testing + # containment. On a brain-surface mesh it claimed 13.5% more points than are + # actually inside, every one verifiably outside the cell it was assigned to. + # Filtering here means only genuinely interior points reach the point mesh, so the + # assertion holds and no phantom observations enter the misfit. num_owned_cells = mesh.topology.index_map(tdim).size_local first_cell = np.full(num_points, -1, dtype=np.int32) if num_owned_cells > 0 and num_points > 0: @@ -199,7 +204,9 @@ def __init__( has_cell = offsets[1:] > offsets[:-1] first_cell[has_cell] = colliding.array[offsets[:-1][has_cell]] - # Break ties by lowest rank, so every rank agrees on a single owner per point. + # Searching owned cells only means a point interior to a cell is found on exactly one + # process; break the remaining ties, on shared facets and vertices, by lowest rank so + # that every process agrees on a single owner per point. candidate_owner = np.where(first_cell >= 0, comm.rank, comm.size).astype(np.int32) owner = np.empty_like(candidate_owner) comm.Allreduce(candidate_owner, owner, op=MPI.MIN) @@ -209,73 +216,40 @@ def __init__( self.num_found = int(self.found.sum()) self.local_indices = np.flatnonzero(owner == comm.rank).astype(np.int32) - self._assemble(padded_points[self.local_indices], first_cell[self.local_indices], gdim) + self._build_matrix(padded_points[self.local_indices]) # -------------------------------------------------------------------- assembly --- - def _assemble(self, points: np.ndarray, cells: np.ndarray, gdim: int) -> None: - """Tabulate the basis functions of the containing cell for every owned point. - - The operator is stored as two dense ``(num_local_rows, num_dofs_per_cell)`` arrays - -- the basis values and the corresponding *ghosted* local DOF indices -- rather - than a sparse matrix. Every row has exactly the same number of entries, so this is - both smaller and faster than a general sparse format, and keeps the package free - of a sparse-matrix dependency. + def _build_matrix(self, local_points: np.ndarray) -> None: + """Build the interpolation matrix from ``V`` onto a point mesh of the owned points. + + The point mesh carries one cell per observation point this process owns, and a + DG-0 space on it has exactly one degree of freedom per point (``block_size`` of + them for a vector space), so the interpolation matrix from ``V`` onto that space + *is* :math:`B`. :mod:`fenicsx_ii` assembles it, including all the cross-process + communication, and its transpose gives the adjoint for free. """ V = self.function_space - dofmap = V.dofmap - bs = dofmap.bs - self._num_local_dofs = (dofmap.index_map.size_local + dofmap.index_map.num_ghosts) * bs - - if len(points) == 0: - num_dofs_per_cell = dofmap.dof_layout.num_dofs - self._basis = np.zeros((0, num_dofs_per_cell)) - self._cols = np.zeros((0, num_dofs_per_cell), dtype=np.int32) - return + gdim = V.mesh.geometry.dim + bs = V.dofmap.bs - reference_points = self._pull_back(points, cells, gdim) - - # The reference basis is shared by all cells, so a single tabulate call suffices. - basis = np.asarray(self._basix_element.tabulate(0, reference_points))[0, :, :, 0] - cell_dofs = np.asarray(dofmap.list)[cells] - - if bs == 1: - self._basis = basis - self._cols = cell_dofs.astype(np.int32) - else: - # Unroll to one row per (point, component): row i*bs + c reads the DOFs of - # component c, weighted by the same scalar basis values. - num_rows, num_dofs_per_cell = basis.shape - self._basis = np.repeat(basis, bs, axis=0) - components = np.arange(bs, dtype=np.int32) - cols = cell_dofs[:, None, :] * bs + components[None, :, None] - self._cols = cols.reshape(num_rows * bs, num_dofs_per_cell).astype(np.int32) - - def _pull_back(self, points: np.ndarray, cells: np.ndarray, gdim: int) -> np.ndarray: - """Map physical points to the reference coordinates of their containing cell.""" - mesh = self.function_space.mesh - geometry_dofmap = _geometry_dofmap(mesh) - geometry_x = mesh.geometry.x - tdim = mesh.topology.dim - dtype = geometry_x.dtype - - if _is_affine_simplex(mesh) and gdim == tdim: - # X = x0 + J xi, so xi = J^{-1} (X - x0) with a constant Jacobian per cell. - coords = geometry_x[np.asarray(geometry_dofmap)[cells]][:, : tdim + 1, :gdim] - jacobian = np.swapaxes(coords[:, 1:, :] - coords[:, :1, :], 1, 2) - offsets = (points[:, :gdim] - coords[:, 0, :])[..., None] - return np.linalg.solve(jacobian, offsets)[..., 0].astype(dtype) - - # Otherwise fall back to the coordinate element's Newton iteration, which handles - # one cell at a time; group the points by cell to call it as rarely as possible. - cmap = _coordinate_map(mesh) - reference_points = np.zeros((len(points), tdim), dtype=dtype) - order = np.argsort(cells, kind="stable") - boundaries = np.flatnonzero(np.diff(cells[order])) + 1 - for group in np.split(order, boundaries): - cell_coords = geometry_x[geometry_dofmap[cells[group[0]]]][:, :gdim] - physical = np.ascontiguousarray(points[group, :gdim], dtype=dtype) - reference_points[group] = cmap.pull_back(physical, cell_coords) - return reference_points + self.point_mesh = dolfinx.mesh.create_point_mesh( + self.comm, np.ascontiguousarray(local_points[:, :gdim], dtype=V.mesh.geometry.x.dtype) + ) + element = ("DG", 0) if bs == 1 else ("DG", 0, (bs,)) + self.observation_space = dolfinx.fem.functionspace(self.point_mesh, element) + + self._matrix, _, _ = create_interpolation_matrix( + V, + self.observation_space, + _PointCloudTrace(self.point_mesh), + tol=self.padding, + use_petsc=True, + ) + # Reusable work vectors, so that apply/apply_transpose do not allocate per call. + self._observation_function = dolfinx.fem.Function(self.observation_space) + self._state_function = dolfinx.fem.Function(V) + index_map = self.observation_space.dofmap.index_map + self._num_local_rows = index_map.size_local * self.observation_space.dofmap.index_map_bs # --------------------------------------------------------------------- actions --- @property @@ -285,23 +259,28 @@ def block_size(self) -> int: @property def num_local_rows(self) -> int: - """Number of rows owned by this rank, including the block-size unrolling.""" - return self._basis.shape[0] + """Number of rows owned by this process, including the block-size unrolling.""" + return self._num_local_rows + + @property + def matrix(self): + """The interpolation matrix :math:`B`, as a distributed ``PETSc.Mat``.""" + return self._matrix def apply(self, u: dolfinx.fem.Function) -> npt.NDArray[np.float64]: - """Evaluate :math:`Bu` for the rows owned by this rank. + """Evaluate :math:`Bu` for the rows owned by this process. Args: u: Function in :attr:`function_space`. Returns: - The point values of the rows owned by this rank, ordered as + The point values of the rows owned by this process, ordered as :attr:`local_indices` (component-fastest for vector spaces). """ u.x.scatter_forward() - if self.num_local_rows == 0: - return np.zeros(0) - return np.einsum("ij,ij->i", self._basis, u.x.array[self._cols]) + self._matrix.mult(u.x.petsc_vec, self._observation_function.x.petsc_vec) + self._observation_function.x.scatter_forward() + return np.asarray(self._observation_function.x.array[: self.num_local_rows], dtype=np.float64) def evaluate(self, u: dolfinx.fem.Function, fill: float = np.nan) -> npt.NDArray[np.float64]: """Evaluate ``u`` at every observation point. @@ -335,13 +314,14 @@ def apply_transpose(self, values: npt.ArrayLike, out: dolfinx.la.Vector | None = if out is None: out = dolfinx.la.vector(V.dofmap.index_map, V.dofmap.bs, dtype=V.mesh.geometry.x.dtype) out.array[:] = 0.0 - if self.num_local_rows > 0: - weights = self._basis * np.asarray(values, dtype=np.float64)[:, None] - contributions = np.bincount( - self._cols.reshape(-1), weights=weights.reshape(-1), minlength=self._num_local_dofs - ) - out.array[:] += contributions - out.scatter_reverse(dolfinx.la.InsertMode.add) + + self._observation_function.x.array[: self.num_local_rows] = np.asarray(values, dtype=np.float64) + self._observation_function.x.scatter_forward() + # PETSc performs the reverse communication itself, so no scatter_reverse is needed. + self._matrix.multTranspose(self._observation_function.x.petsc_vec, self._state_function.x.petsc_vec) + self._state_function.x.scatter_forward() + + out.array[:] += self._state_function.x.array out.scatter_forward() return out @@ -383,22 +363,6 @@ def gather(self, values: npt.ArrayLike, fill: float = np.nan) -> npt.NDArray[np. total.reshape(self.num_points, bs)[missing] = fill return total - def to_scipy(self): - """This rank's block of :math:`B` as a ``scipy.sparse`` CSR matrix. - - Columns index the ghosted local DOF array (``u.x.array``). Requires ``scipy``, - which is not a dependency of this package; :meth:`apply` and - :meth:`apply_transpose` do not need it. - """ - import scipy.sparse - - num_rows, num_dofs_per_cell = self._basis.shape - rows = np.repeat(np.arange(num_rows), num_dofs_per_cell) - return scipy.sparse.csr_matrix( - (self._basis.reshape(-1), (rows, self._cols.reshape(-1))), - shape=(num_rows, self._num_local_dofs), - ) - def point_observation_misfit( u: dolfinx.fem.Function, diff --git a/tests/test_observation.py b/tests/test_observation.py index 753fcff..15ccbe9 100644 --- a/tests/test_observation.py +++ b/tests/test_observation.py @@ -176,6 +176,37 @@ def test_ownership_is_unique_and_agreed_on_by_every_rank(): assert np.array_equal(other, B.owner) +def test_ambiguous_points_are_owned_exactly_once(): + """Points on cell vertices and facets lie in several cells, and often on several processes. + + These are the cases the ownership tie-break exists for, so they are worth pinning down + separately from randomly scattered interior points. + """ + comm = MPI.COMM_WORLD + n = 8 + mesh = unit_square(comm, n) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + axis = np.linspace(0.0, 1.0, n + 1) + vertices = np.stack(np.meshgrid(axis, axis, indexing="ij"), axis=-1).reshape(-1, 2) + midpoints = 0.5 * (vertices[:-1] + vertices[1:]) + points = np.vstack([vertices, midpoints]) + + B = dolfinx_adjoint.PointObservation(V, points) + assert B.found.all() + + claimed = np.zeros(B.num_points, dtype=np.int64) + claimed[B.local_indices] = 1 + total = np.empty_like(claimed) + comm.Allreduce(claimed, total, op=MPI.SUM) + assert np.all(total == 1), "every point must be owned by exactly one process" + + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: 1.0 + 2.0 * x[0] - 3.0 * x[1]) + expected = 1.0 + 2.0 * points[:, 0] - 3.0 * points[:, 1] + assert np.allclose(B.evaluate(u), expected, atol=1e-12) + + def test_transpose_is_the_adjoint_of_apply(): """ == globally, including ghost contributions.""" comm = MPI.COMM_WORLD @@ -341,21 +372,22 @@ def test_distorted_hexahedra_are_handled(): assert np.allclose(B.evaluate(u), expected, atol=1e-11) -def test_affine_cell_types_classification(): - """Only degree-1 simplices take the fast affine pull-back.""" +def test_matrix_is_distributed_over_the_points(): + """The operator is an interpolation matrix onto a point mesh of the observed points.""" comm = MPI.COMM_WORLD - triangles = unit_square(comm, 4) - assert observation._is_affine_simplex(triangles) - assert triangles.topology.cell_type in observation.AFFINE_CELL_TYPES + V = dolfinx.fem.functionspace(unit_square(comm, 8), ("Lagrange", 1)) + points = sample_points(31, seed=91) + B = dolfinx_adjoint.PointObservation(V, points) - quads = dolfinx.mesh.create_rectangle( - comm, - [np.array([0.0, 0.0]), np.array([1.0, 1.0])], - [3, 3], - cell_type=dolfinx.mesh.CellType.quadrilateral, - ) - assert not observation._is_affine_simplex(quads) - assert quads.topology.cell_type not in observation.AFFINE_CELL_TYPES + rows, columns = B.matrix.getSizes() + assert rows[1] == B.num_found # one global row per located point + assert columns[1] == V.dofmap.index_map.size_global * V.dofmap.bs + assert rows[0] == B.num_local_rows + assert comm.allreduce(B.num_local_rows, op=MPI.SUM) == B.num_found + + # The point mesh carries exactly the points this process owns. + owned = B.point_mesh.geometry.x[: B.num_local_rows, :2] + assert np.allclose(owned, points[B.local_indices]) def test_quadrilateral_mesh_uses_the_generic_pull_back(): @@ -376,11 +408,9 @@ def test_quadrilateral_mesh_uses_the_generic_pull_back(): assert np.allclose(B.gather(B.apply(u)), 2.0 * points[:, 0] + 3.0 * points[:, 1], atol=1e-12) -def test_to_scipy_matches_apply(): +def test_matrix_action_matches_apply(): + """`apply` is exactly a matrix-vector product with the interpolation matrix.""" comm = MPI.COMM_WORLD - scipy_sparse = pytest.importorskip("scipy.sparse") - assert scipy_sparse is not None - V = dolfinx.fem.functionspace(unit_square(comm, 6), ("Lagrange", 1)) B = dolfinx_adjoint.PointObservation(V, sample_points(15, seed=89)) @@ -389,10 +419,15 @@ def test_to_scipy_matches_apply(): u.x.array[:] = rng.random(u.x.array.shape) u.x.scatter_forward() - matrix = B.to_scipy() - assert matrix.shape[0] == B.num_local_rows - assert np.allclose(matrix @ u.x.array, B.apply(u)) - assert np.allclose(np.asarray(matrix.sum(axis=1)).reshape(-1), 1.0) + result = dolfinx.fem.Function(B.observation_space) + B.matrix.mult(u.x.petsc_vec, result.x.petsc_vec) + result.x.scatter_forward() + assert np.allclose(result.x.array[: B.num_local_rows], B.apply(u)) + + # Rows form a partition of unity, so the matrix reproduces constants exactly. + ones = dolfinx.fem.Function(V) + ones.x.array[:] = 1.0 + assert np.allclose(B.apply(ones), 1.0) # --------------------------------------------------------------------------- From eaf8175e4a59807e53535e111c930cdbfb4e284d Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 14 Aug 2026 13:08:04 +0000 Subject: [PATCH 03/11] Fix mypy --- demos/point_observations.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/demos/point_observations.py b/demos/point_observations.py index b1ca324..ff43d6b 100644 --- a/demos/point_observations.py +++ b/demos/point_observations.py @@ -203,7 +203,7 @@ def true_source(x): cells, types, geometry = dolfinx.plot.vtk_mesh(V) sensor_cloud = pyvista.PolyData(np.column_stack([sensors, np.zeros(len(sensors))])) - plotter = pyvista.Plotter(shape=(1, 2), window_size=(900, 400)) + plotter = pyvista.Plotter(shape=(1, 2), window_size=[900, 400]) for column, (field, title) in enumerate([(f_true, "True source"), (f_opt, "Recovered source")]): grid = pyvista.UnstructuredGrid(cells, types, geometry) grid.point_data["f"] = field.x.array.real @@ -212,7 +212,7 @@ def true_source(x): plotter.add_text(title, font_size=10) plotter.add_mesh(grid, show_edges=False, clim=[0.0, 5.0]) plotter.add_mesh(sensor_cloud, color="black", point_size=5, render_points_as_spheres=True) - plotter.view_xy() + plotter.view_xy() # type: ignore[call-arg] if pyvista.OFF_SCREEN: plotter.screenshot("point_observations.png") else: From dca0897a6466d8867e19bba9b628246f6d7cca90 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 14 Aug 2026 13:24:36 +0000 Subject: [PATCH 04/11] Ruff --- demos/point_observations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demos/point_observations.py b/demos/point_observations.py index ff43d6b..12992d1 100644 --- a/demos/point_observations.py +++ b/demos/point_observations.py @@ -212,7 +212,7 @@ def true_source(x): plotter.add_text(title, font_size=10) plotter.add_mesh(grid, show_edges=False, clim=[0.0, 5.0]) plotter.add_mesh(sensor_cloud, color="black", point_size=5, render_points_as_spheres=True) - plotter.view_xy() # type: ignore[call-arg] + plotter.view_xy() # type: ignore[call-arg] if pyvista.OFF_SCREEN: plotter.screenshot("point_observations.png") else: From 0715f9185ef11aac5487854d66d5a27298469383 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 14 Aug 2026 13:08:04 +0000 Subject: [PATCH 05/11] Fix mypy --- demos/point_observations.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/demos/point_observations.py b/demos/point_observations.py index b1ca324..ff43d6b 100644 --- a/demos/point_observations.py +++ b/demos/point_observations.py @@ -203,7 +203,7 @@ def true_source(x): cells, types, geometry = dolfinx.plot.vtk_mesh(V) sensor_cloud = pyvista.PolyData(np.column_stack([sensors, np.zeros(len(sensors))])) - plotter = pyvista.Plotter(shape=(1, 2), window_size=(900, 400)) + plotter = pyvista.Plotter(shape=(1, 2), window_size=[900, 400]) for column, (field, title) in enumerate([(f_true, "True source"), (f_opt, "Recovered source")]): grid = pyvista.UnstructuredGrid(cells, types, geometry) grid.point_data["f"] = field.x.array.real @@ -212,7 +212,7 @@ def true_source(x): plotter.add_text(title, font_size=10) plotter.add_mesh(grid, show_edges=False, clim=[0.0, 5.0]) plotter.add_mesh(sensor_cloud, color="black", point_size=5, render_points_as_spheres=True) - plotter.view_xy() + plotter.view_xy() # type: ignore[call-arg] if pyvista.OFF_SCREEN: plotter.screenshot("point_observations.png") else: From 28fcea6518014f8c50afc1f9c95624d100caceee Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 14 Aug 2026 13:24:36 +0000 Subject: [PATCH 06/11] Ruff --- demos/point_observations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demos/point_observations.py b/demos/point_observations.py index ff43d6b..12992d1 100644 --- a/demos/point_observations.py +++ b/demos/point_observations.py @@ -212,7 +212,7 @@ def true_source(x): plotter.add_text(title, font_size=10) plotter.add_mesh(grid, show_edges=False, clim=[0.0, 5.0]) plotter.add_mesh(sensor_cloud, color="black", point_size=5, render_points_as_spheres=True) - plotter.view_xy() # type: ignore[call-arg] + plotter.view_xy() # type: ignore[call-arg] if pyvista.OFF_SCREEN: plotter.screenshot("point_observations.png") else: From 8c69fbd2bf4af129f0e885716ca1c4cf9fff190c Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Tue, 18 Aug 2026 10:32:30 +0000 Subject: [PATCH 07/11] Remove long comment --- src/dolfinx_adjoint/observation.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/dolfinx_adjoint/observation.py b/src/dolfinx_adjoint/observation.py index a138362..947c722 100644 --- a/src/dolfinx_adjoint/observation.py +++ b/src/dolfinx_adjoint/observation.py @@ -185,14 +185,6 @@ def __init__( self.padding = _default_padding(mesh) if padding is None else float(padding) # Locate the points with an exact containment test, over owned cells only. - # - # This is the one piece that is *not* delegated. `fenicsx_ii` locates points with - # `dolfinx.geometry.determine_point_ownership`, and asserts that every point is - # found -- but that routine snaps points to nearby cells rather than testing - # containment. On a brain-surface mesh it claimed 13.5% more points than are - # actually inside, every one verifiably outside the cell it was assigned to. - # Filtering here means only genuinely interior points reach the point mesh, so the - # assertion holds and no phantom observations enter the misfit. num_owned_cells = mesh.topology.index_map(tdim).size_local first_cell = np.full(num_points, -1, dtype=np.int32) if num_owned_cells > 0 and num_points > 0: From 7ff072e0f267de4850d52b678983025a74c36e4f Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Wed, 19 Aug 2026 08:03:56 +0000 Subject: [PATCH 08/11] Fix comments --- demos/point_observations.py | 12 ---- src/dolfinx_adjoint/blocks/observation.py | 39 +++++++---- src/dolfinx_adjoint/observation.py | 80 +++++++++++++++++------ tests/test_observation.py | 55 ++++++++++++++++ 4 files changed, 140 insertions(+), 46 deletions(-) diff --git a/demos/point_observations.py b/demos/point_observations.py index 12992d1..672a435 100644 --- a/demos/point_observations.py +++ b/demos/point_observations.py @@ -219,15 +219,3 @@ def true_source(x): plotter.show() # - - -# ## Going further -# -# * **Sensors outside the domain.** `B.found` flags any points that fall outside the mesh, and -# those are left out of the misfit rather than being treated as measurements of zero. -# * **Unreliable sensors.** `point_observation_misfit(..., weights=...)` reweights individual -# measurements; a zero weight drops one entirely. -# * **Noisier data.** `point_observation_misfit(..., noise_variance=...)` divides the misfit by -# $\sigma^2$, which is what makes it a negative log-likelihood when the regularization term is -# a genuine prior. -# * **Time-dependent data.** Build `B` once and call `point_observation_misfit` once per -# observation time, summing the results -- the operator only depends on the sensor positions. diff --git a/src/dolfinx_adjoint/blocks/observation.py b/src/dolfinx_adjoint/blocks/observation.py index 83bb379..187903a 100644 --- a/src/dolfinx_adjoint/blocks/observation.py +++ b/src/dolfinx_adjoint/blocks/observation.py @@ -48,9 +48,12 @@ def __init__( super().__init__(ad_block_tag=ad_block_tag) self.add_dependency(u) self.observation = observation - self.data = data + # Copied: the tape holds this block for as long as it is needed for differentiation, + # so a caller mutating a data/weights buffer it passed in (for instance reusing one + # local-length array across a time-stepping loop) must not retroactively change it. + self.data = np.array(data, dtype=np.float64, copy=True) self.noise_variance = noise_variance - self.weights = weights + self.weights = None if weights is None else np.array(weights, dtype=np.float64, copy=True) def __str__(self) -> str: return f"point_observation_misfit({self.observation.num_found} points)" @@ -63,13 +66,22 @@ def _apply_weights(self, values: npt.NDArray[np.float64]) -> npt.NDArray[np.floa return values return self.weights * values - def _transpose_action(self, residual: npt.NDArray[np.float64], scale: float) -> _SpecialVector: - """:math:`\\mathrm{scale} \\cdot \\sigma^{-2} B^T W^2 r`, as a DOF vector.""" + def _weighted_row(self, values: npt.NDArray[np.float64], scale: float) -> npt.NDArray[np.float64]: + """:math:`\\mathrm{scale} \\cdot \\sigma^{-2} W^2 v`, in the row layout `apply` uses.""" + # W is applied twice: once to the residual, once from differentiating ||W r||^2. + return self._apply_weights(self._apply_weights(values)) * (scale / self.noise_variance) + + def _transpose(self, weighted: npt.NDArray[np.float64]) -> _SpecialVector: + """:math:`B^T` applied to a row-space vector, as a freshly built DOF vector. + + Built fresh on every call rather than reused: pyadjoint stores the returned vector by + reference on the tape (`BlockVariable.add_adj_output`/`add_hessian_output`), so handing + back the same buffer twice would let a later call silently overwrite a value pyadjoint + is still holding. + """ V = self.observation.function_space - out = _vector(V.dofmap.index_map, V.dofmap.bs, V, dtype=V.mesh.geometry.x.dtype) + out = _vector(V.dofmap.index_map, V.dofmap.bs, V, dtype=self.observation.dtype) out.array[:] = 0.0 - # W is applied twice: once to the residual, once from differentiating ||W r||^2. - weighted = self._apply_weights(self._apply_weights(residual)) * (scale / self.noise_variance) self.observation.apply_transpose(weighted, out=out) return out @@ -81,7 +93,7 @@ def recompute_component(self, inputs, block_variable, idx, prepared=None): def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): adj_input = 1.0 if adj_inputs[0] is None else float(adj_inputs[0]) - return self._transpose_action(self._residual(inputs[0]), adj_input) + return self._transpose(self._weighted_row(self._residual(inputs[0]), adj_input)) def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None): tlm_u = tlm_inputs[0] @@ -105,10 +117,11 @@ def evaluate_hessian_component( hessian_input = 0.0 if hessian_inputs[0] is None else float(hessian_inputs[0]) adj_input = 1.0 if adj_inputs[0] is None else float(adj_inputs[0]) - # Second-order seed, propagated through the first derivative ... - out = self._transpose_action(self._residual(inputs[0]), hessian_input) - # ... plus the curvature of J applied to the TLM direction. + # Second-order seed, propagated through the first derivative, plus the curvature of J + # applied to the TLM direction and summed before B^T is applied, so that B^T (a full + # communication round-trip) runs once per Hessian action instead of twice. + combined = self._weighted_row(self._residual(inputs[0]), hessian_input) tlm_u = block_variable.tlm_value if tlm_u is not None: - out.array[:] += self._transpose_action(self.observation.apply(tlm_u), adj_input).array[:] - return out + combined += self._weighted_row(self.observation.apply(tlm_u), adj_input) + return self._transpose(combined) diff --git a/src/dolfinx_adjoint/observation.py b/src/dolfinx_adjoint/observation.py index 947c722..528f3c0 100644 --- a/src/dolfinx_adjoint/observation.py +++ b/src/dolfinx_adjoint/observation.py @@ -97,7 +97,16 @@ def _default_padding(mesh: dolfinx.mesh.Mesh) -> float: def _pad_points(points: npt.ArrayLike) -> np.ndarray: - """Return points as a contiguous ``(num_points, 3)`` float64 array.""" + """Return points as a contiguous ``(num_points, 3)`` float64 array. + Add zero padding for points in 1D or 2D, so that the bounding-box tree can be + built once and used for all points. + + Args: + points: Input points, shape ``(num_points, dim)`` with ``dim <= 3``. + + Returns: + Padded points, shape ``(num_points, 3)``. + """ padded_input = np.atleast_2d(np.asarray(points, dtype=np.float64)) if padded_input.ndim != 2: raise ValueError(f"Points must be 2D with shape (num_points, dim), got shape {padded_input.shape}") @@ -169,13 +178,18 @@ def __init__( "`V.sub(0).collapse()[0]`, and build the operator on that." ) from exc + # Make sure all points are 3D. For dim < 3, pad with zeros so that the bounding-box + # tree can be built once and used for all points. padded_points = _pad_points(points) num_points = padded_points.shape[0] if comm.allreduce(num_points, op=MPI.MIN) != comm.allreduce(num_points, op=MPI.MAX): raise ValueError("`points` must be replicated on all processes (got differing lengths).") # A checksum costs one scalar reduction and catches the far nastier case of equally # many but *different* points, which would otherwise silently corrupt the operator. - checksum = float(padded_points.sum()) + # Weighting each entry by its position makes the checksum sensitive to permuted rows + # too, not just to changed coordinates. + position_weights = np.arange(1, padded_points.size + 1, dtype=np.float64) + checksum = float(np.dot(padded_points.ravel(), position_weights)) if comm.allreduce(checksum, op=MPI.MIN) != comm.allreduce(checksum, op=MPI.MAX): raise ValueError("`points` must be replicated on all processes (got differing coordinates).") self.num_points = num_points @@ -184,7 +198,7 @@ def __init__( tdim = mesh.topology.dim self.padding = _default_padding(mesh) if padding is None else float(padding) - # Locate the points with an exact containment test, over owned cells only. + # Locate the points in the mesh num_owned_cells = mesh.topology.index_map(tdim).size_local first_cell = np.full(num_points, -1, dtype=np.int32) if num_owned_cells > 0 and num_points > 0: @@ -203,14 +217,35 @@ def __init__( owner = np.empty_like(candidate_owner) comm.Allreduce(candidate_owner, owner, op=MPI.MIN) - self.found = owner < comm.size - self.owner = np.where(self.found, owner, -1).astype(np.int32) - self.num_found = int(self.found.sum()) + self._found = owner < comm.size + self._owner = np.where(self._found, owner, -1).astype(np.int32) + self.num_found = int(self._found.sum()) self.local_indices = np.flatnonzero(owner == comm.rank).astype(np.int32) self._build_matrix(padded_points[self.local_indices]) - # -------------------------------------------------------------------- assembly --- + def __repr__(self) -> str: + return type(self).__name__ + f"(V={self.function_space}, points={self.points}, padding={self.padding})" + + def __str__(self) -> str: + return f"PointObservation({self.num_found} points, {self.function_space})" + + @property + def found(self) -> npt.NDArray[np.bool_]: + """Boolean array of length ``num_points``, ``True`` where the point was located in the mesh on some rank. + + Identical on every rank. + """ + return self._found + + @property + def owner(self) -> npt.NDArray[np.int32]: + """Rank owning each point, ``-1`` where it was not found. + + Identical on every rank. + """ + return self._owner + def _build_matrix(self, local_points: np.ndarray) -> None: """Build the interpolation matrix from ``V`` onto a point mesh of the owned points. @@ -240,8 +275,8 @@ def _build_matrix(self, local_points: np.ndarray) -> None: # Reusable work vectors, so that apply/apply_transpose do not allocate per call. self._observation_function = dolfinx.fem.Function(self.observation_space) self._state_function = dolfinx.fem.Function(V) - index_map = self.observation_space.dofmap.index_map - self._num_local_rows = index_map.size_local * self.observation_space.dofmap.index_map_bs + dm = self.observation_space.dofmap + self._num_local_rows = dm.index_map.size_local * dm.index_map_bs # --------------------------------------------------------------------- actions --- @property @@ -259,6 +294,15 @@ def matrix(self): """The interpolation matrix :math:`B`, as a distributed ``PETSc.Mat``.""" return self._matrix + @property + def dtype(self) -> np.dtype: + """Scalar dtype of the observed function space. + + Independent of the mesh geometry's dtype, which is always real -- a complex-valued + ``V`` still has real geometry. + """ + return self._state_function.x.array.dtype + def apply(self, u: dolfinx.fem.Function) -> npt.NDArray[np.float64]: """Evaluate :math:`Bu` for the rows owned by this process. @@ -304,7 +348,7 @@ def apply_transpose(self, values: npt.ArrayLike, out: dolfinx.la.Vector | None = """ V = self.function_space if out is None: - out = dolfinx.la.vector(V.dofmap.index_map, V.dofmap.bs, dtype=V.mesh.geometry.x.dtype) + out = dolfinx.la.vector(V.dofmap.index_map, V.dofmap.bs, dtype=self.dtype) out.array[:] = 0.0 self._observation_function.x.array[: self.num_local_rows] = np.asarray(values, dtype=np.float64) @@ -322,13 +366,14 @@ def restrict(self, data: npt.ArrayLike) -> npt.NDArray[np.float64]: Args: data: Array of length ``num_points * block_size``, identical on every rank. + + Returns: + Array of length ``num_local_rows``, in the layout produced by :meth:`apply`. """ values = np.asarray(data, dtype=np.float64) bs = self.block_size if values.shape[0] != self.num_points * bs: raise ValueError(f"Expected data of length {self.num_points * bs}, got {values.shape[0]}") - if bs == 1: - return values[self.local_indices] return values.reshape(self.num_points, bs)[self.local_indices].reshape(-1) def gather(self, values: npt.ArrayLike, fill: float = np.nan) -> npt.NDArray[np.float64]: @@ -341,18 +386,11 @@ def gather(self, values: npt.ArrayLike, fill: float = np.nan) -> npt.NDArray[np. bs = self.block_size buffer = np.zeros(self.num_points * bs, dtype=np.float64) local = np.asarray(values, dtype=np.float64) - if bs == 1: - buffer[self.local_indices] = local - else: - buffer.reshape(self.num_points, bs)[self.local_indices] = local.reshape(-1, bs) + buffer.reshape(self.num_points, bs)[self.local_indices] = local.reshape(-1, bs) total = np.empty_like(buffer) self.comm.Allreduce(buffer, total, op=MPI.SUM) if not np.all(self.found): - missing = ~self.found - if bs == 1: - total[missing] = fill - else: - total.reshape(self.num_points, bs)[missing] = fill + total.reshape(self.num_points, bs)[~self.found] = fill return total diff --git a/tests/test_observation.py b/tests/test_observation.py index 15ccbe9..c7b7037 100644 --- a/tests/test_observation.py +++ b/tests/test_observation.py @@ -274,6 +274,37 @@ def test_points_must_match_across_processes(): dolfinx_adjoint.PointObservation(V, points) +def test_permuted_points_are_rejected(): + """Reordering the rows leaves the sum of coordinates unchanged, so a plain-sum checksum + would miss it; the checksum must weight by position to catch this too.""" + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + if comm.size == 1: + pytest.skip("mismatched points require more than one process") + + points = sample_points(6, seed=131) + if comm.rank == 1: + points = points[::-1].copy() + with pytest.raises(ValueError, match="differing coordinates"): + dolfinx_adjoint.PointObservation(V, points) + + +def test_bad_point_shape_on_one_process_does_not_deadlock(): + """A per-process malformed `points` array must raise on every rank, not just the bad one. + + A validation failure on only some ranks would leave those ranks raising while the others + proceed into the collective replication check below and block forever. + """ + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + if comm.size == 1: + pytest.skip("a per-process shape mismatch requires more than one process") + + points = np.zeros((2, 5)) if comm.rank == 1 else sample_points(5, seed=151) + with pytest.raises(ValueError, match="at most 3 components"): + dolfinx_adjoint.PointObservation(V, points) + + def test_padding_does_not_depend_on_the_partition(): """The default padding comes from the global bounding box, not each process's own.""" comm = MPI.COMM_WORLD @@ -499,6 +530,30 @@ def test_mismatched_data_length_raises(): dolfinx_adjoint.point_observation_misfit(u, B, np.zeros(B.num_points + 3)) +def test_reusing_a_data_buffer_does_not_corrupt_earlier_blocks(): + comm = MPI.COMM_WORLD + tape = pyadjoint.get_working_tape() + tape.clear_tape() + V = dolfinx.fem.functionspace(unit_square(comm, 4), ("Lagrange", 1)) + u = dolfinx_adjoint.Function(V, name="u") + u.interpolate(lambda x: x[0]) + + B = dolfinx_adjoint.PointObservation(V, sample_points(5, seed=17)) + buffer = B.restrict(np.zeros(B.num_points)) + + buffer[:] = 1.0 + dolfinx_adjoint.point_observation_misfit(u, B, buffer) + first_block = tape.get_blocks()[-1] + + buffer[:] = 2.0 # mutated after the block was built, as a caller reusing the array would + dolfinx_adjoint.point_observation_misfit(u, B, buffer) + second_block = tape.get_blocks()[-1] + + assert np.allclose(first_block.data, 1.0) + assert np.allclose(second_block.data, 2.0) + tape.clear_tape() + + # --------------------------------------------------------------------------- # Differentiability # --------------------------------------------------------------------------- From 9becd6a2a3128eeed28fe6f26caa6f278da49f17 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Wed, 19 Aug 2026 08:34:16 +0000 Subject: [PATCH 09/11] Fix docstring --- src/dolfinx_adjoint/observation.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/dolfinx_adjoint/observation.py b/src/dolfinx_adjoint/observation.py index 528f3c0..cf9e7db 100644 --- a/src/dolfinx_adjoint/observation.py +++ b/src/dolfinx_adjoint/observation.py @@ -131,10 +131,6 @@ class PointObservation: Attributes: num_points: Total (global) number of input points. - found: Boolean array of length ``num_points``, ``True`` where the point was located - in the mesh on some rank. Identical on every rank. - owner: Rank owning each point, ``-1`` where it was not found. Identical on every - rank. local_indices: Indices into the global point array of the points owned by this rank. This is the row ordering used by :meth:`apply`. From 4d4ada8947858f762d5a9672302e59e55d8b54a0 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Wed, 19 Aug 2026 08:41:40 +0000 Subject: [PATCH 10/11] Check padding on each rank before raising - if not we might get deadlocks --- src/dolfinx_adjoint/observation.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/src/dolfinx_adjoint/observation.py b/src/dolfinx_adjoint/observation.py index cf9e7db..95e15f7 100644 --- a/src/dolfinx_adjoint/observation.py +++ b/src/dolfinx_adjoint/observation.py @@ -117,6 +117,27 @@ def _pad_points(points: npt.ArrayLike) -> np.ndarray: return padded +def _pad_points_collective(comm: MPI.Comm, points: npt.ArrayLike) -> np.ndarray: + """``_pad_points``, but validated on every rank before any rank can raise. + + A bad ``points`` array is exactly the kind of per-rank data bug the shape checks in + ``_pad_points`` exist to catch, and it need not affect every rank alike. Raising locally, + before the replication check below has run, would let one rank exit the constructor while + the others block forever on that check's collective reduction. + """ + try: + padded = _pad_points(points) + message = "" + except ValueError as exc: + padded = np.zeros((0, 3), dtype=np.float64) + message = str(exc) + + failures = [message for message in comm.allgather(message) if message] + if failures: + raise ValueError(failures[0]) + return padded + + class PointObservation: """The operator :math:`B` evaluating a finite element function at a set of points. @@ -176,7 +197,7 @@ def __init__( # Make sure all points are 3D. For dim < 3, pad with zeros so that the bounding-box # tree can be built once and used for all points. - padded_points = _pad_points(points) + padded_points = _pad_points_collective(comm, points) num_points = padded_points.shape[0] if comm.allreduce(num_points, op=MPI.MIN) != comm.allreduce(num_points, op=MPI.MAX): raise ValueError("`points` must be replicated on all processes (got differing lengths).") From 138bc55ad0a57519dd77616c4e3173c2cdd9ef98 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Wed, 19 Aug 2026 08:52:14 +0000 Subject: [PATCH 11/11] Remove commented out line --- demos/point_observations.py | 1 - 1 file changed, 1 deletion(-) diff --git a/demos/point_observations.py b/demos/point_observations.py index 672a435..deb55b2 100644 --- a/demos/point_observations.py +++ b/demos/point_observations.py @@ -199,7 +199,6 @@ def true_source(x): print("Install pyvista to visualize the result") else: - # pyvista.set_jupyter_backend("html") cells, types, geometry = dolfinx.plot.vtk_mesh(V) sensor_cloud = pyvista.PolyData(np.column_stack([sensors, np.zeros(len(sensors))]))