Extend method for ptrace and entanglement

qutip · Mar 19, 2024 · 5c7911c · 5c7911c
1 parent cd0bfa6
commit 5c7911c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/general_functions.jl b/src/general_functions.jl
@@ -173,6 +173,7 @@ function ptrace(QO::QuantumObject{<:AbstractArray{T1},OpType}, sel::Vector{T2})
         return QuantumObject(ρtr, dims=dkeep)
     end
 end
+ptrace(QO::QuantumObject, sel::Int) = ptrace(QO, [sel])
 
 @doc raw"""
     entropy_vn(ρ::QuantumObject; base::Int=0, tol::Real=1e-15)
@@ -241,6 +242,7 @@ function entanglement(QO::QuantumObject{<:AbstractArray{T},OpType}, sel::Vector{
     entropy = entropy_vn(ρ_tr)
     return (entropy > 0) * entropy
 end
+entanglement(QO::QuantumObject, sel::Int) = entanglement(QO, [sel])
 
 @doc raw"""
     expect(O::QuantumObject, ψ::QuantumObject)

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -220,7 +220,7 @@ end
     sol_me = mesolve(H, psi0, t_l, c_ops, e_ops=[sp1 * sm1, sp2 * sm2], progress=false)
     sol_mc = mcsolve(H, psi0, t_l, c_ops, n_traj=500, e_ops=[sp1 * sm1, sp2 * sm2], progress=false)
     @test sum(abs.(sol_mc.expect[1:2, :] .- sol_me.expect[1:2, :])) / length(t_l) < 0.1
-    @test expect(sp1 * sm1, sol_me.states[end]) ≈ expect(sigmap() * sigmam(), ptrace(sol_me.states[end], [1]))
+    @test expect(sp1 * sm1, sol_me.states[end]) ≈ expect(sigmap() * sigmam(), ptrace(sol_me.states[end], 1))
 end
 
 @testset "Dynamical Fock Dimension mesolve" begin
@@ -560,7 +560,7 @@ end
     e = fock(2, 0)
     state = normalize(kron(g, e) + kron(e, g))
     rho = state * state'
-    @test entanglement(state, [1]) / log(2) ≈ 1
+    @test entanglement(state, 1) / log(2) ≈ 1
 end
 
 @testset "Wigner" begin