Makie bloch sphere extension

Project.toml

@@ -17,6 +17,12 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
+[weakdeps]
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+
+[extensions]
+QuantumOpticsBaseMakieExt = ["Makie"]
+
 [compat]
 Adapt = "3.3, 4"
 FFTW = "1.2"
@@ -25,10 +31,11 @@ FastGaussQuadrature = "0.5, 1"
 FillArrays = "1.9"
 LRUCache = "1"
 LinearAlgebra = "1"
+Makie = "0.24"
 QuantumInterface = "0.4"
 Random = "1"
 RecursiveArrayTools = "3.1, 4"
 SparseArrays = "1"
 Strided = "1, 2"
 UnsafeArrays = "1"
-julia = "1.10"
+julia = "1.10"

ext/QuantumOpticsBaseMakieExt.jl

@@ -0,0 +1,148 @@
+module QuantumOpticsBaseMakieExt
+
+import QuantumOpticsBase
+import QuantumOpticsBase: Ket, blochsphereplot, blochsphereplot!, blochsphereplot_axis
+import Makie
+using Makie: Figure, @recipe, Attributes, Axis3
+using Makie: surface!, arrows3d!, lines!, text!, meshscatter!
+using Makie: Point3f, Vec3f
+
+@recipe(BlochSpherePlot, state) do scene
+    Attributes(
+        arrowcolor    = :red,
+        spherecolor   = :lightblue,
+        spherealpha   = 0.15,
+        showwireframe = true,
+        showaxes      = true,
+        showlabels    = true,
+        labelsize     = 18,
+        shaftradius   = 0.018,
+        tipradius     = 0.050,
+        tiplength     = 0.10,
+    )
+end
+
+function Makie.plot!(p::BlochSpherePlot)
+    state_obs = p[1]
+
+    blochvec = Makie.@lift begin
+        s = $state_obs
+        length(s.data) == 2 ||
+            error("BlochSphere requires a 2-level (spin-1/2) state")
+        α, β = s.data
+        Vec3f(
+            Float32(2 * real(conj(α) * β)),
+            Float32(2 * imag(conj(α) * β)),
+            Float32(abs2(α) - abs2(β)),
+        )
+    end
+
+    let npts = 200
+        θ = LinRange(0f0, 2f0π, npts)
+        φ = LinRange(0f0, Float32(π), npts)
+        xs = Float32[cos(t) * sin(q) for t in θ, q in φ]
+        ys = Float32[sin(t) * sin(q) for t in θ, q in φ]
+        zs = Float32[cos(q)          for _  in θ, q in φ]
+        c  = Makie.to_color(p[:spherecolor][])
+        α  = Float32(p[:spherealpha][])
+        rgba = Makie.RGBAf(c.r, c.g, c.b, α)
+        surface!(p, xs, ys, zs;
+            color        = fill(rgba, npts, npts),
+            transparency = true,
+        )
+    end
+
+    if p[:showwireframe][]
+        ncirc = 120
+        θc = LinRange(0f0, 2f0π, ncirc)
+        for pts in (
+            [Point3f( cos(t),  sin(t), 0f0) for t in θc],
+            [Point3f( cos(t), 0f0, sin(t)) for t in θc],
+            [Point3f(0f0, cos(t), sin(t)) for t in θc],
+        )
+            lines!(p, pts; color = (:black, 0.70), linewidth = 1.2)
+        end
+    end
+
+    if p[:showaxes][]
+        r = 1.18f0
+        for (a, b) in (
+            (Point3f(-r, 0, 0), Point3f(r, 0, 0)),
+            (Point3f(0, -r, 0), Point3f(0, r, 0)),
+            (Point3f(0, 0, -r), Point3f(0, 0, r)),
+        )
+            lines!(p, [a, b]; color = :black, linewidth = 1, linestyle = :dash)
+        end
+    end
+
+    arrows3d!(p,
+        [Point3f(0, 0, 0)],
+        Makie.@lift([$blochvec]);
+        shaftradius = p[:shaftradius],
+        tipradius   = p[:tipradius],
+        tiplength   = p[:tiplength],
+        color       = p[:arrowcolor],
+    )
+
+    meshscatter!(p,
+        Makie.@lift([Point3f($blochvec)]);
+        color      = p[:arrowcolor],
+        markersize = 0.06,
+    )
+
+    if p[:showlabels][]
+        ls  = p[:labelsize][]
+        off = 1.40f0
+        for (pos, lbl, align) in (
+            (Point3f( 0f0,   0f0,  off), "|0⟩",  (:center, :bottom)),
+            (Point3f( 0f0,   0f0, -off), "|1⟩",  (:center, :top   )),
+            (Point3f( off,   0f0,  0f0), "|+⟩",  (:left,   :center)),
+            (Point3f(-off,   0f0,  0f0), "|-⟩",  (:right,  :center)),
+            (Point3f( 0f0,   off,  0f0), "|+i⟩", (:left,   :center)),
+            (Point3f( 0f0,  -off,  0f0), "|-i⟩", (:right,  :center)),
+        )
+            text!(p, pos; text = lbl, fontsize = ls, align = align)
+        end
+    end
+
+    return p
+end
+
+
+function QuantumOpticsBase.blochsphereplot_axis(ax::Makie.AbstractAxis, state; limits=1.6, kwargs...)
+    ax.perspectiveness = 0f0
+    lim = Float32(limits)
+    Makie.limits!(ax, -lim, lim, -lim, lim, -lim, lim)
+    blochsphereplot!(ax, state; kwargs...)
+end
+
+function QuantumOpticsBase.blochsphereplot_axis(state; limits=1.6, kwargs...)
+    fig = Figure(size = (700, 700))
+    ax  = Axis3(fig[1, 1];
+        aspect   = :data,
+        viewmode = :fit,
+        perspectiveness = 0f0,
+        xticksvisible      = false,
+        yticksvisible      = false,
+        zticksvisible      = false,
+        xticklabelsvisible = false,
+        yticklabelsvisible = false,
+        zticklabelsvisible = false,
+        xlabelvisible      = false,
+        ylabelvisible      = false,
+        zlabelvisible      = false,
+        xspinesvisible     = false,
+        yspinesvisible     = false,
+        zspinesvisible     = false,
+        xgridvisible       = false,
+        ygridvisible       = false,
+        zgridvisible       = false,
+        xypanelvisible     = false,
+        xzpanelvisible     = false,
+        yzpanelvisible     = false,
+    )
+    plt = QuantumOpticsBase.blochsphereplot_axis(ax, state; limits, kwargs...)
+    return fig, ax, plt
+end
+
+end # module

