fix basis consistency, ptrace, and superops

src/QSymbolicsBase/QSymbolicsBase.jl

@@ -11,7 +11,7 @@ import LinearAlgebra: eigvecs,ishermitian,inv
 import QuantumInterface:
     apply!,
     tensor, ⊗,
-    basis,Basis,SpinBasis,FockBasis,CompositeBasis,
+    basis,Basis,samebases,IncompatibleBases,SpinBasis,FockBasis,CompositeBasis,
     nqubits,
     projector,dagger,tr,ptrace,
     AbstractBra,AbstractKet,AbstractOperator,AbstractSuperOperator
@@ -23,12 +23,12 @@ export SymQObj,QObj,
        apply!,
        express,
        tensor,⊗,
-       dagger,projector,commutator,anticommutator,expand,tr,ptrace,
+       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,
+       N,n̂,Create,âꜛ,Destroy,â,basis,SpinBasis,FockBasis,
        SBra,SKet,SOperator,SHermitianOperator,SUnitaryOperator,SHermitianUnitaryOperator,SSuperOperator,
        @ket,@bra,@op,@superop,
        SAdd,SAddBra,SAddKet,SAddOperator,

src/QSymbolicsBase/basic_ops_homogeneous.jl

@@ -139,9 +139,13 @@ 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
@@ -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,13 +87,20 @@ 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 outer product of a ket and a bra
@@ -101,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

@@ -5,11 +5,10 @@
 """Kraus representation of a quantum channel
 
 ```jldoctest
-julia> @superop ℰ;
-
 julia> @op A₁; @op A₂; @op A₃;
 
-julia> K = kraus(ℰ, A₁, A₂, A₃);
+julia> K = kraus(A₁, A₂, A₃)
+𝒦(A₁,A₂,A₃)
 
 julia> @op ρ;
 
@@ -18,19 +17,17 @@ julia> K*ρ
 ```
 """
 @withmetadata struct KrausRepr <: Symbolic{AbstractSuperOperator}
-    sop
     krausops
 end
 isexpr(::KrausRepr) = true
 iscall(::KrausRepr) = true
-arguments(x::KrausRepr) = [x.sop, x.krausops]
+arguments(x::KrausRepr) = 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))
+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

src/QSymbolicsBase/linalg.jl

@@ -184,6 +184,8 @@ function ptrace(x::Symbolic{AbstractOperator}, s)
     if isa(ex, typeof(x))
         if isa(basis(x), CompositeBasis)
             SPartialTrace(x, s)
+        elseif s==1
+            tr(x)
         else
             throw(ArgumentError("cannot take partial trace of a single quantum system"))
         end
@@ -192,19 +194,37 @@ function ptrace(x::Symbolic{AbstractOperator}, 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"))
+    if isa(basis(x), CompositeBasis)
+        for i in arguments(x)
+            if isexpr(i)
+                if isa(i, SScaledOperator) && operation(i.obj) === ⊗
+                    prod_terms = arguments(i.obj)
+                    coeff = i.coeff
+                elseif operation(i) === ⊗
+                    prod_terms = arguments(i)
+                    coeff = 1
+                else
+                    return SPartialTrace(x,s)
+                end
+            else
+                return SPartialTrace(x,s)
+            end
+            if any(j -> isa(basis(j), CompositeBasis), prod_terms)
+                return SPartialTrace(x,s)
+            else
+                sys_op = coeff*prod_terms[s]
+                new_terms = deleteat!(copy(prod_terms), s)
+                trace = tr(sys_op)
+                isone(length(new_terms)) ? push!(add_terms, trace*first(new_terms)) : push!(add_terms, trace*(⊗)(new_terms...))
+            end  
         end
+        (+)(add_terms...)
+    elseif s==1
+        tr(x)
+    else
+        throw(ArgumentError("cannot take partial trace of a single quantum system"))
     end
-    (+)(add_terms...)
 end
 function ptrace(x::STensorOperator, s)
     ex = qexpand(x)
@@ -213,8 +233,8 @@ function ptrace(x::STensorOperator, s)
     else
         terms = arguments(ex)
         newterms = []
-        if isa(basis(terms[s]), CompositeBasis)
-            SPartial(ex, s)
+        if any(i -> isa(basis(i), CompositeBasis), terms)
+            SPartialTrace(ex, s)
         else 
             sys_op = terms[s]
             new_terms = deleteat!(copy(terms), s)

src/QSymbolicsBase/predefined.jl

@@ -179,7 +179,7 @@ const CPHASE = CPHASEGate()
 
 abstract type AbstractSingleBosonOp <: Symbolic{AbstractOperator} end
 abstract type AbstractSingleBosonGate <: AbstractSingleBosonOp end # TODO maybe an IsUnitaryTrait is a better choice
-isexpr(::AbstractSingleBosonGate) = false
+isexpr(::AbstractSingleBosonOp) = false
 basis(::AbstractSingleBosonGate) = inf_fock_basis
 
 @withmetadata struct NumberOp <: AbstractSingleBosonOp end
@@ -188,6 +188,7 @@ symbollabel(::NumberOp) = "n"
 symbollabel(::CreateOp) = "a†"
 @withmetadata struct DestroyOp <: AbstractSingleBosonOp end
 symbollabel(::DestroyOp) = "a"
+basis(::Union{NumberOp, CreateOp, DestroyOp}) = inf_fock_basis
 
 """Number operator, also available as the constant `n̂`"""
 const N = const n̂ = NumberOp()

src/QSymbolicsBase/utils.jl

@@ -2,7 +2,7 @@ function prefactorscalings(xs)
     terms = []
     coeff = 1::Any
     for x in xs
-        if isa(x, SScaledOperator)
+        if isa(x, SScaled)
             coeff *= x.coeff
             push!(terms, x.obj)
         elseif isa(x, Union{Number, Symbolic{Number}})

test/test_basis_consistency.jl

@@ -1,5 +1,7 @@
 using Test
 using QuantumSymbolics
+using QuantumOptics
+using QuantumInterface: IncompatibleBases
 
 @test express(Z*Z1) == express(Z1)
 @test express(Z*Z2) == -express(Z2)
@@ -11,3 +13,14 @@ using QuantumSymbolics
 @test express(Pp*Z2) == express(Z1)
 @test express(Pm*L0) == express(L1)
 @test express(Pp*L1) == express(L0)
+
+@op A SpinBasis(1//2) ⊗ SpinBasis(1//2); @op B; 
+@ket k; @bra b; @ket l SpinBasis(1//2) ⊗ SpinBasis(1//2);
+
+@test_throws IncompatibleBases A*B
+@test_throws IncompatibleBases commutator(A, B)
+@test_throws IncompatibleBases anticommutator(A, B)
+@test_throws IncompatibleBases A*k
+@test_throws IncompatibleBases b*A
+@test_throws IncompatibleBases l*b
+@test_throws IncompatibleBases b*l

test/test_expand.jl

@@ -26,9 +26,6 @@ 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

@@ -45,7 +45,7 @@ noisy_pair_fromdm = (noiseop ⊗ noiseop) * pure_pair_dm
 @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);
+@superop S; K = kraus(A, B, C);