Implement KrausOperators, conversions, tensor, checks, tests

Author: akirakyle

src/QuantumOpticsBase.jl

@@ -39,8 +39,10 @@ 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, KrausOperators,
+                canonicalize, orthogonalize, make_trace_preserving, is_cptp, is_cptni,
+                is_completely_positive, is_trace_preserving, is_trace_nonincreasing,
+                spre, spost, sprepost, liouvillian, identitysuperoperator,
         #fock
                 FockBasis, number, destroy, create,
                 fockstate, coherentstate, coherentstate!,

src/pauli.jl

@@ -257,9 +257,9 @@ ChiMatrix(ptm::DensePauliTransferMatrix) = ChiMatrix(SuperOperator(ptm))
 
 """Equality for all varieties of superoperators."""
 ==(sop1::T, sop2::T) where T<:Union{DensePauliTransferMatrix, DenseSuperOpType, DenseChiMatrix} = sop1.data == sop2.data
-==(sop1::Union{DensePauliTransferMatrix, DenseSuperOpType, DenseChiMatrix}, sop2::Union{DensePauliTransferMatrix, DenseSuperOpType, DenseChiMatrix}) = false
+==(sop1::Union{DensePauliTransferMatrix, DenseSuperOpType, DenseChiMatrix}, sop2::Union{DensePauliTransferMatrix, DenseChiMatrix}) = false
 
 """Approximate equality for all varieties of superoperators."""
-function isapprox(sop1::T, sop2::T; kwargs...) where T<:Union{DensePauliTransferMatrix, DenseSuperOpType, DenseChiMatrix}
+function isapprox(sop1::T, sop2::T; kwargs...) where T<:Union{DensePauliTransferMatrix, DenseChiMatrix}
     return isapprox(sop1.data, sop2.data; kwargs...)
 end

src/superoperators.jl

@@ -1,7 +1,22 @@
+import Base: isapprox
 import QuantumInterface: AbstractSuperOperator
 import FastExpm: fastExpm
 import Adapt
 
+# TODO: this should belong in QuantumInterface.jl
+abstract type OperatorBasis{BL<:Basis,BR<:Basis} end
+abstract type SuperOperatorBasis{BL<:OperatorBasis,BR<:OperatorBasis} end
+
+"""
+    tensor(E::AbstractSuperOperator, F::AbstractSuperOperator, G::AbstractSuperOperator...)
+
+Tensor product ``\\mathcal{E}⊗\\mathcal{F}⊗\\mathcal{G}⊗…`` of the given super operators.
+"""
+tensor(a::AbstractSuperOperator, b::AbstractSuperOperator) = arithmetic_binary_error("Tensor product", a, b)
+tensor(op::AbstractSuperOperator) = op
+tensor(operators::AbstractSuperOperator...) = reduce(tensor, operators)
+
+
 """
     SuperOperator <: AbstractSuperOperator
 
@@ -70,6 +85,9 @@ sparse(a::SuperOperator) = SuperOperator(a.basis_l, a.basis_r, sparse(a.data))
 
 ==(a::SuperOperator{B1,B2}, b::SuperOperator{B1,B2}) where {B1,B2} = (samebases(a,b) && a.data == b.data)
 ==(a::SuperOperator, b::SuperOperator) = false
+isapprox(a::SuperOperator{B1,B2}, b::SuperOperator{B1,B2}; kwargs...) where {B1,B2} =
+    (samebases(a,b) && isapprox(a.data, b.data; kwargs...))
+isapprox(a::SuperOperator, b::SuperOperator; kwargs...) = false
 
 Base.length(a::SuperOperator) = length(a.basis_l[1])*length(a.basis_l[2])*length(a.basis_r[1])*length(a.basis_r[2])
 samebases(a::SuperOperator, b::SuperOperator) = samebases(a.basis_l[1], b.basis_l[1]) && samebases(a.basis_l[2], b.basis_l[2]) &&
@@ -130,7 +148,7 @@ function spre(op::AbstractOperator)
     if !samebases(op.basis_l, op.basis_r)
         throw(ArgumentError("It's not clear what spre of a non-square operator should be. See issue #113"))
     end
-    SuperOperator((op.basis_l, op.basis_l), (op.basis_r, op.basis_r), tensor(op, identityoperator(op)).data)
+    SuperOperator((op.basis_l, op.basis_l), (op.basis_r, op.basis_r), kron(identityoperator(op).data, op.data))
 end
 
 """
@@ -319,10 +337,14 @@ end
     throw(IncompatibleBases())
 end
 
