Add Wigner function computation and Makie visualization extension

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"

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
+# ═══════════════════════════════════════════════════════════════════════════════
+# 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

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

src/visualization.jl

@@ -24,4 +24,27 @@ 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 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.
+"""

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"

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