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
|
import sys
import time
from mpi4py.futures import MPICommExecutor
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 = maxit
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 contextlib
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")
with contextlib.suppress(Exception):
plt.draw()
plt.pause(2)
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(f"Mandelbrot Set {w}x{h} in {toc - tic:.2f} seconds.")
if len(sys.argv) > 1 and sys.argv[1] == "-plot":
plot(image)
if __name__ == "__main__":
test_mandelbrot()
|
