From 487f6ba6024a07ea71f2a48d0f4d7ea923462a26 Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Thu, 31 Oct 2024 01:12:03 +0100 Subject: [PATCH 1/2] Make clenshaw method as default for wigner --- src/wigner.jl | 116 +++++++++++++++++++++++++++++---------- test/core-test/wigner.jl | 16 +++--- 2 files changed, 94 insertions(+), 38 deletions(-) diff --git a/src/wigner.jl b/src/wigner.jl index 466cbe376..f5be0a7b8 100644 --- a/src/wigner.jl +++ b/src/wigner.jl @@ -12,34 +12,90 @@ end WignerLaguerre(; parallel = false, tol = 1e-14) = WignerLaguerre(parallel, tol) @doc raw""" - wigner(state::QuantumObject, xvec::AbstractVector, yvec::AbstractVector; g::Real=√2, - solver::WignerSolver=WignerLaguerre()) - -Generates the [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) -of `state` at points `xvec + 1im * yvec`. The `g` parameter is a scaling factor related to the value of ``\hbar`` in the -commutation relation ``[x, y] = i \hbar`` via ``\hbar=2/g^2`` giving the default value ``\hbar=1``. - -The `solver` parameter can be either `WignerLaguerre()` or `WignerClenshaw()`. The former uses the Laguerre polynomial -expansion of the Wigner function, while the latter uses the Clenshaw algorithm. The Laguerre expansion is faster for -sparse matrices, while the Clenshaw algorithm is faster for dense matrices. The `WignerLaguerre` solver has an optional -`parallel` parameter which defaults to `true` and uses multithreading to speed up the calculation. + wigner( + state::QuantumObject{DT,OpType}, + xvec::AbstractVector, + yvec::AbstractVector; + g::Real = √2, + method::MySolver = WignerClenshaw(), + ) + +Generates the [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) of `state` at points `xvec + 1im * yvec` in phase space. The `g` parameter is a scaling factor related to the value of ``\hbar`` in the commutation relation ``[x, y] = i \hbar`` via ``\hbar=2/g^2`` giving the default value ``\hbar=1``. + +The `method` parameter can be either `WignerLaguerre()` or `WignerClenshaw()`. The former uses the Laguerre polynomial expansion of the Wigner function, while the latter uses the Clenshaw algorithm. The Laguerre expansion is faster for sparse matrices, while the Clenshaw algorithm is faster for dense matrices. The `WignerLaguerre` method has an optional `parallel` parameter which defaults to `true` and uses multithreading to speed up the calculation. + +# Arguments +- `state::QuantumObject`: The quantum state for which the Wigner function is calculated. It can be either a [`KetQuantumObject`](@ref), [`BraQuantumObject`](@ref), or [`OperatorQuantumObject`](@ref). +- `xvec::AbstractVector`: The x-coordinates of the phase space grid. +- `yvec::AbstractVector`: The y-coordinates of the phase space grid. +- `g::Real`: The scaling factor related to the value of ``\hbar`` in the commutation relation ``[x, y] = i \hbar`` via ``\hbar=2/g^2``. +- `method`: The method used to calculate the Wigner function. It can be either `WignerLaguerre()` or `WignerClenshaw()`, with `WignerClenshaw()` as default. The `WignerLaguerre` method has the optional `parallel` and `tol` parameters, with default values `true` and `1e-14`, respectively. + +# Returns +- `W::Matrix`: The Wigner function of the state at the points `xvec + 1im * yvec` in phase space. + +# Example +``` +julia> ψ = fock(10, 0) + fock(10, 1) |> normalize +Quantum Object: type=Ket dims=[10] size=(10,) +10-element Vector{ComplexF64}: + 0.7071067811865475 + 0.0im + 0.7071067811865475 + 0.0im + 0.0 + 0.0im + 0.0 + 0.0im + 0.0 + 0.0im + 0.0 + 0.0im + 0.0 + 0.0im + 0.0 + 0.0im + 0.0 + 0.0im + 0.0 + 0.0im + +julia> xvec = range(-5, 5, 200) +-5.0:0.05025125628140704:5.0 + +julia> wig = wigner(ψ, xvec, xvec) +200×200 Matrix{Float64}: + 2.63558e-21 4.30187e-21 6.98638e-21 1.12892e-20 1.81505e-20 … 1.50062e-20 9.28736e-21 5.71895e-21 3.50382e-21 + 4.29467e-21 7.00905e-21 1.13816e-20 1.83891e-20 2.9562e-20 2.45173e-20 1.51752e-20 9.3454e-21 5.72614e-21 + 6.96278e-21 1.13621e-20 1.8448e-20 2.98026e-20 4.79043e-20 3.98553e-20 2.46711e-20 1.51947e-20 9.31096e-21 + 1.12314e-20 1.83256e-20 2.97505e-20 4.80558e-20 7.72344e-20 6.4463e-20 3.99074e-20 2.45808e-20 1.50639e-20 + 1.80254e-20 2.94073e-20 4.77351e-20 7.70963e-20 1.23892e-19 1.0374e-19 6.42289e-20 3.95652e-20 2.42491e-20 + ⋮ ⋱ + 1.80254e-20 2.94073e-20 4.77351e-20 7.70963e-20 1.23892e-19 … 1.0374e-19 6.42289e-20 3.95652e-20 2.42491e-20 + 1.12314e-20 1.83256e-20 2.97505e-20 4.80558e-20 7.72344e-20 6.4463e-20 3.99074e-20 2.45808e-20 1.50639e-20 + 6.96278e-21 1.13621e-20 1.8448e-20 2.98026e-20 4.79043e-20 3.98553e-20 2.46711e-20 1.51947e-20 9.31096e-21 + 4.29467e-21 7.00905e-21 1.13816e-20 1.83891e-20 2.9562e-20 2.45173e-20 1.51752e-20 9.3454e-21 5.72614e-21 + 2.63558e-21 4.30187e-21 6.98638e-21 1.12892e-20 1.81505e-20 1.50062e-20 9.28736e-21 5.71895e-21 3.50382e-21 +``` + +or taking advantage of the parallel computation of the `WignerLaguerre` method + +``` +julia> wig = wigner(ρ, xvec, xvec, method=WignerLaguerre(parallel=true)) +200×200 Matrix{Float64}: + 2.63558e-21 4.30187e-21 6.98638e-21 1.12892e-20 1.81505e-20 … 1.50062e-20 9.28736e-21 5.71895e-21 3.50382e-21 + 4.29467e-21 7.00905e-21 1.13816e-20 1.83891e-20 2.9562e-20 2.45173e-20 1.51752e-20 9.3454e-21 5.72614e-21 + 6.96278e-21 1.13621e-20 1.8448e-20 2.98026e-20 4.79043e-20 3.98553e-20 2.46711e-20 1.51947e-20 9.31096e-21 + 1.12314e-20 1.83256e-20 2.97505e-20 4.80558e-20 7.72344e-20 6.4463e-20 3.99074e-20 2.45808e-20 1.50639e-20 + 1.80254e-20 2.94073e-20 4.77351e-20 7.70963e-20 1.23892e-19 1.0374e-19 6.42289e-20 3.95652e-20 2.42491e-20 + ⋮ ⋱ + 1.80254e-20 2.94073e-20 4.77351e-20 7.70963e-20 1.23892e-19 … 1.0374e-19 6.42289e-20 3.95652e-20 2.42491e-20 + 1.12314e-20 1.83256e-20 2.97505e-20 4.80558e-20 7.72344e-20 6.4463e-20 3.99074e-20 2.45808e-20 1.50639e-20 + 6.96278e-21 1.13621e-20 1.8448e-20 2.98026e-20 4.79043e-20 3.98553e-20 2.46711e-20 1.51947e-20 9.31096e-21 + 4.29467e-21 7.00905e-21 1.13816e-20 1.83891e-20 2.9562e-20 2.45173e-20 1.51752e-20 9.3454e-21 5.72614e-21 + 2.63558e-21 4.30187e-21 6.98638e-21 1.12892e-20 1.81505e-20 1.50062e-20 9.28736e-21 5.71895e-21 3.50382e-21 +``` """ function wigner( - state::QuantumObject{<:AbstractArray{T},OpType}, + state::QuantumObject{DT,OpType}, xvec::AbstractVector, yvec::AbstractVector; g::Real = √2, - solver::MySolver = WignerLaguerre(), -) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject},MySolver<:WignerSolver} - if isket(state) - ρ = (state * state').data - elseif isbra(state) - ρ = (state' * state).data - else - ρ = state.data - end + method::MySolver = WignerClenshaw(), +) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject},MySolver<:WignerSolver} + ρ = ket2dm(state).data - return _wigner(ρ, xvec, yvec, g, solver) + return _wigner(ρ, xvec, yvec, g, method) end function _wigner( @@ -47,7 +103,7 @@ function _wigner( xvec::AbstractVector{T}, yvec::AbstractVector{T}, g::Real, - solver::WignerLaguerre, + method::WignerLaguerre, ) where {T<:BlasFloat} g = convert(T, g) X, Y = meshgrid(xvec, yvec) @@ -55,7 +111,7 @@ function _wigner( W = similar(A, T) W .= 0 - return _wigner_laguerre(ρ, A, W, g, solver) + return _wigner_laguerre(ρ, A, W, g, method) end function _wigner( @@ -63,7 +119,7 @@ function _wigner( xvec::AbstractVector{T}, yvec::AbstractVector{T}, g::Real, - solver::WignerClenshaw, + method::WignerClenshaw, ) where {T1<:BlasFloat,T<:BlasFloat} g = convert(T, g) M = size(ρ, 1) @@ -90,11 +146,11 @@ function _wigner( return @. real(W) * exp(-B / 2) * g^2 / 2 / π end -function _wigner_laguerre(ρ::AbstractSparseArray, A::AbstractArray, W::AbstractArray, g::Real, solver::WignerLaguerre) +function _wigner_laguerre(ρ::AbstractSparseArray, A::AbstractArray, W::AbstractArray, g::Real, method::WignerLaguerre) rows, cols, vals = findnz(ρ) B = @. 4 * abs2(A) - if solver.parallel + if method.parallel iter = filter(x -> x[2] >= x[1], collect(zip(rows, cols, vals))) Wtot = similar(B, size(B)..., length(iter)) Threads.@threads for i in eachindex(iter) @@ -122,12 +178,12 @@ function _wigner_laguerre(ρ::AbstractSparseArray, A::AbstractArray, W::Abstract return @. W * g^2 * exp(-B / 2) / 2 / π end -function _wigner_laguerre(ρ::AbstractArray, A::AbstractArray, W::AbstractArray, g::Real, solver::WignerLaguerre) - tol = solver.tol +function _wigner_laguerre(ρ::AbstractArray, A::AbstractArray, W::AbstractArray, g::Real, method::WignerLaguerre) + tol = method.tol M = size(ρ, 1) B = @. 4 * abs2(A) - if solver.parallel + if method.parallel throw(ArgumentError("Parallel version is not implemented for dense matrices")) else for m in 0:M-1 diff --git a/test/core-test/wigner.jl b/test/core-test/wigner.jl index ae30b188e..e12a12c2e 100644 --- a/test/core-test/wigner.jl +++ b/test/core-test/wigner.jl @@ -5,10 +5,10 @@ xvec = LinRange(-3, 3, 300) yvec = LinRange(-3, 3, 300) - wig = wigner(ψ, xvec, yvec, solver = WignerLaguerre(tol = 1e-6)) - wig2 = wigner(ρ, xvec, yvec, solver = WignerLaguerre(parallel = false)) - wig3 = wigner(ρ, xvec, yvec, solver = WignerLaguerre(parallel = true)) - wig4 = wigner(ψ, xvec, yvec, solver = WignerClenshaw()) + wig = wigner(ψ, xvec, yvec, method = WignerLaguerre(tol = 1e-6)) + wig2 = wigner(ρ, xvec, yvec, method = WignerLaguerre(parallel = false)) + wig3 = wigner(ρ, xvec, yvec, method = WignerLaguerre(parallel = true)) + wig4 = wigner(ψ, xvec, yvec, method = WignerClenshaw()) @test sqrt(sum(abs.(wig2 .- wig)) / length(wig)) < 1e-3 @test sqrt(sum(abs.(wig3 .- wig)) / length(wig)) < 1e-3 @@ -22,9 +22,9 @@ @test sqrt(sum(abs.(wig2 .- wig)) / length(wig)) < 0.1 @testset "Type Inference (wigner)" begin - @inferred wigner(ψ, xvec, yvec, solver = WignerLaguerre(tol = 1e-6)) - @inferred wigner(ρ, xvec, yvec, solver = WignerLaguerre(parallel = false)) - @inferred wigner(ρ, xvec, yvec, solver = WignerLaguerre(parallel = true)) - @inferred wigner(ψ, xvec, yvec, solver = WignerClenshaw()) + @inferred wigner(ψ, xvec, yvec, method = WignerLaguerre(tol = 1e-6)) + @inferred wigner(ρ, xvec, yvec, method = WignerLaguerre(parallel = false)) + @inferred wigner(ρ, xvec, yvec, method = WignerLaguerre(parallel = true)) + @inferred wigner(ψ, xvec, yvec, method = WignerClenshaw()) end end From b97f64d541ba2d5b9a3632731b3fb2a2461f6300 Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Thu, 31 Oct 2024 10:45:19 +0100 Subject: [PATCH 2/2] Minor changes --- src/wigner.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/wigner.jl b/src/wigner.jl index f5be0a7b8..c64bf8181 100644 --- a/src/wigner.jl +++ b/src/wigner.jl @@ -17,7 +17,7 @@ WignerLaguerre(; parallel = false, tol = 1e-14) = WignerLaguerre(parallel, tol) xvec::AbstractVector, yvec::AbstractVector; g::Real = √2, - method::MySolver = WignerClenshaw(), + method::WignerSolver = WignerClenshaw(), ) Generates the [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) of `state` at points `xvec + 1im * yvec` in phase space. The `g` parameter is a scaling factor related to the value of ``\hbar`` in the commutation relation ``[x, y] = i \hbar`` via ``\hbar=2/g^2`` giving the default value ``\hbar=1``. @@ -29,7 +29,7 @@ The `method` parameter can be either `WignerLaguerre()` or `WignerClenshaw()`. T - `xvec::AbstractVector`: The x-coordinates of the phase space grid. - `yvec::AbstractVector`: The y-coordinates of the phase space grid. - `g::Real`: The scaling factor related to the value of ``\hbar`` in the commutation relation ``[x, y] = i \hbar`` via ``\hbar=2/g^2``. -- `method`: The method used to calculate the Wigner function. It can be either `WignerLaguerre()` or `WignerClenshaw()`, with `WignerClenshaw()` as default. The `WignerLaguerre` method has the optional `parallel` and `tol` parameters, with default values `true` and `1e-14`, respectively. +- `method::WignerSolver`: The method used to calculate the Wigner function. It can be either `WignerLaguerre()` or `WignerClenshaw()`, with `WignerClenshaw()` as default. The `WignerLaguerre` method has the optional `parallel` and `tol` parameters, with default values `true` and `1e-14`, respectively. # Returns - `W::Matrix`: The Wigner function of the state at the points `xvec + 1im * yvec` in phase space. @@ -91,8 +91,8 @@ function wigner( xvec::AbstractVector, yvec::AbstractVector; g::Real = √2, - method::MySolver = WignerClenshaw(), -) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject},MySolver<:WignerSolver} + method::WignerSolver = WignerClenshaw(), +) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} ρ = ket2dm(state).data return _wigner(ρ, xvec, yvec, g, method)