data_utils.py

import os
import tempfile
import traceback
import streamlit as st
import xarray as xr
from typing import List
import numpy as np

@st.cache_resource
def save_uploaded_files(uploaded_files):
    if 'temp_file_paths' not in st.session_state:
        st.session_state.temp_file_paths = []
    for uploaded_file in uploaded_files:
        suffix = os.path.splitext(uploaded_file.name)[1]
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
        temp_file.write(uploaded_file.read())
        temp_file.close()
        st.session_state.temp_file_paths.append(temp_file.name)
        

@st.cache_resource
def load_dataset(file_paths: List[str]):
    try:
        ds = xr.open_mfdataset(file_paths, combine='by_coords').load()
        return ds
    except Exception:
        st.error("Error loading dataset:")
        st.error(traceback.format_exc())
        return None
    
@st.cache_resource
def load_dataset_pangu(file_path: str):
    try:
        ds = np.load(file_path)
        return ds
    except Exception:
        st.error("Error loading dataset:")
        st.error(traceback.format_exc())
        return None