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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
from DType import DType
from Memory import memset_zero
from Object import object, Attr
from Pointer import DTypePointer, Pointer
from Random import rand
from Range import range
from TargetInfo import dtype_sizeof
from Complex import ComplexSIMD as ComplexGenericSIMD
struct Matrix:
var data: DTypePointer[DType.si64]
var rows: Int
var cols: Int
var rc: Pointer[Int]
fn __init__(self&, cols: Int, rows: Int):
self.data = DTypePointer[DType.si64].alloc(rows * cols)
self.rows = rows
self.cols = cols
self.rc = Pointer[Int].alloc(1)
self.rc.store(1)
fn __copyinit__(self&, other: Self):
other._inc_rc()
self.data = other.data
self.rc = other.rc
self.rows = other.rows
self.cols = other.cols
fn __del__(owned self):
self._dec_rc()
fn _get_rc(self) -> Int:
return self.rc.load()
fn _dec_rc(self):
let rc = self._get_rc()
if rc > 1:
self.rc.store(rc - 1)
return
self._free()
fn _inc_rc(self):
let rc = self._get_rc()
self.rc.store(rc + 1)
fn _free(self):
self.data.free()
self.rc.free()
@always_inline
fn __getitem__(self, col: Int, row: Int) -> SI64:
return self.load[1](col, row)
@always_inline
fn load[nelts:Int](self, col: Int, row: Int) -> SIMD[DType.si64, nelts]:
return self.data.simd_load[nelts](row * self.cols + col)
@always_inline
fn __setitem__(self, col: Int, row: Int, val: SI64):
return self.store[1](col, row, val)
@always_inline
fn store[nelts:Int](self, col: Int, row: Int, val: SIMD[DType.si64, nelts]):
self.data.simd_store[nelts](row * self.cols + col, val)
def to_numpy(self) -> PythonObject:
let np = Python.import_module("numpy")
let numpy_array = np.zeros((self.rows, self.cols), np.uint32)
for col in range(self.cols):
for row in range(self.rows):
numpy_array.itemset((row, col), self[col, row].cast[DType.f32]())
return numpy_array
@register_passable("trivial")
struct Complex:
var real: F32
var imag: F32
fn __init__(real: F32, imag: F32) -> Self:
return Self {real: real, imag: imag}
fn __add__(lhs, rhs: Self) -> Self:
return Self(lhs.real + rhs.real, lhs.imag + rhs.imag)
fn __mul__(lhs, rhs: Self) -> Self:
return Self(
lhs.real * rhs.real - lhs.imag * rhs.imag,
lhs.real * rhs.imag + lhs.imag * rhs.real,
)
fn norm(self) -> F32:
return self.real * self.real + self.imag * self.imag
alias xmin: F32 = -2.25
alias xmax: F32 = 0.75
alias xn = 450
alias ymin: F32 = -1.25
alias ymax: F32 = 1.25
alias yn = 375
# Compute the number of steps to escape.
def mandelbrot_kernel(c: Complex) -> Int:
max_iter = 200
z = c
for i in range(max_iter):
z = z * z + c
if z.norm() > 4:
return i
return max_iter
def compute_mandelbrot() -> Matrix:
# create a matrix. Each element of the matrix corresponds to a pixel
result = Matrix(xn, yn)
dx = (xmax - xmin) / xn
dy = (ymax - ymin) / yn
y = ymin
for j in range(yn):
x = xmin
for i in range(xn):
result[i, j] = mandelbrot_kernel(Complex(x, y))
x += dx
y += dy
return result
def make_plot(m: Matrix):
np = Python.import_module("numpy")
plt = Python.import_module("matplotlib.pyplot")
colors = Python.import_module("matplotlib.colors")
dpi = 32
width = 10
height = 10 * yn // xn
fig = plt.figure(1, [width, height], dpi)
ax = fig.add_axes([0.0, 0.0, 1.0, 1.0], False, 1)
light = colors.LightSource(315, 10, 0, 1, 1, 0)
image = light.shade(m.to_numpy(), plt.cm.hot, colors.PowerNorm(0.3), "hsv", 0, 0, 1.5)
plt.imshow(image)
plt.axis("off")
plt.show()
make_plot(compute_mandelbrot())
print("finished")
|
