trace, partial trace, and superoperators

src/QSymbolicsBase/QSymbolicsBase.jl

@@ -1,7 +1,7 @@
 using Symbolics
 import Symbolics: simplify
 using SymbolicUtils
-import SymbolicUtils: Symbolic,_isone,flatten_term,isnotflat,Chain,Fixpoint,Prewalk
+import SymbolicUtils: Symbolic,_isone,flatten_term,isnotflat,Chain,Fixpoint,Prewalk,sorted_arguments
 using TermInterface
 import TermInterface: isexpr,head,iscall,children,operation,arguments,metadata,maketerm
 
@@ -11,9 +11,9 @@ import LinearAlgebra: eigvecs,ishermitian,inv
 import QuantumInterface:
     apply!,
     tensor, ⊗,
-    basis,Basis,SpinBasis,FockBasis,
+    basis,Basis,SpinBasis,FockBasis,CompositeBasis,
     nqubits,
-    projector,dagger,
+    projector,dagger,tr,ptrace,
     AbstractBra,AbstractKet,AbstractOperator,AbstractSuperOperator
 
 export SymQObj,QObj,
@@ -23,28 +23,30 @@ export SymQObj,QObj,
        apply!,
        express,
        tensor,⊗,
-       dagger,projector,commutator,anticommutator,
+       dagger,projector,commutator,anticommutator,expand,tr,ptrace,
        I,X,Y,Z,σˣ,σʸ,σᶻ,Pm,Pp,σ₋,σ₊,
        H,CNOT,CPHASE,XCX,XCY,XCZ,YCX,YCY,YCZ,ZCX,ZCY,ZCZ,
        X1,X2,Y1,Y2,Z1,Z2,X₁,X₂,Y₁,Y₂,Z₁,Z₂,L0,L1,Lp,Lm,Lpi,Lmi,L₀,L₁,L₊,L₋,L₊ᵢ,L₋ᵢ,
        vac,F₀,F0,F₁,F1,
        N,n̂,Create,âꜛ,Destroy,â,SpinBasis,FockBasis,
-       SBra,SKet,SOperator,SHermitianOperator,SUnitaryOperator,SHermitianUnitaryOperator,
-       @ket,@bra,@op,
+       SBra,SKet,SOperator,SHermitianOperator,SUnitaryOperator,SHermitianUnitaryOperator,SSuperOperator,
+       @ket,@bra,@op,@superop,
        SAdd,SAddBra,SAddKet,SAddOperator,
        SScaled,SScaledBra,SScaledOperator,SScaledKet,
        STensorBra,STensorKet,STensorOperator,
        SZeroBra,SZeroKet,SZeroOperator,
-       SProjector,MixedState,IdentityOp,SInvOperator,
-       SApplyKet,SApplyBra,SMulOperator,SSuperOpApply,SCommutator,SAnticommutator,SDagger,SBraKet,SOuterKetBra,
+       MixedState,IdentityOp,
+       SProjector,SDagger,STrace,SPartialTrace,SInvOperator,
+       SApplyKet,SApplyBra,SMulOperator,SSuperOpApply,SCommutator,SAnticommutator,SBraKet,SOuterKetBra,
        HGate,XGate,YGate,ZGate,CPHASEGate,CNOTGate,
        XBasisState,YBasisState,ZBasisState,
        NumberOp,CreateOp,DestroyOp,
        XCXGate,XCYGate,XCZGate,YCXGate,YCYGate,YCZGate,ZCXGate,ZCYGate,ZCZGate,
        qsimplify,qsimplify_pauli,qsimplify_commutator,qsimplify_anticommutator,
        qexpand,
-       isunitary
-
+       isunitary,
+       KrausRepr,kraus
+
 ##
 # Metadata cache helpers
 ##
@@ -161,6 +163,7 @@ include("utils.jl")
 include("literal_objects.jl")
 include("basic_ops_homogeneous.jl")
 include("basic_ops_inhomogeneous.jl")
+include("basic_superops.jl")
 include("linalg.jl")
 include("predefined.jl")
 include("predefined_CPTP.jl")

src/QSymbolicsBase/basic_ops_homogeneous.jl

