Merge pull request #152 from aarontrowbridge/main

functionality for permuting tensor product order of Qobj

docs/src/api.md

@@ -75,6 +75,7 @@ expect
 LinearAlgebra.kron
 tensor
 ⊗
+permute
 sparse_to_dense
 dense_to_sparse
 vec2mat

src/QuantumToolbox.jl

@@ -1,7 +1,7 @@
 module QuantumToolbox
 
-# Re-export: 
-#   1. basic functions in LinearAlgebra and SparseArrays 
+# Re-export:
+#   1. basic functions in LinearAlgebra and SparseArrays
 #   2. the solvers in ODE and LinearSolve
 import Reexport: @reexport
 @reexport using LinearAlgebra

src/qobj/arithmetic_and_attributes.jl

@@ -8,6 +8,7 @@ export sqrtm, sinm, cosm
 export ptrace
 export tidyup, tidyup!
 export get_data, get_coherence
+export permute
 
 #    Broadcasting
 Base.broadcastable(x::QuantumObject) = x.data
@@ -458,7 +459,7 @@ get_data(A::QuantumObject) = A.data
 Get the coherence value ``\alpha`` by measuring the expectation value of the destruction
 operator ``\hat{a}`` on the state ``\ket{\psi}``.
 
-It returns both ``\alpha`` and the state 
+It returns both ``\alpha`` and the state
 ``\ket{\delta_\psi} = \exp ( \bar{\alpha} \hat{a} - \alpha \hat{a}^\dagger )``. The
 latter corresponds to the quantum fulctuations around the coherent state ``\ket{\alpha}``.
 """
@@ -471,3 +472,51 @@ function get_coherence(
 
     return α, D' * ψ
 end
+
+@doc raw"""
+    permute(A::QuantumObject, order::Vector{Int})
+
+Permute the tensor structure of a [`QuantumObject`](@ref) `A` according to the specified `order` list
+
+Note that this method currently works for [`Ket`](@ref), [`Bra`](@ref), and [`Operator`](@ref) types of [`QuantumObject`](@ref).
+
+# Examples
+
+If `order = [2, 1, 3]`, the Hilbert space structure will be re-arranged: H₁ ⊗ H₂ ⊗ H₃ → H₂ ⊗ H₁ ⊗ H₃.
+
+```
+julia> ψ1 = fock(2, 0)
+julia> ψ2 = fock(3, 1)
+julia> ψ3 = fock(4, 2)
+julia> ψ_123 = tensor(ψ1, ψ2, ψ3)
+julia> permute(ψ_123, [2, 1, 3]) ≈ tensor(ψ2, ψ1, ψ3)
+true
+```
+"""
+function permute(
+    A::QuantumObject{<:AbstractArray{T},ObjType},
+    order::AbstractVector{Int},
+) where {T,ObjType<:Union{KetQuantumObject,BraQuantumObject,OperatorQuantumObject}}
+    (length(order) != length(A.dims)) &&
+        throw(ArgumentError("The order list must have the same length as the number of subsystems (A.dims)"))
+
+    !isperm(order) && throw(ArgumentError("$(order) is not a valid permutation of the subsystems (A.dims)"))
+
+    # obtain the arguments: dims for reshape; perm for PermutedDimsArray
+    dims, perm = _dims_and_perm(A.type, A.dims, order, length(order))
+
+    return QuantumObject(reshape(PermutedDimsArray(reshape(A.data, dims...), perm), size(A)), A.type, A.dims[order])
+end
+
+function _dims_and_perm(
+    ::ObjType,
+    dims::AbstractVector{Int},
+    order::AbstractVector{Int},
+    L::Int,
+) where {ObjType<:Union{KetQuantumObject,BraQuantumObject}}
+    return reverse(dims), reverse!((L + 1) .- order)
+end
+
+function _dims_and_perm(::OperatorQuantumObject, dims::AbstractVector{Int}, order::AbstractVector{Int}, L::Int)
+    return reverse!([dims; dims]), reverse!((2 * L + 1) .- [order; order .+ L])
+end

src/qobj/functions.jl

@@ -4,8 +4,8 @@ Functions which manipulates QuantumObject
 
 export ket2dm
 export expect
-export tensor, ⊗
 export sparse_to_dense, dense_to_sparse
+export tensor, ⊗
 export vec2mat, mat2vec
 
 @doc raw"""
