[Python] Continue colour flow to 3rd strand of a sankey diagram

I have created this code for a sankey diagram, but I can't get the colours to flow through the middle to the right side. The left side has distinguished fuel types and the right side has combinations of these fuels. Each fuel has a colour and it should go through to the combination as well. This I have done with previous sankeys. But now I added a middle strand through which these lines must flow (seasonal use) and it seems like the code basically creates two back-to-back sankeys.

import pandas as pd
import plotly.graph_objects as go

# Load data
df = pd.read_excel('/content/Fuel combination_Winter_Summer.xlsx')

# Clean column headers
df.columns = df.columns.str.strip()

# Melt the data to long format
id_vars = ['Fuel']
df_long = df.melt(id_vars=id_vars, var_name='Household', value_name='Season')

# Drop rows with blank or missing seasons
df_long = df_long[df_long['Season'].notna() & (df_long['Season'].str.strip() != '')]

# Clean strings
df_long['Fuel'] = df_long['Fuel'].astype(str).str.strip()
df_long['Season'] = df_long['Season'].astype(str).str.strip()
df_long['Household'] = df_long['Household'].astype(str).str.strip()

# Split multi-season entries like "Winter Summer" into multiple rows
df_long = df_long.assign(Season=df_long['Season'].str.split())
df_long = df_long.explode('Season')
df_long['Season'] = df_long['Season'].str.strip()

# Create fuel combinations per household and season
combo_df = df_long.groupby(['Household', 'Season'])['Fuel'].apply(
lambda x: ', '.join(sorted(x.unique()))
).reset_index()
combo_df.rename(columns={'Fuel': 'Fuel_Combo'}, inplace=True)

# Merge combo info back into main df
df_long = df_long.merge(combo_df, on=['Household', 'Season'])

# Create unique internal labels
df_long['Fuel_node'] = 'Fuel: ' + df_long['Fuel']
df_long['Season_node'] = 'Season: ' + df_long['Season']
df_long['Combo_node'] = 'Combo: ' + df_long['Fuel_Combo']

# Create list of unique nodes
internal_labels = pd.concat([
df_long['Fuel_node'],
df_long['Season_node'],
df_long['Combo_node']
]).unique().tolist()

# Create label-to-index mapping
label_to_index = {label: i for i, label in enumerate(internal_labels)}

# Map source and target
sources_fuel_season = df_long['Fuel_node'].map(label_to_index)
targets_fuel_season = df_long['Season_node'].map(label_to_index)

sources_season_combo = df_long['Season_node'].map(label_to_index)
targets_season_combo = df_long['Combo_node'].map(label_to_index)

sources = sources_fuel_season.tolist() + sources_season_combo.tolist()
targets = targets_fuel_season.tolist() + targets_season_combo.tolist()
values = [1] * len(sources)

# Display labels without prefixes
display_labels = [label.split(': ', 1)[1] for label in internal_labels]

# Define fuel colors
fuel_colors = {
'Coal': 'rgba(38, 38, 38, 0.5)',
'Wood': 'rgba(139, 69, 19, 0.5)',
'Electricity': 'rgba(0, 123, 255, 0.5)',
'Animal dung': 'rgba(96, 72, 30, 0.5)',
'Paraffin': 'rgba(255, 215, 0, 0.5)',
'Gas': 'rgba(173, 213, 174, 0.4)'
}

# Create a function to calculate a clearer gradient based on fuel proportions
def get_gradient_color(fuel_combination):
fuel_types = ['Coal', 'Wood', 'Electricity', 'Animal dung', 'Paraffin', 'Gas']

# Calculate fuel proportions in the combination
fuel_proportions = {fuel: fuel_combination.count(fuel) / len(fuel_combination) for fuel in fuel_types}

# Build an RGBA string based on proportions and apply opacity based on the fuel type's proportion
rgba_values = []
for fuel in fuel_types:
proportion = fuel_proportions.get(fuel, 0)
# Start darker and gradually lighten the color based on proportion
rgba_values.append(fuel_colors.get(fuel, 'lightgray').replace('0.5', str(0.2 + proportion * 0.3)))

# Adjust the final opacity to ensure clearer differentiation
final_opacity = max([fuel_proportions.get(fuel, 0) for fuel in fuel_types]) * 0.9
# Use this proportion to control how much of the fuel impacts the final color.
return f'rgba(100, 150, 200, {final_opacity})' # This is a placeholder color

# Create combination colors using the gradient function
combination_colors = {combo: get_gradient_color(combo.split(', ')) for combo in df_long['Fuel_Combo'].unique()}

# Assign node colors (fuel and combo nodes)
node_colors = []
for label in internal_labels:
if label.startswith('Fuel: '):
fuel = label.split(': ', 1)[1]
node_colors.append(fuel_colors.get(fuel, 'lightgray'))
elif label.startswith('Combo: '):
combo = label.split(': ', 1)[1]
node_colors.append(combination_colors.get(combo, 'lightgray'))
else:
node_colors.append('lightgray')

# Assign link colors for fuel → season (fuel-based links)
link_colors = []
for src in sources_fuel_season:
fuel_label = internal_labels[src]
fuel = fuel_label.split(': ', 1)[1]
link_colors.append(fuel_colors.get(fuel, 'lightgray'))

# Second set of links (season → combo): Use combo color
link_colors += [combination_colors.get(internal_labels[targets_season_combo].split(': ', 1)[1], 'lightgray') for i in range(len(sources_season_combo))]

# Build Sankey
fig = go.Figure(data=[go.Sankey(
node=dict(
pad=15,
thickness=20,
line=dict(color='gray', width=0.5),
label=display_labels,
color=node_colors
),
link=dict(
source=sources,
target=targets,
value=values,
color=link_colors
)
)])

fig.update_layout(
title_text="Fuel → Season → Household Fuel Combination",
font_size=10,
width=1100,
height=1000
)

fig.show()

Continue reading...