Add Fock distribution bar chart visualization

Project.toml

@@ -17,6 +17,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
+WeakDepHelpers = "7869a13a-7328-4bcf-a489-0f4bb64497c7"
 
 [weakdeps]
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
@@ -46,4 +47,5 @@ SparseArrays = "1"
 SpecialFunctions = "2.1.4"
 Strided = "1, 2"
 UnsafeArrays = "1"
+WeakDepHelpers = "0.1"
 julia = "1.10"

docs/src/api.md

@@ -653,6 +653,9 @@ DenseChiMatrix
 blochsphereplot
 blochsphereplot!
 blochsphereplot_axis
+fockdistributionplot
+fockdistributionplot!
+fockdistributionplot_axis
 ```
 
 ## [Printing](@id API: Printing)

ext/QuantumOpticsBaseMakieExt.jl

@@ -1,11 +1,12 @@
 module QuantumOpticsBaseMakieExt
 
 import QuantumOpticsBase
-import QuantumOpticsBase: Ket, blochsphereplot, blochsphereplot!, blochsphereplot_axis
+import QuantumOpticsBase: Ket, blochsphereplot, blochsphereplot!, blochsphereplot_axis, fockdistributionplot, fockdistributionplot!, fockdistributionplot_axis
 import Makie
 using Makie: Figure, @recipe, Attributes, Axis3
-using Makie: surface!, arrows3d!, lines!, text!, meshscatter!
+using Makie: surface!, arrows3d!, lines!, text!, meshscatter!, barplot!
 using Makie: Point3f, Vec3f
+using LinearAlgebra: diag
 
 @recipe(BlochSpherePlot, state) do scene
     Attributes(
@@ -145,4 +146,47 @@ function QuantumOpticsBase.blochsphereplot_axis(state; limits=1.6, kwargs...)
     return fig, ax, plt
 end
 
+# ---------------------------------------------------------------------------
+# Fock-state distribution
+# ---------------------------------------------------------------------------
+
+# Occupation probabilities P(n): |⟨n|ψ⟩|² for a Ket, ⟨n|ρ|n⟩ for a density operator.
+_fock_probabilities(state::Ket) = abs2.(state.data)
+_fock_probabilities(state)      = real.(diag(state.data))
+
+@recipe(FockDistributionPlot, state) do scene
+    Attributes()
+end
+
+function Makie.plot!(p::FockDistributionPlot)
+    state_obs = p[1]
+
+    probs = Makie.@lift _fock_probabilities($state_obs)
+    xs    = Makie.@lift Float32.(0:(length($probs) - 1))
+    ys    = Makie.@lift Float32.($probs)
+
+    barplot!(p, xs, ys)
+
+    return p
+end
+
+function QuantumOpticsBase.fockdistributionplot_axis(ax::Makie.AbstractAxis, state;
+                                                     unit_y_range=true, kwargs...)
+    plt = fockdistributionplot!(ax, state; kwargs...)
+    N = length(_fock_probabilities(Makie.to_value(state)))
+    Makie.xlims!(ax, -0.5, N)
+    unit_y_range && Makie.ylims!(ax, 0, 1)
+    return plt
+end
+
+function QuantumOpticsBase.fockdistributionplot_axis(state; kwargs...)
+    fig = Figure(size = (700, 500))
+    ax  = Makie.Axis(fig[1, 1];
+        xlabel = "Fock number",
+        ylabel = "Occupation probability",
+    )
+    plt = QuantumOpticsBase.fockdistributionplot_axis(ax, state; kwargs...)
+    return fig, ax, plt
+end
+
 end # module

src/QuantumOpticsBase.jl

@@ -3,6 +3,7 @@ module QuantumOpticsBase
 using SparseArrays, LinearAlgebra, LRUCache, Strided, UnsafeArrays, FillArrays
 import LinearAlgebra: mul!, rmul!
 import RecursiveArrayTools
+using WeakDepHelpers: WeakDepCache, register_weakdep_cache, @declare_method_is_in_extension
 
 import QuantumInterface: dagger, directsum, ⊕, dm, embed, nsubsystems, expect, identityoperator, identitysuperoperator,
         permutesystems, projector, ptrace, reduced, tensor, ⊗, variance, apply!, basis, AbstractSuperOperator
@@ -75,9 +76,14 @@ export Basis, GenericBasis, CompositeBasis, basis,
                 apply!,
 
         #visualizations
-                blochsphereplot, blochsphereplot!, blochsphereplot_axis
+                blochsphereplot, blochsphereplot!, blochsphereplot_axis, 
+                fockdistributionplot, fockdistributionplot!, fockdistributionplot_axis
 
 
+# Cache of weak-dependency method-error hints, populated by
+# `@declare_method_is_in_extension` in visualization.jl and hooked up in __init__.
+const WEAKDEP_METHOD_ERROR_HINTS = WeakDepCache()
+
 include("bases.jl")
 include("states.jl")
 include("operators.jl")
@@ -107,4 +113,8 @@ include("printing.jl")
 include("apply.jl")
 include("visualization.jl")
 
+function __init__()
+        register_weakdep_cache(WEAKDEP_METHOD_ERROR_HINTS)
+    end
+
 end # module

src/visualization.jl

@@ -24,4 +24,33 @@ or plotting onto an existing one.
 
 Requires a Makie backend be already imported.
 """
-function blochsphereplot_axis end
+function blochsphereplot_axis end
+
+@declare_method_is_in_extension WEAKDEP_METHOD_ERROR_HINTS fockdistributionplot (:Makie,) """
+    fockdistributionplot(state; kwargs...)
+
+Visualize the Fock-state (number-state) distribution of a quantum state as a
+bar chart of occupation probabilities P(n).
+
+For a `Ket` the probabilities are `|⟨n|ψ⟩|²`; for a density operator they are the
+real diagonal entries `⟨n|ρ|n⟩`.
+
+Requires a Makie backend be already imported.
+"""
+
+@declare_method_is_in_extension WEAKDEP_METHOD_ERROR_HINTS fockdistributionplot! (:Makie,) """
+    fockdistributionplot!(ax, state; kwargs...)
+
+In-place version of [`fockdistributionplot`](@ref). Plots onto an existing Makie axis.
+
+Requires a Makie backend be already imported.
+"""
+
+@declare_method_is_in_extension WEAKDEP_METHOD_ERROR_HINTS fockdistributionplot_axis (:Makie,) """
+    fockdistributionplot_axis([ax,] state; kwargs...) -> (Figure, Axis, Plot)
+
+Visualize the Fock-state distribution of a quantum state, creating a new Figure
+and Axis or plotting onto an existing one.
+
+Requires a Makie backend be already imported.
+"""

test/test_fock_plotting.jl

@@ -0,0 +1,65 @@
+@testitem "Fock Distribution Plotting" tags=[:plotting] begin
+    using QuantumOpticsBase
+    using CairoMakie
+
+    b = FockBasis(20)
+
+    # ── Return types ──────────────────────────────────────────────────────────
+    @testset "fockdistributionplot_axis returns Figure, Axis, and plot object" begin
+        fig, ax, plt = fockdistributionplot_axis(coherentstate(b, 2.0))
+        @test fig isa Figure
+        @test ax  isa Axis
+        @test plt isa AbstractPlot
+    end
+
+    # ── Render tests: kets ────────────────────────────────────────────────────
+    @testset "coherent state renders without error" begin
+        fig, _, _ = fockdistributionplot_axis(coherentstate(b, 2.0))
+        save("test_fock.png", fig)
+        @test isfile("test_fock.png")
+        rm("test_fock.png")
+    end
+
+    @testset "number state renders without error" begin
+        fig, _, _ = fockdistributionplot_axis(fockstate(b, 3))
+        save("test_fock_number.png", fig)
+        @test isfile("test_fock_number.png")
+        rm("test_fock_number.png")
+    end
+
+    # ── Render test: density operator ─────────────────────────────────────────
+    @testset "thermal density operator renders without error" begin
+        ρ = thermalstate(number(b), 1.0)
+        fig, _, _ = fockdistributionplot_axis(ρ)
+        save("test_fock_thermal.png", fig)
+        @test isfile("test_fock_thermal.png")
+        rm("test_fock_thermal.png")
+    end
+
+    # ── Render tests: custom attributes ──────────────────────────────────────
+    @testset "Custom color, alpha and width" begin
+        fig, _, _ = fockdistributionplot_axis(coherentstate(b, 2.0); color=:purple, alpha=0.8, width=0.5)
+        save("test_fock_custom.png", fig)
+        @test isfile("test_fock_custom.png")
+        rm("test_fock_custom.png")
+    end
+
+    @testset "unit_y_range toggled off" begin
+        fig, _, _ = fockdistributionplot_axis(coherentstate(b, 2.0); unit_y_range=false)
+        save("test_fock_freey.png", fig)
+        @test isfile("test_fock_freey.png")
+        rm("test_fock_freey.png")
+    end
+
+    # ── Observable reactivity ─────────────────────────────────────────────────
+    @testset "Observable state updates reactively" begin
+        using Makie: Observable
+        state_obs = Observable(coherentstate(b, 1.0))
+        fig, ax, _ = fockdistributionplot_axis(state_obs)
+        state_obs[] = coherentstate(b, 3.0)
+        save("test_fock_observable.png", fig)
+        @test isfile("test_fock_observable.png")
+        rm("test_fock_observable.png")
+    end
+end
+