From 7ee6ee883316b93730d4d8e28972fdffb1da616d Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Sat, 9 Nov 2024 20:29:02 +0100
Subject: [PATCH 01/13] Add Wigner plotting via CairoMakie

---
 Project.toml                       |   3 +
 ext/QuantumToolboxCairoMakieExt.jl | 127 +++++++++++++++++++++++++++++
 src/QuantumToolbox.jl              |   1 +
 src/visualization.jl               |   5 ++
 4 files changed, 136 insertions(+)
 create mode 100644 ext/QuantumToolboxCairoMakieExt.jl
 create mode 100644 src/visualization.jl

diff --git a/Project.toml b/Project.toml
index 711c4828..1865fe92 100644
--- a/Project.toml
+++ b/Project.toml
@@ -28,14 +28,17 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 
 [extensions]
 QuantumToolboxCUDAExt = "CUDA"
+QuantumToolboxCairoMakieExt = "CairoMakie"
 
 [compat]
 Aqua = "0.8"
 ArrayInterface = "6, 7"
 CUDA = "5"
+CairoMakie = "0.12.15"
 DiffEqBase = "6"
 DiffEqCallbacks = "4.2.1 - 4"
 DiffEqNoiseProcess = "5"
diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
new file mode 100644
index 00000000..8a9a4b76
--- /dev/null
+++ b/ext/QuantumToolboxCairoMakieExt.jl
@@ -0,0 +1,127 @@
+module QuantumToolboxCairoMakieExt
+
+using QuantumToolbox
+using CairoMakie
+
+function QuantumToolbox.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::String = "2d",
+    fig::Union{Figure,Nothing} = nothing,
+    ax::Union{Axis,Nothing} = nothing,
+    colorbar::Bool = false,
+    kwargs...,
+) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    projection == "2d" || projection == "3d" || throw(ArgumentError("Unsupported projection: $projection"))
+
+    return _plot_wigner(
+        library,
+        state,
+        xvec,
+        yvec,
+        Val(Symbol(projection)),
+        g,
+        method,
+        fig,
+        ax,
+        colorbar;
+        kwargs...
+    )
+end
+
+function _plot_wigner(
+    ::Val{:CairoMakie},
+    state::QuantumObject{<:AbstractArray{T},OpType},
+    xvec::AbstractVector,
+    yvec::AbstractVector,
+    projection::Val{Symbol("2d")},
+    g::Real,
+    method::WignerSolver,
+    fig::Union{Figure,Nothing},
+    ax::Union{Axis,Nothing},
+    colorbar::Bool;
+    kwargs...,
+) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    fig, ax = _getFigAx(fig, ax)
+
+    gridPos = _gridPosFromAx(ax)
+    CairoMakie.delete!(ax)
+
+    lyt = GridLayout(gridPos)
+    ax = Axis(lyt[1, 1])
+
+    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...)
+
+    if colorbar
+        Colorbar(lyt[1, 2], hm)
+    end
+
+    ax.xlabel = L"\Re(\alpha)"
+    ax.ylabel = L"\Im(\alpha)"
+    return fig, ax, hm
+end
+
+function _plot_wigner(
+    ::Val{:CairoMakie},
+    state::QuantumObject{<:AbstractArray{T},OpType},
+    xvec::AbstractVector,
+    yvec::AbstractVector,
+    projection::Val{Symbol("3d")},
+    g::Real,
+    method::WignerSolver,
+    fig::Union{Figure,Nothing},
+    ax::Union{Axis,Nothing},
+    colorbar::Bool;
+    kwargs...,
+) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    fig, ax = _getFigAx(fig, ax)
+
+    gridPos = _gridPosFromAx(ax)
+    CairoMakie.delete!(ax)
+
+    lyt = GridLayout(gridPos)
+    ax = Axis3(lyt[1, 1], azimuth = 1.775pi, elevation = pi / 16, protrusions = (30, 90, 30, 30), viewmode = :stretch)
+
+    wig = wigner(state, xvec, yvec; g = g, method = method)
+    wlim = maximum(abs, wig)
+
+    kwargs = merge(Dict(:colormap => :RdBu, :colorrange => (-wlim, wlim)), kwargs)
+    surf = surface!(ax, xvec, yvec, wig; kwargs...)
+
+    if colorbar
+        Colorbar(lyt[1, 2], surf)
+    end
+
+    ax.xlabel = L"\Re(\alpha)"
+    ax.ylabel = L"\Im(\alpha)"
+    ax.zlabel = "Wigner function"
+    return fig, ax, surf
+end
+
+_getFigAx(fig::Figure, ax::Axis) = fig, ax
+_getFigAx(fig::Figure, ::Nothing) = fig, Axis(fig[1, 1])
+_getFigAx(::Nothing, ax::Axis) = _figFromChildren(ax), ax
+function _getFigAx(::Nothing, ::Nothing)
+    fig = Figure()
+    ax = Axis(fig[1, 1])
+    return fig, ax
+end
+
+_figFromChildren(children) = _figFromChildren(children.parent)
+_figFromChildren(fig::Figure) = fig
+
+function _gridPosFromAx(ax::Axis)
+    content = CairoMakie.Makie.GridLayoutBase.gridcontent(ax)
+    gl, sp, si = content.parent, content.span, content.side
+    return GridPosition(gl, sp, si)
+end
+
+end
\ No newline at end of file
diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
index 6b711223..ad387e26 100644
--- a/src/QuantumToolbox.jl
+++ b/src/QuantumToolbox.jl
@@ -114,6 +114,7 @@ include("arnoldi.jl")
 include("metrics.jl")
 include("negativity.jl")
 include("steadystate.jl")
