diff --git a/Project.toml b/Project.toml index 7a1775c..22ba715 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "FLOYao" uuid = "6d9310a3-f1d0-41b7-8edb-11c1cf57cd2d" authors = ["janlukas.bosse and contributors"] -version = "1.4.0" +version = "1.4.1" [deps] LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/src/auto_diff.jl b/src/auto_diff.jl index eaca00c..d734d5e 100644 --- a/src/auto_diff.jl +++ b/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 diff --git a/src/utils.jl b/src/utils.jl index fa27e81..a769daa 100644 --- a/src/utils.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 8d90dca..4eb3281 100644 --- a/test/runtests.jl +++ b/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)