before embedding check that the destination basis matches the operator basis

src/operators.jl

@@ -6,6 +6,7 @@ export AbstractOperator, DataOperator, length, basis, dagger, ishermitian, tenso
 
 import Base: ==, +, -, *, /, ^, length, one, exp, conj, conj!, transpose
 import LinearAlgebra: tr, ishermitian
+import SparseArrays: sparse
 import ..bases: basis, tensor, ptrace, permutesystems,
             samebases, check_samebases, multiplicable
 import ..states: dagger, normalize, normalize!
@@ -88,22 +89,128 @@ tensor(op::AbstractOperator) = op
 tensor(operators::AbstractOperator...) = reduce(tensor, operators)
 
 
+
+"""
+    check_indices(indices::Array)
+
+Determine whether a collection of indices, written as a list of (integers or lists of integers) is unique.
+This assures that the embedded operators are in non-overlapping subspaces.
+"""
+function check_indices(indices::Array)
+    status = true
+    # Check that no sub-list contains duplicates.
+    for i in filter(x -> typeof(x) <: Array, indices)
+        status &= (length(Set(i)) == length(i))
+    end
+    # Check that there are no duplicates across `indices`
+    for (idx, i) in enumerate(indices[1:end-1])
+        for j in indices[idx+1:end]
+            status &= (length(intersect(i, j)) == 0)
+        end
+    end
+    return status
+end
+
+
 """
     embed(basis1[, basis2], indices::Vector, operators::Vector)
 
 Tensor product of operators where missing indices are filled up with identity operators.
 """
 function embed(basis_l::CompositeBasis, basis_r::CompositeBasis,
-               indices::Vector{Int}, operators::Vector{T}) where T<:AbstractOperator
+               indices::Vector, operators::Vector{T}) where T<:AbstractOperator
+
+    @assert check_indices(indices)
+
     N = length(basis_l.bases)
     @assert length(basis_r.bases) == N
     @assert length(indices) == length(operators)
+
+    # Embed all single-subspace operators.
+    idxop_sb = [x for x in zip(indices, operators) if typeof(x[1]) <: Int64]
+    indices_sb = [x[1] for x in idxop_sb]
+    ops_sb = [x[2] for x in idxop_sb]
+
+    for (idxsb, opsb) in zip(indices_sb, ops_sb)
+        @assert (opsb.basis_l == basis_l.bases[idxsb]) || throw(bases.IncompatibleBases())
+        @assert (opsb.basis_r == basis_r.bases[idxsb]) || throw(bases.IncompatibleBases())
+    end
+
+    embed_op = tensor([i ∈ indices_sb ? ops_sb[indexin(i, indices_sb)[1]] : identityoperator(T, basis_l.bases[i], basis_r.bases[i]) for i=1:N]...)
+
+    # Embed all joint-subspace operators.
+    idxop_comp = [x for x in zip(indices, operators) if typeof(x[1]) <: Array]
+
+    for (idxs, op) in idxop_comp
+        embed_op *= embed(basis_l, basis_r, idxs, op)
+    end
+
+    return embed_op
+end
+
+
+"""
+    embed(basis1[, basis2], indices::Vector, operators::Vector)
+
+Embed operator acting on a joint Hilbert space where missing indices are filled up with identity operators.
+"""
+function embed(basis_l::CompositeBasis, basis_r::CompositeBasis,
+               indices::Vector{Int}, op::T) where T<:AbstractOperator
+    N = length(basis_l.bases)
+    @assert length(basis_r.bases) == N
+
+    @assert reduce(tensor, basis_l.bases[indices]) == op.basis_l || throw(bases.IncompatibleBases())
+    @assert reduce(tensor, basis_r.bases[indices]) == op.basis_r || throw(bases.IncompatibleBases())
+
+    index_order = [idx for idx in 1:length(basis_l.bases) if idx ∉ indices]
+    all_operators = AbstractOperator[identityoperator(T, basis_l.bases[i], basis_r.bases[i]) for i in index_order]
+
+    for idx in indices
+        pushfirst!(index_order, idx)
+    end
+    push!(all_operators, op)
+
     sortedindices.check_indices(N, indices)
