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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
|
import time
from math import cos, sin
import numpy as np
from compyle.api import annotate, elementwise, get_config, wrap
@annotate
def julia(i, z, xa, ya, t):
c0 = 0.7885*cos(t)
c1 = 0.7885*sin(t)
x = xa[i]
y = ya[i]
iters = 0
while (x*x + y*y) < 400 and iters < 50:
xn = x*x - y*y + c0
y = x*y*2.0 + c1
x = xn
iters += 1
z[i] = 1.0 - iters*0.02
def timer(x, y, z):
s = time.perf_counter()
n = 2000
dt = 4*np.pi/n
for i in range(n):
julia(z, x, y, -dt*i)
print("Took", time.perf_counter() - s, "seconds")
def plot(x, y, z, nx, ny):
from mayavi import mlab
mlab.figure(size=(600, 600))
xmin, xmax = np.min(x.data), np.max(x.data)
ymin, ymax = np.min(y.data), np.max(y.data)
s = mlab.imshow(z.data.reshape((nx, ny)),
extent=[xmin, xmax, ymin, ymax, 0, 0],
colormap='jet')
s.scene.z_plus_view()
n = 2000
dt = 4*np.pi/n
for i in range(n):
julia(z, x, y, -dt*i)
z.pull()
s.mlab_source.scalars = z.data.reshape((nx, ny))
if i % 3 == 0:
mlab.process_ui_events()
mlab.show()
def save(x, y, z, gif_path='julia_set.gif'):
import imageio as iio
n = 250
dt = 2*np.pi/n
print(f"Writing {gif_path}")
with iio.get_writer(gif_path, mode='I') as writer:
for i in range(n):
julia(z, x, y, -dt*i)
z.pull()
writer.append_data(
(z.data.reshape((nx, ny))*255).astype(np.uint8)
)
print(f"{i}/{n}", end='\r')
print("Done. ")
try:
from pygifsicle import optimize
optimize(gif_path)
except ImportError:
print("Install pygifsicle for an optimized GIF")
if __name__ == '__main__':
from compyle.utils import ArgumentParser
p = ArgumentParser()
p.add_argument('-n', action='store', type=int, dest='n',
default=512, help='Number of grid points in y.')
p.add_argument(
'--show', action='store_true', dest='show',
default=False, help='Show animation (requires mayavi)'
)
p.add_argument(
'--gif', action='store_true',
default=False, help='Make a gif animation (requires imageio)'
)
cfg = get_config()
cfg.suppress_warnings = True
o = p.parse_args()
julia = elementwise(julia)
ny = o.n
nx = int(4*ny//3)
x, y = np.mgrid[-2:2:nx*1j, -1.5:1.5:ny*1j]
x, y = x.ravel(), y.ravel()
z = np.zeros_like(x)
x, y, z = wrap(x, y, z)
timer(x, y, z)
if o.show:
plot(x, y, z, nx, ny)
if o.gif:
save(x, y, z)
|
