Implement embed functions for Gaussian unitaries, channels, and linear combinations

src/channels.jl

@@ -324,6 +324,96 @@ function _tensor(op1::GaussianChannel{B1,D1,T1}, op2::GaussianChannel{B2,D2,T2})
     return disp′′, transform′′, noise′′
 end
 
+##
+# Embedding Gaussian channels
+##
+
+"""
+    embed(basis::SymplecticBasis, idx::Int, op::GaussianChannel)
+    embed(basis::SymplecticBasis, indices::Vector{<:Int}, op::GaussianChannel)
+
+Embed a smaller Gaussian channel into a larger symplectic basis. The channel
+acts on the modes specified by `idx`/`indices` and leaves other modes unchanged.
+"""
+function embed(
+    basis::QuadPairBasis, index::Int, op::GaussianChannel{<:QuadPairBasis,D,S}
+) where {D,S}
+    return embed(basis, [index], op)
+end
+function embed(
+    basis::QuadPairBasis, indices::Vector{<:Int}, op::GaussianChannel{<:QuadPairBasis,D,S}
+) where {D,S}
+    @assert length(indices) == op.basis.nmodes "Number of indices must match number of modes in the channel"
+    @assert basis.nmodes ≥ length(indices) "Target basis must be large enough"
+
+    total_modes = basis.nmodes
+    sub_modes = op.basis.nmodes
+    disp = zeros(eltype(op.disp), 2 * total_modes)
+    transform = Matrix{eltype(op.transform)}(I, 2 * total_modes, 2 * total_modes)
+    noise = zeros(eltype(op.noise), 2 * total_modes, 2 * total_modes)
+
+    @inbounds for i in Base.OneTo(sub_modes)
+        idx = indices[i]
+        disp[2idx-1] = op.disp[2i-1]
+        disp[2idx] = op.disp[2i]
+    end
+    @inbounds for i in Base.OneTo(sub_modes)
+        idx_i = indices[i]
+        @inbounds for j in Base.OneTo(sub_modes)
+            idx_j = indices[j]
+            transform[2idx_i-1, 2idx_j-1] = op.transform[2i-1, 2j-1]
+            transform[2idx_i-1, 2idx_j] = op.transform[2i-1, 2j]
+            transform[2idx_i, 2idx_j-1] = op.transform[2i, 2j-1]
+            transform[2idx_i, 2idx_j] = op.transform[2i, 2j]
+
+            noise[2idx_i-1, 2idx_j-1] = op.noise[2i-1, 2j-1]
+            noise[2idx_i-1, 2idx_j] = op.noise[2i-1, 2j]
+            noise[2idx_i, 2idx_j-1] = op.noise[2i, 2j-1]
+            noise[2idx_i, 2idx_j] = op.noise[2i, 2j]
+        end
+    end
+    return GaussianChannel(basis, disp, transform, noise; ħ = op.ħ)
+end
+function embed(
+    basis::QuadBlockBasis, index::Int, op::GaussianChannel{<:QuadBlockBasis,D,S}
+) where {D,S}
+    return embed(basis, [index], op)
+end
+function embed(
+    basis::QuadBlockBasis, indices::Vector{<:Int}, op::GaussianChannel{<:QuadBlockBasis,D,S}
+) where {D,S}
+    @assert length(indices) == op.basis.nmodes "Number of indices must match number of modes in the channel"
+    @assert basis.nmodes ≥ length(indices) "Target basis must be large enough"
+
+    total_modes = basis.nmodes
+    sub_modes = op.basis.nmodes
+    disp = zeros(eltype(op.disp), 2 * total_modes)
+    transform = Matrix{eltype(op.transform)}(I, 2 * total_modes, 2 * total_modes)
+    noise = zeros(eltype(op.noise), 2 * total_modes, 2 * total_modes)
+
+    @inbounds for i in Base.OneTo(sub_modes)
+        idx = indices[i]
+        disp[idx] = op.disp[i]
+        disp[idx + total_modes] = op.disp[i + sub_modes]
+    end
+    @inbounds for i in Base.OneTo(sub_modes)
+        idx_i = indices[i]
+        @inbounds for j in Base.OneTo(sub_modes)
+            idx_j = indices[j]
+            transform[idx_i, idx_j] = op.transform[i, j]
+            transform[idx_i, idx_j + total_modes] = op.transform[i, j + sub_modes]
+            transform[idx_i + total_modes, idx_j] = op.transform[i + sub_modes, j]
+            transform[idx_i + total_modes, idx_j + total_modes] = op.transform[i + sub_modes, j + sub_modes]
+
+            noise[idx_i, idx_j] = op.noise[i, j]
+            noise[idx_i, idx_j + total_modes] = op.noise[i, j + sub_modes]
+            noise[idx_i + total_modes, idx_j] = op.noise[i + sub_modes, j]
+            noise[idx_i + total_modes, idx_j + total_modes] = op.noise[i + sub_modes, j + sub_modes]
+        end
+    end
+    return GaussianChannel(basis, disp, transform, noise; ħ = op.ħ)
+end
+
 """
     changebasis(::SymplecticBasis, state::GaussianChannel)
 

src/linearcombinations.jl

@@ -535,6 +535,24 @@ function ptrace(lc::GaussianLinearCombination, indices::AbstractVector{<:Int})
     return result
 end
 
+"""
+    embed(basis::SymplecticBasis, idx::Int, lc::GaussianLinearCombination)
+    embed(basis::SymplecticBasis, indices::Vector{<:Int}, lc::GaussianLinearCombination)
+
+Embed every state in a `GaussianLinearCombination` into a larger symplectic
+basis while preserving coefficients.
+"""
+function embed(basis::SymplecticBasis, index::Int, lc::GaussianLinearCombination)
+    return embed(basis, [index], lc)
+end
+function embed(basis::SymplecticBasis, indices::Vector{<:Int}, lc::GaussianLinearCombination)
+    @assert length(indices) == lc.basis.nmodes "Number of indices must match number of modes in the linear combination"
+    @assert basis.nmodes ≥ length(indices) "Target basis must be large enough"
+
+    embedded_states = [embed(basis, indices, state) for state in lc.states]
+    return GaussianLinearCombination(basis, copy(lc.coeffs), embedded_states)
+end
+
 """
     cross_wigner(state1::GaussianState, state2::GaussianState, x::AbstractVector)
 

