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

Juanjo Gomez Navarro juanjo.gomeznavarro at gmail.com
Mon Jun 8 17:52:29 EDT 2009

Hi all,

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?

--
Juan José Gómez Navarro

Edificio CIOyN, Campus de Espinardo, 30100
Departamento de Física