Add Wigner plotting via CairoMakie

Author: LorenzoFioroni

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 = "2 - 3.1, 3.8, 4"
 DiffEqNoiseProcess = "5"

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

src/QuantumToolbox.jl

@@ -104,6 +104,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,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...)