src/unitaries.jl

@@ -625,6 +625,84 @@ function _tensor(op1::GaussianUnitary{B1,D1,S1}, op2::GaussianUnitary{B2,D2,S2})
     return disp′′, symp′′
 end
 
+##
+# Embedding Gaussian unitaries
+##
+
+"""
+    embed(basis::SymplecticBasis, idx::Int, op::GaussianUnitary)
+    embed(basis::SymplecticBasis, indices::Vector{<:Int}, op::GaussianUnitary)
+
+Embed a smaller Gaussian unitary into a larger symplectic basis. The unitary
+acts on the modes specified by `idx`/`indices` and the identity acts elsewhere.
+"""
+function embed(
+    basis::QuadPairBasis, index::Int, op::GaussianUnitary{<:QuadPairBasis,D,S}
+) where {D,S}
+    return embed(basis, [index], op)
+end
+function embed(
+    basis::QuadPairBasis, indices::Vector{<:Int}, op::GaussianUnitary{<:QuadPairBasis,D,S}
+) where {D,S}
+    @assert length(indices) == op.basis.nmodes "Number of indices must match number of modes in the unitary"
+    @assert basis.nmodes ≥ length(indices) "Target basis must be large enough"
+
+    total_modes = basis.nmodes
+    sub_modes = op.basis.nmodes
+    disp = zeros(eltype(op.disp), 2 * total_modes)
+    symp = Matrix{eltype(op.symplectic)}(I, 2 * total_modes, 2 * total_modes)
+
+    @inbounds for i in Base.OneTo(sub_modes)
+        idx = indices[i]
+        disp[2idx-1] = op.disp[2i-1]
+        disp[2idx] = op.disp[2i]
+    end
+    @inbounds for i in Base.OneTo(sub_modes)
+        idx_i = indices[i]
+        @inbounds for j in Base.OneTo(sub_modes)
+            idx_j = indices[j]
+            symp[2idx_i-1, 2idx_j-1] = op.symplectic[2i-1, 2j-1]
+            symp[2idx_i-1, 2idx_j] = op.symplectic[2i-1, 2j]
+            symp[2idx_i, 2idx_j-1] = op.symplectic[2i, 2j-1]
+            symp[2idx_i, 2idx_j] = op.symplectic[2i, 2j]
+        end
+    end
+    return GaussianUnitary(basis, disp, symp; ħ = op.ħ)
+end
+function embed(
+    basis::QuadBlockBasis, index::Int, op::GaussianUnitary{<:QuadBlockBasis,D,S}
+) where {D,S}
+    return embed(basis, [index], op)
+end
+function embed(
+    basis::QuadBlockBasis, indices::Vector{<:Int}, op::GaussianUnitary{<:QuadBlockBasis,D,S}
+) where {D,S}
+    @assert length(indices) == op.basis.nmodes "Number of indices must match number of modes in the unitary"
+    @assert basis.nmodes ≥ length(indices) "Target basis must be large enough"
+
+    total_modes = basis.nmodes
+    sub_modes = op.basis.nmodes
+    disp = zeros(eltype(op.disp), 2 * total_modes)
+    symp = Matrix{eltype(op.symplectic)}(I, 2 * total_modes, 2 * total_modes)
+
+    @inbounds for i in Base.OneTo(sub_modes)
+        idx = indices[i]
+        disp[idx] = op.disp[i]
+        disp[idx + total_modes] = op.disp[i + sub_modes]
+    end
+    @inbounds for i in Base.OneTo(sub_modes)
+        idx_i = indices[i]
+        @inbounds for j in Base.OneTo(sub_modes)
+            idx_j = indices[j]
+            symp[idx_i, idx_j] = op.symplectic[i, j]
+            symp[idx_i, idx_j + total_modes] = op.symplectic[i, j + sub_modes]
+            symp[idx_i + total_modes, idx_j] = op.symplectic[i + sub_modes, j]
+            symp[idx_i + total_modes, idx_j + total_modes] = op.symplectic[i + sub_modes, j + sub_modes]
+        end
+    end
+    return GaussianUnitary(basis, disp, symp; ħ = op.ħ)
+end
+
 """
     issymplectic(basis::SymplecticBasis, x::T)
 

test/test_channels.jl

@@ -156,4 +156,47 @@
         c1, c2 = coherentstate(qpairbasis, alpha1), coherentstate(qpairbasis, alpha2)
         @test (d1 ⊗ d2) * (v1 ⊗ v2) == c1 ⊗ c2
     end
+
+    @testset "embed" begin
+        α = rand(ComplexF64)
+        θ = rand(Float64)
+        n = rand(1:5)
+
+        for basis in (QuadPairBasis(1), QuadBlockBasis(1))
+            z = zeros(2 * basis.nmodes, 2 * basis.nmodes)
+            op = displace(basis, α, z)
+            full_basis = basis ⊕ basis ⊕ basis
+
+            s1 = coherentstate(basis, rand())
+            s2 = thermalstate(basis, rand(1:5))
+            s3 = vacuumstate(basis)
+            input_state = s1 ⊗ s2 ⊗ s3
+
+            embedded = embed(full_basis, 3, op)
+            expected = s1 ⊗ s2 ⊗ (op * s3)
+            @test embedded.basis == full_basis
+            @test embedded * input_state == expected
+        end
+
+        for basis in (QuadPairBasis(2), QuadBlockBasis(2))
+            z = zeros(2 * basis.nmodes, 2 * basis.nmodes)
+            single_basis = typeof(basis)(1)
+            op_two = attenuator(basis, θ, n)
+            full_basis = basis ⊕ single_basis
+
+            s1 = coherentstate(single_basis, rand())
+            s2 = thermalstate(single_basis, rand(1:5))
+            s3 = vacuumstate(single_basis)
+            input_state = s1 ⊗ s2 ⊗ s3
+
+            embedded_two = embed(full_basis, [1, 3], op_two)
+            transformed_13 = op_two * (s1 ⊗ s3)
+            result = embedded_two * input_state
+
+            @test ptrace(result, 2) == transformed_13
+            @test ptrace(result, [1, 3]) == s2
+        end
+
+        @test_throws AssertionError embed(QuadPairBasis(3), [1, 2, 3, 4], displace(QuadPairBasis(1), α, zeros(2, 2)))
+    end
 end

test/test_linearcombinations.jl

@@ -494,6 +494,47 @@
         @test lc_h1_scaled.ħ == 1
     end
 
+    @testset "Embedding" begin
+        α = rand(ComplexF64)
+
+        for basis in [qpairbasis, qblockbasis]
+            lc = GaussianLinearCombination(basis, [0.7, 0.3], [coherentstate(basis, α), vacuumstate(basis)])
+            big_basis = basis ⊕ basis
+            indices = collect(basis.nmodes + 1:2*basis.nmodes)
+            embedded = embed(big_basis, indices, lc)
+            expected_states = [embed(big_basis, indices, state) for state in lc.states]
+
+            @test embedded.basis == big_basis
+            @test embedded.coeffs == lc.coeffs
+            @test embedded.states == expected_states
+            @test embedded.ħ == lc.ħ
+        end
+
+        small_pair = QuadPairBasis(2)
+        pair_coeffs = [0.6, 0.4]
+        r_vec = rand(Float64, small_pair.nmodes)
+        θ_vec = rand(Float64, small_pair.nmodes)
+        pair_states = [coherentstate(small_pair, rand(ComplexF64, small_pair.nmodes)), squeezedstate(small_pair, r_vec, θ_vec)]
+        lc_pair = GaussianLinearCombination(small_pair, pair_coeffs, pair_states)
+        full_pair = small_pair ⊕ QuadPairBasis(1)
+        embedded_pair = embed(full_pair, [1, 3], lc_pair)
+        expected_pair_states = [embed(full_pair, [1, 3], state) for state in pair_states]
+        @test embedded_pair.coeffs == pair_coeffs
+        @test embedded_pair.states == expected_pair_states
+
+        small_block = QuadBlockBasis(2)
+        block_coeffs = [0.5, 0.5]
+        block_states = [thermalstate(small_block, fill(2, small_block.nmodes)), coherentstate(small_block, rand(ComplexF64, small_block.nmodes))]
+        lc_block = GaussianLinearCombination(small_block, block_coeffs, block_states)
+        full_block = small_block ⊕ QuadBlockBasis(1)
+        embedded_block = embed(full_block, [1, 3], lc_block)
+        expected_block_states = [embed(full_block, [1, 3], state) for state in block_states]
+        @test embedded_block.coeffs == block_coeffs
+        @test embedded_block.states == expected_block_states
+
+        @test_throws AssertionError embed(full_pair, [1], lc_pair)
+    end
+
     @testset "Static arrays compatibility" begin
         coh_static = coherentstate(SVector{2*nmodes}, SMatrix{2*nmodes,2*nmodes}, qpairbasis, 1.0)
         vac_static = vacuumstate(SVector{2*nmodes}, SMatrix{2*nmodes,2*nmodes}, qpairbasis)

test/test_unitaries.jl

@@ -110,4 +110,42 @@
         c1, c2 = coherentstate(qpairbasis, alpha1), coherentstate(qpairbasis, alpha2)
         @test (d1 ⊗ d2) * (v1 ⊗ v2) == c1 ⊗ c2
     end
+
+    @testset "embed" begin
+        α = rand(ComplexF64)
+        r, θ = rand(Float64), rand(Float64)
+
+        for basis in (QuadPairBasis(1), QuadBlockBasis(1))
+            op = displace(basis, α)
+            full_basis = basis ⊕ basis ⊕ basis
+
+            s1, s2, s3 = vacuumstate(basis), vacuumstate(basis), vacuumstate(basis)
+            input_state = s1 ⊗ s2 ⊗ s3
+
+            embedded = embed(full_basis, 2, op)
+            expected = s1 ⊗ (op * s2) ⊗ s3
+            @test embedded.basis == full_basis
+            @test embedded * input_state == expected
+        end
+
+        for basis in (QuadPairBasis(2), QuadBlockBasis(2))
+            single_basis = typeof(basis)(1)
+            op_two = squeeze(basis, r, θ)
+            full_basis = basis ⊕ single_basis
+
+            s1 = coherentstate(single_basis, rand())
+            s2 = thermalstate(single_basis, rand(1:5))
+            s3 = vacuumstate(single_basis)
+            input_state = s1 ⊗ s2 ⊗ s3
+
+            embedded_two = embed(full_basis, [1, 3], op_two)
+            transformed_13 = op_two * (s1 ⊗ s3)
+            result = embedded_two * input_state
+
+            @test ptrace(result, 2) == transformed_13
+            @test ptrace(result, [1, 3]) == s2
+        end
+
+        @test_throws AssertionError embed(QuadPairBasis(3), [1, 2, 3, 4], displace(QuadPairBasis(1), α))
+    end
 end