+include("visualization.jl")
 
 # deprecated functions
 include("deprecated.jl")
diff --git a/src/visualization.jl b/src/visualization.jl
new file mode 100644
index 00000000..c43540a0
--- /dev/null
+++ b/src/visualization.jl
@@ -0,0 +1,5 @@
+export plot_wigner
+
+plot_wigner(library::Val{T}, args...; kwargs...) where {T} =
+    throw(ArgumentError("Unsupported visualization library: $(getVal(library))"))
+plot_wigner(library::Symbol, args...; kwargs...) = plot_wigner(Val(library), args...; kwargs...)
\ No newline at end of file

From e0ceaaee9abe16534d940b10b84093aa423ac46a Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Sun, 10 Nov 2024 23:08:41 +0100
Subject: [PATCH 02/13] CairoMakie default lib

---
 ext/QuantumToolboxCairoMakieExt.jl | 61 ++++++++++++------------------
 src/visualization.jl               |  9 +++--
 2 files changed, 31 insertions(+), 39 deletions(-)

diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
index 8a9a4b76..605f89fc 100644
--- a/ext/QuantumToolboxCairoMakieExt.jl
+++ b/ext/QuantumToolboxCairoMakieExt.jl
@@ -5,31 +5,31 @@ using CairoMakie
 
 function QuantumToolbox.plot_wigner(
     library::Val{:CairoMakie},
-    state::QuantumObject{<:AbstractArray{T},OpType},
+    state::QuantumObject{<:AbstractArray{T},OpType};
     xvec::Union{Nothing,AbstractVector} = nothing,
-    yvec::Union{Nothing,AbstractVector} = nothing;
+    yvec::Union{Nothing,AbstractVector} = nothing,
     g::Real = √2,
     method::WignerSolver = WignerClenshaw(),
-    projection::String = "2d",
-    fig::Union{Figure,Nothing} = nothing,
-    ax::Union{Axis,Nothing} = nothing,
+    projection::Union{Val,Symbol} = Val(:two_dim),
+    location::Union{GridPosition,Nothing} = nothing,
     colorbar::Bool = false,
     kwargs...,
 ) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
-    projection == "2d" || projection == "3d" || throw(ArgumentError("Unsupported projection: $projection"))
+    QuantumToolbox.getVal(projection) == :two_dim ||
+        QuantumToolbox.getVal(projection) == :three_dim ||
+        throw(ArgumentError("Unsupported projection: $projection"))
 
     return _plot_wigner(
         library,
         state,
         xvec,
         yvec,
-        Val(Symbol(projection)),
+        QuantumToolbox.makeVal(projection),
         g,
         method,
-        fig,
-        ax,
+        location,
         colorbar;
-        kwargs...
+        kwargs...,
     )
 end
 
@@ -38,20 +38,17 @@ function _plot_wigner(
     state::QuantumObject{<:AbstractArray{T},OpType},
     xvec::AbstractVector,
     yvec::AbstractVector,
-    projection::Val{Symbol("2d")},
+    projection::Val{:two_dim},
     g::Real,
     method::WignerSolver,
-    fig::Union{Figure,Nothing},
-    ax::Union{Axis,Nothing},
+    location::Union{GridPosition,Nothing},
     colorbar::Bool;
     kwargs...,
 ) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
-    fig, ax = _getFigAx(fig, ax)
+    fig, location = _getFigAndLocation(location)
 
-    gridPos = _gridPosFromAx(ax)
-    CairoMakie.delete!(ax)
+    lyt = GridLayout(location)
 
-    lyt = GridLayout(gridPos)
     ax = Axis(lyt[1, 1])
 
     wig = wigner(state, xvec, yvec; g = g, method = method)
@@ -74,20 +71,17 @@ function _plot_wigner(
     state::QuantumObject{<:AbstractArray{T},OpType},
     xvec::AbstractVector,
     yvec::AbstractVector,
-    projection::Val{Symbol("3d")},
+    projection::Val{:three_dim},
     g::Real,
     method::WignerSolver,
-    fig::Union{Figure,Nothing},
-    ax::Union{Axis,Nothing},
+    location::Union{GridPosition,Nothing},
     colorbar::Bool;
     kwargs...,
 ) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
-    fig, ax = _getFigAx(fig, ax)
+    fig, location = _getFigAndLocation(location)
 
-    gridPos = _gridPosFromAx(ax)
-    CairoMakie.delete!(ax)
+    lyt = GridLayout(location)
 
-    lyt = GridLayout(gridPos)
     ax = Axis3(lyt[1, 1], azimuth = 1.775pi, elevation = pi / 16, protrusions = (30, 90, 30, 30), viewmode = :stretch)
 
     wig = wigner(state, xvec, yvec; g = g, method = method)
