Chaos Game image generator

chaos_game.py

import sys
import random
import math

import numpy as np
from PIL import Image

from python.common import line


def chaos_game(size, n=3, r=1 / 2, randomize_points=False, coloring=False, random_weight=False, name="chaosgame.png"):
    im = Image.new("RGB", (size + 10, size + 10), (255, 255, 255))

    angle = lambda: math.radians((360 / n))

    if randomize_points:
        angle = lambda: math.radians((360 / n) + random.randrange(0, 360))

    length = size / 2
    canvas_half = size / 2
    corner_points = [(math.cos(i * angle()), math.sin(i * angle()))
                     for i in range(n)]
    rot_angle = math.radians(-90)
    normalized_points = [(round((p[0] * math.cos(rot_angle) - p[1] * math.sin(rot_angle)) * length + canvas_half),
                          round((p[0] * math.sin(rot_angle) + p[1] * math.cos(rot_angle)) * length + canvas_half))
                         for p in corner_points]

    print(corner_points)
    print(normalized_points)

    dmin, dmax = sys.maxsize, -sys.maxsize

    start_point = round(size / 2), round(size / 2)

    paint_points = []

    random_nums = [random.randint(1, n + 1) for _ in range(n)]
    probabilities = [x / sum(random_nums) for x in random_nums]

    print(probabilities)

    for i in range(1000000):
        edge_point = random.choice(normalized_points)
        if random_weight:
            index = np.random.choice(n, 1, p=probabilities)[0]
            edge_point = normalized_points[index]

        paint_point = round(edge_point[0] * (1 - r) + start_point[0] * r), \
                      round(edge_point[1] * (1 - r) + start_point[1] * r)
        start_point = paint_point
        if i > 100:
            distances = []
            if coloring:
                for point in normalized_points:
                    distance = line.fromPoints(start_point, point).len()
                    distances.append(distance)
                    dmin = min(dmin, distance)
                    dmax = max(dmax, distance)

            paint_points.append((paint_point, distances))

    for p, distances in paint_points:
        color = [0, 0, 0]

        if coloring:
            span = dmax - dmin
            ratio = 256 / span

            color.clear()
            sorted(distances)

            for d in distances[:3]:
                c = int((d - dmin) * ratio)
                color.append(c)

        im.putpixel(p, tuple(color))

    im.save(name)

# chaos_game(1500, 5, 1 / 3, False, False, True, "1500_pent_random_weights.png")
# chaos_game(1500, 5, 1 / 3, True, False, False, "1500_pent_random_points.png")
# chaos_game(1500, 5, 1 / 3, False, True, False, "1500_pent_colors.png")

# chaos_game(1500, 3, 1 / 2, False, False, True, "triangle_random_weights.png")
# chaos_game(1500, 3, 1 / 2, True, False, False, "triangle_random_points.png")
# chaos_game(1500, 3, 1 / 2, False, True, False, "triangle_colors.png")