@@ -61,6 +61,48 @@ function expect(
     return ishermitian(O) ? real(tr(O * ρ)) : tr(O * ρ)
 end
 
+"""
+    sparse_to_dense(A::QuantumObject)
+
+Converts a sparse QuantumObject to a dense QuantumObject.
+"""
+sparse_to_dense(A::QuantumObject{<:AbstractVecOrMat}) = QuantumObject(sparse_to_dense(A.data), A.type, A.dims)
+sparse_to_dense(A::MT) where {MT<:AbstractSparseMatrix} = Array(A)
+for op in (:Transpose, :Adjoint)
+    @eval sparse_to_dense(A::$op{T,<:AbstractSparseMatrix}) where {T<:BlasFloat} = Array(A)
+end
+sparse_to_dense(A::MT) where {MT<:AbstractArray} = A
+
+function sparse_to_dense(::Type{M}) where {M<:SparseMatrixCSC}
+    T = M
+    par = T.parameters
+    npar = length(par)
+    (2 == npar) || error("Type $M is not supported.")
+    return Matrix{par[1]}
+end
+
+sparse_to_dense(::Type{M}) where {M<:AbstractMatrix} = M
+
+"""
+    dense_to_sparse(A::QuantumObject)
+
+Converts a dense QuantumObject to a sparse QuantumObject.
+"""
+dense_to_sparse(A::QuantumObject{<:AbstractVecOrMat}, tol::Real = 1e-10) =
+    QuantumObject(dense_to_sparse(A.data, tol), A.type, A.dims)
+function dense_to_sparse(A::MT, tol::Real = 1e-10) where {MT<:AbstractMatrix}
+    idxs = findall(@. abs(A) > tol)
+    row_indices = getindex.(idxs, 1)
+    col_indices = getindex.(idxs, 2)
+    vals = getindex(A, idxs)
+    return sparse(row_indices, col_indices, vals, size(A)...)
+end
+function dense_to_sparse(A::VT, tol::Real = 1e-10) where {VT<:AbstractVector}
+    idxs = findall(@. abs(A) > tol)
+    vals = getindex(A, idxs)
+    return sparsevec(idxs, vals, length(A))
+end
+
 @doc raw"""
     kron(A::QuantumObject, B::QuantumObject)
 
@@ -131,12 +173,12 @@ julia> tensor(x_list...)
 Quantum Object:   type=Operator   dims=[2, 2, 2]   size=(8, 8)   ishermitian=true
 8×8 SparseMatrixCSC{ComplexF64, Int64} with 8 stored entries:
      ⋅          ⋅          ⋅      …      ⋅          ⋅      1.0+0.0im
-     ⋅          ⋅          ⋅             ⋅      1.0+0.0im      ⋅    
-     ⋅          ⋅          ⋅         1.0+0.0im      ⋅          ⋅    
-     ⋅          ⋅          ⋅             ⋅          ⋅          ⋅    
-     ⋅          ⋅          ⋅             ⋅          ⋅          ⋅    
-     ⋅          ⋅      1.0+0.0im  …      ⋅          ⋅          ⋅    
-     ⋅      1.0+0.0im      ⋅             ⋅          ⋅          ⋅    
+     ⋅          ⋅          ⋅             ⋅      1.0+0.0im      ⋅
+     ⋅          ⋅          ⋅         1.0+0.0im      ⋅          ⋅
+     ⋅          ⋅          ⋅             ⋅          ⋅          ⋅
+     ⋅          ⋅          ⋅             ⋅          ⋅          ⋅
+     ⋅          ⋅      1.0+0.0im  …      ⋅          ⋅          ⋅
+     ⋅      1.0+0.0im      ⋅             ⋅          ⋅          ⋅
  1.0+0.0im      ⋅          ⋅             ⋅          ⋅          ⋅
 ```
 """
@@ -187,48 +229,6 @@ Quantum Object:   type=Operator   dims=[20, 20]   size=(400, 400)   ishermitian=
 """
 ⊗(A::QuantumObject, B::QuantumObject) = kron(A, B)
 
-"""
-    sparse_to_dense(A::QuantumObject)
-
-Converts a sparse QuantumObject to a dense QuantumObject.
-"""
-sparse_to_dense(A::QuantumObject{<:AbstractVecOrMat}) = QuantumObject(sparse_to_dense(A.data), A.type, A.dims)
-sparse_to_dense(A::MT) where {MT<:AbstractSparseMatrix} = Array(A)
-for op in (:Transpose, :Adjoint)
-    @eval sparse_to_dense(A::$op{T,<:AbstractSparseMatrix}) where {T<:BlasFloat} = Array(A)
-end
-sparse_to_dense(A::MT) where {MT<:AbstractArray} = A
-
-function sparse_to_dense(::Type{M}) where {M<:SparseMatrixCSC}
-    T = M
-    par = T.parameters
-    npar = length(par)
-    (2 == npar) || error("Type $M is not supported.")
-    return Matrix{par[1]}
-end
-
-sparse_to_dense(::Type{M}) where {M<:AbstractMatrix} = M
-
-"""
-    dense_to_sparse(A::QuantumObject)
-
-Converts a dense QuantumObject to a sparse QuantumObject.
-"""
-dense_to_sparse(A::QuantumObject{<:AbstractVecOrMat}, tol::Real = 1e-10) =
-    QuantumObject(dense_to_sparse(A.data, tol), A.type, A.dims)
-function dense_to_sparse(A::MT, tol::Real = 1e-10) where {MT<:AbstractMatrix}
-    idxs = findall(@. abs(A) > tol)
-    row_indices = getindex.(idxs, 1)
-    col_indices = getindex.(idxs, 2)
-    vals = getindex(A, idxs)
-    return sparse(row_indices, col_indices, vals, size(A)...)
-end
-function dense_to_sparse(A::VT, tol::Real = 1e-10) where {VT<:AbstractVector}
-    idxs = findall(@. abs(A) > tol)
-    vals = getindex(A, idxs)
-    return sparsevec(idxs, vals, length(A))
-end
-
 """
     vec2mat(A::AbstractVector)
 

test/quantum_objects.jl

@@ -291,4 +291,39 @@
     @test typeof(SparseMatrixCSC(Md).data) == SparseMatrixCSC{Int64,Int64}
     @test typeof(SparseMatrixCSC(Ms).data) == SparseMatrixCSC{Int64,Int64}
     @test typeof(SparseMatrixCSC{ComplexF64}(Ms).data) == SparseMatrixCSC{ComplexF64,Int64}
+
+    # permute tests
+    ket_a = Qobj(rand(ComplexF64, 2))
+    ket_b = Qobj(rand(ComplexF64, 3))
+    ket_c = Qobj(rand(ComplexF64, 4))
+    ket_d = Qobj(rand(ComplexF64, 5))
+    ket_bdca = permute(tensor(ket_a, ket_b, ket_c, ket_d), [2, 4, 3, 1])
+    bra_a = ket_a'
+    bra_b = ket_b'
+    bra_c = ket_c'
+    bra_d = ket_d'
+    bra_bdca = permute(tensor(bra_a, bra_b, bra_c, bra_d), [2, 4, 3, 1])
+    op_a = Qobj(rand(ComplexF64, 2, 2))
+    op_b = Qobj(rand(ComplexF64, 3, 3))
+    op_c = Qobj(rand(ComplexF64, 4, 4))
+    op_d = Qobj(rand(ComplexF64, 5, 5))
+    op_bdca = permute(tensor(op_a, op_b, op_c, op_d), [2, 4, 3, 1])
+    correct_dims = [3, 5, 4, 2]
+    wrong_order1 = [1]
+    wrong_order2 = [2, 3, 4, 5]
+    @test ket_bdca ≈ tensor(ket_b, ket_d, ket_c, ket_a)
+    @test bra_bdca ≈ tensor(bra_b, bra_d, bra_c, bra_a)
+    @test op_bdca ≈ tensor(op_b, op_d, op_c, op_a)
+    @test ket_bdca.dims == correct_dims
+    @test bra_bdca.dims == correct_dims
+    @test op_bdca.dims == correct_dims
+    @test typeof(ket_bdca.type) <: KetQuantumObject
+    @test typeof(bra_bdca.type) <: BraQuantumObject
+    @test typeof(op_bdca.type) <: OperatorQuantumObject
+    @test_throws ArgumentError permute(ket_bdca, wrong_order1)
+    @test_throws ArgumentError permute(ket_bdca, wrong_order2)
+    @test_throws ArgumentError permute(bra_bdca, wrong_order1)
+    @test_throws ArgumentError permute(bra_bdca, wrong_order2)
+    @test_throws ArgumentError permute(op_bdca, wrong_order1)
+    @test_throws ArgumentError permute(op_bdca, wrong_order2)
 end