fix: use jit for analog hamiltonian example

Author: AbeCoull

examples/analog_hamiltonian_simulation/06_Analog_Hamiltonian_simulation_with_PennyLane.ipynb

@@ -142,7 +142,7 @@
     "Following the convention of PennyLane, we have used megahertz (MHz) as the unit for the amplitude and detuning. Also, we have limited the maximum duration of the program, `time_max`, to be 4$\\mu s$, and introduced the Rydberg interaction constant $C_6$. These constants are specific to QuEra's Aquila, and they can be queried from Braket SDK, as shown in [this AHS example notebook](https://github.com/aws/amazon-braket-examples/blob/main/examples/analog_hamiltonian_simulation/01_Introduction_to_Aquila.ipynb). With these constants, the differentiable driving fields for the Rydberg atom device can be defined via [PennyLane's JAX interface](https://pennylane.ai/qml/demos/tutorial_jax_transformations)."
    ]
   },
-  {
+    {
    "cell_type": "code",
    "execution_count": 4,
    "id": "7ec45e20-ed22-4f1d-9540-db456794e5ed",
@@ -153,18 +153,19 @@
    "source": [
     "import jax\n",
     "import jax.numpy as jnp\n",
+    "from functools import partial\n",
     "\n",
     "# Set to float64 precision and remove jax CPU/GPU warning\n",
     "jax.config.update(\"jax_enable_x64\", True)\n",
     "jax.config.update(\"jax_platform_name\", \"cpu\")\n",
     "jax.config.update(\"jax_debug_nans\", True)\n",
     "\n",
-    "\n",
+    "@jax.jit\n",
     "def ryd_amplitude(p, t):\n",
     "    \"\"\"Parametrized function for amplitude\"\"\"\n",
     "    return p[0] * (1 - (1 - jnp.sin(jnp.pi * t / time_max) ** 2) ** (p[1] / 2))\n",
     "\n",
-    "\n",
+    "@jax.jit\n",
     "def ryd_detuning(p, t):\n",
     "    \"\"\"Parametrized function for detuning\"\"\"\n",
     "    return p[0] * jnp.arctan(p[1] * (t - time_max / 2)) / (jnp.pi / 2)"
@@ -416,7 +417,7 @@
     "With the cost Hamiltonian defined, we can create a variational AHS program, which evolves the Rydberg system followed by measuring the cost function $H_\\text{cost}$.\n"
    ]
   },
-  {
+ {
    "cell_type": "code",
    "execution_count": 12,
    "id": "94cb8760-1ee5-4ed3-8f3d-cd64f6103e12",
@@ -429,9 +430,9 @@
     "dev = qml.device(\"default.qubit.jax\", wires=range(len(coords)))\n",
     "\n",
     "\n",
-    "# Define the qnode that evolves the Rydberg system, followed by calculating the cost function\n",
+    "@partial(jax.jit, static_argnums=(1,))\n",
     "@qml.qnode(dev, interface=\"jax\")\n",
-    "def program_cost(detuning_param, amplitude_param=amplitude_param, ts=ts):\n",
+    "def program_cost(detuning_param, amplitude_param=amplitude_param):\n",
     "    qml.evolve(H_ryd)([amplitude_param, detuning_param], ts)\n",
     "    return qml.expval(H_cost)"
    ]
@@ -511,14 +512,10 @@
     "import optax\n",
     "\n",
     "n_epochs = 10\n",
-    "\n",
-    "# The following block creates a constant schedule of the learning rate\n",
-    "# that increases from 0.1 to 0.5 after 10 epochs\n",
-    "schedule0 = optax.constant_schedule(1e-1)\n",
-    "schedule1 = optax.constant_schedule(5e-1)\n",
-    "schedule = optax.join_schedules([schedule0, schedule1], [10])\n",
-    "optimizer = optax.nadam(learning_rate=schedule)\n",
-    "opt_state = optimizer.init(theta)"
+    "schedule = optax.join_schedules(\n",
+    "    [optax.constant_schedule(1e-1), optax.constant_schedule(5e-1)], [10]\n",
+    ")\n",
+    "optimizer = optax.nadam(learning_rate=schedule)"
    ]
   },
   {
@@ -542,36 +539,36 @@
      "output_type": "stream",
      "text": [
       "epoch  expectation\n",
-      "0     -1.4220786805245293\n",
-      "1     -1.4440255545399934\n",
-      "2     -1.4682669156813517\n",
-      "3     -1.494929936649468\n",
-      "4     -1.52414541861693\n",
-      "5     -1.5561124234490418\n",
-      "6     -1.590274680470023\n",
-      "7     -1.6261105967628833\n",
-      "8     -1.6618436930215126\n",
-      "9     -1.695197967500683\n",
-      "The final parameter for detuning = [12.268508609205842, 0.4353973115231312]\n"
+      "0    -1.422078680524523\n",
+      "1    -1.4552870814551753\n",
+      "2    -1.48565657822032\n",
+      "3    -1.5177182172197772\n",
+      "4    -1.551923494981983\n",
+      "5    -1.5883691608696957\n",
+      "6    -1.626260941139706\n",
+      "7    -1.6636605058546716\n",
+      "8    -1.6978874284223098\n",
+      "9    -1.7255362493082869\n",
+      "The final parameter for detuning = [12.208998163228085, 0.32605421934464207]\n"
      ]
     }
    ],
    "source": [
-    "energy = np.zeros(n_epochs + 1)\n",
-    "energy[0] = program_cost(theta)\n",
-    "gradients = np.zeros(n_epochs)\n",
+    "@jax.jit\n",
+    "def optimization_step(theta, opt_state):\n",
+    "    val, grad = jax.value_and_grad(program_cost)(theta)\n",
+    "    updates, opt_state = optimizer.update(grad, opt_state)\n",
+    "    theta = optax.apply_updates(theta, updates)\n",
+    "    return val, grad, theta, opt_state\n",
+    "\n",
+    "theta = detuning_param\n",
+    "opt_state = optimizer.init(theta)\n",
     "\n",
     "## Optimization loop\n",
     "print(\"epoch  expectation\")\n",
     "for n in range(n_epochs):\n",
-    "    val, grad_program = value_and_grad(theta)\n",
-    "    updates, opt_state = optimizer.update(grad_program, opt_state)\n",
-    "    theta = optax.apply_updates(theta, updates)\n",
-    "\n",
-    "    energy[n + 1] = val\n",
-    "    gradients[n] = np.mean(np.abs(grad_program))\n",
-    "\n",
-    "    print(n, \"   \", val)\n",
+    "    val, grad, theta, opt_state = optimization_step(theta, opt_state)\n",
+    "    print(f\"{n}    {val}\")\n",
     "\n",
     "print(f\"The final parameter for detuning = {[float(i) for i in theta]}\")"
    ]