Create ChoiState type and conversions to/from SuperOperator

Author: akirakyle

src/QuantumOpticsBase.jl

@@ -36,8 +36,8 @@ export Basis, GenericBasis, CompositeBasis, basis,
                 current_time, time_shift, time_stretch, time_restrict, static_operator,
         #superoperators
                 SuperOperator, DenseSuperOperator, DenseSuperOpType,
-                SparseSuperOperator, SparseSuperOpType, spre, spost, sprepost, liouvillian,
-                identitysuperoperator,
+                SparseSuperOperator, SparseSuperOpType, ChoiState, 
+                spre, spost, sprepost, liouvillian, identitysuperoperator,
         #fock
                 FockBasis, number, destroy, create,
                 fockstate, coherentstate, coherentstate!,

src/superoperators.jl

@@ -11,8 +11,9 @@ mutable struct SuperOperator{B1,B2,T} <: AbstractSuperOperator{B1,B2}
     basis_r::B2
     data::T
     function SuperOperator{BL,BR,T}(basis_l::BL, basis_r::BR, data::T) where {BL,BR,T}
-        if length(basis_l[1])*length(basis_l[2]) != size(data, 1) ||
-           length(basis_r[1])*length(basis_r[2]) != size(data, 2)
+        if (length(basis_l) != 2 || length(basis_r) != 2 ||
+            length(basis_l[1])*length(basis_l[2]) != size(data, 1) ||
+            length(basis_r[1])*length(basis_r[2]) != size(data, 2))
             throw(DimensionMismatch("Tried to assign data of size $(size(data)) to Hilbert spaces of sizes $(length.(basis_l)), $(length.(basis_r))"))
         end
         new(basis_l, basis_r, data)
@@ -316,3 +317,63 @@ end
         }
     throw(IncompatibleBases())
 end
+
+"""
+    ChoiState <: AbstractSuperOperator
+
+Superoperator represented as a choi state.
+"""
+mutable struct ChoiState{B1,B2,T} <: AbstractSuperOperator{B1,B2}
+    basis_l::B1
+    basis_r::B2
+    data::T
+    function ChoiState{BL,BR,T}(basis_l::BL, basis_r::BR, data::T) where {BL,BR,T}
+        if (length(basis_l) != 2 || length(basis_r) != 2 ||
+            length(basis_l[1])*length(basis_l[2]) != size(data, 1) ||
+            length(basis_r[1])*length(basis_r[2]) != size(data, 2))
+            throw(DimensionMismatch("Tried to assign data of size $(size(data)) to Hilbert spaces of sizes $(length.(basis_l)), $(length.(basis_r))"))
+        end
+        new(basis_l, basis_r, data)
+    end
+end
+ChoiState{BL,BR}(b1::BL, b2::BR, data::T) where {BL,BR,T} = ChoiState{BL,BR,T}(b1, b2, data)
+ChoiState(b1::BL, b2::BR, data::T) where {BL,BR,T} = ChoiState{BL,BR,T}(b1, b2, data)
+ChoiState(b,data) = ChoiState(b, b, data)
+
+dense(a::ChoiState) = ChoiState(a.basis_l, a.basis_r, Matrix(a.data))
+sparse(a::ChoiState) = ChoiStae(a.basis_l, a.basis_r, sparse(a.data))
+dagger(a::ChoiState) = ChoiState(dagger(SuperOperator(a)))
+*(a::ChoiState, b::ChoiState) = ChoiState(SuperOperator(a)*SuperOperator(b))
+*(a::ChoiState, b::Operator) = SuperOperator(a)*b
+==(a::ChoiState, b::ChoiState) = SuperOperator(a) == SuperOperator(b)
+
+# reshape swaps within systems due to colum major ordering
+# https://docs.qojulia.org/quantumobjects/operators/#tensor_order
+function _super_choi((l1, l2), (r1, r2), data)
+    data = reshape(data, map(length, (l2, l1, r2, r1)))
+    (l1, l2), (r1, r2) = (r2, l2), (r1, l1)
+    data = permutedims(data, (1, 3, 2, 4))
+    data = reshape(data, map(length, (l1⊗l2, r1⊗r2)))
+    return (l1, l2), (r1, r2), data
+end
+
+function _super_choi((r2, l2), (r1, l1), data::SparseMatrixCSC)
+    data = _permutedims(data, map(length, (l2, r2, l1, r1)), (1, 3, 2, 4))
+    data = reshape(data, map(length, (l1⊗l2, r1⊗r2)))
+    # sparse(data) is necessary since reshape of a sparse array returns a
+    # ReshapedSparseArray which is not a subtype of AbstractArray and so
+    # _permutedims fails to acces the ".m" field
+    # https://github.com/qojulia/QuantumOpticsBase.jl/pull/83
+    # https://github.com/JuliaSparse/SparseArrays.jl/issues/24
+    # permutedims in SparseArrays.jl only implements perm (2,1) and so
+    # _permutedims should probably be upstreamed
+    # https://github.com/JuliaLang/julia/issues/26534
+    return (l1, l2), (r1, r2), sparse(data)
+end
+
+ChoiState(op::SuperOperator) = ChoiState(_super_choi(op.basis_l, op.basis_r, op.data)...)
+SuperOperator(op::ChoiState) = SuperOperator(_super_choi(op.basis_l, op.basis_r, op.data)...)
+
+*(a::ChoiState, b::SuperOperator) = SuperOperator(a)*b
+*(a::SuperOperator, b::ChoiState) = a*SuperOperator(b)
+

