build_graph_d3js.py

import json # export data for d3 visualization
import numpy as np
import scipy.cluster.hierarchy as hclust

## Compute histogram based on the Wasserstein distance matrix

def build_json_tree(node, parent):
    """
    Build a tree for the json format
    """
    new_node = dict( node_id = node.id, children = [] )
    parent["children"].append( new_node )

    if node.left:
        build_json_tree(node.left, new_node)
    if node.right:
        build_json_tree(node.right, new_node)

def decorate_json_tree( node, node_names ):
    """
    Modify the json tree to use for D3js
    """
    if len(node["children"]) == 0:
        try:
            leaf_names = [ node_names[ node['node_id'] ] ]
        except KeyError:
            print("Warning: KeyError", node['node_id'])
            pass
    else:
        children_names = reduce(lambda ls, c: np.append(ls, decorate_json_tree(c, node_names)),
                                np.asarray(node['children']), np.empty(0))
        leaf_names = []
        # leaf_names = np.unique( children_names )
        # if len(leaf_names) > 2 or ("internal" in leaf_names):
        #     leaf_names = ["internal"]

    del node["node_id"]
    node["name"] = name = '-'.join(sorted(map(str, leaf_names)))

    return leaf_names


def hierarchical_clustering(D, kls, linkage_method = 'centroid', title = None):
    """
    Perform hierarchical clustering and
    export the result to a json file

    \param D :- NxN distance matrix
    """
    Z = hclust.linkage( D, method = linkage_method )

    ## Transform the tree file to json format and display it with D3js
    T = hclust.to_tree(Z, rd = False)
    d3_json_tree = dict(children = [], name = "root")
    build_json_tree(T, d3_json_tree)
    decorate_json_tree( d3_json_tree["children"][0], kls )

    if title:
        out_fname = '_'.join(["kimia", title, "dendrogram", linkage_method]) + ".json"
    else:
        out_fname = '_'.join(["kimia", "dendrogram", linkage_method]) + ".json"
    json.dump(d3_json_tree, open(out_fname, "w"), sort_keys = True)