geus-glaciology-and-climate / mass_balance Goto Github PK

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

df = pd.DataFrame(columns=['value','source','target','meta'])
df.index.name = 'label'

df.loc['rain:freeze'] = [50/2, 'rain', 'GL', 'A2019']  # Alexander 2019
df.loc['snow:fall'] = [750, 'snow fall', 'GL', 'A2019'] # fudge from A2019
df.loc['DEFICIT'] = [160, 'DEFICIT', 'GL', '']
df.loc['condensation'] = [10, 'condensation', 'GL', 'guess']
df.loc['drifting_out'] = [100, 'GL', 'drift', 'guess']
df.loc['drifting_in'] = [100
                         , 'drift', 'GL', 'guess']

df.loc['rain:runoff'] = [50/2, 'rain', 'runoff', 'A2019']
df.loc['rain:runoff:liquid'] = [50/2, 'runoff', 'liquid', 'A2019']

df.loc['SMB (loss)'] = [400, 'GL', 'SMB (loss)', 'A2019']
df.loc['dynamics'] = [500, 'GL', 'dynamics', '']

df.loc['icebergs'] = [250, 'dynamics', 'icebergs', 'M2020 & E2014'] # Enderlin & Mankoff
df.loc['frontal melt'] = [250, 'dynamics', 'frontal melt', 'E2014']

df.loc['runoff'] = [400, 'SMB (loss)','runoff', '']
df.loc['refreeze'] = [250, 'GL','refreeze', '']
df.loc['refreeze1'] = [250, 'refreeze','GL', '']
# df.loc['runoff'] = [400, 'melt','runoff', '']

#df.loc['basal melt'] = [25, 'basal melt', 'liquid', 'K2020'] # Karlsson
df.loc['basal melt'] = [25, 'GL', 'basal melt', 'K2020'] # Karlsson
df.loc['basal melt1'] = [25, 'basal melt', 'runoff', 'K2020'] # Karlsson
df.loc['basal melt2'] = [25, 'runoff', 'liquid', 'K2020'] # Karlsson
# df.loc['basal melt2'] = [50, 'runoff', 'liquid', 'K2020']

df.loc['evaporation'] = [10, 'GL', 'evaporation', '']
df.loc['sublimation'] = [10, 'GL', 'sublimation', '']

# Sum all the liquid terms to a liquid budget (solid is just 'iceberg')
for i,v in enumerate(['runoff','frontal melt']):
    df.loc[f'liquid_{v}'] = [df.loc[v]['value'], v, 'liquid', '']

    
# QC Check
for s in df['source'].unique():
    if s in df['target'].unique():
        sv = df[df['source'] == s]['value'].sum()
        tv = df[df['target'] == s]['value'].sum()
        if( sv != tv ):
            print(f"Error: {s} inputs != outputs [{sv} != {tv}]")

            
# # Add numeric values in parentheses
for t in np.unique(df['source'].values.astype(str)):
    v = (df[df["source"] == t]["value"].sum()).round().astype(int).astype(str)
    df = df.replace(to_replace=t, value=f'{t} ({v})')

for t in np.unique(df['target'].values.astype(str)):
    if t[-1] == ')': continue
    v = (df[df["target"] == t]["value"].sum()).round().astype(int).astype(str)
    df = df.replace(to_replace=t, value=f'{t} ({v})')

# print(df)

## Plotly
# uniqify columns: Convert strings to numbers. Sequential. Across columns
keys = np.unique(df[['source','target']].values.flatten().astype(str))
df = df.replace(to_replace=keys, value=np.arange(len(keys)))

fig = go.Figure(go.Sankey(
    arrangement = "snap",
    node = {
        "label": keys,
        'color': "#666666",
        'pad':10},  # 10 Pixels
    link = {
        "source": df['source'].values,
        "target": df['target'].values,
        # "color": "#DDDDDD",
        "value": df['value'].values}))

# fig.show()
fig.write_image("plotly.png")
# interactive HTML: drag to re-arrange and clean up, then screenshot.
fig.write_html("plotly.html")