1
0
pwr_project/src/main_matrix_vec_timeit.py

25 lines
515 B
Python
Raw Permalink Normal View History

2024-07-27 13:17:21 +00:00
from mpi4py import MPI
import sys
import timeit
from matrix_mpi import MatrixMPI as Matrix
from vector_mpi import VectorMPI as Vector
# from matrix import Matrix
# from vector import Vector
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
n = int(sys.argv[1])
m_data = [(i / k) for i in range(1, n+1) for k in range(1, n+1)]
A = Matrix(m_data, (n, n))
v = Vector(list(range(1, n+1)))
time = timeit.timeit(lambda: A * v, number=1)
if rank == 0:
print(f"ranks = {size}: time = {time}s")