File: run_mandelbrot.py

package info (click to toggle)
mpi4py 3.0.3-8
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 12,428 kB
  • sloc: python: 18,672; javascript: 9,118; ansic: 7,092; makefile: 567; sh: 183; f90: 158; cpp: 103
file content (69 lines) | stat: -rw-r--r-- 1,414 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
from __future__ import print_function
from __future__ import division
import sys
import time

from mpi4py.futures import MPICommExecutor

try:
    range = xrange
except NameError:
    pass

x0 = -2.0
x1 = +1.0
y0 = -1.0
y1 = +1.0

w = 750
h = 500

dx = (x1 - x0) / w
dy = (y1 - y0) / h

def mandelbrot(x, y, maxit=255):
    c = complex(x, y)
    z = complex(0, 0)
    n = 255
    while abs(z) < 2 and n > 1:
        z = z**2 + c
        n -= 1
    return n

def mandelbrot_line(k):
    line = bytearray(w)
    y = y1 - k * dy
    for j in range(w):
        x = x0 + j * dx
        line[j] = mandelbrot(x, y)
    return line

def plot(image):
    import warnings
    warnings.simplefilter('ignore', UserWarning)
    try:
        from matplotlib import pyplot as plt
    except ImportError:
        return
    plt.figure()
    plt.imshow(image, aspect='equal', cmap='spectral')
    plt.axis('off')
    try:
        plt.draw()
        plt.pause(2)
    except:
        pass

def test_mandelbrot():
    with MPICommExecutor() as executor:
        if executor is None: return # worker process
        tic = time.time()
        image = list(executor.map(mandelbrot_line, range(h), chunksize=10))
        toc = time.time()

    print("%s Set %dx%d in %.2f seconds." % ('Mandelbrot', w, h, toc-tic))
    if len(sys.argv) > 1 and sys.argv[1] == '-plot':
        plot(image)

if __name__ == '__main__':
    test_mandelbrot()