Create a TMAP from a scipy hierarchy

dendrogram.py

import tmap as tm
import numpy as np
from scipy.cluster import hierarchy
from matplotlib import pyplot as plt

N = 1000
CFG = tm.LayoutConfiguration()
# This is the most important parameter to tune when changing the
# number of nodes. More nodes -> smaller number
CFG.node_size = 1 / 10


def get_random_tree():
    x = np.random.rand(N).reshape(int(N / 2), 2)
    Z = hierarchy.linkage(x)
    hierarchy.to_tree(Z)

    return hierarchy.to_tree(Z, rd=True)


def edge_list_from_tree(node, edge_list, parent_id=None):
    if node.left:
        edge_list_from_tree(node.left, edge_list, node.id)
    if node.right:
        edge_list_from_tree(node.right, edge_list, node.id)

    if parent_id:
        edge_list.append([node.id, parent_id, 1.0])


def main():
    root, node_list = get_random_tree()
    edge_list = []
    edge_list_from_tree(root, edge_list)

    x, y, s, t, _ = tm.layout_from_edge_list(
        N - 1, edge_list, config=CFG, create_mst=False
    )

    # Plot the edges
    for i in range(len(s)):
        plt.plot(
            [x[s[i]], x[t[i]]],
            [y[s[i]], y[t[i]]],
            "k-",
            linewidth=0.5,
            alpha=0.5,
            zorder=1,
        )

    # Plot the vertices
    plt.scatter(x, y, c=[n.dist for n in node_list], s=1, cmap="rainbow_r", zorder=2)
    plt.tight_layout()
    plt.savefig("graph.png", dpi=600)


if __name__ == "__main__":
    main()