-    tensor([i ∈ indices ? operators[indexin(i, indices)[1]] : identityoperator(T, basis_l.bases[i], basis_r.bases[i]) for i=1:N]...)
+
+    # Create the operator.
+    permuted_op = tensor(all_operators...)
+
+    # Create a copy to fill with correctly-ordered objects (basis and data).
+    unpermuted_op = copy(permuted_op)
+
+    # Create the correctly ordered basis.
+    unpermuted_op.basis_l = basis_l
+    unpermuted_op.basis_r = basis_r
+
+    # Reorient the matrix to act in the correctly ordered basis.
+    # Get the dimensions necessary for index permuting.
+    dims_l = [b.shape[1] for b in basis_l.bases]
+    dims_r = [b.shape[1] for b in basis_r.bases]
+
+    # Get the order of indices to use in the first reshape. Julia indices go in
+    # reverse order.
+    expand_order = index_order[end:-1:1]
+    # Get the dimensions associated with those indices.
+    expand_dims_l = dims_l[expand_order]
+    expand_dims_r = dims_r[expand_order]
+
+    # Prepare the permutation to the correctly ordered basis.
+    perm_order_l = [indexin(idx, expand_order)[1] for idx in 1:length(dims_l)]
+    perm_order_r = [indexin(idx, expand_order)[1] for idx in 1:length(dims_r)]
+
+    # Perform the index expansion, the permutation, and the index collapse.
+    M = (reshape(permuted_op.data, tuple([expand_dims_l; expand_dims_r]...)) |>
+         x -> permutedims(x, [perm_order_l; perm_order_r .+ length(dims_l)]) |>
+         x -> sparse(reshape(x, (prod(dims_l), prod(dims_r)))))
+
+    unpermuted_op.data = M
+
+    return unpermuted_op
 end
 embed(basis_l::CompositeBasis, basis_r::CompositeBasis, index::Int, op::AbstractOperator) = embed(basis_l, basis_r, Int[index], [op])
 embed(basis::CompositeBasis, index::Int, op::AbstractOperator) = embed(basis, basis, Int[index], [op])
