A decision tree for a superflous task in theoretical physics II

theo-annoying-exercise-decision-tree.py

# Written by Splines, https://github.com/splines
# Change WITH_PUTBACKS to True or False to change if we allow putbacks or not

# In order to run this, you need to install Graphviz first,
#  see the installation guide here:
# https://graphviz.readthedocs.io/en/stable/#installation


from typing import List
import numpy as np
import graphviz
from fractions import Fraction

# Params
balls = ['A', 'B', 'C', 'D']
occurrences = [2, 1, 2, 3]
NUM_LEVELS = 1
WITH_PUTBACKS = True
filename = f'decision-tree-{"with" if WITH_PUTBACKS else "without"}-putbacks-{NUM_LEVELS}-levels'

# for checking purposes
all_last_level_probabilities = []

# Init graph
G = graphviz.Digraph('decision-tree', filename=filename)
G.graph_attr["rankdir"] = "LR"  # horizontal tree layout
# A4 paper size (at least try to fit width)
G.graph_attr["size"] = "8.3!,11.7!"
G.graph_attr["margin"] = "0"  # no margin


def construct_last_level_node(parent_node, probability):
    # Change style for last element (probability of chained event)
    G.attr('node', shape='diamond', style='filled', color='lightgrey')

    # Construct
    last_node = f'{parent_node}-final'
    G.node(last_node, label=f'{probability}')
    G.edge(parent_node, last_node)

    # Reset style
    G.attr('node', shape='ellipse', style='', color='')

    all_last_level_probabilities.append(probability)


def probability_for_ball(i, current_occurrences: List[int]) -> Fraction:
    return Fraction(current_occurrences[i], sum(current_occurrences))


def construct_next_level_from_one_ball(next_level, parent_node,
                                       i, ball, probability,
                                       current_occurrences: List[int]):
    # Probability
    probability_for_current_ball = probability_for_ball(i, current_occurrences)

    # Construct node in next level and edge to that node
    next_node = f'{parent_node}-{ball}'
    # G.node(next_node, label=f'{ball}')
    # G.edge(parent_node, next_node, label=f'{probability_for_current_ball}')

    # Adjust occurrences of elements
    modified_occurrences = current_occurrences
    if not WITH_PUTBACKS:
        modified_occurrences = current_occurrences.copy()
        modified_occurrences[i] -= 1

    # New probability to pass to next level
    new_probability = probability * probability_for_current_ball

    # Step down recursively
    next_recursive(next_level + 1, next_node,
                   modified_occurrences, new_probability)


def next_recursive(next_level, parent_node, current_occurrences: List[int], probability=1):
    # Recursion anchor (last node)
    if next_level == NUM_LEVELS:
        construct_last_level_node(parent_node, probability)
        return

    for i, ball in enumerate(balls):
        # Only for balls that are still available
        if not current_occurrences[i] > 0:
            continue
        construct_next_level_from_one_ball(
            next_level, parent_node, i, ball, probability, current_occurrences)


# Start recursion
next_recursive(0, 'start', occurrences)

# Check Kolmogorov
assert(sum(all_last_level_probabilities) == 1)

# Save graph as PDF
G.render(directory='./').replace('\\', '/')

# Calculate entropy
entropy = - sum([p * np.log(float(p)) for p in all_last_level_probabilities])
print('Num probabilities:', len(all_last_level_probabilities))
for p in all_last_level_probabilities:
    print(f'{p} * ln({p})')
print(f'Entropy S = {entropy}')

See the results:

• with putbacks
• without putbacks