The Mandelbrot Set In TensorFlow

What Is the Mandelbrot Set?

According to Wikipedia, “The Mandelbrot set is a famous example of a fractal in mathematics. The Mandelbrot set is important for chaos theory. The edging of the set shows a self-similarity, which is not perfect because it has deformations. The Mandelbrot set can be explained with the equation zn+1 = zn2 + c. In that equation, c and z are complex numbers, and n is zero or a positive integer (natural number). Starting with z0=0, c is in the Mandelbrot set if the absolute value of zn never becomes larger than a certain number (that number depends on c) no matter how large n gets.”

Visualizing the set has nothing to do with machine learning. It can be thought of as another TensorFlow example for mathematics. So, let’s learn how can we compute the Mandelbrot set in TensorFlow.

Setup For Mandelbrot Set in TensorFlow

For the Mandelbrot set in TensorFlow, you’ll need a few imports to get started.

# Import libraries for simulation
import tensorflow as tf
import numpy as np
# Imports for visualization
import PIL.Image
from io import BytesIO
from IPython.display import Image, display

Now, define a function to actually display the image once you have iteration counts.

def DisplayFractal(a, fmt='jpeg'):
  """Display an array of iteration counts as a
     colorful picture of a fractal."""
  a_cyclic = (6.28*a/20.0).reshape(list(a.shape)+[1])
  img = np.concatenate([10+20*np.cos(a_cyclic),
                        30+50*np.sin(a_cyclic),
                        155-80*np.cos(a_cyclic)], 2)
  img[a==a.max()] = 0
  a = img
  a = np.uint8(np.clip(a, 0, 255))
  f = BytesIO()
  PIL.Image.fromarray(a).save(f, fmt)
  display(Image(data=f.getvalue()))

Session and Variable Initialization in the Mandelbrot Set

Here, an interactive session is used, but a regular session would work as well.

sess = tf.InteractiveSession()

Introducing NumPy with Tensorflow:

# Use NumPy to create a 2D array of complex numbers
Y, X = np.mgrid[-1.3:1.3:0.005, -2:1:0.005]
Z = X+1j*Y

Initializing TensorFlow tensors.

xs = tf.constant(Z.astype(np.complex64))
zs = tf.Variable(xs)
ns = tf.Variable(tf.zeros_like(xs, tf.float32))

Now, as you are already aware, TensorFlow requires you to explicitly declare variables before using them.

tf.global_variables_initializer().run()

Running Computation — TensorFlow Mandelbrot Set

Now, we specify more of the computation for the Mandelbrot set in TensorFlow.

# Compute the new values of z: z^2 + x
zs_ = zs*zs + xs
# Have we diverged with this new value?
not_diverged = tf.abs(zs_) < 4
# Operation to update the zs and the iteration count.
#
# Note: We keep computing zs after they diverge! This
#       is very wasteful! There are better, if a little
#       less simple, ways to do this.
#
step = tf.group(
  zs.assign(zs_),
  ns.assign_add(tf.cast(not_diverged, tf.float32))
  )

and we run it for a couple hundred iterations:
for i in range(200): step.run()

Then, displaying the result using:
DisplayFractal(ns.eval())

Conclusion

We saw the concept of the Mandelbrot Set in TensorFlow, just like the Mandelbrot Set solution. TensorFlow proves to be a great tool for visualizing and understanding not only machine learning concepts but complex mathematical functions including random processes and graph theory. Lastly, we discussed how to run the computation for TensorFlow Mandelbrot Set. Next up is solving PDEs using TensorFlow. If you have any questions regarding the Mandelbrot Set in TensorFlow, feel free to ask in the comments section.