trace, partial trace, and superoperators

src/QSymbolicsBase/QSymbolicsBase.jl

@@ -1,20 +1,20 @@
 using Symbolics
 import Symbolics: simplify
 using SymbolicUtils
-import SymbolicUtils: Symbolic, _isone, flatten_term, isnotflat, Chain, Fixpoint
+import SymbolicUtils: Symbolic,_isone,flatten_term,isnotflat,Chain,Fixpoint,sorted_arguments
 using TermInterface
-import TermInterface: isexpr, head, iscall, children, operation, arguments, metadata
+import TermInterface: isexpr,head,iscall,children,operation,arguments,metadata
 
 using LinearAlgebra
-import LinearAlgebra: eigvecs, ishermitian, inv
+import LinearAlgebra: eigvecs,ishermitian,inv
 
 import QuantumInterface:
     apply!,
     tensor, ⊗,
-    basis, Basis, SpinBasis, FockBasis,
+    basis,Basis,SpinBasis,FockBasis,
     nqubits,
-    projector, dagger,
-    AbstractKet, AbstractOperator, AbstractSuperOperator, AbstractBra
+    projector,dagger,tr,ptrace,
+    AbstractBra,AbstractKet,AbstractOperator,AbstractSuperOperator
 
 export SymQObj,QObj,
        AbstractRepresentation,AbstractUse,
@@ -23,26 +23,28 @@ export SymQObj,QObj,
        apply!,
        express,
        tensor,⊗,
-       dagger,projector,commutator,anticommutator,expand,
+       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,SHermitianOperator,SUnitaryOperator,SHermitianUnitaryOperator,
+       SProjector,MixedState,IdentityOp,SInvOperator,
        SApplyKet,SApplyBra,SMulOperator,SSuperOpApply,SCommutator,SAnticommutator,SDagger,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_flatten,qsimplify_commutator,qsimplify_anticommutator,
-       isunitary
+       isunitary,
+       KrausRepr,
+       kraus
 
 function countmap(samples) # A simpler version of StatsBase.countmap, because StatsBase is slow to import
     counts = Dict{Any,Any}()
