add plot_fock_distribution to CairoMakie extension

Author: TendonFFF

docs/src/users_guide/extensions/cairomakie.md

@@ -18,4 +18,5 @@ The supported plotting functions are listed as follows:
 
 | **Plotting Function** | **Description** |
 |:----------------------|:----------------|
-| [`plot_wigner`](@ref) | [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) |
+| [`plot_wigner`](@ref) | [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) |
+| [`plot_fock_distribution`](@ref) | [Fock state](https://en.wikipedia.org/wiki/Fock_state) distribution |

ext/QuantumToolboxCairoMakieExt.jl

@@ -1,7 +1,8 @@
 module QuantumToolboxCairoMakieExt
 
 using QuantumToolbox
-using CairoMakie: Axis, Axis3, Colorbar, Figure, GridLayout, heatmap!, surface!, GridPosition, @L_str, Reverse
+using CairoMakie:
+    Axis, Axis3, Colorbar, Figure, GridLayout, heatmap!, surface!, barplot!, GridPosition, @L_str, Reverse, ylims!
 
 @doc raw"""
     plot_wigner(
@@ -139,6 +140,84 @@ function _plot_wigner(
     return fig, ax, surf
 end
 
+@doc raw"""
+    plot_fock_distribution(
+        library::Val{:CairoMakie},
+        ρ::QuantumObject{SType};
+        fock_numbers::Union{Nothing, AbstractVector} = nothing,
+        unit_y_range::Bool = true,
+        location::Union{GridPosition,Nothing} = nothing,
+        kwargs...
+    ) where {SType<:Union{KetQuantumObject,OperatorQuantumObject}}
+
+Plot the [Fock state](https://en.wikipedia.org/wiki/Fock_state) distribution of `ρ`. 
+
+# Arguments
+- `library::Val{:CairoMakie}`: The plotting library to use.
+- `ρ::QuantumObject`: The quantum state for which the Fock state distribution is to be plotted. It can be either a [`Ket`](@ref), [`Bra`](@ref), or [`Operator`](@ref).
+- `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`.
+- `fock_numbers::Union{Nothing, AbstractVector}`: list of x ticklabels to represent fock numbers, default is `nothing`.
+- `unit_y_range::Bool`: Set y-axis limits [0, 1] or not, default is `true`.
+- `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.
+
+!!! 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.
+"""
+function QuantumToolbox.plot_fock_distribution(
+    library::Val{:CairoMakie},
+    ρ::QuantumObject{SType};
+    fock_numbers::Union{Nothing,AbstractVector} = nothing,
+    unit_y_range::Bool = true,
+    location::Union{GridPosition,Nothing} = nothing,
+    kwargs...,
+) where {SType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    return _plot_fock_distribution(
+        library,
+        ρ;
+        fock_numbers = fock_numbers,
+        unit_y_range = unit_y_range,
+        location = location,
+        kwargs...,
+    )
+end
+
+function _plot_fock_distribution(
+    ::Val{:CairoMakie},
+    ρ::QuantumObject{SType};
+    fock_numbers::Union{Nothing,AbstractVector} = nothing,
+    unit_y_range::Bool = true,
+    location::Union{GridPosition,Nothing} = nothing,
+    kwargs...,
+) where {SType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+    ρ = ket2dm(ρ)
+    D = prod(ρ.dims)
+    (real(tr(ρ)) > 1) && (@warn "The input ρ should be normalized.")
+
+    xvec = 0:(D-1)
+    isnothing(fock_numbers) && (fock_numbers = string.(collect(xvec)))
+
+    fig, location = _getFigAndLocation(location)
+    lyt = GridLayout(location)
+    ax = Axis(lyt[1, 1])
+
+    barplot!(ax, xvec, real(diag(ρ)); kwargs...)
+
+    ax.xticks = (xvec, fock_numbers)
+    ax.xlabel = "Fock number"
+    ax.ylabel = "Occupation probability"
+    unit_y_range && ylims!(ax, 0, 1)
+
+    return fig, ax
+end
+
 raw"""
     _getFigAndLocation(location::Nothing)
     

src/visualization.jl

@@ -1,4 +1,4 @@
-export plot_wigner
+export plot_wigner, plot_fock_distribution
 
 @doc raw"""
     plot_wigner(
@@ -35,3 +35,39 @@ plot_wigner(
     kwargs...,
 ) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
     throw(ArgumentError("The specified plotting library $T is not available. Try running `using $T` first."))
+
+@doc raw"""
+    plot_fock_distribution(
+        ρ::QuantumObject{SType};
+        library::Union{Val, Symbol} = Val(:CairoMakie),
+        kwargs...
+    ) where {SType<:Union{KetQuantumObject,OperatorQuantumObject}}
+
+Plot the [Fock state](https://en.wikipedia.org/wiki/Fock_state) distribution of `ρ`. 
+    
+The `library` keyword argument specifies the plotting library to use, defaulting to [`CairoMakie`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie). 
+
+# Arguments
+- `ρ::QuantumObject`: The quantum state for which to plot the Fock state 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.
+"""
+plot_fock_distribution(
+    ρ::QuantumObject{SType};
+    library::Union{Val,Symbol} = Val(:CairoMakie),
+    kwargs...,
+) where {SType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
+    plot_fock_distribution(makeVal(library), ρ; kwargs...)
+
+plot_fock_distribution(
+    ::Val{T},
+    ρ::QuantumObject{SType};
+    kwargs...,
+) where {T,SType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} =
+    throw(ArgumentError("The specified plotting library $T is not available. Try running `using $T` first."))

test/ext-test/cairomakie/cairomakie_ext.jl

@@ -5,6 +5,8 @@
 
     @test_throws ArgumentError plot_wigner(ψ; library = :CairoMakie, xvec = xvec, yvec = yvec)
 
+    @test_throws ArgumentError plot_fock_distribution(ψ; library = :CairoMakie)
+
     using CairoMakie
 
     fig, ax, hm = plot_wigner(
@@ -60,4 +62,14 @@
     )
     @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 = plot_fock_distribution(ψ; library = Val(:CairoMakie), 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 = @test_logs (:warn,) plot_fock_distribution(ψ * 2; library = Val(:CairoMakie), location = pos)
 end