Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
172 changes: 65 additions & 107 deletions plots/sankey-basic/implementations/python/highcharts.py
Original file line number Diff line number Diff line change
@@ -1,152 +1,129 @@
""" pyplots.ai
""" anyplot.ai
sankey-basic: Basic Sankey Diagram
Library: highcharts unknown | Python 3.13.11
Quality: 91/100 | Created: 2025-12-23
Library: highcharts unknown | Python 3.13.13
Quality: 88/100 | Updated: 2026-04-30
"""

import os
import tempfile
import time
import urllib.request
from pathlib import Path

import numpy as np
from highcharts_core.chart import Chart
from highcharts_core.options import HighchartsOptions
from PIL import Image
from selenium import webdriver
from selenium.webdriver.chrome.options import Options


# Reproducibility
np.random.seed(42)
# Theme tokens
THEME = os.getenv("ANYPLOT_THEME", "light")
PAGE_BG = "#FAF8F1" if THEME == "light" else "#1A1A17"
INK = "#1A1A17" if THEME == "light" else "#F0EFE8"
INK_SOFT = "#4A4A44" if THEME == "light" else "#B8B7B0"

# Data - Energy flow from sources to sectors (values in TWh - Terawatt-hours)
# Format: [source, target, value]
# Data - U.S. Energy flow from sources to sectors (values in TWh)
flows = [
# Coal flows
["Coal", "Electricity", 150],
["Coal", "Industrial", 80],
# Natural Gas flows
["Natural Gas", "Electricity", 120],
["Natural Gas", "Residential", 90],
["Natural Gas", "Commercial", 60],
["Natural Gas", "Industrial", 50],
# Nuclear flows
["Nuclear", "Electricity", 200],
# Petroleum flows
["Petroleum", "Transportation", 280],
["Petroleum", "Industrial", 70],
["Petroleum", "Residential", 30],
# Renewable flows
["Renewable", "Electricity", 100],
["Renewable", "Transportation", 20],
# Electricity flows to end uses
["Electricity", "Residential", 180],
["Electricity", "Commercial", 160],
["Electricity", "Industrial", 140],
]

# Collect unique nodes
nodes_set = set()
for source, target, _ in flows:
nodes_set.add(source)
nodes_set.add(target)
nodes = list(nodes_set)
# Node colors - Okabe-Ito palette (canonical order for source nodes)
OKABE_ITO = ["#009E73", "#D55E00", "#0072B2", "#CC79A7", "#E69F00", "#56B4E9", "#F0E442"]

# Colorblind-safe colors for nodes - all dark enough for white text contrast
node_colors = {
# Sources (energy sources) - dark tones for white text readability
"Coal": "#1A3A5C", # Dark Blue
"Natural Gas": "#6B4E12", # Darker Goldenrod (darkened for contrast)
"Nuclear": "#5B2E8F", # Dark Purple
"Petroleum": "#0A6B78", # Dark Cyan
"Renewable": "#155415", # Dark Green
# Intermediate - dark for contrast
"Electricity": "#4D2A22", # Dark Brown
# End uses - darker shades for white text visibility
"Residential": "#8B3A6B", # Dark Rose
"Commercial": "#3A3A3A", # Darker Gray
"Industrial": "#5B5C0A", # Dark Olive
"Transportation": "#994D00", # Darker Orange
"Coal": OKABE_ITO[0],
"Natural Gas": OKABE_ITO[1],
"Nuclear": OKABE_ITO[2],
"Petroleum": OKABE_ITO[3],
"Renewable": OKABE_ITO[4],
"Electricity": OKABE_ITO[5],
"Transportation": OKABE_ITO[6],
"Industrial": OKABE_ITO[0],
"Residential": OKABE_ITO[1],
"Commercial": OKABE_ITO[2],
}

# Create nodes data with colors
nodes_data = [{"id": node, "name": node, "color": node_colors.get(node, "#306998")} for node in nodes]
nodes_set = set()
for source, target, _ in flows:
nodes_set.add(source)
nodes_set.add(target)

# Create links data
nodes_data = [{"id": node, "name": node, "color": node_colors.get(node, OKABE_ITO[0])} for node in nodes_set]
links_data = [{"from": source, "to": target, "weight": value} for source, target, value in flows]

# Create chart
# Chart
chart = Chart(container="container")
chart.options = HighchartsOptions()

# Chart configuration with margins to prevent label cutoff at edges
chart.options.chart = {
"type": "sankey",
"width": 4800,
"height": 2700,
"backgroundColor": "#ffffff",
"backgroundColor": PAGE_BG,
"marginLeft": 180,
"marginRight": 180,
"marginTop": 160,
"marginBottom": 80,
"marginBottom": 160,
}

# Title
chart.options.title = {
"text": "sankey-basic · highcharts · pyplots.ai",
"style": {"fontSize": "64px", "fontWeight": "bold", "color": "#333333"},
"text": "sankey-basic · highcharts · anyplot.ai",
"style": {"fontSize": "64px", "fontWeight": "bold", "color": INK},
}

# Subtitle with units info
chart.options.subtitle = {"text": "U.S. Energy Flow (values in TWh)", "style": {"fontSize": "40px", "color": "#666666"}}
chart.options.subtitle = {"text": "U.S. Energy Flow (values in TWh)", "style": {"fontSize": "40px", "color": INK_SOFT}}

# Tooltip with units
chart.options.tooltip = {
"style": {"fontSize": "36px"},
"nodeFormat": "{point.name}: {point.sum} TWh",
"pointFormat": "{point.fromNode.name} → {point.toNode.name}: {point.weight} TWh",
}

# Sankey series configuration
series_config = {
"type": "sankey",
"name": "Energy Flow",
"keys": ["from", "to", "weight"],
"nodes": nodes_data,
"data": links_data,
"dataLabels": {
"enabled": True,
"style": {"fontSize": "36px", "fontWeight": "bold", "color": "#FFFFFF", "textOutline": "3px #333333"},
"nodeFormat": "{point.name}",
},
"nodeWidth": 50,
"nodePadding": 35,
"linkOpacity": 0.5,
"curveFactor": 0.5,
"colorByPoint": True,
"linkColorMode": "from",
}

chart.options.series = [series_config]
chart.options.series = [
{
"type": "sankey",
"name": "Energy Flow",
"keys": ["from", "to", "weight"],
"nodes": nodes_data,
"data": links_data,
"dataLabels": {
"enabled": True,
"style": {"fontSize": "36px", "fontWeight": "bold", "color": "#FFFFFF", "textOutline": "3px #333333"},
"nodeFormat": "{point.name}",
},
"nodeWidth": 50,
"nodePadding": 35,
"linkOpacity": 0.5,
"curveFactor": 0.5,
"colorByPoint": True,
"linkColorMode": "from",
}
]

# Disable legend for sankey (nodes are labeled)
chart.options.legend = {"enabled": False}

# Disable credits
chart.options.credits = {"enabled": False}

# Download Highcharts JS and sankey module
highcharts_url = "https://code.highcharts.com/highcharts.js"
sankey_url = "https://code.highcharts.com/modules/sankey.js"

with urllib.request.urlopen(highcharts_url, timeout=30) as response:
# Download Highcharts JS and sankey module inline (required for headless Chrome)
with urllib.request.urlopen("https://cdn.jsdelivr.net/npm/highcharts@latest/highcharts.js", timeout=30) as response:
highcharts_js = response.read().decode("utf-8")

with urllib.request.urlopen(sankey_url, timeout=30) as response:
with urllib.request.urlopen("https://cdn.jsdelivr.net/npm/highcharts@latest/modules/sankey.js", timeout=30) as response:
sankey_js = response.read().decode("utf-8")

# Generate HTML with inline scripts
html_str = chart.to_js_literal()
html_content = f"""<!DOCTYPE html>
<html>
Expand All @@ -155,30 +132,17 @@
<script>{highcharts_js}</script>
<script>{sankey_js}</script>
</head>
<body style="margin:0;">
<div id="container" style="width: 4800px; height: 2700px;"></div>
<script>{html_str}</script>
</body>
</html>"""

# Save HTML for interactive version (use CDN for standalone) - fixed dimensions like PNG
standalone_html = f"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<script src="https://code.highcharts.com/highcharts.js"></script>
<script src="https://code.highcharts.com/modules/sankey.js"></script>
</head>
<body style="margin:0; overflow:auto;">
<body style="margin:0; background:{PAGE_BG};">
<div id="container" style="width: 4800px; height: 2700px;"></div>
<script>{html_str}</script>
</body>
</html>"""

