Scatter the global numpy array to local array

dolfinx
,
1

Hello All,

I am trying to write the restart code for my optimization problem using MPI, where I have outputs from the global mesh as numpy arrays for the last iteration, which I have to feed into the restart process but it needs to be scattered first to the locally owned processes.

I have simply turned the code to MWE as:

import numpy
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
from dolfinx.fem import (Expression, Function, FunctionSpace, dirichletbc,
                         form, functionspace, locate_dofs_topological, Constant, assemble_scalar, assemble)
from dolfinx.mesh import (CellType, GhostMode, create_box, create_rectangle,
                          locate_entities_boundary, meshtags)

def mpi_print(s):
    print(f"Rank {comm.rank}: {s}")

dtype = PETSc.ScalarType  # type: ignore
comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size

domain = create_rectangle(MPI.COMM_WORLD, [np.array([0, 0]), np.array([2, 2])], [2, 2],
                          CellType.triangle, ghost_mode=GhostMode.shared_facet)

num_elements = domain.topology.index_map(2).size_local
num_nodes = domain.topology.index_map(0).size_local
# Summing up the counts from all processes in MPI_COMM_WORLD
num_elements_global = MPI.COMM_WORLD.allreduce(num_elements, op=MPI.SUM)
num_nodes_global = MPI.COMM_WORLD.allreduce(num_nodes, op=MPI.SUM)

if (rank==0):
    mpi_print(f"Number of local elements: {num_elements}")
    mpi_print(f"Number of local nodes: {num_nodes}")
    mpi_print(f"Number of global elements: {num_elements_global}")
    mpi_print(f"Number of global nodes: {num_nodes_global}")

# Fespace
Vh0 = functionspace(domain, ("DG", 0))

numpy_array = numpy.arange(num_elements_global)  # Global input
uD = Function(Vh0)
vec = uD.vector.getArray()    # local output
print(len(numpy_array), numpy_array)
print(len(vec), vec)

# for finding the local range
E = Function(Vh0)
E.x.scatter_forward()
imap = E.function_space.dofmap.index_map
local_range = np.asarray(imap.local_range, dtype=np.int32) * \
              E.function_space.dofmap.index_map_bs
# Communicate ranges and local data
ranges = comm.gather(local_range, root=0)

if(rank==0):
    for r in ranges:
        vec = numpy_array[r[0]:r[1]]
print(vec)

For which the output is

Rank 0: Number of local elements: 4
Rank 0: Number of local nodes: 4
Rank 0: Number of global elements: 8
Rank 0: Number of global nodes: 9

8 [0 1 2 3 4 5 6 7]
4 [0. 0. 0. 0.]
[4 5 6 7]

8 [0 1 2 3 4 5 6 7]
4 [0. 0. 0. 0.]
[0. 0. 0. 0.]

The expected output for the vec is

[0 1 2 3]
[4 5 6 7]

So, in short my question is simple, is there a simple way to scatter the global numpy array to locally owned array?

This post was helpful from converting the local to global (gather), but not for global to local (scatter)
In parall computation(mpirun), how to communicate(scatter & gather) the data between local processes and some global process?

Thanks!

2

You may want to use GitHub - jorgensd/adios4dolfinx: Extending DOLFINx with checkpointing functionality rather than reinventing the wheel.

If you do want to reinvent the wheel, then I would question why is it that you have a global array in the first place. If your only goal is to restart from a previous iteration, why can’t you simply store the local part of the array? After all, it’s the only part that is needed to restart.

3

@francesco-ballarin,
The reason for storing the global arrays instead of local is:

  1. I have to store it seperately for individual processes and keep the track of the rank and the corresponding number of local elements belonging to that particular rank, which is impossible when I run on, say 64 procs.
  2. Even if I store the local arrays separately, there is no assurance the mesh partition is going to remain intact at restart of optimization process, unless I know to save the mesh partition. For example, when I ran for my test case which had 45,000 global elements. The initial partition was equal for 2 procs (22,500 elements in each), when I reran for restart process, the partition changed (23400 and 21600 elements) resulting in error while the importing the local stored data, as expected.
ValueError: could not broadcast input array from shape (22500,) into shape (23400,)
ValueError: could not broadcast input array from shape (22500,) into shape (21600,)

Feel free to correct me, if I am wrong.

Questions:

  1. Is there something simple as this can be done for scattering global array to local array as this interpolation
import numpy
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
from dolfinx.fem import (Expression, Function, FunctionSpace, dirichletbc,
                         form, functionspace, locate_dofs_topological, Constant, assemble_scalar, assemble)
from dolfinx.mesh import (CellType, GhostMode, create_box, create_rectangle,
                          locate_entities_boundary, meshtags)

dtype = PETSc.ScalarType  # type: ignore
comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size

# Create a box Mesh:
domain = create_rectangle(MPI.COMM_WORLD, [np.array([0, 0]), np.array([2, 2])], [2, 2],
                          CellType.triangle, ghost_mode=GhostMode.shared_facet)


# Assuming `mesh` is the mesh created in the code
num_elements = domain.topology.index_map(2).size_local
num_nodes = domain.topology.index_map(0).size_local
# Summing up the counts from all processes in MPI_COMM_WORLD
num_elements_global = MPI.COMM_WORLD.allreduce(num_elements, op=MPI.SUM)
num_nodes_global = MPI.COMM_WORLD.allreduce(num_nodes, op=MPI.SUM)

# Fespace
Vh0 = functionspace(domain, ("DG", 0))

uD = Function(Vh0)
uD.interpolate(lambda x: x[0])
vec = uD.vector.getArray()
print(len(uD.vector.array), uD.vector.array)

For 2 procs:
Results:

4 [0.33333333 0.66666667 0.33333333 1.33333333]
4 [0.66666667 1.33333333 1.66666667 1.66666667]

whereas, when I try this in addition to the above MWE

uD2 = Function(Vh0)
numpy_array = numpy.arange(num_elements_global)
uD2.vector.array = numpy_array
vec2 = uD2.vector.getArray()
print(len(uD2.vector.array), uD2.vector.array)

I get the error as the partition is done while creating the function

ValueError: could not broadcast input array from shape (8,) into shape (4,)
ValueError: could not broadcast input array from shape (8,) into shape (4,)

Seems like I have to create the local ranges and assign it locally. or is there simple process you suggest?

  1. I haven’t used the adios4dolfinx yet, is there any particular tutorial related to my issue, you can point?

Thanks!

4

There are several examples regarding how to use adios4dolfinx at
http://jsdokken.com/adios4dolfinx/
that shows you have to store a large variety of fields, including meshes, partitioning information, functions from arbitrary order spaces etc.

1 Like