PyCafe - Solara - Altair River Flow Forecast Visualization

app.py
requirements.txt
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import solara
from solara.website.utils import apidoc
import pandas as pd
import altair as alt

async def install_packages():
    import micropip
    await micropip.install("zarr")
    await micropip.install("s3fs")

async def get_forecast_data(river_number, date):
    import zarr
    import s3fs
    
    s3_bucket_url = f's3://geoglows-v2-forecasts/{date}.zarr/'
    s3 = s3fs.S3FileSystem(anon=True)

    try:
        mapper = s3fs.S3Map(root=s3_bucket_url, s3=s3, check=False)
        zarr_group = zarr.open_group(mapper, mode='r')

        rivid_array = zarr_group['rivid'][:]
        if river_number in rivid_array:
            river_index = list(rivid_array).index(river_number)
            qout_array = zarr_group['Qout'][:, :, river_index]
            time_array = zarr_group['time'][:]
            ensemble_array = zarr_group['ensemble'][:]

            forecast_df = pd.DataFrame(qout_array, columns=time_array)
            forecast_df.index = [f"ensemble_{i}" for i in ensemble_array]
            forecast_df.columns = time_array
            forecast_df = forecast_df.transpose()

            return forecast_df
        else:
            return pd.DataFrame()
    except Exception as e:
        print(f"Error accessing data for RiverNumber {river_number} on {date}: {e}")
        return pd.DataFrame()

return_periods = {
    "2_year": 100,
    "5_year": 200,
    "10_year": 300
}

river_number = 110123714
date = '2024040100'

forecast_data = await get_forecast_data(river_number, date)

if forecast_data.empty:
    df = pd.DataFrame(columns=["time", "ensemble", "flow"])
else:
    forecast_data.reset_index(inplace=True)
    forecast_data = forecast_data.melt(id_vars='index', var_name='ensemble', value_name='flow')
    forecast_data.rename(columns={'index': 'time'}, inplace=True)
    forecast_data['time'] = pd.to_datetime(forecast_data['time'])
    df = forecast_data

@solara.component
def Page():
    click_data, set_click_data = solara.use_state(None)
    hover_data, set_hover_data = solara.use_state(None)

    line_chart = alt.Chart(df).mark_line().encode(
        x='time:T',
        y='flow:Q',
        color='ensemble:N'
    ).properties(
        width=800,
        height=400,
        title=f'Forecast data for RiverNumber {river_number} on {date}'
    )

    for period, value in return_periods.items():
        rp_line = alt.Chart(pd.DataFrame({
            'time': [df['time'].min(), df['time'].max()],
            'flow': [value, value]
        })).mark_rule(color='red', strokeDash=[4, 4]).encode(
            x='time:T',
            y='flow:Q'
        ).properties(
            title=period
        )
        line_chart += rp_line

    with solara.Div() as main:
        solara.AltairChart(line_chart, on_click=set_click_data, on_hover=set_hover_data)
        solara.Markdown(
            f"""
Click data:


Hover data:
            """
        )
    return main

__doc__ += apidoc(solara.FigureAltair.f)

await install_packages()
Py.Cafe

altair-river-flow-forecast

Altair River Flow Forecast Visualization
