Rydberg atoms: generating highly-entangled states in large atomic arrays

Generating high-fidelity GHZ states using Q-CTRL pulses

BOULDER OPAL enables you to optimize controls to achieve target states in your Rydberg atomic system. In this Application note, you'll learn to optimize the generation of GHZ states for an array of multiple atoms initially in the ground state. This notebook adopts an integration method using sparse Hamiltonians to deal with the exponentially increasing memory requirements to simulate and optimize the pulses as the number of atoms in the system increases. This choice allows the optimizations in this Application note to be executed on a personal computer.

This notebook contains preview features that are not currently available in BOULDER OPAL. Please contact us if you want a live demonstration for your hardware. Preview features are included through the package qctrlcore, which is not publicly available: cells with qctrlcore will not be executable.

Imports and initialization

# essential imports
import qctrlcore
import time
import json
import jsonpickle
import tensorflow as tf
import numpy as np
import scipy
import tensorflow.python.util.deprecation as deprecation
deprecation._PRINT_DEPRECATION_WARNINGS = False

# plotting imports
from qctrlvisualizer import get_qctrl_style
from qctrlvisualizer import plot_controls
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

# helper functions
def create_initial_state(N_qubits):
    """Create initial state."""

    initial_state = np.zeros([2**N_qubits, 1])
    initial_state[0] = 1.
    initial_state = tf.constant(initial_state, dtype=tf.dtypes.complex128)
    return initial_state

def create_target_state(N_qubits):
    """Create target state."""

    # indices of states |0101...> and |10101...>
    idx_even = sum([2**n for n in range(0, N_qubits, 2)])
    idx_odd = sum([2**n for n in range(1, N_qubits, 2)])

    target_state = np.zeros([2**N_qubits, 1])
    target_state[[idx_even, idx_odd]] = 1./np.sqrt(2.)
    target_state = tf.constant(target_state, dtype=tf.dtypes.complex128)
    return target_state

def get_durations(segments):
    """Get control segment durations."""
    
    return [total_time / segments] * segments

def get_recommended_k(N_qubits, segment_count):
    """
    Return the appropriate Krylov subspace dimension (k) for a given set of parameters, as 
    required for the integration. This is obtained for the Hamiltonian with the largest spectral range.
    """

    with tf.Graph().as_default():
        
        durations = get_durations(segment_count)

        # worst-case-scenario pulses
        omega_segments = qctrlcore.complex_tensor([omega_max]*segment_count)
        delta_segments = qctrlcore.complex_tensor([-delta_range]*segment_count)
        
        # calculate terms and put together the full Hamiltonian, get first value
        H_omega, H_delta, H_fixed = get_rydberg_hamiltonian_terms(N_qubits)
        worst_case_hamiltonian = qctrlcore.create_pwc_hamiltonian_from_sparse_modulated_terms(
            terms=[
                (qctrlcore.TensorPwc(durations, omega_segments), H_omega),
                (qctrlcore.TensorPwc(durations, delta_segments), H_delta),
                (qctrlcore.TensorPwc(durations, tf.ones_like(omega_segments)), H_fixed),
            ]
        ).get_segment(qctrlcore.int_tensor(0))[1]        
        
        # calculate spectral range
        spectral_range = qctrlcore.calculate_hamiltonian_spectral_range(worst_case_hamiltonian)
        
        # calculate the appropriate krylov k parameter
        k = qctrlcore.recommend_lanczos_k(
            error_tolerance=1e-6,
            spectral_range=spectral_range,
            total_time=total_time,
            time_step=total_time/segment_count,
        )
        
        with tf.compat.v1.Session() as session:
            k = session.run(k)

    return k

def load_var(file_name):
    """Return a variable from a json file."""
    
    f = open(file_name, 'r+')
    encoded = f.read()
    decoded = jsonpickle.decode(encoded)
    f.close()
    return decoded

Generating GHZ states using Q-CTRL pulses

