treemap_org.py

"""
Plot a treemap corresponding to an org-mode file with weight being the log of number of characters, and the body of a node displayed on hover.

Dependencies: numpy, pandas, plotly, orgparse

Usage

python treemap_org.py file.org

"""
import argparse
from pathlib import Path

import numpy as np
import pandas as pd

import plotly
import plotly.express as px

from orgparse import load, loads

parser = argparse.ArgumentParser()   
parser.add_argument('file', type=str, nargs='+')
args = parser.parse_args()
# Merge org files if there are several
res = []
if len(args.file)>1:
    for f in args.file:
        name = Path(f).stem
        res.append("* "+name)
        with open(f, "r") as fname:
            text = fname.read()
            text = text.split("\n")
            for line in text:
                if len(line) > 0:
                    if line[0]=="#":
                        pass
                    elif line[0] == '*':
                        res.append("*"+line)
                    else:
                        res.append(line)
else:
    with open(args.file[0], "r") as fname:
        text = fname.read().split("\n")
        for line in text:
            if len(line) > 0:
                if line[0]=="#":
                    pass
                else:
                    res.append(line)

# processing of the tree into something readable by plotly treemaps
tree = loads("\n".join(res))

def get_value(node):
    if len(node.children) == 0:
        return len(node.body)
    else:
        res = len(node.body)
        for child in node.children:
             res += get_value(child)
        return res


node_to_str = lambda node : node.heading + " (" + str(get_value(node))+ ")"
    
values = [get_value(node) for node in tree]
parents = ['']+[ node_to_str(node.get_parent()) for node in tree[1:]]
names = [node_to_str(node) for i,node in enumerate(tree)]
descriptions = [node.body for i,node in enumerate(tree)]


df = pd.DataFrame({"names":names, "parents":parents, "values" : values, "text":descriptions})
df = df.sort_values("values", ascending=False).reset_index(drop=True)

# Uncomment next two lines to limit at 50 cases
# n_nodes = 50
# df = df.iloc[:n_nodes]

df["values"] = np.log(1+df["values"])

# Plot the treemap

fig = px.treemap(
    df,
    names = "names",
    parents = "parents",
    values = "values",
    color_discrete_sequence=px.colors.qualitative.Set2,
)
fig.update_traces(hoverinfo = "label+text",
                  hovertext = df.text.str.wrap(60).apply(lambda x: x.replace('\n', '<br>')), # line of 60 characters for readability
                  root_color="lightgrey")
fig.update_traces(hovertemplate='label=%{label}<br>text=%{hovertext}<extra></extra>',)

fig.update_layout(margin = dict(t=50, l=25, r=25, b=25))

# Uncomment to save to html
# plotly.offline.plot(fig, filename='treemap.html')
fig.show()