Support qudits in YaoToEinsum

lib/YaoBlocks/src/composite/composite.jl

@@ -70,3 +70,4 @@ include("tag/tag.jl")
 include("tag/cache.jl")
 include("tag/dagger.jl")
 include("tag/scale.jl")
+include("tag/onlevels.jl")

lib/YaoBlocks/src/composite/tag/onlevels.jl

@@ -0,0 +1,28 @@
+export OnLevels
+
+"""
+    OnLevels{D, Ds, T <: AbstractBlock{Ds}} <: TagBlock{T, D}
+
+Define a gate that is applied to a subset of levels.
+
+### Fields
+- `gate`: the gate to be applied.
+- `levels`: the levels to apply the gate to.
+"""
+struct OnLevels{D, Ds, T <: AbstractBlock{Ds}} <: TagBlock{T, D}
+    gate::T
+    levels::NTuple{Ds, Int}
+    function OnLevels{D}(gate::T, levels::NTuple{Ds, Int}) where {D, Ds, T <: AbstractBlock{Ds}}
+        @assert nqudits(gate) == 1 "only single qubit gate is supported"
+        new{D, Ds, T}(gate, levels)
+    end
+end
+content(g::OnLevels) = g.gate
+function mat(::Type{T}, g::OnLevels{D, Ds}) where {T, D, Ds}
+    m = mat(T, g.gate)
+    I, J, V = LuxurySparse.findnz(m)
+    return sparse(collect(g.levels[I]), collect(g.levels[J]), V, D, D)
+end
+PropertyTrait(::OnLevels) = PreserveAll()
+Base.adjoint(x::OnLevels{D}) where D = OnLevels{D}(adjoint(x.gate), x.levels)
+Base.copy(x::OnLevels{D}) where D = OnLevels{D}(copy(x.gate), x.levels)

lib/YaoBlocks/src/layout.jl

@@ -295,6 +295,10 @@ print_annotation(io::IO, c::Daggered) =
 print_annotation(io::IO, c::CachedBlock) =
     printstyled(io, "[cached] "; bold = true, color = :yellow)
 
+function print_annotation(io::IO, x::OnLevels)
+    printstyled(io, "[on levels: ", x.levels, "] "; bold = true, color = :yellow)
+end
+
 function print_annotation(io::IO, x::Scale)
     if x.alpha == im
         printstyled(io, "[+im] "; bold = true, color = :yellow)

lib/YaoBlocks/test/composite/tag.jl

@@ -107,4 +107,21 @@ end
     @test gatecount(Val(0.1) * X) == Dict(typeof(Val(0.1)*X)=>1)
     @test gatecount(cache(Val(0.1) * X)) == Dict(typeof(Val(0.1)*X)=>1)
     @test gatecount(Daggered(X)) == Dict(typeof(Daggered(X))=>1)
-end
+end
+
+@testset "OnLevels" begin
+    Z1r = OnLevels{3}(Z, (2, 3))
+    mz = mat(ComplexF64, Z1r)
+    @test YaoBlocks.SparseArrays.nnz(mz) == 2
+    g = OnLevels{3}(Rx(0.5), (2, 3))
+    g2 = copy(g)
+    g2.gate.theta = 0.6
+    @test !ishermitian(g)
+    @test isunitary(g)
+    @test g2.gate.theta == 0.6
+    @test g.gate.theta == 0.5
+    @test parameters(g) == [0.5]
+    @test g' == OnLevels{3}(Rx(-0.5), (2, 3))
+    println(g)
+end
+

lib/YaoToEinsum/src/circuitmap.jl

@@ -1,4 +1,6 @@
-struct EinBuilder{T}
+# T is the element type of the tensor network
+# D is the dimension of the qudits
+struct EinBuilder{T, D}
     slots::Vector{Int}
     labels::Vector{Vector{Int}}
     tensors::Vector{AbstractArray{T}}
@@ -12,29 +14,29 @@ function add_tensor!(eb::EinBuilder{T}, tensor::AbstractArray{T,N}, labels::Vect
     push!(eb.labels, labels)
 end
 
-function EinBuilder(::Type{T}, n::Int) where T
-    EinBuilder(collect(1:n), Vector{Int}[], AbstractArray{T}[], Ref(n))
+function EinBuilder{T, D}(n::Int) where {T, D}
+    EinBuilder{T, D}(collect(1:n), Vector{Int}[], AbstractArray{T}[], Ref(n))
 end
 newlabel!(eb::EinBuilder) = (eb.maxlabel[] += 1; eb.maxlabel[])
 
 function add_gate!(eb::EinBuilder{T}, b::PutBlock{D,C}) where {T,D,C}
     return add_matrix!(eb, C, mat(T, b.content), collect(b.locs))
 end
 # general and diagonal gates
-function add_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector) where T
+function add_matrix!(eb::EinBuilder{T, D}, k::Int, m::AbstractMatrix, locs::Vector) where {T, D}
     if isdiag(m)
