mpitest.py

# Based on https://gist.github.com/wd15/871797

import time
import fipy as fp
import numpy as np
from mpi4py import MPI

np.random.seed(1)
L = 1.
N = 30
m = fp.GmshGrid3D(nx=N, ny=N, nz=N, dx=L / N, dy=L / N, dz=L / N)
v = fp.CellVariable(mesh=m)
v0 = np.random.random(m.numberOfCells)
v[:] = v0.copy()
eqn = fp.TransientTerm() == fp.DiffusionTerm()
eqn.solve(v, dt=1.)
v[:] = v0.copy()
t0 = time.time()
eqn.solve(v, dt=1.)
t0 = time.time() - t0

#for rank in range(MPI.COMM_WORLD.Get_size()):
MPI.COMM_WORLD.Barrier()
if MPI.COMM_WORLD.Get_rank()==0:
  print r'  %d  %f' % (m.numberOfCells, t0)