+# TODO should all of PauliTransferMatrix, ChiMatrix, ChoiState, and KrausOperators subclass AbstractSuperOperator?
 """
     ChoiState <: AbstractSuperOperator
 
 Superoperator represented as a choi state.
+
+The convention is chosen such that the input operators live in `(basis_l[1], basis_r[1])` while
+the output operators live in `(basis_r[2], basis_r[2])`.
 """
 mutable struct ChoiState{B1,B2,T} <: AbstractSuperOperator{B1,B2}
     basis_l::B1
@@ -345,6 +367,27 @@ 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))
+isapprox(a::ChoiState, b::ChoiState; kwargs...) = isapprox(SuperOperator(a), SuperOperator(b); kwargs...)
+
+# Container to hold each of the four bases for a Choi operator when converting it to
+# an operator so that if any are CompositeBases tensor doesn't lossily collapse them
+struct ChoiSubBasis{S,B<:Basis} <: Basis
+    shape::S
+    basis::B
+end
+ChoiSubBasis(b::Basis) = ChoiSubBasis(b.shape, b)
+
+# TODO: decide whether to document and export this
+choi_to_operator(c::ChoiState) = Operator(
+    ChoiSubBasis(c.basis_l[2])⊗ChoiSubBasis(c.basis_l[1]), ChoiSubBasis(c.basis_r[2])⊗ChoiSubBasis(c.basis_r[1]), c.data)
+
+function tensor(a::ChoiState, b::ChoiState)
+    op = choi_to_operator(a) ⊗ choi_to_operator(b)
+    op = permutesystems(op, [1,3,2,4])
+    ChoiState((a.basis_l[1] ⊗ b.basis_l[1], a.basis_l[2] ⊗ b.basis_l[2]),
+              (a.basis_r[1] ⊗ b.basis_r[1], a.basis_r[2] ⊗ b.basis_r[2]), op.data)
+end
+tensor(a::SuperOperator, b::SuperOperator) = SuperOperator(tensor(ChoiState(a), ChoiState(b)))
 
 # reshape swaps within systems due to colum major ordering
 # https://docs.qojulia.org/quantumobjects/operators/#tensor_order
@@ -359,8 +402,218 @@ 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)...)
 
