|
"""Export journal to HTML visualization of tree + code.""" |
|
|
|
import json |
|
import textwrap |
|
from pathlib import Path |
|
|
|
import numpy as np |
|
from igraph import Graph |
|
from ..journal import Journal |
|
|
|
|
|
def get_edges(journal: Journal): |
|
for node in journal: |
|
for c in node.children: |
|
yield (node.step, c.step) |
|
|
|
|
|
def generate_layout(n_nodes, edges, layout_type="rt"): |
|
"""Generate visual layout of graph""" |
|
layout = Graph( |
|
n_nodes, |
|
edges=edges, |
|
directed=True, |
|
).layout(layout_type) |
|
y_max = max(layout[k][1] for k in range(n_nodes)) |
|
layout_coords = [] |
|
for n in range(n_nodes): |
|
layout_coords.append((layout[n][0], 2 * y_max - layout[n][1])) |
|
return np.array(layout_coords) |
|
|
|
|
|
def normalize_layout(layout: np.ndarray): |
|
"""Normalize layout to [0, 1]""" |
|
layout = (layout - layout.min(axis=0)) / (layout.max(axis=0) - layout.min(axis=0)) |
|
layout[:, 1] = 1 - layout[:, 1] |
|
layout[:, 1] = np.nan_to_num(layout[:, 1], nan=0) |
|
layout[:, 0] = np.nan_to_num(layout[:, 0], nan=0.5) |
|
return layout |
|
|
|
|
|
def strip_code_markers(code: str) -> str: |
|
"""Remove markdown code block markers if present.""" |
|
code = code.strip() |
|
if code.startswith("```"): |
|
|
|
first_newline = code.find("\n") |
|
if first_newline != -1: |
|
code = code[first_newline:].strip() |
|
if code.endswith("```"): |
|
code = code[:-3].strip() |
|
return code |
|
|
|
|
|
def cfg_to_tree_struct(cfg, jou: Journal): |
|
edges = list(get_edges(jou)) |
|
layout = normalize_layout(generate_layout(len(jou), edges)) |
|
|
|
|
|
|
|
|
|
|
|
metrics = np.array([0 for n in jou]) |
|
|
|
return dict( |
|
edges=edges, |
|
layout=layout.tolist(), |
|
plan=[textwrap.fill(n.plan, width=80) for n in jou.nodes], |
|
code=[strip_code_markers(n.code) for n in jou], |
|
term_out=[n.term_out for n in jou], |
|
analysis=[n.analysis for n in jou], |
|
exp_name=cfg.exp_name, |
|
metrics=metrics.tolist(), |
|
) |
|
|
|
|
|
def generate_html(tree_graph_str: str): |
|
template_dir = Path(__file__).parent / "viz_templates" |
|
|
|
with open(template_dir / "template.js") as f: |
|
js = f.read() |
|
js = js.replace("<placeholder>", tree_graph_str) |
|
|
|
with open(template_dir / "template.html") as f: |
|
html = f.read() |
|
html = html.replace("<!-- placeholder -->", js) |
|
|
|
return html |
|
|
|
|
|
def generate(cfg, jou: Journal, out_path: Path): |
|
tree_graph_str = json.dumps(cfg_to_tree_struct(cfg, jou)) |
|
html = generate_html(tree_graph_str) |
|
with open(out_path, "w") as f: |
|
f.write(html) |
|
|
