barnsley.coco

from random import choices
import numpy as np

data Point(x=0, y=0):
    def __rmatmul__(self, mat: np.array):
        point_array = np.array([self.x, self.y, 1])
        x, y, *_ = tuple(*(mat @ point_array))
        return Point(x, y)


def chaos_game(Point(initial_location), hutchinson_op, probabilities):
    point = initial_location
    while True:
        yield (point := choices(hutchinson_op, probabilities) @ point)

barnsley_hutchinson = [
    np.array([
        [0., 0., 0.],
        [0., 0.16, 0.],
        [0., 0., 1.],
    ]),
    np.array([
        [0.85, 0.04, 0.],
        [-0.04, 0.85, 1.6],
        [0., 0., 1.],
    ]),
    np.array([
        [0.2, -0.26, 0.],
        [0.23, 0.22, 1.6],
        [0., 0., 1.],
    ]),
    np.array([
        [-0.15, 0.28, 0.],
        [0.26, 0.24, 0.44],
        [0., 0., 1.],
    ]),
]

barnsley_probabilities = [0.01, 0.85, 0.07, 0.07]

if __name__ == '__main__':
    output_gen = chaos_game(Point(0, 0), barnsley_hutchinson, barnsley_probabilities)
    output_points = np.array([*output_gen$[:10000]])
    np.savetxt("out.dat", output_points)