fengwu_utils.py

import streamlit as st
import torch
# from Pangu-Weather import *
import numpy as np
from datetime import datetime
import numpy as np
import onnx
import onnxruntime as ort
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import io

def fengwu_config_data():
    st.subheader("FengWu Model Data Input")

    # Detailed data description section
    st.markdown("""
    **Input Data Requirements (FengWu):**  
    FengWu takes **two consecutive six-hour atmospheric states** as input:
    1. **First Input (input1.npy)**: Atmospheric data at the initial time.  
    2. **Second Input (input2.npy)**: Atmospheric data 6 hours later.

    **Shape & Variables:**  
    Each input is a NumPy array with shape `(69, 721, 1440)`:
    - **Dimension 0 (69 features):**  
      The first 4 features are surface variables:  
      1. U10 (10-meter Eastward Wind)  
      2. V10 (10-meter Northward Wind)  
      3. T2M (2-meter Temperature)  
      4. MSL (Mean Sea Level Pressure)

      These are followed by non-surface variables, each with 13 pressure levels:  
      - Z (Geopotential)  
      - Q (Specific Humidity)  
      - U (Eastward Wind)  
      - V (Northward Wind)  
      - T (Temperature)  

      The 13 vertical levels are: [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000] hPa

      The total count is:  
      - Surface vars: 4  
      - For each non-surface var (Z, Q, U, V, T): 13 levels = 65 vars  
      4 (surface) + 65 (5 vars * 13 levels) = 69 total features.

    **Spatial & Coordinate Details:**  
    - Latitude dimension (721 points) ranges from 90°N to -90°S with ~0.25° spacing.
    - Longitude dimension (1440 points) ranges from 0° to 360°E with ~0.25° spacing.
    - Ensure data is single precision floats (`.astype(np.float32)`).

    **Data Frequency & Forecasting Scheme:**  
    - `input1.npy` corresponds to a given time (e.g., 06:00 UTC Jan 1, 2018).
    - `input2.npy` corresponds to 6 hours later (e.g., 12:00 UTC Jan 1, 2018).
    - The model predicts future states at subsequent 6-hour intervals.

    **Converting Your Data:**  
    - ERA5 `.nc` files or ECMWF `.grib` files can be converted to `.npy` using appropriate Python packages (`netCDF4` or `pygrib`).
    - Ensure you follow the exact variable and level ordering as described.


    """)

    # File uploaders for FengWu input data (two consecutive time steps)
    st.markdown("### Upload Your FengWu Input Data Files")
    input1_file = st.file_uploader(
        "Upload input1.npy (Initial Time)",
        type=["npy"],
        key="fengwu_input1"
    )

    input2_file = st.file_uploader(
        "Upload input2.npy (6 Hours Later)",
        type=["npy"],
        key="fengwu_input2"
    )

    st.markdown("---")
    st.markdown("### References & Resources")
    st.markdown("""
    - **Research Paper:** [FengWu: Pushing the Skillful Global Medium-range Weather Forecast beyond 10 Days Lead](https://arxiv.org/abs/2304.02948)
                
    - **GitHub Source Code:** [Fengwu on GitHub](https://github.com/OpenEarthLab/FengWu?tab=readme-ov-file)
    """)

    return input1_file, input2_file


