def test_mse_nd(a, b):
axis = 0
expected = np.squeeze(a[0, :, :]).copy()
for i in range(np.shape(a)[1]):
for j in range(np.shape(a)[2]):
_a = a[:, i, j]
_b = b[:, i, j]
expected[i, j] = mean_squared_error(_a, _b)
actual = _mse(a, b, weights, axis, skipna)
assert_allclose(actual, expected)
axis = 1
expected = np.squeeze(a[:, 0, :]).copy()
for i in range(np.shape(a)[0]):
for j in range(np.shape(a)[2]):
_a = a[i, :, j]
_b = b[i, :, j]
expected[i, j] = mean_squared_error(_a, _b)
actual = _mse(a, b, weights, axis, skipna)
assert_allclose(actual, expected)
axis = 2
expected = np.squeeze(a[:, :, 0]).copy()
for i in range(np.shape(a)[0]):
for j in range(np.shape(a)[1]):
_a = a[i, j, :]
_b = b[i, j, :]
expected[i, j] = mean_squared_error(_a, _b)
actual = _mse(a, b, weights, axis, skipna)
assert_allclose(actual, expected)
def test_mse_r_xr(a, b, dim):
actual = mse(a, b, dim)
dim, axis = _preprocess(dim)
_a = a.values
_b = b.values
axis = tuple(a.dims.index(d) for d in dim)
res = _mse(_a, _b, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_mse_r_xr_dask(a_dask, b_dask, dim):
actual = mse(a_dask, b_dask, dim)
dim, axis = _preprocess(dim)
_a_dask = a_dask.values
_b_dask = b_dask.values
axis = tuple(a_dask.dims.index(d) for d in dim)
res = _mse(_a_dask, _b_dask, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_mse_r_xr_dask(a_dask, b_dask, dim):
actual = xr.apply_ufunc(_mse,
a_dask,
b_dask,
input_core_dims=[[dim], [dim]],
kwargs={'axis': -1},
dask='parallelized',
output_dtypes=[float])
_a_dask = a_dask.values
_b_dask = b_dask.values
axis = a_dask.dims.index(dim)
res = _mse(_a_dask, _b_dask, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)
def test_mse_r_xr(a, b, dim, weight, weights_ones, weights_latitude):
# Generates subsetted weights to pass in as arg to main function and for the numpy testing.
weights_arg, weights_np = adjust_weights(weight, dim, weights_ones,
weights_latitude)
actual = mse(a, b, dim, weights=weights_arg)
assert actual.chunks is None
dim, axis = _preprocess_dims(dim)
_a = a
_b = b
_weights_np = _preprocess_weights(_a, dim, dim, weights_np)
axis = tuple(a.dims.index(d) for d in dim)
res = _mse(_a.values, _b.values, _weights_np.values, axis)
expected = actual.copy()
expected.values = res
assert_allclose(actual, expected)