docs

Author: LorenzoFioroni

ext/QuantumToolboxCairoMakieExt.jl

@@ -3,6 +3,44 @@ module QuantumToolboxCairoMakieExt
 using QuantumToolbox
 using CairoMakie
 
+@doc raw"""
+    plot_wigner(
+        library::Val{:CairoMakie},
+        state::QuantumObject{<:AbstractArray{T},OpType};
+        xvec::Union{Nothing,AbstractVector} = nothing,
+        yvec::Union{Nothing,AbstractVector} = nothing,
+        g::Real = √2,
+        method::WignerSolver = WignerClenshaw(),
+        projection::Union{Val,Symbol} = Val(:two_dim),
+        location::Union{GridPosition,Nothing} = nothing,
+        colorbar::Bool = false,
+        kwargs...
+    )
+
+Plot the [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) of `state` using the [CairoMakie](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie) plotting library.
+
+Note that CairoMakie must first be imported before using this function.
+
+# Arguments
+- `library::Val{:CairoMakie}`: The plotting library to use.
+- `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. Default is # TODO
+- `yvec::AbstractVector`: The y-coordinates of the phase space grid. Default is # TODO
+- `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.
+- `projection::Union{Val,Symbol}`: Wheather to plot the Wigner function in 2D or 3D. It can be either `Val(:two_dim)` or `Val(:three_dim)`, with `Val(:two_dim)` as default.
+- `location::Union{GridPosition,Nothing}`: The location of the plot in the layout. If `nothing`, the plot is created in a new figure. Default is `nothing`.
+- `colorbar::Bool`: Whether to include a colorbar in the plot. Default is `false`.
+- `kwargs...`: Additional keyword arguments to pass to the plotting function. 
+
+# Returns
+- `fig`: The figure object.
+- `ax`: The axis object.
+- `hm`: Either the heatmap or surface object, depending on the projection.
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `Val(:two_dim)` and `Val(:three_dim)` instead of `:two_dim` and `:three_dim`, respectively. Also, specify the library as `Val(:CairoMakie)` See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
+"""
 function QuantumToolbox.plot_wigner(
     library::Val{:CairoMakie},
     state::QuantumObject{<:AbstractArray{T},OpType};
@@ -54,8 +92,8 @@ function _plot_wigner(
     wig = wigner(state, xvec, yvec; g = g, method = method)
     wlim = maximum(abs, wig)
 
-    kwargs = merge(Dict(:colormap => :RdBu, :colorrange => (-wlim, wlim)), kwargs)
-    hm = heatmap!(ax, xvec, yvec, wig; kwargs...)
+    kwargs = merge(Dict(:colormap => Reverse(:RdBu), :colorrange => (-wlim, wlim)), kwargs)
+    hm = heatmap!(ax, xvec, yvec, wig'; kwargs...)
 
     if colorbar
         Colorbar(lyt[1, 2], hm)
@@ -88,7 +126,7 @@ function _plot_wigner(
     wlim = maximum(abs, wig)
 
     kwargs = merge(Dict(:colormap => :RdBu, :colorrange => (-wlim, wlim)), kwargs)
-    surf = surface!(ax, xvec, yvec, wig; kwargs...)
+    surf = surface!(ax, xvec, yvec, wig'; kwargs...)
 
     if colorbar
         Colorbar(lyt[1, 2], surf)

src/visualization.jl

@@ -1,5 +1,21 @@
 export plot_wigner
 
+@doc raw"""
+    plot_wigner(
+        state::QuantumObject{<:AbstractArray{T},OpType}; 
+        library::Union{Val,Symbol}=Val(:CairoMakie), 
+        kwargs...
+    )
+
+Plot the [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) of `state` using the [`wigner`](@ref) function. 
+    
+The `library` keyword argument specifies the plotting library to use, defaulting to `CairoMakie`. Note that plotting libraries must first be imported before using them with this function.
+
+# Arguments
+- `state::QuantumObject{<:AbstractArray{T},OpType}`: The quantum state for which to plot the Wigner distribution.
+- `library::Union{Val,Symbol}`: The plotting library to use. Default is `Val(:CairoMakie)`.
+- `kwargs...`: Additional keyword arguments to pass to the plotting function. See the documentation for the specific plotting library for more information.
+"""
 plot_wigner(
     state::QuantumObject{<:AbstractArray{T},OpType};
     library::Union{Val,Symbol} = Val(:CairoMakie),