@@ -106,22 +100,17 @@ function _plot_wigner(
     return fig, ax, surf
 end
 
-_getFigAx(fig::Figure, ax::Axis) = fig, ax
-_getFigAx(fig::Figure, ::Nothing) = fig, Axis(fig[1, 1])
-_getFigAx(::Nothing, ax::Axis) = _figFromChildren(ax), ax
-function _getFigAx(::Nothing, ::Nothing)
+function _getFigAndLocation(location::Nothing)
     fig = Figure()
-    ax = Axis(fig[1, 1])
-    return fig, ax
+    return fig, fig[1, 1]
+end
+function _getFigAndLocation(location::GridPosition)
+    fig = _figFromChildren(location.layout)
+    return fig, location
 end
 
 _figFromChildren(children) = _figFromChildren(children.parent)
 _figFromChildren(fig::Figure) = fig
-
-function _gridPosFromAx(ax::Axis)
-    content = CairoMakie.Makie.GridLayoutBase.gridcontent(ax)
-    gl, sp, si = content.parent, content.span, content.side
-    return GridPosition(gl, sp, si)
-end
+_figFromChildren(::Nothing) = throw(ArgumentError("No Figure has been found at the top of the layout hierarchy."))
 
 end
\ No newline at end of file
diff --git a/src/visualization.jl b/src/visualization.jl
index c43540a0..9f64c7e1 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -1,5 +1,8 @@
 export plot_wigner
 
-plot_wigner(library::Val{T}, args...; kwargs...) where {T} =
-    throw(ArgumentError("Unsupported visualization library: $(getVal(library))"))
-plot_wigner(library::Symbol, args...; kwargs...) = plot_wigner(Val(library), args...; kwargs...)
\ No newline at end of file
+plot_wigner(
+    state::QuantumObject{<:AbstractArray{T},OpType};
+    library::Union{Val,Symbol} = Val(:CairoMakie),
+    kwargs...,
+) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
+    plot_wigner(makeVal(library), state; kwargs...)

From af18daea5e19e5c96515032ffc52707283a62fff Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Sat, 16 Nov 2024 17:32:43 +0100
Subject: [PATCH 03/13] docs

---
 ext/QuantumToolboxCairoMakieExt.jl | 44 ++++++++++++++++++++++++++++--
 src/visualization.jl               | 16 +++++++++++
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
index 605f89fc..585da0e7 100644
--- a/ext/QuantumToolboxCairoMakieExt.jl
+++ b/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)
diff --git a/src/visualization.jl b/src/visualization.jl
index 9f64c7e1..1bd48bba 100644
--- a/src/visualization.jl
+++ b/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),

From 956cee1a65e0cfe9d67142822fb96184b40126d0 Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Mon, 9 Dec 2024 12:05:52 +0100
Subject: [PATCH 04/13] formatting

---
 ext/QuantumToolboxCairoMakieExt.jl | 70 +++++++++++++++++++++++++++---
 src/visualization.jl               |  3 ++
 2 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
