1import pytest
2
3pytestmark = pytest.mark.gpu
4
5import dask.array as da
6from dask.array.utils import assert_eq
7
8cupy = pytest.importorskip("cupy")
9cupyx = pytest.importorskip("cupyx")
10
11
12def test_sparse_hstack_vstack_csr():
13    pytest.importorskip("cupyx")
14    x = cupy.arange(24, dtype=cupy.float32).reshape(4, 6)
15
16    sp = da.from_array(x, chunks=(2, 3), asarray=False, fancy=False)
17    sp = sp.map_blocks(cupyx.scipy.sparse.csr_matrix, dtype=cupy.float32)
18
19    y = sp.compute()
20
21    assert cupyx.scipy.sparse.isspmatrix(y)
22    assert_eq(x, y.todense())
23
24
25@pytest.mark.parametrize("axis", [0, 1])
26def test_sparse_concatenate(axis):
27    pytest.importorskip("cupyx")
28
29    rs = da.random.RandomState(RandomState=cupy.random.RandomState)
30    meta = cupyx.scipy.sparse.csr_matrix((0, 0))
31
32    xs = []
33    ys = []
34    for i in range(2):
35        x = rs.random((1000, 10), chunks=(100, 10))
36        x[x < 0.9] = 0
37        xs.append(x)
38        ys.append(x.map_blocks(cupyx.scipy.sparse.csr_matrix, meta=meta))
39
40    z = da.concatenate(ys, axis=axis)
41    z = z.compute()
42
43    if axis == 0:
44        sp_concatenate = cupyx.scipy.sparse.vstack
45    elif axis == 1:
46        sp_concatenate = cupyx.scipy.sparse.hstack
47    z_expected = sp_concatenate(
48        [cupyx.scipy.sparse.csr_matrix(e.compute()) for e in xs]
49    )
50
51    assert (z.toarray() == z_expected.toarray()).all()
52
53
54@pytest.mark.parametrize("sp_format", ["csr", "csc"])
55def test_sparse_dot(sp_format):
56    pytest.importorskip("cupyx")
57
58    if sp_format == "csr":
59        sp_matrix = cupyx.scipy.sparse.csr_matrix
60    elif sp_format == "csc":
61        sp_matrix = cupyx.scipy.sparse.csc_matrix
62    dtype = "f"
63    density = 0.3
64    x_shape, x_chunks = (4, 8), (2, 4)
65    y_shape, y_chunks = (8, 6), (4, 3)
66    x = cupy.random.random(x_shape, dtype=dtype)
67    y = cupy.random.random(y_shape, dtype=dtype)
68    x[x < 1 - density] = 0
69    y[y < 1 - density] = 0
70    z = x.dot(y)
71
72    da_x = da.from_array(x, chunks=x_chunks, asarray=False, fancy=False)
73    da_y = da.from_array(y, chunks=y_chunks, asarray=False, fancy=False)
74    da_x = da_x.map_blocks(sp_matrix, dtype=dtype)
75    da_y = da_y.map_blocks(sp_matrix, dtype=dtype)
76    da_z = da.dot(da_x, da_y).compute()
77
78    assert cupyx.scipy.sparse.isspmatrix(da_z)
79    assert_eq(z, da_z.todense())
80