SparseArrayKit.jl/test/basic.jl at 0435d908f134c77ae9eac28c19dc1e36e6067dc9 · QuantumKitHub/SparseArrayKit.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
module BasicTests

using SparseArrayKit
using Test, TestExtras, LinearAlgebra, Random
using TupleTools

#=
generate a whole bunch of random contractions, compare with the dense result
=#
function randn_sparse(T::Type{<:Number}, sz::Dims, p = 0.5)
    a = SparseArray{T}(undef, sz)
    for I in keys(a)
        if rand() < p
            a[I] = randn(T)
        end
    end
    return a
end

@timedtestset "Basic linear algebra" begin
    MAX_DIM = 20
    MAX_LEGS = 4

    dims = ntuple(l -> rand(1:MAX_DIM), MAX_LEGS)
    ar = randn_sparse(Float64, dims)
    ac = randn_sparse(ComplexF64, dims)

    slice = (1:rand(1:dims[1]), rand(1:dims[2]), rand(1:dims[3]):dims[3], rand(1:dims[4]))
    @test ar[slice...] == Array(ar)[slice...]

    α = randn(ComplexF64)
    β = randn(Float64)
    γ = 2
    @test @constinferred(α * ar) == α * Array(ar)
    @test @constinferred(β * ar) == β * Array(ar)
    @test @constinferred(γ * ar) == γ * Array(ar)
    @test @constinferred(ac * β) == Array(ac) * β
    @test @constinferred(ar + ac) == Array(ar) + Array(ac)
    @test @constinferred(ac - ar) == Array(ac) - Array(ar)
    mar = @constinferred(-ar)
    @test iszero(ar + mar)
    ar2 = copy(ar)
    ar2[first(nonzero_keys(ar2))] = 0
    @test !(ar2 == ar)
    @test @constinferred(zero(ar)) + ac == ac
    @test ac == convert(SparseArray, Array(ac))

    @test norm(ar + @constinferred(α * ar)) ≈ norm(ar) * abs(1 + α)
    @test @constinferred(norm(ac + ac * α)) ≈ norm(ac) * abs(1 + α)
    @test ar * α ≈ α * ar
    @test norm(ar * α) ≈ norm(ar) * abs(α)
    @test ac / β ≈ β \ ac
    @test norm(ac / γ) ≈ norm(ac) / abs(γ)
    @test norm(ac) ≈ sqrt(real(@constinferred(dot(ac, ac))))
    @test @constinferred(dot(ac, ar)) ≈ dot(Array(ac), Array(ar))
    @test dot(ar, ac) ≈ conj(dot(ac, ar))

    @test @constinferred(lmul!(β, copy(ar))) == β * ar == β * Array(ar)
    @test @constinferred(rmul!(copy(ac), α)) == ac * α == Array(ac) * α
    @test @constinferred(rmul!(copy(ac), β)) == ac * β == Array(ac) * β
    @test @constinferred(ldiv!(β, copy(ar))) == ldiv!(β, Array(ar)) ≈ β \ ar
    @test @constinferred(rdiv!(copy(ac), α)) == rdiv!(Array(ac), α) ≈ ac / α
    @test @constinferred(rdiv!(copy(ac), β)) == rdiv!(Array(ac), β) ≈ ac / β
    @test @constinferred(conj!(copy(ac))) == conj!(Array(ac))
    @test_throws InexactError rmul!(copy(ar), α)
    @test_throws InexactError ldiv!(α, copy(ar))

    p = randperm(MAX_LEGS)
    @test permutedims(ar, p) == permutedims(Array(ar), p)

    @test @constinferred(axpy!(α, copy(ar), copy(ac))) == α * ar + ac
    @test @constinferred(axpby!(α, copy(ar), β, copy(ac))) == α * ar + β * ac
end

@timedtestset "Basic matrix algebra" begin
    using SparseArrays
    N = 100
    a = sprandn(ComplexF64, N, N, 0.1)
    b = sprandn(N, N, 0.1)
    aa = @constinferred(SparseArray(a))
    bb = SparseArray(b)
    @test aa == a
    @test @constinferred(aa * bb) ≈ SparseArray(a * b)
    @test aa * bb ≈ Array(aa) * Array(bb)
    @test adjoint(aa) == adjoint(Array(aa))
    @test transpose(aa) == transpose(Array(aa))
    @test aa' * bb ≈ Array(aa)' * Array(bb)
    @test aa' * bb' ≈ Array(aa)' * Array(bb)'
    @test aa * bb' ≈ Array(aa) * Array(bb)'
    @test @constinferred(one(aa)) == one(Array(aa))
    @test norm(one(aa)) ≈ sqrt(N)
    @test one(aa) + zero(aa) == one(bb)
    @test adjoint!(copy(aa), bb) == SparseArray(adjoint(bb)) == adjoint(Array(bb))
    @test transpose!(copy(aa), bb) == SparseArray(transpose(bb)) == transpose(Array(bb))

    @test SparseArray(I, (5, 5)) == I
    @test SparseArray{Float64}(I, (5, 5)) == I
    @test size(SparseArray(I, 3, 8)) == (3, 8)
end

@timedtestset "Index manipulations" begin
    dims = (2, 3, 4)
    A = randn_sparse(Float64, dims)
    @test @constinferred(reindexdims(A, (2, 1, 3))) == permutedims(A, (2, 1, 3))

    A_expanded = @constinferred reindexdims(A, (1, 1, 2, 3))
    @test size(A_expanded) == TupleTools.getindices(size(A), (1, 1, 2, 3))
    @test norm(A_expanded) ≈ norm(A)
    @test reindexdims(A_expanded, (1, 3, 4)) == A

    A_reduced = @constinferred reindexdims(A, (1, 2))
    @test size(A_reduced) == TupleTools.getindices(size(A), (1, 2))
    @test Array(A_reduced) ≈ sum(Array(A); dims = 3)
end

end