diff --git a/CHANGELOG.md b/CHANGELOG.md index 51f3946c5..a09a9ee93 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main) - Add support for `OperatorKet` state input for `mesolve` and `smesolve`. ([#423]) +- Introduce `plot_fock_distribution` to plot the population of a state (ket, bra, or density matrix) in its basis (assumed to be Fock basis). ([#428]) ## [v0.28.0] Release date: 2025-02-22 @@ -174,3 +175,4 @@ Release date: 2024-11-13 [#420]: https://github.com/qutip/QuantumToolbox.jl/issues/420 [#421]: https://github.com/qutip/QuantumToolbox.jl/issues/421 [#423]: https://github.com/qutip/QuantumToolbox.jl/issues/423 +[#428]: https://github.com/qutip/QuantumToolbox.jl/issues/428 diff --git a/docs/src/resources/api.md b/docs/src/resources/api.md index bb3e0b361..eb768caca 100644 --- a/docs/src/resources/api.md +++ b/docs/src/resources/api.md @@ -313,4 +313,5 @@ meshgrid ```@docs plot_wigner +plot_fock_distribution ``` diff --git a/docs/src/users_guide/extensions/cairomakie.md b/docs/src/users_guide/extensions/cairomakie.md index e847dd36c..b96c92d64 100644 --- a/docs/src/users_guide/extensions/cairomakie.md +++ b/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) | \ No newline at end of file +| [`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 | \ No newline at end of file diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl index 23be6452d..950859e3f 100644 --- a/ext/QuantumToolboxCairoMakieExt.jl +++ b/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. +- `bp`: The barplot object. + +!!! 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) + isapprox(tr(ρ), 1, atol = 1e-4) || (@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]) + + bp = 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, bp +end + raw""" _getFigAndLocation(location::Nothing) diff --git a/src/visualization.jl b/src/visualization.jl index bd21582a0..0e1fbd3fb 100644 --- a/src/visualization.jl +++ b/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.")) diff --git a/test/ext-test/cairomakie/cairomakie_ext.jl b/test/ext-test/cairomakie/cairomakie_ext.jl index 2571e02d1..0c73dc363 100644 --- a/test/ext-test/cairomakie/cairomakie_ext.jl +++ b/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