{"spec_id":"sankey-basic","library":"altair","language":"python","code":"\"\"\" anyplot.ai\nsankey-basic: Basic Sankey Diagram\nLibrary: altair 6.2.2 | Python 3.13.14\nQuality: 89/100 | Updated: 2026-07-25\n\"\"\"\n\nimport os\n\nimport altair as alt\nimport pandas as pd\nfrom PIL import Image\n\n\n# Theme tokens\nTHEME = os.getenv(\"ANYPLOT_THEME\", \"light\")\nPAGE_BG = \"#FAF8F1\" if THEME == \"light\" else \"#1A1A17\"\nELEVATED_BG = \"#FFFDF6\" if THEME == \"light\" else \"#242420\"\nINK = \"#1A1A17\" if THEME == \"light\" else \"#F0EFE8\"\nINK_SOFT = \"#4A4A44\" if THEME == \"light\" else \"#B8B7B0\"\n\n# Data - Energy flow from sources to sectors\nflows = [\n    {\"source\": \"Coal\", \"target\": \"Residential\", \"value\": 20},\n    {\"source\": \"Coal\", \"target\": \"Commercial\", \"value\": 15},\n    {\"source\": \"Coal\", \"target\": \"Industrial\", \"value\": 45},\n    {\"source\": \"Gas\", \"target\": \"Residential\", \"value\": 35},\n    {\"source\": \"Gas\", \"target\": \"Commercial\", \"value\": 25},\n    {\"source\": \"Gas\", \"target\": \"Industrial\", \"value\": 30},\n    {\"source\": \"Nuclear\", \"target\": \"Residential\", \"value\": 15},\n    {\"source\": \"Nuclear\", \"target\": \"Commercial\", \"value\": 20},\n    {\"source\": \"Nuclear\", \"target\": \"Industrial\", \"value\": 15},\n    {\"source\": \"Renewable\", \"target\": \"Residential\", \"value\": 25},\n    {\"source\": \"Renewable\", \"target\": \"Commercial\", \"value\": 20},\n    {\"source\": \"Renewable\", \"target\": \"Transport\", \"value\": 10},\n]\n\ndf = pd.DataFrame(flows)\n\n# Canvas: 620x320 inner view @ scale_factor=4.0 — vl-convert pads title/legend\n# outside this view, landing near the 3200x1800 target (see prompts/library/altair.md).\nwidth = 620\nheight = 320\nnode_width = 18\nnode_padding = 6\nfill_ratio = 0.84  # leaves headroom so the last node never brushes the canvas edge\n\n# Compute node positions\nsources = df[\"source\"].unique().tolist()\ntargets = df[\"target\"].unique().tolist()\n\nsource_totals = df.groupby(\"source\")[\"value\"].sum().to_dict()\ntarget_totals = df.groupby(\"target\")[\"value\"].sum().to_dict()\ntotal_flow = df[\"value\"].sum()\n\ntop_margin = 40\nbottom_margin = 34\navailable_height = height - top_margin - bottom_margin\n\n# Position source nodes on left, vertically centered\nsource_total_height = sum(source_totals.values()) / total_flow * available_height * fill_ratio\nsource_total_with_padding = source_total_height + node_padding * (len(sources) - 1)\nstart_y_sources = top_margin + (available_height - source_total_with_padding) / 2\n\nsource_positions = {}\ncurrent_y = start_y_sources\nfor src in sources:\n    node_height = (source_totals[src] / total_flow) * available_height * fill_ratio\n    source_positions[src] = {\"y\": current_y, \"height\": node_height}\n    current_y += node_height + node_padding\n\n# Position target nodes on right, vertically centered\ntarget_total_height = sum(target_totals.values()) / total_flow * available_height * fill_ratio\ntarget_total_with_padding = target_total_height + node_padding * (len(targets) - 1)\nstart_y_targets = top_margin + (available_height - target_total_with_padding) / 2\n\ntarget_positions = {}\ncurrent_y = start_y_targets\nfor tgt in targets:\n    node_height = (target_totals[tgt] / total_flow) * available_height * fill_ratio\n    target_positions[tgt] = {\"y\": current_y, \"height\": node_height}\n    current_y += node_height + node_padding\n\n# Imprint palette positions 1-4 for source colors — distinct, colorblind-safe\nsource_colors = {\n    \"Coal\": \"#009E73\",  # Imprint #1 (brand green) — ALWAYS first series\n    \"Gas\": \"#C475FD\",  # Imprint #2 (lavender)\n    \"Nuclear\": \"#4467A3\",  # Imprint #3 (blue)\n    \"Renewable\": \"#BD8233\",  # Imprint #4 (ochre)\n}\n\n# Build node rectangles data\nnodes_data = []\nfor src in sources:\n    pos = source_positions[src]\n    nodes_data.append(\n        {\n            \"name\": src,\n            \"x\": 0,\n            \"y\": pos[\"y\"],\n            \"x2\": node_width,\n            \"y2\": pos[\"y\"] + pos[\"height\"],\n            \"color\": source_colors[src],\n            \"label_x\": node_width + 8,\n            \"label_y\": pos[\"y\"] + pos[\"height\"] / 2,\n            \"total\": source_totals[src],\n            \"side\": \"source\",\n        }\n    )\n\nfor tgt in targets:\n    pos = target_positions[tgt]\n    nodes_data.append(\n        {\n            # Target nodes use the theme-adaptive neutral anchor, not a categorical\n            # color — they are aggregation/baseline blocks, not their own data series.\n            \"name\": tgt,\n            \"x\": width - node_width,\n            \"y\": pos[\"y\"],\n            \"x2\": width,\n            \"y2\": pos[\"y\"] + pos[\"height\"],\n            \"color\": INK,\n            \"label_x\": width - node_width - 8,\n            \"label_y\": pos[\"y\"] + pos[\"height\"] / 2,\n            \"total\": target_totals[tgt],\n            \"side\": \"target\",\n        }\n    )\n\nnodes_df = pd.DataFrame(nodes_data)\n\n# Generate smoothstep S-curve polygon points for each flow band\nsource_y_offsets = {src: source_positions[src][\"y\"] for src in sources}\ntarget_y_offsets = {tgt: target_positions[tgt][\"y\"] for tgt in targets}\n\nall_flow_data = []\nnum_curve_points = 40\n\nfor _, row in df.iterrows():\n    src = row[\"source\"]\n    tgt = row[\"target\"]\n    val = row[\"value\"]\n\n    src_height = (val / source_totals[src]) * source_positions[src][\"height\"]\n    tgt_height = (val / target_totals[tgt]) * target_positions[tgt][\"height\"]\n\n    src_y_top = source_y_offsets[src]\n    src_y_bottom = src_y_top + src_height\n    tgt_y_top = target_y_offsets[tgt]\n    tgt_y_bottom = tgt_y_top + tgt_height\n\n    source_y_offsets[src] += src_height\n    target_y_offsets[tgt] += tgt_height\n\n    x_start = node_width\n    x_end = width - node_width\n\n    top_points = []\n    for i in range(num_curve_points):\n        t = i / (num_curve_points - 1)\n        x = x_start + t * (x_end - x_start)\n        bezier_t = t * t * (3 - 2 * t)\n        y = src_y_top + bezier_t * (tgt_y_top - src_y_top)\n        top_points.append((x, y))\n\n    bottom_points = []\n    for i in range(num_curve_points - 1, -1, -1):\n        t = i / (num_curve_points - 1)\n        x = x_start + t * (x_end - x_start)\n        bezier_t = t * t * (3 - 2 * t)\n        y = src_y_bottom + bezier_t * (tgt_y_bottom - src_y_bottom)\n        bottom_points.append((x, y))\n\n    all_points = top_points + bottom_points\n    for pt_idx, (x, y) in enumerate(all_points):\n        all_flow_data.append(\n            {\"flow_id\": f\"{src}-{tgt}\", \"source\": src, \"target\": tgt, \"value\": val, \"x\": x, \"y\": y, \"order\": pt_idx}\n        )\n\nflows_df = pd.DataFrame(all_flow_data)\n\n# Flow polygons colored by source — symbolType=\"square\" so the legend swatch\n# reads as a filled band rather than the line-mark default\nlinks_chart = (\n    alt.Chart(flows_df)\n    .mark_line(filled=True, opacity=0.55, strokeWidth=0)\n    .encode(\n        x=alt.X(\"x:Q\", scale=alt.Scale(domain=[0, width]), axis=None),\n        y=alt.Y(\"y:Q\", scale=alt.Scale(domain=[0, height]), axis=None),\n        color=alt.Color(\n            \"source:N\",\n            scale=alt.Scale(domain=list(source_colors.keys()), range=list(source_colors.values())),\n            legend=alt.Legend(\n                title=\"Energy Source\",\n                titleFontSize=10,\n                labelFontSize=10,\n                orient=\"right\",\n                symbolType=\"square\",\n                symbolSize=120,\n            ),\n        ),\n        detail=\"flow_id:N\",\n        order=\"order:Q\",\n    )\n)\n\n# Node rectangles\nnodes_chart = (\n    alt.Chart(nodes_df)\n    .mark_rect(stroke=INK_SOFT, strokeWidth=1.5)\n    .encode(\n        x=alt.X(\"x:Q\", scale=alt.Scale(domain=[0, width])),\n        y=alt.Y(\"y:Q\", scale=alt.Scale(domain=[0, height])),\n        x2=\"x2:Q\",\n        y2=\"y2:Q\",\n        color=alt.Color(\"color:N\", scale=None),\n        opacity=alt.condition(alt.datum.side == \"target\", alt.value(0.8), alt.value(1.0)),\n        tooltip=[alt.Tooltip(\"name:N\", title=\"Node\"), alt.Tooltip(\"total:Q\", title=\"Total Flow (units)\")],\n    )\n)\n\n# Source labels (right of left nodes)\nsource_labels_df = nodes_df[nodes_df[\"side\"] == \"source\"]\nsource_labels = (\n    alt.Chart(source_labels_df)\n    .mark_text(fontSize=12, fontWeight=\"bold\", align=\"left\", baseline=\"middle\")\n    .encode(\n        x=alt.X(\"label_x:Q\", scale=alt.Scale(domain=[0, width])),\n        y=alt.Y(\"label_y:Q\", scale=alt.Scale(domain=[0, height])),\n        text=\"name:N\",\n        color=alt.value(INK),\n    )\n)\n\n# Target labels (left of right nodes)\ntarget_labels_df = nodes_df[nodes_df[\"side\"] == \"target\"]\ntarget_labels = (\n    alt.Chart(target_labels_df)\n    .mark_text(fontSize=12, fontWeight=\"bold\", align=\"right\", baseline=\"middle\")\n    .encode(\n        x=alt.X(\"label_x:Q\", scale=alt.Scale(domain=[0, width])),\n        y=alt.Y(\"label_y:Q\", scale=alt.Scale(domain=[0, height])),\n        text=\"name:N\",\n        color=alt.value(INK),\n    )\n)\n\n# Compose all layers with theme-adaptive chrome\nchart = (\n    alt.layer(links_chart, nodes_chart, source_labels, target_labels)\n    .properties(\n        width=width,\n        height=height,\n        background=PAGE_BG,\n        title=alt.Title(\n            text=\"sankey-basic · python · altair · anyplot.ai\",\n            subtitle=\"Energy Flow from Sources to Sectors\",\n            fontSize=16,\n            subtitleFontSize=13,\n            anchor=\"middle\",\n            color=INK,\n            subtitleColor=INK_SOFT,\n        ),\n    )\n    .configure_view(strokeWidth=0, fill=PAGE_BG)\n    .configure_legend(\n        padding=10, cornerRadius=4, fillColor=ELEVATED_BG, strokeColor=INK_SOFT, labelColor=INK_SOFT, titleColor=INK\n    )\n)\n\n# Save outputs — PNG padded to the exact 3200x1800 target (see altair.md \"Canvas\")\nchart.save(f\"plot-{THEME}.png\", scale_factor=4.0)\n\nTW, TH = 3200, 1800\n_img = Image.open(f\"plot-{THEME}.png\").convert(\"RGB\")\n_w, _h = _img.size\nif _w > TW or _h > TH:\n    raise SystemExit(\n        f\"altair vl-convert produced {_w}x{_h}, exceeds target {TW}x{TH}. \"\n        f\"Shrink chart .properties(width=, height=) values and re-render.\"\n    )\nif _w < TW or _h < TH:\n    _canvas = Image.new(\"RGB\", (TW, TH), PAGE_BG)\n    _canvas.paste(_img, ((TW - _w) // 2, (TH - _h) // 2))\n    _canvas.save(f\"plot-{THEME}.png\")\n\nchart.save(f\"plot-{THEME}.html\")\n"}