Merge pull request #10 from QuantumBFS/jg/improve-sparsematrix-operat…

…ions

Fast overlap and fast add of sparse matrices

Project.toml

@@ -1,7 +1,7 @@
 name = "FLOYao"
 uuid = "6d9310a3-f1d0-41b7-8edb-11c1cf57cd2d"
 authors = ["janlukas.bosse <[email protected]> and contributors"]
-version = "1.4.0"
+version = "1.4.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/auto_diff.jl

@@ -38,7 +38,7 @@ function Yao.AD.backward_params!(st::Tuple{<:MajoranaReg,<:MajoranaReg},
     out, outδ = st
     ham = Yao.AD.generator(block)
     majoranaham = yaoham2majoranasquares(ham)
-    g = outδ.state ⋅ (majoranaham * out.state) / 4
+    g = fast_overlap(outδ.state, majoranaham, out.state) / 4
     pushfirst!(collector, g)
     return nothing
 end
@@ -48,7 +48,7 @@ function Yao.AD.backward_params!(st::Tuple{<:MajoranaReg,<:MajoranaReg},
     out, outδ = st
     ham = block.H
     majoranaham = yaoham2majoranasquares(ham)
-    g = outδ.state ⋅ (majoranaham * out.state) / 2
+    g = fast_overlap(outδ.state, majoranaham, out.state) / 2
     pushfirst!(collector, g)
     return nothing
 end

src/utils.jl

@@ -252,15 +252,16 @@ Convert a hamiltonian written as a YaoBlock into the corresponding
 """
 function yaoham2majoranasquares(::Type{T}, yaoham::Add{2}) where {T<:Real}
     ham = zeros(T, 2nqubits(yaoham), 2nqubits(yaoham))
-    for k in yaoham
+    @inbounds for k in yaoham
         if k isa Scale
-            ham += k.alpha * yaoham2majoranasquares(T, k.content)
+            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)
             ham[i1,i2] += 2
             ham[i2,i1] -= 2
         elseif k isa Add
-            ham += yaoham2majoranasquares(T, k)
+            fast_add!(ham, yaoham2majoranasquares(T, k))
         elseif k isa PutBlock{2,1,ZGate} 
             i1, i2 = kron2majoranaindices(k)
             ham[i1,i2] += 2
@@ -348,3 +349,49 @@ function random_orthogonal_matrix(::Type{T}, n) where {T}
 end
 
 random_orthogonal_matrix(n) = random_orthogonal_matrix(Float64, n)
+
+
+# -------------------------------------------
+# Utilities for fast sparse matrix operations
+# -------------------------------------------
+"""
+    fast_add!(A::AbstractMatrix, B::SparseMatrixCSC)
+
+Fast implementation of `A .+= B` for sparse `B`.
+"""
+function fast_add!(A::AbstractMatrix, B::SparseMatrixCSC)
+    @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
+    end
+    return A
+end
+
+function fast_add!(A::AbstractMatrix, B::AbstractMatrix)
+    return A .+= B
+end
+
+"""
+    fast_overlap(y::AbstractVecOrMat, A::SparseMatrixCSC, x::AbstractVecOrMat)
+
+Fast implementation of `tr(y' * A * x)` for sparse `A`.
+"""
+function fast_overlap(y::AbstractVecOrMat{T1}, A::SparseMatrixCSC{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
+        end
+    end
+    return g
+end
+
+function fast_overlap(y::AbstractVecOrMat{T1}, A::AbstractMatrix{T2}, x::AbstractVecOrMat{T3}) where {T1,T2,T3}
+    return y ⋅ (A * x)
+end

test/runtests.jl

@@ -42,6 +42,25 @@ function ising_hamiltonian(nq, J, h)
     return hamiltonian
 end
 
+@testset "fast_overlap" begin
+    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)
+end
+
+@testset "fast_add!" begin
+    A = randn(10, 10)
+    B = FLOYao.sprand(10, 10, 0.2)
+    C = copy(A)
+    @test FLOYao.fast_add!(C, B) ≈ A + B
+    C = copy(A)
+    @test FLOYao.fast_add!(C, Matrix(B)) ≈ A + B
+end
+
 @testset "MajoranaRegister" begin
     nq = 2
     mreg = FLOYao.zero_state(nq)