Save the result of mpirun in a numpy array

Here is a simple example that gathers the data on the first process:

from dolfin import *
import numpy as np
mesh = UnitSquareMesh(2, 2)
V = FunctionSpace(mesh, "CG", 1)
u = project(Expression("x[0]", degree=1), V)
vertex_values = u.compute_vertex_values()
coordinates = mesh.coordinates()

from mpi4py import MPI
gathered_coordinates = MPI.COMM_WORLD.gather(coordinates, root=0)
gathered_values = MPI.COMM_WORLD.gather(vertex_values, root=0)
global_indices = MPI.COMM_WORLD.gather(mesh.topology().global_indices(0))
if MPI.COMM_WORLD.rank == 0:
    num_vertices = mesh.num_entities_global(0)
    all_coordinates = np.zeros((num_vertices, coordinates.shape[1]))
    all_values = np.zeros(num_vertices)
    for coord, value, indices in zip(gathered_coordinates, gathered_values, global_indices):
        all_values[indices] = value
        all_coordinates[indices] = coord

    for (coord, value) in zip(all_coordinates, all_values):
        print(coord, value)
4 Likes