Add Wigner plotting via CairoMakie

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

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"
 DiffEqBase = "6"
 DiffEqCallbacks = "4.2.1 - 4"
 DiffEqNoiseProcess = "5"
@@ -62,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"]

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",

docs/src/resources/api.md

@@ -281,3 +281,9 @@ convert_unit
 row_major_reshape
 meshgrid
 ```
+
+## [Visualization](@id doc-API:Visualization)
+
+```@docs
+plot_wigner
+```

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
 ```

ext/QuantumToolboxCairoMakieExt.jl

@@ -0,0 +1,214 @@
+module QuantumToolboxCairoMakieExt
+
+using QuantumToolbox
+using CairoMakie
+
+@doc raw"""
+    plot_wigner(
+        library::Val{:CairoMakie},
+        state::QuantumObject{DT,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...
+    ) 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.
+
+# 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. 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}`: 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. 
+
+# 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{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,
+    method::WignerSolver = WignerClenshaw(),
+    projection::Union{Val,Symbol} = Val(:two_dim),
+    location::Union{GridPosition,Nothing} = nothing,
+    colorbar::Bool = false,
+    kwargs...,
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    QuantumToolbox.getVal(projection) == :two_dim ||
+        QuantumToolbox.getVal(projection) == :three_dim ||
+        throw(ArgumentError("Unsupported projection: $projection"))
+
+    return _plot_wigner(
+        library,
+        state,
+        xvec,
+        yvec,
+        QuantumToolbox.makeVal(projection),
+        g,
+        method,
+        location,
+        colorbar;
+        kwargs...,
+    )
+end
+
+function _plot_wigner(
+    ::Val{:CairoMakie},
+    state::QuantumObject{DT,OpType},
+    xvec::AbstractVector,
+    yvec::AbstractVector,
+    projection::Val{:two_dim},
+    g::Real,
+    method::WignerSolver,
+    location::Union{GridPosition,Nothing},
+    colorbar::Bool;
+    kwargs...,
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    fig, location = _getFigAndLocation(location)
+
+    lyt = GridLayout(location)
+
+    ax = Axis(lyt[1, 1])
+
+    wig = wigner(state, xvec, yvec; g = g, method = method)
+    wlim = maximum(abs, wig)
+
+    kwargs = merge(Dict(:colormap => Reverse(: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{DT,OpType},
+    xvec::AbstractVector,
+    yvec::AbstractVector,
+    projection::Val{:three_dim},
+    g::Real,
+    method::WignerSolver,
+    location::Union{GridPosition,Nothing},
+    colorbar::Bool;
+    kwargs...,
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    fig, location = _getFigAndLocation(location)
+
+    lyt = GridLayout(location)
+
+    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
+
+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
+
+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
+
+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)
+
+raw"""
+    _figFromChildren(fig::Figure)
+
+    Return the figure object
+
+    # Arguments
+    - `fig::Figure`
+
+    # Returns
+    - `fig`: The figure object.
+"""
+_figFromChildren(fig::Figure) = fig
+
+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

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")

src/visualization.jl

@@ -0,0 +1,34 @@
+export plot_wigner
+
+@doc raw"""
+    plot_wigner(
+        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{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.
+
+!!! 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{DT,OpType};
+    library::Union{Val,Symbol} = Val(:CairoMakie),
+    kwargs...,
+) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
+    plot_wigner(makeVal(library), state; kwargs...)
+
+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."))