import xarray as xr from datetime import datetime, timedelta import pandas as pd import numpy as np import sys import os import torch import random from torch.utils import data import torch.nn.functional as F import matplotlib.pyplot as plt sfc_vars_1 = ['SSTK', 'TCW', 'TCWV', 'CP', 'MSL', 'TCC', 'U10M', 'V10M', 'T2M', 'TP', 'SKT'] sfc_vars_2 = ['sst', 'tcw', 'tcwv', 'cp', 'msl', 'tcc', 'u10', 'v10', 't2m', 'tp', 'skt'] pl_vars_1 = ["Z", "T", "Q", "W", "D", "U", "V"] pl_vars_2 = ["z", "t", "q", "w", "d", "u", "v"] var_map = {} for var1, var2 in zip(sfc_vars_1+pl_vars_1, sfc_vars_2+pl_vars_2): var_map[var1] = var2 class Aurora_CDF_Dataset_china(data.Dataset): """Dataset class for the era5 upper and surface variables.""" def __init__(self, nc_path='', seed=1234, startDate='2010', endDate='2020', freq='12h', horizon = 12, surface = ["2m_temperature","10m_u_component_of_wind","10m_v_component_of_wind", "mean_sea_level_pressure", "total_precipitation_6hr"], upper = ["temperature", "u_component_of_wind", "v_component_of_wind", "relative_humidity", "geopotential"], level = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000], h = 240, w = 240, degree = 0.25, num_points = 5, evaluate = False, ): """Initialize.""" self.nc_path = nc_path """ To do if start and end is valid date, if the date can be found in the downloaded files, length >= 0 """ # Prepare the datetime objects for training, validation, and test self.freq = int(freq[:-1]) self.surface_variables = surface self.upper_variables = upper self.levels = level self.horizon = horizon self.h = h self.w = w self.degree = degree self.num_points = num_points self.keys = list(pd.date_range(start=startDate, end=endDate, freq=freq))[:-1] random.seed(seed) def __getitem__(self, index): """Return input frames, target frames, and its corresponding time steps.""" key = self.keys[index] input_surfaces = [] input_uppers = [] time_points = [] for p in range(self.num_points): time_str = datetime.strftime(key, '%Y%m%d%H') time_points.append(key.timestamp()) data = np.load(os.path.join(self.nc_path,f"{time_str}.npy")).astype(np.float32) input_surface_variables = data[:,0] input_upper_variables = data[:,1:] input_surfaces.append(input_surface_variables[np.newaxis, ...]) input_uppers.append(input_upper_variables[np.newaxis, ...]) key = key + timedelta(hours=self.horizon) return np.concatenate(input_surfaces,axis=0), np.concatenate(input_uppers, axis=0), time_points def __len__(self): return len(self.keys) def __repr__(self): return self.__class__.__name__ class Dataset_Postprocessing(data.Dataset): def __init__(self, data_path='', seed=1234, start_date='1998-01-01', end_date='2015-12-31', val = False, surface = ['SSTK', 'TCW', 'TCWV', 'CP', 'MSL', 'TCC', 'U10M', 'V10M', 'T2M', 'TP', 'SKT'], upper = ["Z", "T", "Q", "W", "D", "U", "V"], levels = [500, 850], target_surface= ["T2M",'U10M', 'V10M'], target_upper = ["Z", "T"], target_level = [500, 850], H = 361, W = 720, full = False, ): """Initialize.""" self.data_path = data_path self.surface_variables = surface self.upper_variables = upper self.levels = levels self.target_sfc_variables = target_surface self.target_pl_variables = target_upper self.target_pl_level = target_level self.H = 361 self.W = 720 self.full = full time_range = slice(start_date, end_date) if self.full: self.ensemble_path = os.path.join(data_path, "ensemble") self.T = 10 self.sfc_value_range = {"sst": (260., 305.), "tcw": (0., 60.), "tcwv": (0., 60.), "cp": (0., 0.04), "msl": (97000., 1.1e5), "tcc": (0., 1.0), "u10": (-13., 11.), "v10": (-10., 15.), "t2m":(218, 304), "tp": (0., 0.07), "skt": (210., 310.)} self.pl_value_range = [{"z": (48200, 58000), "t": (230, 269), "q": (0., 4e-3), "w": (-0.7, 1.4), "d": (-5e-5, 8e-5), "u": (-7., 27.), "v": (-7., 7.)}, {"z": (10000, 15500), "t":(240, 299), "q": (0., 1.5e-2), "w": (-1.2, 1.8), "d": (-1.9e-4, 1.6e-4), "u": (-16., 17.5), "v": (-10., 16.)}] else: self.T = 2 self.sfc_ens_mean_normalized = xr.open_dataset(os.path.join(data_path, "ENS10_sfc_mean_normalized.nc"), engine="h5netcdf").sel(time=time_range) self.sfc_ens_std_normalized = xr.open_dataset(os.path.join(data_path, "ENS10_sfc_std_normalized.nc"), engine="h5netcdf").sel(time=time_range) self.sfc_ens_mean = xr.open_dataset(os.path.join(data_path, "ENS10_sfc_mean.nc"), engine="h5netcdf").sel(time=time_range) self.sfc_ens_std = xr.open_dataset(os.path.join(data_path, "ENS10_sfc_std.nc"), engine="h5netcdf").sel(time=time_range) self.pl_ens_mean_normalized = [xr.open_dataset(os.path.join(data_path, f"ENS10_pl_mean_{str(l)}_normalized.nc"), engine="h5netcdf").sel(time=time_range) for l in levels] self.pl_ens_std_normalized = [xr.open_dataset(os.path.join(data_path, f"ENS10_pl_std_{str(l)}_normalized.nc"), engine="h5netcdf").sel(time=time_range) for l in levels] self.pl_ens_mean = [xr.open_dataset(os.path.join(data_path, f"ENS10_pl_mean_{str(l)}.nc"), engine="h5netcdf").sel(time=time_range) for l in levels] self.pl_ens_std = [xr.open_dataset(os.path.join(data_path, f"ENS10_pl_std_{str(l)}.nc"), engine="h5netcdf").sel(time=time_range) for l in levels] self.era5 = xr.open_dataset(os.path.join(data_path, "ERA5.nc"), engine="h5netcdf").sel(time=time_range) if val: self.keys = self.sfc_ens_mean.drop_sel(time="2017-01-02").time.values # self.keys = [self.sfc_ens_mean.sel(time="2017-12-01").time.values] else: self.keys = self.sfc_ens_mean.time.values self.era5_sfc_scale = {} random.seed(seed) def __getitem__(self, index): """Return input frames, target frames, and its corresponding time steps.""" time_points = [] key = self.keys[index] time_points.append(key.item()) sfc_inputs = np.zeros((self.T, len(self.surface_variables), self.H, self.W)).astype(np.float32) pl_inputs = np.zeros((self.T, len(self.upper_variables), len(self.levels), self.H, self.W)).astype(np.float32) ds_targets = self.era5.sel(time=key) time_str = (pd.to_datetime(key)-timedelta(days=2)).strftime("%Y%m%d") sfc_targets = np.zeros((len(self.target_sfc_variables), self.H, self.W)).astype(np.float32) sfc_scale = np.zeros((2, len(self.target_sfc_variables), self.H, self.W)).astype(np.float32) if self.target_sfc_variables: for i in range(len(self.target_sfc_variables)): sfc_targets[i] = ds_targets[self.target_sfc_variables[i]].values.astype(np.float32) sfc_scale[0,i] = self.sfc_ens_mean[self.target_sfc_variables[i]].sel(time=key).values.astype(np.float32) sfc_scale[1,i] = self.sfc_ens_std[self.target_sfc_variables[i]].sel(time=key).values.astype(np.float32) pl_targets = np.zeros((len(self.target_pl_variables), self.H, self.W)).astype(np.float32) pl_scale = np.zeros((2, len(self.target_pl_variables), self.H, self.W)).astype(np.float32) if self.target_pl_variables: for i in range(len(self.target_pl_variables)): pl_targets[i] = ds_targets[self.target_pl_variables[i]].values.astype(np.float32)[0] pl_scale[0,i] = self.pl_ens_mean[self.levels.index(self.target_pl_level[i])][self.target_pl_variables[i]].sel(time=key,plev=self.target_pl_level[i]*1e2).values.astype(np.float32) pl_scale[1,i] = self.pl_ens_std[self.levels.index(self.target_pl_level[i])][self.target_pl_variables[i]].sel(time=key,plev=self.target_pl_level[i]*1e2).values.astype(np.float32) if self.full: sfc_ds = xr.open_dataset(os.path.join(self.ensemble_path, f"output.sfc.{time_str}.grib"), backend_kwargs={"indexpath":""}).fillna(9999.0) pl_ds = xr.open_dataset(os.path.join(self.ensemble_path, f"output.pl.{time_str}.grib"), backend_kwargs={"indexpath":""}).sel(isobaricInhPa=self.levels).fillna(9999.0) for i, var in enumerate(self.surface_variables): value = sfc_ds[var_map[var]].values.astype(np.float32) minval, maxval = self.sfc_value_range[var_map[var]] sfc_inputs[:,i] = (value - minval) / (maxval - minval) for i, var in enumerate(self.upper_variables): value = pl_ds[var_map[var]].values.astype(np.float32) for j in range(len(self.levels)): minval, maxval = self.pl_value_range[j][var_map[var]] pl_inputs[:,i,j] = (value[:,j] - minval) / (maxval - minval) else: for i, var in enumerate(self.surface_variables): sfc_inputs[1,i] = self.sfc_ens_mean_normalized[var].sel(time=key).values.astype(np.float32) sfc_inputs[0,i] = self.sfc_ens_std_normalized[var].sel(time=key).values.astype(np.float32) for i, var in enumerate(self.upper_variables): for j, l in enumerate(self.levels): pl_inputs[1,i,j] = self.pl_ens_mean_normalized[j][var].sel(time=key).values[0].astype(np.float32) pl_inputs[0,i,j] = self.pl_ens_std_normalized[j][var].sel(time=key).values[0].astype(np.float32) return sfc_inputs, pl_inputs, sfc_scale[:,:,:-1], sfc_targets[:,:-1], pl_scale[:,:,:-1],pl_targets[:,:-1], time_points def __len__(self): return len(self.keys) def __repr__(self): return self.__class__.__name__ class Aurora_downscale(data.Dataset): """Dataset class for the era5 upper and surface variables.""" def __init__(self, nc_path='', seed=1234, startDate='1999', endDate='2020', freq='6h', surface_prefix = ["2m_temperature","10m_u_component_of_wind","10m_v_component_of_wind"], upper_prefix = ["temperature", "u_component_of_wind", "v_component_of_wind", "specific_humidity", "geopotential"], surface = ["t2m","u10","v10"], upper = ["t", "u", "v", "q", "z"], levels = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000], num_points = 2, degree = ("5.625", "1.40625"), lr_dir = None, hr_dir = None, lr_degree = None, hr_degree = None, spatial_multiple = 4, ): """Initialize.""" self.nc_path = nc_path """ To do if start and end is valid date, if the date can be found in the downloaded files, length >= 0 """ # Prepare the datetime objects for training, validation, and test self.freq = int(freq[:-1]) self.surface_variables = surface self.upper_variables = upper self.levels = levels self.spatial_multiple = int(spatial_multiple) self.num_points = num_points self.keys = list(pd.date_range(start=startDate, end=endDate, freq=freq))[:-2] self.lr_datasets = {} self.hr_datasets = {} years = [datetime.strftime(key, '%Y') for key in list(pd.date_range(start=startDate[:4], end=endDate[:4], freq="YE"))] if not years: raise ValueError("No years were resolved from startDate/endDate.") if lr_degree is None: lr_degree = degree[0] if hr_degree is None: hr_degree = degree[1] # Resolve actual folder/file naming in datasets (e.g., 1.40625 vs 1.5_nc). sample_prefix = surface_prefix[0] sample_year = years[0] lr_dir, lr_degree = self._resolve_dataset_layout( nc_root=self.nc_path, prefix=sample_prefix, year=sample_year, requested_dir=lr_dir, requested_degree=lr_degree, candidates=(("5.625", "5.625"), ("5.625_nc", "5.625")), kind="LR", ) hr_dir, hr_degree = self._resolve_dataset_layout( nc_root=self.nc_path, prefix=sample_prefix, year=sample_year, requested_dir=hr_dir, requested_degree=hr_degree, candidates=( ("1.40625", "1.40625"), ("1.40625_nc", "1.40625"), ("1.5_nc", "1.5"), ("1.5", "1.5"), ), kind="HR", ) self.lr_dir = lr_dir self.hr_dir = hr_dir self.lr_degree = lr_degree self.hr_degree = hr_degree print(f"[Aurora_downscale] LR={self.lr_dir} ({self.lr_degree}deg), HR={self.hr_dir} ({self.hr_degree}deg)") for i, var in enumerate(self.surface_variables): for year in years: self.lr_datasets[f"{var}_{year}"] = xr.open_dataset( os.path.join(self.nc_path, self.lr_dir, surface_prefix[i], f'{surface_prefix[i]}_{year}_{self.lr_degree}deg.nc') ) self.hr_datasets[f"{var}_{year}"] = xr.open_dataset( os.path.join(self.nc_path, self.hr_dir, surface_prefix[i], f'{surface_prefix[i]}_{year}_{self.hr_degree}deg.nc') ) for i, var in enumerate(self.upper_variables): for year in years: self.lr_datasets[f"{var}_{year}"] = xr.open_dataset( os.path.join(self.nc_path, self.lr_dir, upper_prefix[i], f'{upper_prefix[i]}_{year}_{self.lr_degree}deg.nc') ) self.hr_datasets[f"{var}_{year}"] = xr.open_dataset( os.path.join(self.nc_path, self.hr_dir, upper_prefix[i], f'{upper_prefix[i]}_{year}_{self.hr_degree}deg.nc') ) random.seed(seed) def _trim_width_to_multiple(self, value): """Trim the last longitude columns if width is not divisible by `self.spatial_multiple`.""" if self.spatial_multiple <= 0: return value w = value.shape[-1] w_new = (w // self.spatial_multiple) * self.spatial_multiple if w_new == 0: return value if w_new == w: return value return value[..., :w_new] @staticmethod def _lat_name(da): if "latitude" in da.coords: return "latitude" if "lat" in da.coords: return "lat" raise KeyError("Latitude coordinate not found.") @staticmethod def _lon_name(da): if "longitude" in da.coords: return "longitude" if "lon" in da.coords: return "lon" raise KeyError("Longitude coordinate not found.") @staticmethod def _level_name(da): for c in ("level", "pressure_level", "plev"): if c in da.dims or c in da.coords: return c raise KeyError("Level coordinate not found for upper variable.") def _read_surface(self, ds, var, key): da = ds[var].sel(time=key) lat_name = self._lat_name(da) lon_name = self._lon_name(da) lat = da[lat_name].values if lat[0] < lat[-1]: da = da.isel({lat_name: slice(None, None, -1)}) da = da.transpose(lat_name, lon_name) value = da.values.astype(np.float32) value = self._trim_width_to_multiple(value) return value def _read_upper(self, ds, var, key): da = ds[var].sel(time=key) lev_name = self._level_name(da) lat_name = self._lat_name(da) lon_name = self._lon_name(da) lat = da[lat_name].values if lat[0] < lat[-1]: da = da.isel({lat_name: slice(None, None, -1)}) da = da.transpose(lev_name, lat_name, lon_name) value = da.values.astype(np.float32) value = self._trim_width_to_multiple(value) return value @staticmethod def _resolve_dataset_layout( nc_root, prefix, year, requested_dir, requested_degree, candidates, kind, ): if requested_dir is not None and requested_degree is None: for cand_dir, cand_degree in candidates: if cand_dir != requested_dir: continue expected = os.path.join(nc_root, requested_dir, prefix, f"{prefix}_{year}_{cand_degree}deg.nc") if os.path.exists(expected): return requested_dir, cand_degree if requested_degree is not None and requested_dir is None: for cand_dir, cand_degree in candidates: if cand_degree != requested_degree: continue expected = os.path.join(nc_root, cand_dir, prefix, f"{prefix}_{year}_{requested_degree}deg.nc") if os.path.exists(expected): return cand_dir, requested_degree if requested_dir is not None and requested_degree is not None: expected = os.path.join(nc_root, requested_dir, prefix, f"{prefix}_{year}_{requested_degree}deg.nc") if os.path.exists(expected): return requested_dir, requested_degree raise FileNotFoundError( f"[{kind}] Requested path not found: {expected}. " f"Check --{kind.lower()}_data_dir/--{kind.lower()}_degree_tag." ) for cand_dir, cand_degree in candidates: expected = os.path.join(nc_root, cand_dir, prefix, f"{prefix}_{year}_{cand_degree}deg.nc") if os.path.exists(expected): return cand_dir, cand_degree raise FileNotFoundError( f"[{kind}] Could not auto-resolve dataset layout for prefix={prefix}, year={year}. " f"Tried: {[(d, deg) for d, deg in candidates]}" ) def __getitem__(self, index): """Return input frames, target frames, and its corresponding time steps.""" key = self.keys[index] input_surfaces = [] input_uppers = [] target_surfaces = [] target_uppers = [] time_points = [] for _ in range(self.num_points): time_points.append(key.timestamp()) year = datetime.strftime(key, '%Y') input_surface_variables = [] input_upper_variables = [] for var in self.surface_variables: value = self._read_surface(self.lr_datasets[f"{var}_{year}"], var, key) input_surface_variables.append(value[np.newaxis, ...]) for var in self.upper_variables: value = self._read_upper(self.lr_datasets[f"{var}_{year}"], var, key) input_upper_variables.append(value[np.newaxis, ...]) input_surfaces.append(np.concatenate(input_surface_variables, axis=0)[np.newaxis, ...]) input_uppers.append(np.concatenate(input_upper_variables, axis=0)[np.newaxis, ...]) key = key + timedelta(hours=6) key = key - timedelta(hours=6) year = datetime.strftime(key, '%Y') for var in self.surface_variables: value = self._read_surface(self.hr_datasets[f"{var}_{year}"], var, key) target_surfaces.append(value[np.newaxis, ...]) for var in self.upper_variables: value = self._read_upper(self.hr_datasets[f"{var}_{year}"], var, key) target_uppers.append(value[np.newaxis, ...]) return np.concatenate(input_surfaces,axis=0), np.concatenate(input_uppers, axis=0), np.concatenate(target_surfaces,axis=0), np.concatenate(target_uppers, axis=0), time_points def __len__(self): return len(self.keys) def __repr__(self): return self.__class__.__name__ class AuroraFactorizedST(data.Dataset): """Dataset for factorized spatial/temporal SR with commutativity training.""" def __init__( self, nc_path="", seed=1234, startDate="1999", endDate="2020", freq="6h", surface_prefix=("2m_temperature", "10m_u_component_of_wind", "10m_v_component_of_wind"), upper_prefix=("temperature", "u_component_of_wind", "v_component_of_wind", "specific_humidity", "geopotential"), surface=("t2m", "u10", "v10"), upper=("t", "u", "v", "q", "z"), levels=(50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000), lr6h_dir="5.625", hr6h_dir="1.5_nc", hr1h_dir="1.5_1h_nc", lr1h_dir="", lr6h_degree="5.625", hr_degree="1.5", include_endpoints=True, derive_hr6h_from_hr1h=True, derive_lr1h_from_hr1h=True, ): self.nc_path = nc_path self.surface_prefix = list(surface_prefix) self.upper_prefix = list(upper_prefix) self.surface_variables = list(surface) self.upper_variables = list(upper) self.levels = tuple(levels) self.include_endpoints = include_endpoints self.derive_hr6h_from_hr1h = bool(derive_hr6h_from_hr1h) self.derive_lr1h_from_hr1h = bool(derive_lr1h_from_hr1h) self.lr1h_available = bool(lr1h_dir) and (not self.derive_lr1h_from_hr1h) if self.derive_hr6h_from_hr1h and (not self.include_endpoints): raise ValueError( "derive_hr6h_from_hr1h=True requires include_endpoints=True to access 6h endpoints." ) self.keys = list(pd.date_range(start=startDate, end=endDate, freq=freq))[:-2] self.hr_offsets = list(range(0, 7)) if include_endpoints else list(range(1, 6)) self.ds_lr6h = {} self.ds_hr6h = {} self.ds_hr1h = {} self.ds_lr1h = {} self._time_pos_cache = {} start_year = pd.Timestamp(startDate).year end_year = pd.Timestamp(endDate).year years = [str(y) for y in range(start_year, end_year + 1)] for year in years: for i, var in enumerate(self.surface_variables): pref = self.surface_prefix[i] self.ds_lr6h[f"{var}_{year}"] = xr.open_dataset( os.path.join(nc_path, lr6h_dir, pref, f"{pref}_{year}_{lr6h_degree}deg.nc") ) if not self.derive_hr6h_from_hr1h: self.ds_hr6h[f"{var}_{year}"] = xr.open_dataset( os.path.join(nc_path, hr6h_dir, pref, f"{pref}_{year}_{hr_degree}deg.nc") ) self.ds_hr1h[f"{var}_{year}"] = xr.open_dataset( os.path.join(nc_path, hr1h_dir, pref, f"{pref}_{year}_{hr_degree}deg.nc") ) if self.lr1h_available: self.ds_lr1h[f"{var}_{year}"] = xr.open_dataset( os.path.join(nc_path, lr1h_dir, pref, f"{pref}_{year}_{lr6h_degree}deg.nc") ) for i, var in enumerate(self.upper_variables): pref = self.upper_prefix[i] self.ds_lr6h[f"{var}_{year}"] = xr.open_dataset( os.path.join(nc_path, lr6h_dir, pref, f"{pref}_{year}_{lr6h_degree}deg.nc") ) if not self.derive_hr6h_from_hr1h: self.ds_hr6h[f"{var}_{year}"] = xr.open_dataset( os.path.join(nc_path, hr6h_dir, pref, f"{pref}_{year}_{hr_degree}deg.nc") ) self.ds_hr1h[f"{var}_{year}"] = xr.open_dataset( os.path.join(nc_path, hr1h_dir, pref, f"{pref}_{year}_{hr_degree}deg.nc") ) if self.lr1h_available: self.ds_lr1h[f"{var}_{year}"] = xr.open_dataset( os.path.join(nc_path, lr1h_dir, pref, f"{pref}_{year}_{lr6h_degree}deg.nc") ) random.seed(seed) @staticmethod def _lat_name(da): if "latitude" in da.coords: return "latitude" if "lat" in da.coords: return "lat" raise KeyError("Latitude coordinate not found.") @staticmethod def _lon_name(da): if "longitude" in da.coords: return "longitude" if "lon" in da.coords: return "lon" raise KeyError("Longitude coordinate not found.") @staticmethod def _level_name(da): for c in ("level", "pressure_level", "plev"): if c in da.dims or c in da.coords: return c raise KeyError("Level coordinate not found for upper variable.") @staticmethod def _year_str(key): return datetime.strftime(pd.Timestamp(key).to_pydatetime(), "%Y") def _dataset_for_time(self, ds_dict, var, key): year = self._year_str(key) ds_key = f"{var}_{year}" if ds_key not in ds_dict: raise KeyError(f"Dataset key {ds_key} is not available.") return ds_dict[ds_key] def _time_pos(self, ds, key): ds_id = id(ds) if ds_id not in self._time_pos_cache: if "time" not in ds.indexes: raise KeyError("Dataset does not have `time` index.") time_index = pd.DatetimeIndex(ds.indexes["time"]) self._time_pos_cache[ds_id] = {int(ts.value): i for i, ts in enumerate(time_index)} key_ns = int(pd.Timestamp(key).value) pos = self._time_pos_cache[ds_id].get(key_ns, None) if pos is None: raise KeyError(f"time={pd.Timestamp(key)} not found in dataset time index.") return pos def _standardize_lat(self, da): lat_name = self._lat_name(da) lat = da[lat_name].values if lat[0] < lat[-1]: da = da.isel({lat_name: slice(None, None, -1)}) return da def _read_surface(self, ds, var, key): t_pos = self._time_pos(ds, key) da = ds[var].isel(time=t_pos) da = self._standardize_lat(da) lat_name = self._lat_name(da) lon_name = self._lon_name(da) da = da.transpose(lat_name, lon_name) return da.values.astype(np.float32) def _read_upper(self, ds, var, key): t_pos = self._time_pos(ds, key) da = ds[var].isel(time=t_pos) da = self._standardize_lat(da) lev_name = self._level_name(da) lat_name = self._lat_name(da) lon_name = self._lon_name(da) if lev_name in da.coords: da = da.sel({lev_name: list(self.levels)}) da = da.transpose(lev_name, lat_name, lon_name) return da.values.astype(np.float32) def _interp_surface_to_lr(self, da_hr, lat_lr, lon_lr): da_hr = self._standardize_lat(da_hr) lat_name = self._lat_name(da_hr) lon_name = self._lon_name(da_hr) da_lr = da_hr.interp( {lat_name: xr.DataArray(lat_lr, dims=(lat_name,)), lon_name: xr.DataArray(lon_lr, dims=(lon_name,))}, method="linear", ) da_lr = da_lr.transpose(lat_name, lon_name) return da_lr.values.astype(np.float32) def _interp_upper_to_lr(self, da_hr, lat_lr, lon_lr): da_hr = self._standardize_lat(da_hr) lev_name = self._level_name(da_hr) lat_name = self._lat_name(da_hr) lon_name = self._lon_name(da_hr) da_lr = da_hr.interp( {lat_name: xr.DataArray(lat_lr, dims=(lat_name,)), lon_name: xr.DataArray(lon_lr, dims=(lon_name,))}, method="linear", ) da_lr = da_lr.transpose(lev_name, lat_name, lon_name) return da_lr.values.astype(np.float32) def __getitem__(self, index): key = self.keys[index] t0 = key t6 = key + timedelta(hours=6) times_6h = [t0, t6] times_1h = [t0 + timedelta(hours=h) for h in self.hr_offsets] x_lr6h_surface = [] y_hr6h_surface = [] y_hr1h_surface = [] y_lr1h_surface = [] if self.lr1h_available else None x_lr6h_upper = [] y_hr6h_upper = [] y_hr1h_upper = [] y_lr1h_upper = [] if self.lr1h_available else None for var in self.surface_variables: lr_stack = [self._read_surface(self._dataset_for_time(self.ds_lr6h, var, t), var, t) for t in times_6h] hr1_stack = [self._read_surface(self._dataset_for_time(self.ds_hr1h, var, t), var, t) for t in times_1h] if self.derive_hr6h_from_hr1h: hr6_stack = [hr1_stack[0], hr1_stack[-1]] else: hr6_stack = [self._read_surface(self._dataset_for_time(self.ds_hr6h, var, t), var, t) for t in times_6h] if self.lr1h_available: lr1_stack = [self._read_surface(self._dataset_for_time(self.ds_lr1h, var, t), var, t) for t in times_1h] else: lr1_stack = None x_lr6h_surface.append(np.stack(lr_stack, axis=0)) y_hr6h_surface.append(np.stack(hr6_stack, axis=0)) y_hr1h_surface.append(np.stack(hr1_stack, axis=0)) if y_lr1h_surface is not None: y_lr1h_surface.append(np.stack(lr1_stack, axis=0)) for var in self.upper_variables: lr_stack = [self._read_upper(self._dataset_for_time(self.ds_lr6h, var, t), var, t) for t in times_6h] hr1_stack = [self._read_upper(self._dataset_for_time(self.ds_hr1h, var, t), var, t) for t in times_1h] if self.derive_hr6h_from_hr1h: hr6_stack = [hr1_stack[0], hr1_stack[-1]] else: hr6_stack = [self._read_upper(self._dataset_for_time(self.ds_hr6h, var, t), var, t) for t in times_6h] if self.lr1h_available: lr1_stack = [self._read_upper(self._dataset_for_time(self.ds_lr1h, var, t), var, t) for t in times_1h] else: lr1_stack = None x_lr6h_upper.append(np.stack(lr_stack, axis=0)) y_hr6h_upper.append(np.stack(hr6_stack, axis=0)) y_hr1h_upper.append(np.stack(hr1_stack, axis=0)) if y_lr1h_upper is not None: y_lr1h_upper.append(np.stack(lr1_stack, axis=0)) # [T, V, ...] layout x_lr6h_surface = np.stack(x_lr6h_surface, axis=1).astype(np.float32) y_hr6h_surface = np.stack(y_hr6h_surface, axis=1).astype(np.float32) y_hr1h_surface = np.stack(y_hr1h_surface, axis=1).astype(np.float32) if y_lr1h_surface is not None: y_lr1h_surface = np.stack(y_lr1h_surface, axis=1).astype(np.float32) x_lr6h_upper = np.stack(x_lr6h_upper, axis=1).astype(np.float32) y_hr6h_upper = np.stack(y_hr6h_upper, axis=1).astype(np.float32) y_hr1h_upper = np.stack(y_hr1h_upper, axis=1).astype(np.float32) if y_lr1h_upper is not None: y_lr1h_upper = np.stack(y_lr1h_upper, axis=1).astype(np.float32) time_unix = np.asarray([t.timestamp() for t in times_1h], dtype=np.int64) out = { "x_lr6h_surface": x_lr6h_surface, "x_lr6h_upper": x_lr6h_upper, "y_hr6h_surface": y_hr6h_surface, "y_hr6h_upper": y_hr6h_upper, "y_hr1h_surface": y_hr1h_surface, "y_hr1h_upper": y_hr1h_upper, "time_unix": time_unix, } if y_lr1h_surface is not None: out["y_lr1h_surface"] = y_lr1h_surface if y_lr1h_upper is not None: out["y_lr1h_upper"] = y_lr1h_upper return out def __len__(self): return len(self.keys) def __repr__(self): return self.__class__.__name__