Learn how you can become a Python programmer in just 12 weeks.

    We respect your privacy. Unsubscribe at anytime.

    NumPy: Compute Mandelbrot set by Vectorization

    What will we cover in this tutorial?

    • Understand what the Mandelbrot set it and why it is so fascinating.
    • Master how to make images in multiple colors of the Mandelbrot set.
    • How to implement it using NumPy vectorization.

    Step 1: What is Mandelbrot?

    Mandelbrot is a set of complex numbers for which the function f(z) = z^2 + c does not converge when iterated from z=0 (from wikipedia).

    Take a complex number, c, then you calculate the sequence for N iterations:

    z_(n+1) = z_n + c for n = 0, 1, …, N-1

    If absolute(z_(N-1)) < 2, then it is said not to diverge and is part of the Mandelbrot set.

    The Mandelbrot set is part of the complex plane, which is colored by numbers part of the Mandelbrot set and not.

    Mandelbrot set.

    This only gives a block and white colored image of the complex plane, hence often the images are made more colorful by giving it colors by the iteration number it diverged. That is if z_4 diverged for a point in the complex plane, then it will be given the color 4. That is how you end up with colorful maps like this.

    Mandelbrot set (made by program from this tutorial).

    Step 2: Understand the code of the non-vectorized approach to compute the Mandelbrot set

    To better understand the images from the Mandelbrot set, think of the complex numbers as a diagram, where the real part of the complex number is x-axis and the imaginary part is y-axis (also called the Argand diagram).

    Argand diagram

    Then each point is a complex number c. That complex number will be given a color depending on which iteration it diverges (if it is not part of the Mandelbrot set).

    Now the pseudocode for that should be easy to digest.

    for x in [-2, 2] do:
      for y in [-1.5, 1.5] do:
        c = x + i*y
        z = 0
        N = 0
        while absolute(z) < 2 and N < MAX_ITERATIONS:
          z = z^2 + c
        set color for x,y to N

    Simple enough to understand. That is some of the beauty of it. The simplicity.

    Step 3: Make a vectorized version of the computations

    Now we understand the concepts behind we should translate that into to a vectorized version. If you are new to vectorization we can recommend you read this tutorial first.

    What do we achieve with vectorization? That we compute all the complex numbers simultaneously. To understand that inspect the initialization of all the points here.

    import numpy as np
    def mandelbrot(height, width, x_from=-2, x_to=1, y_from=-1.5, y_to=1.5, max_iterations=100):
        x = np.linspace(x_from, x_to, width).reshape((1, width))
        y = np.linspace(y_from, y_to, height).reshape((height, 1))
        c = x + 1j * y

    You see that we initialize all the x-coordinates at once using the linespace. It will create an array with numbers from x_from to x_to in width points. The reshape is to fit the plane.

    The same happens for y.

    Then all the complex numbers are created in c = x + 1j*y, where 1j is the imaginary part of the complex number.

    This leaves us to the full implementation.

    There are two things we need to keep track of in order to make a colorful Mandelbrot set. First, in which iteration the point diverged. Second, to achieve that, we need to remember when a point diverged.

    import numpy as np
    import matplotlib.pyplot as plt
    def mandelbrot(height, width, x=-0.5, y=0, zoom=1, max_iterations=100):
        # To make navigation easier we calculate these values
        x_width = 1.5
        y_height = 1.5*height/width
        x_from = x - x_width/zoom
        x_to = x + x_width/zoom
        y_from = y - y_height/zoom
        y_to = y + y_height/zoom
        # Here the actual algorithm starts
        x = np.linspace(x_from, x_to, width).reshape((1, width))
        y = np.linspace(y_from, y_to, height).reshape((height, 1))
        c = x + 1j * y
        # Initialize z to all zero
        z = np.zeros(c.shape, dtype=np.complex128)
        # To keep track in which iteration the point diverged
        div_time = np.zeros(z.shape, dtype=int)
        # To keep track on which points did not converge so far
        m = np.full(c.shape, True, dtype=bool)
        for i in range(max_iterations):
            z[m] = z[m]**2 + c[m]
            diverged = np.greater(np.abs(z), 2, out=np.full(c.shape, False), where=m) # Find diverging
            div_time[diverged] = i      # set the value of the diverged iteration number
            m[np.abs(z) > 2] = False    # to remember which have diverged
        return div_time
    # Default image of Mandelbrot set
    plt.imshow(mandelbrot(800, 1000), cmap='magma')
    # The image below of Mandelbrot set
    # plt.imshow(mandelbrot(800, 1000, -0.75, 0.0, 2, 200), cmap='magma')
    # The image below of below of Mandelbrot set
    # plt.imshow(mandelbrot(800, 1000, -1, 0.3, 20, 500), cmap='magma')

    Notice that z[m] = z[m]**2 + c[m] only computes updates on values that are still not diverged.

    I have added the following two images from above (the one not commented out is above in previous step.

    Mandelbrot set from the program above.
    Mandelbrot set from the code above.
    Also check out the tutorial on Julia sets.

    Python for Finance: Unlock Financial Freedom and Build Your Dream Life

    Discover the key to financial freedom and secure your dream life with Python for Finance!

    Say goodbye to financial anxiety and embrace a future filled with confidence and success. If you’re tired of struggling to pay bills and longing for a life of leisure, it’s time to take action.

    Imagine breaking free from that dead-end job and opening doors to endless opportunities. With Python for Finance, you can acquire the invaluable skill of financial analysis that will revolutionize your life.

    Make informed investment decisions, unlock the secrets of business financial performance, and maximize your money like never before. Gain the knowledge sought after by companies worldwide and become an indispensable asset in today’s competitive market.

    Don’t let your dreams slip away. Master Python for Finance and pave your way to a profitable and fulfilling career. Start building the future you deserve today!

    Python for Finance a 21 hours course that teaches investing with Python.

    Learn pandas, NumPy, Matplotlib for Financial Analysis & learn how to Automate Value Investing.

    “Excellent course for anyone trying to learn coding and investing.” – Lorenzo B.

    2 thoughts on “NumPy: Compute Mandelbrot set by Vectorization”

    1. Thanks for the great post, Helped me to understand better numpy’s vectorization.

      In working through the code I noticed an improvement that could be made, so here it is:

      The last line of the inner loop is `m[np.abs(z) > 2] = False`. That’s doing quite a lot of calculation.

      You can simply replace it with `m[diverged] = False`.

      The boolean array `diverged` already holds `True` values for each cell/pixel/z-sequence that diverged
      on the current iteration.

      This cut the inner loop computation time to about 61% of the original on my simple tests.


    Leave a Comment