trace, partial trace, and superoperators

docs/src/index.md

@@ -202,4 +202,4 @@ express(MixedState(X1)/2+SProjector(Z1)/2, CliffordRepr())
 
 !!! warning "Stabilizer state expressions"
 
-    The state written as $\frac{|Z₁⟩⊗|Z₁⟩+|Z₂⟩⊗|Z₂⟩}{√2}$ is a well known stabilizer state, namely a Bell state. However, automatically expressing it as a stabilizer is a prohibitively expensive computational operation in general. We do not perform that computation automatically. If you want to ensure that states you define can be automatically converted to tableaux for Clifford simulations, avoid using sumation of kets. On the other hand, in all of our Clifford Monte-Carlo simulations, `⊗` is fully supported, as well as [`SProjector`](@ref), [`MixedState`](@ref), [`StabilizerState`](@ref), and sumation of density matrices.
+    The state written as $\frac{|Z₁⟩⊗|Z₁⟩+|Z₂⟩⊗|Z₂⟩}{√2}$ is a well known stabilizer state, namely a Bell state. However, automatically expressing it as a stabilizer is a prohibitively expensive computational operation in general. We do not perform that computation automatically. If you want to ensure that states you define can be automatically converted to tableaux for Clifford simulations, avoid using summation of kets. On the other hand, in all of our Clifford Monte-Carlo simulations, `⊗` is fully supported, as well as [`SProjector`](@ref), [`MixedState`](@ref), [`StabilizerState`](@ref), and summation of density matrices.

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,samebases,IncompatibleBases,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,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,
+       N,n̂,Create,âꜛ,Destroy,â,basis,SpinBasis,FockBasis,
+       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

@@ -139,15 +139,19 @@ children(x::SMulOperator) = [:*;x.terms]
 function Base.:(*)(xs::Symbolic{AbstractOperator}...) 
     zero_ind = findfirst(x->iszero(x), xs)
     if isnothing(zero_ind)
-        terms = flattenop(*, collect(xs))
-        coeff, cleanterms = prefactorscalings(terms)
-        coeff * SMulOperator(cleanterms)
+        if any(x->!(samebases(basis(x),basis(first(xs)))),xs)
+            throw(IncompatibleBases())
+        else
+            terms = flattenop(*, collect(xs))
+            coeff, cleanterms = prefactorscalings(terms)
+            coeff * SMulOperator(cleanterms)
+        end
     else
         SZeroOperator()
     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)
 
@@ -216,8 +220,12 @@ operation(x::SCommutator) = commutator
 head(x::SCommutator) = :commutator
 children(x::SCommutator) = [:commutator, x.op1, x.op2]
 function commutator(o1::Symbolic{AbstractOperator}, o2::Symbolic{AbstractOperator})
-    coeff, cleanterms = prefactorscalings([o1 o2])
-    cleanterms[1] === cleanterms[2] ? SZeroOperator() : coeff * SCommutator(cleanterms...)   
+    if !(samebases(basis(o1),basis(o2)))
+        throw(IncompatibleBases())
+    else
+        coeff, cleanterms = prefactorscalings([o1 o2])
+        cleanterms[1] === cleanterms[2] ? SZeroOperator() : coeff * SCommutator(cleanterms...)  
+    end 
 end
 commutator(o1::SZeroOperator, o2::Symbolic{AbstractOperator}) = SZeroOperator()
 commutator(o1::Symbolic{AbstractOperator}, o2::SZeroOperator) = SZeroOperator()
@@ -245,8 +253,12 @@ operation(x::SAnticommutator) = anticommutator
 head(x::SAnticommutator) = :anticommutator
 children(x::SAnticommutator) = [:anticommutator, x.op1, x.op2]
 function anticommutator(o1::Symbolic{AbstractOperator}, o2::Symbolic{AbstractOperator})
-    coeff, cleanterms = prefactorscalings([o1 o2])
-    coeff * SAnticommutator(cleanterms...)
+    if !(samebases(basis(o1),basis(o2)))
+        throw(IncompatibleBases())
+    else
+        coeff, cleanterms = prefactorscalings([o1 o2])
+        coeff * SAnticommutator(cleanterms...)
+    end
 end
 anticommutator(o1::SZeroOperator, o2::Symbolic{AbstractOperator}) = SZeroOperator()
 anticommutator(o1::Symbolic{AbstractOperator}, o2::SZeroOperator) = SZeroOperator()

src/QSymbolicsBase/basic_ops_inhomogeneous.jl

@@ -14,18 +14,21 @@ A|k⟩
 @withmetadata struct SApplyKet <: Symbolic{AbstractKet}
     op
     ket
-    function SApplyKet(o, k)
-        coeff, cleanterms = prefactorscalings([o k])
-        coeff*new(cleanterms...)
-    end
 end
 isexpr(::SApplyKet) = true
 iscall(::SApplyKet) = true
 arguments(x::SApplyKet) = [x.op,x.ket]
 operation(x::SApplyKet) = *
 head(x::SApplyKet) = :*
 children(x::SApplyKet) = [:*,x.op,x.ket]
-Base.:(*)(op::Symbolic{AbstractOperator}, k::Symbolic{AbstractKet}) = SApplyKet(op,k)
+function Base.:(*)(op::Symbolic{AbstractOperator}, k::Symbolic{AbstractKet})
+    if !(samebases(basis(op),basis(k)))
+        throw(IncompatibleBases())
+    else
+        coeff, cleanterms = prefactorscalings([op k])
+        coeff*SApplyKet(cleanterms...)
+    end
+end
 Base.:(*)(op::SZeroOperator, k::Symbolic{AbstractKet}) = SZeroKet()
 Base.:(*)(op::Symbolic{AbstractOperator}, k::SZeroKet) = SZeroKet()
 Base.:(*)(op::SZeroOperator, k::SZeroKet) = SZeroKet()