@@ -171,10 +173,10 @@ propsequal(x,y) = all(n->isequal(getproperty(x,n),getproperty(y,n)), propertynam
 ##
 
 include("literal_objects.jl")
+include("predefined_CPTP.jl")
 include("basic_ops_homogeneous.jl")
 include("basic_ops_inhomogeneous.jl")
 include("predefined.jl")
-include("predefined_CPTP.jl")
 
 ##
 # Symbolic and simplification rules

src/QSymbolicsBase/basic_ops_homogeneous.jl

@@ -45,7 +45,7 @@
 end
 const SScaledOperator = SScaled{AbstractOperator}
 function Base.show(io::IO, x::SScaledOperator)
-    if x.coeff isa Number
+    if x.coeff isa Real
         print(io, "$(x.coeff)$(x.obj)")
     else
         print(io, "($(x.coeff))$(x.obj)")
@@ -93,6 +93,11 @@
 Base.:(+)(xs::Vararg{Symbolic{<:QObj},0}) = 0 # to avoid undefined type parameters issue in the above method
 basis(x::SAdd) = basis(first(x.dict).first)
 
+const SAddBra = SAdd{AbstractBra}
+function Base.show(io::IO, x::SAddBra)
+    ordered_terms = sort([repr(i) for i in arguments(x)])
+    print(io, "("*join(ordered_terms,"+")::String*")") # type assert to help inference
+end
 const SAddKet = SAdd{AbstractKet}
 function Base.show(io::IO, x::SAddKet)
     ordered_terms = sort([repr(i) for i in arguments(x)])
@@ -103,11 +108,6 @@
     ordered_terms = sort([repr(i) for i in arguments(x)])
     print(io, "("*join(ordered_terms,"+")::String*")") # type assert to help inference
 end
-const SAddBra = SAdd{AbstractBra}
-function Base.show(io::IO, x::SAddBra)
-    ordered_terms = sort([repr(i) for i in arguments(x)])
-    print(io, "("*join(ordered_terms,"+")::String*")") # type assert to help inference
-end
 
 """Symbolic application of operator on operator
 
@@ -121,7 +121,7 @@
 @withmetadata struct SMulOperator <: Symbolic{AbstractOperator}
     terms
     function SMulOperator(terms)
-        coeff, cleanterms = prefactorscalings(terms)
+        coeff, cleanterms = prefactorscalings(terms, scalar=true)
         coeff*new(cleanterms)
     end
 end
@@ -149,7 +149,7 @@
 julia> @op A; @op B;
 
 julia> A ⊗ B 
-A⊗B
+(A⊗B)
 ```
 """
 @withmetadata struct STensor{T<:QObj} <: Symbolic{T}
@@ -164,21 +164,21 @@
 arguments(x::STensor) = x.terms
 operation(x::STensor) = ⊗
 head(x::STensor) = :⊗
-children(x::STensor) = pushfirst!(x.terms,:⊗)
+children(x::STensor) = [:⊗; x.terms]
 function ⊗(xs::Symbolic{T}...) where {T<:QObj}
     zero_ind = findfirst(x->iszero(x), xs)
     isnothing(zero_ind) ? STensor{T}(collect(xs)) : SZero{T}()
 end
 basis(x::STensor) = tensor(basis.(x.terms)...)
 
+const STensorBra = STensor{AbstractBra}
+Base.show(io::IO, x::STensorBra) = print(io, join(map(string, arguments(x)),""))
 const STensorKet = STensor{AbstractKet}
 Base.show(io::IO, x::STensorKet) = print(io, join(map(string, arguments(x)),""))
 const STensorOperator = STensor{AbstractOperator}
-Base.show(io::IO, x::STensorOperator) = print(io, join(map(string, arguments(x)),"⊗"))
+Base.show(io::IO, x::STensorOperator) = print(io, "("*join(map(string, arguments(x)),"⊗")*")")
 const STensorSuperOperator = STensor{AbstractSuperOperator}
-Base.show(io::IO, x::STensorSuperOperator) = print(io, join(map(string, arguments(x)),"⊗"))
-const STensorBra = STensor{AbstractBra}
-Base.show(io::IO, x::STensorBra) = print(io, join(map(string, arguments(x)),""))
+Base.show(io::IO, x::STensorSuperOperator) = print(io, "("*join(map(string, arguments(x)),"⊗")*")")
 
 """Symbolic commutator of two operators
 
@@ -188,6 +188,9 @@
 julia> commutator(A, B)
 [A,B]
 
+julia> expand(commutator(A, B))
+(-1BA+AB)
+
 julia> commutator(A, A)
 𝟎
 ```
@@ -221,6 +224,9 @@
 
 julia> anticommutator(A, B)
 {A,B}
+
+julia> expand(anticommutator(A, B))
+(AB+BA)
 ```
 """
 @withmetadata struct SAnticommutator <: Symbolic{AbstractOperator}

src/QSymbolicsBase/basic_ops_inhomogeneous.jl

@@ -86,9 +86,17 @@
 Base.:(*)(b::Symbolic{AbstractBra}, k::SZeroKet) = 0
 Base.:(*)(b::SZeroBra, k::SZeroKet) = 0
 Base.show(io::IO, x::SBraKet) = begin print(io,x.bra); print(io,x.ket) end
+Base.hash(x::SBraKet, h::UInt) = hash((head(x), arguments(x)), h)
 Base.isequal(x::SBraKet, y::SBraKet) = isequal(x.bra, y.bra) && isequal(x.ket, y.ket)
 
-"""Symbolic application of a superoperator on an operator"""
+"""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
@@ -102,8 +110,10 @@
 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
+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)
 
 """Symbolic outer product of a ket and a bra

src/QSymbolicsBase/literal_objects.jl

@@ -67,15 +67,29 @@
 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
@@ -92,4 +106,4 @@
 basis(x::SZero) = nothing
 
 Base.show(io::IO, x::SZero) = print(io, symbollabel(x))
-Base.iszero(x::SymQObj) = isa(x, SZero)
+Base.iszero(x::Union{SymQObj, Symbolic{Number}, Symbolic{Complex}}) = isa(x, SZero)

src/QSymbolicsBase/predefined.jl

@@ -286,7 +286,109 @@
     print(io,x.obj)
     print(io,"†")
 end
-symbollabel(x::SDagger) = symbollabel(x.obj)
+
+"""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 A; @op B;
+
+julia> ptrace(A, 1)
+tr1(A)
+
+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{Complex}
+    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]
+Base.show(io::IO, x::SPartialTrace) = print(io, "tr$(x.sys)($(x.obj))")
+ptrace(x::Symbolic{AbstractOperator}, s) = SPartialTrace(x, s)
+function ptrace(x::STensorOperator, s)
+    terms = arguments(x)
+    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
+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
 
 """Inverse Operator
 
@@ -368,4 +470,4 @@
 isunitary(::IdentityOp) = true
 
 """Identity operator in qubit basis"""
-const I = IdentityOp(qubit_basis)
+const I = IdentityOp(qubit_basis)   

src/QSymbolicsBase/predefined_CPTP.jl

@@ -47,3 +47,37 @@
     print(io, x.gate)
     print(io, "]")
 end
+
+##
+# Representations of CPTP maps
+##
+
+"""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))