Upgrade luxury sparse (#555)

* Fix tests

* update Project.toml

Author: GiggleLiu

Project.toml

@@ -24,7 +24,7 @@ CuYao = "CUDA"
 BitBasis = "0.8, 0.9"
 CUDA = "4, 5"
 LinearAlgebra = "1"
-LuxurySparse = "0.7"
+LuxurySparse = "0.8"
 Reexport = "1"
 YaoAPI = "0.4"
 YaoArrayRegister = "0.9"

lib/YaoArrayRegister/Project.toml

@@ -22,7 +22,7 @@ Adapt = "3, 4"
 BitBasis = "0.8, 0.9"
 DocStringExtensions = "0.8, 0.9"
 LegibleLambdas = "0.3"
-LuxurySparse = "0.7"
+LuxurySparse = "0.8"
 MLStyle = "0.4"
 StaticArrays = "1"
 StatsBase = "0.33 - 0.34"

lib/YaoArrayRegister/src/YaoArrayRegister.jl

@@ -12,7 +12,8 @@ using BitBasis
 using LinearAlgebra
 using LegibleLambdas
 using StatsBase, Random
-using LuxurySparse, StaticArrays
+using StaticArrays
+import LuxurySparse: fastkron, IMatrix, PermMatrix, Diagonal, SparseMatrixCSC
 using DocStringExtensions
 
 export AbstractArrayReg,

lib/YaoArrayRegister/src/utils.jl

@@ -441,12 +441,11 @@ function _wrap_identity(
     nlevel
 ) where {T<:AbstractMatrix}
     length(num_bit_list) == length(data_list) + 1 || throw(ArgumentError())
-    ⊗ = kron
     reduce(
         zip(data_list, num_bit_list[2:end]);
         init = IMatrix(nlevel ^ num_bit_list[1]),
     ) do x, y
-        x ⊗ y[1] ⊗ IMatrix(nlevel ^ y[2])
+        fastkron(fastkron(x, y[1]), IMatrix(nlevel ^ y[2]))
     end
 end
 

lib/YaoBlocks/Project.toml

@@ -30,7 +30,7 @@ InteractiveUtils = "1"
 KrylovKit = "0.5, 0.6, 0.7, 0.8, 0.9"
 LegibleLambdas = "0.2, 0.3"
 LinearAlgebra = "1"
-LuxurySparse = "0.7"
+LuxurySparse = "0.8"
 MLStyle = "0.3, 0.4"
 Random = "1"
 SparseArrays = "1"

lib/YaoBlocks/src/YaoBlocks.jl

@@ -9,11 +9,11 @@ using YaoAPI
 using LinearAlgebra
 using YaoArrayRegister
 using YaoArrayRegister: @λ, matvec, diff, autostatic, rot_mat
-using BitBasis, LuxurySparse
+using BitBasis, LuxurySparse, SparseArrays
+using LuxurySparse: fastkron
 using StatsBase, TupleTools, InteractiveUtils
 using MLStyle: @match
 using LinearAlgebra: eigen!
-using SparseArrays, LuxurySparse
 using Random, CacheServers
 import KrylovKit: exponentiate
 using DocStringExtensions

lib/YaoBlocks/src/autodiff/outerproduct_and_projection.jl

@@ -80,8 +80,7 @@ unsafe_projection!(y::Diagonal, op::OuterProduct) = (y.diag .+= op.left .* op.ri
 unsafe_projection!(y::Matrix, adjy, v) = y .+= adjy .* v
 
 @inline function unsafe_projection!(y::AbstractSparseMatrix, m::AbstractMatrix)
-    is, js, vs = findnz(y)
-    for (k, (i, j)) in enumerate(zip(is, js))
+    for (k, (i, j, _)) in enumerate(IterNz(y))
         @inbounds y.nzval[k] += m[i, j]
     end
     y

lib/YaoBlocks/src/composite/kron.jl

@@ -172,7 +172,7 @@ function mat(::Type{T}, k::KronBlock{D,M}) where {T,D,M}
         Iterators.reverse(zip(subblocks(k), num_bit_list)),
         init = IMatrix{T}(D^ntrail),
     ) do x, y
-        kron(x, mat(T, y[1]), IMatrix(D^y[2]))
+        fastkron(fastkron(x, mat(T, y[1])), IMatrix(D^y[2]))
     end
 end
 

lib/YaoBlocks/src/composite/tag/onlevels.jl

@@ -20,8 +20,15 @@ end
 content(g::OnLevels) = g.gate
 function mat(::Type{T}, g::OnLevels{D, Ds}) where {T, D, Ds}
     m = mat(T, g.gate)
-    I, J, V = LuxurySparse.findnz(m)
-    return sparse(collect(g.levels[I]), collect(g.levels[J]), V, D, D)
+    iter = IterNz(m)
+    nnz = length(iter)
+    is, js, vs = Vector{Int}(undef, nnz), Vector{Int}(undef, nnz), Vector{T}(undef, nnz)
+    for (k, (i, j, v)) in enumerate(iter)
+        is[k] = g.levels[i]
+        js[k] = g.levels[j]
+        vs[k] = v
+    end
+    return sparse(is, js, vs, D, D)
 end
 PropertyTrait(::OnLevels) = PreserveAll()
 Base.adjoint(x::OnLevels{D}) where D = OnLevels{D}(adjoint(x.gate), x.levels)

lib/YaoBlocks/src/outerproduct_and_projection.jl

@@ -80,8 +80,7 @@ unsafe_projection!(y::Diagonal, op::OuterProduct) = (y.diag .+= op.left .* op.ri
 unsafe_projection!(y::Matrix, adjy, v) = y .+= adjy .* v
 
 @inline function unsafe_projection!(y::AbstractSparseMatrix, m::AbstractMatrix)
-    is, js, vs = findnz(y)
-    for (k, (i, j)) in enumerate(zip(is, js))
+    for (k, (i, j, _)) in enumerate(IterNz(y))
         @inbounds y.nzval[k] += m[i, j]
     end
     y