with open("plot.html", "w", encoding="utf-8") as f:
f.write(standalone_html)
# Save HTML artifact
with open(f"plot-{THEME}.html", "w", encoding="utf-8") as f:
f.write(html_content)

# Write temp HTML and take screenshot
# Screenshot via headless Chrome
with tempfile.NamedTemporaryFile(mode="w", suffix=".html", delete=False, encoding="utf-8") as f:
f.write(html_content)
temp_path = f.name
Expand All @@ -188,18 +152,12 @@
chrome_options.add_argument("--no-sandbox")
chrome_options.add_argument("--disable-dev-shm-usage")
chrome_options.add_argument("--disable-gpu")
chrome_options.add_argument("--window-size=4800,2900")
chrome_options.add_argument("--window-size=4800,2700")

driver = webdriver.Chrome(options=chrome_options)
driver.get(f"file://{temp_path}")
time.sleep(5) # Wait for chart to render
driver.save_screenshot("plot_raw.png")
time.sleep(5)
driver.save_screenshot(f"plot-{THEME}.png")
driver.quit()

# Crop to exact 4800x2700 dimensions
img = Image.open("plot_raw.png")
img_cropped = img.crop((0, 0, 4800, 2700))
img_cropped.save("plot.png")
Path("plot_raw.png").unlink()

Path(temp_path).unlink() # Clean up temp file
Path(temp_path).unlink()
Loading
Loading