def get_obs( obs: str, train_period: slice, variables: Union[str, List[str]], chunking_approach: Optional[str] = None, cache_within_rechunk: Optional[bool] = True, ) -> xr.Dataset: if obs == 'ERA5': ds_obs = open_era5(variables=variables, time_period=train_period) else: raise NotImplementedError( 'only ERA5 is available as observation dataset right now') if chunking_approach is None: return ds_obs if cache_within_rechunk: path_dict = { 'obs': obs, 'train_period': train_period, 'variables': variables, } rechunked_path = make_rechunked_obs_path( chunking_approach=chunking_approach, **path_dict, ) else: rechunked_path = None ds_obs_rechunked = rechunk_zarr_array_with_caching( zarr_array=ds_obs, chunking_approach=chunking_approach, output_path=rechunked_path, ) return ds_obs_rechunked
def return_gcm_train_full_time(coarse_obs_full_time_ds: xr.Dataset, gcm: str, scenario: str, variable: str, train_period: slice, predict_period: slice, bbox: BBox, **kwargs) -> xr.Dataset: """Returns GCM training rechunked dataset in full time. Parameters ---------- coarse_obs_full_time_ds : xr.Dataset Output coarse observation dataset rechunked in full_time gcm : str Input GCM scenario: str Input GCM scenario variable: str The variable included in the dataset. train_period: slice Start and end year slice of training/historical period. Ex: slice('1990','1990') predict_period: slice Start and end year slice of predict period. Ex: slice('2020','2020') bbox : BBox Bounding box including latmin,latmax,lonmin,lonmax. **kwargs : dict, optional Returns ------- xr.Dataset x_train rechunked dataset in full time. """ gcm_train_ds = load_cmip(source_ids=gcm, variable_ids=[variable], return_type='xr') gcm_train_ds_180 = lon_to_180(gcm_train_ds) gcm_train_ds_subset = subset_dataset( gcm_train_ds_180, variable, train_period, bbox, chunking_schema={ 'time': 365, 'lat': 150, 'lon': 150 }, ) # this call was to force the timestamps for the cmip data to use the friendlier era5 timestamps. (i forget which dataset used which time formats). i could picture this introducing a tricky bug though (for instance if gcm timestamp didn't align for some reason) so we could use another conversion system if that is better. Perhaps datetime equivilence test. gcm_train_ds_subset['time'] = coarse_obs_full_time_ds.time.values gcm_train_subset_full_time_ds = rechunk_zarr_array_with_caching( gcm_train_ds_subset, chunking_approach='full_time', max_mem='1GB', ) return gcm_train_subset_full_time_ds
def maca_coarse_bias_correction_task( ds_gcm: xr.Dataset, ds_obs: xr.Dataset, train_period_start: str, train_period_end: str, variables: Union[str, List[str]], chunking_approach: str, batch_size: Optional[int] = 15, buffer_size: Optional[int] = 15, **kwargs, ): """ Task that implements the coarse scale bias correction in MACA. The historical GCM is mapped to historical coarsened observation in the bias correction. Rechunks the GCM data to match observation data because the bias correction model in skdownscale requires these datasets to have the same chunks/blocks. ds_gcm: xr.Dataset GCM dataset ds_obs: xr.Dataset Observation dataset train_period_start: str Start year of training/historical period train_period_end: str End year of training/historical period variables: List[str] Names of the variables used in obs and gcm dataset (including features and label) chunking_approach: str 'full_space', 'full_time', 'matched' or None batch_size: Optional[int] The batch size in terms of day of year to bias correct together buffer_size: Optional[int] The buffer size in terms of day of year to include in the bias correction Returns ------- bias_corrected: xr.Dataset Bias corrected GCM dataset """ # TODO: test if this is needed if both ds_gcm and ds_obs are chunked in full time ds_gcm_rechunked = rechunk_zarr_array_with_caching( zarr_array=ds_gcm, template_chunk_array=ds_obs) historical_period = slice(train_period_start, train_period_end) bias_corrected = maca_bias_correction( ds_gcm=ds_gcm_rechunked, ds_obs=ds_obs, historical_period=historical_period, variables=variables, batch_size=batch_size, buffer_size=buffer_size, ) return bias_corrected
def prep_gard_input_task( obs: str, train_period: slice, predict_period: slice, variables: List[str], X_train: xr.Dataset, X_pred: xr.Dataset, gcm_identifier: str, bias_correction_method: str, bbox, ): # get observation data in the same chunking scheme as ds_obs = return_obs(obs=obs, train_period=train_period, variable=variables, bbox=bbox) rechunked_obs_path = make_rechunked_obs_path( obs=obs, train_period=train_period, variable=variables, bbox=bbox, chunking_approach='matched', ) y_train_rechunked = rechunk_zarr_array_with_caching( zarr_array=ds_obs, template_chunk_array=X_train, output_path=rechunked_obs_path) rechunked_gcm_path = make_bias_corrected_gcm_path( gcm_identifier=gcm_identifier, method=bias_correction_method, chunking_approach='matched') X_pred_rechunked = rechunk_zarr_array_with_caching( zarr_array=X_pred.sel(time=predict_period), template_chunk_array=X_train, output_path=rechunked_gcm_path, ) return X_train, y_train_rechunked, X_pred_rechunked
def rechunker_task( zarr_array: xr.Dataset, chunking_approach: Optional[str] = None, template_chunk_array: Optional[xr.Dataset] = None, naming_func: Optional[Callable] = None, **kwargs, ): """ Task to rechunk a dataset Parameters ---------- zarr_array : zarr or xarray dataset Dataset you want to rechunk. chunking_approach : str, optional Has to be one of `full_space` or `full_time`. If `full_space`, the data will be rechunked such that the space dimensions are contiguous (i.e. each chunk will contain full maps). If `full_time`, the data will be rechunked such that the time dimension is contiguous (i.e. each chunk will contain full time series). Either the chunking approach or the template chunk array must be provided. template_chunk_array: zarr or xarray dataset, optional A template dataset with the desired chunksizes. Either the chunking approach or the template chunk array must be provided. naming_func: callable, optional A function that returns a string that represents the output caching location that the rechunk task should save to. The input arguments of this naming func must be provided as kwargs to this method Returns ------- rechunked_ds : xr.Dataset Rechunked dataset """ if naming_func is not None: output_path = naming_func(chunking_approach=chunking_approach, **kwargs) else: output_path = None rechunked = rechunk_zarr_array_with_caching( zarr_array=zarr_array, chunking_approach=chunking_approach, template_chunk_array=template_chunk_array, output_path=output_path, ) return rechunked
def return_coarse_obs_full_time(coarse_obs_ds: xr.Dataset, gcm: str, scenario: str, variable: str, train_period: slice, predict_period: slice, bbox: BBox, **kwargs) -> xr.Dataset: """ Return coarse observation dataset that has been chunked in time. Parameters ---------- coarse_obs_ds : xr.Dataset Input coarse observation dataset gcm : str Input GCM scenario: str Input GCM scenario variable: str The variable included in the dataset. train_period: slice Start and end year slice of training/historical period. Ex: slice('1990','1990') predict_period: slice Start and end year slice of predict period. Ex: slice('2020','2020') bbox : BBox Bounding box including latmin,latmax,lonmin,lonmax. **kwargs : dict, optional Returns ------- xr.Dataset coarse_obs_full_time_ds rechunked dataset """ coarse_obs_full_time_ds = rechunk_zarr_array_with_caching( coarse_obs_ds, chunking_approach='full_time', max_mem='1GB', ) return coarse_obs_full_time_ds
def interpolate_gcm_task( obs: str, gcm: str, scenario: str, train_period: slice, predict_period: slice, variables: Union[str, List[str]], chunking_approach: str, bbox, ): """ Interpolate the GCM dataset to the grid of the observation dataset. Rechunk the final output according to chunking approach. Parameters ---------- obs: str Name of obs dataset gcm: str Name of the GCM model scenario: str Name of the emission scenario training_period: slice Training/historical period bounds predict_period: slice Prediction (historical and/or future) period bounds variables: List[str] List of variables to get in obs dataset chunking_approach: str 'full_space', 'full_time', or None Returns ------- ds_gcm_interpolated_rechunked: xr.Dataset The GCM dataset that has been interpolated to the obs grid then rechunked. """ # get gcm in full space chunks ds_gcm_full_space = get_gcm( gcm=gcm, scenario=scenario, variables=variables, train_period=train_period, predict_period=predict_period, chunking_approach="full_space", cache_within_rechunk=False, bbox=bbox, ) # get obs as a template ds_obs = return_obs(obs=obs, train_period=train_period, variable=variables, bbox=bbox) # interpolate gcm to obs resolution ds_gcm_interpolated = regrid_ds( ds=ds_gcm_full_space, target_grid_ds=ds_obs.isel(time=0).load(), chunking_approach="full_space", ) # rechunked to final output chunking approach if needed ds_gcm_interpolated_rechunked = rechunk_zarr_array_with_caching( zarr_array=ds_gcm_interpolated, output_path=None, chunking_approach=chunking_approach, ) return ds_gcm_interpolated_rechunked
def coarsen_and_interpolate_obs_task(obs, train_period, predict_period, variables, gcm, scenario, chunking_approach, bbox, **kwargs): """ Coarsen the observation dataset to the grid of the GCM model specified in inputs then interpolate back into the observation grid. Rechunk the final output according to chunking approach. Parameters ---------- obs: str Name of obs dataset gcm: str Name of GCM model training_period_start: str Start year of training/historical period training_period_end: str End year of training/historical period variables: List[str] List of variables to get in obs dataset chunking_approach: str 'full_space', 'full_time', or None **kwargs: Dict Other arguments to be used in generating the target path Returns ------- ds_obs_interpolated_rechunked: xr.Dataset An observation dataset that has been coarsened, interpolated back to original grid, and then rechunked. """ # get obs ds_obs = return_obs(obs=obs, train_period=train_period, variable=variables, bbox=bbox) # regrid to coarse scale ds_obs_coarse = get_coarse_obs( obs_ds=ds_obs, gcm=gcm, scenario=scenario, variable=variables, train_period=train_period, predict_period=predict_period, bbox=bbox, **kwargs, ) # interpolate to fine scale again ds_obs_interpolated = regrid_ds( ds=ds_obs_coarse, target_grid_ds=ds_obs.isel(time=0).chunk({ 'lat': -1, 'lon': -1 }), chunking_approach="full_space", ) # rechunked to final output chunking approach if needed ds_obs_interpolated_rechunked = rechunk_zarr_array_with_caching( zarr_array=ds_obs_interpolated, output_path=None, chunking_approach=chunking_approach, ) return ds_obs_interpolated_rechunked
def fit_and_predict(gcm_train_subset_full_time_ds: xr.Dataset, coarse_obs_full_time_ds: xr.Dataset, gcm_predict_rechunked_ds: xr.Dataset, gcm: str, scenario: str, variable: str, train_period: slice, predict_period: slice, bbox: BBox, dim: str = "time", **kwargs) -> xr.Dataset: """Fit bcsd model on prepared CMIP data with obs at corresponding spatial scale. Then predict for a set of CMIP data (likely future). Parameters ---------- gcm_train_subset_full_time_ds : xr.Dataset GCM training dataset chunked along space coarse_obs_full_time_ds : xr.Dataset Obs training dataset chunked along space gcm_predict_rechunked_ds : xr.Dataset GCM prediction dataset chunked along space. gcm : str Input GCM scenario: str Input GCM scenario variable: str The variable included in the dataset. train_period: slice Start and end year slice of training/historical period. Ex: slice('1990','1990') predict_period: slice Start and end year slice of predict period. Ex: slice('2020','2020') bbox : BBox Bounding box including latmin,latmax,lonmin,lonmax. dim : str, optional dimension on which you want to do the modelling, by default "time" **kwargs : dict, optional Returns ------- bias_corrected_ds : xr.Dataset Bias-corrected dataset """ if variable in ABSOLUTE_VARS: bcsd_model = BcsdTemperature(return_anoms=False) elif variable in RELATIVE_VARS: bcsd_model = BcsdPrecipitation(return_anoms=False) pointwise_model = PointWiseDownscaler(model=bcsd_model, dim=dim) coarse_obs_rechunked_validated_ds = rechunk_zarr_array_with_caching( coarse_obs_full_time_ds, template_chunk_array=gcm_train_subset_full_time_ds) pointwise_model.fit(gcm_train_subset_full_time_ds[variable], coarse_obs_rechunked_validated_ds[variable]) bias_corrected_da = pointwise_model.predict( gcm_predict_rechunked_ds[variable]) bias_corrected_ds = bias_corrected_da.astype('float32').to_dataset( name=variable) return bias_corrected_ds
def get_gcm( gcm: str, scenario: str, variables: Union[str, List[str]], train_period: slice, predict_period: slice, bbox, chunking_approach: Optional[str] = None, cache_within_rechunk: Optional[bool] = True, ) -> xr.Dataset: """ Load and combine historical and future GCM into one dataset. Parameters ---------- gcm: str Name of GCM scenario: str Name of scenario variables: str or list Name of variable(s) to load train_period_start: str Start year of train/historical period train_period_end: str End year of train/historical period predict_period_start: str Start year of predict/future period predict_period_end: str End year of predict/future period chunking_approach: Optional[str] 'full_space', 'full_time', or None Returns ------- ds_gcm: xr.Dataset A dataset containing both historical and future period of GCM data """ historical_gcm = load_cmip( activity_ids='CMIP', experiment_ids='historical', source_ids=gcm, variable_ids=variables, return_type='xr', ) future_gcm = load_cmip( activity_ids='ScenarioMIP', experiment_ids=scenario, source_ids=gcm, variable_ids=variables, return_type='xr', ) ds_gcm = xr.combine_by_coords([historical_gcm, future_gcm], combine_attrs='drop_conflicts') ds_gcm_train = subset_dataset( ds=ds_gcm, variable=variables[0], time_period=train_period, bbox=bbox, ) ds_gcm_predict = subset_dataset( ds=ds_gcm, variable=variables[0], time_period=predict_period, bbox=bbox, ) ds_gcm = xr.combine_by_coords([ds_gcm_train, ds_gcm_predict], combine_attrs='drop_conflicts') ds_gcm = ds_gcm.reindex(time=sorted(ds_gcm.time.values)) if chunking_approach is None: return ds_gcm if cache_within_rechunk: path_dict = { 'gcm': gcm, 'scenario': scenario, 'train_period': train_period, 'predict_period': predict_period, 'variables': variables, } rechunked_path = make_rechunked_gcm_path(chunking_approach=chunking_approach, **path_dict) else: rechunked_path = None ds_gcm_rechunked = rechunk_zarr_array_with_caching( zarr_array=ds_gcm, chunking_approach=chunking_approach, output_path=rechunked_path, ) return ds_gcm_rechunked