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/ext/QuantumOpticsBaseMakieExt.jl b/ext/QuantumOpticsBaseMakieExt.jl index 926b1f47..36157c03 100644 --- a/ext/QuantumOpticsBaseMakieExt.jl +++ b/ext/QuantumOpticsBaseMakieExt.jl @@ -1,12 +1,18 @@ module QuantumOpticsBaseMakieExt import QuantumOpticsBase -import QuantumOpticsBase: Ket, blochsphereplot, blochsphereplot!, blochsphereplot_axis +import QuantumOpticsBase: Ket, wigner, + blochsphereplot, blochsphereplot!, blochsphereplot_axis, + wignerplot, wignerplot!, wignerplot_axis import Makie -using Makie: Figure, @recipe, Attributes, Axis3 -using Makie: surface!, arrows3d!, lines!, text!, meshscatter! +using Makie: Figure, @recipe, Attributes, Axis, Axis3, Colorbar, DataAspect +using Makie: surface!, arrows3d!, lines!, text!, meshscatter!, heatmap! using Makie: Point3f, Vec3f +# ═══════════════════════════════════════════════════════════════════════════════ +# Bloch sphere recipe +# ═══════════════════════════════════════════════════════════════════════════════ + @recipe(BlochSpherePlot, state) do scene Attributes( arrowcolor = :red, @@ -108,7 +114,6 @@ function Makie.plot!(p::BlochSpherePlot) return p end - function QuantumOpticsBase.blochsphereplot_axis(ax::Makie.AbstractAxis, state; limits=1.6, kwargs...) ax.perspectiveness = 0f0 lim = Float32(limits) @@ -145,4 +150,71 @@ function QuantumOpticsBase.blochsphereplot_axis(state; limits=1.6, kwargs...) return fig, ax, plt end -end # module \ No newline at end of file +# ═══════════════════════════════════════════════════════════════════════════════ +# Wigner plot recipe +# ═══════════════════════════════════════════════════════════════════════════════ + +@recipe(WignerPlot, state) do scene + Attributes( + xrange = (-5.0, 5.0), + prange = (-5.0, 5.0), + npoints = 100, + colormap = :RdBu, + ) +end + +function Makie.plot!(p::WignerPlot) + state_obs = p[1] + + grid = Makie.@lift begin + s = $state_obs + s.basis isa QuantumOpticsBase.FockBasis || + error("wignerplot requires a FockBasis state, got $(typeof(s.basis))") + xmin, xmax = p[:xrange][] + pmin, pmax = p[:prange][] + n = p[:npoints][] + xvec = collect(LinRange(Float64(xmin), Float64(xmax), n)) + pvec = collect(LinRange(Float64(pmin), Float64(pmax), n)) + W = wigner(s, xvec, pvec) + mx = max(abs(minimum(W)), abs(maximum(W))) + (xvec, pvec, W, mx) + end + + xs = Makie.@lift $grid[1] + pvs = Makie.@lift $grid[2] + Ws = Makie.@lift $grid[3] + clim = Makie.@lift (-$grid[4], $grid[4]) + + heatmap!(p, xs, pvs, Ws; + colormap = p[:colormap], + colorrange = clim, + ) + + return p +end + +function QuantumOpticsBase.wignerplot_axis(ax, state; kwargs...) + wignerplot!(ax, state; kwargs...) +end + +function QuantumOpticsBase.wignerplot_axis(state; kwargs...) + fig = Figure(size = (600, 500)) + ax = Axis(fig[1, 1]; + xlabel = "x", + ylabel = "p", + aspect = DataAspect(), + xgridvisible = false, + ygridvisible = false, + backgroundcolor = :white, + ) + plt = wignerplot_axis(ax, state; kwargs...) + + Colorbar(fig[1, 2], plt; + label = "W(x,p)", + width = 15, + ) + + return fig, ax, plt +end + +end # module diff --git a/src/QuantumOpticsBase.jl b/src/QuantumOpticsBase.jl index f289952e..3699babf 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 +import WeakDepHelpers: WeakDepCache, @declare_method_is_in_extension, register_weakdep_cache import QuantumInterface: dagger, directsum, ⊕, dm, embed, nsubsystems, expect, identityoperator, identitysuperoperator, permutesystems, projector, ptrace, reduced, tensor, ⊗, variance, apply!, basis, AbstractSuperOperator @@ -75,9 +76,12 @@ export Basis, GenericBasis, CompositeBasis, basis, apply!, #visualizations + wignerplot, wignerplot!, wignerplot_axis, blochsphereplot, blochsphereplot!, blochsphereplot_axis +const WEAKDEP_METHOD_ERROR_HINTS = WeakDepCache() + include("bases.jl") include("states.jl") include("operators.jl") @@ -107,4 +111,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..5e45e22c 100644 --- a/src/visualization.jl +++ b/src/visualization.jl @@ -24,4 +24,27 @@ 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 wignerplot (:Makie,) """ + wignerplot(state; kwargs...) + +Visualize the Wigner quasi-probability distribution of a quantum state as a heatmap. + +Requires a Makie backend be already imported. +""" +@declare_method_is_in_extension WEAKDEP_METHOD_ERROR_HINTS wignerplot! (:Makie,) """ + wignerplot!(ax, state; kwargs...) + +In-place version of [`wignerplot`](@ref). Plots onto an existing Makie axis. + +Requires a Makie backend be already imported. +""" +@declare_method_is_in_extension WEAKDEP_METHOD_ERROR_HINTS wignerplot_axis (:Makie,) """ + wignerplot_axis([ax,] state; kwargs...) -> (Figure, Axis, Plot) + +Visualize the Wigner quasi-probability 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/Project.toml b/test/Project.toml index e30c1813..32727195 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -7,6 +7,7 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" OrdinaryDiffEqLowOrderRK = "1344f307-1e59-4825-a18e-ace9aa3fa4c6" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" QuantumInterface = "5717a53b-5d69-4fa3-b976-0bf2f97ca1e5" diff --git a/test/test_wigner_plotting.jl b/test/test_wigner_plotting.jl new file mode 100644 index 00000000..c871d532 --- /dev/null +++ b/test/test_wigner_plotting.jl @@ -0,0 +1,69 @@ +@testitem "Wigner Plot" tags=[:plotting] begin + using QuantumOpticsBase + using QuantumOptics + using CairoMakie + + b = FockBasis(10) # truncated Fock space — adequate for coherent α≤2, Fock n≤3 + + # ── Return types ────────────────────────────────────────────────────────── + @testset "wignerplot_axis returns Figure, Axis, and plot object" begin + ψ = coherentstate(b, 1.0) + fig, ax, plt = wignerplot_axis(ψ) + @test fig isa Figure + @test ax isa Axis + @test plt isa AbstractPlot + end + + # ── Render tests ────────────────────────────────────────────────────────── + @testset "coherent state renders without error" begin + ψ = coherentstate(b, 2.0) + fig, _, _ = wignerplot_axis(ψ) + save("test_wigner_coherent.png", fig) + @test isfile("test_wigner_coherent.png") + rm("test_wigner_coherent.png") + end + + @testset "Fock state renders without error (Wigner can go negative)" begin + ψ = fockstate(b, 3) + fig, _, _ = wignerplot_axis(ψ) + save("test_wigner_fock.png", fig) + @test isfile("test_wigner_fock.png") + rm("test_wigner_fock.png") + end + + # ── Custom attributes ───────────────────────────────────────────────────── + @testset "Custom xrange, prange, npoints" begin + ψ = coherentstate(b, 0.5) + fig, _, _ = wignerplot_axis(ψ; xrange=(-3.0, 3.0), prange=(-3.0, 3.0), npoints=50) + save("test_wigner_custom_range.png", fig) + @test isfile("test_wigner_custom_range.png") + rm("test_wigner_custom_range.png") + end + + @testset "Custom colormap" begin + ψ = coherentstate(b, 1.0) + fig, _, _ = wignerplot_axis(ψ; colormap=:bwr) + save("test_wigner_custom_colormap.png", fig) + @test isfile("test_wigner_custom_colormap.png") + rm("test_wigner_custom_colormap.png") + end + + # ── Observable reactivity ───────────────────────────────────────────────── + @testset "Observable state updates reactively" begin + using Makie: Observable + state_obs = Observable(coherentstate(b, 1.0)) + fig, ax, _ = wignerplot_axis(state_obs) + state_obs[] = coherentstate(b, -1.0) # shift coherent peak to opposite side + save("test_wigner_observable.png", fig) + @test isfile("test_wigner_observable.png") + rm("test_wigner_observable.png") + end + + # ── Error handling ──────────────────────────────────────────────────────── + @testset "Wrong basis type throws error" begin + # wigner is only defined for FockBasis states + b_spin = SpinBasis(1//2) + ψ_spin = spinup(b_spin) + @test_throws "FockBasis" wignerplot_axis(ψ_spin) + end +end