I am using mpi4py to parallise my code in python. comm.Gather gathers numpy arrays of all cores to one array automatically. I have used this and it works, but when I try to gather arrays of different shapes (which could be easily concatenated) the code fails. My code:
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
a = np.zeros((2 if rank==1 else 5, 3),dtype=float)+rank #size=3: shape[0] is 5,2,5 for ranks 0,1,2
#a = np.zeros((6,3),dtype=float)+rank #size=2: shape[0] is 6,6 for ranks 0,1 (this works)
print(rank,a)
b = np.zeros((12,3),dtype=float)-1
comm.Gather(a,b,root=0)
if rank==0:
print(b)
My question now is how to write the code to concatenate the arrays of different shapes so that it works. What I don't want is to split up the arrays so that every array has the same shape and then comm.send the rest manually and concatenate manually.