Sankey Diagram

# === IMPORTS ===
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px

# === USER-EDITABLE PARAMETERS ===
# Change: Title of the Sankey diagram
title = "Energy Flow from Sources to Consumers (TWh/year)"

# Change: Color palette for the Sankey diagram
source_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']  # Sources
intermediate_colors = ['#e377c2', '#7f7f7f']  # Intermediate stages
consumer_colors = ['#17becf', '#bcbd22', '#ff9896', '#c5b0d5']  # Consumers

# Change: Node padding and thickness
node_padding = 10  # Change: Space between nodes (pixels)
node_thickness = 20  # Change: Thickness of nodes (pixels)

# Change: Font sizes
title_fontsize = 18
label_fontsize = 12

# === GENERATE EXAMPLE DATASET ===
# Create a realistic energy flow dataset
data = {
    'source': [
        'Solar', 'Solar', 'Wind', 'Wind', 'Coal', 'Coal', 
        'Natural Gas', 'Natural Gas', 'Nuclear', 'Nuclear', 'Hydro', 'Hydro'
    ],
    'target': [
        'Electric Grid', 'Direct Use', 'Electric Grid', 'Direct Use', 
        'Electric Grid', 'Industrial', 'Electric Grid', 'Direct Use',
        'Electric Grid', 'Industrial', 'Electric Grid', 'Direct Use'
    ],
    'value': [
        120, 15, 180, 25, 350, 80, 420, 150, 280, 40, 95, 35
    ]
}

# Add flows from intermediate to consumers
intermediate_to_consumer = {
    'source': [
        'Electric Grid', 'Electric Grid', 'Electric Grid', 'Electric Grid',
        'Direct Use', 'Direct Use', 'Direct Use', 'Industrial'
    ],
    'target': [
        'Residential', 'Commercial', 'Industrial', 'Transportation',
        'Residential', 'Commercial', 'Transportation', 'Industrial'
    ],
    'value': [
        450, 380, 520, 180, 85, 120, 95, 120
    ]
}

# Combine all data
all_data = {
    'source': data['source'] + intermediate_to_consumer['source'],
    'target': data['target'] + intermediate_to_consumer['target'],
    'value': data['value'] + intermediate_to_consumer['value']
}

df = pd.DataFrame(all_data)

# Get unique nodes
all_nodes = list(set(df['source'].unique().tolist() + df['target'].unique().tolist()))

# Create node to index mapping
node_to_index = {node: i for i, node in enumerate(all_nodes)}

# Prepare source, target, and value arrays for Sankey
source_indices = [node_to_index[source] for source in df['source']]
target_indices = [node_to_index[target] for target in df['target']]
values = df['value'].tolist()

# Create color array for nodes
node_colors = []
for node in all_nodes:
    if node in ['Solar', 'Wind', 'Coal', 'Natural Gas', 'Nuclear', 'Hydro']:
        node_colors.append(source_colors[all_nodes[:6].index(node)] if node in all_nodes[:6] else '#1f77b4')
    elif node in ['Electric Grid', 'Direct Use']:
        node_colors.append(intermediate_colors[0] if node == 'Electric Grid' else intermediate_colors[1])
    else:
        node_colors.append(consumer_colors[all_nodes[-4:].index(node)] if node in all_nodes[-4:] else '#17becf')

# === CREATE SANKEY DIAGRAM ===
fig = go.Figure(data=[go.Sankey(
    node=dict(
        pad=node_padding,
        thickness=node_thickness,
        line=dict(color='#000000', width=0.5),
        label=all_nodes,
        color=node_colors
    ),
    link=dict(
        source=source_indices,
        target=target_indices,
        value=values,
        hovertemplate='%{source.label} → %{target.label}<br>Energy: %{value} TWh/year<extra></extra>'
    )
)])

# Update layout
fig.update_layout(
    title=dict(
        text=title,
        x=0.5,
        y=0.98,
        font=dict(size=title_fontsize)
    ),
    font=dict(size=label_fontsize),
    margin=dict(t=130, b=80, l=80, r=80),
    hovermode='x'
)

# Print summary statistics
total_energy = df['value'].sum()
print(f"\nTotal Energy Flow: {total_energy} TWh/year")
print(f"Number of Energy Sources: {len([n for n in all_nodes if n in ['Solar', 'Wind', 'Coal', 'Natural Gas', 'Nuclear', 'Hydro']])}")
print(f"Number of Consumer Sectors: {len([n for n in all_nodes if n in ['Residential', 'Commercial', 'Industrial', 'Transportation']])}")
print(f"\nTop 3 Energy Sources:")
source_totals = df[df['source'].isin(['Solar', 'Wind', 'Coal', 'Natural Gas', 'Nuclear', 'Hydro'])].groupby('source')['value'].sum().sort_values(ascending=False)
for source, value in source_totals.head(3).items():
    print(f"  {source}: {value} TWh/year")

# Display the plot
fig.show()
# END-OF-CODE

Total Energy Input: 98 Quad BTU
Electricity Generation: 38 Quad BTU
Energy Loss: 24 Quad BTU (24.5%)
End Use: 123 Quad BTU