@@ -147,7 +147,7 @@ function Base.:(*)(xs::Symbolic{AbstractOperator}...)
     end
 end
 Base.show(io::IO, x::SMulOperator) = print(io, join(map(string, arguments(x)),""))
-basis(x::SMulOperator) = basis(x.terms)
+basis(x::SMulOperator) = basis(first(x.terms))
 
 """Tensor product of quantum objects (kets, operators, or bras)
 

src/QSymbolicsBase/basic_ops_inhomogeneous.jl

@@ -90,24 +90,6 @@ Base.hash(x::SBraKet, h::UInt) = hash((head(x), arguments(x)), h)
 maketerm(::Type{<:SBraKet}, f, a, t, m) = f(a...)
 Base.isequal(x::SBraKet, y::SBraKet) = isequal(x.bra, y.bra) && isequal(x.ket, y.ket)
 
-"""Symbolic application of a superoperator on an operator"""
-@withmetadata struct SSuperOpApply <: Symbolic{AbstractOperator}
-    sop
-    op
-end
-isexpr(::SSuperOpApply) = true
-iscall(::SSuperOpApply) = true
-arguments(x::SSuperOpApply) = [x.sop,x.op]
-operation(x::SSuperOpApply) = *
-head(x::SSuperOpApply) = :*
-children(x::SSuperOpApply) = [:*,x.sop,x.op]
-Base.:(*)(sop::Symbolic{AbstractSuperOperator}, op::Symbolic{AbstractOperator}) = SSuperOpApply(sop,op)
-Base.:(*)(sop::Symbolic{AbstractSuperOperator}, op::SZeroOperator) = SZeroOperator()
-Base.:(*)(sop::Symbolic{AbstractSuperOperator}, k::Symbolic{AbstractKet}) = SSuperOpApply(sop,SProjector(k))
-Base.:(*)(sop::Symbolic{AbstractSuperOperator}, k::SZeroKet) = SZeroKet()
-Base.show(io::IO, x::SSuperOpApply) = begin print(io, x.sop); print(io, x.op) end
-basis(x::SSuperOpApply) = basis(x.op)
-
 """Symbolic outer product of a ket and a bra
 ```jldoctest 
 julia> @bra b; @ket k;

src/QSymbolicsBase/basic_superops.jl

@@ -0,0 +1,64 @@
+##
+# Superoperator representations
+##
+
+"""Kraus representation of a quantum channel
+
+```jldoctest
+julia> @superop ℰ;
+
+julia> @op A₁; @op A₂; @op A₃;
+
+julia> K = kraus(ℰ, A₁, A₂, A₃);
+
+julia> @op ρ;
+
+julia> K*ρ
+(A₁ρA₁†+A₂ρA₂†+A₃ρA₃†)
+```
+"""
+@withmetadata struct KrausRepr <: Symbolic{AbstractSuperOperator}
+    sop
+    krausops
+end
+isexpr(::KrausRepr) = true
+iscall(::KrausRepr) = true
+arguments(x::KrausRepr) = [x.sop, x.krausops]
+operation(x::KrausRepr) = kraus
+head(x::KrausRepr) = :kraus
+children(x::KrausRepr) = [:kraus, x.sop, x.krausops]
+kraus(s::Symbolic{AbstractSuperOperator}, xs::Symbolic{AbstractOperator}...) = KrausRepr(s,collect(xs))
+symbollabel(x::KrausRepr) = symbollabel(x.sop)
+basis(x::KrausRepr) = basis(x.sop)
+Base.show(io::IO, x::KrausRepr) = print(io, symbollabel(x))
+
+##
+# Superoperator operations
+##
+
+"""Symbolic application of a superoperator on an operator
+
+```jldoctest
+julia> @op A; @superop S;
+
+julia> S*A
+S[A]
+"""
+@withmetadata struct SSuperOpApply <: Symbolic{AbstractOperator}
+    sop
+    op
+end
+isexpr(::SSuperOpApply) = true
+iscall(::SSuperOpApply) = true
+arguments(x::SSuperOpApply) = [x.sop,x.op]
+operation(x::SSuperOpApply) = *
+head(x::SSuperOpApply) = :*
+children(x::SSuperOpApply) = [:*,x.sop,x.op]
+Base.:(*)(sop::Symbolic{AbstractSuperOperator}, op::Symbolic{AbstractOperator}) = SSuperOpApply(sop,op)
+Base.:(*)(sop::Symbolic{AbstractSuperOperator}, op::SZeroOperator) = SZeroOperator()
+Base.:(*)(sop::Symbolic{AbstractSuperOperator}, k::Symbolic{AbstractKet}) = SSuperOpApply(sop,SProjector(k))
+Base.:(*)(sop::Symbolic{AbstractSuperOperator}, k::SZeroKet) = SZeroOperator()
+Base.:(*)(sop::KrausRepr, op::Symbolic{AbstractOperator}) = (+)((i*op*dagger(i) for i in sop.krausops)...)
+Base.:(*)(sop::KrausRepr, k::Symbolic{AbstractKet}) = (+)((i*SProjector(k)*dagger(i) for i in sop.krausops)...)
+Base.show(io::IO, x::SSuperOpApply) = print(io, "$(x.sop)[$(x.op)]")
+basis(x::SSuperOpApply) = basis(x.op)

src/QSymbolicsBase/linalg.jl

@@ -89,6 +89,140 @@ function Base.show(io::IO, x::SDagger)
     print(io,"†")
 end
 
+"""Trace of an operator
+
+```jldoctest
+julia> @op A; @op B;
+
+julia> tr(A)
+tr(A)
+
+julia> tr(commutator(A, B))
+0
+
+julia> @bra b; @ket k;
+
+julia> tr(k*b)
+⟨b||k⟩
+```
+"""
+@withmetadata struct STrace <: Symbolic{Complex}
+    op::Symbolic{AbstractOperator}
+end
+isexpr(::STrace) = true
+iscall(::STrace) = true
+arguments(x::STrace) = [x.op]
+sorted_arguments(x::STrace) = arguments(x)
+operation(x::STrace) = tr
+head(x::STrace) = :tr
+children(x::STrace) = [:tr, x.op]
+Base.show(io::IO, x::STrace) = print(io, "tr($(x.op))")
+tr(x::Symbolic{AbstractOperator}) = STrace(x)
+tr(x::SScaled{AbstractOperator}) = x.coeff*tr(x.obj)
+tr(x::SAdd{AbstractOperator}) = (+)((tr(i) for i in arguments(x))...)
+tr(x::SOuterKetBra) = x.bra*x.ket
+tr(x::SCommutator) = 0
+tr(x::STensorOperator) = (*)((tr(i) for i in arguments(x))...)
+Base.hash(x::STrace, h::UInt) = hash((head(x), arguments(x)), h)
+Base.isequal(x::STrace, y::STrace) = isequal(x.op, y.op)
+
+"""Partial trace over system i of a composite quantum system
+
+```jldoctest
+julia> @op 𝒪 SpinBasis(1//2)⊗SpinBasis(1//2);
+
+julia> op = ptrace(𝒪, 1)
+tr1(𝒪)
+
+julia> QuantumSymbolics.basis(op)
+Spin(1/2)
+
+julia> @op A; @op B;
+
+julia> ptrace(A⊗B, 1)
+(tr(A))B
+
+julia> @ket k; @bra b;
+
+julia> pure_state = A ⊗ (k*b)
+(A⊗|k⟩⟨b|)
+
+julia> ptrace(pure_state, 1)
+(tr(A))|k⟩⟨b|
+
+julia> ptrace(pure_state, 2)
+(⟨b||k⟩)A
+
+julia> mixed_state = (A⊗(k*b)) + ((k*b)⊗B)
+((A⊗|k⟩⟨b|)+(|k⟩⟨b|⊗B))
+
+julia> ptrace(mixed_state, 1)
+((0 + ⟨b||k⟩)B+(tr(A))|k⟩⟨b|)
+
+julia> ptrace(mixed_state, 2)
+((0 + ⟨b||k⟩)A+(tr(B))|k⟩⟨b|)
+```
+"""
+@withmetadata struct SPartialTrace <: Symbolic{AbstractOperator}
+    obj
+    sys::Int
+end
+isexpr(::SPartialTrace) = true
+iscall(::SPartialTrace) = true
+arguments(x::SPartialTrace) = [x.obj, x.sys]
+operation(x::SPartialTrace) = ptrace
+head(x::SPartialTrace) = :ptrace
+children(x::SPartialTrace) = [:ptrace, x.obj, x.sys]
+function basis(x::SPartialTrace)
+    obj_bases = collect(basis(x.obj).bases)
+    new_bases = deleteat!(copy(obj_bases), x.sys)
+    tensor(new_bases...)
+end
+Base.show(io::IO, x::SPartialTrace) = print(io, "tr$(x.sys)($(x.obj))")
+function ptrace(x::Symbolic{AbstractOperator}, s) 
+    ex = isexpr(x) ? qexpand(x) : x
+    if isa(ex, typeof(x))
+        if isa(basis(x), CompositeBasis)
+            SPartialTrace(x, s)
+        else
+            throw(ArgumentError("cannot take partial trace of a single quantum system"))
+        end
+    else
+        ptrace(ex, s)
+    end
+end
+function ptrace(x::SAddOperator, s)
+    terms = arguments(x)
+    add_terms = []
+    for i in terms
+        if isexpr(i) && operation(i) === ⊗
+            isa(i, SScaledOperator) ? prod_terms = arguments(i.obj) : prod_terms = arguments(i)
+            sys_op = prod_terms[s]
+            new_terms = deleteat!(copy(prod_terms), s)
+            isone(length(new_terms)) ? push!(add_terms, tr(sys_op)*first(new_terms)) : push!(add_terms, tr(sys_op)*STensorOperator(new_terms))
+        else
+            throw(ArgumentError("cannot take partial trace of a single quantum system"))
+        end
+    end
+    (+)(add_terms...)
+end
+function ptrace(x::STensorOperator, s)
+    ex = qexpand(x)
+    if isa(ex, SAddOperator)
+        ptrace(ex, s)
+    else
+        terms = arguments(ex)
+        newterms = []
+        if isa(basis(terms[s]), CompositeBasis)
+            SPartial(ex, s)
+        else 
+            sys_op = terms[s]
+            new_terms = deleteat!(copy(terms), s)
+            isone(length(new_terms)) ? tr(sys_op)*first(new_terms) : tr(sys_op)*STensorOperator(new_terms)
+        end
+    end
+end
+
 """Inverse Operator
 
 ```jldoctest

src/QSymbolicsBase/literal_objects.jl

@@ -67,15 +67,29 @@ SHermitianUnitaryOperator(name) = SHermitianUnitaryOperator(name, qubit_basis)
 ishermitian(::SHermitianUnitaryOperator) = true
 isunitary(::SHermitianUnitaryOperator) = true
 
-const SymQ = Union{SKet,SBra,SOperator,SHermitianOperator,SUnitaryOperator,SHermitianUnitaryOperator}
+struct SSuperOperator <: Symbolic{AbstractSuperOperator}
+    name::Symbol
+    basis::Basis
+end
+SSuperOperator(name) = SSuperOperator(name, qubit_basis)
+macro superop(name, basis)
+    :($(esc(name)) = SSuperOperator($(Expr(:quote, name)), $(basis)))
+end
+macro superop(name)
+    :($(esc(name)) = SSuperOperator($(Expr(:quote, name))))
+end
+ishermitian(x::SSuperOperator) = false
+isunitary(x::SSuperOperator) = false
+
+const SymQ = Union{SKet,SBra,SOperator,SHermitianOperator,SUnitaryOperator,SHermitianUnitaryOperator,SSuperOperator}
 isexpr(::SymQ) = false
 metadata(::SymQ) = nothing
 symbollabel(x::SymQ) = x.name
 basis(x::SymQ) = x.basis
 
 Base.show(io::IO, x::SKet) = print(io, "|$(symbollabel(x))⟩")
 Base.show(io::IO, x::SBra) = print(io, "⟨$(symbollabel(x))|")
-Base.show(io::IO, x::Union{SOperator,SHermitianOperator,SUnitaryOperator,SHermitianUnitaryOperator}) = print(io, "$(symbollabel(x))")
+Base.show(io::IO, x::Union{SOperator,SHermitianOperator,SUnitaryOperator,SHermitianUnitaryOperator,SSuperOperator}) = print(io, "$(symbollabel(x))")
 Base.show(io::IO, x::SymQObj) = print(io, symbollabel(x)) # fallback that probably is not great
 
 struct SZero{T<:QObj} <: Symbolic{T} end

src/QSymbolicsBase/rules.jl

@@ -102,9 +102,10 @@ RULES_EXPAND = [
     @rule(+(~~ops) ⊗ ~o1 => +(map(op -> op ⊗ ~o1, ~~ops)...)),
     @rule(~o1 * +(~~ops) => +(map(op -> ~o1 * op, ~~ops)...)),
     @rule(+(~~ops) * ~o1 => +(map(op -> op * ~o1, ~~ops)...)),
-    @rule(+(~~ops) * ~o1 => +(map(op -> op * ~o1, ~~ops)...)),
+    @rule(⊗(~~ops1::_vecisa(Symbolic{AbstractBra})) * ⊗(~~ops2::_vecisa(Symbolic{AbstractKet})) => *(map(*, ~~ops1, ~~ops2)...)),
     @rule(⊗(~~ops1::_vecisa(Symbolic{AbstractOperator})) * ⊗(~~ops2::_vecisa(Symbolic{AbstractOperator})) => ⊗(map(*, ~~ops1, ~~ops2)...)),
-    @rule(⊗(~~ops1::_vecisa(Symbolic{AbstractBra})) * ⊗(~~ops2::_vecisa(Symbolic{AbstractKet})) => *(map(*, ~~ops1, ~~ops2)...))
+    @rule(~o1::_isa(Symbolic{AbstractOperator}) * ⊗(~~ops) => ⊗(map(op -> ~o1 * op, ~~ops)...)),
+    @rule(⊗(~~ops) * ~o1::_isa(Symbolic{AbstractOperator}) => ⊗(map(op -> op * ~o1, ~~ops)...)),
 ]
 
 # 

test/runtests.jl

@@ -35,6 +35,7 @@ println("Starting tests with $(Threads.nthreads()) threads out of `Sys.CPU_THREA
 @doset "anticommutator"
 @doset "dagger"
 @doset "zero_obj"
