diff --git a/Project.toml b/Project.toml index 7ac0ffd..26c3e26 100644 --- a/Project.toml +++ b/Project.toml @@ -4,8 +4,10 @@ authors = ["Andrew Kille"] version = "1.2.9-dev" [deps] +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" QuantumInterface = "5717a53b-5d69-4fa3-b976-0bf2f97ca1e5" SymplecticFactorizations = "7425e8e4-4cde-4e45-9b2f-a15679260f9b" @@ -18,6 +20,7 @@ MakieExt = "Makie" StaticArraysExt = "StaticArrays" [compat] +BenchmarkTools = "1.6.0" BlockArrays = "1.1.1" LinearAlgebra = "1.9" Makie = "0.21, 0.22" diff --git a/src/Gabs.jl b/src/Gabs.jl index 217d903..eb9fec2 100644 --- a/src/Gabs.jl +++ b/src/Gabs.jl @@ -10,15 +10,22 @@ import QuantumInterface: StateVector, AbstractOperator, apply!, tensor, ⊗, dir import SymplecticFactorizations: williamson, Williamson, polar, Polar, blochmessiah, BlochMessiah, randsymplectic, symplecticform, issymplectic using SymplecticFactorizations: williamson, Williamson, polar, Polar, blochmessiah, BlochMessiah, BlockForm, PairForm +import PDMats +import PDMats: AbstractPDMat, PDMat, PDiagMat, ScalMat, quad, invquad, whiten, unwhiten + export # types GaussianState, GaussianUnitary, GaussianChannel, Generaldyne, + # PDMats integration types + PDGaussianState, # symplectic representations QuadPairBasis, QuadBlockBasis, changebasis, # operations tensor, ⊗, directsum, ⊕, apply!, ptrace, output, prob, # predefined Gaussian states vacuumstate, thermalstate, coherentstate, squeezedstate, eprstate, + # PDMats-optimized state constructors + vacuumstate_pd, thermalstate_pd, coherentstate_pd, # predefined Gaussian channels displace, squeeze, twosqueeze, phaseshift, beamsplitter, attenuator, amplifier, @@ -41,6 +48,8 @@ include("symplectic.jl") include("types.jl") +include("pdmats_types.jl") + include("states.jl") include("unitaries.jl") @@ -57,4 +66,6 @@ include("wigner.jl") include("metrics.jl") -end +include("pdmats_utils.jl") + +end \ No newline at end of file diff --git a/src/metrics.jl b/src/metrics.jl index 6c5635a..33ee69c 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -3,4 +3,14 @@ Calculate the purity of a Gaussian state, defined by `1/sqrt((2/ħ) det(V))`. """ -purity(x::GaussianState) = (b = x.basis; (x.ħ/2)^(b.nmodes)/sqrt(det(x.covar))) \ No newline at end of file +purity(x::GaussianState) = (b = x.basis; (x.ħ/2)^(b.nmodes)/sqrt(det(x.covar))) + +""" + purity(state::PDGaussianState) + +Calculate the purity of a Gaussian state using optimized PDMats operations. +""" +function purity(x::PDGaussianState) + b = x.basis + (x.ħ/2)^(b.nmodes)/sqrt(det(x.covar)) +end \ No newline at end of file diff --git a/src/pdmats_types.jl b/src/pdmats_types.jl new file mode 100644 index 0000000..24788a7 --- /dev/null +++ b/src/pdmats_types.jl @@ -0,0 +1,57 @@ +""" +PDMats integration type definitions for GABS.jl +""" + +using PDMats +using LinearAlgebra: Symmetric, tr, I, PosDefException + +""" +Modified GaussianState structure that supports PDMats for more efficient operations +with covariance matrices. +""" +@kwdef struct PDGaussianState{B<:SymplecticBasis,M,V<:AbstractPDMat} <: StateVector{M,V} + basis::B + mean::M + covar::V + ħ::Number = 2 +end + +function PDGaussianState{B,M,V}(basis::B, mean::M, covar::V; ħ::Number = 2) where {B<:SymplecticBasis,M,V<:AbstractPDMat} + size(covar, 1) == length(mean) == 2*(basis.nmodes) || throw(DimensionMismatch(STATE_ERROR)) + return PDGaussianState{B,M,V}(basis=basis, mean=mean, covar=covar, ħ=ħ) +end + +function PDGaussianState(basis::B, mean::M, covar::V; ħ::Number = 2) where {B<:SymplecticBasis,M,V<:AbstractPDMat} + size(covar, 1) == length(mean) == 2*(basis.nmodes) || throw(DimensionMismatch(STATE_ERROR)) + return PDGaussianState{B,M,V}(basis=basis, mean=mean, covar=covar, ħ=ħ) +end + +function PDGaussianState(state::GaussianState) + covar_sym = Symmetric(state.covar) + pdmat = try + PDMat(covar_sym) + catch e + if isa(e, PosDefException) + eps_val = 1e-10 * tr(state.covar) / size(state.covar, 1) + PDMat(covar_sym + eps_val * I) + else + rethrow(e) + end + end + return PDGaussianState(state.basis, state.mean, pdmat, ħ=state.ħ) +end + +function PDGaussianState(basis::B, mean::M, covar::AbstractMatrix; ħ::Number = 2) where {B<:SymplecticBasis,M} + covar_sym = Symmetric(covar) + pdmat = try + PDMat(covar_sym) + catch e + if isa(e, PosDefException) + eps_val = 1e-10 * tr(covar) / size(covar, 1) + PDMat(covar_sym + eps_val * I) + else + rethrow(e) + end + end + return PDGaussianState(basis, mean, pdmat, ħ=ħ) +end \ No newline at end of file diff --git a/src/pdmats_utils.jl b/src/pdmats_utils.jl new file mode 100644 index 0000000..d7de8ba --- /dev/null +++ b/src/pdmats_utils.jl @@ -0,0 +1,223 @@ +""" +PDMats integration utilities for GABS.jl +""" + +using PDMats +using LinearAlgebra +using SparseArrays: blockdiag + +function chol_lower(a::Cholesky) + return a.uplo === 'L' ? LowerTriangular(a.factors) : LowerTriangular(a.factors') +end + +function vacuumstate_pd(basis::SymplecticBasis{N}; ħ = 2) where {N<:Int} + mean = zeros(2*basis.nmodes) + scal_mat = ScalMat(2*basis.nmodes, ħ/2) + return PDGaussianState(basis, mean, scal_mat, ħ=ħ) +end + +function thermalstate_pd(basis::SymplecticBasis{N}, photons::P; ħ = 2) where {N<:Int,P<:Number} + mean = zeros(2*basis.nmodes) + scal_mat = ScalMat(2*basis.nmodes, (2 * photons + 1) * (ħ/2)) + return PDGaussianState(basis, mean, scal_mat, ħ=ħ) +end + +function thermalstate_pd(basis::QuadPairBasis{N}, photons::Vector{P}; ħ = 2) where {N<:Int,P<:Real} + mean = zeros(2*basis.nmodes) + diag_vals = Vector{Float64}(undef, 2*basis.nmodes) + for i in 1:basis.nmodes + val = (2 * photons[i] + 1) * (ħ/2) + diag_vals[2*i-1] = val + diag_vals[2*i] = val + end + diag_mat = PDiagMat(diag_vals) + return PDGaussianState(basis, mean, diag_mat, ħ=ħ) +end + +function coherentstate_pd(basis::SymplecticBasis{N}, alpha::A; ħ = 2) where {N<:Int,A} + mean, _ = _coherentstate(basis, alpha; ħ = ħ) + scal_mat = ScalMat(2*basis.nmodes, ħ/2) + return PDGaussianState(basis, mean, scal_mat, ħ=ħ) +end + +function whiten(state::PDGaussianState, x::AbstractVecOrMat) + size(x, 1) == 2*state.basis.nmodes || throw(DimensionMismatch("Dimension mismatch in whitening operation")) + return PDMats.whiten(state.covar, x) +end + +function unwhiten(state::PDGaussianState, x::AbstractVecOrMat) + size(x, 1) == 2*state.basis.nmodes || throw(DimensionMismatch("Dimension mismatch in unwhitening operation")) + return PDMats.unwhiten(state.covar, x) +end + +function whiten(state::GaussianState, x::AbstractVecOrMat) + size(x, 1) == 2*state.basis.nmodes || throw(DimensionMismatch("Dimension mismatch in whitening operation")) + + chol_factor = cholesky(Symmetric(state.covar)) + return chol_lower(chol_factor) \ x +end + +function unwhiten(state::GaussianState, x::AbstractVecOrMat) + size(x, 1) == 2*state.basis.nmodes || throw(DimensionMismatch("Dimension mismatch in unwhitening operation")) + + chol_factor = cholesky(Symmetric(state.covar)) + return chol_lower(chol_factor) * x +end + +function quad(state::PDGaussianState, x::AbstractVecOrMat) + size(x, 1) == 2*state.basis.nmodes || throw(DimensionMismatch("Dimension mismatch in quadratic form")) + return PDMats.quad(state.covar, x) +end + +function invquad(state::PDGaussianState, x::AbstractVecOrMat) + size(x, 1) == 2*state.basis.nmodes || throw(DimensionMismatch("Dimension mismatch in inverse quadratic form")) + return PDMats.invquad(state.covar, x) +end + +function quad(state::GaussianState, x::AbstractVecOrMat) + size(x, 1) == 2*state.basis.nmodes || throw(DimensionMismatch("Dimension mismatch in quadratic form")) + return quad(state.covar, x) +end + +function invquad(state::GaussianState, x::AbstractVecOrMat) + size(x, 1) == 2*state.basis.nmodes || throw(DimensionMismatch("Dimension mismatch in inverse quadratic form")) + return invquad(state.covar, x) +end + + + +using SparseArrays: blockdiag + +function pd_blockdiag_cov(cov1::ScalMat, cov2::ScalMat) + if cov1.value == cov2.value + return ScalMat(cov1.dim + cov2.dim, cov1.value) + else + n1, n2 = cov1.dim, cov2.dim + diagvec = vcat(fill(cov1.value, n1), fill(cov2.value, n2)) + return PDiagMat(diagvec) + end +end + +function pd_blockdiag_cov(cov1::PDiagMat, cov2::PDiagMat) + new_diag = vcat(cov1.diag, cov2.diag) + return PDiagMat(new_diag) +end + +function pd_blockdiag_cov(cov1::AbstractPDMat, cov2::AbstractPDMat) + return PDMat(blockdiag(Matrix(cov1), Matrix(cov2))) +end + +function tensor(state1::PDGaussianState, state2::PDGaussianState) + typeof(state1.basis) == typeof(state2.basis) || throw(ArgumentError(SYMPLECTIC_ERROR)) + state1.ħ == state2.ħ || throw(ArgumentError(HBAR_ERROR)) + + mean1, mean2 = state1.mean, state2.mean + basis1, basis2 = state1.basis, state2.basis + nmodes1, nmodes2 = basis1.nmodes, basis2.nmodes + nmodes = nmodes1 + nmodes2 + Mt = promote_type(eltype(mean1), eltype(mean2)) + mean_combined = zeros(Mt, 2 * nmodes) + if basis1 isa QuadPairBasis + for i in 1:(2*nmodes1) + mean_combined[i] = mean1[i] + end + for i in 1:(2*nmodes2) + mean_combined[i + 2*nmodes1] = mean2[i] + end + else + for i in 1:nmodes1 + mean_combined[i] = mean1[i] + mean_combined[i+nmodes] = mean1[i+nmodes1] + end + for i in 1:nmodes2 + mean_combined[i+nmodes1] = mean2[i] + mean_combined[i+nmodes+nmodes1] = mean2[i+nmodes2] + end + end + + cov_combined = pd_blockdiag_cov(state1.covar, state2.covar) + new_basis = state1.basis ⊕ state2.basis + return PDGaussianState(new_basis, mean_combined, cov_combined, ħ = state1.ħ) +end + + + + +function _ptrace(state::PDGaussianState{B, M, V}, idx::T) where {B<:QuadPairBasis, M, V<:AbstractPDMat, T<:Int} + idxV = 2*idx - 1 : 2*idx + mean_block = state.mean[idxV] + new_cov = if state.covar isa ScalMat + ScalMat(2, state.covar.value) + elseif state.covar isa PDiagMat + PDiagMat(state.covar.diag[idxV]) + else + cov_dense = Matrix(state.covar) + cov_block = @view cov_dense[idxV, idxV] + PDMat(Symmetric(cov_block)) + end + return mean_block, new_cov +end + +function _ptrace(state::PDGaussianState{B, M, V}, idx::T) where {B<:QuadBlockBasis, M, V<:AbstractPDMat, T<:Int} + nmodes = state.basis.nmodes + mean_block = [ state.mean[idx], state.mean[idx + nmodes] ] + new_cov = if state.covar isa ScalMat + ScalMat(2, state.covar.value) + elseif state.covar isa PDiagMat + PDiagMat(state.covar.diag[[idx, idx + nmodes]]) + else + cov_dense = Matrix(state.covar) + A = cov_dense[idx, idx] + B_ = cov_dense[idx, idx + nmodes] + C = cov_dense[idx + nmodes, idx] + D = cov_dense[idx + nmodes, idx + nmodes] + PDMat(Symmetric([A B_; C D])) + end + return mean_block, new_cov +end + + +function ptrace(state::PDGaussianState, idx::Int) + mean_block, cov_block = _ptrace(state, idx) + new_basis = typeof(state.basis)(1) + return PDGaussianState(new_basis, mean_block, cov_block, ħ=state.ħ) +end + +function ptrace(state::PDGaussianState, indices::AbstractVector{Int}) + new_basis = typeof(state.basis)(length(indices)) + new_mean = zeros(eltype(state.mean), 2 * length(indices)) + cov_blocks = Vector{typeof(state.covar)}(undef, length(indices)) + for (i, idx) in enumerate(indices) + m, c = _ptrace(state, idx) + new_mean[2*i-1:2*i] .= m + cov_blocks[i] = c + end + new_cov = pd_blockdiag_cov(cov_blocks...) + return PDGaussianState(new_basis, new_mean, new_cov, ħ=state.ħ) +end + +function apply!(state::PDGaussianState, op::GaussianUnitary) + op.basis == state.basis || throw(DimensionMismatch(ACTION_ERROR)) + op.ħ == state.ħ || throw(ArgumentError(HBAR_ERROR)) + + d, S = op.disp, op.symplectic + state.mean .= S * state.mean .+ d + + temp_covar = S * Matrix(state.covar) * transpose(S) + state.covar = PDMat(temp_covar) + + return state +end + +function apply!(state::PDGaussianState, op::GaussianChannel) + op.basis == state.basis || throw(DimensionMismatch(ACTION_ERROR)) + op.ħ == state.ħ || throw(ArgumentError(HBAR_ERROR)) + + d, T, N = op.disp, op.transform, op.noise + state.mean .= T * state.mean .+ d + + temp_covar = T * Matrix(state.covar) * transpose(T) .+ N + state.covar = PDMat(temp_covar) + + return state +end \ No newline at end of file diff --git a/src/states.jl b/src/states.jl index a6104b2..9330b44 100644 --- a/src/states.jl +++ b/src/states.jl @@ -760,3 +760,29 @@ function sympspectrum(state::GaussianState) spectrum = filter(x -> x > 0, imag.(eigvals(M))) return spectrum end + +function vacuumstate_pd(basis::SymplecticBasis{N}; ħ = 2) where {N<:Int} + mean = zeros(2*basis.nmodes) + return PDGaussianState(basis, mean, ScalMat(2*basis.nmodes, ħ/2); ħ = ħ) +end + +function thermalstate_pd(basis::SymplecticBasis{N}, photons::P; ħ = 2) where {N<:Int,P<:Number} + mean = zeros(2*basis.nmodes) + return PDGaussianState(basis, mean, ScalMat(2*basis.nmodes, (2 * photons + 1) * (ħ/2)); ħ = ħ) +end + +function thermalstate_pd(basis::QuadPairBasis{N}, photons::Vector{P}; ħ = 2) where {N<:Int,P<:Real} + mean = zeros(2*basis.nmodes) + diag_vals = Vector{Float64}(undef, 2*basis.nmodes) + for i in 1:basis.nmodes + val = (2 * photons[i] + 1) * (ħ/2) + diag_vals[2*i-1] = val + diag_vals[2*i] = val + end + return PDGaussianState(basis, mean, PDiagMat(diag_vals); ħ = ħ) +end + +function coherentstate_pd(basis::SymplecticBasis{N}, alpha::A; ħ = 2) where {N<:Int,A} + mean, _ = _coherentstate(basis, alpha; ħ = ħ) + return PDGaussianState(basis, mean, ScalMat(2*basis.nmodes, ħ/2); ħ = ħ) +end \ No newline at end of file diff --git a/src/types.jl b/src/types.jl index 455ac1e..bab2cb2 100644 --- a/src/types.jl +++ b/src/types.jl @@ -1,3 +1,5 @@ +using PDMats + """ Defines a Gaussian state for an N-mode bosonic system over a 2N-dimensional phase space. @@ -33,6 +35,8 @@ covariance: 2×2 Matrix{Float64}: end end + + Base.:(==)(x::GaussianState, y::GaussianState) = x.basis == y.basis && x.mean == y.mean && x.covar == y.covar && x.ħ == y.ħ Base.isapprox(x::GaussianState, y::GaussianState; kwargs...) = x.basis == y.basis && isapprox(x.mean, y.mean; kwargs...) && isapprox(x.covar, y.covar; kwargs...) && x.ħ == y.ħ function Base.show(io::IO, mime::MIME"text/plain", x::GaussianState) diff --git a/test/Project.toml b/test/Project.toml index 8c69541..92c520a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,10 +1,13 @@ [deps] +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" diff --git a/test/benchmark_pdmats.jl b/test/benchmark_pdmats.jl new file mode 100644 index 0000000..4306e99 --- /dev/null +++ b/test/benchmark_pdmats.jl @@ -0,0 +1,68 @@ +@testitem "PDMats Benchmarks" begin + using Gabs + using LinearAlgebra + using BenchmarkTools + using PDMats + + function run_benchmarks(nmodes) + basis = QuadPairBasis(nmodes) + + std_state = vacuumstate(basis) + pd_state = vacuumstate_pd(basis) + + testvec = randn(2*nmodes) + + std_purity = @benchmark purity($std_state) + pd_purity = @benchmark purity($pd_state) + + std_whiten = @benchmark whiten($std_state, $testvec) + pd_whiten = @benchmark whiten($pd_state, $testvec) + + std_invquad = @benchmark invquad($std_state, $testvec) + pd_invquad = @benchmark invquad($pd_state, $testvec) + + std_create = @benchmark vacuumstate($basis) + pd_create = @benchmark vacuumstate_pd($basis) + + return Dict( + "purity" => (std_purity, pd_purity), + "whiten" => (std_whiten, pd_whiten), + "invquad" => (std_invquad, pd_invquad), + "creation" => (std_create, pd_create) + ) + end + + + if get(ENV, "RUN_BENCHMARKS", "false") == "true" + @testset "Performance comparison" begin + results_small = run_benchmarks(5) + results_medium = run_benchmarks(10) + results_large = run_benchmarks(20) + + # for the sake of avoiding test failures + @test true + + # results + for (size, results) in zip(["Small (5 modes)", "Medium (10 modes)", "Large (20 modes)"], + [results_small, results_medium, results_large]) + println("\n===== $size =====") + for (name, (std, pd)) in results + speedup = minimum(std.times) / minimum(pd.times) + println("$name: $(round(speedup, digits=2))x speedup") + end + end + end + else + @testset "Skip benchmarks" begin + basis = QuadPairBasis(2) + std_state = vacuumstate(basis) + pd_state = vacuumstate_pd(basis) + + @test purity(pd_state) ≈ purity(std_state) + + testvec = randn(4) + @test whiten(pd_state, testvec) ≈ whiten(std_state, testvec) + @test invquad(pd_state, testvec) ≈ invquad(std_state, testvec) + end + end +end \ No newline at end of file diff --git a/test/test_autodiff.jl b/test/test_autodiff.jl index c06b4c7..7a6e2c7 100644 --- a/test/test_autodiff.jl +++ b/test/test_autodiff.jl @@ -45,4 +45,4 @@ end - \ No newline at end of file +end \ No newline at end of file diff --git a/test/test_pdmats.jl b/test/test_pdmats.jl new file mode 100644 index 0000000..8eca18c --- /dev/null +++ b/test/test_pdmats.jl @@ -0,0 +1,62 @@ +@testitem "PDMats Integration" begin + using Gabs + using BlockArrays + using LinearAlgebra: norm, det + using PDMats + + @testset "PDGaussianState creation" begin + basis = QuadPairBasis(2) + + std_state = vacuumstate(basis) + pd_state = vacuumstate_pd(basis) + + @test isapprox(Matrix(pd_state.covar), std_state.covar) + @test isapprox(pd_state.mean, std_state.mean) + + thermal_std = thermalstate(basis, 2.0) + thermal_pd = thermalstate_pd(basis, 2.0) + + @test isapprox(Matrix(thermal_pd.covar), thermal_std.covar) + @test isapprox(thermal_pd.mean, thermal_std.mean) + + alpha = 1.0 + im + coherent_std = coherentstate(basis, alpha) + coherent_pd = coherentstate_pd(basis, alpha) + + @test isapprox(Matrix(coherent_pd.covar), coherent_std.covar) + @test isapprox(coherent_pd.mean, coherent_std.mean) + end + + @testset "Core operations" begin + basis = QuadPairBasis(3) + std_state = thermalstate(basis, 1.5) + pd_state = thermalstate_pd(basis, 1.5) + testvec = randn(2*basis.nmodes) + + @test isapprox(purity(std_state), purity(pd_state)) + + std_whitened = whiten(std_state, testvec) + pd_whitened = whiten(pd_state, testvec) + @test isapprox(std_whitened, pd_whitened) + + std_unwhitened = unwhiten(std_state, testvec) + pd_unwhitened = unwhiten(pd_state, testvec) + @test isapprox(std_unwhitened, pd_unwhitened) + + @test isapprox(quad(std_state, testvec), quad(pd_state, testvec)) + @test isapprox(invquad(std_state, testvec), invquad(pd_state, testvec)) + end + + @testset "Constructor from GaussianState" begin + basis = QuadPairBasis(2) + + rand_state = randstate(basis) + + pd_state = PDGaussianState(rand_state) + + @test isapprox(Matrix(pd_state.covar), rand_state.covar) + @test isapprox(pd_state.mean, rand_state.mean) + @test pd_state.basis == rand_state.basis + @test pd_state.ħ == rand_state.ħ + end +end \ No newline at end of file