Mpi4py: printing and plotting during execution

Question

I recently start working with mpi in order to use it to accelerate some code (a sph/gravity simulation).

So far it seem to be working. The position of my particle have change between the beginning and the end of the program, task manager show that several python thread are working...

But i have two problem:

1/ i can't print text during the execution of the program. The text is only print once it's finished

2/ I'm unable to create a graph using matplotlib.

In both case, neither python nor mpi return any error. My guess, for the text at least, is that its print when the execution end with mpiexec.

thanks for any insight !

here the code i run it with !mpiexec -n 8 python mpi4py_test.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from mpi4py import MPI

class particule:

    def __init__(self, h, pos, vel = [0, 0], m = 1, P = 0, density = 0, acc = [0, 0]):
        
        self.x, self.y = pos
        self.u, self.v = vel
        self.acc_x, self.acc_y = acc
        self.m = m
        self.h = h # kernel lenght
        self.P = P # Pressure
        self.density  = density
    
def kernel(part_a, part_b ):
    " monaghan cubic spline "
    cst = 1 / (np.pi * part_a.h**3) 
    dist = np.sqrt((part_b.x-part_a.x)**2 + (part_b.y-part_a.y)**2 )
    r = dist / h
    tmp = 0
    if r < 1:
        tmp = cst * (1 - 3/2* r**2 + 3/4*r**3)
    elif r < 2:
        tmp = cst * (1/4*(2-r)**3)
    
    return tmp
        

def grad_kernel(part_a, part_b):
    cst = 1 / (np.pi * part_a.h**3) 
    dist = np.sqrt((part_b.x-part_a.x)**2 + (part_b.y-part_a.y)**2 )
    r = dist / part_a.h
    tmp = 0
    if r < 1:
        tmp = cst * (9/4 * r**2 - 3*r)
    elif r < 2:
        tmp = cst * (-3/4*(2-r)**2)
        
    return tmp
 
class hash_grid:
    
    def __init__(self, cell_size):
        self.cell_size = cell_size
        self.cell = {}
        self.part_list = []
        
    def key(self, part):
        return (part.x//self.cell_size,
                part.y//self.cell_size)
    
    def add(self, part):        
        tmp = self.key(part)
        self.cell.setdefault(tmp, []).append(part)
        self.part_list.append(part)

    def neighbours(self, part):
        idx = self.key(part)
        cell_N = None
        cell_S = None
        cell_E = None
        cell_W = None        
        cell_NE = None
        cell_NW = None
        cell_SE = None
        cell_SW = None
        if (idx[0], idx[1]+1) in self.cell: cell_N = (idx[0], idx[1]+1)
        if (idx[0], idx[1]-1) in self.cell: cell_S = (idx[0], idx[1]-1)
        if (idx[0]+1, idx[1]) in self.cell: cell_E = (idx[0]+1, idx[1])
        if (idx[0]-1, idx[1]) in self.cell: cell_W = (idx[0]-1, idx[1])
       
        if (idx[0]+1, idx[1]+1) in self.cell: cell_NE = (idx[0]+1, idx[1]+1)
        if (idx[0]-1, idx[1]+1) in self.cell: cell_NW = (idx[0]-1, idx[1]+1)
        if (idx[0]+1, idx[1]-1) in self.cell: cell_SE = (idx[0]+1, idx[1]-1)
        if (idx[0]-1, idx[1]-1) in self.cell: cell_SW = (idx[0]-1, idx[1]-1)

        return [value for cel in (idx, cell_N, cell_S, cell_E, cell_W, cell_NE, cell_NW, cell_SE, cell_SW) if cel!=None for value in self.cell.get(cel) ]
    
def split(to_split, nb_chunk):
    "take a list and split it most evenly possible"
    result = []
    q, r = divmod(len(to_split), nb_chunk) 
    curr = 0
    last = 0
    for i in range(nb_chunk):
        if r>0:
            last  = curr + q + 1
            result.append(to_split[curr: last])
            r = r-1
            curr = last
        else:
            last  = curr + q 
            result.append(to_split[curr: last])
            curr = last
    
    return result   

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()   

if rank == 0:
    n_part = 2000
    h = 2
    grid = hash_grid(h)
    points = np.zeros((n_part, 2))
    sim_range = 20
    dt = 0.005
    n_iter = 2500
    k = 5

    for _i in range(n_part):
        pos = np.random.uniform(-sim_range, sim_range, 2)
        vel = np.random.uniform(-.2, .2, 2)
        p = particule(h, pos, vel)
        grid.add(p)
        points[_i, :] = pos

    for part in grid.part_list:
        part.density = 0
        part.P = 0
        for p in grid.neighbours(part):
                part.density =  part.density + p.m * kernel(part, p)
        part.P = k * ( part.density )**2 # - density_0)
    data = split(grid.part_list,8)

for t in range(10):

    if rank == 0 :
        "1 - verlet 1/2 ; serial"
        for i, part in enumerate(grid.part_list):          
            part.u, part.v = part.u + part.acc_x * dt/2, part.v + part.acc_y * dt/2
            part.x, part.y = part.x + part.u * dt, part.y + part.v * dt
            
        "2 - update grid ; serial"
        jnk = grid.part_list
        del grid
        grid=hash_grid(h)        
        for p in jnk:
            grid.add(p)
                    
        grid_bcast = grid
        data_b = grid
        chunk_of_part_list = split(grid.part_list,8)
    else:
        grid_bcast=None
        chunk_of_part_list = None
        data_b=None
    
    grid_bcast = comm.bcast(grid_bcast, root=0)
    chunk_of_part_list = comm.scatter(chunk_of_part_list, root=0)
    
    "3 - get acc ; parallel"
    
    for part in chunk_of_part_list: 
        part.acc_x = 0
        part.acc_y = 0
        for p in grid_bcast.neighbours(part):
            if p != part:
                r = np.sqrt((p.x-part.x)**2 + (p.y-part.y)**2)
                if r==0: pass
                else:   
                    part.acc_x = part.acc_x - p.m * (part.P/part.density**2 + p.P/p.density**2) * grad_kernel(part, p) * (p.x - part.x)/r
                    part.acc_y = part.acc_y - p.m * (part.P/part.density**2 + p.P/p.density**2) * grad_kernel(part, p) * (p.y - part.y)/r
        
        dist = np.sqrt(part.x**2+part.y**2) 
        part.acc_x = part.acc_x - .5 * part.x -1*part.u
        part.acc_y = part.acc_y - .5 * part.y -1*part.v
    
    chunk_of_part_list = comm.gather(chunk_of_part_list,root=0)
    
    "4 - verlet 2/2 ; serial"
    if rank == 0:
        grid.part_list = list(matplotlib.cbook.flatten(chunk_of_part_list))
        for i, part in enumerate(grid.part_list): 
            part.u, part.v = part.u + part.acc_x * dt/2, part.v + part.acc_y * dt/2
            points[i,0], points[i,1] = part.x, part.y # for the figure
    
    if rank==0:
        print(t, flush=True)
    
if rank == 0:
    print('point 0', points[0,:])
    fig = plt.figure()
    ax = fig.add_subplot()
    sc = ax.scatter(points[:, 0], points[:, 1],s=3)
    ax.set_aspect('equal', 'box')
    ax.set_xlim(-sim_range,sim_range)
    ax.set_ylim(-sim_range,sim_range)