mpi4py broadcasting numpy arrays between functions

user_17313150 0
Viewed 178

I'm a beginner with MPI and I'm trying to make a simple example program where the master sends a numpy array to the workers. The only catch is these have to work between distinct master and worker functions. This is a practice problem to figure out why a similar setup in a larger project is not currently working. All of the examples I've found for using Bcast do not include any tips for how to make it work between different modules or function calls.

Here is the code:

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.rank

def master():
    data = np.arange(4.0)
    comm.Bcast([data, MPI.DOUBLE], root=0)
    print('rank',rank,data)
    return()

def worker():
    data = np.zeros(4)
    data = comm.Bcast([data, MPI.DOUBLE], root=0)
    print('rank',rank,data)
    return()

if rank == 0:
    master()
else:
    worker()

Here is the current output:

rank 1 None
rank 2 None
rank 0 [0. 1. 2. 3.]
rank 3 None

I'm not sure why this version doesn't work. The almost identical version that used a list instead works just fine.

mpi4py array-broadcasting mpi python
1 Answers

I would like to thank commenters for responding so quickly. To make it more visible, this is what the worker function should look like:

def worker():
    data = np.zeros(4)
    #do not set data equal to Bcast call
    comm.Bcast([data, MPI.DOUBLE], root=0)
    print('rank',rank,data)
    return()

Thank you Gilles Gouaillardet for helping. This was a simple solution, and it even looks like my working version of the module for Fortran but I got caught up with recreating examples while trying to write the project in Python.

user_173131500
Related