Mandelbrot set drawer

mandel.py

from sys import argv
from math import sqrt, pi, tan, atan2

def mandelbrot(r, i):
    limit = 256
    threshold = 1e10
    zr = 0.0
    zi = 0.0

    for n in range(limit):
        zr_n = zr * zr - zi * zi + r
        zi_n = 2.0 * zr * zi + i
        zr = zr_n
        zi = zi_n
        if zr * zr + zi * zi > threshold:
            return n, zr, zi

    return -1, zr, zi

def rgb_color(hue, brightness):
    colorfulness = 1
    h = tan(hue)
    if brightness < 0:
        brightness = 0
    elif brightness > 255:
        brightness = 255

    if 0.0 <= hue < pi * 2.0 / 3.0:
        blue = 0
        green = int(2 * brightness * h / (3 * h + sqrt(3)))
        if green > brightness:
            green = brightness
        elif green < 0:
            green = 0
        red = brightness - green

    elif - pi * 2.0 / 3.0 <= hue < 0.0:
        green = 0
        blue = int(2 * brightness * h / (3 * h - sqrt(3)))
        if blue > brightness:
            blue = brightness
        elif blue < 0:
            blue = 0
        red = brightness - blue

    else:
        red = 0
        green = int(brightness * (sqrt(3) - h)/(2 * sqrt(3)))
        if green > brightness:
            green = brightness
        elif green < 0:
            green = 0
        blue = brightness - green

    return 255 - red, 255 - green, 255 - blue

if __name__ == '__main__':
    min_r = -2.1
    max_r = 0.6
    min_i = -1.2
    max_i = 1.2

    resolution = 0.0005

    width = int((max_r - min_r) / resolution)
    height= int((max_i - min_i) / resolution)

    f = open(argv[1], 'w')
    print('P3', file=f)
    print('# Mandelbrot set', file=f)
    print(width, height, file=f)
    print('255', file=f)

    for y in reversed(range(height)):
        for x in range(width):
            r = min_r + resolution * x
            i = min_i + resolution * y
            n, zr, zi = mandelbrot(r, i)

            if n < 0:
                r, g, b = 0, 0, 0
            else:
                rad = atan2(zi, zr)
                r, g, b = rgb_color(rad, 8*n)

            print(r, g, b, file=f)

    f.close()