(use_dask: bool)
| 2011 | |
| 2012 | @pytest.mark.parametrize("use_dask", [True, False]) |
| 2013 | def test_dot(use_dask: bool) -> None: |
| 2014 | if use_dask and not has_dask: |
| 2015 | pytest.skip("test for dask.") |
| 2016 | |
| 2017 | a = np.arange(30 * 4).reshape(30, 4) |
| 2018 | b = np.arange(30 * 4 * 5).reshape(30, 4, 5) |
| 2019 | c = np.arange(5 * 60).reshape(5, 60) |
| 2020 | da_a = xr.DataArray(a, dims=["a", "b"], coords={"a": np.linspace(0, 1, 30)}) |
| 2021 | da_b = xr.DataArray(b, dims=["a", "b", "c"], coords={"a": np.linspace(0, 1, 30)}) |
| 2022 | da_c = xr.DataArray(c, dims=["c", "e"]) |
| 2023 | if use_dask: |
| 2024 | da_a = da_a.chunk({"a": 3}) |
| 2025 | da_b = da_b.chunk({"a": 3}) |
| 2026 | da_c = da_c.chunk({"c": 3}) |
| 2027 | actual = xr.dot(da_a, da_b, dim=["a", "b"]) |
| 2028 | assert actual.dims == ("c",) |
| 2029 | assert (actual.data == np.einsum("ij,ijk->k", a, b)).all() |
| 2030 | assert isinstance(actual.variable.data, type(da_a.variable.data)) |
| 2031 | |
| 2032 | actual = xr.dot(da_a, da_b) |
| 2033 | assert actual.dims == ("c",) |
| 2034 | assert (actual.data == np.einsum("ij,ijk->k", a, b)).all() |
| 2035 | assert isinstance(actual.variable.data, type(da_a.variable.data)) |
| 2036 | |
| 2037 | # for only a single array is passed without dims argument, just return |
| 2038 | # as is |
| 2039 | actual = xr.dot(da_a) |
| 2040 | assert_identical(da_a, actual) |
| 2041 | |
| 2042 | # test for variable |
| 2043 | actual = xr.dot(da_a.variable, da_b.variable) |
| 2044 | assert actual.dims == ("c",) |
| 2045 | assert (actual.data == np.einsum("ij,ijk->k", a, b)).all() |
| 2046 | assert isinstance(actual.data, type(da_a.variable.data)) |
| 2047 | |
| 2048 | if use_dask: |
| 2049 | da_a = da_a.chunk({"a": 3}) |
| 2050 | da_b = da_b.chunk({"a": 3}) |
| 2051 | actual = xr.dot(da_a, da_b, dim=["b"]) |
| 2052 | assert actual.dims == ("a", "c") |
| 2053 | assert (actual.data == np.einsum("ij,ijk->ik", a, b)).all() |
| 2054 | assert isinstance(actual.variable.data, type(da_a.variable.data)) |
| 2055 | |
| 2056 | actual = xr.dot(da_a, da_b, dim=["b"]) |
| 2057 | assert actual.dims == ("a", "c") |
| 2058 | assert (actual.data == np.einsum("ij,ijk->ik", a, b)).all() |
| 2059 | |
| 2060 | actual = xr.dot(da_a, da_b, dim="b") |
| 2061 | assert actual.dims == ("a", "c") |
| 2062 | assert (actual.data == np.einsum("ij,ijk->ik", a, b)).all() |
| 2063 | |
| 2064 | actual = xr.dot(da_a, da_b, dim="a") |
| 2065 | assert actual.dims == ("b", "c") |
| 2066 | assert (actual.data == np.einsum("ij,ijk->jk", a, b)).all() |
| 2067 | |
| 2068 | actual = xr.dot(da_a, da_b, dim="c") |
| 2069 | assert actual.dims == ("a", "b") |
| 2070 | assert (actual.data == np.einsum("ij,ijk->ij", a, b)).all() |
nothing calls this directly
no test coverage detected
searching dependent graphs…