{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Junctions and optimizers\n",
        "\n",
        "The service architecture of the QLM provides all the tools to describe quite complicated sequences of compilation steps and post-processings, like so:\n",
        "\n",
        "    stack = plugin1 | plugin2 | qpu\n",
        "\n",
        "However, the information flow along a stack composed of Plugins and a QPU remains linear: the job goes down the stack, reaches the QPU and a result comes back up the stack. Many applications require to be able to iterate some step of job evaluation in an adaptive manner. One may think of, for example, any variational eigensolving procedure.\n",
        "\n",
        "    \n",
        "\n",
        "\n",
        "That's exactly what the ``Junction`` object is here for. A ``Junction`` can be seen as a Plugin that will adaptively keep sending jobs down the stack and analyze the corresponding results until it decides to stop and returns a final result. For instance, a variational optimizer will optimize variational circuits until some criterion on the variational energy is reached.\n",
        "\n",
        "\n",
        "Junctions can be used to compose stacks with the same pipe syntax as Plugins and QPUs:\n",
        "\n",
        "    stack = plugin1 | somejunction | plugin2  | qpu\n",
        "\n",
        "\n",
        "Let us try and program some ``Junction`` called `IterativeExplorer` that will:\n",
        " - take a parametrized job with a single abstract parameter,\n",
        " - explore the energies resulting from the execution of the job for various values spread between 0 and 2$\\pi$ (with some fixed number of steps),\n",
        " - return the best encountered value, and the best parameter value."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from qat.plugins import Junction\n",
        "from qat.core import Result, Job\n",
        "from qat.comm.exceptions.ttypes import PluginException\n",
        "\n",
        "# Junctions are built by inheriting from the Junction class\n",
        "class IterativeExplorer(Junction):\n",
        "    r\"\"\" \n",
        "    Iteratively explores the (0, 2$\\pi$) range for the incoming job's parameter\n",
        "    \n",
        "    Args:\n",
        "        nsteps (int): the number of values to try\n",
        "    \"\"\"\n",
        "    def __init__(self, nsteps=100):\n",
        "        self.nsteps = nsteps\n",
        "        # Here the 'collective' parameter tells the junction that we would like to handle jobs one by one\n",
        "        # If set to True, we would have to process the full incoming Batch in one go. Let us keep things simple.\n",
        "        super(IterativeExplorer, self).__init__(collective=False)\n",
        "        \n",
        "    # Junctions are abstract classes that require you to implement the following method\n",
        "    def run(self, qlm_object: Job, meta_data: dict) -> Result:\n",
        "        parameters = qlm_object.get_variables() # This returns the list of variables of the job\n",
        "        if len(parameters) != 1:\n",
        "            raise PluginException(message=\"Can't handle Jobs with more than 1 variable\")\n",
        "        vname = parameters.pop() # getting the first and only variable name\n",
        "        angles = np.linspace(0, 2 * np.pi, self.nsteps)\n",
        "        values = []\n",
        "        for angle in angles:\n",
        "            # We bind the value of the variable to `angle` \n",
        "            job = qlm_object(**{vname: angle})\n",
        "            # We evaluate the job using the `self.execute` method\n",
        "            result = self.execute(job)\n",
        "            # We extract the energy and push it to our result list\n",
        "            values.append(result.value)\n",
        "        # Extracting the best value\n",
        "        best_value = min(values)\n",
        "        # and the best angle\n",
        "        best_angle = angles[values.index(best_value)]\n",
        "        # We need to return a QLM Result object\n",
        "        return Result(\n",
        "            value=best_value,\n",
        "            meta_data={\"best_angle\": best_angle}\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As you can see, Junctions are built by inheriting from the Junction class and implementing its `run` method.\n",
        "This method will receive the incoming job object, together with some optional meta data, from the higher part of the stack and has access to the lower part of the stack via the `execute` method.\n",
        "\n",
        "Let us now run our junction!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# first we need a qpu\n",
        "from qat.qpus import get_default_qpu\n",
        "qpu = get_default_qpu()\n",
        "# and build a stack\n",
        "stack = IterativeExplorer(nsteps=25) | qpu"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Now lets build a simple parametrized Job\n",
        "from qat.core import Observable\n",
        "from qat.lang.AQASM import Program, RY\n",
        "\n",
        "prog = Program()\n",
        "qbits = prog.qalloc(1)\n",
        "theta = prog.new_var(float, \"\\\\theta\")\n",
        "prog.apply(RY(theta), qbits)\n",
        "circuit = prog.to_circ()\n",
        "job = circuit.to_job(observable=Observable.sigma_z(0, 1)) # Z on qbit 0\n",
        "circuit.display()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Run!\n",
        "result = stack.submit(job)\n",
        "print(\"Final energy:\", result.value, \" | best angle:\", result.meta_data[\"best_angle\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# What about optimizers?\n",
        "\n",
        "In the previous example, we managed to achieve the expected behavior for our simple optimizer. However, some line of codes are purely here for administrative purposes.\n",
        "This is why the `Optimizer` class exists. It provides a very similar interface to the ``Junction`` class, but takes care of some of the administrative burden.\n",
        "\n",
        "Let us see how the previous example can be rewritten using an ``Optimizer``:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from qat.plugins import Optimizer\n",
        "\n",
        "class IterativeExplorer(Optimizer):\n",
        "    r\"\"\" \n",
        "    Iteratively explores the (0, 2$\\pi$) range for the incoming job's parameter\n",
        "    \n",
        "    Args:\n",
        "        nsteps (int): the number of values to try\n",
        "    \"\"\"\n",
        "    def __init__(self, nsteps=100):\n",
        "        self.nsteps = nsteps\n",
        "        super(IterativeExplorer, self).__init__(collective=False)\n",
        "        \n",
        "    # The run method changed name and is now called `optimize`\n",
        "    def optimize(self, variables: list) -> tuple:\n",
        "        # the argument `variables` contains the list of variables of the job\n",
        "        if len(variables) != 1:\n",
        "            raise PluginException(message=\"Can't handle Jobs with more than 1 variable\")\n",
        "        vname = variables.pop() # getting the first and only variable name\n",
        "        angles = np.linspace(0, 2 * np.pi, self.nsteps)\n",
        "        values = []\n",
        "        for angle in angles:\n",
        "            values.append(self.evaluate({vname: angle}))\n",
        "        best_value = min(values)\n",
        "        best_angle = angles[values.index(best_value)]\n",
        "        return best_value, [best_angle], \"hello :)\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As you can see, we need to implement a method called `optimize` and have access to a method called `evaluate` that takes a value map and returns a float.\n",
        "The `optimize` method receives the list of variables of the job and should return:\n",
        "- the final energy value\n",
        "- the set of parameters corresponding to this value\n",
        "- optionally, any object (it will be stringified and stored in the meta data of the result)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "stack = IterativeExplorer(nsteps=25) | qpu\n",
        "result = stack.submit(job)\n",
        "print(\"Final energy:\", result.value, \" | best angle:\", result.meta_data[\"parameters\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As you can see the final result is quite similar as in the first implementation.\n",
        "We can have a deep look at its meta data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "for key, value in result.meta_data.items():\n",
        "    print(key, \":\", value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The entry \"optimizer_data\" contains our third returned value.\n",
        "Most importantly, the optimizer kept track of the different evaluation and returned this trace in the \"optimization_trace\" entry.\n",
        "\n",
        "## Going further\n",
        "\n",
        "When we implemented our Junction (or our Optimizer), we didn't make any use of the `meta_data` parameter. In order to improve our plugin, we could make it so the job itself could control the resolution of our exploration.\n",
        "\n",
        "To do so, we will assume that the user will transmit the number of steps via the \"IterativeExplorer_nsteps\" entry of the meta data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Here most of the code is the same as in the first cell\n",
        "class IterativeExplorer(Junction):\n",
        "    r\"\"\" \n",
        "    Iteratively explores the (0, 2$\\pi$) range for the incoming job's parameter\n",
        "    \n",
        "    Args:\n",
        "        nsteps (int): the number of values to try\n",
        "    \"\"\"\n",
        "    def __init__(self, nsteps=100):\n",
        "        self.nsteps = nsteps # This is now a default value\n",
        "        super(IterativeExplorer, self).__init__(collective=False)\n",
        "        \n",
        "    def run(self, qlm_object: Job, meta_data: dict) -> Result:\n",
        "        nsteps = meta_data.get(\"IterativeExplorer_nsteps\", None)\n",
        "        if nsteps is not None:\n",
        "            nsteps = int(nsteps)\n",
        "        else:\n",
        "            nsteps = self.nsteps\n",
        "        print(self.__class__.__name__, \": using\", nsteps, 'steps')\n",
        "        parameters = qlm_object.get_variables() \n",
        "        if len(parameters) != 1:\n",
        "            raise PluginException(message=\"Can't handle Jobs with more than 1 variable\")\n",
        "        vname = parameters.pop() \n",
        "        angles = np.linspace(0, 2 * np.pi, nsteps)\n",
        "        values = []\n",
        "        for angle in angles:\n",
        "            job = qlm_object(**{vname: angle})\n",
        "            result = self.execute(job)\n",
        "            values.append(result.value)\n",
        "        best_value = min(values)\n",
        "        best_angle = angles[values.index(best_value)]\n",
        "        return Result(\n",
        "            value=best_value,\n",
        "            meta_data={\"best_angle\": best_angle}\n",
        "        )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "stack = IterativeExplorer(nsteps=1) | qpu\n",
        "result = stack.submit(job)\n",
        "print(\"Final energy:\", result.value, \" | best angle:\", result.meta_data[\"best_angle\"])\n",
        "result = stack.submit(job, meta_data={\"IterativeExplorer_nsteps\": \"32\"})\n",
        "print(\"Final energy:\", result.value, \" | best angle:\", result.meta_data[\"best_angle\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "authors": [
      "Simon Martiel"
    ],
    "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.12.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}