AWS Braket: VQE

1import strangeworks.braket
2
3# general imports
4import numpy as np
5from scipy.optimize import minimize
6import matplotlib.pyplot as plt
7import time
8from datetime import datetime
9
10# Ensure consistent results
11np.random.seed(0)
12
13from braket.circuits import Circuit, Gate, AngledGate, QubitSet
14from braket.devices import LocalSimulator
15from braket.aws import AwsDevice, AwsSession
16
17# Set up device: Local Simulator
18device = LocalSimulator()
19## example code for other backends
20## choose the cloud-based managed simulator to run your circuit
21# device = AwsDevice("arn:aws:braket:::device/quantum-simulator/amazon/sv1")
22## choose the Rigetti device to run your circuit
23# device = AwsDevice("arn:aws:braket:::device/qpu/rigetti/Aspen-8")
24## choose the Ionq device to run your circuit
25# device = AwsDevice("arn:aws:braket:::device/qpu/ionq/ionQdevice")
26
27# Please enter the S3 bucket you created during onboarding in the code below
28my_bucket = f"amazon-braket-Your-Bucket-Name" # the name of the bucket
29my_prefix = "Your-Folder-Name" # the name of the folder in the bucket
30s3_folder = (my_bucket, my_prefix)
31
32# helper function to set up interaction term
33def get_ising_interactions(n_qubits):
34    """
35    function to setup Ising interaction term
36    """
37    # set number of qubits
38    ising = np.zeros((n_qubits, n_qubits))
39    # set nearest-neighbour interactions to nonzero values only
40    for ii in range(0, n_qubits-1):
41        ising[ii][ii+1] = -1
42    # add periodic boundary conditions
43    ising[0][n_qubits-1] = -1
44    print('Ising matrix:\n', ising)
45    
46    return ising
47
48    # function to build the VQE ansatz
49def circuit(params, n_qubits):
50    """
51    function to return full VQE circuit ansatz
52    input: parameter list with three parameters 
53    """
54
55    # instantiate circuit object
56    circuit = Circuit()
57    
58    # add Hadamard gate on first qubit
59    circuit.rz(0, params[0]).ry(0, params[1])
60
61    # apply series of CNOT gates
62    for ii in range(1, n_qubits):
63        circuit.cnot(control=0, target=ii)
64
65    # add parametrized single-qubit rotations around y
66    for qubit in range(n_qubits):
67        gate = Circuit().ry(qubit, params[2])
68        circuit.add(gate)
69
70    return circuit
71
72
73# function that computes cost function for given params
74def objective_function(params, b_field, n_qubits, n_shots, s3_folder, verbose=False):
75    """
76    objective function takes a list of variational parameters as input,
77    and returns the cost associated with those parameters
78    """
79    
80    global CYCLE
81    CYCLE += 1
82    
83    if verbose:
84        print('==================================' * 2)
85        print('Calling the quantum circuit. Cycle:', CYCLE)
86    
87    # obtain a quantum circuit instance from the parameters
88    vqe_circuit = circuit(params, n_qubits)
89    
90    # Computations are independent of one another, so can be triggered in parallel
91    # run the circuit on appropriate device
92    if device.name == 'DefaultSimulator':
93        task_zz = device.run(vqe_circuit, shots=n_shots)
94    else: 
95        task_zz = device.run(vqe_circuit, shots=n_shots, 
96                             s3_destination_folder=s3_folder, poll_timeout_seconds=3*24*60*60)
97
98    # Hb term: construct the circuit for measuring in the X-basis
99    H_on_all = Circuit().h(range(0, n_qubits))
100    vqe_circuit.add(H_on_all)
101    
102    # run the circuit (with H rotation at end)
103    if device.name == 'DefaultSimulator':
104        task_b = device.run(vqe_circuit, shots=n_shots)
105    else:
106        task_b = device.run(vqe_circuit, shots=n_shots, 
107                     s3_destination_folder=s3_folder, poll_timeout_seconds=3*24*60*60)
108
109    # Collect results from devices (wait for results, if necessary)
110    result_zz = task_zz.result()
111    result_b = task_b.result()
112    
113    # Compute Hzz term
114    # convert results (0 and 1) to ising (1 and -1)
115    meas_ising = result_zz.measurements
116    meas_ising[meas_ising == 1] = -1
117    meas_ising[meas_ising == 0] = 1
118    
119    # Hzz term: get all energies (for every shot): (n_shots, 1) vector
120    all_energies_zz = np.diag(np.dot(meas_ising, np.dot(J, np.transpose(meas_ising))))
121    
122    # Hzz term: get approx energy expectation value (factor 1/4 for Pauli vs spin 1/2) 
123    energy_expect_zz = 0.25*np.sum(all_energies_zz)/n_shots
124    
125    # Compute Hb term
126    # convert results (0 and 1) to ising (1 and -1)
127    meas_ising = result_b.measurements
128    meas_ising[meas_ising == 1] = -1
129    meas_ising[meas_ising == 0] = 1
130    
131    # Hb term: get all energies (for every shot): (n_shots, 1) vector
132    # factor 1/2 for Pauli vs spin 1/2
133    energy_expect_b = -1*b_field/2*meas_ising.sum()/n_shots
134    
135    # get total energy expectation value
136    energy_expect = energy_expect_zz + energy_expect_b
137    # per site
138    energy_expect_ind = energy_expect/n_qubits
139    
140    # print energy expectation value
141    if verbose:
142        print('Energy expectation value:', energy_expect)
143        print('Energy expectation value (per particle):', energy_expect_ind)
144
145    return energy_expect
146
147
148# The function to execute the training: run classical minimization.
149def train(b_field, options, n_qubits, n_shots, n_initial=10, s3_folder=None, verbose=False):
150    """
151    function to run VQE algorithm with several random seeds for initialization
152    """
153    print('Starting the training.')
154    
155    if verbose:
156        print('==================================' * 3)
157        print('Running VQE OPTIMIZATION.')
158    
159    # initialize vectors for results per random seed
160    cost_energy = []
161    angles = []
162    
163    # optimize for different random initializations: avoid local optima
164    for ii in range(n_initial):
165        
166        #print counter
167        if verbose:
168            run_init = ii+1
169            print('Running VQE OPTIMIZATION for random seed NUMBER', run_init)
170        
171        # randomly initialize variational parameters within appropriate bounds
172        params0 = np.random.uniform(0, 2 * np.pi, 3).tolist()
173        # set bounds for search space
174        bnds = [(0, 2 * np.pi) for _ in range(int(len(params0)))]
175
176        # run classical optimization
177        result = minimize(objective_function, params0, args=(b_field, n_qubits, n_shots, s3_folder, verbose), 
178                          options=options, method='SLSQP', bounds=bnds)
179
180        # store result of classical optimization
181        result_energy = result.fun
182        cost_energy.append(result_energy)
183        result_angle = result.x
184        angles.append(result_angle)
185        if verbose:
186            print('Optimal avg energy:', result_energy)
187            print('Optimal angles:', result_angle)
188        
189        # reset cycle count 
190        global CYCLE
191        CYCLE = 0
192    
193    # store energy minimum (over different initial configurations)
194    energy_min = np.min(cost_energy)
195    optim_angles = angles[np.argmin(cost_energy)]
196    if verbose:
197        print('Energy per initial seeds:', cost_energy)
198        print('Minimal energy:', energy_min)
199        print('Optimal variational angles:', optim_angles)
200    
201    return energy_min
202
203# visualize VQE circuit example
204N = 4
205params = [0.1, 0.2, 0.3]
206
207vqe_circuit = circuit(params, N)
208
209print('1. Printing VQE test circuit:')
210print(vqe_circuit)
211
212# Hb term: construct the circuit for measuring in the X-basis
213print('')
214print('2. Apply Hadamard to measure in x-basis:')
215H_on_all = Circuit().h(range(0, N))
216vqe_circuit.add(H_on_all)
217print(vqe_circuit)
218
219# set up the problem
220SHOTS = 1000
221N = 10 # number of qubits
222n_initial = 30 # number of random seeds to explore optimization landscape
223verbose = False # control amount of print output
224
225# set counters
226CYCLE = 0
227
228# set up Ising matrix with nearest neighbour interactions and PBC
229J = get_ising_interactions(N)
230
231# set options for classical optimization
232if verbose:
233    options = {'disp': True}
234else: 
235    options = {}
236
237# kick off training
238start = time.time()
239
240# parameter scan
241stepsize = 0.5
242xvalues = np.arange(0, 2+stepsize, stepsize)
243results = []
244results_site = []
245
246for bb in xvalues:
247    b_field = bb
248    print('Strength of magnetic field:', b_field)
249    energy_min = train(b_field, options=options, n_qubits=N, n_shots=SHOTS, n_initial=n_initial, 
250                       s3_folder=s3_folder, verbose=verbose)
251    results.append(energy_min)
252    results_site.append(energy_min/N)
253    
254    # reset counters
255    CYCLE = 0
256
257end = time.time()
258# print execution time
259print('Code execution time [sec]:', end - start)
260
261# print optimized results
262print('Optimal energies:', results)
263print('Optimal energies (per site):', results_site)
264
265fig = plt.figure()
266plt.plot(xvalues, results_site, 'm--o')
267plt.xlabel('transverse field $B [J]$')
268plt.ylabel('groundstate energy per site $E_{0}/N [J]$')
269strangeworks.plot(fig,transparent = False)
270
271# helper function to numerically solve for gs energy of TIM
272def num_integrate_gs(B):
273    """
274    numerically integrate exact band to get gs energy of TIM
275    this should give -E_0/(N*J) by Pfeufy
276    Here set J=1 (units of energy)
277    """
278    # lamba_ratio (setting J=1): compare thesis
279    ll = 1/(2*B)
280    
281    # set energy
282    gs_energy = 0
283    
284    # numerical integration
285    step_size = 0.0001
286    k_values = np.arange(0, np.pi, step_size)
287    integration_values = [step_size*np.sqrt(1 + ll**2 + 2*ll*np.cos(kk)) for kk in k_values]
288    integral = np.sum(integration_values)
289    gs_energy = 1*integral/(4*np.pi*ll)
290    
291    return gs_energy
292
293
294plt.clf()
295fig2 = plt.figure()
296# plot exact gs energy of TIM vs VQE results
297x = np.arange(0.01, 2.2, 0.2)
298y = [-1*num_integrate_gs(xx) for xx in x]
299# plot exact results
300plt.plot(x, y, 'b--s')
301# plot vqe results
302plt.plot(xvalues, results_site, 'm--o')
303plt.xlabel('magnetic field $B [J]$')
304plt.ylabel('groundstate energy $E_{0}/N [J]$')
305plt.tight_layout();
306strangeworks.plot(fig2,transparent=False)