Update braket static demos; Update torch dependency (#1148)

**Summary:** * Update `ahs_aquila` and `braket_parallel_gradients` demos to be up to date with PennyLane 0.37. * Update `pyproject.toml` to install torch 2.1.2. The old version (1.13) was quite old. I tried changing the version to 2.3.1 which requires `typing-extensions>=4.8`, but that caused a dependency conflict with tensorflow which required `typing-extensions<4.6` **Relevant references:** **Possible Drawbacks:** **Related GitHub Issues:** Co-authored-by: Astral Cai <[email protected]>
PennyLaneAI · Jul 5, 2024 · 13fc290 · 13fc290
1 parent bc07b81
commit 13fc290
Show file tree

Hide file tree

Showing 4 changed files with 126 additions and 204 deletions.
diff --git a/demonstrations/ahs_aquila.py b/demonstrations/ahs_aquila.py
@@ -254,7 +254,7 @@
 #      {'area': {'width': Decimal('0.000075'), 'height': Decimal('0.000076')},
 #       'geometry': {'spacingRadialMin': Decimal('0.000004'),
 #        'spacingVerticalMin': Decimal('0.000004'),
-#        'positionResolution': Decimal('1E-7'),
+#        'positionResolution': Decimal('1E-8'),
 #        'numberSitesMax': 256}}
 #
 # We can see that the atom field has a width of :math:`75 \, \mu m` and a height of :math:`76 \, \mu m`.
@@ -289,6 +289,7 @@
 plt.scatter([x for x, y in coordinates], [y for x, y in coordinates])
 plt.xlabel("μm")
 plt.ylabel("μm")
+plt.show()
 
 ##############################################################################
 # .. figure:: ../_static/demonstration_assets/ahs_aquila/rydberg_blockade_coordinates.png
@@ -347,16 +348,17 @@
 #       'rydbergGlobal': {'rabiFrequencyRange': (Decimal('0.0'),
 #         Decimal('15800000.0')),
 #        'rabiFrequencyResolution': Decimal('400.0'),
-#        'rabiFrequencySlewRateMax': Decimal('250000000000000.0'),
+#        'rabiFrequencySlewRateMax': Decimal('400000000000000.0'),
 #        'detuningRange': (Decimal('-125000000.0'), Decimal('125000000.0')),
 #        'detuningResolution': Decimal('0.2'),
-#        'detuningSlewRateMax': Decimal('2500000000000000.0'),
+#        'detuningSlewRateMax': Decimal('6000000000000000.0'),
 #        'phaseRange': (Decimal('-99.0'), Decimal('99.0')),
 #        'phaseResolution': Decimal('5E-7'),
 #        'timeResolution': Decimal('1E-9'),
 #        'timeDeltaMin': Decimal('5E-8'),
 #        'timeMin': Decimal('0.0'),
-#        'timeMax': Decimal('0.000004')}}
+#        'timeMax': Decimal('0.000004')},
+#       'rydbergLocal': None}
 #
 # It is important to note that these quantities are in radians per second rather than Hz where relevant, and
 # are all in SI units. This means that for amplitude and detuning, we will need to convert from angular
@@ -436,7 +438,7 @@ def angular_SI_to_MHz(angular_SI):
 
 
 def gaussian_fn(p, t):
-    return p[0] * jnp.exp(-((t-p[1])**2) / (2*p[2]**2))
+    return p[0] * jnp.exp(-((t - p[1]) ** 2) / (2 * p[2] ** 2))
 
 
 # Visualize pulse, time in μs
@@ -453,6 +455,7 @@ def gaussian_fn(p, t):
 plt.ylabel("Amplitude [MHz]")
 
 plt.plot(time, y)
+plt.show()
 
 ##############################################################################
 #
@@ -467,10 +470,7 @@ def gaussian_fn(p, t):
 # We can then define our drive using via :func:`~pennylane.pulse.rydberg_drive`:
 #
 
-global_drive = qml.pulse.rydberg_drive(amplitude=gaussian_fn,
-                                       phase=0,
-                                       detuning=0,
-                                       wires=[0, 1, 2])
+global_drive = qml.pulse.rydberg_drive(amplitude=gaussian_fn, phase=0, detuning=0, wires=[0, 1, 2])
 
 ######################################################################
 # With only amplitude as non-zero, the overall driven Hamiltonian in this case simplifies to:
@@ -606,7 +606,7 @@ def circuit(params):
 start_val = amplitude[0]
 stop_val = amplitude[-1]
 max_val = np.max(amplitude)
-max_rate = np.max([(amplitude[i + 1] - amplitude[i]) / timestep for i in range(len(times)-1)])
+max_rate = np.max([(amplitude[i + 1] - amplitude[i]) / timestep for i in range(len(times) - 1)])
 
 print(f"start value: {start_val:.3} MHz")
 print(f"stop value: {stop_val:.3} MHz")
@@ -636,10 +636,7 @@ def circuit(params):
 #
 
 amp_fn = qml.pulse.rect(gaussian_fn, windows=[0.01, 1.749])
-global_drive = qml.pulse.rydberg_drive(amplitude=amp_fn,
-                                       phase=0,
-                                       detuning=0,
-                                       wires=[0, 1, 2])
+global_drive = qml.pulse.rydberg_drive(amplitude=amp_fn, phase=0, detuning=0, wires=[0, 1, 2])
 
 ######################################################################
 # At this point we could skip directly to defining a ``qnode`` using the ``aquila`` device and running our
@@ -679,6 +676,7 @@ def circuit(params):
 plt.xlabel("μm")
 plt.ylabel("μm")
 plt.legend()
+plt.show()
 
 ##############################################################################
 #
@@ -705,8 +703,7 @@ def circuit(params):
 # values for plotting the function defined in PennyLane for amplitude
 input_times = np.linspace(*ts, 1000)
 input_amplitudes = [
-    qml.pulse.rect(gaussian_fn, windows=[0.01, 1.749])(amplitude_params, _t)
-    for _t in input_times
+    qml.pulse.rect(gaussian_fn, windows=[0.01, 1.749])(amplitude_params, _t) for _t in input_times
 ]
 
 # plot PL input and hardware setpoints for comparison

diff --git a/demonstrations/braket-parallel-gradients.py b/demonstrations/braket-parallel-gradients.py
@@ -95,7 +95,6 @@
 Let's load the SV1 simulator in PennyLane with 25 qubits by specifying the device ARN.
 """
 
-
 device_arn = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
 
 ##############################################################################
@@ -193,8 +192,8 @@ def circuit(params):
 #
 #  .. code-block:: none
 #
-#      Execution time on remote device (seconds): 3.5898206680030853
-#      Execution time on local device (seconds): 23.50668462700196
+#      Execution time on remote device (seconds): 7.571744918823242
+#      Execution time on local device (seconds): 27.32159185409546
 #
 # Nice! These timings highlight the advantage of using the Amazon Braket SV1 device for simulations
 # with large qubit numbers. In general, simulation times scale exponentially with the number of
@@ -235,7 +234,7 @@ def circuit(params):
 #
 #  .. code-block:: none
 #
-#      Gradient calculation time on remote device (seconds): 20.92005863400118
+#      Gradient calculation time on remote device (seconds): 6.4775872230529785
 #
 # Now, the local device:
 #
@@ -260,9 +259,9 @@ def circuit(params):
 #
 #  .. code-block:: none
 #
-#      Gradient calculation time on local device (seconds): 941.8518133479993
+#      Gradient calculation time on local device (seconds): 181.5902488231659
 #
-# Wow, the local device needs around 15 minutes or more! Compare this to less than a minute spent
+# Wow, the local device needs around 3 minutes! Compare this to less around 6 seconds spent
 # calculating the gradient on SV1. This provides a powerful lesson in parallelization.
 #
 # What if we had run on SV1 with ``parallel=False``? It would have taken around 3 minutes—still
@@ -394,37 +393,37 @@ def circuit(params, **kwargs):
 #
 #  .. code-block:: none
 #
-#    Initial cost: -29.98570234095951
+#    Initial cost: -29.98570234095953
 #    Completed iteration 1
-#    Time to complete iteration: 93.96246099472046 seconds
-#    Cost at step 1: -27.154071768632154
+#    Time to complete iteration: 11.028863906860352 seconds
+#    Cost at step 1: -29.995107300982195
 #    Completed iteration 2
-#    Time to complete iteration: 84.80994844436646 seconds
-#    Cost at step 2: -29.98726230006233
+#    Time to complete iteration: 10.132628917694092 seconds
+#    Cost at step 2: -29.99981056829394
 #    Completed iteration 3
-#    Time to complete iteration: 83.13504934310913 seconds
-#    Cost at step 3: -29.999163153600062
+#    Time to complete iteration: 9.985284090042114 seconds
+#    Cost at step 3: -30.00233965117029
 #    Completed iteration 4
-#    Time to complete iteration: 85.61391234397888 seconds
-#    Cost at step 4: -30.002158646044307
+#    Time to complete iteration: 10.059210062026978 seconds
+#    Cost at step 4: -30.01206799710033
 #    Completed iteration 5
-#    Time to complete iteration: 86.70688223838806 seconds
-#    Cost at step 5: -30.012058444011906
+#    Time to complete iteration: 9.966963052749634 seconds
+#    Cost at step 5: -30.048670540725286
 #    Completed iteration 6
-#    Time to complete iteration: 83.26341080665588 seconds
-#    Cost at step 6: -30.063709712612443
+#    Time to complete iteration: 11.216133832931519 seconds
+#    Cost at step 6: -30.13240170461342
 #    Completed iteration 7
-#    Time to complete iteration: 85.25566911697388 seconds
-#    Cost at step 7: -30.32522304705352
+#    Time to complete iteration: 10.09446096420288 seconds
+#    Cost at step 7: -30.25935979157505
 #    Completed iteration 8
-#    Time to complete iteration: 83.55433392524719 seconds
-#    Cost at step 8: -31.411030331978186
+#    Time to complete iteration: 10.020287990570068 seconds
+#    Cost at step 8: -30.41499293575487
 #    Completed iteration 9
-#    Time to complete iteration: 84.08745908737183 seconds
-#    Cost at step 9: -33.87153965616938
+#    Time to complete iteration: 10.345153093338013 seconds
+#    Cost at step 9: -30.587353834845057
 #    Completed iteration 10
-#    Time to complete iteration: 87.4032838344574 seconds
-#    Cost at step 10: -36.05424874438809
+#    Time to complete iteration: 10.07306981086731 seconds
+#    Cost at step 10: -30.76855713404487
 #    Parameters saved to params.npy
 #
 # This example shows us that a 20-qubit QAOA problem can be trained within around 1-2 minutes per
@@ -592,7 +591,7 @@ def cost_function(params, **kwargs):
 #     :target: javascript:void(0);
 
 ##############################################################################
-# Great! The loss function gets lower with each iteration, reaching a value of approximately -36.5.
+# Great! The loss function gets lower with each iteration, reaching a value of approximately -37.8.
 #
 # Conclusion
 # ----------------------------------