Mandelbrot set in Python

mandel.py

import os
import math
import click
import datetime
import colorsys
import itertools
import multiprocessing as mp
from PIL import Image
from tqdm import tqdm
from time import time
from yaspin import yaspin
from functools import partial


def log_color(distance, base, const, scale):
    color = -1 * math.log(distance, base)
    rgb = colorsys.hsv_to_rgb(const + scale * color, 0.8, 0.9)
    return tuple(round(i * 255) for i in rgb)


def power_color(distance, exp, const, scale):
    color = distance**exp
    rgb = colorsys.hsv_to_rgb(const + scale * color, 1 - 0.6 * color, 0.9)
    return tuple(round(i * 255) for i in rgb)


def foo(min_x, max_y, x_range, y_range, width, height, precision, color_map, pos_tuple):
    row, col = pos_tuple
    x = min_x + col * x_range / width
    y = max_y - row * y_range / height
    oldX = x
    oldY = y
    rgb = 0
    for i in range(precision + 1):
        a = x*x - y*y   # real component of z^2
        b = 2 * x * y   # imaginary component of z^2
        x = a + oldX    # real component of new z
        y = b + oldY    # imaginary component of new z
        if x*x + y*y > 4:
            break
    if i < precision:
        distance = (i + 1) / (precision + 1)
        if color_map == 'log':
            rgb = log_color(distance, 0.2, 0.27, 1.0)
        elif color_map == 'power':
            rgb = power_color(distance, 0.2, 0.27, 1.0)
    return rgb


@click.command(context_settings={'help_option_names': ['-h', '--help']})
@click.option('--width', '-w', 'width', type=click.INT, default=1000, help='The width of the image. [1000]')
@click.option('--aspect', '-a', 'aspect_ratio', type=click.FLOAT, default=4/3, help='The aspect ratio of the image. [4/3]')
@click.option('--precision', '-r', 'precision', type=click.INT, default=500, help='The precision. [500]')
@click.option('--color', '-c', 'color_map', type=click.Choice(['log', 'power']), default='power',
              help='Image generation color map scheme. [power]')
@click.option('--parallel', '-p', 'parallel', is_flag=True, help='Enables parallel processing.')
def cli(width, aspect_ratio, precision, color_map, parallel):
    t = time()

    # frame parameters
    x, y = -0.65, 0
    x_range = 3.4

    height = round(width / aspect_ratio)
    y_range = x_range / aspect_ratio
    min_x = x - x_range / 2
    max_x = x + x_range / 2
    min_y = y - y_range / 2
    max_y = y + y_range / 2

    img = Image.new('RGB', (width, height), color='black')
    pixels = img.load()

    # Parallel processing
    if parallel:
        with yaspin(text='Computing...'):
            pool = mp.Pool(mp.cpu_count()-1)
            row = range(height)
            col = range(width)
            paramlist = list(itertools.product(row, col))
            unary = partial(foo, min_x, max_y, x_range, y_range,
                            width, height, precision, color_map)
            res = pool.map(unary, paramlist)
            for idx, pixel in enumerate(res):
                row, col = paramlist[idx]
                pixels[col, row] = pixel
    # Normal processing
    else:
        bar_format = '{l_bar}{bar}| [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
        for row in tqdm(range(height), leave=None, bar_format=bar_format):
            for col in tqdm(range(width), leave=None, bar_format=bar_format):
                pixels[col, row] = foo(
                    min_x, max_y, x_range, y_range, width, height, precision, color_map, (row, col))

    elapsed = time() - t
    print(f'Elapsed time:  {str(datetime.timedelta(seconds=elapsed))}')
    img.save('output.png')
    os.system('open output.png')


if __name__ == '__main__':
    cli()

$ python mandel.py -h

Usage: mandel.py [OPTIONS]

Options:
  -w, --width INTEGER      The width of the image. [1000]
  -a, --aspect FLOAT       The aspect ratio of the image. [4/3]
  -r, --precision INTEGER  The precision. [500]
  -c, --color [log|power]  Image generation color map scheme. [power]
  -p, --parallel           Enables parallel processing.
  -h, --help               Show this message and exit.