Skip to content

Make clenshaw method as default for wigner #279

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Oct 31, 2024
Merged

Make clenshaw method as default for wigner #279

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
487f6ba
Make clenshaw method as default for wigner
albertomercurio Oct 31, 2024
b97f64d
Minor changes
albertomercurio Oct 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 86 additions & 30 deletions src/wigner.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,58 +12,114 @@ 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(
ρ::AbstractArray,
xvec::AbstractVector{T},
yvec::AbstractVector{T},
g::Real,
solver::WignerLaguerre,
method::WignerLaguerre,
) where {T<:BlasFloat}
g = convert(T, g)
X, Y = meshgrid(xvec, yvec)
A = g / 2 * (X + 1im * Y)
W = similar(A, T)
W .= 0

return _wigner_laguerre(ρ, A, W, g, solver)
return _wigner_laguerre(ρ, A, W, g, method)
end

function _wigner(
ρ::AbstractArray{T1},
xvec::AbstractVector{T},
yvec::AbstractVector{T},
g::Real,
solver::WignerClenshaw,
method::WignerClenshaw,
) where {T1<:BlasFloat,T<:BlasFloat}
g = convert(T, g)
M = size(ρ, 1)
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down
16 changes: 8 additions & 8 deletions test/core-test/wigner.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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