Merge pull request #16 from QuantumBFS/jg/speedup

Improve the performance

Project.toml

@@ -1,12 +1,11 @@
 name = "FLOYao"
 uuid = "6d9310a3-f1d0-41b7-8edb-11c1cf57cd2d"
 authors = ["janlukas.bosse <[email protected]> and contributors"]
-version = "1.4.3"
+version = "1.4.4"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Yao = "5872b779-8223-5990-8dd0-5abbb0748c8c"
 
 [compat]

src/FLOYao.jl

@@ -39,8 +39,8 @@ module FLOYao
 
 using LinearAlgebra
 using Random
-using SparseArrays
 using Yao
+using Yao.YaoBlocks.LuxurySparse: SparseMatrixCOO
 
 export MajoranaReg
 

src/apply_composite.jl

@@ -96,23 +96,27 @@ end
 
 # Defined in FLOYao.jl
 # const RGate = RotationGate{2,<:Real,<:PauliKronBlock}
-function YaoBlocks.unsafe_apply!(reg::MajoranaReg, b::RGate)
-    i1, i2 = kron2majoranaindices(b.block)
-    s, c = sincos(b.theta)
-    ψ1, ψ2 = reg.state[i1,:], reg.state[i2,:]
-    reg.state[i1,:] .= c .* ψ1 .+ s .* ψ2
-    reg.state[i2,:] .= c .* ψ2 .- s .* ψ1
+function YaoBlocks.unsafe_apply!(reg::MajoranaReg, ψ2::RGate)
+    i1, i2 = kron2majoranaindices(ψ2.block)
+    s, c = sincos(ψ2.theta)
+    for k in 1:size(reg.state, 2)
+        ψ1, ψ2 = reg.state[i1,k], reg.state[i2,k]
+        reg.state[i1,k] = c * ψ1 + s * ψ2
+        reg.state[i2,k] = c * ψ2 - s * ψ1
+    end
     return reg
 end
 
-function YaoBlocks.unsafe_apply!(reg::MajoranaReg, b::PutBlock{2,N,<:RGate}) where {N}
-    areconsecutive(b.locs) || throw(NonFLOException("$(b.blocks) on $(b.locs) is not a FLO gate"))
+function YaoBlocks.unsafe_apply!(reg::MajoranaReg, ψ2::PutBlock{2,N,<:RGate}) where {N}
+    areconsecutive(ψ2.locs) || throw(NonFLOException("$(ψ2.blocks) on $(ψ2.locs) is not a FLO gate"))
     # goddamnit, 1-based indexing
-    i1, i2 = 2 * (minimum(b.locs) - 1) .+ kron2majoranaindices(b.content.block)
-    s, c = sincos(b.content.theta)
-    ψ1, ψ2 = reg.state[i1,:], reg.state[i2,:]
-    reg.state[i1,:] .= c .* ψ1 .+ s .* ψ2
-    reg.state[i2,:] .= c .* ψ2 .- s .* ψ1
+    i1, i2 = 2 * (minimum(ψ2.locs) - 1) .+ kron2majoranaindices(ψ2.content.block)
+    s, c = sincos(ψ2.content.theta)
+    for k in 1:size(reg.state, 2)
+        ψ1, ψ2 = reg.state[i1, k], reg.state[i2, k]
+        reg.state[i1, k] = c * ψ1 + s * ψ2
+        reg.state[i2, k] = c * ψ2 - s * ψ1
+    end
     return reg
 end
 

src/auto_diff.jl

@@ -24,10 +24,10 @@ function Yao.AD.expect_g(op::AbstractAdd, in::MajoranaReg)
     ham = yaoham2majoranasquares(op)
     inδ = copy(in)
     inδ.state .= ham * inδ.state
-    for i in 1:nqudits(in) 
-        ψ1, ψ2 = inδ.state[:,2i-1], inδ.state[:,2i]
-        inδ.state[:,2i-1] .= -ψ2
-        inδ.state[:,2i] .= ψ1
+    @inbounds for i in 1:nqudits(in)
+        for k in 1:size(inδ.state, 1)
+            inδ.state[k,2i-1], inδ.state[k,2i] = -inδ.state[k,2i], inδ.state[k,2i-1]
+        end
     end
     inδ.state[:,1:end] .*= -1
     return inδ

src/instruct.jl

@@ -42,9 +42,11 @@ for (G, SC) in zip([:T, :Tdag], [(1/√2, 1/√2), (-1/√2, 1/√2)])
     @eval function Yao.instruct!(reg::MajoranaReg, ::Val{$(QuoteNode(G))}, locs::Tuple)
         loc = locs[1]
         s, c = $SC
