|
| 1 | +""" anyplot.ai |
| 2 | +sankey-basic: Basic Sankey Diagram |
| 3 | +Library: pygal 3.1.0 | Python 3.13.13 |
| 4 | +Quality: 85/100 | Created: 2026-04-30 |
| 5 | +""" |
| 6 | + |
| 7 | +import os |
| 8 | +import sys |
| 9 | + |
| 10 | + |
| 11 | +# Pop script dir so this file (pygal.py) doesn't shadow the installed pygal package |
| 12 | +_script_dir = sys.path.pop(0) |
| 13 | +import cairosvg # noqa: E402 |
| 14 | +from pygal.style import Style # noqa: E402 |
| 15 | + |
| 16 | + |
| 17 | +sys.path.insert(0, _script_dir) |
| 18 | + |
| 19 | +# Theme tokens |
| 20 | +THEME = os.getenv("ANYPLOT_THEME", "light") |
| 21 | +PAGE_BG = "#FAF8F1" if THEME == "light" else "#1A1A17" |
| 22 | +INK = "#1A1A17" if THEME == "light" else "#F0EFE8" |
| 23 | +INK_SOFT = "#4A4A44" if THEME == "light" else "#B8B7B0" |
| 24 | +INK_MUTED = "#6B6A63" if THEME == "light" else "#A8A79F" |
| 25 | + |
| 26 | +OKABE_ITO = ("#009E73", "#D55E00", "#0072B2", "#CC79A7", "#E69F00", "#56B4E9", "#F0E442") |
| 27 | + |
| 28 | +# pygal Style is the single source of truth for all visual properties |
| 29 | +chart_style = Style( |
| 30 | + background=PAGE_BG, |
| 31 | + plot_background=PAGE_BG, |
| 32 | + foreground=INK, |
| 33 | + foreground_strong=INK, |
| 34 | + foreground_subtle=INK_MUTED, |
| 35 | + colors=OKABE_ITO, |
| 36 | + title_font_size=48, |
| 37 | + label_font_size=38, |
| 38 | + value_font_size=28, |
| 39 | + font_family="sans-serif", |
| 40 | +) |
| 41 | + |
| 42 | +# Read all visual tokens from the Style object — single source of truth |
| 43 | +BG = chart_style.background |
| 44 | +FG = chart_style.foreground |
| 45 | +FG_SUBTLE = chart_style.foreground_subtle |
| 46 | +PALETTE = chart_style.colors |
| 47 | +TITLE_SIZE = chart_style.title_font_size |
| 48 | +LABEL_SIZE = chart_style.label_font_size |
| 49 | +VALUE_SIZE = chart_style.value_font_size |
| 50 | +FONT = chart_style.font_family |
| 51 | + |
| 52 | +# Canvas |
| 53 | +WIDTH = 4800 |
| 54 | +HEIGHT = 2700 |
| 55 | +MARGIN_L = 400 |
| 56 | +MARGIN_R = 440 |
| 57 | +MARGIN_T = 220 |
| 58 | +MARGIN_B = 130 |
| 59 | +NODE_W = 52 |
| 60 | +NODE_GAP = 40 |
| 61 | + |
| 62 | +# Dominant flows get higher opacity to direct attention to key pathways |
| 63 | +ALPHA_DOMINANT = 0.72 |
| 64 | +ALPHA_DEFAULT = 0.38 |
| 65 | +DOMINANT_THRESHOLD = 20 # TWh |
| 66 | + |
| 67 | +# Data — energy flow in TWh (sources → end-use sectors) |
| 68 | +node_labels = [ |
| 69 | + "Coal", |
| 70 | + "Natural Gas", |
| 71 | + "Nuclear", |
| 72 | + "Renewables", |
| 73 | + "Residential", |
| 74 | + "Commercial", |
| 75 | + "Industrial", |
| 76 | + "Transportation", |
| 77 | +] |
| 78 | +N_SRC = 4 # first 4 are sources; rest are targets |
| 79 | + |
| 80 | +flows = [ |
| 81 | + (0, 4, 5), # Coal → Residential |
| 82 | + (0, 5, 8), # Coal → Commercial |
| 83 | + (0, 6, 25), # Coal → Industrial ← dominant |
| 84 | + (1, 4, 22), # Gas → Residential ← dominant |
| 85 | + (1, 5, 18), # Gas → Commercial |
| 86 | + (1, 6, 15), # Gas → Industrial |
| 87 | + (1, 7, 3), # Gas → Transportation |
| 88 | + (2, 4, 12), # Nuclear → Residential |
| 89 | + (2, 5, 10), # Nuclear → Commercial |
| 90 | + (2, 6, 8), # Nuclear → Industrial |
| 91 | + (3, 4, 8), # Renewables → Residential |
| 92 | + (3, 5, 6), # Renewables → Commercial |
| 93 | + (3, 6, 5), # Renewables → Industrial |
| 94 | + (3, 7, 4), # Renewables → Transportation |
| 95 | +] |
| 96 | + |
| 97 | +# Compute per-node totals |
| 98 | +node_total = [0] * len(node_labels) |
| 99 | +for src, tgt, val in flows: |
| 100 | + node_total[src] += val |
| 101 | + node_total[tgt] += val |
| 102 | + |
| 103 | +# Layout: vertical scale so the taller column fills available height |
| 104 | +avail_h = HEIGHT - MARGIN_T - MARGIN_B |
| 105 | +n_src_gaps = N_SRC - 1 |
| 106 | +n_tgt_gaps = len(node_labels) - N_SRC - 1 |
| 107 | +scale = (avail_h - max(n_src_gaps, n_tgt_gaps) * NODE_GAP) / sum(node_total[:N_SRC]) |
| 108 | + |
| 109 | +# Node y positions |
| 110 | +node_x = [] |
| 111 | +node_y0 = [] |
| 112 | +node_y1 = [] |
| 113 | + |
| 114 | +# Source nodes (left column) |
| 115 | +src_block_h = sum(node_total[i] * scale for i in range(N_SRC)) + n_src_gaps * NODE_GAP |
| 116 | +y = MARGIN_T + (avail_h - src_block_h) / 2 |
| 117 | +for i in range(N_SRC): |
| 118 | + h = node_total[i] * scale |
| 119 | + node_x.append(MARGIN_L) |
| 120 | + node_y0.append(y) |
| 121 | + node_y1.append(y + h) |
| 122 | + y += h + NODE_GAP |
| 123 | + |
| 124 | +# Target nodes (right column) |
| 125 | +tgt_indices = list(range(N_SRC, len(node_labels))) |
| 126 | +tgt_block_h = sum(node_total[i] * scale for i in tgt_indices) + n_tgt_gaps * NODE_GAP |
| 127 | +y = MARGIN_T + (avail_h - tgt_block_h) / 2 |
| 128 | +for i in tgt_indices: |
| 129 | + h = node_total[i] * scale |
| 130 | + node_x.append(WIDTH - MARGIN_R - NODE_W) |
| 131 | + node_y0.append(y) |
| 132 | + node_y1.append(y + h) |
| 133 | + y += h + NODE_GAP |
| 134 | + |
| 135 | +# Link paths (cubic bezier ribbons) |
| 136 | +src_cursor = list(node_y0[:N_SRC]) |
| 137 | +tgt_cursor = list(node_y0[N_SRC:]) |
| 138 | +link_data = [] |
| 139 | +for src, tgt, val in flows: |
| 140 | + h = val * scale |
| 141 | + x1 = node_x[src] + NODE_W |
| 142 | + y1t = src_cursor[src] |
| 143 | + y1b = y1t + h |
| 144 | + src_cursor[src] += h |
| 145 | + |
| 146 | + tgt_local = tgt - N_SRC |
| 147 | + x2 = node_x[tgt] |
| 148 | + y2t = tgt_cursor[tgt_local] |
| 149 | + y2b = y2t + h |
| 150 | + tgt_cursor[tgt_local] += h |
| 151 | + |
| 152 | + cx = (x1 + x2) / 2 |
| 153 | + path = ( |
| 154 | + f"M {x1:.1f},{y1t:.1f} " |
| 155 | + f"C {cx:.1f},{y1t:.1f} {cx:.1f},{y2t:.1f} {x2:.1f},{y2t:.1f} " |
| 156 | + f"L {x2:.1f},{y2b:.1f} " |
| 157 | + f"C {cx:.1f},{y2b:.1f} {cx:.1f},{y1b:.1f} {x1:.1f},{y1b:.1f} Z" |
| 158 | + ) |
| 159 | + c = PALETTE[src] # color drawn from Style object palette |
| 160 | + r, g, b = int(c[1:3], 16), int(c[3:5], 16), int(c[5:7], 16) |
| 161 | + alpha = ALPHA_DOMINANT if val >= DOMINANT_THRESHOLD else ALPHA_DEFAULT |
| 162 | + dominant = val >= DOMINANT_THRESHOLD |
| 163 | + # Ribbon midpoint for annotation placement |
| 164 | + ribbon_mid_y = (y1t + y1b + y2t + y2b) / 4 |
| 165 | + link_data.append((f"rgba({r},{g},{b},{alpha})", path, dominant, cx, ribbon_mid_y, val)) |
| 166 | + |
| 167 | + |
| 168 | +# Build SVG string |
| 169 | +parts = [ |
| 170 | + f'<svg xmlns="http://www.w3.org/2000/svg" width="{WIDTH}" height="{HEIGHT}" viewBox="0 0 {WIDTH} {HEIGHT}">', |
| 171 | + f'<rect width="{WIDTH}" height="{HEIGHT}" fill="{BG}"/>', |
| 172 | + # Title — font size from chart_style.title_font_size |
| 173 | + f'<text x="{WIDTH // 2}" y="{MARGIN_T // 2}" text-anchor="middle" ' |
| 174 | + f'dominant-baseline="middle" font-family="{FONT}" font-size="{TITLE_SIZE}" ' |
| 175 | + f'font-weight="600" fill="{FG}">' |
| 176 | + f"Energy Distribution · sankey-basic · pygal · anyplot.ai</text>", |
| 177 | + '<g id="links">', |
| 178 | +] |
| 179 | + |
| 180 | +# Non-dominant flows drawn first (background layer) |
| 181 | +for fill, path, dominant, _cx, _ribbon_mid_y, _val in link_data: |
| 182 | + if not dominant: |
| 183 | + parts.append(f'<path d="{path}" fill="{fill}" stroke="none"/>') |
| 184 | + |
| 185 | +# Dominant flows drawn on top with annotation showing their magnitude |
| 186 | +for fill, path, dominant, cx, ribbon_mid_y, val in link_data: |
| 187 | + if dominant: |
| 188 | + parts.append(f'<path d="{path}" fill="{fill}" stroke="none"/>') |
| 189 | + parts.append( |
| 190 | + f'<text x="{cx:.1f}" y="{ribbon_mid_y:.1f}" text-anchor="middle" ' |
| 191 | + f'dominant-baseline="middle" font-family="{FONT}" font-size="{VALUE_SIZE}" ' |
| 192 | + f'font-weight="700" fill="{FG}" opacity="0.80">{val} TWh</text>' |
| 193 | + ) |
| 194 | + |
| 195 | +parts.append("</g>") |
| 196 | + |
| 197 | +# Nodes |
| 198 | +parts.append('<g id="nodes">') |
| 199 | +for i in range(len(node_labels)): |
| 200 | + color = PALETTE[i] if i < N_SRC else INK_SOFT |
| 201 | + x = node_x[i] |
| 202 | + y0 = node_y0[i] |
| 203 | + h = node_y1[i] - node_y0[i] |
| 204 | + parts.append(f'<rect x="{x:.1f}" y="{y0:.1f}" width="{NODE_W}" height="{h:.1f}" fill="{color}" rx="5"/>') |
| 205 | +parts.append("</g>") |
| 206 | + |
| 207 | +# Labels — font sizes from chart_style.label_font_size / chart_style.value_font_size |
| 208 | +parts.append('<g id="labels">') |
| 209 | +for i in range(len(node_labels)): |
| 210 | + y_mid = (node_y0[i] + node_y1[i]) / 2 |
| 211 | + label = node_labels[i] |
| 212 | + val_str = f"{node_total[i]} TWh" |
| 213 | + if i < N_SRC: |
| 214 | + tx = node_x[i] - 24 |
| 215 | + anchor = "end" |
| 216 | + else: |
| 217 | + tx = node_x[i] + NODE_W + 24 |
| 218 | + anchor = "start" |
| 219 | + parts.append( |
| 220 | + f'<text x="{tx:.1f}" y="{y_mid - 22:.1f}" text-anchor="{anchor}" ' |
| 221 | + f'dominant-baseline="middle" font-family="{FONT}" font-size="{LABEL_SIZE}" ' |
| 222 | + f'font-weight="500" fill="{FG}">{label}</text>' |
| 223 | + ) |
| 224 | + parts.append( |
| 225 | + f'<text x="{tx:.1f}" y="{y_mid + 26:.1f}" text-anchor="{anchor}" ' |
| 226 | + f'dominant-baseline="middle" font-family="{FONT}" font-size="{VALUE_SIZE}" ' |
| 227 | + f'fill="{FG_SUBTLE}">{val_str}</text>' |
| 228 | + ) |
| 229 | +parts.append("</g>") |
| 230 | +parts.append("</svg>") |
| 231 | + |
| 232 | +svg_content = "\n".join(parts) |
| 233 | + |
| 234 | +# Save HTML (pygal-style interactive output) |
| 235 | +html_content = ( |
| 236 | + f'<!DOCTYPE html><html><head><meta charset="utf-8">' |
| 237 | + f"<title>sankey-basic · pygal · anyplot.ai</title>" |
| 238 | + f"<style>body{{margin:0;background:{BG}}}</style></head>" |
| 239 | + f"<body>{svg_content}</body></html>" |
| 240 | +) |
| 241 | +with open(f"plot-{THEME}.html", "w", encoding="utf-8") as fh: |
| 242 | + fh.write(html_content) |
| 243 | + |
| 244 | +# Save PNG via cairosvg (same pipeline pygal.render_to_png uses internally) |
| 245 | +cairosvg.svg2png(bytestring=svg_content.encode("utf-8"), write_to=f"plot-{THEME}.png") |
0 commit comments