{"spec_id":"sankey-basic","library":"bokeh","language":"python","code":"\"\"\" anyplot.ai\nsankey-basic: Basic Sankey Diagram\nLibrary: bokeh 3.9.1 | Python 3.13.14\nQuality: 89/100 | Updated: 2026-07-25\n\"\"\"\n\nimport os\nimport sys\nimport time\nfrom pathlib import Path\n\n\n_script_dir = os.path.dirname(os.path.abspath(__file__))\nsys.path = [p for p in sys.path if os.path.abspath(p or \".\") != _script_dir]\n\nimport numpy as np\nfrom bokeh.io import output_file, save\nfrom bokeh.models import ColumnDataSource, HoverTool, Label\nfrom bokeh.plotting import figure\nfrom selenium import webdriver\nfrom selenium.webdriver.chrome.options import Options\n\n\n# Theme tokens\nTHEME = os.getenv(\"ANYPLOT_THEME\", \"light\")\nPAGE_BG = \"#FAF8F1\" if THEME == \"light\" else \"#1A1A17\"\nINK = \"#1A1A17\" if THEME == \"light\" else \"#F0EFE8\"\nINK_SOFT = \"#4A4A44\" if THEME == \"light\" else \"#B8B7B0\"\nFLOW_ALPHA = 0.45 if THEME == \"light\" else 0.65  # dark bg needs more opacity to keep ribbons visible\n\n# Imprint palette — first source always #009E73\nIMPRINT_PALETTE = [\"#009E73\", \"#C475FD\", \"#4467A3\", \"#BD8233\", \"#AE3030\"]\nNEUTRAL = INK  # theme-adaptive anchor for sector (target) nodes — structural, not categorical\n\n# Data - Energy flow from sources to sectors (TWh)\nflows = [\n    {\"source\": \"Coal\", \"target\": \"Industrial\", \"value\": 25},\n    {\"source\": \"Coal\", \"target\": \"Residential\", \"value\": 10},\n    {\"source\": \"Gas\", \"target\": \"Residential\", \"value\": 30},\n    {\"source\": \"Gas\", \"target\": \"Commercial\", \"value\": 20},\n    {\"source\": \"Gas\", \"target\": \"Industrial\", \"value\": 15},\n    {\"source\": \"Nuclear\", \"target\": \"Industrial\", \"value\": 18},\n    {\"source\": \"Nuclear\", \"target\": \"Commercial\", \"value\": 12},\n    {\"source\": \"Hydro\", \"target\": \"Residential\", \"value\": 8},\n    {\"source\": \"Hydro\", \"target\": \"Commercial\", \"value\": 7},\n    {\"source\": \"Solar\", \"target\": \"Residential\", \"value\": 5},\n    {\"source\": \"Solar\", \"target\": \"Commercial\", \"value\": 6},\n]\n\n# Extract unique sources and targets (preserve order)\nsources = []\ntargets = []\nfor f in flows:\n    if f[\"source\"] not in sources:\n        sources.append(f[\"source\"])\n    if f[\"target\"] not in targets:\n        targets.append(f[\"target\"])\n\n# Source colors: Imprint palette in canonical order (encounter order, independent\n# of the crossing-minimized visual stacking order computed below)\nsource_colors = {s: IMPRINT_PALETTE[i] for i, s in enumerate(sources)}\n\n# Calculate totals for node sizing\nsource_totals = {s: sum(f[\"value\"] for f in flows if f[\"source\"] == s) for s in sources}\ntarget_totals = {t: sum(f[\"value\"] for f in flows if f[\"target\"] == t) for t in targets}\n\n\n# Crossing-minimization: reorder the vertical stacking of source/target nodes via\n# a barycenter heuristic (each node's position converges toward the weighted-average\n# position of the nodes it connects to) so ribbons cross less and the flow reads\n# with a clearer focal point.\ndef _barycenter_reorder(this_side, other_order, connections):\n    other_index = {name: i for i, name in enumerate(other_order)}\n    scores = {}\n    for name in this_side:\n        conns = connections[name]\n        total = sum(v for _, v in conns)\n        scores[name] = sum(other_index[o] * v for o, v in conns) / total if total else other_index.get(name, 0)\n    return sorted(this_side, key=lambda n: scores[n])\n\n\nsource_to_targets = {s: [(f[\"target\"], f[\"value\"]) for f in flows if f[\"source\"] == s] for s in sources}\ntarget_to_sources = {t: [(f[\"source\"], f[\"value\"]) for f in flows if f[\"target\"] == t] for t in targets}\n\nordered_sources = list(sources)\nordered_targets = list(targets)\nfor _ in range(4):\n    ordered_targets = _barycenter_reorder(ordered_targets, ordered_sources, target_to_sources)\n    ordered_sources = _barycenter_reorder(ordered_sources, ordered_targets, source_to_targets)\n\n# Layout parameters (data-space percent units, independent of canvas pixels)\nleft_x = 0\nright_x = 100\nnode_width = 8\nnode_gap = 3\ntotal_height = 100\npadding_y = 5\n\n# Calculate node positions for sources (left side)\nsource_height_total = sum(source_totals.values())\nscale_src = (total_height - 2 * padding_y - (len(sources) - 1) * node_gap) / source_height_total\n\nsource_nodes = {}\ncurrent_y = padding_y\nfor s in ordered_sources:\n    height = source_totals[s] * scale_src\n    source_nodes[s] = {\"x\": left_x, \"y\": current_y, \"height\": height, \"value\": source_totals[s]}\n    current_y += height + node_gap\n\n# Calculate node positions for targets (right side)\ntarget_height_total = sum(target_totals.values())\nscale_tgt = (total_height - 2 * padding_y - (len(targets) - 1) * node_gap) / target_height_total\n\ntarget_nodes = {}\ncurrent_y = padding_y\nfor t in ordered_targets:\n    height = target_totals[t] * scale_tgt\n    target_nodes[t] = {\"x\": right_x - node_width, \"y\": current_y, \"height\": height, \"value\": target_totals[t]}\n    current_y += height + node_gap\n\n# Track flow offsets for stacking flows at each node\nsource_offsets = dict.fromkeys(sources, 0.0)\ntarget_offsets = dict.fromkeys(targets, 0.0)\n\n# Build flow ribbons as bezier patches, collected into a ColumnDataSource so\n# HoverTool can read per-flow source/target/value on mouseover.\nflow_xs, flow_ys, flow_source, flow_target, flow_value, flow_color = [], [], [], [], [], []\nfor f in flows:\n    src = f[\"source\"]\n    tgt = f[\"target\"]\n    value = f[\"value\"]\n\n    src_node = source_nodes[src]\n    tgt_node = target_nodes[tgt]\n\n    src_flow_height = (value / source_totals[src]) * src_node[\"height\"]\n    tgt_flow_height = (value / target_totals[tgt]) * tgt_node[\"height\"]\n\n    x0 = src_node[\"x\"] + node_width\n    y0_bottom = src_node[\"y\"] + source_offsets[src]\n    y0_top = y0_bottom + src_flow_height\n\n    x1 = tgt_node[\"x\"]\n    y1_bottom = tgt_node[\"y\"] + target_offsets[tgt]\n    y1_top = y1_bottom + tgt_flow_height\n\n    source_offsets[src] += src_flow_height\n    target_offsets[tgt] += tgt_flow_height\n\n    t = np.linspace(0, 1, 60)\n    cx0 = x0 + (x1 - x0) * 0.4\n    cx1 = x0 + (x1 - x0) * 0.6\n\n    x_path = (1 - t) ** 3 * x0 + 3 * (1 - t) ** 2 * t * cx0 + 3 * (1 - t) * t**2 * cx1 + t**3 * x1\n    y_bottom = (1 - t) * y0_bottom + t * y1_bottom\n    y_top = (1 - t) * y0_top + t * y1_top\n\n    flow_xs.append(list(x_path) + list(x_path[::-1]))\n    flow_ys.append(list(y_top) + list(y_bottom[::-1]))\n    flow_source.append(src)\n    flow_target.append(tgt)\n    flow_value.append(value)\n    flow_color.append(source_colors[src])\n\nflow_cds = ColumnDataSource(\n    data={\n        \"xs\": flow_xs,\n        \"ys\": flow_ys,\n        \"flow_source\": flow_source,\n        \"flow_target\": flow_target,\n        \"flow_value\": flow_value,\n        \"flow_color\": flow_color,\n    }\n)\n\n# Nodes (sources + sectors) as a single ColumnDataSource for hover + rendering\nnode_name = list(ordered_sources) + list(ordered_targets)\nnode_role = [\"Source\"] * len(ordered_sources) + [\"Sector\"] * len(ordered_targets)\nnode_left = [source_nodes[s][\"x\"] for s in ordered_sources] + [target_nodes[t][\"x\"] for t in ordered_targets]\nnode_right = [source_nodes[s][\"x\"] + node_width for s in ordered_sources] + [\n    target_nodes[t][\"x\"] + node_width for t in ordered_targets\n]\nnode_bottom = [source_nodes[s][\"y\"] for s in ordered_sources] + [target_nodes[t][\"y\"] for t in ordered_targets]\nnode_top = [source_nodes[s][\"y\"] + source_nodes[s][\"height\"] for s in ordered_sources] + [\n    target_nodes[t][\"y\"] + target_nodes[t][\"height\"] for t in ordered_targets\n]\nnode_value = [source_nodes[s][\"value\"] for s in ordered_sources] + [target_nodes[t][\"value\"] for t in ordered_targets]\nnode_color = [source_colors[s] for s in ordered_sources] + [NEUTRAL] * len(ordered_targets)\n\nnodes_cds = ColumnDataSource(\n    data={\n        \"name\": node_name,\n        \"role\": node_role,\n        \"left\": node_left,\n        \"right\": node_right,\n        \"bottom\": node_bottom,\n        \"top\": node_top,\n        \"value\": node_value,\n        \"color\": node_color,\n    }\n)\n\n# Plot — canonical 3200x1800 landscape canvas\np = figure(\n    width=3200,\n    height=1800,\n    title=\"sankey-basic · python · bokeh · anyplot.ai\",\n    # Generous L/R x_range margin — Label text clips at the frame/range boundary,\n    # not the canvas edge, so overflow room must live in x_range, not min_border_*.\n    x_range=(-40, 150),\n    y_range=(-4, 100),\n    tools=\"\",\n    toolbar_location=None,  # bokeh's default toolbar adds ~30-50px above the canvas\n    min_border_bottom=60,\n    min_border_left=60,\n    min_border_top=110,\n    min_border_right=60,\n)\n\nflow_renderer = p.patches(\n    \"xs\",\n    \"ys\",\n    source=flow_cds,\n    fill_color=\"flow_color\",\n    fill_alpha=FLOW_ALPHA,\n    line_color=\"flow_color\",\n    # More opaque, slightly thicker stroke than the fill so a ribbon's own edge\n    # stays traceable through alpha-blended crossings instead of dissolving into\n    # a blended hue.\n    line_alpha=0.9,\n    line_width=1.5,\n)\n\nnode_renderer = p.quad(\n    left=\"left\",\n    right=\"right\",\n    bottom=\"bottom\",\n    top=\"top\",\n    source=nodes_cds,\n    fill_color=\"color\",\n    fill_alpha=0.92,\n    line_color=PAGE_BG,\n    line_width=2,\n)\n\np.add_tools(\n    HoverTool(\n        renderers=[flow_renderer], tooltips=[(\"Flow\", \"@flow_source → @flow_target\"), (\"Volume\", \"@flow_value TWh\")]\n    )\n)\np.add_tools(HoverTool(renderers=[node_renderer], tooltips=[(\"Node\", \"@name (@role)\"), (\"Total\", \"@value TWh\")]))\n\n# Node labels — source nodes left-aligned outward, target nodes right-aligned outward\nfor s in ordered_sources:\n    node = source_nodes[s]\n    label = Label(\n        x=node[\"x\"] - 1.5,\n        y=node[\"y\"] + node[\"height\"] / 2,\n        text=f\"{s} ({node['value']} TWh)\",\n        text_font_size=\"26pt\",\n        text_align=\"right\",\n        text_baseline=\"middle\",\n        text_color=INK,\n        text_font=\"helvetica\",\n    )\n    p.add_layout(label)\n\nfor t in ordered_targets:\n    node = target_nodes[t]\n    label = Label(\n        x=node[\"x\"] + node_width + 1.5,\n        y=node[\"y\"] + node[\"height\"] / 2,\n        text=f\"{t} ({node['value']} TWh)\",\n        text_font_size=\"26pt\",\n        text_align=\"left\",\n        text_baseline=\"middle\",\n        text_color=INK,\n        text_font=\"helvetica\",\n    )\n    p.add_layout(label)\n\n# Style — theme-adaptive chrome\np.title.text_font_size = \"50pt\"\np.title.text_color = INK\np.title.align = \"center\"\np.title.text_font = \"helvetica\"\n\np.xaxis.visible = False\np.yaxis.visible = False\np.xgrid.visible = False\np.ygrid.visible = False\np.outline_line_color = None\n\np.background_fill_color = PAGE_BG\np.border_fill_color = PAGE_BG\n\n# Save — write the interactive HTML, then screenshot it with headless Chrome.\n# bokeh.io.export_png is avoided here (unreliable chromedriver resolution);\n# Selenium + CDP viewport pinning matches the exact 3200x1800 canvas.\noutput_file(f\"plot-{THEME}.html\")\nsave(p)\n\nW, H = 3200, 1800\nopts = Options()\nfor arg in (\n    \"--headless=new\",\n    \"--no-sandbox\",\n    \"--disable-dev-shm-usage\",\n    \"--disable-gpu\",\n    f\"--window-size={W},{H}\",\n    \"--hide-scrollbars\",\n):\n    opts.add_argument(arg)\ndriver = webdriver.Chrome(options=opts)\ndriver.set_window_size(W, H)\ndriver.get(f\"file://{Path(f'plot-{THEME}.html').resolve()}\")\ndriver.execute_cdp_cmd(\n    \"Emulation.setDeviceMetricsOverride\", {\"width\": W, \"height\": H, \"deviceScaleFactor\": 1, \"mobile\": False}\n)\n# Zero out the default page margin/background so no stray edge pixel of the\n# browser's default white page bleeds through around the themed canvas.\ndriver.execute_script(\n    f\"document.documentElement.style.background='{PAGE_BG}';\"\n    f\"document.body.style.background='{PAGE_BG}';\"\n    \"document.body.style.margin='0';\"\n    \"document.body.style.overflow='hidden';\"\n)\ntime.sleep(3)  # let bokeh's JS render the canvas\ndriver.save_screenshot(f\"plot-{THEME}.png\")\ndriver.quit()\n"}