-        ψ1, ψ2 = reg.state[2loc-1,:], reg.state[2loc,:]
-        reg.state[2loc-1,:] .= c .* ψ1 .+ s .* ψ2
-        reg.state[2loc,:] .= c .* ψ2 .- s .* ψ1
+        @inbounds for k = 1:size(reg.state, 2)
+            ψ1, ψ2 = reg.state[2loc-1,k], reg.state[2loc,k]
+            reg.state[2loc-1,k] = c * ψ1 + s * ψ2
+            reg.state[2loc,k] = -s * ψ1 + c * ψ2
+        end
         return reg
     end
 end
@@ -53,9 +55,11 @@ for G in [:Rz, :Shift]
     @eval function Yao.instruct!(reg::MajoranaReg, ::Val{$(QuoteNode(G))}, locs::Tuple, theta)
         loc = locs[1]
         s, c = sincos(theta)
-        ψ1, ψ2 = reg.state[2loc-1,:], reg.state[2loc,:]
-        reg.state[2loc-1,:] .= c .* ψ1 .+ s .* ψ2
-        reg.state[2loc,:] .= c .* ψ2 .- s .* ψ1
+        for k = 1:size(reg.state, 2)
+            ψ1, ψ2 = reg.state[2loc-1,k], reg.state[2loc,k]
+            reg.state[2loc-1,k] = c * ψ1 + s * ψ2
+            reg.state[2loc,k] = -s * ψ1 + c * ψ2
+        end
         return reg
     end
 end

src/utils.jl

@@ -195,10 +195,10 @@ end
 
 Get the indices of majorana operators from a kronecker product of pauli operators
 """
-function kron2majoranaindices(k::KronBlock)
-    locs = collect(Iterators.flatten(k.locs))
-    perm = sortperm(locs)
-    areconsecutive(locs[perm]) || throw(NonFLOException("$k is not acting on consecutive qubits"))
+function kron2majoranaindices(k::KronBlock{2, M, <:NTuple{M, PauliGate}}) where {M}
+    locs = first.(k.locs)
+    perm = YaoBlocks.TupleTools.sortperm(locs)
+    areconsecutive(ntuple(i->locs[perm[i]], M)) || throw(NonFLOException("$k is not acting on consecutive qubits"))
 
     # these are the indices in the majorana hamiltonian corresponding to this
     # kronecker product. swap accumulates the overall sign
@@ -258,7 +258,7 @@ function yaoham2majoranasquares(::Type{T}, yaoham::Add{2}) where {T<:Real}
     ham = zeros(T, 2nqubits(yaoham), 2nqubits(yaoham))
     @inbounds for k in yaoham
         if k isa Scale
-            fast_add!(ham, rmul!(yaoham2majoranasquares(T, k.content), k.alpha))
+            fast_add!(ham, _rmul!(yaoham2majoranasquares(T, k.content), k.alpha))
             #ham += k.alpha * yaoham2majoranasquares(T, k.content)
         elseif k isa KronBlock
             i1, i2 = kron2majoranaindices(k)
@@ -283,28 +283,19 @@ function yaoham2majoranasquares(::Type{T}, yaoham::Add{2}) where {T<:Real}
 end
 
 function yaoham2majoranasquares(::Type{T}, yaoham::KronBlock{2}) where {T<:Real}
-    ham = spzeros(T, 2nqubits(yaoham), 2nqubits(yaoham))
     i1, i2 = kron2majoranaindices(yaoham)
-    ham[i1,i2] = 2
-    ham[i2,i1] = -2
-    return ham
+    return SparseMatrixCOO([i1, i2], [i2, i1], T[2, -2], 2nqubits(yaoham), 2nqubits(yaoham))
 end
 
 function yaoham2majoranasquares(::Type{T}, yaoham::PutBlock{2,1,ZGate}) where {T<:Real}
-    ham = spzeros(T, 2nqubits(yaoham), 2nqubits(yaoham))
     i1, i2 = 2yaoham.locs[1]-1, 2yaoham.locs[1]
-    ham[i1,i2] = 2
-    ham[i2,i1] = -2
-    return ham
+    return SparseMatrixCOO([i1, i2], [i2, i1], T[2, -2], 2nqubits(yaoham), 2nqubits(yaoham))
 end
 
 function yaoham2majoranasquares(::Type{T}, yaoham::PutBlock{2,N,<:PauliKronBlock}) where {T<:Real, N}
     areconsecutive(yaoham.locs) || throw(NonFLOException("$(yaoham.content) contains terms that are not the product of two Majoranas"))
-    ham = spzeros(T, 2nqubits(yaoham), 2nqubits(yaoham))
     i1, i2 = 2 * (minimum(yaoham.locs) - 1) .+ kron2majoranaindices(yaoham.content)
-    ham[i1,i2] = 2
-    ham[i2,i1] = -2
-    return ham
+    return SparseMatrixCOO([i1, i2], [i2, i1], T[2, -2], 2nqubits(yaoham), 2nqubits(yaoham))
 end
 
 function yaoham2majoranasquares(::Type{T}, yaoham::Scale) where {T<:Real}
@@ -358,17 +349,18 @@ random_orthogonal_matrix(n) = random_orthogonal_matrix(Float64, n)
 # -------------------------------------------
 # Utilities for fast sparse matrix operations
 # -------------------------------------------
+_rmul!(A::SparseMatrixCOO, α::Real) = SparseMatrixCOO(A.is, A.js, rmul!(A.vs, α), A.m, A.n)
+_rmul!(A::AbstractMatrix, α::Real) = rmul!(A, α)
+
 """
