# Mandelbrot using GPU and two CPU variants

Calculate the Mandelbrot set, display it in a Tkinter window. You can choose to use the GPU (ElementwiseKernel) for the calculation or two numpy CPU solutions. Timing results are shown in the code.

 License of this example: GPL Date: July 2010 PyCUDA version: 0.94

```   1 # Mandelbrot calculate using GPU, Serial numpy and faster numpy
2 # Use to show the speed difference between CPU and GPU calculations
3 # ian@ianozsvald.com July 2010
4
5 # Based on vegaseat's TKinter/numpy example code from 2006
6 # http://www.daniweb.com/code/snippet216851.html#
7 # with minor changes to move to numpy from the obsolete Numeric
8
9 import sys
10 import numpy as nm
11
12 import Tkinter as tk
13 import Image          # PIL
14 import ImageTk        # PIL
15
16 import pycuda.driver as drv
17 import pycuda.tools
18 import pycuda.autoinit
19 from pycuda.compiler import SourceModule
20 import pycuda.gpuarray as gpuarray
21
22 # set width and height of window, more pixels take longer to calculate
23 w = 1000
24 h = 1000
25
26 from pycuda.elementwise import ElementwiseKernel
27 complex_gpu = ElementwiseKernel(
28         "pycuda::complex<float> *z, pycuda::complex<float> *q, int *iteration, int maxiter",
29             "for (int n=0; n < maxiter; n++) {z[i] = (z[i]*z[i])+q[i]; if (abs(z[i]) > 2.0f) {iteration[i]=n; z[i] = pycuda::complex<float>(); q[i] = pycuda::complex<float>();};}",
30         "complex5",
31         preamble="#include <pycuda-complex.hpp>",)
32
33 def calculate_z_gpu(q, maxiter, z):
34     output = nm.resize(nm.array(0,), q.shape)
35     q_gpu = gpuarray.to_gpu(q.astype(nm.complex64))
36     z_gpu = gpuarray.to_gpu(z.astype(nm.complex64))
37     iterations_gpu = gpuarray.to_gpu(output)
38     # the for loop and complex calculations are all done on the GPU
39     # we bring the iterations_gpu array back to determine pixel colours later
40     complex_gpu(z_gpu, q_gpu, iterations_gpu, maxiter)
41
42     iterations = iterations_gpu.get()
43     return iterations
44
45
46 def calculate_z_numpy_gpu(q, maxiter, z):
47     """Calculate z using numpy on the GPU via gpuarray"""
48     outputg = gpuarray.to_gpu(nm.resize(nm.array(0,), q.shape).astype(nm.int32))
49     zg = gpuarray.to_gpu(z.astype(nm.complex64))
50     qg = gpuarray.to_gpu(q.astype(nm.complex64))
51     # 2.0 as an array
52     twosg = gpuarray.to_gpu(nm.array([2.0]*zg.size).astype(nm.float32))
53     # 0+0j as an array
54     cmplx0sg = gpuarray.to_gpu(nm.array([0+0j]*zg.size).astype(nm.complex64))
55     # for abs_zg > twosg result
56     comparison_result = gpuarray.to_gpu(nm.array([False]*zg.size).astype(nm.bool))
57     # we'll add 1 to iterg after each iteration
58     iterg = gpuarray.to_gpu(nm.array([0]*zg.size).astype(nm.int32))
59
60     for iter in range(maxiter):
61         zg = zg*zg + qg
62
63         # abs returns a complex (rather than a float) from the complex
64         # input where the real component is the absolute value (which
65         # looks like a bug) so I take the .real after abs()
66         abs_zg = abs(zg).real
67
68         comparison_result = abs_zg > twosg
69         qg = gpuarray.if_positive(comparison_result, cmplx0sg, qg)
70         zg = gpuarray.if_positive(comparison_result, cmplx0sg, zg)
71         outputg = gpuarray.if_positive(comparison_result, iterg, outputg)
72         iterg = iterg + 1
73     output = outputg.get()
74     return output
75
76
77 def calculate_z_numpy(q, maxiter, z):
78     # calculate z using numpy, this is the original
79     # routine from vegaseat's URL
80     # NOTE this routine was faster using a default of double-precision complex nbrs
81     # rather than the current single precision
82     output = nm.resize(nm.array(0,), q.shape).astype(nm.int32)
83     for iter in range(maxiter):
84         z = z*z + q
85         done = nm.greater(abs(z), 2.0)
86         q = nm.where(done,0+0j, q)
87         z = nm.where(done,0+0j, z)
88         output = nm.where(done, iter, output)
89     return output
90
91 def calculate_z_serial(q, maxiter, z):
92     # calculate z using pure python with numpy arrays
93     # this routine unrolls calculate_z_numpy as an intermediate
94     # step to the creation of calculate_z_gpu
95     # it runs slower than calculate_z_numpy
96     output = nm.resize(nm.array(0,), q.shape).astype(nm.int32)
97     for i in range(len(q)):
98         if i % 100 == 0:
99             # print out some progress info since it is so slow...
100             print "%0.2f%% complete" % (1.0/len(q) * i * 100)
101         for iter in range(maxiter):
102             z[i] = z[i]*z[i] + q[i]
103             if abs(z[i]) > 2.0:
104                 q[i] = 0+0j
105                 z[i] = 0+0j
106                 output[i] = iter
107     return output
108
109
110 show_instructions = False
111 if len(sys.argv) == 1:
112     show_instructions = True
113
114 if len(sys.argv) > 1:
115     if sys.argv[1] not in ['gpu', 'gpuarray', 'numpy', 'python']:
116         show_instructions = True
117
118 if show_instructions:
119     print "Usage: python mandelbrot.py [gpu|gpuarray|numpy|python]"
120     print "Where:"
121     print " gpu is a pure CUDA solution on the GPU"
122     print " gpuarray uses a numpy-like CUDA wrapper in Python on the GPU"
123     print " numpy is a pure Numpy (C-based) solution on the CPU"
124     print " python is a pure Python solution on the CPU with numpy arrays"
125     sys.exit(0)
126
127 routine = {'gpuarray':calculate_z_numpy_gpu,
128             'gpu':calculate_z_gpu,
129             'numpy':calculate_z_numpy,
130             'python':calculate_z_serial}
131
132 calculate_z = routine[sys.argv[1]]
133 ##if sys.argv[1] == 'python':
134 #    import psyco
135 #    psyco.full()
136
137 # Using a WinXP Intel Core2 Duo 2.66GHz CPU (1 CPU used)
138 # with a 9800GT GPU I get the following timings (smaller is better).
139 # With 200x200 problem with max iterations set at 300:
140 # calculate_z_gpu: 0.03s
141 # calculate_z_serial: 8.7s
142 # calculate_z_numpy: 0.3s
143 #
144 # Using WinXP Intel 2.9GHz CPU (1 CPU used)
145 # with a GTX 480 GPU I get the following using 1000x1000 plot with 1000 max iterations:
146 # gpu: 0.07s
147 # gpuarray: 3.4s
148 # numpy: 43.4s
149 # python (serial): 1605.6s
150
151 class Mandelbrot(object):
152     def __init__(self):
153         # create window
154         self.root = tk.Tk()
155         self.root.title("Mandelbrot Set")
156         self.create_image()
157         self.create_label()
158         # start event loop
159         self.root.mainloop()
160
161
162     def draw(self, x1, x2, y1, y2, maxiter=300):
163         # draw the Mandelbrot set, from numpy example
164         xx = nm.arange(x1, x2, (x2-x1)/w*2)
165         yy = nm.arange(y2, y1, (y1-y2)/h*2) * 1j
166         # force yy, q and z to use 32 bit floats rather than
167         # the default 64 doubles for nm.complex for consistency with CUDA
168         yy = yy.astype(nm.complex64)
169         q = nm.ravel(xx+yy[:, nm.newaxis]).astype(nm.complex64)
170         z = nm.zeros(q.shape, nm.complex64)
171
172         start_main = drv.Event()
173         end_main  = drv.Event()
174         start_main.record()
175         output = calculate_z(q, maxiter, z)
176
177         end_main.record()
178         end_main.synchronize()
179         secs = start_main.time_till(end_main)*1e-3
180         print "Main took", secs
181
182         output = (output + (256*output) + (256**2)*output) * 8
183         # convert output to a string
184         self.mandel = output.tostring()
185
186     def create_image(self):
187         """"
188         create the image from the draw() string
189         """
190         self.im = Image.new("RGB", (w/2, h/2))
191         # you can experiment with these x and y ranges
192         self.draw(-2.13, 0.77, -1.3, 1.3, 1000)
193         self.im.fromstring(self.mandel, "raw", "RGBX", 0, -1)
194
195     def create_label(self):
196         # put the image on a label widget
197         self.image = ImageTk.PhotoImage(self.im)
198         self.label = tk.Label(self.root, image=self.image)
199         self.label.pack()
200
201 # test the class
202 if __name__ == '__main__':
203     test = Mandelbrot()
```