+
+"""
+    KrausOperators <: AbstractSuperOperator
+
+Superoperator represented as a list of Kraus operators.
+
+Note that KrausOperators can only represent linear maps taking density operators to other
+(potentially unnormalized) density operators.
+In contrast the `SuperOperator` or `ChoiState` representations can represent arbitrary linear maps
+taking arbitrary operators defined on ``H_A \\to H_B`` to ``H_C \\to H_D``.
+In otherwords, the Kraus representation is only defined for completely positive linear maps of the form
+``(H_A \\to H_A) \\to (H_B \\to H_B)``.
+Thus converting from `SuperOperator` or `ChoiState` to `KrausOperators` will throw an exception if the
+map cannot be faithfully represented up to the specificed tolerance `tol`.
+
+----------------------------
+Old text:
+Note unlike the SuperOperator or ChoiState types where it is possible to have
+`basis_l[1] != basis_l[2]` and `basis_r[1] != basis_r[2]`
+which allows representations of maps between general linear operators defined on ``H_A \\to H_B``,
+a quantum channel can only act on valid density operators which live in ``H_A \\to H_A``.
+Thus the Kraus representation is only defined for quantum channels which map
+``(H_A \\to H_A) \\to (H_B \\to H_B)``.
+"""
+mutable struct KrausOperators{B1,B2,T} <: AbstractSuperOperator{B1,B2}
+    basis_l::B1
+    basis_r::B2
+    data::Vector{T}
+    function KrausOperators{BL,BR,T}(basis_l::BL, basis_r::BR, data::Vector{T}) where {BL,BR,T}
+        if (any(!samebases(basis_r, M.basis_r) for M in data) ||
+            any(!samebases(basis_l, M.basis_l) for M in data))
+            throw(DimensionMismatch("Tried to assign data with incompatible bases"))
+        end
+
+        new(basis_l, basis_r, data)
+    end
+end
+KrausOperators{BL,BR}(b1::BL,b2::BR,data::Vector{T}) where {BL,BR,T} = KrausOperators{BL,BR,T}(b1,b2,data)
+KrausOperators(b1::BL,b2::BR,data::Vector{T}) where {BL,BR,T} = KrausOperators{BL,BR,T}(b1,b2,data)
+
+dense(a::KrausOperators) = KrausOperators(a.basis_l, a.basis_r, [dense(op) for op in a.data])
+sparse(a::KrausOperators) = KrausOperators(a.basis_l, a.basis_r, [sparse(op) for op in a.data])
+dagger(a::KrausOperators) = KrausOperators(a.basis_r, a.basis_l, [dagger(op) for op in a.data])
+*(a::KrausOperators{B1,B2}, b::KrausOperators{B2,B3}) where {B1,B2,B3} =
+    KrausOperators(a.basis_l, b.basis_r, [A*B for A in a.data for B in b.data])
+*(a::KrausOperators, b::KrausOperators) = throw(IncompatibleBases())
+*(a::KrausOperators{BL,BR}, b::Operator{BR,BR}) where {BL,BR} = sum(op*b*dagger(op) for op in a.data)
+==(a::KrausOperators, b::KrausOperators) = (SuperOperator(a) == SuperOperator(b))
+isapprox(a::KrausOperators, b::KrausOperators; kwargs...) = isapprox(SuperOperator(a), SuperOperator(b); kwargs...)
+tensor(a::KrausOperators, b::KrausOperators) =
+    KrausOperators(a.basis_l ⊗ b.basis_l, a.basis_r ⊗ b.basis_r,
+                   [A ⊗ B for A in a.data for B in b.data])
+
+"""
+    orthogonalize(kraus::KrausOperators; tol=√eps)
+
+Orthogonalize the set kraus operators by performing a qr decomposition on their vec'd operators.
+Note that this is different than `canonicalize` which returns a kraus decomposition such
+that the kraus operators are Hilbert–Schmidt orthorgonal.
+
+If the input dimension is d and output dimension is d' then the number of kraus
+operators returned is guaranteed to be no greater than dd', however it may be greater
+than the Kraus rank.
+
+`orthogonalize` should always be much faster than canonicalize as it avoids an explicit eigendecomposition
+and thus also preserves sparsity if the kraus operators are sparse.
+"""
+function orthogonalize(kraus::KrausOperators; tol=_get_tol(kraus))
+    bl, br = kraus.basis_l, kraus.basis_r
+    dim = length(bl)*length(br)
+
+    A = stack(reshape(op.data, dim) for op in kraus.data; dims=1)
+    F = qr(A; tol=tol)
+    # rank(F) for some reason doesn't work but should
+    rank = maximum(findnz(F.R)[1]) 
+    # Sanity checks that help illustrate what qr() returns:
+    # @assert (F.R ≈ (sparse(F.Q') * A[F.prow,F.pcol])[1:dim,:])
+    # @assert (all(iszero(F.R[rank+1:end,:])))
+
+    ops = [Operator(bl, br, copy(reshape( # copy materializes reshaped view
+        F.R[i,invperm(F.pcol)], (length(bl), length(br))))) for i=1:rank]
+    return KrausOperators(bl, br, ops)
+end
+
+"""
+    canonicalize(kraus::KrausOperators; tol=√eps)
+
+Transform the quantum channel into canonical form such that the kraus operators ``{A_k}``
+are Hilbert–Schmidt orthorgonal:
+
+```math
+\\Tr A_i^\\dagger A_j \\sim \\delta_{i,j}
+```
+
+If the input dimension is d and output dimension is d' then the number of kraus
+operators returned is guaranteed to be no greater than dd' and will furthermore
+be equal the Kraus rank of the channel up to numerical imprecision controlled by `tol`.  
+"""
+canonicalize(kraus::KrausOperators; tol=_get_tol(kraus)) = KrausOperators(ChoiState(kraus); tol=tol)
+
+# TODO: document
+function make_trace_preserving(kraus; tol=_get_tol(kraus))
+    m = I - sum(dagger(M)*M for M in kraus.data).data
+    if isa(_positive_eigen(m, tol), Number)
+        throw(ArgumentError("Channel must be trace nonincreasing"))
+    end
+    K = Operator(kraus.basis_l, kraus.basis_r, sqrt(Matrix(m)))
+    return KrausOperators(kraus.basis_l, kraus.basis_r, [kraus.data; K])
+end
+
+SuperOperator(kraus::KrausOperators) =
+    SuperOperator((kraus.basis_l, kraus.basis_l), (kraus.basis_r, kraus.basis_r),
+                  (sum(conj(op)⊗op for op in kraus.data)).data)
+
+ChoiState(kraus::KrausOperators) =
+    ChoiState((kraus.basis_r, kraus.basis_l), (kraus.basis_r, kraus.basis_l),
+              (sum((M=op.data; reshape(M, (length(M), 1))*reshape(M, (1, length(M))))
+                   for op in kraus.data)))
+
+_choi_state_maps_density_ops(choi::ChoiState) = (samebases(choi.basis_l[1], choi.basis_r[1]) &&
+                                                samebases(choi.basis_l[2], choi.basis_r[2]))
+
+# TODO: consider using https://github.com/jlapeyre/IsApprox.jl
+_is_hermitian(M, tol) = ishermitian(M) || isapprox(M, M', atol=tol)
+_is_identity(M, tol) = isapprox(M, I, atol=tol)
+
+# TODO: document
+# data must be Hermitian!
+function _positive_eigen(data, tol)
+    # LinearAlgebra's eigen returns eigenvals sorted smallest to largest for Hermitian matrices
+    vals, vecs = eigen(Hermitian(Matrix(data)))
+    vals[1] < -tol && return vals[1]
+    ret = [(val, vecs[:,i]) for (i, val) in enumerate(vals) if val > tol]
+    return ret
+end
+
+function KrausOperators(choi::ChoiState; tol=_get_tol(choi))
+    if !_choi_state_maps_density_ops(choi)
+        throw(DimensionMismatch("Tried to convert Choi state of something that isn't a quantum channel mapping density operators to density operators"))
+    end
+    if !_is_hermitian(choi.data, tol)
+        throw(ArgumentError("Tried to convert nonhermitian Choi state"))
+    end
+    bl, br = choi.basis_l[2], choi.basis_l[1]
+    eigs = _positive_eigen(choi.data, tol)
+    if isa(eigs, Number)
+        throw(ArgumentError("Tried to convert a non-positive semidefinite Choi state,"*
+            "failed for smallest eigval $(eigs), consider increasing tol=$(tol)"))
+    end
+
+    ops = [Operator(bl, br, sqrt(val)*reshape(vec, length(bl), length(br))) for (val, vec) in eigs]
+    return KrausOperators(bl, br, ops)
+end
+
+KrausOperators(op::SuperOperator; tol=_get_tol(op)) = KrausOperators(ChoiState(op); tol=tol)
+
+# TODO: document superoperator representation precident: everything of mixed type returns SuperOperator
 *(a::ChoiState, b::SuperOperator) = SuperOperator(a)*b
 *(a::SuperOperator, b::ChoiState) = a*SuperOperator(b)
+*(a::KrausOperators, b::SuperOperator) = SuperOperator(a)*b
+*(a::SuperOperator, b::KrausOperators) = a*SuperOperator(b)
+*(a::KrausOperators, b::ChoiState) = SuperOperator(a)*SuperOperator(b)
+*(a::ChoiState, b::KrausOperators) = SuperOperator(a)*SuperOperator(b)
+
+_get_tol(kraus::KrausOperators) = sqrt(eps(real(eltype(eltype(fieldtypes(typeof(kraus))[3])))))
+_get_tol(super::SuperOperator) = sqrt(eps(real(eltype(fieldtypes(typeof(super))[3]))))
+_get_tol(super::ChoiState) = sqrt(eps(real(eltype(fieldtypes(typeof(super))[3]))))
+
+# TODO: document this
+is_completely_positive(choi::KrausOperators; tol=_get_tol(choi)) = true
+
+function is_completely_positive(choi::ChoiState; tol=_get_tol(choi))
+    _choi_state_maps_density_ops(choi) || return false
+    _is_hermitian(choi.data, tol) || return false
+    isa(_positive_eigen(Hermitian(choi.data), tol), Number) && return false
+    return true
+end
+
+is_completely_positive(super::SuperOperator; tol=_get_tol(super)) =
+    is_completely_positive(ChoiState(super); tol=tol)
+
+# TODO: document this
+is_trace_preserving(kraus::KrausOperators; tol=_get_tol(kraus)) =
+    _is_identity(sum(dagger(M)*M for M in kraus.data).data, tol)
+
+is_trace_preserving(choi::ChoiState; tol=_get_tol(choi)) =
+    _is_identity(ptrace(choi_to_operator(choi), 1).data, tol)
+
+is_trace_preserving(super::SuperOperator; tol=_get_tol(super)) =
+    is_trace_preserving(ChoiState(super); tol=tol)
+
+# TODO: document this
+function is_trace_nonincreasing(kraus::KrausOperators; tol=_get_tol(kraus))
+    m = I - sum(dagger(M)*M for M in kraus.data).data
+    _is_hermitian(m, tol) || return false
+    return !isa(_positive_eigen(Hermitian(m), tol), Number)
+end
+
+function is_trace_nonincreasing(choi::ChoiState; tol=_get_tol(choi))
+    m = I - ptrace(choi_to_operator(choi), 1).data
+    _is_hermitian(m, tol) || return false
+    return !isa(_positive_eigen(Hermitian(m), tol), Number)
+end
+
+is_trace_nonincreasing(super::SuperOperator; tol=_get_tol(super)) =
+    is_trace_nonincreasing(ChoiState(super); tol=tol)
+
+# TODO: document this
+is_cptp(sop; tol=_get_tol(sop)) = is_completely_positive(sop; tol=tol) && is_trace_preserving(sop; tol=tol)
+
+# TODO: document this
+is_cptni(sop; tol=_get_tol(sop)) = is_completely_positive(sop; tol=tol) && is_trace_nonincreasing(sop; tol=tol)
 
 # GPU adaptation
 Adapt.adapt_structure(to, x::SuperOperator) = SuperOperator(x.basis_l, x.basis_r, Adapt.adapt(to, x.data))