1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
|
#!/usr/bin/env python
#
# Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
# Copyright (c) 1997-2016 California Institute of Technology.
# Copyright (c) 2016-2025 The Uncertainty Quantification Foundation.
# License: 3-clause BSD. The full license text is available at:
# - https://github.com/uqfoundation/pathos/blob/master/LICENSE
"""example: parallelism in python with mp4py
Run with:
> $ mpirun -np 2 scatter_gather.py
"""
import numpy as np
from mpi4py import MPI
comm = MPI.COMM_WORLD
nodes = comm.size #2
my_N = 3
N = my_N * nodes
# a print function that prints only to rank 0
def pprint(str="", end="\n", comm=comm):
"""Print for MPI parallel programs: Only rank 0 prints *str*."""
if comm.rank == 0:
print(str, end=end)
# set up the target arrays
if comm.rank == 0:
x = np.arange(N, dtype=np.float64)
else:
x = np.empty(N, dtype=np.float64)
my_x = np.empty(my_N, dtype=np.float64)
# scatter data into arrays on each node
comm.Scatter( [x, MPI.DOUBLE], [my_x, MPI.DOUBLE] )
# print the input to screen
pprint("Input:")
for r in range(nodes):
if comm.rank == r:
print(" [node %d] %s" % (comm.rank, my_x))
comm.Barrier()
# take the sin squared of all data
pprint("Running on %d cores..." % nodes)
my_x = np.sin(my_x)**2
# gather data into the head node
comm.Gather( [my_x, MPI.DOUBLE], [x, MPI.DOUBLE] )
# print the ouput to screen
pprint("Output:\n %s" % x)
# EOF
|
