GARD post processing

random_effects.py

def add_random_effect(ds, ds_rand, pop_thresh=None, t_root=1., p_root=3.):
    '''Add random effects to dataset `ds`

    Parameters
    ----------
    ds : xarray.Dataset
        Input dataset with variables `tas` and/or `pr`
    ds_rand : xarray.Dataset
        Input dataset with random fields with variable names `t_rand` and `p_rand`
    pop_thresh : float
        Threshold value for POP, default: 0.5
    t_root : float
        Root used to transform temperature
    p_root : float
        Root used to transform precipitation

    Returns
    -------
    ds : xarray.Dataset
        Updated dataset (`ds`) including variables with random effects added in.
    '''

    ntimes = ds.dims['time']

    # Temperature
    if 'tas' in ds:
        trand = ds_rand.t_rand.isel(time=slice(0, ntimes)).values
        ds['tas_with_errors'] = (
            (ds['tas'] ** (1./t_root)) + (ds['tas_error'] * trand)) ** t_root

    # Precipitation
    if 'pr' in ds:
        prand = ds_rand.p_rand.isel(time=slice(0, ntimes)).values
        da_pop = ds['pr_exceedence_probability']
        dap = ((ds['pr'] ** (1./p_root)) + (ds['pr_error'] * prand)) ** p_root

        if pop_thresh is not None:
            # mask where POP is less than the threshold
            dap = ds['pr_with_errors'].where(da_pop > pop_thresh)
        else:
            # mask where POP is less than a uniform transform of prand
            p_rand_uniform = ds_rand.p_rand_uniform.isel(time=slice(0, ntimes)).values
            dap = dap.where(da_pop > p_rand_uniform)
        # mask where precip is negative
        dap = dap.where(dap > 0)
        # Fill in masked cells with zeros
        ds['pr_with_errors'] = dap.fillna(0)

    return ds