Method test_dot

xarray/tests/test_dataarray.py:4347–4391

(self)

            full_like(da, fill_value=True, dtype={"x": bool})

    def test_dot(self) -> None:
        x = np.linspace(-3, 3, 6)
        y = np.linspace(-3, 3, 5)
        z = range(4)
        da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4))
        da = DataArray(da_vals, coords=[x, y, z], dims=["x", "y", "z"])

        dm_vals1 = range(4)
        dm1 = DataArray(dm_vals1, coords=[z], dims=["z"])

        # nd dot 1d
        actual1 = da.dot(dm1)
        expected_vals1 = np.tensordot(da_vals, dm_vals1, (2, 0))
        expected1 = DataArray(expected_vals1, coords=[x, y], dims=["x", "y"])
        assert_equal(expected1, actual1)

        # all shared dims
        actual2 = da.dot(da)
        expected_vals2 = np.tensordot(da_vals, da_vals, axes=([0, 1, 2], [0, 1, 2]))
        expected2 = DataArray(expected_vals2)
        assert_equal(expected2, actual2)

        # multiple shared dims
        dm_vals3 = np.arange(20 * 5 * 4).reshape((20, 5, 4))
        j = np.linspace(-3, 3, 20)
        dm3 = DataArray(dm_vals3, coords=[j, y, z], dims=["j", "y", "z"])
        actual3 = da.dot(dm3)
        expected_vals3 = np.tensordot(da_vals, dm_vals3, axes=([1, 2], [1, 2]))
        expected3 = DataArray(expected_vals3, coords=[x, j], dims=["x", "j"])
        assert_equal(expected3, actual3)

        # Ellipsis: all dims are shared
        actual4 = da.dot(da, dim=...)
        expected4 = da.dot(da)
        assert_equal(expected4, actual4)

        # Ellipsis: not all dims are shared
        actual5 = da.dot(dm3, dim=...)
        expected5 = da.dot(dm3, dim=("j", "x", "y", "z"))
        assert_equal(expected5, actual5)

        with pytest.raises(NotImplementedError):
            da.dot(dm3.to_dataset(name="dm"))
        with pytest.raises(TypeError):
            da.dot(dm3.values)  # type: ignore[type-var]

    def test_dot_align_coords(self) -> None:
        # GH 3694