(
n: int, dim: str | None, array_tuples: tuple[xr.DataArray, xr.DataArray]
)
| 1660 | @pytest.mark.parametrize("n", [0, 1, 2]) |
| 1661 | @pytest.mark.parametrize("dim", [None, "time"]) |
| 1662 | def test_corr( |
| 1663 | n: int, dim: str | None, array_tuples: tuple[xr.DataArray, xr.DataArray] |
| 1664 | ) -> None: |
| 1665 | da_a, da_b = array_tuples[n] |
| 1666 | |
| 1667 | if dim is not None: |
| 1668 | |
| 1669 | def np_corr_ind(ts1, ts2, a, x): |
| 1670 | # Ensure the ts are aligned and missing values ignored |
| 1671 | ts1, ts2 = broadcast(ts1, ts2) |
| 1672 | valid_values = ts1.notnull() & ts2.notnull() |
| 1673 | |
| 1674 | ts1 = ts1.where(valid_values) |
| 1675 | ts2 = ts2.where(valid_values) |
| 1676 | |
| 1677 | return np.ma.corrcoef( |
| 1678 | np.ma.masked_invalid(ts1.sel(a=a, x=x).data.flatten()), |
| 1679 | np.ma.masked_invalid(ts2.sel(a=a, x=x).data.flatten()), |
| 1680 | )[0, 1] |
| 1681 | |
| 1682 | expected = np.zeros((3, 4)) |
| 1683 | for a in [0, 1, 2]: |
| 1684 | for x in [0, 1, 2, 3]: |
| 1685 | expected[a, x] = np_corr_ind(da_a, da_b, a=a, x=x) |
| 1686 | actual = xr.corr(da_a, da_b, dim) |
| 1687 | assert_allclose(actual, expected) |
| 1688 | |
| 1689 | else: |
| 1690 | |
| 1691 | def np_corr(ts1, ts2): |
| 1692 | # Ensure the ts are aligned and missing values ignored |
| 1693 | ts1, ts2 = broadcast(ts1, ts2) |
| 1694 | valid_values = ts1.notnull() & ts2.notnull() |
| 1695 | |
| 1696 | ts1 = ts1.where(valid_values) |
| 1697 | ts2 = ts2.where(valid_values) |
| 1698 | |
| 1699 | return np.ma.corrcoef( |
| 1700 | np.ma.masked_invalid(ts1.data.flatten()), |
| 1701 | np.ma.masked_invalid(ts2.data.flatten()), |
| 1702 | )[0, 1] |
| 1703 | |
| 1704 | expected = np_corr(da_a, da_b) |
| 1705 | actual = xr.corr(da_a, da_b, dim) |
| 1706 | assert_allclose(actual, expected) |
| 1707 | |
| 1708 | |
| 1709 | @pytest.mark.parametrize("n", range(9)) |
nothing calls this directly
no test coverage detected
searching dependent graphs…