TensorCast ftw

Author: akirakyle

Project.toml

@@ -15,6 +15,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
+TensorCast = "02d47bb6-7ce6-556a-be16-bb1710789e2b"
 UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
 [compat]
@@ -30,5 +31,6 @@ Random = "1"
 RecursiveArrayTools = "3"
 SparseArrays = "1"
 Strided = "1, 2"
+TensorCast = "0.4"
 UnsafeArrays = "1"
 julia = "1.10"

src/superoperators.jl

@@ -1,4 +1,5 @@
 import QuantumInterface: KetBraBasis, ChoiBasis
+using TensorCast
 
 const SuperKetType = Ket{<:KetBraBasis}
 
@@ -17,45 +18,39 @@ const PauliTransferType = Operator{<:ChoiBasis,<:ChoiBasis}
 #const ChoiStateType = Operator{ChoiBasisType,ChoiBasisType}
 #const ChoiStateType = Union{Operator{CompositeBasis{ChoiBasis,S},CompositeBasis{ChoiBasis,T}} where {S,T}, Operator{CompositeBasis{ChoiBasis{BL},S},CompositeBasis{ChoiBasis{BR},T}} where {BL,BR,S,T}}
 
-vec(op::Operator) = Ket(KetBraBasis(basis_l(op), basis_r(op)), reshape(op.data, dimension(op.data)))
+vec(op::Operator) = Ket(KetBraBasis(basis_l(op), basis_r(op)), vec(op.data))
 function unvec(k::SuperKetType)
     bl, br = basis_l(basis(k)), basis_r(basis(k))
-    return Operator(bl, br, reshape(k.data, dimension(bl), dimension(br)))
+    @cast A[n,m] |= k.data[(n,m)] (n ∈ 1:dimension(bl), m ∈ 1:dimension(br))
+    return Operator(bl, br, A)
 end
 
 function spre(op::Operator)
     multiplicable(op, op) || throw(ArgumentError("It's not clear what spre of a non-square operator should be. See issue #113"))
-    Operator(KetBraBasis(basis_l(op)), KetBraBasis(basis_r(op)), tensor(op, identityoperator(op)).data)
+    sprepost(op, identityoperator(op))
 end
 
 function spost(op::Operator)
     multiplicable(op, op) || throw(ArgumentError("It's not clear what spost of a non-square operator should be. See issue #113"))
-    Operator(KetBraBasis(basis_r(op)), KetBraBasis(basis_l(op)), kron(permutedims(op.data), identityoperator(op).data))
+    sprepost(identityoperator(op), op)
 end
 
-sprepost(A::Operator, B::Operator) = Operator(KetBraBasis(A.basis_l, B.basis_r), KetBraBasis(A.basis_r, B.basis_l), kron(permutedims(B.data), A.data))
-
-function _super_choi((l1, l2), (r1, r2), data)
-    data = reshape(data, map(dimension, (l2, l1, r2, r1)))
-    (l1, l2), (r1, r2) = (r2, l2), (r1, l1)
-    data = PermutedDimsArray(data, (1, 3, 2, 4))
-    data = reshape(data, map(dimension, (l1⊗l2, r1⊗r2)))
-    return (l1, l2), (r1, r2), data
+function sprepost(A::Operator, B::Operator)
+    @cast C[(ν,μ), (n,m)] |= A.data[ν,n] * B.data[m,μ]
+    Operator(KetBraBasis(basis_l(A), basis_r(B)), KetBraBasis(basis_r(A), basis_l(B)), C)
 end
 
-function choi(op::SuperOperatorType)
-    bl, br = basis_l(op), basis_r(op)
-    (l1, l2), (r1, r2) = (basis_l(bl), basis_r(bl)), (basis_l(br), basis_r(br))
-    (l1, l2), (r1, r2), data = _super_choi((l1, l2), (r1, r2), op.data)
-    return Operator(ChoiBasis(l1,l2), ChoiBasis(r1,r2), data)
+# Sec IV.A. of  https://arxiv.org/abs/1111.6950
+function _super_choi(basis_fn, op)
+    l1, l2 = basis_l(basis_l(op)), basis_r(basis_l(op))
+    r1, r2 = basis_l(basis_r(op)), basis_r(basis_r(op))
+    d1, d2, d3, d4 = map(dimension, (l1, l2, r1, r2))
+    @cast A[(ν,μ), (n,m)] |= op.data[(m,μ), (n,ν)] (m ∈ 1:d1, μ ∈ 1:d2, n ∈ 1:d3, ν ∈ 1:d4)
+    return Operator(basis_fn(r2, l2), basis_fn(r1, l1), A)
 end
 
-function super(op::ChoiStateType)
-    bl, br = basis_l(op), basis_r(op)
-    (l1, l2), (r1, r2) = (basis_l(bl), basis_r(bl)), (basis_l(br), basis_r(br))
-    (l1, l2), (r1, r2), data = _super_choi((l1, l2), (r1, r2), op.data)
-    return Operator(KetBraBasis(l1,l2), KetBraBasis(r1,r2), data)
-end
+choi(op::SuperOperatorType) = _super_choi(ChoiBasis, op)
+super(op::ChoiStateType) =  _super_choi(KetBraBasis, op)
 
 dagger(a::ChoiStateType) = choi(dagger(super(a)))
 

test/test_superoperators.jl

@@ -140,12 +140,12 @@ op2 = DenseOperator(spinbasis, [0.2+0.1im 0.1+2.3im; 0.8+4.0im 0.3+1.4im])
 @test tracedistance(spost(op1)*op2, op2*op1) < 1e-12
 @test tracedistance(choi(spost(op1))*op2, op2*op1) < 1e-12
 
-@test spre(sparse(op1))*op2 == op1*op2
-@test choi(spre(sparse(op1)))*op2 == op1*op2
-@test spost(sparse(op1))*op2 == op2*op1
-@test choi(spost(sparse(op1)))*op2 == op2*op1
-@test spre(sparse(dagger(op1)))*op2 == dagger(op1)*op2
-@test choi(spre(sparse(dagger(op1))))*op2 == dagger(op1)*op2
+@test spre(sparse(op1))*op2 ≈ op1*op2
+@test choi(spre(sparse(op1)))*op2 ≈ op1*op2
+@test spost(sparse(op1))*op2 ≈ op2*op1
+@test choi(spost(sparse(op1)))*op2 ≈ op2*op1
+@test spre(sparse(dagger(op1)))*op2 ≈ dagger(op1)*op2
+@test choi(spre(sparse(dagger(op1))))*op2 ≈ dagger(op1)*op2
 @test spre(dense(dagger(op1)))*op2 ≈ dagger(op1)*op2
 @test choi(spre(dense(dagger(op1))))*op2 ≈ dagger(op1)*op2
 @test sprepost(sparse(op1), op1)*op2 ≈ op1*op2*op1