1#!/usr/bin/env python
2
3# If you want VampirTrace to log MPI calls, you have to add the two
4# lines below at the very beginning of your main bootstrap script.
5import mpi4py
6mpi4py.rc.threads = False
7mpi4py.profile('vt', logfile='cpilog')
8
9# Import the MPI extension module
10from mpi4py import MPI
11
12# Import the 'array' module
13from array import array
14
15# This is just to make the logging
16# output a bit more interesting
17from time import sleep
18
19comm = MPI.COMM_WORLD
20nprocs = comm.Get_size()
21myrank = comm.Get_rank()
22
23n  = array('i', [0])
24pi = array('d', [0])
25mypi = array('d', [0])
26
27def comp_pi(n, myrank=0, nprocs=1):
28    h = 1.0 / n;
29    s = 0.0;
30    for i in range(myrank + 1, n + 1, nprocs):
31        x = h * (i - 0.5);
32        s += 4.0 / (1.0 + x**2);
33    return s * h
34
35comm.Barrier()
36
37for N in [10000]*10:
38
39    if myrank == 0:
40        n[0] = N
41
42    comm.Bcast([n, MPI.INT], root=0)
43
44    mypi[0] = comp_pi(n[0], myrank, nprocs)
45
46    comm.Reduce([mypi, MPI.DOUBLE],
47                [pi, MPI.DOUBLE],
48                op=MPI.SUM, root=0)
49
50    comm.Barrier()
51
52    sleep(0.01)
53