def autocorr(ds, lag=1, dim='time', return_p=False):
"""
Calculated lagged correlation of a xr.Dataset.
Parameters
----------
ds : xarray dataset/dataarray
lag : int (default 1)
number of time steps to lag correlate.
dim : str (default 'time')
name of time dimension/dimension to autocorrelate over
return_p : boolean (default False)
if false, return just the correlation coefficient.
if true, return both the correlation coefficient and p-value.
Returns
-------
r : Pearson correlation coefficient
p : (if return_p True) p-value
"""
return st.autocorr(ds, lag=lag, dim=dim, return_p=return_p)
def test_autocorr(control_3d_NA):
"""Check autocorr results with scipy."""
ds = control_3d_NA.isel(x=5, y=5)
actual = autocorr(ds)
expected = correlate(ds, ds)
np.allclose(actual, expected)
def test_corr_autocorr(PM_da_control_3d):
res1 = corr(PM_da_control_3d, PM_da_control_3d, lag=1, return_p=True)
res2 = autocorr(PM_da_control_3d, return_p=True)
for i in [0, 1]:
assert_allclose(res1[i], res2[i])