Generic external point-evaluation data input in Firedrake

firedrake_data_input.py

from firedrake import *
from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD

# Here are some data points
if comm.rank == 0:
    data_coords = np.array(
        [
            [0.0, 0.0],
            [0.0, 0.5],
            [0.0, 1.0],
            [0.5, 0.0],
            [0.5, 0.5],
            [0.5, 1.0],
            [1.0, 0.0],
            [1.0, 0.5],
            [1.0, 1.0],
        ]
    )
    data_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, -1.0, -2.0, -3.0])
else:
    data_coords = np.array([]).reshape(0, 2)
    data_vals = np.array([])

# Here is a mesh which has vertices at the data coordinates. This is a
# contrived example, but a utility that produces a triangulation given a list
# of vertices is not too disimilar to mesh.plex_from_cell_list
m = UnitSquareMesh(2, 2, comm=comm)

if m.comm.size == 1:
    assert len(m.coordinates.dat.data_ro) == len(data_coords)  # sanity check

# Input the data using a vertex-only mesh, using the input_ordering property
# to ensure we respect the mesh parallel decomposition
vm = VertexOnlyMesh(m, data_coords, redundant=False)
P0DG = FunctionSpace(vm, "DG", 0)
P0DG_input_ordering = FunctionSpace(vm.input_ordering, "DG", 0)
f_point_data_input_ordering = Function(P0DG_input_ordering)
f_point_data_input_ordering.dat.data_wo[:] = data_vals  # This is now safe to do
f_point_data = interpolate(f_point_data_input_ordering, P0DG)

# I now have a function on the vertex-only mesh, but I need to work out how to
# assign the values back to the original mesh. This mostly means on-rank
# coordinate reordering, but because of the vertex-only mesh voting algorithm,
# which might decide that points are in cells on the other side of a parallel
# boundary than our mesh vertices, we also have some cross-rank reordering. We
# know that this is a permution operation, since we have the same number of
# data points as mesh vertices. The trick is to find this permutation:

# We firstly create a function space on m which has point evaluation nodes at
# the mesh vertices (i.e. the data locations) - this is a P1CG function space
V = FunctionSpace(m, "CG", 1)
# The interpolation from this space onto P0DG on the vertex only mesh will
# perform the permutation from the P1CG nodes on m to the P0DG nodes on vm.
I = Interpolator(TestFunction(V), P0DG)
# But we want the permution from the P0DG nodes on vm to the P1CG nodes on m,
# i.e. the inverse of I, which, since permutation matrices are orthogonal, is
# just its transpose.
f = I.interpolate(f_point_data, transpose=True)
# NOTE: the above is a bit naughty since f_point_data ought to be a cofunction
# but is actually a function. It works since P0DG on a vertex-only mesh is
# 'self-dual', i.e. cofunction basis coefficients are the transpose-conjugate
# of function basis coefficients. Since Firedrake only supports real-valued
# basis functions, so there's no need to worry about complex conjugation. In
# some hypothetical complex mode, I would need to (a) create a cofunction from
# f_point_data, (b) supply that to the transpose/adjoint interpolator, then (c)
# convert that back to a primal function. This is the same as what
# _technically_ I ought to do in real mode.

# Now we can check that the data are in the right place using .at. (Using a
# vertex-only mesh again would feel like marking our own homework!)


# NOTE: we need to gather all the data from all the ranks, since .at requires
# the same list of coordinates on all ranks
def allgather(comm, array_togather):
    """Gather all coordinates from all ranks."""
    array_togather = array_togather.copy()
    array_togather = comm.allgather(array_togather)
    array_togather = np.concatenate(array_togather)
    return array_togather


data_coords_all_ranks = allgather(comm, data_coords)
data_vals_all_ranks = allgather(comm, data_vals)
assert np.allclose(np.asarray(f.at(data_coords_all_ranks)), data_vals_all_ranks)

# We can now do whatever we want with f, e.g. interpolate it to a CG2 function
# on a different mesh with quadrilateral cells
m2 = UnitSquareMesh(3, 3)
V2 = FunctionSpace(m2, "CG", 2)
f2 = interpolate(f, V2)

# check our data again
assert np.allclose(np.asarray(f2.at(data_coords_all_ranks)), data_vals_all_ranks)

firedrake_data_input_2_ranks.py