Here you'll learn to set up and optimize controls to prepare GHZ states for different numbers of atoms.

Characterizing the dynamics of the Rydberg system

Consider a chain of $N$ atoms, each one modeled as a qubit, with $|0\rangle$ representing the ground state and $|1\rangle$ representing a Rydberg state. The total Hamiltonian of the system, as described by A. Omran et al., is given by

$$ H = \frac{\Omega(t)}{2} \sum_{i=1}^N \sigma_x^{(i)} - \sum_{i=1}^N (\Delta(t) +\delta_i) n_i + \sum_{i<j}^N \frac{V}{\vert i-j \vert^6}n_i n_j , $$

with $\sigma_x^{(i)} = \vert 0 \rangle \langle 1 \vert_i + \vert 1 \rangle \langle 0 \vert_i$, and $n_i = \vert 1\rangle\langle 1 \vert_i$.

The fixed system parameters are the interaction strength between excited atoms, $V= 2\pi \times 24$ MHz, and the local energy shifts $\delta_i$,

$$ \delta_i = 2 \pi \times \begin{cases} - 4.5 \ \mathrm{MHz} \quad \mathrm{for} \ i=1,N \\ 0 \quad \mathrm{otherwise} \end{cases} \qquad \mathrm{for} \ N \leq 8 , \\ \delta_i = 2 \pi \times \begin{cases} - 6.0 \ \mathrm{MHz} \quad \mathrm{for} \ i=1,N \\ - 1.5 \ \mathrm{MHz} \quad \mathrm{for} \ i=4,N-3 \\ 0 \quad \mathrm{otherwise} \end{cases} \qquad \mathrm{for} \ N > 8 . $$

The control parameters are the effective coupling strength to the Rydberg state and the detuning, $\Omega(t)$ and $\Delta(t)$ respectively. These can be manipulated during the system evolution, which is specified to take place over a total time $T = 1.1$ µs. For convenience, the Hamiltonian terms can be grouped as follows:

$$ H = \Omega(t) h_\Omega + \Delta(t) h_\Delta + H_{\rm fixed} $$

with

$$ h_\Omega = \frac{1}{2}\sum_{i=1}^N \sigma_x^{(i)} \qquad h_\Delta = - \sum_{i=1}^N n_i \qquad H_{\rm fixed} = \sum_{i<j}^N \frac{V}{\vert i-j \vert^6}n_i n_j - \sum_{i=1}^N \delta_i n_i . $$

In the following cell the system parameters are specified, along with a function that implements these Hamiltonian terms.

# total time for the evolution
total_time = 1.1e-6  # s

# interaction strength
interaction_strength = 24.e6 * (2*np.pi)  # Hz

# individual light shift for each qubit
def get_local_shifts(N_qubits):
    local_shifts = np.zeros((N_qubits,))
    if N_qubits <= 8:
        local_shifts[0] = - 4.5e6 * (2.*np.pi)  # Hz
        local_shifts[-1] = - 4.5e6 * (2.*np.pi)  # Hz
    else:
        local_shifts[0] = - 6.0e6 * (2.*np.pi)  # Hz
        local_shifts[3] = - 1.5e6 * (2.*np.pi)  # Hz
        local_shifts[-4] = - 1.5e6 * (2.*np.pi)  # Hz
        local_shifts[-1] = - 6.0e6 * (2.*np.pi)  # Hz
    return local_shifts