+@doset "trace"
 @doset "expand"
 
 VERSION >= v"1.9" && @doset "doctests"

test/test_expand.jl

@@ -26,6 +26,9 @@ using Test
     @test isequal(qexpand((B+C+D)*A), B*A + C*A + D*A)
     @test isequal(qexpand(commutator(A, B) * C), A*B*C - B*A*C)
 
+    @test isequal(qexpand(A*(B⊗C⊗D)), (A*B)⊗(A*C)⊗(A*D))
+    @test isequal(qexpand((B⊗C⊗D)*A), (B*A)⊗(C*A)⊗(D*A))
+
     @test isequal(qexpand((A⊗B)*(C⊗D)), (A*C)⊗(B*D))
     @test isequal(qexpand((b₁⊗b₂)*(k₁⊗k₂)), (b₁*k₁)*(b₂*k₂))
 end

test/test_superop.jl

@@ -44,6 +44,19 @@ noisy_pair_fromdm = (noiseop ⊗ noiseop) * pure_pair_dm
 @test tr(express(noisy_pair)) ≈ tr(express(noisy_pair_fromdm)) ≈ 1
 @test express(noisy_pair) ≈ express(noisy_pair_fromdm)
 
+@op A; @op B; @op C; @op O; @ket k;
+@superop S; K = kraus(S, A, B, C);
+
+
+
+@testset "symbolic superoperator tests" begin
+    @test isequal(S*SZeroOperator(), SZeroOperator())
+    @test isequal(S*SZeroKet(), SZeroOperator())
+    @test isequal(S*k, S*projector(k))
+    @test isequal(K*O, A*O*dagger(A) + B*O*dagger(B) + C*O*dagger(C))
+    @test isequal(K*k, A*projector(k)*dagger(A) + B*projector(k)*dagger(B) + C*projector(k)*dagger(C))
+end
+
 # TODO
 # test against depolarization
 # Depolarization over two qubits is different from depolarizing each separately (see related tutorial)