Pythagoras' Tree

Idea from Mathologer's video about Fibonacci numbers and Pythagoras' theorem.

Libraries and utility function

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle
%matplotlib inline
import ipywidgets as widgets


def progress_bar(current, total, bar_length=20):
    fraction = current / total

    arrow = int(fraction * bar_length - 1) * '-' + '>'
    padding = int(bar_length - len(arrow)) * ' '

    ending = '\n' if current == total else '\r'

    print(f'Progress: [{arrow}{padding}] {int(fraction*100)}%', end=ending)  

Tree nodes and builder

class Node():
    def __init__(
        self, parent: Node, left: Node, right: Node,
        depth: int, pos: np.ndarray(2), angle: float, size: float
    ):
        # pointers to other nodes
        self.parent = parent
        self.left = left
        self.right = right
        # plotting parameters
        self.p = pos
        self.a = angle #in radians
        self.l = size
        # tree parameters
        self.d = depth
    
    def move_left(self, d_angle: float):
        if self.d % 2 == 0:
            dp = self.parent.l * np.array([-np.sin(self.parent.a), np.cos(self.parent.a)])
        else:
            dp = self.parent.l * np.array([-np.sin(self.parent.a), np.cos(self.parent.a)])
        self.p += dp
        self.a += d_angle
        return
    
    def square(self):
        return Rectangle( (self.p[0],self.p[1]), self.l, self.l, angle=self.a*180/np.pi )

    def mapColor(self, N):
        return (self.d/N)**2

#####################################à######################################################

def build_tree(N, angle):
    root = Node(None,None,None, 0, np.zeros(2), 0.0, 1.0)
    current = root
    points = [current.p]
    squares = [current.square()]
    colors = [current.mapColor(N)]
    total_nodes = 2**N -1
    count = 1
    while(current.left is None or current.right is None or current.parent is not None):
        # if maximum depth is reached, go back
        if current.d==N-1:
            current = current.parent
            continue
        # try to go deeper, left first...
        if current.left is None:
            # every two depths, switch the symmetry
            if current.d%2 == 0:
                dp = current.l * np.array([-np.sin(current.a), np.cos(current.a)])
                a = angle
            else:
                dp = current.l * np.array([-np.sin(current.a), np.cos(current.a)])
                a = np.pi/2 - angle
            current.left = Node(current,None,None,
                                current.d+1,
                                current.p+dp,
                                current.a+a,
                                current.l * np.cos(a)
                               )
            current = current.left
            points.append(current.p)
            squares.append( current.square() )
            colors.append( current.mapColor(N) )
            count += 1
            progress_bar(count,total_nodes)
        # ... else go right...
        elif current.right is None:
            # every two depths, switch the symmetry
            if current.d%2 == 0:
                a = angle
                dp = current.l * np.array([-np.sin(current.a), np.cos(current.a)])
                dp+= current.l * np.cos(a) * np.array([-np.sin(current.a -np.pi/2+a), np.cos(current.a -np.pi/2+a)])
            else:
                a = np.pi/2 - angle
                dp = current.l * np.array([-np.sin(current.a), np.cos(current.a)])
                dp+= current.l * np.cos(a) * np.array([-np.sin(current.a -np.pi/2+a), np.cos(current.a -np.pi/2+a)])
            current.right = Node(current,None,None,
                                current.d+1,
                                current.p + dp,
                                current.a -np.pi/2+a,
                                current.l * np.sin(a)
                                )
            current = current.right
            points.append(current.p)
            squares.append( current.square() )
            colors.append( current.mapColor(N) )
            count += 1
            progress_bar(count,total_nodes)
        # ... if right was already visited too, go back (if possible)
        elif current.parent is not None:
            current = current.parent
    return squares, colors, points

Rendering

def render(N=5, angle = np.pi/6, colormap='winter'):
    # build the tree
    squares, colors, points = build_tree(N, angle)
    
    pc = PatchCollection(squares, cmap=plt.colormaps[colormap], alpha=0.7)
    pc.set_array(colors)
    
    fig,ax = plt.subplots(figsize=(8,8))
    ax.add_collection(pc)
    ax.plot([],[]) #empty plot
    #ax.plot([p[0] for  p in points], [p[1] for p in points],'r.')
    #ax.set_title(f"Pythagoras' Tree\n with angle $\\alpha_0$={angle}, N={N} iterations")
    ax.axis('equal')
    ax.axis('off')
    plt.tight_layout()
    #plt.savefig('tree.png')
    #plt.savefig('tree.eps')
    plt.show()

dict_args = {
    "N": widgets.IntSlider(min=1, max=16, step=1, value=10),
    "angle": widgets.FloatSlider(min=0.01, max=np.pi/2, step=.01, value=0.5),
    "colormap": widgets.Dropdown(options=['winter','cool','coolwarm','RdYlGn','twilight_shifted','hsv','rainbow','gist_earth'], value='winter')
}

widgets.interact_manual(render, **dict_args);