# Source code for higra.hierarchy.random_hierarchy

############################################################################
# Copyright ESIEE Paris (2018)                                             #
#                                                                          #
# Contributor(s) : Benjamin Perret                                         #
#                                                                          #
#                                                                          #
# The full license is in the file LICENSE, distributed with this software. #
############################################################################

import higra as hg
import numpy as np

[docs]def random_binary_partition_tree(num_leaves, asymmetry_probability):
"""
Random binary partition tree with a controlled amount of asymmetry/unbalancedness.

The tree is grown from the root to the leaves.
At each step, the algorithm randomly select one of the *growable* leaf node of the current tree.
Two children are added to the selected node; the number of leaf nodes is hence increased by one.
Then,

- with probability :math:1-asymmetry\_probability, both new children are marked as *growable*
- with probability :math:asymmetry\_probability, only one of the children is marked as *growable*

The altitudes of the returned hierarchy are obtained with :func:~higra.attribute_regular_altitudes:
*The regular altitudes is comprised between 0 and 1 and is inversely proportional to the depth of a node*.

A valid minimal connected graph (a tree) is associated to the leaves of the tree.

:param num_leaves: expected number of leaves in the generated tree
:param asymmetry_probability: real value between 0 and 1. At 0 the tree is perfectly unbalanced, at 1 it is
perfectly balanced (if :attr:num_leaves is  a power of 2)
:return: a tree (Concept :class:~higra.CptBinaryHierarchy) and its node altitudes
"""
import random
import math

assert (0 <= asymmetry_probability <= 1)
num_leaves = int(num_leaves)
assert (num_leaves > 0)

parents = np.zeros((num_leaves * 2 - 1,), dtype=np.int64)

n = 1
root = {}
leaves = []
leaves.append(root)

all_nodes = [root]

i = parents.size - 1
root["parent"] = i

while n != 2 * num_leaves - 1:

ni = random.randint(0, math.floor(asymmetry_probability * (len(leaves) - 1)))
node = leaves[ni]
del leaves[ni]

node["i"] = i
node["left"] = {"parent": i}
node["right"] = {"parent": i}
i -= 1
all_nodes.append(node["left"])
all_nodes.append(node["right"])
n += 2

if random.random() <= asymmetry_probability:
if random.random() >= 0.5:
leaves.append(node["right"])
else:
leaves.append(node["left"])
else:
leaves.append(node["left"])
leaves.append(node["right"])

k = 0
for node in all_nodes:
if "i" not in node:
node["i"] = k
k += 1
parents[node["i"]] = node["parent"]

tree = hg.Tree(parents)

altitudes = hg.attribute_regular_altitudes(tree)

def _get_associated_mst(tree, altitudes):
"""
Create a valid edge mst for the given tree (returns an edge weighted undirected graph)
"""
nb = tree.num_leaves()

g = hg.UndirectedGraph(nb)
edge_weights = np.zeros((nb - 1,), np.float32)
for r in tree.leaves_to_root_iterator(include_leaves=False):