def parcorr_z(field: xr.DataArray, ts: np.ndarray, z: pd.DataFrame, lag_z: int = 0): ''' Regress out influence of 1-d timeseries z. if lag_z==0, dates of z will match dates of field. Note, lag_z >= 1 probably only makes sense when using subseasonal data (more then 1 value per year). Parameters ---------- field : xr.DataArray (time, lat, lon) field. ts : np.ndarray Target timeseries. z : pd.DataFrame 1-d timeseries. Returns ------- corr_vals : np.ndarray pvals : np.ndarray ''' # if more then one year is filled with NaNs -> no corr value calculated. dates = pd.to_datetime(field.time.values) field, ts = check_NaNs(field, ts) x = np.ma.zeros(field.shape[1]) corr_vals = np.array(x) pvals = np.array(x) fieldnans = np.array([np.isnan(field[:, i]).any() for i in range(x.size)]) nonans_gc = np.arange(0, fieldnans.size)[fieldnans == False] # ts = np.expand_dims(ts[:], axis=1) # adjust to shape (samples, dimension) and remove first datapoints if # lag_z != 0. y = np.expand_dims(ts[lag_z:], axis=1) if len(z.values.squeeze().shape) == 1: z = np.expand_dims(z.loc[dates].values.squeeze(), axis=1) else: z = z.loc[dates].values.squeeze() if lag_z >= 1: z = z[:-lag_z] # last values are 'removed' for i in nonans_gc: cond_ind_test = ParCorr() field_i = np.expand_dims(field[lag_z:, i], axis=1) a, b = cond_ind_test.run_test_raw(field_i, y, z) corr_vals[i] = a pvals[i] = b # restore original nans corr_vals[fieldnans] = np.nan return corr_vals, pvals
def parcorr_map_time(field: xr.DataArray, ts: np.ndarray, lag_y=0, lag_x=0): ''' Only works for subseasonal data (more then 1 datapoint per year). Lag must be >= 1 Parameters ---------- field : xr.DataArray (time, lat, lon) field. ts : np.ndarray Target timeseries. lag : int, optional DESCRIPTION. The default is 1. target : TYPE, optional DESCRIPTION. The default is True. precursor : TYPE, optional DESCRIPTION. The default is True. Returns ------- corr_vals : np.ndarray pvals : np.ndarray ''' # field = precur_train.sel(time=dates_lag) ; ts = RV_ts.values.squeeze() if type(lag_y) is int: lag_y = [lag_y] if type(lag_x) is int: lag_x = [lag_x] max_lag = max(max(lag_y), max(lag_x)) assert max_lag > 0, 'lag_x or lag_y must be >= 1' # if more then one year is filled with NaNs -> no corr value calculated. field, ts = check_NaNs(field, ts) x = np.ma.zeros(field.shape[1]) corr_vals = np.array(x) pvals = np.array(x) fieldnans = np.array([np.isnan(field[:, i]).any() for i in range(x.size)]) nonans_gc = np.arange(0, fieldnans.size)[fieldnans == False] if max(lag_y) > 0: zy = [ np.expand_dims(ts[max_lag - l:-l], axis=1) for l in lag_y if l != 0 ] zy = np.concatenate(zy, axis=1) y = np.expand_dims(ts[max_lag:], axis=1) for i in nonans_gc: cond_ind_test = ParCorr() if max(lag_x) > 0: zx = [ np.expand_dims(field[max_lag - l:-l, i], axis=1) for l in lag_x if l != 0 ] zx = np.concatenate(zx, axis=1) if max(lag_x) > 0 and max(lag_y) > 0: # both zy and zx defined z = np.concatenate((zy, zx), axis=1) elif max(lag_x) > 0 and max(lag_y) == 0: # only zx defined z = zx elif max(lag_x) == 0 and max(lag_y) > 0: z = zy field_i = np.expand_dims(field[max_lag:, i], axis=1) a, b = cond_ind_test.run_test_raw(field_i, y, z) corr_vals[i] = a pvals[i] = b # restore original nans corr_vals[fieldnans] = np.nan return corr_vals, pvals