In [1]:
import ipyparallel as ipp
c = ipp.Client(profile='mpi')
print(c.ids)
view = c[:]
view.activate()


[0, 1, 2, 3, 4, 5]

In [2]:
%%px --block 
from mpi4py import MPI
try:
    import numpy
except ImportError:
    raise SystemExit

# --------------------------------------------------------------------

n = 5 * MPI.COMM_WORLD.Get_size()

# compute number of processes and myrank
p = MPI.COMM_WORLD.Get_size()
myrank = MPI.COMM_WORLD.Get_rank()

# compute size of local block
m = n//p
if myrank < (n - p * m):
    m = m + 1

#compute neighbors
if myrank == 0:
    left = MPI.PROC_NULL
else:
    left = myrank - 1
if myrank == p - 1:
    right = MPI.PROC_NULL
else:
    right = myrank + 1

# allocate local arrays
A = numpy.empty((n+2, m+2), dtype='d', order='F')
B = numpy.empty((n, m),     dtype='d', order='F')

A.fill(1)
A[0, :] = A[-1, :] = 0
A[:, 0] = A[:, -1] = 0

# main loop
converged = False
#while not converged:
for i in range(10):
    # compute,  B = 0.25 * ( N + S + E + W)
    N, S = A[:-2, 1:-1], A[2:, 1:-1]
    E, W = A[1:-1, :-2], A[1:-1, 2:]
    numpy.add(N, S, B)
    numpy.add(E, B, B)
    numpy.add(W, B, B)
    B *= 0.25
    A[1:-1, 1:-1] = B
    # communicate
    tag = 0
    MPI.COMM_WORLD.Sendrecv([B[:, -1], MPI.DOUBLE], right, tag,
                            [A[:,  0], MPI.DOUBLE], left,  tag)
    MPI.COMM_WORLD.Sendrecv((B[:,  0], MPI.DOUBLE), left,  tag,
                            (A[:, -1], MPI.DOUBLE), right, tag)
    # convergence
    myconv = numpy.allclose(B, 0)
    loc_conv = numpy.asarray(myconv, dtype='i')
    glb_conv = numpy.asarray(0, dtype='i')
    MPI.COMM_WORLD.Allreduce([loc_conv, MPI.INT],
                             [glb_conv, MPI.INT],
                             op=MPI.LAND)
    converged = bool(glb_conv)
print("OK")


[stdout:0] OK
[stdout:1] OK
[stdout:2] OK
[stdout:3] OK
[stdout:4] OK
[stdout:5] OK

In [6]:
view['A'][3]


Out[6]:
array([[ 0.4308033 ,  0.        ,  0.        ,  0.        ,  0.        ,
         0.        ,  0.4308033 ],
       [ 0.66515636,  0.4308033 ,  0.34909916,  0.33295631,  0.34909916,
         0.4308033 ,  0.66515636],
       [ 0.79619503,  0.66515636,  0.60678768,  0.58583832,  0.60678768,
         0.66515636,  0.79619503],
       [ 0.8567028 ,  0.79619503,  0.76002121,  0.75399208,  0.76002121,
         0.79619503,  0.8567028 ],
       [ 0.88446808,  0.8567028 ,  0.84567165,  0.84032059,  0.84567165,
         0.8567028 ,  0.88446808],
       [ 0.89346218,  0.88446808,  0.88132858,  0.88274765,  0.88132858,
         0.88446808,  0.89346218],
       [ 0.89640236,  0.89346218,  0.89592743,  0.8973608 ,  0.89592743,
         0.89346218,  0.89640236],
       [ 0.89698792,  0.89640236,  0.89987755,  0.9023056 ,  0.89987755,
         0.89640236,  0.89698792],
       [ 0.89710331,  0.89698792,  0.90092373,  0.90334511,  0.90092373,
         0.89698792,  0.89710331],
       [ 0.8971138 ,  0.89710331,  0.90108109,  0.90354538,  0.90108109,
         0.89710331,  0.8971138 ],
       [ 0.89711475,  0.8971138 ,  0.90110111,  0.90356541,  0.90110111,
         0.8971138 ,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89711475,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89711475],
       [ 0.89662647,  0.89711475,  0.90110207,  0.90356636,  0.90110207,
         0.89711475,  0.89662647],
       [ 0.89516258,  0.89662647,  0.90110111,  0.90356541,  0.90110111,
         0.89662647,  0.89516258],
       [ 0.88882637,  0.89516258,  0.90060234,  0.90354538,  0.90060234,
         0.89516258,  0.88882637],
       [ 0.87429333,  0.88882637,  0.89904308,  0.90245628,  0.89904308,
         0.88882637,  0.87429333],
       [ 0.83873463,  0.87429333,  0.89187813,  0.89891052,  0.89187813,
         0.87429333,  0.83873463],
       [ 0.77518272,  0.83873463,  0.87468719,  0.88465405,  0.87468719,
         0.83873463,  0.77518272],
       [ 0.66700363,  0.77518272,  0.83161068,  0.85178757,  0.83161068,
         0.77518272,  0.66700363],
       [ 0.51828098,  0.66700363,  0.75326061,  0.77874184,  0.75326061,
         0.66700363,  0.51828098],
       [ 0.33533573,  0.51828098,  0.62118053,  0.65669823,  0.62118053,
         0.51828098,  0.33533573],
       [ 0.15079021,  0.33533573,  0.44155121,  0.4740448 ,  0.44155121,
         0.33533573,  0.15079021],
       [ 0.        ,  0.15079021,  0.22486877,  0.24841309,  0.22486877,
         0.15079021,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ,
         0.        ,  0.        ]])

In [5]:
view['glb_conv'][3]


Out[5]:
array(0, dtype=int32)

In [ ]: