Now Code

Contour Plotting with Numpy¶

Overview

Teaching: 0 min
Exercises: 90 min

Questions

How can I use numpy and numerical python to speed up my code?
How can I use numpy to produce pretty plots of numerical data?

Objectives

Use numpy to create/populate 1D and 2D arrays
Create a function which updates one of the array
Output the data as a contour plot
Time and optimise performance with numerical python

Getting Started¶

Open Jupyter Notebooks and change to this directory. Then open a new note notebook.

In this exercise you will write code to generate a Mandelbrot figure like this, clicking on the picture should take you to a higher definition of the image:

The aim of the exercise is to explore using numpy arrays before using them together with matplotlib to generate plots. Finally we will introduce profiling your code and show how you can quickly improve its performance.

1: Set up your arrays

Import numpy in the standard way and generate two 1D arrays, x and y of equally spaced values using specified limits e.g. xmin, xmax (inclusive) and the number of values. Also create a 2D array, z populated with zeros which is the size of y and x respectively in each dimension.

Using small arrays, print the arrays and allocate values to the z to verify that you are creating arrays as requested and to understand why we specify the size of y before x.

Hint: You may find the numpy function linspace useful.

2: Create functions to populate your array

We will now create 2 functions to populate the array z. The x and y arrays will be the 'coordinates' for which each value of z will be generated. The first function should take a single pair of values and a complex number c as required arguments, and implement the following pseudo code:

generate a complex number from the given coordinates in the form $z = x+iy$
loop up to a default maximum of 100 iterations
- if the absolute value of the complex number is greater than a default value of 2, return the number of iterations
- otherwise update the value of the complex number: $z = z \times z + c$
end loop
otherwise return 0

The second function should, take as arguments:

2 $\times$ 1D arrays: x, and y
the 2D array: z
a complex number the value to be used as c in the first function

It should call the first function for all pairs of values in x and y and populate z with the returned value.

Do not worry if you are not familiar with complex numbers, we are using them so that we can create a pretty figure. You can create a complex number out of two values with the function complex(xvalue,yvalue) and maths can be performed as if the were 'normal' numbers.

Set the limits of x and y to $\pm1$ so that they are spaced by $0.1$ and call your function with the complex number c in the equation above, set to -1.

3: Plot your Mandelbrot

Using a filled (contour plot) or otherwise, plot your Mandelbrot, the array z, making sure to add sensible tick marks. You should see a blocky image something like the plot above. Slowly improve the definition by increasing the number of values in x and y, remembering to re-execute all code cells that you change.

Sandpit: Performance and Appearance

You can time your function by prefixing its use with the Python magic %timeit <your_function>. Find out how long your functions takes to run as you increase the definition. You should find that is taking several seconds to run, once you have around 1000 values in each direction. At this point it becomes a bit boring to sit around waiting for our code to run, and these are still quite small images with low definition, also we would like to increase the maximum number of iterations we can run our code for.

There are many ways to optimise the performance of code, which you can explore in this blog post. This simplest and one of the most effective is to import one function and add two lines of code to our functions.

We need to import from the numba library: from numba import jit and add the line @jit before each of your two functions. Re-run your function with these changes and see how long the code now takes to run.

If you explore the link above and the original blog article it references the plotting function: plt.imshow(mandel,cmap='hot', norm=colors.PowerNorm(0.3), interpolation='none'), which you can use to plot a more attractive plot, to use this as written you will also need to import the colors library from matplotlib.

Key Points

numpy contains a rich range of functions to make our lifes easier when peforming numerical work
generally it is quicker than using list, especially when arrays are large(, and multi-dimensional)
The Python magic %timeit can help you profile your code to see where optimisation may help
The numba library jit can provide improved performance without re-writing your code.
To produce attractive plots make use of extensive documentation and examples!