# function to generate the Hamiltonian terms as sparse matrices
def get_rydberg_hamiltonian_terms(N_qubits):
    
    local_shifts = get_local_shifts(N_qubits)

    h_omega_values = []
    h_omega_indices = []
    h_delta_values = []
    h_delta_indices = []
    h_fixed_values = []
    h_fixed_indices = []

    # for each computational basis state
    for j in range(2**N_qubits):
        # get binary representation
        s = np.binary_repr(j, width=N_qubits)
        state = np.array([int(m) for m in s])

        # global shift term (big delta)
        n_exc = np.sum(state)  # number of excited states
        h_delta_values.append(- n_exc)
        h_delta_indices.append([j,j])

        # omega term
        for i in range(N_qubits):
            # couple "state" and "state with qubit i flipped"
            state_p = state.copy()
            state_p[i] = 1 - state_p[i]
            h_omega_values.append(0.5)
            h_omega_indices.append([j, state_p.dot(2**np.arange(N_qubits)[::-1])])
            
        # local shift term (small deltas)
        h_fixed = - np.sum(local_shifts * state)

        # interaction term
        for d in range(1, N_qubits):
            # number of excited qubits at distance d
            n_d = np.dot(((state[:-d] - state[d:]) == 0), state[:-d])
            h_fixed += interaction_strength * n_d/d**6

        if h_fixed != 0.0:
            h_fixed_values.append(h_fixed)
            h_fixed_indices.append([j,j])


    H_omega = scipy.sparse.coo_matrix(
        (h_omega_values, np.array(h_omega_indices).T),
        shape=[2**N_qubits, 2**N_qubits]
    )

    H_delta = scipy.sparse.coo_matrix(
        (h_delta_values, np.array(h_delta_indices).T),
        shape=[2**N_qubits, 2**N_qubits]
    )
    
    H_fixed = scipy.sparse.coo_matrix(
        (h_fixed_values, np.array(h_fixed_indices).T),
        shape=[2**N_qubits, 2**N_qubits]
    )
    
    return H_omega, H_delta, H_fixed

Optimizing Q-CTRL pulses for a system with 6 atoms

Using the Hamiltonian described in the previous section, you can now proceed with the pulse optimization. With the system initially in its ground state,

$$ |\psi(t=0)\rangle = |0000\ldots 0\rangle , $$

the optimization seeks to maximize the fidelity:

$$ {\mathcal F} = \Big|\langle \psi_\mathrm{target}|\psi(t=T)\rangle\Big|^2 , $$

which is determined by the final occupation of the target state

$$ |\psi_\mathrm{target}\rangle = \frac{1}{\sqrt{2}} \Big( |0101\ldots 1\rangle + |1010\ldots 0\rangle \Big) . $$

In order to parametrize the optimization, the controls $\Omega (t)$ and $\Delta (t)$ are specified to be piecewise constant functions: they take constant values over segments with equal durations. The optimization is defined in the following cell. This procedure involves specifying the control variables, calculating the full Hamiltonian for the system, and calculating the resulting infidelity of the process.

# constraints for the control variables
omega_max = 5.0e6 * (2.*np.pi)  # Hz
delta_range = 15.0e6 * (2.*np.pi)  # Hz