-embed(basis::CompositeBasis, indices::Vector{Int}, operators::Vector{T}) where {T<:AbstractOperator} = embed(basis, basis, indices, operators)
+embed(basis::CompositeBasis, indices::Vector, operators::Vector{T}) where {T<:AbstractOperator} = embed(basis, basis, indices, operators)
+embed(basis::CompositeBasis, indices::Vector{Int}, op::AbstractOperator) = embed(basis, basis, indices, op)
 
 """
     embed(basis1[, basis2], operators::Dict)

src/operators_sparse.jl

@@ -130,6 +130,7 @@ function operators.permutesystems(rho::SparseOperator{B1,B2}, perm::Vector{Int})
 end
 
 operators.identityoperator(::Type{T}, b1::Basis, b2::Basis) where {T<:SparseOperator} = SparseOperator(b1, b2, sparse(ComplexF64(1)*I, length(b1), length(b2)))
+operators.identityoperator(::Type{T}, b1::Basis, b2::Basis) where {T<:DataOperator} = SparseOperator(b1, b2, sparse(ComplexF64(1)*I, length(b1), length(b2)))
 operators.identityoperator(b1::Basis, b2::Basis) = identityoperator(SparseOperator, b1, b2)
 operators.identityoperator(b::Basis) = identityoperator(b, b)
 

test/test_operators.jl

@@ -13,10 +13,11 @@ end
 
 Random.seed!(0)
 
-b1 = GenericBasis(5)
-b2 = GenericBasis(3)
+b1 = GenericBasis(3)
+b2 = GenericBasis(2)
 b = b1 ⊗ b2
 op1 = randoperator(b1)
+op2 = randoperator(b2)
 op = randoperator(b, b)
 op_test = test_operators(b, b, op.data)
 op_test2 = test_operators(b1, b, randoperator(b1, b).data)
@@ -25,7 +26,7 @@ op_test3 = test_operators(b1 ⊗ b2, b2 ⊗ b1, randoperator(b, b).data)
 ρ = randoperator(b)
 
 @test basis(op1) == b1
-@test length(op1) == length(op1.data) == 25
+@test length(op1) == length(op1.data) == length(b1)^2
 
 @test_throws ArgumentError op_test*op_test
 @test_throws ArgumentError -op_test
@@ -59,8 +60,55 @@ op_test3 = test_operators(b1 ⊗ b2, b2 ⊗ b1, randoperator(b, b).data)
 @test_throws ArgumentError tensor(op_test, op_test)
 @test_throws ArgumentError permutesystems(op_test, [1, 2])
 
-@test embed(b, b, 1, op) == embed(b, 1, op)
+@test embed(b, b, [1,2], op) == embed(b, [1,2], op)
 @test embed(b, Dict{Vector{Int}, SparseOperator}()) == identityoperator(b)
+@test_throws bases.IncompatibleBases embed(b1⊗b2, [2], [op1])
+
+b_comp = b⊗b
+@test embed(b_comp, [1,[3,4]], [op1,op]) == dense(op1 ⊗ one(b2) ⊗ op)
+@test embed(b_comp, [[1,2],4], [op,op2]) == dense(op ⊗ one(b1) ⊗ op2)
+@test_throws bases.IncompatibleBases embed(b_comp, [[1,2],3], [op,op2])
+@test_throws bases.IncompatibleBases embed(b_comp, [[1,3],4], [op,op2])
+
+function basis_vec(n, N)
+    x = zeros(Complex{Float64}, N)
+    x[n+1] = 1
+    return x
+end
+function basis_maker(dims...)
+    function bm(ns...)
+        bases = [basis_vec(n, dim) for (n, dim) in zip(ns, dims)][end:-1:1]
+        return reduce(kron, bases)
+    end
+end
+
+embed_op = embed(b_comp, [1,4], op)
+bv = basis_maker(3,2,3,2)
+all_idxs = [(idx, jdx) for (idx, jdx) in [Iterators.product(0:1, 0:2)...]]
+
+m11 = reshape([Bra(b_comp, bv(0,idx,jdx,0)) * embed_op * Ket(b_comp, bv(0,kdx,ldx,0))
+              for ((idx, jdx), (kdx, ldx)) in Iterators.product(all_idxs, all_idxs)], (6,6))
+@test isapprox(m11 / op.data[1, 1], diagm(0=>ones(Complex{Float64}, 6)))
+
+m21 = reshape([Bra(b_comp, bv(1,idx,jdx,0)) * embed_op * Ket(b_comp, bv(0,kdx,ldx,0))
+              for ((idx, jdx), (kdx, ldx)) in Iterators.product(all_idxs, all_idxs)], (6,6))
+@test isapprox(m21 / op.data[2,1], diagm(0=>ones(Complex{Float64}, 6)))
+
+m12 = reshape([Bra(b_comp, bv(0,idx,jdx,0)) * embed_op * Ket(b_comp, bv(1,kdx,ldx,0))
+              for ((idx, jdx), (kdx, ldx)) in Iterators.product(all_idxs, all_idxs)], (6,6))
+@test isapprox(m12 / op.data[1,2], diagm(0=>ones(Complex{Float64}, 6)))
+
+
+b_comp = b_comp⊗b_comp
+OP_test1 = dense(tensor([op1,one(b2),op,one(b1),one(b2),op1,one(b2)]...))
+OP_test2 = embed(b_comp, [1,[3,4],7], [op1,op,op1])
+@test isapprox(OP_test1.data, OP_test2.data)
+
+b8 = b2⊗b2⊗b2
+cnot = [1 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0]
+op_cnot = DenseOperator(b2⊗b2, cnot)
+OP_cnot = embed(b8, [1,3], op_cnot)
+@test ptrace(OP_cnot, [2])/2. == op_cnot
 
 @test_throws ErrorException QuantumOptics.operators.gemm!()
 @test_throws ErrorException QuantumOptics.operators.gemv!()