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..deb55b2 --- /dev/null +++ b/demos/point_observations.py @@ -0,0 +1,220 @@ +# # 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: + 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() # type: ignore[call-arg] + if pyvista.OFF_SCREEN: + plotter.screenshot("point_observations.png") + else: + plotter.show() + +# - 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/__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..187903a --- /dev/null +++ b/src/dolfinx_adjoint/blocks/observation.py @@ -0,0 +1,127 @@ +"""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 + # 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 = 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)" + + 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 _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=self.observation.dtype) + out.array[:] = 0.0 + 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(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] + 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, 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: + 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 new file mode 100644 index 0000000..95e15f7 --- /dev/null +++ b/src/dolfinx_adjoint/observation.py @@ -0,0 +1,496 @@ +"""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 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 + +from .blocks.observation import PointObservationBlock + +__all__ = ["PointObservation", "point_observation_misfit"] + + +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. + """ + + 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 __str__(self) -> str: + return f"PointCloudTrace({self._mesh})" + + +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. + 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}") + 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 + + +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. + + 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. + 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: + 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 + + # 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_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).") + # 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. + # 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 + self.points = padded_points + + tdim = mesh.topology.dim + self.padding = _default_padding(mesh) if padding is None else float(padding) + + # 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: + 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]] + + # 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) + + 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]) + + 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. + + 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 + gdim = V.mesh.geometry.dim + bs = V.dofmap.bs + + 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) + dm = self.observation_space.dofmap + self._num_local_rows = dm.index_map.size_local * dm.index_map_bs + + # --------------------------------------------------------------------- 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 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 + + @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. + + Args: + u: Function in :attr:`function_space`. + + Returns: + The point values of the rows owned by this process, ordered as + :attr:`local_indices` (component-fastest for vector spaces). + """ + u.x.scatter_forward() + 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. + + 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=self.dtype) + out.array[:] = 0.0 + + 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 + + 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. + + 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]}") + 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) + 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): + total.reshape(self.num_points, bs)[~self.found] = fill + return total + + +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..c7b7037 --- /dev/null +++ b/tests/test_observation.py @@ -0,0 +1,764 @@ +"""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_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 + 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_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 + 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_matrix_is_distributed_over_the_points(): + """The operator is an interpolation matrix onto a point mesh of the observed points.""" + comm = MPI.COMM_WORLD + V = dolfinx.fem.functionspace(unit_square(comm, 8), ("Lagrange", 1)) + points = sample_points(31, seed=91) + B = dolfinx_adjoint.PointObservation(V, points) + + 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(): + """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_matrix_action_matches_apply(): + """`apply` is exactly a matrix-vector product with the interpolation matrix.""" + comm = MPI.COMM_WORLD + 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() + + 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) + + +# --------------------------------------------------------------------------- +# 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)) + + +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 +# --------------------------------------------------------------------------- + + +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()