{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1a5e2004",
   "metadata": {},
   "source": [
    "# How to reuse graph definitions in different calculations\n",
    "**Reapply graph nodes for multiple applications**\n",
    "\n",
    "In Boulder Opal, graph definitions can be easily carried over between different graphs.\n",
    "For example, the nodes used to create a Hamiltonian can first be used in an optimization task, and then later used in simulation.\n",
    "This is especially useful if one is changing properties of the graph that could not be easily batched together, for example, the dimension of the Fock space or the structure of the Hamiltonian.\n",
    "If instead you want to perform a simulation using the results of an optimization in a single calculation, see [this user guide](https://docs.q-ctrl.com/boulder-opal/toolkit/design/calculate-with-graphs/how-to-perform-optimization-and-simulation-in-the-same-calculation).\n",
    "\n",
    "Below we show an example of reusing nodes for optimization and simulation in separate graphs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53c13704-9f42-47a9-98c0-5ebfc0cfc56d",
   "metadata": {},
   "source": [
    "## Summary workflow\n",
    "### 1. Define a function for the nodes to be reused\n",
    "For convenience, start by creating a function that contains the nodes to be reused. \n",
    "One of the arguments of the function should be `graph` so that the function can be used to add nodes on to the graph.\n",
    "The function should also be passed any parameters required to define the nodes.\n",
    "After these nodes are added to graph, any nodes required to compute the graph (for example the Hamiltonian) should be returned.\n",
    "\n",
    "For example, one might define a function to create an infidelity node for a given Hamiltonian and some target:\n",
    "```python\n",
    "def create_infidelity(graph, hamiltonian):\n",
    "    target = graph.target(operator=graph.pauli_matrix(\"Y\"))\n",
    "    return graph.infidelity_pwc(\n",
    "        hamiltonian=hamiltonian,\n",
    "        target=target,\n",
    "        name=\"infidelity\",\n",
    "    )\n",
    "```\n",
    "### 2. Define and execute the computational graph\n",
    "Set up a [graph object](https://docs.q-ctrl.com/boulder-opal/toolkit/design/calculate-with-graphs/how-to-represent-quantum-systems-using-graphs) to carry out some task, be it to perform a [control optimization](https://docs.q-ctrl.com/boulder-opal/toolkit/design/design-model-based-controls/how-to-optimize-controls-in-arbitrary-quantum-systems-using-graphs) or [simulation](https://docs.q-ctrl.com/boulder-opal/toolkit/design/simulate-quantum-systems/how-to-simulate-closed-noiseless-systems), using the previously defined function.\n",
    "With the graph object created, an optimization can be run using the `boulderopal.run_optimization` function or a simulation can be run using the `boulderopal.execute_graph`.\n",
    "\n",
    "### 3. Define and execute other graphs\n",
    "Reusing the convenience function, define a new graph that carries out a related task.\n",
    "For example, if the first calculation is an optimization, then the new graph might be a [noisy simulation](https://docs.q-ctrl.com/boulder-opal/toolkit/design/simulate-quantum-systems/how-to-simulate-quantum-dynamics-subject-to-noise) or a [quasi-static parameter scan](https://docs.q-ctrl.com/boulder-opal/toolkit/design/design-error-robust-quantum-logic-gates/how-to-evaluate-control-susceptibility-to-quasi-static-noise) with the optimized results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5d370d7-1fd2-4b1b-9720-601b3c54acec",
   "metadata": {},
   "source": [
    "## Example: Reusing nodes for simulation and optimization of a qubit\n",
    "In this example we will implement a Y gate in a noisy single-qubit system. \n",
    "The system is described by a Hamiltonian of the form,\n",
    "$$\n",
    "H(t) = \\alpha(t) \\sigma_{z} + \\frac{1}{2}\\left(\\gamma(t)\\sigma_{-} + \\gamma^*(t)\\sigma_{+}\\right) + \\delta \\sigma_{z} + \\beta(t) \\sigma_{z}\n",
    "$$\n",
    "where,\n",
    "$\\sigma_{z}$ is the Pauli Z operator,\n",
    "$\\sigma_{\\pm}$ are the qubit ladder operators,\n",
    "$\\alpha(t)$ and $\\gamma(t)$ are, respectively, real and complex optimizable time-dependent controls,\n",
    "$\\delta$ is the qubit detuning, and,\n",
    "$\\beta(t)$ is a dephasing noise process.\n",
    "\n",
    "First we will optimize $\\alpha(t)$ and $\\gamma(t)$ to produce a Y gate in the absence of any noise, $\\beta(t)=0$.\n",
    "Then we will optimize the controls in the presence of a dephasing amplitude on the order of $\\beta(t)=\\beta_0$.\n",
    "Finally we will test how both sets of controls perform in the presence of real stochastic noise sampled from a normal distribution centered at $\\beta_0$. We see that for the robust controls optimized in the presence of constant noise, the associated infidelity is significantly lower than for the controls optimized without any noise. \n",
    "\n",
    "To reduce the amount of code required nodes are reused by defining functions to create the optimiziable controls, compute the Hamiltonian and to compute infidelity.\n",
    "In total three graphs are used, two for optimizating the controls and another for simulation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "75cf8709-2bbd-4ede-b803-85952328298a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import boulderopal as bo\n",
    "\n",
    "\n",
    "def create_optimizable_controls(graph):\n",
    "    \"\"\"Create the optimizable controls α(t) and γ(t).\"\"\"\n",
    "    # Maximum value for |α(t)|.\n",
    "    alpha_max = 2 * np.pi * 0.25e6  # rad/s\n",
    "\n",
    "    # Real PWC signal representing α(t).\n",
    "    alpha = graph.real_optimizable_pwc_signal(\n",
    "        segment_count=segment_count,\n",
    "        duration=duration,\n",
    "        minimum=-alpha_max,\n",
    "        maximum=alpha_max,\n",
    "        name=\"$\\\\alpha$\",\n",
    "    )\n",
    "\n",
    "    # Maximum value for |γ(t)|.\n",
    "    gamma_max = 2 * np.pi * 0.5e6  # rad/s\n",
    "\n",
    "    # Complex PWC signal representing γ(t)\n",
    "    gamma = graph.complex_optimizable_pwc_signal(\n",
    "        segment_count=segment_count,\n",
    "        duration=duration,\n",
    "        maximum=gamma_max,\n",
    "        name=\"$\\\\gamma$\",\n",
    "    )\n",
    "    return alpha, gamma\n",
    "\n",
    "\n",
    "def create_noiseless_hamiltonian(graph, alpha_control, gamma_control):\n",
    "    \"\"\"Create the noiseless Hamiltonian from the controls α(t) and γ(t).\"\"\"\n",
    "    # Detuning δ.\n",
    "    delta = 2 * np.pi * 0.25e6  # rad/s\n",
    "    detuning_hamiltonian = delta * graph.pauli_matrix(\"Z\")\n",
    "\n",
    "    # Control Hamiltonian.\n",
    "    control_hamiltonian = alpha_control * graph.pauli_matrix(\n",
    "        \"Z\"\n",
    "    ) + graph.hermitian_part(gamma_control * graph.pauli_matrix(\"M\"))\n",
    "\n",
    "    # Total Hamiltonian.\n",
    "    return detuning_hamiltonian + control_hamiltonian\n",
    "\n",
    "\n",
    "def create_infidelity(graph, hamiltonian, noise_operators=None):\n",
    "    \"\"\"\n",
    "    Create the infidelity node with the option to pass a noise operator.\n",
    "    If a noise operator is passed then the operational infidelity plus\n",
    "    filter function values is calculated.\n",
    "    \"\"\"\n",
    "    # Target operation node.\n",
    "    target = graph.target(operator=graph.pauli_matrix(\"Y\"))\n",
    "\n",
    "    # Compute infidelity.\n",
    "    graph.infidelity_pwc(\n",
    "        hamiltonian=hamiltonian,\n",
    "        noise_operators=noise_operators,\n",
    "        target=target,\n",
    "        name=\"infidelity\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dd728a25-b331-4268-8b61-240f3c87e103",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your task (action_id=\"1829137\") has started.\n",
      "Your task (action_id=\"1829137\") has completed.\n",
      "Optimized noiseless cost: 4.885e-15\n"
     ]
    }
   ],
   "source": [
    "# Pulse parameters.\n",
    "segment_count = 50\n",
    "duration = 10e-6  # s\n",
    "\n",
    "# Dephasing noise amplitude.\n",
    "beta_0 = 2 * np.pi * 20e3  # rad/s\n",
    "\n",
    "# Generate noiseless optimized controls.\n",
    "graph = bo.Graph()\n",
    "\n",
    "# Use convenience functions to create the computational graph.\n",
    "alpha, gamma = create_optimizable_controls(graph)\n",
    "hamiltonian = create_noiseless_hamiltonian(graph, alpha, gamma)\n",
    "create_infidelity(graph, hamiltonian)\n",
    "\n",
    "optimization_result = bo.run_optimization(\n",
    "    graph=graph,\n",
    "    cost_node_name=\"infidelity\",\n",
    "    output_node_names=[\"$\\\\alpha$\", \"$\\\\gamma$\"],\n",
    "    optimization_count=4,\n",
    ")\n",
    "\n",
    "print(f\"Optimized noiseless cost: {optimization_result['cost']:.3e}\")\n",
    "\n",
    "# Retrieve values of the PWC controls α(t) and γ(t).\n",
    "alpha_values_noiseless = optimization_result[\"output\"][\"$\\\\alpha$\"][\"values\"]\n",
    "gamma_values_noiseless = optimization_result[\"output\"][\"$\\\\gamma$\"][\"values\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5d72a6bc-ffbe-480d-8a83-887a6710778f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your task (action_id=\"1829138\") is queued.\n",
      "Your task (action_id=\"1829138\") has started.\n",
      "Your task (action_id=\"1829138\") has completed.\n",
      "Optimized robust cost: 5.564e-13\n"
     ]
    }
   ],
   "source": [
    "# Generate robust optimized controls.\n",
    "graph = bo.Graph()\n",
    "\n",
    "# Use convenience functions to create the computational graph.\n",
    "alpha, gamma = create_optimizable_controls(graph)\n",
    "hamiltonian = create_noiseless_hamiltonian(graph, alpha, gamma)\n",
    "\n",
    "# This time passing a constant dephasing term with amplitude beta\n",
    "# to `compute_infidelity`.\n",
    "noise_operators = [beta_0 * graph.pauli_matrix(\"Z\")]\n",
    "create_infidelity(graph, hamiltonian, noise_operators)\n",
    "\n",
    "optimization_result = bo.run_optimization(\n",
    "    graph=graph,\n",
    "    cost_node_name=\"infidelity\",\n",
    "    output_node_names=[\"$\\\\alpha$\", \"$\\\\gamma$\"],\n",
    "    optimization_count=4,\n",
    ")\n",
    "\n",
    "print(f\"Optimized robust cost: {optimization_result['cost']:.3e}\")\n",
    "\n",
    "# Retrieve values of the robust PWC controls α(t) and γ(t).\n",
    "alpha_values_robust = optimization_result[\"output\"][\"$\\\\alpha$\"][\"values\"]\n",
    "gamma_values_robust = optimization_result[\"output\"][\"$\\\\gamma$\"][\"values\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "eee67dda-167e-4532-ab95-04aee0b3f46a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your task (action_id=\"1829139\") has started.\n",
      "Your task (action_id=\"1829139\") has completed.\n",
      "Infidelity of pulse optimized without noise\n",
      "\tMean: 0.534\n",
      "\tStd : 0.0189\n",
      "\n",
      "Infidelity of pulse optimized with noise\n",
      "\tMean: 0.0222\n",
      "\tStd : 0.00352\n"
     ]
    }
   ],
   "source": [
    "# Run this number of simulations with different noise trajectories.\n",
    "simulation_count = 100\n",
    "\n",
    "# Create a new Boulder Opal graph.\n",
    "graph = bo.Graph()\n",
    "\n",
    "# Create a batch of real PWC signals representing α(t) using the optimized values.\n",
    "alpha_values = np.array([alpha_values_noiseless, alpha_values_robust])\n",
    "alpha = graph.pwc_signal(values=alpha_values, duration=duration)\n",
    "\n",
    "# Create a batch of complex PWC signals representing γ(t) using the optimized values.\n",
    "gamma_values = np.array([gamma_values_noiseless, gamma_values_robust])\n",
    "gamma = graph.pwc_signal(values=gamma_values, duration=duration)\n",
    "\n",
    "# System Hamiltonian.\n",
    "system_hamiltonian = create_noiseless_hamiltonian(graph, alpha, gamma)\n",
    "\n",
    "# Noise Hamiltonian.\n",
    "noise = graph.random.normal(\n",
    "    shape=(simulation_count, 2, segment_count),\n",
    "    mean=beta_0,\n",
    "    standard_deviation=2e4,\n",
    "    seed=0,\n",
    ")\n",
    "noise_pwc = graph.pwc_signal(noise, duration=duration, name=\"noise_pwc\")\n",
    "noise_hamiltonian = noise_pwc * graph.pauli_matrix(\"Z\")\n",
    "\n",
    "# Combine the system and the noise Hamiltonian to get the total Hamiltonian.\n",
    "hamiltonian = system_hamiltonian + noise_hamiltonian\n",
    "\n",
    "create_infidelity(graph, hamiltonian)\n",
    "\n",
    "result = bo.execute_graph(graph=graph, output_node_names=\"infidelity\")\n",
    "\n",
    "infidelity = result[\"output\"][\"infidelity\"][\"value\"]\n",
    "\n",
    "print(\"Infidelity of pulse optimized without noise\")\n",
    "print(f\"\\tMean: {infidelity[:, 0].mean():.3}\")\n",
    "print(f\"\\tStd : {infidelity[:, 0].std():.3}\")\n",
    "\n",
    "print(\"\\nInfidelity of pulse optimized with noise\")\n",
    "print(f\"\\tMean: {infidelity[:, 1].mean():.3}\")\n",
    "print(f\"\\tStd : {infidelity[:, 1].std():.3}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}