from firedrake import *
from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD

# Here are some data points
if comm.rank == 0:
    data_coords = np.array(
        [
            [0.0, 0.0],
            [0.0, 0.5],
            [0.0, 1.0],
            [0.5, 0.0],
            [0.5, 0.5],
            [0.5, 1.0],
            [1.0, 0.0],
            [1.0, 0.5],
        ]
    )
    data_vals = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, -1.0, -2.0])
elif comm.rank == 1:
    data_coords = np.array(
        [
            [1.0, 1.0],
        ]
    )
    data_vals = np.array([-3.0])
else:
    data_coords = np.array([]).reshape(0, 2)
    data_vals = np.array([])

# Here is a mesh which has vertices at the data coordinates. This is a
# contrived example, but a utility that produces a triangulation given a list
# of vertices is not too disimilar to mesh.plex_from_cell_list
m = UnitSquareMesh(2, 2, comm=comm)

if comm.size < 2:
    raise RuntimeError("This example requires at least two ranks")

# Input the data using a vertex-only mesh, using the input_ordering property
# to ensure we respect the mesh parallel decomposition
vm = VertexOnlyMesh(m, data_coords, redundant=False)
P0DG = FunctionSpace(vm, "DG", 0)
P0DG_input_ordering = FunctionSpace(vm.input_ordering, "DG", 0)
f_point_data_input_ordering = Function(P0DG_input_ordering)
f_point_data_input_ordering.dat.data_wo[:] = data_vals  # This is now safe to do
f_point_data = interpolate(f_point_data_input_ordering, P0DG)

# I now have a function on the vertex-only mesh, but I need to work out how to
# assign the values back to the original mesh. This mostly means on-rank
# coordinate reordering, but because of the vertex-only mesh voting algorithm,
# which might decide that points are in cells on the other side of a parallel
# boundary than our mesh vertices, we also have some cross-rank reordering. We
# know that this is a permution operation, since we have the same number of
# data points as mesh vertices. The trick is to find this permutation:

# We firstly create a function space on m which has point evaluation nodes at
# the mesh vertices (i.e. the data locations) - this is a P1CG function space
V = FunctionSpace(m, "CG", 1)
# The interpolation from this space onto P0DG on the vertex only mesh will
# perform the permutation from the P1CG nodes on m to the P0DG nodes on vm.
I = Interpolator(TestFunction(V), P0DG)
# But we want the permution from the P0DG nodes on vm to the P1CG nodes on m,
# i.e. the inverse of I, which, since permutation matrices are orthogonal, is
# just its transpose.
f = I.interpolate(f_point_data, transpose=True)
# NOTE: the above is a bit naughty since f_point_data ought to be a cofunction
# but is actually a function. It works since P0DG on a vertex-only mesh is
# 'self-dual', i.e. cofunction basis coefficients are the transpose-conjugate
# of function basis coefficients. Since Firedrake only supports real-valued
# basis functions, so there's no need to worry about complex conjugation. In
# some hypothetical complex mode, I would need to (a) create a cofunction from
# f_point_data, (b) supply that to the transpose/adjoint interpolator, then (c)
# convert that back to a primal function. This is the same as what
# _technically_ I ought to do in real mode.

# Now we can check that the data are in the right place using .at. (Using a
# vertex-only mesh again would feel like marking our own homework!)


# NOTE: we need to gather all the data from all the ranks, since .at requires
# the same list of coordinates on all ranks
def allgather(comm, array_togather):
    """Gather all coordinates from all ranks."""
    array_togather = array_togather.copy()
    array_togather = comm.allgather(array_togather)
    array_togather = np.concatenate(array_togather)
    return array_togather


data_coords_all_ranks = allgather(comm, data_coords)
data_vals_all_ranks = allgather(comm, data_vals)
assert np.allclose(np.asarray(f.at(data_coords_all_ranks)), data_vals_all_ranks)

# We can now do whatever we want with f, e.g. interpolate it to a CG2 function
# on a different mesh with quadrilateral cells
m2 = UnitSquareMesh(3, 3)
V2 = FunctionSpace(m2, "CG", 2)
f2 = interpolate(f, V2)

# check our data again
assert np.allclose(np.asarray(f2.at(data_coords_all_ranks)), data_vals_all_ranks)

Note this requires firedrakeproject/firedrake#3067 to work