-    fast_add!(A::AbstractMatrix, B::SparseMatrixCSC)
+    fast_add!(A::AbstractMatrix, B::AbstractMatrix)
 
-Fast implementation of `A .+= B` for sparse `B`.
+Fast implementation of `A .+= B` optimised for sparse `B`.
 """
-function fast_add!(A::AbstractMatrix, B::SparseMatrixCSC)
+function fast_add!(A::AbstractMatrix, B::SparseMatrixCOO)
     @assert size(A, 1) == size(B, 1) && size(A, 2) == size(B, 2) "Dimension mismatch"
-    for j = 1:size(B, 2)
-        for k in SparseArrays.nzrange(B, j)
-            @inbounds A[B.rowval[k],j] += B.nzval[k]
-        end
+    for (i, j, v) in zip(B.is, B.js, B.vs)
+        @inbounds A[i, j] += v
     end
     return A
 end
@@ -378,19 +370,16 @@ function fast_add!(A::AbstractMatrix, B::AbstractMatrix)
 end
 
 """
-    fast_overlap(y::AbstractVecOrMat, A::SparseMatrixCSC, x::AbstractVecOrMat)
+    fast_overlap(y::AbstractVecOrMat, A::AbstractMatrix, x::AbstractVecOrMat)
 
-Fast implementation of `tr(y' * A * x)` for sparse `A`.
+Fast implementation of `tr(y' * A * x)` optimised for sparse `A`.
 """
-function fast_overlap(y::AbstractVecOrMat{T1}, A::SparseMatrixCSC{T2}, x::AbstractVecOrMat{T3}) where {T1,T2,T3}
+function fast_overlap(y::AbstractVecOrMat{T1}, A::SparseMatrixCOO{T2}, x::AbstractVecOrMat{T3}) where {T1,T2,T3}
     @assert size(x, 1) == size(A, 2) && size(y, 1) == size(A, 1) && size(x, 2) == size(y, 2) "Dimension mismatch"
     g = zero(promote_type(T1, T2, T3))
-    @inbounds for j = 1:size(A, 2)
-        for k in SparseArrays.nzrange(A, j)
-            i = A.rowval[k]
-            for m = 1:size(y, 2)
-                g += conj(y[i, m]) * A.nzval[k] * x[j, m]
-            end
+    @inbounds for (i, j, v) in zip(A.is, A.js, A.vs)
+        for m = 1:size(y, 2)
+            g += conj(y[i, m]) * v * x[j, m]
         end
     end
     return g

test/Project.toml

@@ -2,7 +2,6 @@
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Yao = "5872b779-8223-5990-8dd0-5abbb0748c8c"

test/runtests.jl

@@ -17,7 +17,7 @@
 
 using Test
 using FLOYao
-using YaoAPI # needed for the doctest tests
+using Yao.YaoAPI # needed for the doctest tests
 using Yao
 using StatsBase
 using Documenter
@@ -46,19 +46,23 @@ end
     nq = 4
     x = randn(ComplexF64, 10, 10)
     y = randn(ComplexF64, 10, 10)
-    A = FLOYao.sprand(ComplexF64, 10, 10, 0.3)
-    r = FLOYao.fast_overlap(y, A, x)
-    @test isapprox(r, tr(y' * A * x), atol=1e-7)
-    @test isapprox(r, y ⋅ (A * x), atol=1e-7)
+    dmat = [rand() < 0.4 ? randn() : 0.0 for i=1:10, j=1:10]
+    B = FLOYao.SparseMatrixCOO(dmat)
+    r2 = FLOYao.fast_overlap(y, B, x)
+    @test isapprox(r2, tr(y' * B * x), atol=1e-7)
+    @test isapprox(r2, y ⋅ (B * x), atol=1e-7)
 end
 
 @testset "fast_add!" begin
     A = randn(10, 10)
-    B = FLOYao.sprand(10, 10, 0.2)
+    dmat = [rand() < 0.4 ? randn() : 0.0 for i=1:10, j=1:10]
+    D = FLOYao.SparseMatrixCOO(dmat)
+
     C = copy(A)
-    @test FLOYao.fast_add!(C, B) ≈ A + B
+    @test FLOYao.fast_add!(C, D) ≈ A + D
     C = copy(A)
-    @test FLOYao.fast_add!(C, Matrix(B)) ≈ A + B
+    @test FLOYao.fast_add!(C, Matrix(D)) ≈ A + D
+
 end
 
 @testset "MajoranaRegister" begin