FractalTrees.py

import numpy as np
from math import sin, cos
import random as rand
from pyglet.gl import *
import sys

AMOUNT_TO_SHRINK = rand.uniform(0.50, 0.75)
# Becareful of setting this too high as it will take longer to create the tree the higher you put it.
# At values higher than 15(2^15 branches) is where you will notice this and it will probably hang for quite some time.
TREE_DEPTH = rand.randint(10, 15)

SIN_MEMOIZED_VALUES = {}
COS_MEMOIZED_VALUES = {}

# Change these RGB colors to your liking to create BEAUTIFUL colored trees.
BRANCH_COLOUR = (101, 67, 33, 101, 67, 33)
BRANCH_LEAF_COLUR = (0, 100, 0, 0, 100, 0)


def memoizedSin(degree):
    if degree not in SIN_MEMOIZED_VALUES:
        SIN_MEMOIZED_VALUES[degree] = sin(np.deg2rad(degree))
    return SIN_MEMOIZED_VALUES[degree]


def memoizedCos(degree):
    if degree not in COS_MEMOIZED_VALUES:
        COS_MEMOIZED_VALUES[degree] = cos(np.deg2rad(degree))
    return COS_MEMOIZED_VALUES[degree]


def rotateVector(vector, degree):
    cosAlpha = memoizedCos(degree)
    sinAlpha = memoizedSin(degree)
    return np.matmul(vector, [[cosAlpha, -sinAlpha], [sinAlpha, cosAlpha]])  # Rotational counter-clockwise matrix


class Branch:
    def __init__(self, begin, end, color):
        self.begin = np.array(begin)
        self.end = np.array(end)
        self.vertices = pyglet.graphics.vertex_list(2,
                                                    ('v2f', (self.begin[0], self.begin[1], self.end[0], self.end[1])),
                                                    ('c3B', color)
                                                    )

    def branch(self, degree, color):
        dir = self.end - self.begin
        dir = rotateVector(dir, degree);
        dir = dir * AMOUNT_TO_SHRINK
        newEnd = self.end + dir
        branch = Branch(self.end, newEnd, color)
        return branch

    def displayBranch(self):
        glLineWidth(2.0)
        self.vertices.draw(GL_LINES)


class FractalTree:
    def __init__(self, height):
        self.branches = []
        self.branches.append(Branch([0, -(height / height)], [0, 0], BRANCH_COLOUR))
        self.totalBranchesToVisit = int(pow(2, TREE_DEPTH - 1)) - 1
        self.currBranchIndex = 0

    def createTree(self):
        degree = rand.randrange(30, 61)
        self.branches.append(self.branches[self.currBranchIndex].branch(-degree, BRANCH_COLOUR))
        self.branches.append(self.branches[self.currBranchIndex].branch(degree, BRANCH_COLOUR))
        self.currBranchIndex += 1

    def displayTree(self):
        for branch in self.branches:
            branch.displayBranch()


class Window(pyglet.window.Window):
    def __init__(self, *args, **kwargs):
        if (sys.version_info > (3, 0)):
            super().__init__(*args, **kwargs)
        else:
            super(Window, self).__init__(*args, **kwargs)

        self.set_minimum_size(640, 480)
        glClearColor(0.5, 0.5, 0.5, 1.0)
        glScalef(0.4, 0.4, 0.4)
        glTranslatef(0.0, -1.3, 0.0)

        windowSize = self.get_size()

        self.tree = FractalTree(windowSize[1])  # We want the height of the window

    def on_draw(self):
        self.clear()
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)

        self.tree.displayTree()

    def on_mouse_press(self,x, y, button, modifiers):
        if self.tree.currBranchIndex < self.tree.totalBranchesToVisit:
            self.tree.createTree()

            if self.tree.currBranchIndex == self.tree.totalBranchesToVisit:
                totalBranches = len(self.tree.branches)

                for branchIndex in range(self.tree.currBranchIndex, totalBranches):
                    self.tree.branches[branchIndex].vertices.colors = BRANCH_LEAF_COLUR

    def on_resize(self, width, height):
        glViewport(0, 0, width, height)


if __name__ == "__main__":
    window = Window(640, 480, "Fractal Trees Demonstration", resizable=True)
    pyglet.app.run()