index 585da0e7..784f2d57 100644
--- a/ext/QuantumToolboxCairoMakieExt.jl
+++ b/ext/QuantumToolboxCairoMakieExt.jl
@@ -24,8 +24,8 @@ 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
+- `xvec::AbstractVector`: The x-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points.
+- `yvec::AbstractVector`: The y-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points.
 - `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.
@@ -44,8 +44,8 @@ Note that CairoMakie must first be imported before using this function.
 function QuantumToolbox.plot_wigner(
     library::Val{:CairoMakie},
     state::QuantumObject{<:AbstractArray{T},OpType};
-    xvec::Union{Nothing,AbstractVector} = nothing,
-    yvec::Union{Nothing,AbstractVector} = nothing,
+    xvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200),
+    yvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200),
     g::Real = √2,
     method::WignerSolver = WignerClenshaw(),
     projection::Union{Val,Symbol} = Val(:two_dim),
@@ -138,17 +138,77 @@ function _plot_wigner(
     return fig, ax, surf
 end
 
+@doc raw"""
+    _getFigAndLocation(location::Nothing)
+    
+    Create a new figure and return it, together with the GridPosition object pointing to the first cell.
+
+    # Arguments
+    - `location::Nothing`
+
+    # Returns
+    - `fig`: The figure object.
+    - `location`: The GridPosition object pointing to the first cell.
+"""
 function _getFigAndLocation(location::Nothing)
     fig = Figure()
     return fig, fig[1, 1]
 end
+
+@doc raw"""
+    _getFigAndLocation(location::GridPosition)
+    
+    Compute which figure does the location belong to and return it, together with the location itself.
+
+    # Arguments
+    - `location::GridPosition`
+
+    # Returns
+    - `fig`: The figure object.
+    - `location`: The GridPosition object.
+"""
 function _getFigAndLocation(location::GridPosition)
     fig = _figFromChildren(location.layout)
     return fig, location
 end
 
+@doc raw"""
+    _figFromChildren(children::GridLayout)
+
+    Recursively find the figure object from the children layout.
+
+    # Arguments
+    - `children::GridLayout`
+
+    # Returns
+    - Union{Nothing, Figure, GridLayout}: The children's parent object.
+"""
 _figFromChildren(children) = _figFromChildren(children.parent)
+
+@doc raw"""
+    _figFromChildren(fig::Figure)
+
+    Return the figure object
+
+    # Arguments
+    - `fig::Figure`
+
+    # Returns
+    - `fig`: The figure object.
+"""
 _figFromChildren(fig::Figure) = fig
+
+@doc raw"""
+    _figFromChildren(::Nothing)
+
+    Throw an error if no figure has been found.
+
+    # Arguments
+    - `::Nothing`
+
+    # Throws
+    - `ArgumentError`: If no figure has been found.
+"""
 _figFromChildren(::Nothing) = throw(ArgumentError("No Figure has been found at the top of the layout hierarchy."))
 
-end
\ No newline at end of file
+end
diff --git a/src/visualization.jl b/src/visualization.jl
index 1bd48bba..02cc8f36 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -15,6 +15,9 @@ The `library` keyword argument specifies the plotting library to use, defaulting
 - `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.
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `Val(:CairoMakie)` instead of `:CairoMakie` as the plotting library. 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.
 """
 plot_wigner(
     state::QuantumObject{<:AbstractArray{T},OpType};

From 0224a76390e8d5849122ca1f3446a425dcc2358d Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Mon, 9 Dec 2024 12:16:30 +0100
Subject: [PATCH 05/13] Fix typo

---
 ext/QuantumToolboxCairoMakieExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
index 784f2d57..8d7255f5 100644
--- a/ext/QuantumToolboxCairoMakieExt.jl
+++ b/ext/QuantumToolboxCairoMakieExt.jl
@@ -28,7 +28,7 @@ Note that CairoMakie must first be imported before using this function.
 - `yvec::AbstractVector`: The y-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points.
 - `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.
+- `projection::Union{Val,Symbol}`: Whether 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. 

From e2761b1199012e9dd5cc2791fc6ab0dbecbc7159 Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Mon, 9 Dec 2024 18:07:33 +0100
Subject: [PATCH 06/13] Add error handling for unavailable plotting libraries
 in plot_wigner function

---
 src/visualization.jl | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/visualization.jl b/src/visualization.jl
index 02cc8f36..4eeacac2 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -25,3 +25,13 @@ plot_wigner(
     kwargs...,
 ) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
     plot_wigner(makeVal(library), state; kwargs...)
+
+plot_wigner(
+    ::Val{T},
+    state::QuantumObject{<:AbstractArray{T1},OpType};
+    kwargs...,
+) where {T,T1,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} = throw(
+    ArgumentError(
+        "The specified plotting library $T is not available. Try running `using $T` first.",
+    ),
+)
\ No newline at end of file

From 7abe4383e0848f588de3dd91f67ff7227c55b5c9 Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Mon, 9 Dec 2024 18:52:31 +0100
Subject: [PATCH 07/13] Docs

---
 docs/make.jl                       |  8 +++++++-
 docs/src/resources/api.md          |  6 ++++++
 ext/QuantumToolboxCairoMakieExt.jl | 26 +++++++++++++-------------
 src/visualization.jl               | 14 +++++++-------
 4 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/docs/make.jl b/docs/make.jl
index 9e281d9c..ecbd1bae 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -7,6 +7,9 @@ using DocumenterVitepress
 using DocumenterCitations
 using Changelog
 
+# Load of packages required to compile the extension documentation
+using CairoMakie
+
 DocMeta.setdocmeta!(QuantumToolbox, :DocTestSetup, :(using QuantumToolbox); recursive = true)
 
 # some options for `makedocs`
@@ -76,7 +79,10 @@ const PAGES = [
 ]
 
 makedocs(;
-    modules = [QuantumToolbox],
+    modules = [
+        QuantumToolbox, 
+        Base.get_extension(QuantumToolbox, :QuantumToolboxCairoMakieExt),
+        ],
     authors = "Alberto Mercurio and Yi-Te Huang",
     repo = Remotes.GitHub("qutip", "QuantumToolbox.jl"),
     sitename = "QuantumToolbox.jl",
diff --git a/docs/src/resources/api.md b/docs/src/resources/api.md
index d20fc680..c606e182 100644
--- a/docs/src/resources/api.md
+++ b/docs/src/resources/api.md
@@ -281,3 +281,9 @@ convert_unit
 row_major_reshape
 meshgrid
 ```
+
+## [Visualization](@id doc-API:Visualization)
+
+```@docs
+plot_wigner
+```
diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
index 8d7255f5..722ca150 100644
--- a/ext/QuantumToolboxCairoMakieExt.jl
+++ b/ext/QuantumToolboxCairoMakieExt.jl
@@ -6,7 +6,7 @@ using CairoMakie
 @doc raw"""
     plot_wigner(
         library::Val{:CairoMakie},
-        state::QuantumObject{<:AbstractArray{T},OpType};
+        state::QuantumObject{DT,OpType};
         xvec::Union{Nothing,AbstractVector} = nothing,
         yvec::Union{Nothing,AbstractVector} = nothing,
         g::Real = √2,
@@ -15,7 +15,7 @@ using CairoMakie
         location::Union{GridPosition,Nothing} = nothing,
         colorbar::Bool = false,
         kwargs...
-    )
+    ) where {DT,OpType}
 
 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.
 
