Generator of Mandelbrot set and Julia fractals

madelbrot.py

import sys
import math
import multiprocessing as mp
import time
import os

from PIL import Image


def mandelbrot_p(width, steps=25, zoom=((-2, -1), (1, 1)), c=None, name="fractal.png", number_processes=mp.cpu_count()):
    q = mp.Queue()
    s1 = time.time()
    w = round(width / number_processes)

    zooms = []
    part_x = abs((zoom[1][0] - zoom[0][0]) / number_processes)

    for i in range(number_processes):
        diff = i * part_x
        diff_2 = (1 + i) * part_x
        zooms.append(
            ((zoom[0][0] + diff, zoom[0][1]), (zoom[0][0] + diff_2, zoom[1][1])))

    processes = [mp.Process(target=mandelbrot_set, name=str(x + 1), args=(q, w, steps, zooms[x], c)) for x in
                 range(number_processes)]

    for p in processes:
        p.start()

    min_s = sys.maxsize
    max_s = 0

    arrays = {}
    for p in processes:
        res = q.get()
        arrays[res[0]] = res[1]
        min_s = min(min_s, res[2][0])
        max_s = max(max_s, res[2][1])

    for p in processes:
        p.join()
    print("Calculation time:", str(time.time() - s1))
    q.close()
    s1 = time.time()

    imageprocesses = [mp.Process(target=color_image, args=(arrays[z], min_s, max_s, z)) for z in arrays.keys()]

    for p in imageprocesses:
        p.start()

    for p in imageprocesses:
        p.join()

    x = 0
    result = None
    for z in zooms:
        timg = Image.open("tmp/" + str(z) + ".png", "r")
        if result is None:
            result = Image.new("RGB", (width, timg.size[1]))
        result.paste(timg, (x, 0))
        x += w
        os.remove("tmp/" + str(z) + ".png")

    print("Image generation time:", str(time.time() - s1))
    result.save(name)
    # result.show()


def color_image(array, _min, _max, _id):
    im = Image.new("RGB", (len(array), len(array[0])))
    for x in range(len(array)):
        for y in range(len(array[x])):
            val = array[x][y]
            col = (0, 0, 0)
            if val != -1:
                col = rgb(_min, _max, val)
            im.putpixel((x, y), col)
    im.save("tmp/" + str(_id) + ".png")


def mandelbrot_set(queue, width, steps, zoom, c=None):
    real_denom = abs(zoom[1][0] - zoom[0][0])
    imag_denom = abs(zoom[1][1] - zoom[0][1])

    box_ratio = imag_denom / real_denom

    height = round(width * box_ratio)

    finalSteps = [[0.0 for y in range(height)] for x in range(width)]

    escaperadius = 20000000000

    max_s = 0
    min_s = sys.maxsize

    for x in range(width):
        for y in range(height):
            real = x / (width / real_denom) + zoom[0][0]
            imag = y / (height / imag_denom) + zoom[0][1]
            z, zn, curr_steps = 0, 0, 0
            if c is None:
                z = complex(real, imag)
            else:
                z = c
                zn = complex(real, imag)
            while curr_steps <= steps and abs(zn) < escaperadius:
                zn = zn ** 2 + z
                curr_steps += 1

            zn = zn ** 2 + z
            curr_steps += 1
            zn = zn ** 2 + z
            curr_steps += 1

            smooth = -1.0

            min_s = min(min_s, curr_steps)
            max_s = max(max_s, curr_steps)

            if abs(zn) >= 2:
                smooth = curr_steps + 1 - \
                         math.log(math.log(abs(zn))) / math.log(2)
            finalSteps[x][y] = smooth
    queue.put((zoom, finalSteps, (min_s, max_s)))


def rgb(minimum, maximum, value):
    minimum, maximum = float(minimum), float(maximum)
    ratio = 4 * (value - minimum) / (maximum - minimum)
    b = int(max(0, 255 * (1 - ratio)))
    r = int(max(0, 255 * (ratio - 1)))
    g = 255 - b - r
    return r, g, b


def timeit(title, function):
    print(title)
    s = time.time()
    exec(function)
    e = time.time() - s
    print("total time: ", e)


if __name__ == '__main__':
    # timeit("seahorse 1 process", "mandelbrot_p(1000, 200, zoom=((-.76, 0.1),(-0.74,0.12)), number_processes=1)")
    # timeit("seahorse 4 processes", "mandelbrot_p(1000, 200, zoom=((-.76, 0.1),(-0.74,0.12)), number_processes=4)")
    # timeit("mandelbrot 6 processes", "mandelbrot_p(1000, 60, number_processes=6)")
    # timeit("mandelbrot 5 processes", "mandelbrot_p(1000, 60, number_processes=5)")
    # timeit("mandelbrot 4 processes", "mandelbrot_p(1000, 60, number_processes=4)")
    # timeit("mandelbrot 3 processes", "mandelbrot_p(1000, 60, number_processes=3)")
    # timeit("mandelbrot 2 processes", "mandelbrot_p(1000, 60, number_processes=2)")
    # timeit("mandelbrot 1 process", "mandelbrot_p(1000, 60, number_processes=1)")

    # mandelbrot_p(1500, 30, c=(0.123-0.745j), zoom=((-2, -2), (2, 2)), name="juliaset.png")
    # mandelbrot_p(1500, 70, c=(-.75), zoom=((-2, -2), (2, 2)), name="juliaset_1.png")
    # mandelbrot_p(1500, 60, c=(-.391 - .587j), zoom=((-2, -2), (2, 2)), name="juliaset_2.1.png")
    pass