Make clenshaw method as default for wigner (#279)

* Make clenshaw method as default for wigner

* Minor changes

Author: albertomercurio

src/wigner.jl

@@ -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::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``.
+
+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::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.
+
+# 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::WignerSolver = WignerClenshaw(),
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    ρ = 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)
@@ -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

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