<jupyter_start><jupyter_text>Import required modules<jupyter_code>import sys
sys.path.append('..')
import numpy as np
import ot
import networkx as nx
import tw<jupyter_output><empty_output><jupyter_text>Generate a tree<jupyter_code>num_node = 1000
G = nx.generators.random_tree(num_node)
edges = [(fr, to, 1)
       for (fr, to) in list(G.edges())]<jupyter_output><empty_output><jupyter_text>Set two probabilities on nodes of the tree<jupyter_code>first_prob = np.random.rand(num_node)
first_prob = first_prob / first_prob.sum()
second_prob = np.random.rand(num_node)
second_prob = second_prob / second_prob.sum()<jupyter_output><empty_output><jupyter_text>Build the distance matrix of the tree metric<jupyter_code>adj_dict = dict(nx.all_pairs_shortest_path_length(G))
metric = np.array([[adj_dict[i][j]
                  for i in range(num_node)] for j in range(num_node)])<jupyter_output><empty_output><jupyter_text>Measure the time to calculate the wasserstein distance with our module<jupyter_code>%%time
tw.distance(first_prob, second_prob, edges)<jupyter_output>CPU times: user 6.35 ms, sys: 339 µs, total: 6.69 ms
Wall time: 8.54 ms<jupyter_text>Measure the time to calculate the wasserstein distance with [POT](https://pot.readthedocs.io/en/stable/) module<jupyter_code>%%time
ot.lp.emd2(first_prob, second_prob, metric)<jupyter_output>CPU times: user 107 ms, sys: 24.7 ms, total: 132 ms
Wall time: 148 ms
{'fraction_non_alphanumeric': 0.11323425336164189, 'fraction_numerical': 0.021231422505307854, 'mean_word_length': 6.602150537634409, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 23, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '198618', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 545, 'n_tokens_neox': 512, 'n_words': 119, 'path': 'examples/example.ipynb', 'repo_name': 'InkToYou/TreeWasserstein', 'size': '1413'}