Calculates Newton's Fractal

newtonsfractal.py

#!/usr/bin/env python3
# Author: @xPMo
# Based on previous work by: Jordan Atlas
#
# Purpose: creates a fractal using the Newton-Raphson
# root-finding method. The polynomial is derived from the
# given roots in commandline args, or the roots of unity
# are used
#
# Based on listing 7.1 from "Beginning Python Visualization"
# by Shai Vaingast

from PIL import Image
from cmath import *
from optparse import OptionParser
from numpy import convolve, core

parser = OptionParser()
parser.usage = "%prog root [root [...]] [options]"
parser.description = "Generate a Newton's Fractal image by repeated application of Newton's method.  This uses the polynomial with the given single roots. (Put the root in twice for a double root, etc.)"
parser.epilog = "Note: Roots will be `eval()`ed, so you probably want to use the form \"a±bj\", where \"a\" and \"b\" are a numeral type. The option parser will attempt to take a leading \"-\" as an option/flag, so enter \"-1-2j\" as \"0-1-2j\" or \" -1-2j\" (note the space) instead. Any cmath functions are supported as well, so feel free to input \"cos(pi/4)-1j*sin(pi/4)\"."
parser.add_option("-d","--delta", dest="delta", type="float", default=1e-6, help="Convergance criterion: within this distance, the root is found. Default is 10**-6 (1e-6).")
parser.add_option("-p","--ppu", dest="pres", type="int", default=0, help="Number of pixels per unit length. Default is half of the smaller of x-res and y-res.")
parser.add_option("-x","--xres", dest="xres", type="int", default=0, help="X-resolution of the final image. If only one of x or y is provided, the image will be square at that resolution. Default is 1500.")
parser.add_option("-y","--yres", dest="yres", type="int", default=0, help="Y-resolution of the final image. Default is 1500.")
parser.add_option("-c","--center","--offset", metavar="X+Yj", dest="center", default="0", help="The coordinate of the center of the image. See the note below on leading \"-\"s, it applies here as well. Default is \"0 + 0j\".")
parser.add_option("-i","--iterations", dest="iters", type="int", default=30, help="Number of iterations. Default is 30.")
parser.add_option("-u","--unity", metavar="N", dest="unity", type="int", default=0, help="Add the N roots of unity to the provided roots.")
parser.add_option("-v","--verbose", action="store_true", default=False, help="Print out progress and other information.")
parser.add_option("-f","--file-string", metavar="FORMAT", dest="file", action="store", default="fractal_{0}iters_{1}_{2}x{3}_{4}.png", help="Specify the default filestring to be used.")

(opts, args) = parser.parse_args()

colors = [
    # red green blue
    (1,0,0), (0,1,0), (0,0,1),
    # yellow cyan purple
    (1,1,0),(0,0.75,0.75) ,(0.75,0,0.75),
    # orange teal indigo
    (0.9,0.6,0),(0,0.9,0.6) ,(0.6,0,0.9),
    # ygreen gblue magenta
    (0.6,0.9,0),(0,0.6,0.9) ,(0.9,0,0.6),
]
# grey
color_fallback = (0.5,0.5,0.5)

# =================
# Preperation
# =================

# Set variables from opts
delta = opts.delta
iters = opts.iters
unity = opts.unity
verbose = opts.verbose
center = eval(opts.center)

xres = opts.xres
yres = opts.yres
if xres == 0:
    xres = yres
    if xres == 0:
        (xres, yres) = (1500, 1500)
elif yres == 0:
    yres = xres

pres = opts.pres
if pres == 0:
    pres = yres // 2 if xres > yres else xres // 2

f = [1]
roots = []
df = []
# set roots and function from unity opt
if unity >= 1:
    roots = [cos((2*n)*pi/unity)+1j*sin((2*n)*pi/unity) for n in range(unity)]
    f = [1] + [0] * (unity - 1) + [-1]
# set roots and function from args
for arg in args:
    root = eval(arg)
    roots.append(root)
    f = convolve(f, [1, -root])
# set derivative from function
for i in range(len(f) - 1):
    df.append(f[i] * (len(f) - i -1))

imgfile = opts.file.format(iters, pres, xres, yres, int(abs(log10(delta))))

# Print out info on roots, function, derivative, and image location
if verbose:
    def human_readable(f):
        if len(f) > 2:
            print("{1}z^{0}".format(len(f) - 1, f[0]), end="")
            for i in range(1, len(f) - 2):
                if f[i] != 0:
                    print(" + {1}z^{0}".format(len(f) - i - 1, f[i]), end="")
            if f[len(f) - 2] != 0:
                print(" + {0}z".format(f[len(f) - 2]), end="")
        else:
            print("{0}z".format(f[len(f) - 2]), end="")
        if f[len(f) - 1] != 0:
            print(" + {0}".format(f[len(f) - 1]))
        else:
            print()

    print("Roots:     ", end="")
    for root in roots:
        print(" {0:4g},".format(root), end="")
    print("\nFunction:   ", end="")
    human_readable(f)
    print("Derivative: ", end="")
    human_readable(df)
    print("\n\nSaving image to {0}\n".format(imgfile))

# create an image to draw on, paint it black
img = Image.new("RGB", (xres, yres), (0,0,0))

# optimization: Compute static arrays' lengths only once
len_f = len(f)
len_df = len(df)
len_roots = len(roots)
len_colors = len(colors)

# optimization: compute bottom corner (starting value) of image
startx = center.real - xres/2/pres
starty = center.imag - yres/2/pres

# =================
# Main computation loop
# =================
for re in range(xres):
    # optimization: compute re_adjusted only once per loop
    re_adjusted = re / pres + startx
    if verbose:
        print("On line {0} of {1}".format(re, xres), end='\r')
    for im in range(yres):
        z = re_adjusted +  1j * (im / pres + starty)
        for i in range(iters):
            # inlining: df_of_z = polyval(df, z)
            df_of_z = 0
            for term in range(len_df):
                df_of_z = df_of_z * z + df[term]
            if df_of_z == 0:
                # prevent division by zero
                continue

            # inlining: f_of_z = polyval(f, z)
            f_of_z = 0
            for term in range(len_f):
                f_of_z = f_of_z * z + f[term]
            z -= f_of_z/df_of_z
            try:
                if (abs(f_of_z) < delta):
                    break
            except OverflowError:
                break

        # color depth is a function of the number of iterations
        color_depth = (iters-i)*255.0/iters

        # find to which solution this guess converged to
        err_index = len_roots - 1
        err = abs(z - roots[len_roots - 1])
        for i in range(len_roots - 1):
            new_err = abs(z - roots[i])
            if new_err < err:
                err_index = i
                err = new_err

        # select the color associated with the solution
        color = color_fallback
        if err_index < len_colors:
            color = colors[err_index]

        # (re=0, im=0) in bottom left of image
        img.putpixel((re, yres - im - 1), tuple([int(i * color_depth) for i in color]))

# save the image to the working directory
if verbose:
    print("\n\nSaving image to {0}".format(imgfile))
img.save(imgfile, dpi = (150, 150))