scripts/sankey.py

# -*- coding: utf-8 -*-
# import json
# import urllib

import plotly.graph_objects as go
from dash import Dash, Input, Output, dcc, html
from dash.exceptions import PreventUpdate

app = Dash(__name__)

app.layout = html.Div(
    [
        html.H1("Sankey diagram"),
        html.P("Data"),
        dcc.Input(id="SM1", placeholder="SM1"),
        dcc.Input(id="SM2", placeholder="SM2"),
        dcc.Input(id="SM3", placeholder="SM3"),
        dcc.Input(id="SM4", placeholder="SM4"),
        dcc.Input(id="SM5", placeholder="SM5"),
        dcc.Input(id="BS1", placeholder="BS1"),
        dcc.Input(id="BS2", placeholder="BS2"),
        dcc.Input(id="BS3", placeholder="BS3"),
        dcc.Input(id="BS4", placeholder="BS4"),
        dcc.Input(id="BS5", placeholder="BS5"),
        dcc.Input(id="rainfall", placeholder="Rainfall"),
        dcc.Input(id="et", placeholder="ET"),
        html.Button("Calculate", id="submit"),
        html.H4("Water conservation diagram"),
        dcc.Graph(id="graph"),
        html.P("Opacity"),
        dcc.Slider(id="slider", min=0, max=1, value=0.5, step=0.1),
    ]
)


def get_excess_clipped_moist_ETb(moisture, bucket_size):
    excess = 0 if moisture <= bucket_size else moisture - bucket_size
    moist_out = moisture if moisture <= bucket_size else bucket_size
    moist_out = moist_out if moisture >= 0 else 0
    ETb = 0 if moisture >= 0 else -moisture
    return ETb, moist_out, excess


def model(prev_state, R, ET, buckets):
    old = {x + "_prev": y for x, y in prev_state.items()}

    # layer 0 to 5 cm
    SM1 = 0.2 * old["SM1_prev"] + R - 0.5 * ET
    Exc1, SM1, ETb1 = get_excess_clipped_moist_ETb(SM1, buckets["BS1"])

    # layer 5 to 15 cm
    SM2 = 0.95 * old["SM2_prev"] + 0.8 * old["SM1_prev"] - 0.2 * ETb1 + Exc1
    Exc2, SM2, ETb2 = get_excess_clipped_moist_ETb(SM2, buckets["BS2"])

    # Run Off
    if (SM1 > buckets["BS1"]) * (SM2 > buckets["BS2"]):
        Exc2 = 0
        run_off = Exc2
    else:
        run_off = 0

    # layer 15 to 30 cm
    SM3 = 0.95 * old["SM3_prev"] + 0.05 * old["SM2_prev"] - 0.15 * ETb2 + Exc2
    Exc3, SM3, ETb3 = get_excess_clipped_moist_ETb(SM3, buckets["BS3"])

    # layer 30 to 60 cm
    SM4 = 0.95 * old["SM4_prev"] + 0.05 * old["SM3_prev"] - 0.1 * ETb3 + Exc3
    Exc4, SM4, ETb4 = get_excess_clipped_moist_ETb(SM4, buckets["BS4"])

    # layer 60 to 100 cm
    SM5 = 0.99 * old["SM5_prev"] + 0.01 * old["SM4_prev"] - 0.05 * ETb4 + Exc4
    Exc5, SM5, ETb5 = get_excess_clipped_moist_ETb(SM5, buckets["BS5"])

    DD = 0.01 * old["SM5_prev"] + Exc5

    return [
        {
            "SM1": SM1,
            "SM2": SM2,
            "SM3": SM3,
            "SM4": SM4,
            "SM5": SM5,
            "DD": DD,
        },
        {
            "ETb1": ETb1,
            "ETb2": ETb2,
            "ETb3": ETb3,
            "ETb4": ETb4,
            "ETb5": ETb5,
        },
        {
            "Exc1": Exc1,
            "Exc2": Exc2,
            "Exc3": Exc3,
            "Exc4": Exc4,
            "Exc5": Exc5,
        },
        run_off,
    ]


@app.callback(
    Output("graph", "figure"),
    Input("slider", "value"),
    Input("SM1", "value"),
    Input("SM2", "value"),
    Input("SM3", "value"),
    Input("SM4", "value"),
    Input("SM5", "value"),
    Input("BS1", "value"),
    Input("BS2", "value"),
    Input("BS3", "value"),
    Input("BS4", "value"),
    Input("BS5", "value"),
    Input("rainfall", "value"),
    Input("et", "value"),
    Input("submit", "n_clicks"),
)
def display_sankey(
    opacity, SM1, SM2, SM3, SM4, SM5, BS1, BS2, BS3, BS4, BS5, R, ET, submit
):
    if submit == 0 or submit is None:
        raise PreventUpdate
    buckets = {
        "BS1": BS1,
        "BS2": BS2,
        "BS3": BS3,
        "BS4": BS4,
        "BS5": BS5,
    }
    prev_state = {
        "SM1": SM1,
        "SM2": SM2,
        "SM3": SM3,
        "SM4": SM4,
        "SM5": SM5,
    }
    SMs, ETbs, Excs, run_off = model(prev_state, R, ET, buckets)

    node = dict(
        pad=15,
        thickness=20,
        line=dict(color="black", width=0.5),
        label=[
            "SM1_prev",
            "SM2_prev",
            "SM3_prev",
            "SM4_prev",
            "SM5_prev",  # 0-4
            "SM1",
            "SM2",
            "SM3",
            "SM4",
            "SM5",  # 5-9
            "ETb1",
            "ETb2",
            "ETb3",
            "ETb4",
            "ETb5",  # 10-14
            "Exc1",
            "Exc2",
            "Exc3",
            "Exc4",
            "Exc5",  # 15-19
            "R",
            "ET",
            "DD",
            "runoff",
        ],  # 20, 21, 22, 23
        color="blue",
    )

    link = dict(
        source=[0, 0, 1, 1, 2, 2, 3, 3, 4, 4],  # indices correspond to labels
        target=[5, 6, 6, 7, 7, 8, 8, 9, 9, 22],
        value=[],
    )

    fig = go.Figure(go.Sankey(link=link, node=node))
    fig.update_layout(font_size=10)
    return fig


app.run_server(debug=True)