{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "88dc868f",
      "metadata": {},
      "source": [
        "## Leveraging parallelism through map\n",
        "\n",
        "One major advantage in workflow systems is the ease of scaling computation horizontally.\n",
        "Data-parallel tasks can act independently; In tierkreis this can simply be achieved through the `map` function.\n",
        "Each map element will receive exactly one sets of inputs and can therefore be immediately dispatched.\n",
        "In this example we will observe the speedup by running multiple independent graphs in parallel.\n",
        "\n",
        "First we define a simple graph that will run a circuit in two version:\n",
        "1. Using the qiskit aer simulator\n",
        "2. Using the [qulacs](https://github.com/qulacs/qulacs) simulator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ee0cabe0",
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install tierkreis pytket qiskit-aer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9e6fae3d",
      "metadata": {},
      "outputs": [],
      "source": [
        "from typing import Literal, NamedTuple\n",
        "from tierkreis.builder import GraphBuilder\n",
        "from tierkreis.controller.data.models import TKR, OpaqueType\n",
        "from tierkreis.builtins import untuple\n",
        "from tierkreis.aer_worker import (\n",
        "    get_compiled_circuit as aer_compile,\n",
        "    run_circuit as aer_run,\n",
        ")\n",
        "from tierkreis.qulacs_worker import (\n",
        "    get_compiled_circuit as qulacs_compile,\n",
        "    run_circuit as qulacs_run,\n",
        ")\n",
        "\n",
        "type BackendResult = OpaqueType[\"pytket.backends.backendresult.BackendResult\"]  # noqa: F821\n",
        "type Circuit = OpaqueType[\"pytket._tket.circuit.Circuit\"]  # noqa: F821\n",
        "\n",
        "\n",
        "class SimulateJobInputsSingle(NamedTuple):\n",
        "    simulator_name: TKR[Literal[\"aer\", \"qulacs\"]]\n",
        "    circuit_shots: TKR[tuple[Circuit, int]]\n",
        "    compilation_optimisation_level: TKR[int]\n",
        "\n",
        "\n",
        "def aer_simulate_single():\n",
        "    g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult])\n",
        "    circuit_shots = g.task(untuple(g.inputs.circuit_shots))\n",
        "\n",
        "    compiled_circuit = g.task(\n",
        "        aer_compile(\n",
        "            circuit=circuit_shots.a,\n",
        "            optimisation_level=g.inputs.compilation_optimisation_level,\n",
        "        )\n",
        "    )\n",
        "    res = g.task(aer_run(compiled_circuit, circuit_shots.b))\n",
        "    g.outputs(res)\n",
        "    return g\n",
        "\n",
        "\n",
        "def qulacs_simulate_single():\n",
        "    g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult])\n",
        "    circuit_shots = g.task(untuple(g.inputs.circuit_shots))\n",
        "\n",
        "    compiled_circuit = g.task(\n",
        "        qulacs_compile(\n",
        "            circuit=circuit_shots.a,\n",
        "            optimisation_level=g.inputs.compilation_optimisation_level,\n",
        "        )\n",
        "    )\n",
        "    res = g.task(qulacs_run(compiled_circuit, circuit_shots.b))\n",
        "    g.outputs(res)\n",
        "    return g"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f1b7e67e",
      "metadata": {},
      "source": [
        "So far these are regular graphs that compile and simulate a single circuit.\n",
        "We are going to combine them into a single graph taking a parameter to decide which simulator to run using `ifelse`\n",
        "Although we we will have two similar subgraphs in the evaluation, this is not a performance detriment as `ifelse` only evaluates lazily. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0a71ddad",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.builtins import str_eq\n",
        "\n",
        "\n",
        "def compile_simulate_single():\n",
        "    g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult])\n",
        "\n",
        "    aer_res = g.eval(aer_simulate_single(), g.inputs)\n",
        "    qulacs_res = g.eval(qulacs_simulate_single(), g.inputs)\n",
        "    res = g.ifelse(\n",
        "        g.task(str_eq(g.inputs.simulator_name, g.const(\"aer\"))), aer_res, qulacs_res\n",
        "    )\n",
        "\n",
        "    g.outputs(res)\n",
        "    return g"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9bdbe246",
      "metadata": {},
      "source": [
        "To make this parallel over multiple circuits we are using the `map` feature in a new graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ae2ee842",
      "metadata": {},
      "outputs": [],
      "source": [
        "class SimulateJobInputs(NamedTuple):\n",
        "    simulator_name: TKR[Literal[\"aer\", \"qulacs\"]]\n",
        "    circuits: TKR[list[Circuit]]\n",
        "    n_shots: TKR[list[int]]\n",
        "    compilation_optimisation_level: TKR[int]\n",
        "\n",
        "\n",
        "g = GraphBuilder(SimulateJobInputs, TKR[list[BackendResult]])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1052dd30",
      "metadata": {},
      "source": [
        "Each of the `SimulateJobInputsSingle` expects a tuple `(Circuit, n_shots)` which we generate by zipping"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d0c9907a",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.builtins import tkr_zip\n",
        "\n",
        "circuits_shots = g.task(tkr_zip(g.inputs.circuits, g.inputs.n_shots))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5b28a6ba",
      "metadata": {},
      "source": [
        "A convenient way to aggregate the inputs is using a map over a lambda"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "aa29dd13",
      "metadata": {},
      "outputs": [],
      "source": [
        "job_inputs = g.map(\n",
        "    lambda x: SimulateJobInputsSingle(\n",
        "        simulator_name=g.inputs.simulator_name,\n",
        "        circuit_shots=x,\n",
        "        compilation_optimisation_level=g.inputs.compilation_optimisation_level,\n",
        "    ),\n",
        "    circuits_shots,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f15c2943",
      "metadata": {},
      "source": [
        "and finally we can map over the jobs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "685a5d47",
      "metadata": {},
      "outputs": [],
      "source": [
        "res = g.map(compile_simulate_single(), job_inputs)\n",
        "\n",
        "g.outputs(res)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d8b7cabf",
      "metadata": {},
      "source": [
        "preparing the storage, executor and inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b7c2baf2",
      "metadata": {},
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "from uuid import UUID\n",
        "\n",
        "from pytket.qasm.qasm import circuit_from_qasm\n",
        "\n",
        "from tierkreis.consts import PACKAGE_PATH\n",
        "from tierkreis.storage import FileStorage\n",
        "from tierkreis.executor import UvExecutor\n",
        "\n",
        "circuit = circuit_from_qasm(Path().parent / \"data\" / \"ghz_state_n23.qasm\")\n",
        "circuits = [circuit] * 3\n",
        "n_shots = 1024\n",
        "\n",
        "storage = FileStorage(UUID(int=107), do_cleanup=True)\n",
        "executor = UvExecutor(PACKAGE_PATH / \"..\" / \"tierkreis_workers\", storage.logs_path)\n",
        "inputs = {\n",
        "    \"circuits\": circuits,\n",
        "    \"n_shots\": [n_shots] * len(circuits),\n",
        "    \"compilation_optimisation_level\": 2,\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1bca6c67",
      "metadata": {},
      "source": [
        "we can now benchmark aer by setting the `simulator_name` input"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "93217435",
      "metadata": {},
      "outputs": [],
      "source": [
        "import time\n",
        "from tierkreis.controller import run_graph\n",
        "\n",
        "inputs[\"simulator_name\"] = \"aer\"\n",
        "print(\"Simulating using aer...\")\n",
        "start = time.time()\n",
        "run_graph(storage, executor, g, inputs, polling_interval_seconds=0.1)\n",
        "print(f\"time taken: {time.time() - start}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "606a9c18",
      "metadata": {},
      "source": [
        "and"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6f6346a6",
      "metadata": {},
      "outputs": [],
      "source": [
        "inputs[\"simulator_name\"] = \"qulacs\"\n",
        "\n",
        "print(\"Simulating using qulacs...\")\n",
        "storage.clean_graph_files()\n",
        "start = time.time()\n",
        "run_graph(storage, executor, g, inputs, polling_interval_seconds=0.1)\n",
        "print(f\"time taken: {time.time() - start}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c34a812e",
      "metadata": {},
      "source": [
        "compared against running the same graph three times:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "39cf4dc4",
      "metadata": {},
      "outputs": [],
      "source": [
        "start = time.time()\n",
        "for circuit in circuits:\n",
        "    inputs = {\n",
        "        \"circuit_shots\": (circuit, n_shots),\n",
        "        \"compilation_optimisation_level\": 2,\n",
        "        \"simulator_name\": \"aer\",\n",
        "    }\n",
        "    storage.clean_graph_files()\n",
        "    run_graph(\n",
        "        storage,\n",
        "        executor,\n",
        "        compile_simulate_single(),\n",
        "        inputs,\n",
        "        polling_interval_seconds=0.1,\n",
        "    )\n",
        "print(f\"time taken: {time.time() - start}\")"
      ]
    }
  ],
  "metadata": {
    "execution": {
      "timeout": 120
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.13.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}