Chaos game

chaos_game.py

"""Create a fractal using chaos game."""

import argparse
import os
import random
from collections import namedtuple

# Matplotlib made me do this
import matplotlib
matplotlib.use(os.environ.get('MPL_BACKEND', 'WXAgg'))

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import PatchCollection
from PIL import Image
from tqdm import tqdm

IMG_SIZE = (1000, 1000)
GOALS = ((200, 200), (200, 800), (800, 200), (800, 800))
STEP_SIZE = 1/2
NUM_STEPS = 10_000
POINT = ('X', 'Y')
Point = namedtuple('Point', 'x y')


def take_step(start, goal, step_size=STEP_SIZE):
    """Take a step towards the given goal and return the new position.

    Method from https://math.stackexchange.com/a/1630886. Where t is the ratio of distances
    (dt/d = step_size), then the point (xt,yt)=(((1-t)x0+tx1),((1-t)y0+ty1)).

    Args:
        start (namedtuple[float, float]): Current x, y coordinates.
        goal (namedtuple[float, float]): Destination x, y coordinates (position of chosen goal).
        step_size (float): Value between 0 and 1 representing the proportion of the distance
            towards the goal to travel.

    Returns:
        namedtuple[float, float]: New x, y coordinates.

    """
    newx = (1 - step_size) * start.x + step_size * goal.x
    newy = (1 - step_size) * start.y + step_size * goal.y

    return Point(newx, newy)


def tuple_of_tuples(s):
    """Convert a string of numbers into a tuple of integer pair tuples."""
    vals = s.split()
    try:
        return tuple((int(vals[i]), int(vals[i + 1])) for i in range(0, len(vals), 2))
    except IndexError:
        raise TypeError("List has an odd number of values")


def parse_args():
    """Parse command-line args."""
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--image_size', type=int, nargs=2, metavar=POINT, default=IMG_SIZE,
                        help='Output image size in pixels.')
    parser.add_argument('--step_size', type=float, default=STEP_SIZE,
                        help='Size of each step as a proportion of the distance to the goal.')
    parser.add_argument('--num_steps', type=int, default=NUM_STEPS,
                        help='Total number of steps.')
    parser.add_argument('--output_file',
                        help='Path to an output image file. If not given, result will be shown '
                        'as a plot.')
    parser.add_argument('--goals', type=tuple_of_tuples, default=GOALS,
                        help='Coordinates of the goals to be used in the game, entered as a string '
                        '"X Y X Y ...".')
    parser.add_argument('--marker_size', type=int, default=10,
                        help='Draw a marker for each goal, a cross with the given size in pixels. '
                        'Enter zero to remove markers.')
    return parser.parse_args()


def _add_point(x, y, arr, lines, **kwargs):
    if arr is not None:
        arr[y][x] = 0
    else:
        lines.append(mpatches.CirclePolygon((x, y), radius=2, **kwargs))


def main():
    args = parse_args()
    image_size = Point(*args.image_size)

    # TODO validation
    # check goals are within image bounds
    # 0<step size<1

    # Initialize output, depending on whether saving an image or just plotting
    if args.output_file:
        # reverse x and y values for correct array size
        arr = np.ones(image_size[::-1])
    else:
        arr = None
    kwargs = {'color': 'r', 'linewidth': 5, 'fill': True, 'alpha': 0.5}
    lines = []
    if args.marker_size:
        for goal in args.goals:
            pt = Point(*goal)
            for x in range(pt.x - args.marker_size, pt.x + args.marker_size):
                _add_point(x, pt.y, arr, lines, **kwargs)
            for y in range(pt.y - args.marker_size, pt.y + args.marker_size):
                _add_point(pt.x, y, arr, lines, **kwargs)

    prev_goal = None
    cur_pt = Point(random.randint(0, image_size.x), random.randint(0, image_size.y))
    for _ in tqdm(range(args.num_steps), total=args.num_steps, ascii=True):
        # Current strategy is cannot go to same goal twice in a row
        # TODO more strategies
        poss_goals = [g for g in args.goals if prev_goal is None or prev_goal != g]
        goal = Point(*random.choice(poss_goals))
        cur_pt = take_step(cur_pt, goal, step_size=args.step_size)
        prev_goal = goal
        _add_point(int(cur_pt.x), int(cur_pt.y), arr, lines, **kwargs)

    if args.output_file:
        # flip upside down, because Image origin is top left
        img = Image.fromarray(arr[::-1])
        img.save(args.output_file)
    else:
        fig, ax = plt.subplots(1, 1, figsize=(6, 5))
        ax.add_collection(PatchCollection(lines, match_original=True))
        ax.autoscale(True)
        fig.tight_layout()
        plt.show()


if __name__ == '__main__':
    main()

A Python script to create fractal images using chaos game - https://en.wikipedia.org/wiki/Chaos_game