Crossing a phase transition with Rydberg atoms in a distributed environment: a DMPSTraj emulation¶
We aim at reproducing numerically the experiment of Bernien et al, Nature 551, 579 (2017).
A 1D chain of Rydberg atoms is described by the hamiltonian: $$ H = \frac{\Omega}{2} \sum_{k=1}^N X_k - \Delta(t) \sum_{k=1}^N n_k + U \sum_{k=1}^{N-1} n_k n_{k+1} $$ with $n_k = (1 - Z_k) / 2$ the number operators. The dynamics is controlled by a Rabi frequency $\Omega = 2\pi \times 2 \text{MHz}$, a varying detuning $\Delta(t)$, and a nearest-neighboor Rydberg interaction $U = 2\pi \times 24 \text{MHz}$.
In addition, the atoms are submitted to a dephasing noise, i.e. jump operators $\sqrt{\Gamma}Z_k$ on each qubit $k$, with $\Gamma$ the dephasing rate.
This system undergo a ferromagnetic to antiferromagnetic phase transition when $\Delta$ is increased from $\Delta \ll -\Omega$ to $\Delta \gg \Omega$.
Going smoothly through the phase transition allows to prepare the antiferromagnetic state and to study the phase transition.
This is the experiment we reproduce in this notebook with the help of DMPSTraj and Qaptiva HPC.
%%writefile dmpstraj_rydberg.py
import numpy as np
import matplotlib.pyplot as plt
from qat.core import Observable, Term, Variable, Schedule, Job
from qat.core.variables import heaviside
from qat.hardware import HardwareModel
from qat.qpus import DMPSTraj
import json
Writing dmpstraj_rydberg.py
Definition of the problem¶
Let's first define the physical parameters of the problem. For this demo we will limit ourself to $N=7$ atoms (but feel free to experiment with more atoms after). This reproduces roughly the Fig. 3 of Bernien et al., showing the preparation of an antiferromagnetic ground state by crossing the phase transition.
%%writefile -a dmpstraj_rydberg.py
### physics parameters
N = 7 # number of qubits
Omega = 2 * 2 * np.pi # MHz
Delta = 10 * 2 * np.pi # MHz
U = 12 * 2 * np.pi # MHz
Gamma = 0.01 # MHz
T = 0.5 # us
duration_steps = [T, 4 * T, T]
Appending to dmpstraj_rydberg.py
and some utility functions:
%%writefile -a dmpstraj_rydberg.py
def plot_error(x, y, e, **kwargs):
assert np.all(e >= 0.0)
line = plt.plot(x, y, **kwargs)
color = line[0].get_color()
plt.fill_between(x, y - e, y + e, lw=0, alpha=0.3, color=color)
def smooth_transition_wf(val_i, val_f, t_i, t_f, n=3):
"""Waveform used to smoothly cross the transition"""
assert t_f > t_i
assert n % 2 == 1
t_var = 2 * (Variable("t") - t_i) / (t_f - t_i) - 1.
return val_i + 0.5 * (val_f - val_i) * (t_var**n + 1.)
Appending to dmpstraj_rydberg.py
As usually with analog quantum computing we define a Schedule corresponding to the time dependent hamiltonian applied to the qubit register.
%%writefile -a dmpstraj_rydberg.py
X = [Observable(N, pauli_terms=[Term(1., "X", [i])]) for i in range(N)]
n = [0.5 * Observable(N, pauli_terms=[Term(1., "I", [i]), Term(-1., "Z", [i])]) for i in range(N)]
inter = sum(n[i] * n[i+1] for i in range(N-1))
t = Variable('t')
t1, t2, t3 = np.cumsum(duration_steps)
delta_wf = -Delta * heaviside(t, 0, t1)
delta_wf += smooth_transition_wf(-Delta, Delta, t1, t2) * heaviside(t, t1, t2)
delta_wf += Delta * heaviside(t, t2, t3)
rabi_wf = Omega
sched = Schedule(drive=[(rabi_wf, sum(X) / 2.), (delta_wf, -sum(n)), (U, inter)], tmax=sum(duration_steps))
Appending to dmpstraj_rydberg.py
Building the QPU¶
Let's build a DMPSTraj QPU and introduce noise by defining a HardwareModel.
DMPSTraj requires a bond dimension $\chi$ to limit the memory footprint of MPS and thus the total computation time.
In this setup, a low bond dimension $\chi=4$ is enough to capture the dynamics.
In general, the bond dimension should be chosen such that the entanglement entropy $S$ such that $\chi \gg e^S$
A number of trajectories n_samples is also required to simulate the effect of noise.
As a rule of thumb, the more noise, the more trajectories are needed, but it depends ultimately on the observables, the schedule, and the type of noise.
A few hundreds is a good starting point, which can then be refined by looking at the statistical error provided in the result.
This error is obtained from the standard deviation between trajectories, and should decrease as $1/\sqrt{N_{\rm traj}}$.
Trajectories can be run in parallel on a number of processors given by n_procs.
A last preparation step before launching the calculation is the following. One can optionnaly provide a "monitor function", i.e. a function to be called along the time evolution which can compute different quantities of the state and return them in a dictionnary. This is only useful to get quantities that are not observables. Here we use it to get the probability of some basis states as a function of time.
Important REMARK: DMPSTraj contains a crucial attribute that must be defined everytime the class is to be used. It is n_nodes. Indeed, in order to distribute computations accros the n number of nodes. The reason why it was not defined here, is because it will done in the sbatch file below.
%%writefile -a dmpstraj_rydberg.py
hm = HardwareModel(jump_operators=[np.sqrt(Gamma) * Observable.z(i, N) for i in range(N)])
def monitor_function(state):
"""
This function will be called at regular time intervals to gather information about
the current state beyond observables such as probabilities.
"""
N = len(state)
out = {}
out["proba0"] = abs(state(0))**2
out["proba1"] = abs(state(1 + 2**(N-1)))**2
idx_antiferro = sum(2**k for k in range(0, N, 2))
out["proba_antiferro"] = [abs(state(idx_antiferro))**2, abs(state(idx_antiferro // 2))**2]
return out
qpu = DMPSTraj(hardware_model=hm,
bond_dimension=4,
n_steps=int(200 * sum(duration_steps)), # number of time steps in the trotterization
n_samples=10, # number of samples (i.e. trajectories)
n_procs=64, # number of parallel processes
monitor_times=np.linspace(0, np.sum(duration_steps), 200), # when to evaluate observables and the optional monitor function
monitor_function=monitor_function,
sim_method='number_continuous',
)
Appending to dmpstraj_rydberg.py
We build a job with a list of observables, here the Rydberg occupation on each qubit.
Observables are computed at the end of the schedule, but also (optionnaly) during the time evolution at times given to the QPU as monitor_times.
%%writefile -a dmpstraj_rydberg.py
obs = Observable.z(0, N)
obs_all = [0.5 * (1. - Observable.z(k, N)) for k in range(N)]
job = sched.to_job(job_type='OBS',
observable=Observable.z(0, N), # needs to be defined but will not be used
observables=[0.5 * (1. - Observable.z(k, N)) for k in range(N)],
psi_0='0' * N)
Appending to dmpstraj_rydberg.py
Simulation¶
The cell below is the one that that performs the simulation.
%%writefile -a dmpstraj_rydberg.py
res = qpu.submit(job)
print("run time:", res.meta_data["simulation_time"])
print("Observables:", res.values)
Appending to dmpstraj_rydberg.py
Time-dependent data is stored in the meta_data of the result, which must be deserialized. We also transform them to numpy arrays for convenience:
%%writefile -a dmpstraj_rydberg.py
res.meta_data = {key: json.loads(val) for key, val in res.meta_data.items()}
times = res.meta_data['times']
data = {key: np.squeeze(val) for key, val in res.meta_data['monitored_data'].items()}
error = {key: np.squeeze(val) for key, val in res.meta_data['error_data'].items()}
Appending to dmpstraj_rydberg.py
Plotting the results¶
In here, you create the plots to show the results of the simulation. Since they are executed along the rest of the code in the file dmpstraj_rydberg.py, plots are saved and will be shown later.
We plot the different terms of the Schedule $\Omega$, $\Delta(t)$ and $U$. $\Delta(t)$ has been made to reproduce approximately the experiment schedule, as seen in Fig. 3(a).
%%writefile -a dmpstraj_rydberg.py
for t in np.cumsum(duration_steps):
plt.axvline(t, c='k', ls=':', alpha=0.3)
for k in range(N):
plot_error(times, data["job_observables"][:, k], 2*error["job_observables"][:, k], label=f"k={k}")
plt.xlabel('time')
plt.ylabel(r'$\langle n_k\rangle$')
plt.legend()
plt.savefig('angle_time.png')
plt.close()
Appending to dmpstraj_rydberg.py
%%writefile -a dmpstraj_rydberg.py
plt.bar(np.arange(7), data["job_observables"][-1, :], yerr=2*error["job_observables"][-1, :])
plt.xlabel('qubit')
plt.ylabel(r'$\langle n \rangle$')
plt.savefig('angle_qubit.png')
plt.close()
Appending to dmpstraj_rydberg.py
%%writefile -a dmpstraj_rydberg.py
for t in np.cumsum(duration_steps):
plt.axvline(t, c='k', ls=':', alpha=0.3)
plot_error(times, data["proba0"], 2*error["proba0"], label='|0000000>')
plot_error(times, data["proba1"], 2*error["proba1"], label='|1000001>')
plot_error(times, data["proba_antiferro"][:, 0], 2*error["proba_antiferro"][:, 0], label='|1010101>')
plot_error(times, data["proba_antiferro"][:, 1], 2*error["proba_antiferro"][:, 1], label='|0101010>')
plt.legend()
plt.xlabel('time')
plt.ylabel('probability of state')
plt.savefig('prob_state_time.png')
Appending to dmpstraj_rydberg.py
Once the python script is written, it needs to be executed. And it is done in a distributed environment. The cell below is the one creating a bash script to be executed. What makes this bash script special is that it contains #SBATCH commands for the execution in a slurm environment. All the lines that start with a #SBATCH are the ones that are used for the configuration the slurm environment.
Important REMARK: As mentioned above, the argumet n_nodes is implicitly attributed when defining the number of nodes with the #SBATCH --nodes command. Once the number of nodes is defined using SBATCH it becomes an environment variable that is fetched by DMPSTraj.
Another command that might be puzzling is module load qaptiva-hpc. This command installs the required packages of the qaptiva-hpc library. Finally, the python3 dmpstraj_rydberg.py is the one executing the python script that was created all along the previous cells of this notebook.
%%file submit.sh
#!/bin/sh
#SBATCH --nodes 1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=64
#SBATCH --time 00:02:00
#SBATCH --job-name dmpstraj_example
#SBATCH --output slurm_output/dmpstraj.out
#SBATCH --error slurm_output/dmpstraj.err
#SBATCH --exclusive
# Load the python3 module to be used on the cluster
# module load python3
echo "Loading the qaptiva-hpc module environment"
module load qaptiva-hpc
# Load other modules if necessary such as openmpi
# module load openmpi
echo "Running the DMPSTraj simulation"
python3 dmpstraj_rydberg.py
Writing submit.sh
The file submit.sh is called a sbatch file. In order to execute it properly, it needs to be done by calling the sbatch command. This is what the cell below does. This command is specific to slurm environments. The --wait instruction allows to not exit before the job terminates. On the side, since there is a time limit to the execution of the job (see sbatch command in submit.sh with the command #SBATCH --time), the instruction --deadline=now+03minutes removes the job if the job did not end.
%%bash
sbatch --wait --deadline=now+03minutes submit.sh
Submitted batch job 1
Once the job is executed, the outputs (e.g. print statement of the script) are kept in the file dmpstraj.out (see the #SBATCH --output command in submit.sh).
%%bash
cat slurm_output/dmpstraj.out
Loading the qaptiva-hpc module environment Running the DMPSTraj simulation
Showing the saved plots from the executable¶
First, the time evolution of the Rydberg probability of each atom. Error bars show the statistical error obtained from the standard deviation over trajectories (here we use 2$\sigma$ error bars).
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
img = mpimg.imread('angle_time.png')
plt.imshow(img)
plt.axis('off')
plt.show()
We can also look at the Rydberg probability at the end of the schedule, and we clearly see the emergence of an antiferromagnetic state.
img = mpimg.imread('angle_qubit.png')
plt.imshow(img)
plt.axis('off')
plt.show()
Now we look at the probability of some of the most important canonical basis states. Observe how, from the two antiferromagnetic states ($|0101010\rangle$ and $|1010101\rangle$), the latter is favored due to boundary conditions.
img = mpimg.imread('prob_state_time.png')
plt.imshow(img)
plt.axis('off')
plt.show()
