import os import xarray as xr import pandas as pd from compute_et0_adjusted import compute_et0 # Mapping of variable names to metadata (title, unit, and NetCDF variable key) VARIABLE_MAPPING = { 'surface_downwelling_shortwave_radiation': ('Surface Downwelling Shortwave Radiation', 'W/m²', 'rsds'), 'moisture_in_upper_portion_of_soil_column': ('Moisture in Upper Portion of Soil Column', 'kg m-2', 'mrsos'), 'precipitation': ('Precipitation', 'kg m-2 s-1', 'pr'), 'near_surface_relative_humidity': ('Relative Humidity', '%', 'hurs'), 'evaporation_including_sublimation_and_transpiration': ('Evaporation (including sublimation and transpiration)', 'kg m-2 s-1', 'evspsbl'), 'total_runoff': ('Total Runoff', 'kg m-2 s-1', 'mrro'), 'daily_minimum_near_surface_air_temperature': ('Daily Minimum Near Surface Air Temperature', '°C', 'tasmin'), 'daily_maximum_near_surface_air_temperature': ('Daily Maximum Near Surface Air Temperature', '°C', 'tasmax'), 'near_surface_wind_speed': ('Near Surface Wind Speed', 'm/s', 'sfcWind'), 'near_surface_air_temperature': ('Near Surface Air Temperature', '°C', 'tas'), } # Function to load data for a given variable from the dataset at the nearest latitude and longitude def load_data(variable: str, ds: xr.Dataset, latitude: float, longitude: float) -> xr.DataArray: """ Load data for a given variable from the dataset at the nearest latitude and longitude. Args: variable (str): The variable to extract from the dataset. ds (xr.Dataset): The xarray dataset containing climate data. latitude(float): Latitude for nearest data point. longitude (float): Longitude for nearest data point. Returns: xr.DataArray: The data array containing the variable values for the specified location. """ try: data = ds[variable].sel(lat=latitude, lon=longitude, method="nearest") # Convert temperature from Kelvin to Celsius for specific variables if variable in ["tas", "tasmin", "tasmax"]: data = data - 273.15 return data except Exception as e: print(f"Error loading {variable}: {e}") return None # Function to load forecast datasets from NetCDF files based on variable mapping def get_forecast_datasets(climate_sub_files: list) -> dict: """ Get the forecast datasets by loading NetCDF files for each variable. Args: climate_sub_files (list): List of file paths to the NetCDF files. Returns: dict: Dictionary with variable names as keys and xarray datasets as values. """ datasets = {} for file_path in climate_sub_files: filename = os.path.basename(file_path) for long_name, (title, unit, var_key) in VARIABLE_MAPPING.items(): if var_key in filename: if var_key in ["tas", "tasmax", "tasmin"]: if f"_{var_key}_" in f"_{filename}_" or filename.endswith(f"_{var_key}.nc"): datasets[long_name] = xr.open_dataset(file_path, engine="netcdf4") else: datasets[long_name] = xr.open_dataset(file_path, engine="netcdf4") return datasets # Function to extract climate data from forecast datasets and convert to a DataFrame def get_forecast_data(latitude: float, longitude: float, scenario: str, shading_coef: float = 0) -> pd.DataFrame: """ Extract climate data from the forecast datasets for a given location and convert to a DataFrame. Args: latitude(float): Latitude of the location to extract data for. longitude (float): Longitude of the location to extract data for. scenario (str): The scenario to extract data for. shading_coef (float, optional): Shading coefficient to use. Defaults to 0 (for no shading).. Returns: pd.DataFrame: A DataFrame containing time series data for each variable. """ assert scenario in ["moderate", "pessimist"] assert 0 <= shading_coef <= 1 # Define the directory to parse folder_to_parse = f"data/climate_data_{scenario}/" # Retrieve the subfolders and files to parse climate_sub_folder = [os.path.join(folder_to_parse, e) for e in os.listdir(folder_to_parse) if os.path.isdir(os.path.join(folder_to_parse, e))] climate_sub_files = [os.path.join(e, i) for e in climate_sub_folder for i in os.listdir(e) if i.endswith('.nc')] # Load the forecast datasets datasets = get_forecast_datasets(climate_sub_files) time_series_data = {'time': []} for long_name, (title, unit, variable) in VARIABLE_MAPPING.items(): print(f"Processing {long_name} ({title}, {unit}, {variable})...") data = load_data(variable, datasets[long_name], latitude, longitude) if data is not None: time_series_data['time'] = data.time.values column_name = f"{title} ({unit})" time_series_data[column_name] = data.values forecast_data = pd.DataFrame(time_series_data) forecast_data = preprocess_forectast_data(forecast_data, latitude, longitude, shading_coef) return forecast_data def preprocess_forectast_data(df: pd.DataFrame, latitude, longitude, shading_coef) -> pd.DataFrame: assert 0 <= shading_coef <= 1 preprocessed_data = df.copy() preprocessed_data["irradiance"] = preprocessed_data['Surface Downwelling Shortwave Radiation (W/m²)'] * (1 - shading_coef) preprocessed_data["air_temperature_min"] = preprocessed_data['Daily Minimum Near Surface Air Temperature (°C)'] preprocessed_data["air_temperature_max"] = preprocessed_data['Daily Maximum Near Surface Air Temperature (°C)'] preprocessed_data["relative_humidity_min"] = preprocessed_data['Relative Humidity (%)'] preprocessed_data["relative_humidity_max"] = preprocessed_data['Relative Humidity (%)'] preprocessed_data["wind_speed"] = preprocessed_data['Near Surface Wind Speed (m/s)'] # Convert 'time' to datetime and calculate Julian day preprocessed_data['time'] = pd.to_datetime(preprocessed_data['time'], errors='coerce') preprocessed_data['month'] = preprocessed_data['time'].dt.month preprocessed_data['day_of_year'] = preprocessed_data['time'].dt.dayofyear # Compute ET0 et0 = compute_et0(preprocessed_data, latitude, longitude) preprocessed_data['Evaporation (mm/day)'] = et0.clip(lower=0) # Convert Precipitation from kg/m²/s to mm/day preprocessed_data['Precipitation (mm/day)'] = 86400 * preprocessed_data['Precipitation (kg m-2 s-1)'] # Calculate Water Deficit: Water Deficit = ET0 - P + M preprocessed_data['Water Deficit (mm/day)'] = ( preprocessed_data['Evaporation (mm/day)'] - preprocessed_data['Precipitation (mm/day)'] # + preprocessed_data['Moisture in Upper Portion of Soil Column (kg m-2)']) ) return preprocessed_data # Main processing workflow def main(): # Get the forecast data for a specific latitude and longitude latitude, longitude = 47.0, 5.0 # Example coordinates scenario = "pessimist" shading_coef = 0 forecast_data = get_forecast_data(latitude, longitude, scenario=scenario, shading_coef=shading_coef) # Display the resulting DataFrame print(forecast_data.head()) print(forecast_data.columns) return forecast_data # Run the main function if __name__ == "__main__": main()