[Numpy-discussion] How to remove fortran-like loops with numpy?

Stéfan van der Walt stefan at sun.ac.za
Mon Jun 8 21:39:03 EDT 2009

Hi Juan

2009/6/8 Juanjo Gomez Navarro <juanjo.gomeznavarro at gmail.com>:
> I'm new in numpy. Actually, I'm new in Python. In order to learn a bit, I
> want to create a program to plot the Mandelbrot set. This program is quite
> simple, and I have already programmed it. The problem is that I come from
> fortran, so I use to think in "for" loops. I know that it is not the best
> way to use Python and in fact the performance of the program is more than
> poor.
>
> Here is the program:
>
>> #!/usr/bin/python
>>
>> import numpy as np
>> import matplotlib.pyplot as plt
>>
>> # Some parameters
>> Xmin=-1.5
>> Xmax=0.5
>> Ymin=-1
>> Ymax=1
>>
>> Ds = 0.01
>>
>> # Initialization of varibles
>> X = np.arange(Xmin,Xmax,Ds)
>> Y = np.arange(Ymax,Ymin,-Ds)
>>
>> N = np.zeros((X.shape,Y.shape),'f')
>>
>> ############## Here are inefficient the calculations ################
>> for i in range(X.shape):
>>   for j in range(Y.shape):
>>     z= complex(0.0, 0.0)
>>     c = complex(X[i], Y[j])
>>     while N[i, j] < 30 and abs(z) < 2:
>>       N[i, j] += 1
>>       z = z**2 + c
>>     if N[i, j] == 29:
>>       N[i, j]=0
>> ####################################################
>>
>> # And now, just for ploting...
>> N = N.transpose()
>> fig = plt.figure()
>> plt.imshow(N,cmap=plt.cm.Blues)
>> plt.title('Mandelbrot set')
>> plt.xticks([]); plt.yticks([])
>> plt.show()
>> fig.savefig('test.png')
>
>
> As you can see, it is very simple, but it takes several seconds running just
> to create a 200x200 plot. Fortran takes same time to create a 2000x2000
> plot, around 100 times faster... So the question is, do you know how to
> programme this in a Python-like fashion in order to improve seriously the
> performance?

Here is another version, similar to Robert's, that I wrote up for the
documentation project last year:

http://mentat.za.net/numpy/intro/intro.html

We never used it, but I still like the pretty pictures :-)

Cheers
Stéfan