Write a spatial sampler
A spatial sampler brings the model’s fields to where a measurement is taken: it interpolates
them onto a measurement mesh, evaluates them at points, integrates them over a surface, and so
on. This page walks through a small, complete one, fenicsx.points, which evaluates a field
at given points (kipe.fenicsx.points). A new sampler follows the same steps. Read
Measurements and Architecture first: they explain what a sampler is
for and how kipe uses it.
Start without kipe
Write the observation first as a plain function, outside kipe: a dolfinx.fem.Function in, a
numpy array out. Try it on the solver’s mesh, with a field whose values you know (e.g., a
linear function, which P1 represents exactly), on one rank and on several. Turning it into a
sampler is then a matter of the steps below.
def observe(u: dolfinx.fem.Function) -> numpy.ndarray:
"""The values of u where the measurement is taken."""
What a sampler must do
A sampler is a subclass of kipe.sampling.SpatialSampler:
The sampled vector has
sizeentries over all ranks. Each rank returns some of them, at the positionsglobal_indices, and each entry is returned by exactly one rank. The filter sums over the entries of all ranks, so an entry returned twice would count twice.The positions do not depend on the number of ranks. The data are stored in this order, so data written on one rank must be readable on eight.
fenicsx.meshuses the rows of the measurement mesh file,fenicsx.pointsthe order of the points in the study file.sample()gets the observed fields as numpy arrays: for a finite element field, this rank’s owned degrees of freedom. It returns this rank’s entries, in the order ofglobal_indices, as a new array.
Everything that does not change between calls (searching the points, building interpolation
data) belongs in the constructor: sample runs for every sigma point at every measurement
time.
1. The options
The options say what a study file can set for the sampler: a frozen dataclass in
kipe.options, with a literal type. The name carries the backend, fenicsx.. The
docstrings are the documentation of the study file, here and in the API pages:
@dataclass(frozen=True, config=_CONFIG)
class PointSamplingOptions:
"""Evaluate the observed field at given points (FEniCSx backend).
The field must be a distributed field whose function space the forward solver provides
(:class:`kipe.fenicsx.forward_solver.SupportsFunctionSpaces`). The data are in the order of
``points``; for a vector field, the components of each point one after another.
"""
type: Literal["fenicsx.points"]
"""Kind of spatial sampling."""
points: Annotated[list[list[float]], Field(min_length=1)]
"""Coordinates of the measurement points, one list per point, e.g., ``[[0.5, 0.5], [0.2,
0.8]]`` in 2D. Each point must lie in the solver's mesh."""
It joins the choices of spatial_sampling, the union that pydantic selects from by
type:
type SpatialSamplingOptions = Annotated[
ArraySamplingOptions | MeshSamplingOptions | PointSamplingOptions,
Field(discriminator="type"),
]
2. The sampler
The class gets everything it needs as arguments, not as options, so that it can be tested without a study file:
class PointSampler(SpatialSampler):
"""Evaluates a field at given points.
Each point is evaluated by one rank: the lowest rank that owns a cell containing it. A point
on the boundary between the cells of two ranks lies in cells of both.
Args:
field: name of the observed field
source: the field's function space, on the solver's mesh
points: coordinates of the points, one row per point
comm: communicator of the run
Raises:
StudyFileError: if the points do not have the coordinates of the mesh's dimension, or a
point lies outside the mesh
"""
def __init__(
self,
field: str,
source: fem.FunctionSpace,
points: list[list[float]],
comm: MPI.Comm = MPI.COMM_WORLD,
) -> None:
mesh = source.mesh
gdim = mesh.geometry.dim
if any(len(point) != gdim for point in points):
raise StudyFileError(f"fenicsx.points: each point needs {gdim} coordinates")
x = np.zeros((len(points), 3)) # dolfinx takes points in 3D
x[:, :gdim] = points
# the cell of this rank that contains each point, -1 if none; owned cells only, so
# that a cell is not found by its owner and by the ranks that have it as a ghost
tdim = mesh.topology.dim
n_owned = mesh.topology.index_map(tdim).size_local
tree = geometry.bb_tree(mesh, tdim, np.arange(n_owned, dtype=np.int32))
candidates = geometry.compute_collisions_points(tree, x)
colliding = geometry.compute_colliding_cells(mesh, candidates, x)
cells = np.full(len(x), -1, dtype=np.int32)
for i in range(len(x)):
if len(colliding.links(i)):
cells[i] = colliding.links(i)[0]
# each point to the lowest rank that found it; comm.size: no rank did
holder = np.where(cells >= 0, comm.rank, comm.size).astype(np.int32)
comm.Allreduce(MPI.IN_PLACE, holder, op=MPI.MIN)
outside = np.flatnonzero(holder == comm.size)
if len(outside):
raise StudyFileError(
f"fenicsx.points: {len(outside)} of {len(x)} points lie outside the solver's "
f"mesh, e.g., {list(points[outside[0]])}"
)
mine = np.flatnonzero(holder == comm.rank)
self._field = field
self._function = fem.Function(source)
self._n_owned = source.dofmap.index_map.size_local * source.dofmap.index_map_bs
self._points, self._cells = x[mine], cells[mine]
# the sampled vector: the values of the points, components of a point one after another
value_size = int(np.prod(source.value_shape))
self.size = len(x) * value_size
self.global_indices = (mine[:, np.newaxis] * value_size + np.arange(value_size)).ravel()
def sample(self, fields: Mapping[str, NDArray_f64]) -> NDArray_f64:
"""Evaluate the field at this rank's points.
Args:
fields: the observed state field, this rank's owned degrees of freedom
Returns:
the values at this rank's points (a new array), at the positions
``global_indices``
"""
self._function.x.array[: self._n_owned] = fields[self._field]
# the values of the ghost degrees of freedom, which cells at the partition boundary use
self._function.x.scatter_forward()
values: NDArray_f64 = self._function.eval(self._points, self._cells)
return values.ravel()
The constructor:
Which rank evaluates which point. A rank searches its owned cells for the points; searching its ghost cells too would find a cell on two ranks. A point on the boundary between the cells of two ranks is still found by both, so the lowest of them takes it: the
MINreduction over all ranks. A point that no rank found lies outside the mesh, which is an error. Without that check, it would get no value, or a wrong one, silently.The positions. The values of point i are at position i, or
i * value_size + cfor component c of a vector field. They depend only on the order of the points in the study file.
sample:
The field’s values come as this rank’s owned degrees of freedom. A cell at the partition boundary also uses degrees of freedom owned by another rank, the ghosts:
scatter_forwardfetches their values before the evaluation.
3. Building it from the options
A build function turns the options into the sampler. It takes what it needs from the solver:
here, the observed field and its function space, from
observed_function_space(), which every FEniCSx sampler can use:
def build_point_sampler(
options: PointSamplingOptions,
fields: list[str],
solver: ForwardSolver,
comm: MPI.Comm = MPI.COMM_WORLD,
) -> PointSampler:
"""Construct a :class:`PointSampler` from the ``spatial_sampling`` options of a measurement.
Args:
options: ``spatial_sampling`` options of the measurement
fields: names of the fields the measurement observes; exactly one
solver: the forward solver, which provides the field's function space
comm: communicator of the run
Returns:
the sampler
"""
field, source = observed_function_space("fenicsx.points", fields, solver)
return PointSampler(field, source, options.points, comm)
kipe calls it from kipe.sampling.build_sampler(), which has one case per kind of
sampler. A backend’s sampler is imported inside its case, so that kipe works without dolfinx
as long as no study uses it:
def build_sampler(
options: ArraySamplingOptions | MeshSamplingOptions | PointSamplingOptions,
fields: list[str],
solver: ForwardSolver,
comm: MPI.Comm = MPI.COMM_WORLD,
) -> SpatialSampler:
"""Construct the sampler given by the ``spatial_sampling`` options of a measurement.
Args:
options: ``spatial_sampling`` options of the measurement
fields: names of the fields the measurement observes
solver: the forward solver; backend samplers take what they need from it, e.g., the
function spaces
comm: communicator of the run
Returns:
the spatial sampler
Raises:
StudyFileError: if a backend sampler is used but its backend is not installed
"""
match options:
case ArraySamplingOptions():
return ArraySampler(fields, solver.state_spec, comm)
case MeshSamplingOptions():
try: # dolfinx only when a study uses it
from kipe.fenicsx.sampling import build_mesh_sampler
except ImportError as err:
raise StudyFileError(f"fenicsx.mesh needs dolfinx: {err}") from err
return build_mesh_sampler(options, fields, solver, comm)
case PointSamplingOptions():
try: # dolfinx only when a study uses it
from kipe.fenicsx.points import build_point_sampler
except ImportError as err:
raise StudyFileError(f"fenicsx.points needs dolfinx: {err}") from err
return build_point_sampler(options, fields, solver, comm)
case _:
assert_never(options)
4. Tests
The tests use a field whose values are known exactly, and run on one rank and on several:
def test_values_at_the_points():
V = Heat(n=12).function_space("u")
sampler = PointSampler("u", V, POINTS)
values = _gathered(sampler, sampler.sample({"u": _owned(V, _linear)}))
np.testing.assert_allclose(values, _linear(np.array(POINTS).T), rtol=0, atol=1e-12)
def test_each_point_on_exactly_one_rank():
sampler = PointSampler("u", Heat(n=12).function_space("u"), POINTS)
indices = np.concatenate(COMM.allgather(sampler.global_indices))
np.testing.assert_array_equal(np.sort(indices), np.arange(len(POINTS)))
_gathered collects the entries of all ranks by their global_indices, so the result is
the same on any number of ranks. The test file also checks a vector field, the path from the
study’s options through build_sampler, and the errors
(tests/fenicsx/test_point_sampler.py):
python -m pytest tests/fenicsx/test_point_sampler.py
mpirun -n 2 python -m pytest tests/fenicsx/test_point_sampler.py
mpirun -n 3 python -m pytest tests/fenicsx/test_point_sampler.py
With more ranks, some ranks hold few or none of the points: a case worth testing too.
5. Use it
In a study file:
measurements:
u:
fields: [u]
times: {start: 0.05, stop: 0.5, step: 0.05}
data: {type: numpy, path: data/u.npz}
noise: {stddev: 0.001, seed: 0}
spatial_sampling:
type: fenicsx.points
points: [[0.25, 0.25], [0.5, 0.5], [0.75, 0.75]]
Before a pull request:
ruff check,ruff format --checkandmypypass (pre-commit run --all-files);the tests pass on one rank and on several;
the new options appear in Measurements, with an example.