-        add_tensor!(eb, reshape(Vector{T}(diag(m)), fill(2, k)...), eb.slots[locs])
+        add_tensor!(eb, reshape(Vector{T}(diag(m)), fill(D, k)...), eb.slots[locs])
     elseif m isa YaoBlocks.OuterProduct  # low rank
         nlabels = [newlabel!(eb) for _=1:k]
         K = rank(m)
         if K == 1  # projector
-            add_tensor!(eb, reshape(Vector{T}(m.right), fill(2, k)...), [eb.slots[locs]...])
-            add_tensor!(eb, reshape(Vector{T}(m.left), fill(2, k)...), [nlabels...])
+            add_tensor!(eb, reshape(Vector{T}(m.right), fill(D, k)...), [eb.slots[locs]...])
+            add_tensor!(eb, reshape(Vector{T}(m.left), fill(D, k)...), [nlabels...])
             eb.slots[locs] .= nlabels
         else
             midlabel = newlabel!(eb)
-            add_tensor!(eb, reshape(Matrix{T}(m.right), fill(2, k)..., K), [eb.slots[locs]..., midlabel])
-            add_tensor!(eb, reshape(Matrix{T}(m.left), fill(2, k)..., K), [nlabels..., midlabel])
+            add_tensor!(eb, reshape(Matrix{T}(m.right), fill(D, k)..., K), [eb.slots[locs]..., midlabel])
+            add_tensor!(eb, reshape(Matrix{T}(m.left), fill(D, k)..., K), [nlabels..., midlabel])
             eb.slots[locs] .= nlabels
         end
     else
@@ -45,31 +47,31 @@ function add_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector)
     return eb
 end
 # swap gate
-function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,2,ConstGate.SWAPGate}) where {T}
+function add_gate!(eb::EinBuilder{T, 2}, b::PutBlock{2,2,ConstGate.SWAPGate}) where {T}
     lj = eb.slots[b.locs[2]]
     eb.slots[b.locs[2]] = eb.slots[b.locs[1]]
     eb.slots[b.locs[1]] = lj
     return eb
 end
 
 # projection gate, todo: generalize to arbitrary low rank gate
-function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,1,ConstGate.P0Gate}) where {T}
+function add_gate!(eb::EinBuilder{T, 2}, b::PutBlock{2,1,ConstGate.P0Gate}) where {T}
     add_matrix!(eb, 1, YaoBlocks.OuterProduct(T[1, 0], T[1, 0]), collect(b.locs))
     return eb
 end
 
 # projection gate, todo: generalize to arbitrary low rank gate
-function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,1,ConstGate.P1Gate}) where {T}
+function add_gate!(eb::EinBuilder{T, 2}, b::PutBlock{2,1,ConstGate.P1Gate}) where {T}
     add_matrix!(eb, 1, YaoBlocks.OuterProduct(T[0, 1], T[0, 1]), collect(b.locs))
     return eb
 end
 
 
 # control gates
-function add_gate!(eb::EinBuilder{T}, b::ControlBlock{BT,C,M}) where {T, BT,C,M}
+function add_gate!(eb::EinBuilder{T, 2}, b::ControlBlock{BT,C,M}) where {T, BT,C,M}
     return add_controlled_matrix!(eb, M, mat(T, b.content), collect(b.locs), collect(b.ctrl_locs), collect(b.ctrl_config))
 end
-function add_controlled_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector, control_locs, control_vals) where T
+function add_controlled_matrix!(eb::EinBuilder{T, 2}, k::Int, m::AbstractMatrix, locs::Vector, control_locs, control_vals) where T
     if length(control_locs) == 0
         return add_matrix!(eb, k, m, locs)
     end
