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 pangu_config_data(): st.subheader("Pangu-Weather Model Data Input") # Detailed data description section st.markdown(""" **Input Data Requirements:** Pangu-Weather uses two NumPy arrays to represent initial atmospheric conditions: 1. **Surface Data (input_surface.npy)** - Shape: `(4, 721, 1440)` - Variables: MSLP, U10, V10, T2M in this exact order. - **MSLP:** Mean Sea Level Pressure - **U10:** 10-meter Eastward Wind - **V10:** 10-meter Northward Wind - **T2M:** 2-meter Temperature 2. **Upper-Air Data (input_upper.npy)** - Shape: `(5, 13, 721, 1440)` - Variables (first dim): Z, Q, T, U, V in this exact order - **Z:** Geopotential (Note: if your source provides geopotential height, multiply by 9.80665 to get geopotential) - **Q:** Specific Humidity - **T:** Temperature - **U:** Eastward Wind - **V:** Northward Wind - Pressure Levels (second dim): 1000hPa, 925hPa, 850hPa, 700hPa, 600hPa, 500hPa, 400hPa, 300hPa, 250hPa, 200hPa, 150hPa, 100hPa, 50hPa in this exact order. **Spatial & Coordinate Details:** - Latitude dimension (721 points) ranges from 90°N to -90°S with a 0.25° spacing. - Longitude dimension (1440 points) ranges from 0° to 359.75°E with a 0.25° spacing. - Data should be single precision floats (`.astype(np.float32)`). **Supported Data Sources:** - ERA5 initial fields (strongly recommended). - ECMWF initial fields (e.g., HRES forecast) can be used, but may result in a slight accuracy drop. - Other types of initial fields are not currently supported due to potentially large discrepancies in data fields. **Converting Your Data:** - ERA5 `.nc` files can be converted to `.npy` using the `netCDF4` Python package. - ECMWF `.grib` files can be converted to `.npy` using the `pygrib` Python package. - Ensure the order of variables and pressure levels is exactly as described above. """) # File uploaders for surface and upper data separately st.markdown("### Upload Your Input Data Files") input_surface_file = st.file_uploader( "Upload input_surface.npy", type=["npy"], key="pangu_input_surface" ) input_upper_file = st.file_uploader( "Upload input_upper.npy", type=["npy"], key="pangu_input_upper" ) st.markdown("---") st.markdown("### References & Resources") st.markdown(""" - **Research Paper:** [Accurate medium-range global weather forecasting with 3D neural networks](https://www.nature.com/articles/s41586-023-06185-3) - [Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast](https://arxiv.org/abs/2211.02556) - **GitHub Source Code:** [Pangu-Weather on GitHub](https://github.com/198808xc/Pangu-Weather?tab=readme-ov-file) """) return input_surface_file, input_upper_file def inference_24hrs(input, input_surface): model_24 = onnx.load('Pangu-Weather/pangu_weather_24.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_24 = ort.InferenceSession('Pangu-Weather/pangu_weather_24.onnx', sess_options=options, providers=['CPUExecutionProvider']) # Run the inference session output, output_surface = ort_session_24.run(None, {'input':input, 'input_surface':input_surface}) return output, output_surface @st.cache_resource def inference_6hrs(input, input_surface): model_6 = onnx.load('Pangu-Weather/pangu_weather_6.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('Pangu-Weather/pangu_weather_6.onnx', sess_options=options, providers=['CPUExecutionProvider']) # Run the inference session output, output_surface = ort_session_6.run(None, {'input':input, 'input_surface':input_surface}) return output, output_surface @st.cache_resource def inference_1hr(input, input_surface): model_1 = onnx.load('Pangu-Weather/pangu_weather_1.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_1 = ort.InferenceSession('Pangu-Weather/pangu_weather_1.onnx', sess_options=options, providers=['CPUExecutionProvider']) # Run the inference session output, output_surface = ort_session_1.run(None, {'input':input, 'input_surface':input_surface}) return output, output_surface @st.cache_resource def inference_3hrs(input, input_surface): model_3 = onnx.load('Pangu-Weather/pangu_weather_3.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_3 = ort.InferenceSession('Pangu-Weather/pangu_weather_3.onnx', sess_options=options, providers=['CPUExecutionProvider']) # Run the inference session output, output_surface = ort_session_3.run(None, {'input':input, 'input_surface':input_surface}) return output, output_surface @st.cache_resource def inference_custom_hrs(input, input_surface, forecast_hours): # Ensure forecast_hours is a multiple of 24 if forecast_hours % 24 != 0: raise ValueError("forecast_hours must be a multiple of 24.") # Load the 24-hour model model_24 = onnx.load('Pangu-Weather/pangu_weather_24.onnx') # Configure ONNX Runtime session options = ort.SessionOptions() options.enable_cpu_mem_arena = False options.enable_mem_pattern = False options.enable_mem_reuse = False options.intra_op_num_threads = 1 # Using CPUExecutionProvider for simplicity ort_session_24 = ort.InferenceSession('Pangu-Weather/pangu_weather_24.onnx', sess_options=options, providers=['CPUExecutionProvider']) # Calculate how many 24-hour steps we need steps = forecast_hours // 24 # Run the 24-hour model repeatedly for i in range(steps): output, output_surface = ort_session_24.run(None, {'input': input, 'input_surface': input_surface}) input, input_surface = output, output_surface # Return the final predictions after completing all steps return input, input_surface def plot_pangu_output(upper_data, surface_data, out_upper, out_surface): # Coordinate setup lat = np.linspace(90, -90, 721) # Latitude grid lon = np.linspace(0, 360, 1440) # Longitude grid # Variable and level names upper_vars = ["Z (Geopotential)", "Q (Specific Humidity)", "T (Temperature)", "U (Eastward Wind)", "V (Northward Wind)"] upper_levels = ["1000hPa", "925hPa", "850hPa", "700hPa", "600hPa", "500hPa", "400hPa", "300hPa", "250hPa", "200hPa", "150hPa", "100hPa", "50hPa"] # Extract numeric hPa values for selection upper_hpa_values = [int(l.replace("hPa", "")) for l in upper_levels] surface_vars = ["MSLP", "U10", "V10", "T2M"] # --- Initial Data Visualization --- st.subheader("Initial Data Visualization") init_col1, init_col2 = st.columns([1,1]) with init_col1: init_data_choice = st.selectbox("Data Source", ["Upper-Air Data", "Surface Data"], key="init_data_choice") with init_col2: if init_data_choice == "Upper-Air Data": init_var = st.selectbox("Variable", upper_vars, key="init_upper_var") else: init_var = st.selectbox("Variable", surface_vars, key="init_surface_var") if init_data_choice == "Upper-Air Data": selected_level_hpa_init = st.select_slider( "Select Pressure Level (hPa)", options=upper_hpa_values, value=850, # Default to 850hPa help="Select the pressure level in hPa.", key="init_level_hpa_slider" ) # Find the corresponding index from the selected hPa value selected_level_index_init = upper_hpa_values.index(selected_level_hpa_init) selected_var_index_init = upper_vars.index(init_var) data_to_plot_init = upper_data[selected_var_index_init, selected_level_index_init, :, :] title_init = f"Initial Upper-Air: {init_var} at {selected_level_hpa_init}hPa" else: selected_var_index_init = surface_vars.index(init_var) data_to_plot_init = surface_data[selected_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") pred_col1, pred_col2 = st.columns([1,1]) with pred_col1: pred_data_choice = st.selectbox("Data Source", ["Upper-Air Data", "Surface Data"], key="pred_data_choice") with pred_col2: if pred_data_choice == "Upper-Air Data": pred_var = st.selectbox("Variable", upper_vars, key="pred_upper_var") else: pred_var = st.selectbox("Variable", surface_vars, key="pred_surface_var") if pred_data_choice == "Upper-Air Data": selected_level_hpa_pred = st.select_slider( "Select Pressure Level (hPa)", options=upper_hpa_values, value=850, # Default to 850hPa help="Select the pressure level in hPa.", key="pred_level_hpa_slider" ) selected_level_index_pred = upper_hpa_values.index(selected_level_hpa_pred) selected_var_index_pred = upper_vars.index(pred_var) data_to_plot_pred = out_upper[selected_var_index_pred, selected_level_index_pred, :, :] title_pred = f"Predicted Upper-Air: {pred_var} at {selected_level_hpa_pred}hPa" else: selected_var_index_pred = surface_vars.index(pred_var) data_to_plot_pred = out_surface[selected_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 Data") # Convert out_upper and out_surface to binary format for download buffer_upper = io.BytesIO() np.save(buffer_upper, out_upper) buffer_upper.seek(0) buffer_surface = io.BytesIO() np.save(buffer_surface, out_surface) buffer_surface.seek(0) st.download_button(label="Download Predicted Upper-Air Data", data=buffer_upper, file_name="predicted_upper.npy", mime="application/octet-stream") st.download_button(label="Download Predicted Surface Data", data=buffer_surface, file_name="predicted_surface.npy", mime="application/octet-stream")