add quantum states, add lines between quantum states on bloch sphere similar to qutip

Author: Fe-r-oz

docs/src/users_guide/plotting_the_bloch_sphere.md

@@ -49,16 +49,32 @@ add_vectors!(b, vecs)
 fig, _ = render(b)
 fig
 ```
+
+# Add Arc, Line, and Vector
+
 ```@example Bloch_sphere_rendering
 clear!(b)
 vec = [[1, 0, 0], [0, 1, 0], [0, 0, 1]];
 add_vectors!(b, vec);
-add_line!(b, [1,0,0], [0,1,0]);
+add_line!(b, [1,0,0], [0,1,0])
 add_arc!(b, [1, 0, 0], [0, 1, 0], [0, 0, 1])
 fig, _ = render(b)
 fig
 ```
 
+# Add Quantum States
+
+```@example Bloch_sphere_rendering
+clear!(b)
+x = basis(2, 0) + basis(2, 1)
+y = basis(2, 0) - im * basis(2, 1)
+z = basis(2, 0)
+b = Bloch()
+add_states!(b, [x, y, z])
+fig, _ = render(b)
+fig
+```
+
 ## Add Multiple Points Around the Equator
 
 ```@example Bloch_sphere_rendering

ext/QuantumToolboxMakieExt.jl

@@ -361,8 +361,40 @@ function _render_bloch_makie(bloch_vec::Vector{Float64}; location = nothing, kwa
     return fig, ax
 end
 
-function add_line!(
-    b::Bloch,
+raw"""
+    add_line!(
+        b::QuantumToolbox.Bloch,
+        start::QuantumToolbox.QuantumObject{<:Union{QuantumToolbox.Ket, QuantumToolbox.Bra, QuantumToolbox.Operator}},
+        endp::QuantumToolbox.QuantumObject{<:Union{QuantumToolbox.Ket, QuantumToolbox.Bra, QuantumToolbox.Operator}};
+        fmt = "k",
+        kwargs...,
+    )
+
+Add a line between two quantum states or operators on the Bloch sphere visualization.
+
+# Arguments
+
+- `b::Bloch`: The Bloch sphere object to modify.
+- `start::QuantumObject`: The starting quantum state or operator. Can be a `Ket`, `Bra`, or `Operator`.
+- `endp::QuantumObject`: The ending quantum state or operator. Can be a `Ket`, `Bra`, or `Operator`.
+- `fmt::String="k"`: (optional) A format string specifying the line style and color (default is black `"k"`).
+- `kwargs...`: Additional keyword arguments forwarded to the underlying line drawing function.
+
+# Description
+
+This function converts the given quantum objects (states or operators) into their Bloch vector representations and adds a line between these two points on the Bloch sphere visualization. 
+
+# Example
+
+```julia
+b = Bloch()
+ψ₁ = normalize(basis(2, 0) + basis(2, 1))
+ψ₂ = normalize(basis(2, 0) - im * basis(2, 1))
+add_line!(b, ψ₁, ψ₂; fmt = "r--")
+```
+"""
+function QuantumToolbox.add_line!(
+    b::QuantumToolbox.Bloch,
     start::QuantumObject{<:Union{Ket,Bra,Operator}},
     endp::QuantumObject{<:Union{Ket,Bra,Operator}};
     fmt = "k",
@@ -381,6 +413,29 @@ function add_line!(
     return add_line!(b, p1, p2; fmt = fmt, kwargs...)
 end
 
+raw"""
+    QuantumToolbox.add_states!(b::Bloch, states::QuantumObject...)
+
+Add one or more quantum states to the Bloch sphere visualization by converting them into Bloch vectors.
+
+# Arguments
+- `b::Bloch`: The Bloch sphere object to modify
+- `states::QuantumObject...`: One or more quantum states (Ket, Bra, or Operator)
+
+"""
+function QuantumToolbox.add_states!(b::Bloch, states::Vector{<:QuantumObject})
+    vecs = Vector{Vector{Float64}}()
+    for state in states
+        vec = if isket(state) || isbra(state)
+            _state_to_bloch(state)
+        else
+            _dm_to_bloch(state)
+        end
+        push!(vecs, vec)
+    end
+    return add_vectors!(b, vecs)
+end
+
 function QuantumToolbox.render(b::QuantumToolbox.Bloch; location = nothing)
     fig, location = _getFigAndLocation(location)
     if isnothing(b.fig) || b.fig !== fig
@@ -459,17 +514,23 @@ function QuantumToolbox.render(b::QuantumToolbox.Bloch; location = nothing)
     end
     for (line, fmt, kwargs) in b.lines
         x, y, z = line
-        color = Base.get(kwargs, :color, fmt == "k" ? Makie.RGBAf(0.1, 0.1, 0.1, 0.8) : fmt)
-        linewidth = Base.get(kwargs, :linewidth, 1.0)
-        Makie.lines!(
-            b.ax,
-            x,
-            y,
-            z;
-            color = color,
-            linewidth = linewidth,
-            linestyle = Base.get(kwargs, :linestyle, :solid),
-        )
+        color_map =
+            Dict("k" => :black, "r" => :red, "g" => :green, "b" => :blue, "c" => :cyan, "m" => :magenta, "y" => :yellow)
+        c = get(color_map, first(fmt), :black)
+        ls = nothing
+        if occursin("--", fmt)
+            ls = :dash
+        elseif occursin(":", fmt)
+            ls = :dot
+        elseif occursin("-.", fmt)
+            ls = :dashdot
+        else
+            ls = :solid
+        end
+        color = get(kwargs, :color, c)
+        linewidth = get(kwargs, :linewidth, 1.0)
+        linestyle = get(kwargs, :linestyle, ls)
+        Makie.lines!(b.ax, x, y, z; color = color, linewidth = linewidth, linestyle = linestyle)
     end
     for arc_pts in b.arcs
         if length(arc_pts) >= 2

src/visualization.jl

@@ -1,5 +1,14 @@
 export plot_wigner,
-    plot_fock_distribution, plot_bloch, Bloch, add_points!, add_vectors!, add_line!, add_arc!, clear!, render
+    plot_fock_distribution,
+    plot_bloch,
+    Bloch,
+    add_points!,
+    add_vectors!,
+    add_line!,
+    add_arc!,
+    clear!,
+    render,
+    add_states!
 
 @doc raw"""
     plot_wigner(
@@ -311,6 +320,13 @@ function add_arc!(b::Bloch, p1::Vector{<:Real}, p2::Vector{<:Real}, p3::Vector{<
     return push!(b.arcs, [convert(Vector{Float64}, p1), convert(Vector{Float64}, p2), convert(Vector{Float64}, p3)])
 end
 
+@doc raw"""
+    QuantumToolbox.add_states!(b::Bloch, states::QuantumObject...)
+
+Add one or more quantum states to the Bloch sphere visualization by converting them into Bloch vectors.
+"""
+function add_states! end
+
 function clear!(b::Bloch)
     empty!(b.points)
     empty!(b.vectors)