@@ -43,7 +43,7 @@ Note that CairoMakie must first be imported before using this function.
 """
 function QuantumToolbox.plot_wigner(
     library::Val{:CairoMakie},
-    state::QuantumObject{<:AbstractArray{T},OpType};
+    state::QuantumObject{DT,OpType};
     xvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200),
     yvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200),
     g::Real = √2,
@@ -52,7 +52,7 @@ function QuantumToolbox.plot_wigner(
     location::Union{GridPosition,Nothing} = nothing,
     colorbar::Bool = false,
     kwargs...,
-) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
     QuantumToolbox.getVal(projection) == :two_dim ||
         QuantumToolbox.getVal(projection) == :three_dim ||
         throw(ArgumentError("Unsupported projection: $projection"))
@@ -73,7 +73,7 @@ end
 
 function _plot_wigner(
     ::Val{:CairoMakie},
-    state::QuantumObject{<:AbstractArray{T},OpType},
+    state::QuantumObject{DT,OpType},
     xvec::AbstractVector,
     yvec::AbstractVector,
     projection::Val{:two_dim},
@@ -82,7 +82,7 @@ function _plot_wigner(
     location::Union{GridPosition,Nothing},
     colorbar::Bool;
     kwargs...,
-) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
     fig, location = _getFigAndLocation(location)
 
     lyt = GridLayout(location)
@@ -106,7 +106,7 @@ end
 
 function _plot_wigner(
     ::Val{:CairoMakie},
-    state::QuantumObject{<:AbstractArray{T},OpType},
+    state::QuantumObject{DT,OpType},
     xvec::AbstractVector,
     yvec::AbstractVector,
     projection::Val{:three_dim},
@@ -115,7 +115,7 @@ function _plot_wigner(
     location::Union{GridPosition,Nothing},
     colorbar::Bool;
     kwargs...,
-) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
     fig, location = _getFigAndLocation(location)
 
     lyt = GridLayout(location)
@@ -138,7 +138,7 @@ function _plot_wigner(
     return fig, ax, surf
 end
 
-@doc raw"""
+raw"""
     _getFigAndLocation(location::Nothing)
     
     Create a new figure and return it, together with the GridPosition object pointing to the first cell.
@@ -155,7 +155,7 @@ function _getFigAndLocation(location::Nothing)
     return fig, fig[1, 1]
 end
 
-@doc raw"""
+raw"""
     _getFigAndLocation(location::GridPosition)
     
     Compute which figure does the location belong to and return it, together with the location itself.
@@ -172,7 +172,7 @@ function _getFigAndLocation(location::GridPosition)
     return fig, location
 end
 
-@doc raw"""
+raw"""
     _figFromChildren(children::GridLayout)
 
     Recursively find the figure object from the children layout.
@@ -185,7 +185,7 @@ end
 """
 _figFromChildren(children) = _figFromChildren(children.parent)
 
-@doc raw"""
+raw"""
     _figFromChildren(fig::Figure)
 
     Return the figure object
@@ -198,7 +198,7 @@ _figFromChildren(children) = _figFromChildren(children.parent)
 """
 _figFromChildren(fig::Figure) = fig
 
-@doc raw"""
+raw"""
     _figFromChildren(::Nothing)
 
     Throw an error if no figure has been found.
diff --git a/src/visualization.jl b/src/visualization.jl
index 4eeacac2..2f6c1205 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -2,17 +2,17 @@ export plot_wigner
 
 @doc raw"""
     plot_wigner(
-        state::QuantumObject{<:AbstractArray{T},OpType}; 
+        state::QuantumObject{DT,OpType}; 
         library::Union{Val,Symbol}=Val(:CairoMakie), 
         kwargs...
-    )
+    ) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}
 
 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.
