add kwargs kind for add_states!

Author: ytdHuang

src/visualization.jl

@@ -381,13 +381,14 @@ function add_arc!(
 end
 
 @doc raw"""
-    add_states!(b::Bloch, states::Vector{QuantumObject})
+    add_states!(b::Bloch, states::Vector{QuantumObject}; kind::Symbol = :vector, kwargs...)
 
 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::Vector{QuantumObject}`: One or more quantum states ([`Ket`](@ref), [`Bra`](@ref), or [`Operator`](@ref))
+- `kind::Symbol`: Type of object to plot (can be either `:vector` or `:point`). Default: `:vector`
 
 # Example
 
@@ -399,16 +400,18 @@ b = Bloch();
 add_states!(b, [x, y, z])
 ```
 """
-function add_states!(b::Bloch, states::Vector{<:QuantumObject})
+function add_states!(b::Bloch, states::Vector{<:QuantumObject}; kind::Symbol = :vector, kwargs...)
     vecs = map(state -> _state_to_bloch(state), states)
-    append!(b.vectors, vecs)
-    return b.vectors
-end
-
-function add_states!(b::Bloch, state::QuantumObject)
-    push!(b.vectors, _state_to_bloch(state))
-    return b.vectors
+    if kind == :vector
+        add_vectors!(b, vecs)
+    elseif kind == :point
+        add_points!(b, hcat(vecs...), kwargs...)
+    else
+        throw(ArgumentError("Invalid kind = :$kind"))
+    end
+    return nothing
 end
+add_states!(b::Bloch, state::QuantumObject; kind::Symbol = :vector, kwargs...) = add_states!(b, [state], kind = kind, kwargs...)
 
 _state_to_bloch(state::QuantumObject{Ket}) = _ket_to_bloch(state)
 _state_to_bloch(state::QuantumObject{Bra}) = _ket_to_bloch(state')

test/ext-test/cpu/makie/makie_ext.jl

@@ -136,6 +136,7 @@ end
     @test isempty(b.vectors)
     @test isempty(b.lines)
     @test isempty(b.arcs)
+
     b = Bloch()
     add_points!(b, hcat([1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]))
     add_vectors!(b, [[1, 1, 0], [0, 1, 1]])
@@ -149,6 +150,9 @@ end
         @test false
         @info "Render threw unexpected error" exception=e
     end
+    b.vector_arrowsize = [0.7, 0.8, 1.5] # 1.5 is the length of arrow head (too long)
+    @test_throws ArgumentError render(b)
+
     b = Bloch()
     ψ₁ = normalize(basis(2, 0) + basis(2, 1))
     ψ₂ = normalize(basis(2, 0) - im * basis(2, 1))
@@ -167,14 +171,20 @@ end
     Pauli_Ops = [sigmax(), sigmay(), sigmaz()]
     ψ = rand_ket(2)
     ρ = rand_dm(2)
+    states = [ψ, ρ]
     x = basis(2, 0) + basis(2, 1)             # unnormalized Ket
     ρ1 = 0.3 * rand_dm(2) + 0.4 * rand_dm(2)  # unnormalized density operator
     ρ2 = Qobj(rand(ComplexF64, 2, 2))         # unnormalized and non-Hermitian Operator
-    add_states!(b, [ψ, ρ])
-    @test_logs (:warn,) (:warn,) (:warn,) (:warn,) add_states!(b, [x, ρ1, ρ2])
+    add_states!(b, states, kind = :vector)
+    add_states!(b, states, kind = :point)
+    @test length(b.vectors) == 2
+    @test length(b.points) == 1
     @test all(expect(Pauli_Ops, ψ) .≈ (b.vectors[1]))
     @test all(expect(Pauli_Ops, ρ) .≈ (b.vectors[2]))
+    @test all([b.vectors[j][k] ≈ b.points[1][k,j] for j in (1, 2) for k in (1, 2, 3)])
+    @test_logs (:warn,) (:warn,) (:warn,) (:warn,) add_states!(b, [x, ρ1, ρ2])
     @test length(b.vectors) == 5
+    @test_throws ArgumentError add_states!(b, states, kind = :wrong)
 
     th = range(0, 2π; length = 20)
     xp = cos.(th);