Copyright 2019 The TensorFlow Authors.



In [ ]:

    
#@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

View source on GitHub

Note: This is an archived TF1 notebook. These are configured to run in TF2's compatbility mode but will run in TF1 as well. To use TF1 in Colab, use the magic.

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. (We may end up providing a more elaborate implementation down the line to produce more truly beautiful images.)

Basic setup

You'll need a few imports to get started.



In [ ]:

    
# Import libraries for simulation
import tensorflow.compat.v1 as tf

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 [ ]:

    
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

For playing around like this, an interactive session is often used, but a regular session would work as well.



In [ ]:

    
sess = tf.InteractiveSession()

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



In [ ]:

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

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

TensorFlow requires that you explicitly initialize variables before using them.



In [ ]:

    
tf.global_variables_initializer().run()

Defining and running the computation

Now you specify more of the computation...



In [ ]:

    
# 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 run it for a couple hundred steps



In [ ]:

    
for i in range(200): step.run()

Let's see what you've got.



In [ ]:

    
DisplayFractal(ns.eval())

Not bad!