+- `state::QuantumObject{DT,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.
 
@@ -20,17 +20,17 @@ The `library` keyword argument specifies the plotting library to use, defaulting
     If you want to keep type stability, it is recommended to use `Val(:CairoMakie)` instead of `:CairoMakie` as the plotting library. 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.
 """
 plot_wigner(
-    state::QuantumObject{<:AbstractArray{T},OpType};
+    state::QuantumObject{DT,OpType};
     library::Union{Val,Symbol} = Val(:CairoMakie),
     kwargs...,
-) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
     plot_wigner(makeVal(library), state; kwargs...)
 
 plot_wigner(
     ::Val{T},
-    state::QuantumObject{<:AbstractArray{T1},OpType};
+    state::QuantumObject{DT,OpType};
     kwargs...,
-) where {T,T1,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} = throw(
+) where {T,DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} = throw(
     ArgumentError(
         "The specified plotting library $T is not available. Try running `using $T` first.",
     ),

From 82553ca8203619a7d42a5f5ebc43d29b5c3cd378 Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Mon, 9 Dec 2024 19:36:20 +0100
Subject: [PATCH 08/13] Update CairoMakie dependency and add CairoMakie
 extension tests

---
 Project.toml                               |  5 +-
 test/ext-test/cairomakie/cairomakie_ext.jl | 63 ++++++++++++++++++++++
 test/runtests.jl                           |  8 +++
 3 files changed, 74 insertions(+), 2 deletions(-)
 create mode 100644 test/ext-test/cairomakie/cairomakie_ext.jl

diff --git a/Project.toml b/Project.toml
index 1865fe92..9b30f61b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -38,7 +38,7 @@ QuantumToolboxCairoMakieExt = "CairoMakie"
 Aqua = "0.8"
 ArrayInterface = "6, 7"
 CUDA = "5"
-CairoMakie = "0.12.15"
+CairoMakie = "0.12"
 DiffEqBase = "6"
 DiffEqCallbacks = "4.2.1 - 4"
 DiffEqNoiseProcess = "5"
@@ -65,8 +65,9 @@ julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "JET", "Test"]
+test = ["Aqua", "CairoMakie", "JET", "Test"]
diff --git a/test/ext-test/cairomakie/cairomakie_ext.jl b/test/ext-test/cairomakie/cairomakie_ext.jl
new file mode 100644
index 00000000..978aae09
--- /dev/null
+++ b/test/ext-test/cairomakie/cairomakie_ext.jl
@@ -0,0 +1,63 @@
+@testset "CairoMakie Extension" verbose = true begin
+    ψ = normalize(coherent(50, 5.0) + coherent(50, -5.0))
+    xvec = yvec = -15.0:0.1:15.0
+    wig = wigner(ψ, xvec, yvec)'
+    
+    @test_throws ArgumentError plot_wigner(ψ; library = :CairoMakie, xvec = xvec, yvec = yvec)
+
+    using CairoMakie
+
+    fig, ax, hm = plot_wigner(
+        ψ;
+        library = Val(:CairoMakie),
+        xvec = xvec,
+        yvec = yvec,
+        projection = Val(:two_dim),
+        colorbar = true,
+    )
+    @test fig isa Figure
+    @test ax isa Axis
+    @test hm isa Heatmap
+    @test all(isapprox.(hm[3].val, wig, atol = 1e-6))
+
+    fig, ax, surf = plot_wigner(
+        ψ;
+        library = Val(:CairoMakie),
+        xvec = xvec,
+        yvec = yvec,
+        projection = Val(:three_dim),
+        colorbar = true,
+    )
+    @test fig isa Figure
+    @test ax isa Axis3
+    @test surf isa Surface
+    @test all(isapprox.(surf[3].val, wig, atol = 1e-6))
+
+    fig = Figure()
+    pos = fig[2, 3]
+    fig1, ax, hm = plot_wigner(
+        ψ;
+        library = Val(:CairoMakie),
+        xvec = xvec,
+        yvec = yvec,
+        projection = Val(:two_dim),
+        colorbar = true,
+        location = pos,
+    )
+    @test fig1 === fig
+    @test fig[2, 3].layout.content[1].content[1, 1].layout.content[1].content === ax
+
+    fig = Figure()
+    pos = fig[2, 3]
+    fig1, ax, surf = plot_wigner(
+        ψ;
+        library = Val(:CairoMakie),
+        xvec = xvec,
+        yvec = yvec,
+        projection = Val(:three_dim),
+        colorbar = true,
+        location = pos,
+    )
+    @test fig1 === fig
+    @test fig[2, 3].layout.content[1].content[1, 1].layout.content[1].content === ax
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 20f72b05..cc52f355 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -26,6 +26,10 @@ core_tests = [
     "wigner.jl",
 ]
 
+ext_tests = [
+    joinpath("cairomakie", "cairomakie_ext.jl"),
+]
+
 if (GROUP == "All") || (GROUP == "Code-Quality")
     using QuantumToolbox
     using Aqua, JET
@@ -44,6 +48,10 @@ if (GROUP == "All") || (GROUP == "Core")
     for test in core_tests
         include(joinpath(testdir, "core-test", test))
     end
+
+    for test in ext_tests
+        include(joinpath(testdir, "ext-test", test))
+    end
 end
 
 if (GROUP == "CUDA_Ext")# || (GROUP == "All")

From 90906efdab0a1cf6095f98815f6bbfc733596fa8 Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Mon, 9 Dec 2024 20:01:27 +0100
Subject: [PATCH 09/13] Update documentation to reflect changes in Wigner
 function plotting and computation

---
 docs/src/tutorials/logo.md | 24 ++++++------------------
 1 file changed, 6 insertions(+), 18 deletions(-)

diff --git a/docs/src/tutorials/logo.md b/docs/src/tutorials/logo.md
index b68386bc..3084378a 100644
--- a/docs/src/tutorials/logo.md
+++ b/docs/src/tutorials/logo.md
@@ -67,25 +67,16 @@ Next, we construct the triangular cat state as a normalized superposition of thr
 normalize!(ψ)
 ```
 
-### Defining the Grid and calculating the Wigner function
+### Defining the Grid and plotting the Wigner function
 
-We define the grid for the Wigner function and calculate it using the [`wigner`](@ref) function. We shift the grid in the imaginary direction to ensure that the Wigner function is centered around the origin of the figure. The [`wigner`](@ref) function also supports the `g` scaling factor, which we put here equal to ``2``.
+We define the grid for the Wigner function and plot it using the [`plot_wigner`](@ref) function. This, internally calls the [`wigner`](@ref) function for the computation. We shift the grid in the imaginary direction to ensure that the Wigner function is centered around the origin of the figure. The [`wigner`](@ref) function also supports the `g` scaling factor, which we put here equal to ``2``.
 
 ```@example logo
 xvec = range(-ρ, ρ, 500) .* 1.5
 yvec = xvec .+ (abs(imag(α1)) - abs(imag(α2))) / 2
 
-wig = wigner(ψ, xvec, yvec, g = 2)
-```
-
-### Plotting the Wigner function
-
-Finally, we plot the Wigner function using the `heatmap` function from the `CairoMakie` package.
-
-```@example logo
 fig = Figure(size = (250, 250), figure_padding = 0)
-ax = Axis(fig[1, 1])
-heatmap!(ax, xvec, yvec, wig', colormap = :RdBu, interpolate = true, rasterize = 1)
+fig, ax, hm = plot_wigner(ψ, library = Val(:CairoMakie), xvec = xvec, yvec = yvec, g = 2, location = fig[1,1])
 hidespines!(ax)
 hidexdecorations!(ax)
 hideydecorations!(ax)
@@ -118,12 +109,8 @@ nothing # hide
 And the Wigner function becomes more uniform:
 
 ```@example logo
-wig = wigner(sol.states[end], xvec, yvec, g = 2)
-
 fig = Figure(size = (250, 250), figure_padding = 0)
