42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
from mpi4py import MPI
|
|
import numpy as np
|
|
import math
|
|
import sys
|
|
|
|
def dot_product(a, x):
|
|
result = 0
|
|
for i in range(len(a)):
|
|
result += a[i] * x[i]
|
|
return result
|
|
|
|
def row(i):
|
|
row = []
|
|
for j in range(1, n+1):
|
|
row.append(i / j)
|
|
return row
|
|
|
|
comm = MPI.COMM_WORLD
|
|
rank = comm.Get_rank()
|
|
size = comm.Get_size()
|
|
|
|
n = int(sys.argv[1])
|
|
|
|
x = list(range(1, n+1))
|
|
|
|
# each rank should compute almost the same amount of matrix.row * x
|
|
# for this split the list [1,...,n+1] in sublists containing the rownumbers every rank has to compute
|
|
# if n is a multiple of size all ranks have the same amount to compute, if not, the first (n % size) ranks compute each one more
|
|
chunks = np.array_split(list(range(1,n+1)), size)
|
|
|
|
chunk = chunks[rank]
|
|
sub_b = []
|
|
for i in chunk:
|
|
sub_b.append(dot_product(row(i), x)) # every rank computes its delegated rows times x
|
|
|
|
comm.barrier()
|
|
|
|
b = comm.gather(sub_b)
|
|
|
|
if rank == 0:
|
|
b = np.concatenate(b).tolist() # b is a list of lists, np.concatenate 'flattens' the list of lists into an np.ndarray and tolist() to get an python list
|