''' QUANTILE DELTA MAPPING: GCM TO ERA-5 The code bias-corrects CMIP6 GCM data to an ERA-5 reference dataset using quantile delta mapping. See the README for more information. The script is set up to run in parallel to speed up processing time. A single dataset takes ~15 min to process on 64 CPUs (16 workers with 4 threads per worker) or ~80 minutes without parallelization. Note that these are large datasets. If working with small amounts of RAM, you may need to edit the lat/ lon chunk sizes. @author: clairevwilson ''' # import libraries import os, time, gc import numpy as np import xarray as xr import cftime import dask from dask.distributed import Client, LocalCluster, wait from scipy.stats import rankdata import logging logging.getLogger("distributed").setLevel(logging.ERROR) # ==================== QDM OPTIONS ==================== # define variables to run VARIABLES = ['tas','pr'] # define reference period REF_START = '1980-01-01' REF_END = '2019-12-31' # number of years to subset timeseries LOOP_YEARS = 30 # flag to test this code on a smaller data subset TESTING = False # ==================== PARALLELIZATION ==================== PARALLELIZE = False # split up on parallel CPUs? if not PARALLELIZE: N_WORKERS = THREADS_PER_WORKER = 1 MEM_LIMIT = '128GB' # will run out of memory if this is too low else: # set up manually for the number of available CPUs N_WORKERS = 16 THREADS_PER_WORKER = 4 MEM_LIMIT = '30GB' # ==================== INPUT DATA ==================== # setup file names DATA_FP = '/trace/group/rounce/cvwilson/climate_data/GCMs/' ERA_FILES = { # example format given here for temperature and precip # 'CMIP variable name': ('ERA-5 file name', 'ERA-5 variable name') 'tas': ('ERA5_temp_monthly.nc', 't2m'), 'pr': ('ERA5_totalprecip_monthly.nc', 'tp'), } # list of GCMs based on the contents of DATA_FP GCMS = [d for d in os.listdir(DATA_FP) if os.path.isdir(os.path.join(DATA_FP, d)) and d != 'bias_corrected'] # list of SSPs for each GCM SSPS = ['126','370','534-over','585'] # ==================== INITIALIZE ==================== # subset GCMs and SSPs if testing if TESTING: GCMS = [GCMS[0]] SSPS = [SSPS[0]] lat_chunks = lon_chunks = 1 time_chunks = 1 else: # if using small amounts of RAM, need smaller chunks lat_chunks = lon_chunks = 50 time_chunks = -1 # define chunk dictionaries chunks_reftime = {'ref_time':time_chunks, 'lat':lat_chunks, 'lon':lon_chunks} chunks = {'time':time_chunks, 'lat':lat_chunks, 'lon':lon_chunks} # ==================== MAPPING FUNCTION ==================== def qdm_core_logic(subset, hist_gcm_sorted, hist_era_sorted, precip_threshold=0.001): """ subset: GCM Projected (FP) hist_gcm: GCM Historic (FH) hist_era: ERA5 Historic (FR) precip_threshold: threshold below which to perform additive correction in meters / day """ # calculate the rank (percentile q(t)) for every value in the current window q_t = (rankdata(subset, method='max') - 0.5) / len(subset) q_t = np.clip(q_t, 0.00001, 0.99999) # Protect extremes # get idx corresponding to each percentile and grab from the sorted data idx = (q_t * (hist_gcm_sorted.shape[-1] - 1)).astype(int) fh_inv = hist_gcm_sorted[idx] fr_inv = hist_era_sorted[idx] # cap fh_inv so don't have 0 in the denominator # only matters for precip: temp will never be below 1e-6 K eps = 1e-6 fh_inv_safe = np.maximum(fh_inv, eps) # future / historical change delta_t = subset / fh_inv_safe # calculate Bias Adjusted Climate x(t) = FR * Delta(t) x_t = fr_inv * delta_t # special additive handling for dry days # only matters for precip: temp will never be below 0.001 K dry = subset < precip_threshold delta_t[dry] = subset[dry] - fh_inv[dry] x_t[dry] = fr_inv[dry] + delta_t[dry] # ensure values remain positive x_t[x_t < 0] = 0 return x_t.astype(np.float32) # ==================== TEMPORAL SUBSET FUNCTION ==================== def get_windows(gcm_all, window_size=LOOP_YEARS): """ Breaks up GCM data into LOOP_YEARS windows. """ # get information on time in the GCM data end_year = int(gcm_all.time[-1].item().year) curr_start = int(gcm_all.time[0].item().year) windows = [] # build windows while curr_start <= end_year: curr_end = curr_start + window_size - 1 if curr_end > end_year: # the tail: back-shift to get percentiles # corresponding to a full LOOP_YEARS period # e.g. period 2080 to 2100 will use distribution # 2070 to 2100 and then crop the data to 2080 dist_s = end_year - window_size + 1 dist_e = end_year windows.append((dist_s, dist_e, curr_start)) break else: # standard 30-year block windows.append((curr_start, curr_end, curr_start)) curr_start += window_size return windows # ==================== MAIN FUNCTION ==================== def run_qdm_for_var(var): global LOOP_YEARS # get ERA file path and variable name era_fn, era_var = ERA_FILES[var] era_path = os.path.join(DATA_FP, era_fn) # main loop for gcm in GCMS: for ssp in SSPS: # ==================== SETUP ==================== # get the output directory out_dir = os.path.join(DATA_FP, f'bias_corrected/{gcm}') os.makedirs(out_dir, exist_ok=True) # get the output filepath and check if it's already built out_fn = f'bias_corrected/{gcm}/{gcm}_ssp{ssp}_r1i1p1f1_{var}_era5corrected_TEST.nc' out_path = os.path.join(DATA_FP, out_fn) if os.path.exists(out_path): print(out_path, 'already exists, skipping...', flush=True) continue # get GCM file name gcm_fn = f'{gcm}/{gcm}_ssp{ssp}_r1i1p1f1_{var}.nc' gcm_path = os.path.join(DATA_FP, gcm_fn) if not os.path.exists(gcm_path): if ssp != '534-over': print(gcm_path, 'not found, skipping...', flush=True) # else, 534-over does not exist for this GCM (expected so don't need to print) continue # set up cores such that sufficient RAM is available cluster = LocalCluster( n_workers=N_WORKERS, threads_per_worker=THREADS_PER_WORKER, memory_limit=MEM_LIMIT, dashboard_address=":0" ) # initiate cluster with Client(cluster) as client: start_time = time.time() # ==================== ERA-5 DATA PREPARATION ==================== # load ERA-5 data into memory ds_era = (xr.open_dataset(era_path, chunks={}) .rename({'longitude': 'lon', 'latitude': 'lat'})) # select the reference period era_ref = (ds_era[era_var].sel(time=slice(REF_START, REF_END), expver=1)) # for testing choose a smaller subset if TESTING: era_ref = era_ref.isel(lat=slice(2), lon=slice(2), time=slice(48)) # process for ufunc era_ref = (era_ref .drop_vars(['time','expver']) .rename({'time': 'ref_time'}) .chunk(chunks_reftime) .map_blocks(lambda x: x.copy( data=np.sort(x.data, axis=x.get_axis_num('ref_time')) )) .persist() ) wait(era_ref) # ==================== GCM DATA PREPARATION ==================== # open GCM and get the attributes to store in the output file ds_gcm = xr.open_dataset(gcm_path, decode_times=False, chunks={}) ds_gcm = xr.decode_cf(ds_gcm, use_cftime=True) gcm_attrs = ds_gcm.attrs var_attrs = ds_gcm[var].attrs # check if units are known original_gcm_units = ds_gcm[var].attrs['units'] addstr = '\nAdd unit conversion code manually' assert original_gcm_units in ['K','kg m-2 s-1'], f'! {gcm} has unknown units: {original_gcm_units}{addstr}' gcm_calendar = ds_gcm.time.encoding.get('calendar', ds_gcm.time.attrs.get('calendar', 'standard')) # transform precip units if original_gcm_units == 'kg m-2 s-1': # need to convert kg m-2 s-1 to m w.e. / day SECONDS_IN_DAY = 24 * 3600 DENSITY_WATER = 1000 # kg m-3 ds_gcm = ds_gcm / DENSITY_WATER * SECONDS_IN_DAY # move to ERA-5 grid using nearest neighbors gcm_interp = (ds_gcm[var] .interp(lat=era_ref.lat, lon=era_ref.lon, method='nearest', kwargs={"fill_value": "extrapolate"})) # handle chunking after interpolation gcm_all = gcm_interp.chunk(chunks).persist() wait(gcm_all) # get GCM reference period gcm_ref = (gcm_all.sel(time=slice(REF_START, REF_END))) # for testing use a smaller subset if TESTING: gcm_ref = gcm_ref.isel(lat=slice(2), lon=slice(2), time=slice(48)) gcm_all = gcm_all.isel(lat=slice(2), lon=slice(2), time=slice(48)) LOOP_YEARS = 4 # process for ufunc gcm_ref = (gcm_ref .drop_vars('time') .rename({'time': 'ref_time'}) .chunk(chunks_reftime) .map_blocks(lambda x: x.copy( data=np.sort(x.data, axis=x.get_axis_num('ref_time')) )) .persist() ) # ==================== QUANTILE DELTA MAPPING ==================== # get LOOP_YEARS windows to subset data windows = get_windows(gcm_all, LOOP_YEARS) # initialize storage for subsets temp_files = [] # loop through subsets for per_start, per_end, keep_s in windows: # get filename for temporary storage temp_fn = f'temp_{gcm}_{ssp}_{per_start}.nc' temp_path = os.path.join(DATA_FP, temp_fn) # get temporal subset keep_start = f'{keep_s}-01-01' start = f'{per_start}-01-01' end = f'{per_end}-12-31' gcm_subset = gcm_all.sel(time=slice(start, end)) # map the subset subset_corrected = xr.apply_ufunc( qdm_core_logic, gcm_subset, gcm_ref, era_ref, input_core_dims=[['time'], ['ref_time'], ['ref_time']], output_core_dims=[['time']], vectorize=True, dask='parallelized', output_dtypes=[gcm_subset.dtype], ) # if needed, transform units back if original_gcm_units == 'kg m-2 s-1': subset_corrected = (subset_corrected * DENSITY_WATER) / SECONDS_IN_DAY gcm_attrs['units'] = 'kg m-2 s-1' # clip to the keep period (only matters for tail) subset_final = (subset_corrected .sel(time=slice(keep_start, end)) .to_dataset(name=var)) # ensure time is in cftime subset_final = subset_final.assign_coords(time=[ cftime.datetime(t.year, t.month, t.day, t.hour, t.minute) for t in subset_final.time.values ]) # define temporal encoding so all subset files are consistent temp_encoding = { var: {'zlib': False}, 'time': { 'units': 'days since 1850-01-01', 'calendar': gcm_calendar, 'dtype': 'float64' } } # store this temporal subset (will be deleted later) subset_final.to_netcdf(temp_path, encoding=temp_encoding, engine='h5netcdf', compute=True) temp_files.append(temp_path) # close the subset dataset subset_final.close() # ==================== STORE AND CLEANUP ==================== # concatenate temp files with xr.open_mfdataset(temp_files, combine='nested', decode_times=False, engine='h5netcdf', concat_dim='time') as qm_result: # get time in cftime and chunk the results qm_result = xr.decode_cf(qm_result, use_cftime=True) ds_out = qm_result.chunk(chunks) # reformat dataarray to dataset and add back all attributes ds_out.attrs = gcm_attrs ds_out[var].attrs = var_attrs ds_out = ds_out.transpose('time', 'lat', 'lon') # make sure time coordinate is a consistent data type ds_out = ds_out.assign_coords(time=ds_out.time.values.astype(object)) # apply the encoding with compression and chunking ds_out.time.encoding = {} encoding = {var: {'zlib': True, 'complevel': 1, 'chunksizes': (100, chunks['lat'], chunks['lon'])}, 'time': { 'units': 'days since 1850-01-01', 'calendar': gcm_calendar, 'dtype': 'float64' }} # store the output ds_out.to_netcdf(out_path, engine='h5netcdf', encoding=encoding) # delete the temporary files for f in temp_files: if os.path.exists(f): os.remove(f) # print the timer time_elapsed = (time.time() - start_time) / 60 print(f'Finished {out_fn} in {time_elapsed:.1f} minutes', flush=True) print('======================================================================================', flush=True) # ==================== EXECUTE QUANTILE DELTA MAPPING ==================== if '__main__' in __name__: # loop through both variables for var in VARIABLES: print('======================================================================================', flush=True) run_qdm_for_var(var)