def optimize_GHZ_state(
    N_qubits,
    segment_count=40,
    optimization_run_count=1,
    seed=None,
):

    # time the execution of the optimization
    start_time = time.time()

    # generate random seed if none is passed
    if seed == None:
        seed = np.random.randint(100000)

    # print parameters
    print("\n Running optimization with:")
    print(f"\t{N_qubits} qubits")
    print(f"\t{segment_count} segments")
    print(f"\t{optimization_run_count} optimization runs")
    print(f"\trandom seed = {seed}")
    
    # get appropriate krylov k parameter
    lanczos_k = get_recommended_k(
        N_qubits=N_qubits,
        segment_count=segment_count,
    )

    # define specification to optimize
    def create_optimization_specification(variable_factory):

        # physical parameters
        durations = get_durations(segment_count)
        local_shifts = get_local_shifts(N_qubits)

        # pulses in omega and delta
        omega_segments = variable_factory.create_bounded(
            count=segment_count,
            lower_bound=0.,
            upper_bound=omega_max)
        delta_segments = variable_factory.create_bounded(
            count=segment_count,
            lower_bound=-delta_range,
            upper_bound=+delta_range)

        # calculate Hamiltonian terms
        H_omega, H_delta, H_fixed = get_rydberg_hamiltonian_terms(N_qubits)
        
        # put together the full Hamiltonian
        hamiltonian = qctrlcore.create_pwc_hamiltonian_from_sparse_modulated_terms(
            terms=[
                (qctrlcore.TensorPwc(durations, omega_segments), H_omega),
                (qctrlcore.TensorPwc(durations, delta_segments), H_delta),
                (qctrlcore.TensorPwc(durations, tf.ones_like(omega_segments)), H_fixed),
            ]
        )

        # create initial and target states
        initial_state = create_initial_state(N_qubits)
        target_state = create_target_state(N_qubits)
        
        # evolve the initial state
        evolved_state = qctrlcore.pwc_lanczos_state_evolution(
            initial_state=initial_state,
            pwc_hamiltonian=hamiltonian,
            segment_subdivision_count=1,
            krylov_dimension=lanczos_k,
        )
        
        # calculate infidelity
        infidelity = qctrlcore.create_state_infidelity(
            evolved_state, qctrlcore.complex_tensor(target_state)
        )[0]        


        # create return output dict
        def output_callable(session): return {
            'N_qubits': N_qubits,
            'segment_count': segment_count,
            'optimization_run_count': optimization_run_count,
            'seed': seed,
            'infidelity': session.run(infidelity),
            'omega_segments': [{'duration': d, 'value': v}
                               for d, v in zip(durations,
                                               session.run(omega_segments))],
            'delta_segments': [{'duration': d, 'value': v}
                               for d, v in zip(durations,
                                               session.run(delta_segments))],
        }

        return qctrlcore.OptimizationSpecification(
            cost=infidelity,
            output_callable=output_callable,
        )

    # run optimization
    parameters = qctrlcore.optimize(
        create_optimization_specification,
        number_of_optimizations=optimization_run_count,
        seed=seed,
    )
    
    execution_time = time.time() - start_time

    # print the execution time and best infidelity obtained
    print(f"Execution time: {execution_time:.0f} seconds")
    print(f"Infidelity reached: {parameters['output']['infidelity']:.1e}")

    # return dictionary with optimized results
    return parameters['output']

The optimization can now be called in a single line for any number of atoms $N$. You can also use pre-saved optimization data by setting the flag use_saved_data=True. To examine the optimized pulses and their performance, consider the case of 6 atoms:

# run optimization or load data
use_saved_data = True
if use_saved_data == False:
    result = optimize_GHZ_state(N_qubits=6)
else:
    result = load_var('./resources/rydberg-atoms-generating-highly-entangled-states-in-large-atomic-arrays/Rydberg_6atoms_1optimization.json')
    print("\n Optimization results with:")
    print(f"\t{result['N_qubits']} qubits")
    print(f"\t{result['segment_count']} segments")
    print(f"\t{result['optimization_run_count']} optimization runs")
    print(f"\trandom seed = {result['seed']}")
    print(f"Infidelity reached: {result['infidelity']:.1e}")
 Optimization results with:
	6 qubits
	40 segments
	1 optimization runs
	random seed = 41814
Infidelity reached: 1.6e-05
plot_controls(plt.figure(), {"$\Omega$": result['omega_segments'], "$\Delta$": result['delta_segments']})
plt.show()

The above plots show the optimized pulses $\Omega (t)$ and $\Delta (t)$ that produce the final state.

The next cell sets up a function to plot the optimized population dynamics and the final state. This involves simulating the dynamics using the full Hamiltonian, including optimized control terms.