-ax = Axis(fig[1, 1])
-
-img_wig = heatmap!(ax, xvec, yvec, wig', colormap = :RdBu, interpolate = true, rasterize = 1)
+fig, ax, hm = plot_wigner(sol.states[end], library = Val(:CairoMakie), xvec = xvec, yvec = yvec, g = 2, location = fig[1,1])
 hidespines!(ax)
 hidexdecorations!(ax)
 hideydecorations!(ax)
@@ -135,7 +122,7 @@ At this stage, we have finished to use the `QuantumToolbox` package. From now on
 
 ### Custom Colormap
 
-We define a custom colormap that changes depending on the Wigner function and spatial coordinates. Indeed, we want the three different colormaps, in the regions corresponding to the three coherent states, to match the colors of the Julia logo. We also want the colormap change to be smooth, so we use a Gaussian function to blend the colors. We introduce also a Wigner function dependent transparency to make the logo more appealing.
+We define a custom colormap that changes depending on the Wigner function and spatial coordinates. Indeed, we want the three different colormaps, in the regions corresponding to the three coherent states, to match the colors of the Julia logo. We also want the colormap change to be smooth, so we use a Gaussian function to blend the colors. We introduce also a Wigner function dependent transparency to make the logo more appealing. In order to do so, we are going to need the value of the wigner function at each point of the grid, rather than its plot. We will thus call the [`wigner`](@ref) function directly.
 
 ```@example logo
 function set_color_julia(x, y, wig::T, α1, α2, α3, cmap1, cmap2, cmap3, δ) where {T}
@@ -156,6 +143,7 @@ function set_color_julia(x, y, wig::T, α1, α2, α3, cmap1, cmap2, cmap3, δ) w
     return RGBAf(c_tot.r, c_tot.g, c_tot.b, alpha)
 end
 
+wig = wigner(sol.states[end], xvec, yvec, g = 2)
 X, Y = meshgrid(xvec, yvec)
 δ = 1.25 # Smoothing parameter for the Gaussian functions
 ```

From 0f475a1d514b8da426c69e43def628424bb90c6d Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Mon, 9 Dec 2024 20:04:23 +0100
Subject: [PATCH 10/13] formatting

---
 src/visualization.jl                       | 7 ++-----
 test/ext-test/cairomakie/cairomakie_ext.jl | 2 +-
 test/runtests.jl                           | 4 +---
 3 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/src/visualization.jl b/src/visualization.jl
index 2f6c1205..bf81401a 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -30,8 +30,5 @@ plot_wigner(
     ::Val{T},
     state::QuantumObject{DT,OpType};
     kwargs...,
-) where {T,DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} = throw(
-    ArgumentError(
-        "The specified plotting library $T is not available. Try running `using $T` first.",
-    ),
-)
\ No newline at end of file
+) where {T,DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
+    throw(ArgumentError("The specified plotting library $T is not available. Try running `using $T` first."))
diff --git a/test/ext-test/cairomakie/cairomakie_ext.jl b/test/ext-test/cairomakie/cairomakie_ext.jl
index 978aae09..36dfee55 100644
--- a/test/ext-test/cairomakie/cairomakie_ext.jl
+++ b/test/ext-test/cairomakie/cairomakie_ext.jl
@@ -2,7 +2,7 @@
     ψ = normalize(coherent(50, 5.0) + coherent(50, -5.0))
     xvec = yvec = -15.0:0.1:15.0
     wig = wigner(ψ, xvec, yvec)'
-    
+
     @test_throws ArgumentError plot_wigner(ψ; library = :CairoMakie, xvec = xvec, yvec = yvec)
 
     using CairoMakie
diff --git a/test/runtests.jl b/test/runtests.jl
index cc52f355..74ad194e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -26,9 +26,7 @@ core_tests = [
     "wigner.jl",
 ]
 
-ext_tests = [
-    joinpath("cairomakie", "cairomakie_ext.jl"),
-]
+ext_tests = [joinpath("cairomakie", "cairomakie_ext.jl")]
 
 if (GROUP == "All") || (GROUP == "Code-Quality")
     using QuantumToolbox

From e24f00329e67a4ce2be422df97c653d8f846ca35 Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <alberto.mercurio96@gmail.com>
Date: Mon, 9 Dec 2024 23:09:43 +0100
Subject: [PATCH 11/13] Update changelog

---
 CHANGELOG.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4b61c2ca..c2451997 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
 
+- Introduce `plot_wigner` function for easy plotting of Wigner functions. ([#86], [#292])
 
 
 ## [v0.23.1]
@@ -44,7 +45,9 @@ Release date: 2024-11-13
 [v0.22.0]: https://github.com/qutip/QuantumToolbox.jl/releases/tag/v0.22.0
 [v0.23.0]: https://github.com/qutip/QuantumToolbox.jl/releases/tag/v0.23.0
 [v0.23.1]: https://github.com/qutip/QuantumToolbox.jl/releases/tag/v0.23.1
+[#86]: https://github.com/qutip/QuantumToolbox.jl/issues/86
 [#139]: https://github.com/qutip/QuantumToolbox.jl/issues/139
+[#292]: https://github.com/qutip/QuantumToolbox.jl/issues/292
 [#306]: https://github.com/qutip/QuantumToolbox.jl/issues/306
 [#309]: https://github.com/qutip/QuantumToolbox.jl/issues/309
 [#311]: https://github.com/qutip/QuantumToolbox.jl/issues/311

From f8f335df45037878ce812304b44585218c2916b9 Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Tue, 10 Dec 2024 09:48:53 +0100
Subject: [PATCH 12/13] Improve docs

---
 docs/src/tutorials/logo.md         | 4 ++--
 ext/QuantumToolboxCairoMakieExt.jl | 9 +++++----
 src/visualization.jl               | 5 ++++-
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/docs/src/tutorials/logo.md b/docs/src/tutorials/logo.md
index 3084378a..12be5156 100644
--- a/docs/src/tutorials/logo.md
+++ b/docs/src/tutorials/logo.md
@@ -76,7 +76,7 @@ xvec = range(-ρ, ρ, 500) .* 1.5
 yvec = xvec .+ (abs(imag(α1)) - abs(imag(α2))) / 2
 
 fig = Figure(size = (250, 250), figure_padding = 0)
-fig, ax, hm = plot_wigner(ψ, library = Val(:CairoMakie), xvec = xvec, yvec = yvec, g = 2, location = fig[1,1])
+fig, ax, hm = plot_wigner(ψ, xvec = xvec, yvec = yvec, g = 2, library = Val(:CairoMakie), location = fig[1,1])
 hidespines!(ax)
 hidexdecorations!(ax)
 hideydecorations!(ax)
@@ -110,7 +110,7 @@ And the Wigner function becomes more uniform:
 
 ```@example logo
 fig = Figure(size = (250, 250), figure_padding = 0)
-fig, ax, hm = plot_wigner(sol.states[end], library = Val(:CairoMakie), xvec = xvec, yvec = yvec, g = 2, location = fig[1,1])
+fig, ax, hm = plot_wigner(sol.states[end], xvec = xvec, yvec = yvec, g = 2, library = Val(:CairoMakie), location = fig[1,1])
 hidespines!(ax)
 hidexdecorations!(ax)
 hideydecorations!(ax)
diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
index 722ca150..9e86e761 100644
--- a/ext/QuantumToolboxCairoMakieExt.jl
+++ b/ext/QuantumToolboxCairoMakieExt.jl
@@ -17,13 +17,11 @@ using CairoMakie
         kwargs...
     ) where {DT,OpType}
 
-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.
+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.
 
 # 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).
+- `state::QuantumObject`: The quantum state for which the Wigner function is calculated. It can be either a [`Ket`](@ref), [`Bra`](@ref), or [`Operator`](@ref).
 - `xvec::AbstractVector`: The x-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points.
 - `yvec::AbstractVector`: The y-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points.
 - `g::Real`: The scaling factor related to the value of ``\hbar`` in the commutation relation ``[x, y] = i \hbar`` via ``\hbar=2/g^2``.
@@ -38,6 +36,9 @@ Note that CairoMakie must first be imported before using this function.
 - `ax`: The axis object.
 - `hm`: Either the heatmap or surface object, depending on the projection.
 
+!!! note "Import library first"
+    [`CairoMakie`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie) must first be imported before using this function.
+
 !!! 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.
 """
diff --git a/src/visualization.jl b/src/visualization.jl
index bf81401a..92ca2f67 100644
--- a/src/visualization.jl
+++ b/src/visualization.jl
@@ -9,13 +9,16 @@ export plot_wigner
 
 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.
+The `library` keyword argument specifies the plotting library to use, defaulting to [`CairoMakie`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie). 
 
 # Arguments
 - `state::QuantumObject{DT,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.
 
+!!! note "Import library first"
+    The plotting libraries must first be imported before using them with this function.
+
 !!! warning "Beware of type-stability!"
     If you want to keep type stability, it is recommended to use `Val(:CairoMakie)` instead of `:CairoMakie` as the plotting library. 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.
 """

From c10ac2b5605f1b0034f3b456150e68c78858b0c6 Mon Sep 17 00:00:00 2001
From: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Date: Wed, 11 Dec 2024 09:03:25 +0100
Subject: [PATCH 13/13] Update CairoMakie imports and change axis labels

---
 ext/QuantumToolboxCairoMakieExt.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
index 9e86e761..16de321c 100644
--- a/ext/QuantumToolboxCairoMakieExt.jl
+++ b/ext/QuantumToolboxCairoMakieExt.jl
@@ -1,7 +1,7 @@
 module QuantumToolboxCairoMakieExt
 
 using QuantumToolbox
-using CairoMakie
+using CairoMakie: Axis, Axis3, Colorbar, Figure, GridLayout, heatmap!, surface!, GridPosition, @L_str, Reverse
 
 @doc raw"""
     plot_wigner(
@@ -100,8 +100,8 @@ function _plot_wigner(
         Colorbar(lyt[1, 2], hm)
     end
 
-    ax.xlabel = L"\Re(\alpha)"
-    ax.ylabel = L"\Im(\alpha)"
+    ax.xlabel = L"\textrm{Re}(\alpha)"
+    ax.ylabel = L"\textrm{Im}(\alpha)"
     return fig, ax, hm
 end
 
@@ -133,8 +133,8 @@ function _plot_wigner(
         Colorbar(lyt[1, 2], surf)
     end
 
-    ax.xlabel = L"\Re(\alpha)"
-    ax.ylabel = L"\Im(\alpha)"
+    ax.xlabel = L"\textrm{Re}(\alpha)"
+    ax.ylabel = L"\textrm{Im}(\alpha)"
     ax.zlabel = "Wigner function"
     return fig, ax, surf
 end