diff --git a/demos/point_observations.py b/demos/point_observations.py index b1ca324..12992d1 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: 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) # ---------------------------------------------------------------------------