def plot_result(parameters):
    
    # segment subsampling for plots
    segment_subdivision_count = 10

    # get physical parameters for the simulation
    N_qubits = parameters['N_qubits']
    segment_count = parameters['segment_count']
    durations = get_durations(segment_count)
    local_shifts = get_local_shifts(N_qubits)
    omega_signal = parameters['omega_segments']
    delta_signal = parameters['delta_segments']
    omega_segments = np.array([sig['value'] for sig in omega_signal])
    delta_segments = np.array([sig['value'] for sig in delta_signal])

    # calculate terms and put together the full Hamiltonian
    H_omega, H_delta, H_fixed = get_rydberg_hamiltonian_terms(N_qubits)
    hamiltonian = qctrlcore.create_pwc_hamiltonian_from_sparse_modulated_terms(
        terms=[
            (qctrlcore.TensorPwc(durations, qctrlcore.complex_tensor(omega_segments)), H_omega),
            (qctrlcore.TensorPwc(durations, qctrlcore.complex_tensor(delta_segments)), H_delta),
            (qctrlcore.TensorPwc(durations, tf.ones_like(qctrlcore.complex_tensor(omega_segments))), H_fixed),
        ]
    )
    
    # get appropriate krylov k parameter 
    lanczos_k = get_recommended_k(
        N_qubits=N_qubits,
        segment_count=segment_count*segment_subdivision_count,
    )

    # create initial and target states
    initial_state = create_initial_state(N_qubits)
    target_state = create_target_state(N_qubits)
    
    # run simulation and store intermediate times and states
    with tf.compat.v1.Session() as session:
        times, states = session.run(
            qctrlcore.solve_lanczos_with_pwc_hamiltonian(
                pwc_hamiltonian=hamiltonian,
                initial_state=qctrlcore.complex_tensor(initial_state),
                segment_subdivision_count=segment_subdivision_count,
                krylov_dimension=lanczos_k,
                return_steps=True,
            )
        )
        target_state = session.run(target_state)

    # data to plot
    idx_even = sum([2**n for n in range(0, N_qubits, 2)])
    idx_odd = sum([2**n for n in range(1, N_qubits, 2)])
    densities_t = abs(np.squeeze(states))**2
    initial_density_t = densities_t[:, 0]
    final_density_t = densities_t[:, idx_odd] + densities_t[:, idx_even]
    target_density = np.abs(np.squeeze(target_state))**2

    # style and label definitions
    plt.style.use(get_qctrl_style())
    bar_labels = ['']*2**N_qubits
    bar_labels[idx_odd] = '|' + '10'*int(N_qubits/2) + '>'
    bar_labels[idx_even] = '|' + '01'*int(N_qubits/2) + '>'

    def bar_plot(ax, density, title):
        ax.bar(np.arange(2**N_qubits), density, edgecolor="#680CE9",
               color="#680CE94C", tick_label=bar_labels, linewidth=2)
        ax.set_title(title, fontsize=14)
        ax.set_ylabel('Population', fontsize=14)
        ax.tick_params(labelsize=14)

    gs = gridspec.GridSpec(2, 2, hspace=0.3)
    fig = plt.figure()
    fig.set_figheight(2 * 5)
    fig.set_figwidth(15)
    fig.suptitle('Optimized state preparation', fontsize=16, y=0.95)
    
    # plot optimized final state and target state
    ax = fig.add_subplot(gs[0, 0])
    bar_plot(ax, densities_t[-1], 'Optimized final state')
    
    ax = fig.add_subplot(gs[0, 1])
    bar_plot(ax, target_density, 'Target GHZ state')

    # plot time evolution of basis state population
    ax = fig.add_subplot(gs[1, :])
    ax.plot(times, densities_t, "#AAAAAA", linewidth=0.5,)
    ax.plot(times, initial_density_t, "#FB00A5", linewidth=1.5,)
    ax.plot(times, final_density_t, "#680CE9", linewidth=2,)
    ax.set_xticks(np.linspace(0, 1e-6, 6))
    ax.set_xticklabels(['0', '200 n', '400 n', '600 n', '800 n', '1 µ'])
    ax.tick_params(labelsize=14)
    ax.set_xlabel('Time (s)', fontsize=14)
    ax.set_ylabel('Population', fontsize=14)
    ax.set_title('Basis states (gray), ground state (pink), and target state (purple)', fontsize=14)