test/test_superoperators.jl

@@ -153,29 +153,43 @@ J = [Ja, Jc]
 ρ₀ = dm(Ψ₀)
 
 @test identitysuperoperator(spinbasis)*sx == sx
+@test ChoiState(identitysuperoperator(spinbasis))*sx == sx
 @test identitysuperoperator(sparse(spre(sx)))*sx == sparse(sx)
+@test ChoiState(identitysuperoperator(sparse(spre(sx))))*sx == sparse(sx)
 @test identitysuperoperator(dense(spre(sx)))*sx == dense(sx)
+@test ChoiState(identitysuperoperator(dense(spre(sx))))*sx == dense(sx)
 
 op1 = DenseOperator(spinbasis, [1.2+0.3im 0.7+1.2im;0.3+0.1im 0.8+3.2im])
 op2 = DenseOperator(spinbasis, [0.2+0.1im 0.1+2.3im; 0.8+4.0im 0.3+1.4im])
 @test tracedistance(spre(op1)*op2, op1*op2) < 1e-12
+@test tracedistance(ChoiState(spre(op1))*op2, op1*op2) < 1e-12
 @test tracedistance(spost(op1)*op2, op2*op1) < 1e-12
+@test tracedistance(ChoiState(spost(op1))*op2, op2*op1) < 1e-12
 
 @test spre(sparse(op1))*op2 == op1*op2
+@test ChoiState(spre(sparse(op1)))*op2 == op1*op2
 @test spost(sparse(op1))*op2 == op2*op1
+@test ChoiState(spost(sparse(op1)))*op2 == op2*op1
 @test spre(sparse(dagger(op1)))*op2 == dagger(op1)*op2
+@test ChoiState(spre(sparse(dagger(op1))))*op2 == dagger(op1)*op2
 @test spre(dense(dagger(op1)))*op2 ≈ dagger(op1)*op2
+@test ChoiState(spre(dense(dagger(op1))))*op2 ≈ dagger(op1)*op2
 @test sprepost(sparse(op1), op1)*op2 ≈ op1*op2*op1
+@test ChoiState(sprepost(sparse(op1), op1))*op2 ≈ op1*op2*op1
 
 @test spre(sparse(op1))*sparse(op2) == sparse(op1*op2)
+@test ChoiState(spre(sparse(op1)))*sparse(op2) == sparse(op1*op2)
 @test spost(sparse(op1))*sparse(op2) == sparse(op2*op1)
+@test ChoiState(spost(sparse(op1)))*sparse(op2) == sparse(op2*op1)
 @test sprepost(sparse(op1), sparse(op1))*sparse(op2) ≈ sparse(op1*op2*op1)
+@test ChoiState(sprepost(sparse(op1), sparse(op1)))*sparse(op2) ≈ sparse(op1*op2*op1)
 
 @test sprepost(op1, op2) ≈ spre(op1)*spost(op2)
 b1 = FockBasis(1)
 b2 = FockBasis(5)
 op = fockstate(b1, 0) ⊗ dagger(fockstate(b2, 0))
 @test sprepost(dagger(op), op)*dm(fockstate(b1, 0)) == dm(fockstate(b2, 0))
+@test ChoiState(sprepost(dagger(op), op))*dm(fockstate(b1, 0)) == dm(fockstate(b2, 0))
 @test_throws ArgumentError spre(op)
 @test_throws ArgumentError spost(op)
 
@@ -185,12 +199,14 @@ for j=J
     ρ .+= j*ρ₀*dagger(j) - 0.5*dagger(j)*j*ρ₀ - 0.5*ρ₀*dagger(j)*j
 end
 @test tracedistance(L*ρ₀, ρ) < 1e-10
+@test tracedistance(ChoiState(L)*ρ₀, ρ) < 1e-10
 
 # tout, ρt = timeevolution.master([0.,1.], ρ₀, H, J; reltol=1e-7)
 # @test tracedistance(exp(dense(L))*ρ₀, ρt[end]) < 1e-6
 # @test tracedistance(exp(sparse(L))*ρ₀, ρt[end]) < 1e-6
 
 @test dense(spre(op1)) == spre(op1)
+@test dense(ChoiState(spre(op1))) == ChoiState(spre(op1))
 
 @test L/2.0 == 0.5*L == L*0.5
 @test -L == SparseSuperOperator(L.basis_l, L.basis_r, -L.data)