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
|
from mpi4py import MPI
try:
import numpy
except ImportError:
raise SystemExit from None
# send a upper triangular matrix
N = 10
a = numpy.empty((N, N), dtype=float, order="c")
b = numpy.zeros((N, N), dtype=float, order="c")
a.flat = numpy.arange(a.size, dtype=float)
# compute start and size of each row
i = numpy.arange(N)
blocklen = N - i
disp = N * i + i
# create datatype for upper triangular part
upper = MPI.DOUBLE.Create_indexed(blocklen, disp)
upper.Commit()
# send and recv matrix
myrank = MPI.COMM_WORLD.Get_rank()
MPI.COMM_WORLD.Sendrecv((a, 1, upper), myrank, 0, [b, 1, upper], myrank, 0)
assert numpy.allclose(numpy.triu(b), numpy.triu(a))
assert numpy.allclose(numpy.tril(b, -1), numpy.zeros((N, N)))
upper.Free()
|
