From f6d6cdec8d202b8c6f08e6cf0a7882ed8b47980b Mon Sep 17 00:00:00 2001
From: Hong-Ye Hu <hongyehu.physics@gmail.com>
Date: Thu, 17 Oct 2024 15:13:52 -0400
Subject: [PATCH] update Rydbergchain system and add test

---
 src/quantum_system_templates/rydberg.jl | 28 +++++++++++++++++--------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/quantum_system_templates/rydberg.jl b/src/quantum_system_templates/rydberg.jl
index e73efbd..cfc5132 100644
--- a/src/quantum_system_templates/rydberg.jl
+++ b/src/quantum_system_templates/rydberg.jl
@@ -28,13 +28,15 @@ function lift(x::Char,i::Int, N::Int)
     qubits[i] = x
     return join(qubits)
 end
-
 @doc raw"""
     RydbergChainSystem(;
+        N::Int=3, # number of atoms
         C::Float64=862690*2π,
         distance::Float64=10.0, # μm
         cutoff_order::Int=2, # 1 is nearest neighbor, 2 is next-nearest neighbor, etc.
         local_detune::Bool=false, # If true, include one local detuning pattern.
+        all2all::Bool=true, # If true, include all-to-all interactions.
+        ignore_Y_drive::Bool=false, # If true, ignore the Y drive. (In the experiments, X&Y drives are implemented by Rabi amplitude and its phase.)
     ) -> QuantumSystem
 
 Returns a `QuantumSystem` object for the Rydberg atom chain in the spin basis
@@ -46,10 +48,13 @@ H = \sum_i 0.5*\Omega_i(t)\cos(\phi_i(t)) \sigma_i^x - 0.5*\Omega_i(t)\sin(\phi_
 
 
 # Keyword Arguments
+- `N`: Number of atoms.
 - `C`: The Rydberg interaction strength in MHz*μm^6.
 - `distance`: The distance between atoms in μm.
 - `cutoff_order`: Interaction range cutoff, 1 is nearest neighbor, 2 is next nearest neighbor.
 - `local_detune`: If true, include one local detuning pattern.
+- `all2all`: If true, include all-to-all interactions.
+- `ignore_Y_drive`: If true, ignore the Y drive. (In the experiments, X&Y drives are implemented by Rabi amplitude and its phase.)
 """
 function RydbergChainSystem(;
     N::Int=3, # number of atoms
@@ -60,23 +65,24 @@ function RydbergChainSystem(;
     all2all::Bool=true,
     ignore_Y_drive::Bool=false,
 )
-    PAULIS = Dict("I" => [1 0; 0 1], "X" => [0 1; 1 0], "Y" => [0 -im; im 0], "Z" => [1 0; 0 -1], "n" => [0 0; 0 1])
+    PAULIS = Dict(:I => [1 0; 0 1], :X => [0 1; 1 0], :Y => [0 -im; im 0], :Z => [1 0; 0 -1], :n => [0 0; 0 1])
+    
     if all2all
         H_drift = zeros(ComplexF64, 2^N, 2^N)
         for gap in 0:N-2
             for i in 1:N-gap-1
                 H_drift += C * operator_from_string(
                     generate_pattern_with_gap(N, i, gap),
-                    PAULIS
+                    lookup = PAULIS
                 ) / ((gap + 1) * distance)^6
             end
         end
     else
         if cutoff_order == 1
-            H_drift = sum([C*operator_from_string(generate_pattern(N,i),PAULIS)/(distance^6) for i in 1:N-1])
+            H_drift = sum([C*operator_from_string(generate_pattern(N,i),lookup=PAULIS)/(distance^6) for i in 1:N-1])
         elseif cutoff_order == 2
-            H_drift = sum([C*operator_from_string(generate_pattern(N,i),PAULIS)/(distance^6) for i in 1:N-1])
-            H_drift += sum([C*operator_from_string(generate_pattern_with_gap(N,i,1),PAULIS)/((2*distance)^6) for i in 1:N-2])
+            H_drift = sum([C*operator_from_string(generate_pattern(N,i),lookup=PAULIS)/(distance^6) for i in 1:N-1])
+            H_drift += sum([C*operator_from_string(generate_pattern_with_gap(N,i,1),lookup=PAULIS)/((2*distance)^6) for i in 1:N-2])
         else
             error("Higher cutoff order not supported")
         end
@@ -85,17 +91,17 @@ function RydbergChainSystem(;
     H_drives = Matrix{ComplexF64}[]
 
     # Add global X drive
-    Hx = sum([0.5*operator_from_string(lift('X',i,N), PAULIS) for i in 1:N])
+    Hx = sum([0.5*operator_from_string(lift('X',i,N), lookup=PAULIS) for i in 1:N])
     push!(H_drives, Hx)
 
     if !ignore_Y_drive
         # Add global Y drive
-        Hy = sum([0.5*operator_from_string(lift('Y',i,N), PAULIS) for i in 1:N])
+        Hy = sum([0.5*operator_from_string(lift('Y',i,N), lookup=PAULIS) for i in 1:N])
         push!(H_drives, Hy)
     end
 
     # Add global detuning
-    H_detune = -sum([operator_from_string(lift('n',i,N), PAULIS) for i in 1:N])
+    H_detune = -sum([operator_from_string(lift('n',i,N), lookup=PAULIS) for i in 1:N])
     push!(H_drives, H_detune)
 
     params = Dict{Symbol, Any}(
@@ -114,3 +120,7 @@ function RydbergChainSystem(;
         params=params,
     )
 end
+
+@testitem "Rydberg system test" begin
+    @test RydbergChainSystem(N=3,cutoff_order=2,all2all=false) isa QuantumSystem
+end
\ No newline at end of file