In [0]:
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Mandelbrot set

Visualizing the Mandelbrot set doesn't have anything to do with machine learning, but it makes for a fun example of how one can use TensorFlow for general mathematics. This is actually a pretty naive implementation of the visualization, but it makes the point. (A more elaborate implementation may be provided down the line to produce more truly beautiful images.)

Basic setup

You'll need a few imports to get started.


In [0]:
# Import libraries for simulation
import tensorflow as tf
assert tf.__version__.startswith('2')
import numpy as np

# Imports for visualization
import PIL.Image
from io import BytesIO
from IPython.display import clear_output, Image, display

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


In [0]:
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()))

Variable initialization

It's handy that you can mix NumPy and TensorFlow.


In [0]:
# 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

Now you define and initialize TensorFlow tensors.


In [0]:
xs = tf.constant(Z.astype(np.complex64))
zs = tf.Variable(xs)
not_diverged =tf.Variable(tf.zeros_like(xs, tf.bool))
ns = tf.Variable(tf.zeros_like(xs, tf.float32))

Defining and running the computation

Now you specify the computation and run it for a couple hundred steps


In [0]:
for i in range(200): 
  # 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 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.
  #
  ns = ns + tf.cast( not_diverged, tf.float32)

Convert from tensor to numpy array for visualization


In [0]:
ns = ns.numpy()

Let's see what you've got.


In [0]:
DisplayFractal(ns)

Not bad!