Repository: dask/dask
File: dask/array/tests/test_routines.py (Python)

Code before:
    e = da.from_array(y, chunks=(4,))

    assert_eq(da.cov(d, e), np.cov(x, y))
    assert_eq(da.cov(e, d), np.cov(y, x))

    with pytest.raises(ValueError):
        da.cov(d, ddof=1.5)


@pytest.mark.skipif(
    not NUMPY_GE_220, reason="fweights is not an kwarg prior to numpy 2.2"
)
def test_cov_fweights():
    x1 = da.array([[0, 2], [1, 1], [2, 0]]).T
    res1 = da.array([[1.0, -1.0], [-1.0, 1.0]])
    x2 = da.array([0.0, 1.0, 2.0], ndmin=2)
    frequencies = da.array([1, 4, 1])
    x2_repeats = da.array([[0.0], [1.0], [1.0], [1.0], [1.0], [2.0]]).T
    res2 = da.array([[0.4, -0.4], [-0.4, 0.4]])
    unit_frequencies = np.ones(3, dtype=np.int_)

    assert allclose(da.cov(x2, fweights=frequencies), da.cov(x2_repeats))
    assert allclose(da.cov(x1, fweights=frequencies), res2)
    assert allclose(da.cov(x1, fweights=unit_frequencies), res1)

    nonint = frequencies + 0.5
    with pytest.raises(TypeError):
        da.cov(x1, fweights=nonint)

    f = da.ones((2, 3), dtype=np.int_)
    with pytest.raises(RuntimeError):
        da.cov(x1, fweights=f)

    f = da.ones(2, dtype=np.int_)
    with pytest.raises(RuntimeError):
        da.cov(x1, fweights=f)

    f = -1 * np.ones(3, dtype=np.int_)
    with pytest.raises(ValueError):
        da.cov(x1, fweights=f)


@pytest.mark.skipif(
    not NUMPY_GE_220, reason="aweights is not an kwarg prior to numpy 2.2"
)
def test_cov_aweights():
    x1 = da.array([[0, 2], [1, 1], [2, 0]]).T
    res1 = da.array([[1.0, -1.0], [-1.0, 1.0]])

    # Test basic functionality with aweights
    aweights = da.array([0.5, 2.0, 0.5])  # Analytical weights

Review comment (suggestion):
Would something like this be possible, to ensure we match NumPy?

```
result = da.cov(x2, fweights=frequencies)
expected = np.cov(x2.compute(), fweights=frequencies.compute())
assert_eq(result, expected)
```

Code after:
    e = da.from_array(y, chunks=(4,))

    assert_eq(da.cov(d, e), np.cov(x, y))
    assert_eq(da.cov(e, d), np.cov(y, x))

    with pytest.raises(ValueError):
        da.cov(d, ddof=1.5)


@pytest.mark.skipif(
    not NUMPY_GE_220, reason="fweights is not an kwarg prior to numpy 2.2"
)
def test_cov_fweights():
    x1 = da.array([[0, 2], [1, 1], [2, 0]]).T
    res1 = da.array([[1.0, -1.0], [-1.0, 1.0]])
    x2 = da.array([0.0, 1.0, 2.0], ndmin=2)
    frequencies = da.array([1, 4, 1])
    x2_repeats = da.array([[0.0], [1.0], [1.0], [1.0], [1.0], [2.0]]).T
    res2 = da.array([[0.4, -0.4], [-0.4, 0.4]])
    unit_frequencies = np.ones(3, dtype=np.int_)

    result = da.cov(x2, fweights=frequencies)
    expected = np.cov(x2.compute(), fweights=frequencies.compute())
    assert_eq(result, expected)

    assert allclose(da.cov(x2, fweights=frequencies), da.cov(x2_repeats))
    assert allclose(da.cov(x1, fweights=frequencies), res2)
    assert allclose(da.cov(x1, fweights=unit_frequencies), res1)

    f = da.ones((2, 3), dtype=np.int_)
    with pytest.raises(RuntimeError):
        da.cov(x1, fweights=f)

    f = da.ones(2, dtype=np.int_)
    with pytest.raises(RuntimeError):
        da.cov(x1, fweights=f)


@pytest.mark.skipif(
    not NUMPY_GE_220, reason="aweights is not an kwarg prior to numpy 2.2"
)
def test_cov_aweights():
    x1 = da.array([[0, 2], [1, 1], [2, 0]]).T
    res1 = da.array([[1.0, -1.0], [-1.0, 1.0]])

    # Test basic functionality with aweights
    aweights = da.array([0.5, 2.0, 0.5])  # Analytical weights
    unit_aweights = np.ones(3, dtype=np.float64)

    # With unit weights, should match unweighted result
    assert allclose(da.cov(x1, aweights=unit_aweights), res1)