Last active
August 30, 2023 14:39
-
-
Save ReubenHill/d4c44bbfb40f33cb5a5e0572613bba27 to your computer and use it in GitHub Desktop.
Generic external point-evaluation data input in Firedrake
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Note this requires firedrakeproject/firedrake#3067 to work