From 1f71706c7c35a02b0d5674f821702bf92a9a39e0 Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Fri, 22 Nov 2024 23:38:55 +0800
Subject: [PATCH] new on-levels

---
 lib/YaoBlocks/src/composite/composite.jl    |  1 +
 lib/YaoBlocks/src/composite/tag/onlevels.jl | 28 +++++++++++++++++++++
 lib/YaoBlocks/src/layout.jl                 |  4 +++
 lib/YaoBlocks/test/composite/tag.jl         | 19 +++++++++++++-
 lib/YaoToEinsum/test/circuitmap.jl          | 21 ++++++++++++++++
 src/EasyBuild/hamiltonians.jl               | 23 +++++++++++------
 test/easybuild/hamiltonians.jl              |  5 ++++
 7 files changed, 92 insertions(+), 9 deletions(-)
 create mode 100644 lib/YaoBlocks/src/composite/tag/onlevels.jl

diff --git a/lib/YaoBlocks/src/composite/composite.jl b/lib/YaoBlocks/src/composite/composite.jl
index 59509e63..d8533f10 100644
--- a/lib/YaoBlocks/src/composite/composite.jl
+++ b/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")
\ No newline at end of file
diff --git a/lib/YaoBlocks/src/composite/tag/onlevels.jl b/lib/YaoBlocks/src/composite/tag/onlevels.jl
new file mode 100644
index 00000000..0b29b5ff
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/lib/YaoBlocks/src/layout.jl b/lib/YaoBlocks/src/layout.jl
index 9cfa764d..4753a497 100644
--- a/lib/YaoBlocks/src/layout.jl
+++ b/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)
diff --git a/lib/YaoBlocks/test/composite/tag.jl b/lib/YaoBlocks/test/composite/tag.jl
index cf6a2fbf..c33b12b4 100644
--- a/lib/YaoBlocks/test/composite/tag.jl
+++ b/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
\ No newline at end of file
+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
+
diff --git a/lib/YaoToEinsum/test/circuitmap.jl b/lib/YaoToEinsum/test/circuitmap.jl
index f75ea28c..5fababde 100644
--- a/lib/YaoToEinsum/test/circuitmap.jl
+++ b/lib/YaoToEinsum/test/circuitmap.jl
@@ -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
diff --git a/src/EasyBuild/hamiltonians.jl b/src/EasyBuild/hamiltonians.jl
index 616f528e..9ef0fe4a 100644
--- a/src/EasyBuild/hamiltonians.jl
+++ b/src/EasyBuild/hamiltonians.jl
@@ -30,17 +30,23 @@ 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
@@ -48,4 +54,5 @@ function rydberg_chain(nbits::Int; Ω::Number=0.0, Δ::Real=0.0, V::Real=0.0, r:
     # Raman term
     !isapprox(r, 0; atol=1e-12) && push!(h, r * sum([put(nbits, i=>X01) for i=1:nbits]))
     return h
-end
\ No newline at end of file
+end
+
diff --git a/test/easybuild/hamiltonians.jl b/test/easybuild/hamiltonians.jl
index 2389d848..ec1c13de 100644
--- a/test/easybuild/hamiltonians.jl
+++ b/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