@st.cache_resource
def inference_6hrs_fengwu(input1, input2):
    model_6 = onnx.load('FengWu/fengwu_v2.onnx')

    # Set the behavier of onnxruntime
    options = ort.SessionOptions()
    options.enable_cpu_mem_arena=False
    options.enable_mem_pattern = False
    options.enable_mem_reuse = False
    # Increase the number for faster inference and more memory consumption
    options.intra_op_num_threads = 1

    # Set the behavier of cuda provider
    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}

    # Initialize onnxruntime session for Pangu-Weather Models
    ort_session_6 = ort.InferenceSession('FengWu/fengwu_v2.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])


    data_mean = np.load("FengWu/data_mean.npy")[:, np.newaxis, np.newaxis]
    data_std = np.load("FengWu/data_std.npy")[:, np.newaxis, np.newaxis]

    input1_after_norm = (input1 - data_mean) / data_std
    input2_after_norm = (input2 - data_mean) / data_std
    input = np.concatenate((input1_after_norm, input2_after_norm), axis=0)[np.newaxis, :, :, :]
    input = input.astype(np.float32)

    output = ort_session_6.run(None, {'input':input})[0]
    output = (output[0, :69] * data_std) + data_mean

    return output


@st.cache_resource
def inference_12hrs_fengwu(input1, input2):
    model_6 = onnx.load('FengWu/fengwu_v2.onnx')

    # Set the behavier of onnxruntime
    options = ort.SessionOptions()
    options.enable_cpu_mem_arena=False
    options.enable_mem_pattern = False
    options.enable_mem_reuse = False
    # Increase the number for faster inference and more memory consumption
    options.intra_op_num_threads = 1

    # Set the behavier of cuda provider
    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}

    # Initialize onnxruntime session for Pangu-Weather Models
    ort_session_6 = ort.InferenceSession('FengWu/fengwu_v2.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])


    data_mean = np.load("FengWu/data_mean.npy")[:, np.newaxis, np.newaxis]
    data_std = np.load("FengWu/data_std.npy")[:, np.newaxis, np.newaxis]

    input1_after_norm = (input1 - data_mean) / data_std
    input2_after_norm = (input2 - data_mean) / data_std
    input = np.concatenate((input1_after_norm, input2_after_norm), axis=0)[np.newaxis, :, :, :]
    input = input.astype(np.float32)

    for i in range(2):
        output = ort_session_6.run(None, {'input':input})[0]
        input = np.concatenate((input[:, 69:], output[:, :69]), axis=1)
        output = (output[0, :69] * data_std) + data_mean
        # print(output.shape)

    return output

@st.cache_resource
def inference_custom_hrs_fengwu(input1, input2, forecast_hours):
    model_6 = onnx.load('FengWu/fengwu_v2.onnx')

    # Set the behavier of onnxruntime
    options = ort.SessionOptions()
    options.enable_cpu_mem_arena=False
    options.enable_mem_pattern = False
    options.enable_mem_reuse = False
    # Increase the number for faster inference and more memory consumption
    options.intra_op_num_threads = 1

    # Set the behavier of cuda provider
    cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}

    # Initialize onnxruntime session for Pangu-Weather Models
    ort_session_6 = ort.InferenceSession('FengWu/fengwu_v2.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])


    data_mean = np.load("FengWu/data_mean.npy")[:, np.newaxis, np.newaxis]
    data_std = np.load("FengWu/data_std.npy")[:, np.newaxis, np.newaxis]

    input1_after_norm = (input1 - data_mean) / data_std
    input2_after_norm = (input2 - data_mean) / data_std
    input = np.concatenate((input1_after_norm, input2_after_norm), axis=0)[np.newaxis, :, :, :]
    input = input.astype(np.float32)

    for i in range(forecast_hours/6):
        output = ort_session_6.run(None, {'input':input})[0]
        input = np.concatenate((input[:, 69:], output[:, :69]), axis=1)
        output = (output[0, :69] * data_std) + data_mean
        # print(output.shape)

    return output

def plot_fengwu_output(initial_data, predicted_data):
    """
    Plot initial and predicted Fengwu model outputs.
    
    Parameters:
    - initial_data: np.ndarray of shape (69, 721, 1440) representing the initial or input state.
    - predicted_data: np.ndarray of shape (69, 721, 1440) representing the predicted state by Fengwu.
    """
    # Coordinate setup
    lat = np.linspace(90, -90, 721)   # Latitude from 90N to 90S
    lon = np.linspace(0, 360, 1440)   # Longitude from 0E to 360E

    # Surface and upper-level variable definitions
    surface_vars = ["U10", "V10", "T2M", "MSL"]
    upper_vars = ["Z (Geopotential)", "Q (Specific Humidity)", "U (Eastward Wind)", "V (Northward Wind)", "T (Temperature)"]
    upper_levels = [50,100,150,200,250,300,400,500,600,700,850,925,1000]

    # Mapping of upper variable groups to their starting indices
    # Each group has 13 levels, so indices shift by 13 for each subsequent group.
    var_group_start = {
        "Z (Geopotential)": 4,   # Z starts at index 4
        "Q (Specific Humidity)": 17,  # Q = 4+13=17
        "U (Eastward Wind)": 30, # U = 17+13=30
        "V (Northward Wind)": 43,# V = 30+13=43
        "T (Temperature)": 56    # T = 43+13=56
    }

    # --- Initial Data Visualization ---
    st.subheader("Initial Data Visualization (Fengwu)")
    init_col1, init_col2 = st.columns([1,1])

    with init_col1:
        init_data_choice = st.selectbox("Data Source", ["Upper-Air Data", "Surface Data"], key="fengwu_init_data_choice")
    with init_col2:
        if init_data_choice == "Upper-Air Data":
            init_var = st.selectbox("Variable", upper_vars, key="fengwu_init_upper_var")
        else:
            init_var = st.selectbox("Variable", surface_vars, key="fengwu_init_surface_var")

    # Select the data slice for initial data
    if init_data_choice == "Upper-Air Data":
        selected_level_hpa_init = st.select_slider(
            "Select Pressure Level (hPa)",
            options=upper_levels,
            value=850,  # Default to 850hPa
            help="Select the pressure level in hPa.",
            key="fengwu_init_level_hpa_slider"
        )
        level_index_init = upper_levels.index(selected_level_hpa_init)
        start_index_init = var_group_start[init_var]
        data_index_init = start_index_init + level_index_init
        data_to_plot_init = initial_data[data_index_init, :, :]
        title_init = f"Initial Upper-Air: {init_var} at {selected_level_hpa_init}hPa"
    else:
        # Surface variable
        var_index_init = surface_vars.index(init_var)
        data_to_plot_init = initial_data[var_index_init, :, :]
        title_init = f"Initial Surface: {init_var}"

    # Plot initial data
    fig_init, ax_init = plt.subplots(figsize=(10, 5), subplot_kw={'projection': ccrs.PlateCarree()})
    ax_init.set_title(title_init)
    im_init = ax_init.imshow(data_to_plot_init, extent=[lon.min(), lon.max(), lat.min(), lat.max()],
                             origin='lower', cmap='coolwarm', transform=ccrs.PlateCarree())
    ax_init.coastlines()
    plt.colorbar(im_init, ax=ax_init, orientation='horizontal', pad=0.05)
    st.pyplot(fig_init)

    # --- Predicted Data Visualization ---
    st.subheader("Predicted Data Visualization (Fengwu)")
    pred_col1, pred_col2 = st.columns([1,1])

    with pred_col1:
        pred_data_choice = st.selectbox("Data Source", ["Upper-Air Data", "Surface Data"], key="fengwu_pred_data_choice")
    with pred_col2:
        if pred_data_choice == "Upper-Air Data":
            pred_var = st.selectbox("Variable", upper_vars, key="fengwu_pred_upper_var")
        else:
            pred_var = st.selectbox("Variable", surface_vars, key="fengwu_pred_surface_var")

    # Select the data slice for predicted data
    if pred_data_choice == "Upper-Air Data":
        selected_level_hpa_pred = st.select_slider(
            "Select Pressure Level (hPa)",
            options=upper_levels,
            value=850,  # Default to 850hPa
            help="Select the pressure level in hPa.",
            key="fengwu_pred_level_hpa_slider"
        )
        level_index_pred = upper_levels.index(selected_level_hpa_pred)
        start_index_pred = var_group_start[pred_var]
        data_index_pred = start_index_pred + level_index_pred
        data_to_plot_pred = predicted_data[data_index_pred, :, :]
        title_pred = f"Predicted Upper-Air: {pred_var} at {selected_level_hpa_pred}hPa"
    else:
        # Surface variable for predicted data
        var_index_pred = surface_vars.index(pred_var)
        data_to_plot_pred = predicted_data[var_index_pred, :, :]
        title_pred = f"Predicted Surface: {pred_var}"

    # Plot predicted data
    fig_pred, ax_pred = plt.subplots(figsize=(10, 5), subplot_kw={'projection': ccrs.PlateCarree()})
    ax_pred.set_title(title_pred)
    im_pred = ax_pred.imshow(data_to_plot_pred, extent=[lon.min(), lon.max(), lat.min(), lat.max()],
                             origin='lower', cmap='coolwarm', transform=ccrs.PlateCarree())
    ax_pred.coastlines()
    plt.colorbar(im_pred, ax=ax_pred, orientation='horizontal', pad=0.05)
    st.pyplot(fig_pred)

    # --- Download Buttons ---
    st.subheader("Download Predicted Fengwu Data")

    # Convert predicted_data to binary format for download
    buffer_pred = io.BytesIO()
    np.save(buffer_pred, predicted_data)
    buffer_pred.seek(0)

    st.download_button(label="Download Predicted Fengwu Data",
                       data=buffer_pred,
                       file_name="predicted_fengwu.npy",
                       mime="application/octet-stream")