src/QuantumOpticsBase.jl

@@ -70,9 +70,13 @@ export Basis, GenericBasis, CompositeBasis, basis,
         PauliBasis, PauliTransferMatrix, DensePauliTransferMatrix,
                 ChiMatrix, DenseChiMatrix, avg_gate_fidelity,
         SumBasis, directsum, ⊕, LazyDirectSum, getblock, setblock!,
-        qfunc, wigner, coherentspinstate, qfuncsu2, wignersu2
+        qfunc, wigner, coherentspinstate, qfuncsu2, wignersu2,
         #apply
-                apply!
+                apply!,
+
+        #visualizations
+                blochsphereplot, blochsphereplot!, blochsphereplot_axis
+
 
 include("bases.jl")
 include("states.jl")
@@ -101,5 +105,6 @@ include("spinors.jl")
 include("phasespace.jl")
 include("printing.jl")
 include("apply.jl")
+include("visualization.jl")
 
 end # module

src/visualization.jl

@@ -0,0 +1,27 @@
+"""
+    blochsphereplot(state::Ket; kwargs...)
+
+Visualize a pure qubit state as an arrow on a Bloch sphere.
+
+Requires a Makie backend be already imported.
+"""
+function blochsphereplot end
+
+"""
+    blochsphereplot!(ax, state::Ket; kwargs...)
+
+In-place version of [`blochsphereplot`](@ref). Plots onto an existing Makie axis.
+
+Requires a Makie backend be already imported.
+"""
+function blochsphereplot! end
+
+"""
+    blochsphereplot_axis([ax,] state::Ket; kwargs...) -> (Figure, Axis3, Plot)
+
+Visualize a pure qubit state on a Bloch sphere, creating a new Figure and Axis3
+or plotting onto an existing one.
+
+Requires a Makie backend be already imported.
+"""
+function blochsphereplot_axis end

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

test/test_plotting.jl

@@ -0,0 +1,63 @@
+@testitem "Bloch Sphere Plotting" tags=[:plotting] begin
+    using QuantumOpticsBase
+    using CairoMakie   
+
+    b = SpinBasis(1//2)
+
+    # ── Return types ──────────────────────────────────────────────────────────
+    @testset "blochsphereplot_axis returns Figure, Axis3, and plot object" begin
+        fig, ax, plt = blochsphereplot_axis(spinup(b))
+        @test fig isa Figure
+        @test ax  isa Axis3
+        @test plt isa AbstractPlot
+    end
+
+    # ── Render test ───────────────────────────────────────────────────────────
+    @testset "arbitrary state renders without error" begin
+        ψ = cos(π/6)*spinup(b) + exp(im*π/4)*sin(π/6)*spindown(b)
+        fig, _, _ = blochsphereplot_axis(ψ)
+        save("test_bloch.png", fig)
+        @test isfile("test_bloch.png")
+        rm("test_bloch.png")
+    end
+
+    # ── Render tests: custom attributes ──────────────────────────────────────
+    @testset "Custom arrowcolor and spherealpha" begin
+        fig, _, _ = blochsphereplot_axis(spinup(b); arrowcolor=:blue, spherealpha=0.3)
+        save("test_bloch_custom_color.png", fig)
+        @test isfile("test_bloch_custom_color.png")
+        rm("test_bloch_custom_color.png")
+    end
+
+    @testset "Custom spherecolor" begin
+        fig, _, _ = blochsphereplot_axis(spindown(b); spherecolor=:pink, spherealpha=0.2)
+        save("test_bloch_custom_sphere.png", fig)
+        @test isfile("test_bloch_custom_sphere.png")
+        rm("test_bloch_custom_sphere.png")
+    end
+
+    @testset "Wireframe, labels and axes toggled off" begin
+        fig, _, _ = blochsphereplot_axis(spinup(b); showwireframe=false, showlabels=false, showaxes=false)
+        save("test_bloch_minimal.png", fig)
+        @test isfile("test_bloch_minimal.png")
+        rm("test_bloch_minimal.png")
+    end
+
+    # ── Error handling ────────────────────────────────────────────────────────
+    @testset "Wrong dimension state throws error" begin
+        b3  = SpinBasis(1)
+        ψ_3 = basisstate(b3, 1)
+        @test_throws ErrorException blochsphereplot_axis(ψ_3)
+    end
+
+    # ── Observable reactivity ─────────────────────────────────────────────────
+    @testset "Observable state updates reactively" begin
+        using Makie: Observable
+        state_obs = Observable(spinup(b))
+        fig, ax, _ = blochsphereplot_axis(state_obs)
+        state_obs[] = spindown(b)
+        save("test_bloch_observable.png", fig)
+        @test isfile("test_bloch_observable.png")
+        rm("test_bloch_observable.png")
+    end
+end