{"spec_id":"sankey-basic","library":"plotnine","language":"python","code":"\"\"\" anyplot.ai\nsankey-basic: Basic Sankey Diagram\nLibrary: plotnine 0.15.7 | Python 3.13.14\nQuality: 91/100 | Updated: 2026-07-25\n\"\"\"\n\nimport os\nimport sys\n\n\nsys.path = [p for p in sys.path if os.path.abspath(p) != os.path.dirname(os.path.abspath(__file__))]\n\nimport numpy as np\nimport pandas as pd\nfrom plotnine import (\n    aes,\n    annotate,\n    coord_cartesian,\n    element_blank,\n    element_rect,\n    element_text,\n    geom_polygon,\n    geom_rect,\n    geom_text,\n    ggplot,\n    labs,\n    scale_fill_manual,\n    theme,\n    theme_minimal,\n)\n\n\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\"\n\n# Imprint palette (canonical order) for source categories\nIMPRINT = [\"#009E73\", \"#C475FD\", \"#4467A3\", \"#BD8233\"]\n\n# Data - Energy flow from sources to sectors\nflows = pd.DataFrame(\n    {\n        \"source\": [\"Coal\", \"Coal\", \"Gas\", \"Gas\", \"Gas\", \"Nuclear\", \"Nuclear\", \"Renewables\", \"Renewables\"],\n        \"target\": [\n            \"Industrial\",\n            \"Residential\",\n            \"Industrial\",\n            \"Commercial\",\n            \"Residential\",\n            \"Commercial\",\n            \"Residential\",\n            \"Commercial\",\n            \"Residential\",\n        ],\n        \"value\": [35, 15, 25, 20, 15, 18, 12, 8, 12],\n    }\n)\n\n# Define node positions\nsources = [\"Coal\", \"Gas\", \"Nuclear\", \"Renewables\"]\ntargets = [\"Industrial\", \"Commercial\", \"Residential\"]\n\n# X positions with margins for labels\nx_left = 0.10\nx_right = 0.90\nnode_width = 0.06\nnode_gap = 0.03\n\n# Calculate node sizes based on total flow\nsource_totals = flows.groupby(\"source\")[\"value\"].sum().to_dict()\ntarget_totals = flows.groupby(\"target\")[\"value\"].sum().to_dict()\ntotal_flow = flows[\"value\"].sum()\n\n# Calculate source node positions (left side)\nsource_positions = {}\ncurrent_y = 1.0\nfor src in sources:\n    height = source_totals[src] / total_flow * 0.8\n    source_positions[src] = {\n        \"x\": x_left,\n        \"y_top\": current_y,\n        \"y_bottom\": current_y - height,\n        \"height\": height,\n        \"flow_offset\": 0,\n    }\n    current_y = current_y - height - node_gap\n\n# Calculate target node positions (right side)\ntarget_positions = {}\ncurrent_y = 1.0\nfor tgt in targets:\n    height = target_totals[tgt] / total_flow * 0.8\n    target_positions[tgt] = {\n        \"x\": x_right,\n        \"y_top\": current_y,\n        \"y_bottom\": current_y - height,\n        \"height\": height,\n        \"flow_offset\": 0,\n    }\n    current_y = current_y - height - node_gap\n\n# Imprint colors for sources; theme-adaptive neutral for targets\nsource_colors_map = {\"Coal\": IMPRINT[0], \"Gas\": IMPRINT[1], \"Nuclear\": IMPRINT[2], \"Renewables\": IMPRINT[3]}\ntarget_colors_map = {\"Industrial\": INK_SOFT, \"Commercial\": INK_SOFT, \"Residential\": INK_SOFT}\n\n# Build node rectangles dataframe\nnode_data = []\nfor src in sources:\n    pos = source_positions[src]\n    node_data.append(\n        {\n            \"name\": src,\n            \"xmin\": pos[\"x\"],\n            \"xmax\": pos[\"x\"] + node_width,\n            \"ymin\": pos[\"y_bottom\"],\n            \"ymax\": pos[\"y_top\"],\n            \"label_x\": pos[\"x\"] - 0.02,\n            \"label_y\": (pos[\"y_top\"] + pos[\"y_bottom\"]) / 2,\n            \"side\": \"source\",\n            \"node_color\": src,\n        }\n    )\n\nfor tgt in targets:\n    pos = target_positions[tgt]\n    node_data.append(\n        {\n            \"name\": tgt,\n            \"xmin\": pos[\"x\"] - node_width,\n            \"xmax\": pos[\"x\"],\n            \"ymin\": pos[\"y_bottom\"],\n            \"ymax\": pos[\"y_top\"],\n            \"label_x\": pos[\"x\"] + 0.02,\n            \"label_y\": (pos[\"y_top\"] + pos[\"y_bottom\"]) / 2,\n            \"side\": \"target\",\n            \"node_color\": tgt,\n        }\n    )\n\nnodes_df = pd.DataFrame(node_data)\n\n# Build flow polygons (curved paths between nodes)\nflow_polygons = []\nflow_labels = []\nflow_x_left = x_left + node_width\nflow_x_right = x_right - node_width\nn_points = 50\n\n# Stagger the label sample point along each source's outgoing flows so\n# labels for flows sharing a source/target pair don't collide, while still\n# sitting exactly on that flow's own curve (not an arbitrary offset).\nfor src in sources:\n    src_flows = flows[flows[\"source\"] == src].reset_index(drop=True)\n    n_src_flows = len(src_flows)\n    for i, row in src_flows.iterrows():\n        tgt = row[\"target\"]\n        val = row[\"value\"]\n\n        flow_height = val / total_flow * 0.8\n\n        src_pos = source_positions[src]\n        src_y_top = src_pos[\"y_top\"] - src_pos[\"flow_offset\"]\n        src_y_bottom = src_y_top - flow_height\n        src_pos[\"flow_offset\"] += flow_height\n\n        tgt_pos = target_positions[tgt]\n        tgt_y_top = tgt_pos[\"y_top\"] - tgt_pos[\"flow_offset\"]\n        tgt_y_bottom = tgt_y_top - flow_height\n        tgt_pos[\"flow_offset\"] += flow_height\n\n        # Smooth cubic Hermite interpolation for flow curves\n        t = np.linspace(0, 1, n_points)\n        x_top = flow_x_left + (flow_x_right - flow_x_left) * t\n        y_top = src_y_top + (tgt_y_top - src_y_top) * (3 * t**2 - 2 * t**3)\n\n        x_bottom = flow_x_right + (flow_x_left - flow_x_right) * t\n        y_bottom = tgt_y_bottom + (src_y_bottom - tgt_y_bottom) * (3 * t**2 - 2 * t**3)\n\n        x_polygon = np.concatenate([x_top, x_bottom])\n        y_polygon = np.concatenate([y_top, y_bottom])\n\n        for j in range(len(x_polygon)):\n            flow_polygons.append({\"x\": x_polygon[j], \"y\": y_polygon[j], \"flow_id\": f\"{src}_{tgt}\", \"source\": src})\n\n        # Sample the label position from a point that actually lies on this\n        # flow's ribbon, staggered by index so co-sourced flows don't overlap.\n        t_label = 0.35 + (i / max(n_src_flows - 1, 1)) * 0.3 if n_src_flows > 1 else 0.5\n        label_idx = int(round(t_label * (n_points - 1)))\n        label_x = x_top[label_idx]\n        label_y = (y_top[label_idx] + y_bottom[n_points - 1 - label_idx]) / 2\n        flow_labels.append({\"x\": label_x, \"y\": label_y, \"value\": str(val), \"flow_height\": flow_height})\n\nflows_df = pd.DataFrame(flow_polygons)\nflow_labels_df = pd.DataFrame(flow_labels)\n\n# Create the plot\nplot = (\n    ggplot()\n    # Flow polygons with transparency\n    + geom_polygon(flows_df, aes(x=\"x\", y=\"y\", group=\"flow_id\", fill=\"source\"), alpha=0.44)\n    # Node rectangles\n    + geom_rect(\n        nodes_df, aes(xmin=\"xmin\", xmax=\"xmax\", ymin=\"ymin\", ymax=\"ymax\", fill=\"node_color\"), color=\"white\", size=0.5\n    )\n    # Flow value labels (only for larger flows to avoid clutter)\n    + geom_text(\n        flow_labels_df[flow_labels_df[\"flow_height\"] >= 0.05],\n        aes(x=\"x\", y=\"y\", label=\"value\"),\n        ha=\"center\",\n        va=\"center\",\n        size=7,\n        color=INK,\n        fontweight=\"bold\",\n    )\n    # Source labels (right-aligned)\n    + geom_text(\n        nodes_df[nodes_df[\"side\"] == \"source\"],\n        aes(x=\"label_x\", y=\"label_y\", label=\"name\"),\n        ha=\"right\",\n        size=9,\n        color=INK,\n        fontweight=\"bold\",\n    )\n    # Target labels (left-aligned)\n    + geom_text(\n        nodes_df[nodes_df[\"side\"] == \"target\"],\n        aes(x=\"label_x\", y=\"label_y\", label=\"name\"),\n        ha=\"left\",\n        size=9,\n        color=INK,\n        fontweight=\"bold\",\n    )\n    + scale_fill_manual(values={**source_colors_map, **target_colors_map})\n    + labs(title=\"Energy Flow · sankey-basic · python · plotnine · anyplot.ai\", x=\"\", y=\"\")\n    + coord_cartesian(xlim=(-0.02, 1.02))\n    + theme_minimal()\n    + theme(\n        figure_size=(8, 4.5),\n        plot_background=element_rect(fill=PAGE_BG, color=PAGE_BG),\n        panel_background=element_rect(fill=PAGE_BG),\n        plot_title=element_text(size=12, ha=\"center\", weight=\"bold\", color=INK),\n        axis_text=element_blank(),\n        axis_ticks=element_blank(),\n        panel_grid=element_blank(),\n        legend_position=\"none\",\n    )\n    + annotate(\"text\", x=x_left + node_width / 2, y=-0.05, label=\"Sources\", size=8, color=INK_SOFT, fontweight=\"bold\")\n    + annotate(\"text\", x=x_right - node_width / 2, y=-0.05, label=\"Sectors\", size=8, color=INK_SOFT, fontweight=\"bold\")\n)\n\nplot.save(f\"plot-{THEME}.png\", dpi=400, width=8, height=4.5, units=\"in\", verbose=False)\n"}