@@ -44,18 +47,21 @@ julia> b*A
 @withmetadata struct SApplyBra <: Symbolic{AbstractBra}
     bra
     op
-    function SApplyBra(b, o)
-        coeff, cleanterms = prefactorscalings([b o])
-        coeff*new(cleanterms...)
-    end
 end
 isexpr(::SApplyBra) = true
 iscall(::SApplyBra) = true
 arguments(x::SApplyBra) = [x.bra,x.op]
 operation(x::SApplyBra) = *
 head(x::SApplyBra) = :*
 children(x::SApplyBra) = [:*,x.bra,x.op]
-Base.:(*)(b::Symbolic{AbstractBra}, op::Symbolic{AbstractOperator}) = SApplyBra(b,op)
+function Base.:(*)(b::Symbolic{AbstractBra}, op::Symbolic{AbstractOperator}) 
+    if !(samebases(basis(b),basis(op)))
+        throw(IncompatibleBases())
+    else
+        coeff, cleanterms = prefactorscalings([b op])
+        coeff*SApplyBra(cleanterms...)
+    end
+end
 Base.:(*)(b::SZeroBra, op::Symbolic{AbstractOperator}) = SZeroBra()
 Base.:(*)(b::Symbolic{AbstractBra}, op::SZeroOperator) = SZeroBra()
 Base.:(*)(b::SZeroBra, op::SZeroOperator) = SZeroBra()
@@ -81,33 +87,22 @@ arguments(x::SBraKet) = [x.bra,x.ket]
 operation(x::SBraKet) = *
 head(x::SBraKet) = :*
 children(x::SBraKet) = [:*,x.bra,x.ket]
-Base.:(*)(b::Symbolic{AbstractBra}, k::Symbolic{AbstractKet}) = SBraKet(b,k)
+function Base.:(*)(b::Symbolic{AbstractBra}, k::Symbolic{AbstractKet}) 
+    if !(samebases(basis(b),basis(k)))
+        throw(IncompatibleBases())
+    else
+        coeff, cleanterms = prefactorscalings([b k])
+        coeff == 1 ? SBraKet(cleanterms...) : coeff*SBraKet(cleanterms...)
+    end
+end
 Base.:(*)(b::SZeroBra, k::Symbolic{AbstractKet}) = 0
 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)
-maketerm(::Type{<:SBraKet}, f, a, t, m) = f(a...)
+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;
@@ -119,18 +114,21 @@ julia> k*b
 @withmetadata struct SOuterKetBra <: Symbolic{AbstractOperator}
     ket
     bra
-    function SOuterKetBra(k, b)
-        coeff, cleanterms = prefactorscalings([k b])
-        coeff*new(cleanterms...)
-    end
 end
 isexpr(::SOuterKetBra) = true
 iscall(::SOuterKetBra) = true
 arguments(x::SOuterKetBra) = [x.ket,x.bra]
 operation(x::SOuterKetBra) = *
 head(x::SOuterKetBra) = :*
 children(x::SOuterKetBra) = [:*,x.ket,x.bra]
-Base.:(*)(k::Symbolic{AbstractKet}, b::Symbolic{AbstractBra}) = SOuterKetBra(k,b)
+function Base.:(*)(k::Symbolic{AbstractKet}, b::Symbolic{AbstractBra})
+    if !(samebases(basis(k),basis(b)))
+        throw(IncompatibleBases())
+    else
+        coeff, cleanterms = prefactorscalings([k b])
+        coeff*SOuterKetBra(cleanterms...)
+    end
+end
 Base.:(*)(k::SZeroKet, b::Symbolic{AbstractBra}) = SZeroOperator()
 Base.:(*)(k::Symbolic{AbstractKet}, b::SZeroBra) = SZeroOperator()
 Base.:(*)(k::SZeroKet, b::SZeroBra) = SZeroOperator()

src/QSymbolicsBase/basic_superops.jl

@@ -0,0 +1,61 @@
+##
+# Superoperator representations
+##
+
+"""Kraus representation of a quantum channel
+
+```jldoctest
+julia> @op A₁; @op A₂; @op A₃;
+
+julia> K = kraus(A₁, A₂, A₃)
+𝒦(A₁,A₂,A₃)
+
+julia> @op ρ;
+
+julia> K*ρ
+(A₁ρA₁†+A₂ρA₂†+A₃ρA₃†)
+```
+"""
+@withmetadata struct KrausRepr <: Symbolic{AbstractSuperOperator}
+    krausops
+end
+isexpr(::KrausRepr) = true
+iscall(::KrausRepr) = true
+arguments(x::KrausRepr) = x.krausops
+operation(x::KrausRepr) = kraus
+head(x::KrausRepr) = :kraus
+children(x::KrausRepr) = [:kraus; x.krausops]
+kraus(xs::Symbolic{AbstractOperator}...) = KrausRepr(collect(xs))
+basis(x::KrausRepr) = basis(first(x.krausops))
+Base.show(io::IO, x::KrausRepr) = print(io, "𝒦("*join([symbollabel(i) for i in x.krausops], ",")*")")
+
+##
+# 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)