diff --git a/Project.toml b/Project.toml index 3bb147f0..f63341c4 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/api.md b/docs/src/api.md index 37ad4fa3..eeee979d 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -653,6 +653,9 @@ DenseChiMatrix blochsphereplot blochsphereplot! blochsphereplot_axis +fockdistributionplot +fockdistributionplot! +fockdistributionplot_axis ``` ## [Printing](@id API: Printing) diff --git a/ext/QuantumOpticsBaseMakieExt.jl b/ext/QuantumOpticsBaseMakieExt.jl index 926b1f47..31abec2c 100644 --- a/ext/QuantumOpticsBaseMakieExt.jl +++ b/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 \ No newline at end of file diff --git a/src/QuantumOpticsBase.jl b/src/QuantumOpticsBase.jl index f289952e..4170f727 100644 --- a/src/QuantumOpticsBase.jl +++ b/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 diff --git a/src/visualization.jl b/src/visualization.jl index d11c80b7..a2d8c209 100644 --- a/src/visualization.jl +++ b/src/visualization.jl @@ -24,4 +24,33 @@ or plotting onto an existing one. Requires a Makie backend be already imported. """ -function blochsphereplot_axis end \ No newline at end of file +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. +""" \ No newline at end of file diff --git a/test/test_fock_plotting.jl b/test/test_fock_plotting.jl new file mode 100644 index 00000000..02857a09 --- /dev/null +++ b/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 + \ No newline at end of file