plot_result(result)

The plots show: (top) Final state obtained with the optimized pulses (left) and target state (right). These states are indistinguishable, with infidelity below $10^{-4}$. (bottom) Time evolution of the populations during the optimized pulse sequence, for all basis states (gray), the ground state (pink), and the target state (purple). Note that since the target GHZ state is a superposition of two basis states, these two also show in gray as having a population of 0.5 at the end of the pulse.

Optimizing Q-CTRL pulses for up to 12 atoms

You can proceed in the same way as in the previous section to optimize the GHZ generation for different numbers of atoms.

# run optimization or load data
use_saved_data = True

# list to store the results of all runs
optimization_results = []

# for different number of qubits...
for n in [4,6,8,10,12]:
    
    # run or load an optimization
    if use_saved_data == False:
        result = optimize_GHZ_state(
            N_qubits=n,
            segment_count=40,
            optimization_run_count=3,
        )
    else:
        result = load_var('./resources/rydberg-atoms-generating-highly-entangled-states-in-large-atomic-arrays/Rydberg_'+str(n)+'atoms_3optimizations.json')
        print("\n Optimization results with:")
        print(f"\t{result['N_qubits']} qubits")
        print(f"\t{result['segment_count']} segments")
        print(f"\t{result['optimization_run_count']} optimization runs")
        print(f"\trandom seed = {result['seed']}")
        print(f"Infidelity reached: {result['infidelity']:.1e}")
    
    # store in the list to plot below
    optimization_results.append(result)
 Optimization results with:
	4 qubits
	40 segments
	3 optimization runs
	random seed = 31307
Infidelity reached: 4.0e-07

 Optimization results with:
	6 qubits
	40 segments
	3 optimization runs
	random seed = 64911
Infidelity reached: 3.7e-05

 Optimization results with:
	8 qubits
	40 segments
	3 optimization runs
	random seed = 99656
Infidelity reached: 1.3e-03

 Optimization results with:
	10 qubits
	40 segments
	3 optimization runs
	random seed = 65687
Infidelity reached: 7.3e-03

 Optimization results with:
	12 qubits
	40 segments
	3 optimization runs
	random seed = 85798
Infidelity reached: 1.2e-02

Plots of the final states and populations are similar to those for $N=6$ above, so here just the final state fidelities are displayed for different atom numbers.

# get infidelity and number of qubits of each run
infidelities = []
Nqs = []
for parameters in optimization_results:
    Nq, infid = parameters["N_qubits"], parameters["infidelity"]
    Nqs.append(Nq)
    infidelities.append(infid)
Nqs = np.array(Nqs)
infidelities = np.array(infidelities)

# plot fidelities/infidelities vs number of qubits
gs = gridspec.GridSpec(1, 2)
fig = plt.figure()
fig.set_figheight(5)
fig.set_figwidth(15)
fig.suptitle('Optimization performance', fontsize=16, y=1.)

ax = fig.add_subplot(gs[0, 0])
ax.plot(Nqs, 1-infidelities, 'o-', linewidth=2)
ax.set_title('Fidelities', fontsize=14)
ax.set_xlabel('Number of atoms', fontsize=14)
ax.tick_params(labelsize=14)

ax = fig.add_subplot(gs[0, 1])
ax.semilogy(Nqs, infidelities, 'o-', linewidth=2)
ax.set_title('Infidelities (log scale)', fontsize=14)
ax.set_xlabel('Number of atoms', fontsize=14)
ax.tick_params(labelsize=14)
plt.show()

The above figures display the GHZ state generation fidelities (left) and infidelities (right) achieved for different numbers of atoms. Although it is generally harder to obtain better infidelities for larger atomic arrays, using Q-CTRL optimized pulses you can obtain fidelities of over 0.99.