@@ -169,24 +171,24 @@ Read-write complexity: 2^6.0
 """
 function yao2einsum(circuit::AbstractBlock{D}; initial_state::Dict=Dict{Int,Int}(), final_state::Dict=Dict{Int,Int}(), optimizer=TreeSA()) where {D}
     T = promote_type(ComplexF64, dict_regtype(initial_state), dict_regtype(final_state), YaoBlocks.parameters_eltype(circuit))
-    vec_initial_state = Dict{keytype(initial_state),ArrayReg{D,T}}([k=>render_single_qubit_state(T, v) for (k, v) in initial_state])
-    vec_final_state = Dict{keytype(final_state),ArrayReg{D,T}}([k=>render_single_qubit_state(T, v) for (k, v) in final_state])
+    vec_initial_state = Dict{keytype(initial_state),ArrayReg{D,T}}([k=>render_single_qudit_state(T, D, v) for (k, v) in initial_state])
+    vec_final_state = Dict{keytype(final_state),ArrayReg{D,T}}([k=>render_single_qudit_state(T, D, v) for (k, v) in final_state])
     yao2einsum(circuit, vec_initial_state, vec_final_state, optimizer)
 end
 dict_regtype(d::Dict) = promote_type(_regtype.(values(d))...)
 _regtype(::ArrayReg{D,VT}) where {D,VT} = VT
 _regtype(::Int) = ComplexF64
-render_single_qubit_state(::Type{T}, x::Int) where T = x == 0 ? zero_state(T, 1) : product_state(T, bit"1")
-render_single_qubit_state(::Type{T}, x::ArrayReg) where T = ArrayReg(collect(T, statevec(x)))
+render_single_qudit_state(::Type{T}, D, x::Int) where T = product_state(T, DitStr{D}([x]))
+render_single_qudit_state(::Type{T}, D, x::ArrayReg) where T = ArrayReg{D}(collect(T, statevec(x)))
 
 function yao2einsum(circuit::AbstractBlock{D}, initial_state::Dict{Int,<:ArrayReg{D,T}}, final_state::Dict{Int,<:ArrayReg{D,T}}, optimizer) where {D,T}
     v_initial_state = Dict{Vector{Int}, ArrayReg{D,T}}([[k]=>v for (k, v) in initial_state])
     v_final_state = Dict{Vector{Int}, ArrayReg{D, T}}([[k]=>v for (k, v) in final_state])
     yao2einsum(circuit, v_initial_state, v_final_state, optimizer)
 end
 function yao2einsum(circuit::AbstractBlock{D}, initial_state::Dict{Vector{Int},<:ArrayReg{D,T}}, final_state::Dict{Vector{Int},<:ArrayReg{D,T}}, optimizer) where {D,T}
-    n = nqubits(circuit)
-    eb = EinBuilder(T, n)
+    n = nqudits(circuit)
+    eb = EinBuilder{T, D}(n)
     openindices = add_states!(eb, initial_state)
     add_gate!(eb, circuit)
     openindices2 = add_states!(eb, final_state; conjugate=true)
@@ -199,7 +201,7 @@ function check_state_spec(state::Dict{Vector{Int},<:ArrayReg{D,T}}, n::Int) wher
     @assert all(1 .<= iks .<= n) "state qubit indices must be in the range 1 to $n"
     return iks
 end
-function add_states!(eb::EinBuilder{T}, states::Dict; conjugate=false) where {T}
+function add_states!(eb::EinBuilder{T, D}, states::Dict; conjugate=false) where {T, D}
     n = nqubits(eb)
     unique_indices = check_state_spec(states, n)
     openindices = eb.slots[setdiff(1:n, unique_indices)]

lib/YaoToEinsum/test/circuitmap.jl

@@ -317,7 +317,7 @@ end
     inner = (2,3)
     focus!(reg, inner)
     for final_state in [Dict([i=>rand_state(1) for i in inner]), Dict([i=>1 for i in inner])]
-        freg = join(YaoToEinsum.render_single_qubit_state(ComplexF64, final_state[3]), YaoToEinsum.render_single_qubit_state(ComplexF64, final_state[2]))
+        freg = join(YaoToEinsum.render_single_qudit_state(ComplexF64, 2, final_state[3]), YaoToEinsum.render_single_qudit_state(ComplexF64, 2, final_state[2]))
         net = yao2einsum(c; initial_state=initial_state, final_state=final_state, optimizer=TreeSA(nslices=3))
         println(net)
         @test vec(contract(net)) ≈ vec(statevec(freg)' * state(reg))
@@ -347,4 +347,25 @@ end
     code, xs = yao2einsum(c; initial_state=Dict([1, 2]=>reg1, [3, 4]=>reg2), final_state=Dict([1]=>reg3, [2,3,4]=>reg4))
     @test code(xs...; size_info=uniformsize(code, 2))[] ≈ join(reg4, reg3)' * join(reg2, reg1)
 end
+
+@testset "multi-level" begin
+    N2 = OnLevels{3}(ConstGate.P1, (2, 3))
+    X01 = OnLevels{3}(ConstGate.P1, (1, 2))
+    X12 = OnLevels{3}(ConstGate.P1, (2, 3))
+    function qaoa_circuit(nbits::Int, depth::Int)
+        n2 = chain([kron(nbits, i=>N2, i+1=>N2) for i=1:nbits-1])
+        x01 = chain([put(nbits, i=>X01) for i=1:nbits])
+        x12 = chain([put(nbits, i=>X12) for i=1:nbits])
+        return chain(repeat([n2, x01, x12], depth))
+    end
+
+    c = qaoa_circuit(5, 2)
+    op = repeat(5, X01)
+    extc = chain(c, op, c') 
+
+    res = yao2einsum(extc, initial_state=Dict(zip(1:5, zeros(Int, 5))), final_state=Dict(zip(1:5, zeros(Int, 5))))
+    @test res isa TensorNetwork
+    expected = expect(op, zero_state(ComplexF64, 5; nlevel=3) |> c)
+    @test res.code(res.tensors...; size_info=uniformsize(res.code, 3))[] ≈ expected
+end
 end

src/EasyBuild/hamiltonians.jl

@@ -30,22 +30,29 @@ function transverse_ising(nbit::Int, h::Number; periodic::Bool=true)
     ising_term + h*sum(map(i->put(nbit,i=>X), 1:nbit))
 end
 
-# a 3 level hamiltonian
+"""
+    rydberg_chain(nbits::Int; Ω::Number=0.0, Δ::Real=0.0, V::Real=0.0, r::Real=0.0)
+
+Rydberg chain hamiltonian defined as:
+```math
+\\sum_{i=1}^{n} \\Delta |r_i\\rangle\\langle r_i| + (\\Omega |1\\rangle\\langle r_i| + h.c.) + V n_i n_{i+1} + r (|0\\rangle\\langle 1| + h.c.)
+```
+where ``n`` is specified by `nbits`.
+"""
 function rydberg_chain(nbits::Int; Ω::Number=0.0, Δ::Real=0.0, V::Real=0.0, r::Real=0.0)
-    Pr = matblock(sparse([3], [3], [1.0+0im], 3, 3); nlevel=3, tag="|r⟩⟨r|")
-    Z1r = matblock(sparse([2, 3], [2, 3], [1.0+0im, -1.0], 3, 3); nlevel=3)
-    X1r = matblock(sparse([2, 3], [3, 2], [1.0+0im, 1.0], 3, 3); nlevel=3, tag="|1⟩⟨r| + |r⟩⟨1|")
-    X01 = matblock(sparse([1, 2], [2, 1], [1.0+0im, 1.0], 3, 3); nlevel=3, tag="|1⟩⟨0| + |0⟩⟨1|")
-    Y1r = matblock(sparse([3, 2], [2, 3], [1.0im, -1.0im], 3, 3); nlevel=3, tag="i|1⟩⟨0| - i|0⟩⟨1|")
+    Pr = OnLevels{3}(ConstGate.P1, (2, 3))
+    X1r = OnLevels{3}(X, (2, 3))
+    X01 = OnLevels{3}(X, (1, 2))
+    Y1r = OnLevels{3}(Y, (2, 3))
     # single site term in {|1>, |r>}.
     h = Add(nbits; nlevel=3)
     !isapprox(Δ, 0; atol=1e-12) && push!(h, (-Δ) * sum([put(nbits, i=>Pr) for i=1:nbits]))
-    #!iszero(Δ) && push!(h, (-Δ/2) * sum([put(nbits, i=>Z1r) for i=1:nbits]))
     !isapprox(real(Ω), 0; atol=1e-12) && push!(h, real(Ω)/2 * sum([put(nbits, i=>X1r) for i=1:nbits]))
     !isapprox(imag(Ω), 0; atol=1e-12) && push!(h, imag(Ω)/2 * sum([put(nbits, i=>Y1r) for i=1:nbits]))
     # interaction
     !isapprox(V, 0; atol=1e-12) && nbits > 1 && push!(h, V * sum([put(nbits, (i,i+1)=>kron(Pr, Pr)) for i=1:nbits-1]))
     # Raman term
     !isapprox(r, 0; atol=1e-12) && push!(h, r * sum([put(nbits, i=>X01) for i=1:nbits]))
     return h
-end
+end
+

test/easybuild/hamiltonians.jl

@@ -13,6 +13,11 @@ using YaoArrayRegister.SparseArrays
     @test vec(h[:,bit"01001000"] |> cleanup) ≈ mat(h)[:, buffer(bit"01001000")+1]
 end
 
+@testset "rydberg_chain" begin
+    h = rydberg_chain(3, Ω=1.0, Δ=1.0, V=1.0, r=1.0)
+    @test mat(ComplexF64, h) isa SparseMatrixCSC
+end
+
 # https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.123.170503
 # Acknology: Jonathan Wurtz and Madelyn Cain for extremely helpful discussion!
 @testset "Levine Pichler pulse" begin