{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "# Resource Estimation for simulating a 2D Heisenberg Model Hamiltonian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "In this Python + Q# notebook we demonstrate how to *efficienty* estimate the resources for simulating a Heisenberg model Hamiltonian on an $N \\times N$ 2D\n",
    "lattice using a *fourth-order Trotter Suzuki product formula* assuming a 2D planar qubit architecture with nearest-neighbor connectivity.\n",
    "\n",
    "First, we load the necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "from qdk import qsharp\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Background: 2D Heisenberg Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "The quantum Heisenberg model is a statistical mechanical model used in the study of magnetic systems. The family of Heisenberg model Hamiltonians considered in this notebook consist of three types of interaction terms between adjacent lattice sites: $\\{X \\otimes X, Y \\otimes Y, Z \\otimes Z\\}$. For our purposes we consider that the interaction strength $J$ is the same for each term. Formally, the Heisenberg model Hamiltonian on an $N \\times N$ lattices divided into sets of commuting terms is formulated as:\n",
    "$$ H = \\underbrace{J \\sum_{i, j} X_i \\otimes X_j}_{A} + \\underbrace{J \\sum_{i, j} Z_i \\otimes Z_j}_{B} + \\underbrace{J \\sum_{i, j} Y_i \\otimes Y_j}_{C}.$$\n",
    "\n",
    "TThe time evolution $e^{-iHt}$ for the Hamiltonian is simulated with the fourth-order Trotter-Suzuki product formula so that any errors in simulation are sufficiently small. Essentially, this is done by simulating the evolution for small slices of time $\\Delta$ and repeating this for `nSteps` $= \\lceil t/\\Delta \\rceil$ to obtain the full time evolution. The Trotter-Suzuki formula for higher orders can be recursively defined using a *fractal decomposition* as discussed in Section 3 of [Hatanao and Suziki's survey](https://link.springer.com/chapter/10.1007/11526216_2). Then the fourth order formula $U_4(\\Delta)$ can be constructed using the second-order one $U_2(\\Delta)$ as follows.\n",
    "$$\n",
    "\\begin{aligned}\n",
    "U_2(\\Delta) & = e^{-iA\\Delta/2} e^{-iB\\Delta/2} e^{-iC\\Delta} e^{-iB\\Delta/2} e^{-iA\\Delta/2}; \\\\\n",
    "U_4(\\Delta) & = U_2(p\\Delta)U_2(p\\Delta)U_2((1 - 4p)\\Delta)U_2(p\\Delta)U_2(p\\Delta); \\\\\n",
    "p & = (4 - 4^{1/3})^{-1}.\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "For the rest of the notebook, we will present the code that computes the time evolution in a step by step fashion."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Implementation\n",
    "\n",
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "Note that expanding $U_4(\\Delta)$ to express it in terms of $A, B, C$ gives:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "U_4(\\Delta) & = e^{-iAp\\Delta/2} \\ T_1 \\ T_2 \\ T_1 \\ e^{-iAp\\Delta/2}, \\\\\n",
    "\\text{where } T_1 & = e^{-iBp\\Delta/2} e^{-iCp\\Delta} e^{-iBp\\Delta/2} e^{-iAp\\Delta} e^{-iBp\\Delta/2} e^{-iCp\\Delta} e^{-iBp\\Delta/2}, \\\\\n",
    "\\text{and } T_2 & = e^{-iA(1 - 3p)\\Delta/2} e^{-iB(1-4p)\\Delta/2} e^{-iC(1-4p)\\Delta} e^{-iB(1-4p)\\Delta/2} e^{-iA(1 - 3p)\\Delta/2}.\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "The above equation has $21$ exponential terms and works for one time step. For `nSteps` $> 1$ time steps, some adjacent terms can be merged to give $20\\lceil t/\\Delta \\rceil+1$ exponential terms for time evolution $e^{-iHt}$.\n",
    "\n",
    "The function below sets the rotation angles used to apply the exponentials involving $A, B$ and $C$ in the formula above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "microsoft": {
     "language": "qsharp"
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    },
    "vscode": {
     "languageId": "qsharp"
    }
   },
   "outputs": [],
   "source": [
    "%%qsharp\n",
    "function SetAngleSequence(p : Double, dt : Double, J : Double) : Double[] {\n",
    "\n",
    "    let len1 = 8;\n",
    "    let len2 = 4;\n",
    "    let valLength = 2*len1 + len2;\n",
    "\n",
    "    mutable values = [0.0, size=valLength];\n",
    "\n",
    "    let val1 = -J*p*dt/2.0;\n",
    "    let val2 = -J*(1.0-3.0*p)*dt/2.0;\n",
    "    let val3 = -J*(1.0-4.0*p)*dt/2.0;\n",
    "\n",
    "    // Angles bookending Term2\n",
    "    set values w/= len1 <- val2;\n",
    "    set values w/= len1+len2 <- val2;\n",
    "\n",
    "    for i in 0..len1-1 {\n",
    "\n",
    "        if (i % 2 == 0) {\n",
    "            set values w/= i <- val1*2.0;\n",
    "        }\n",
    "        else {\n",
    "            set values w/= i <- val1;\n",
    "        }\n",
    "    }\n",
    "\n",
    "    for i in len1+1..len1+len2-1 {\n",
    "\n",
    "        if (i % 2 == 0) {\n",
    "            set values w/= i <- val3*2.0;\n",
    "        }\n",
    "        else {\n",
    "            set values w/= i <- val3;\n",
    "        }\n",
    "    }\n",
    "\n",
    "    for i in len1+len2+1..valLength-1 {\n",
    "\n",
    "        if (i % 2 == 0) {\n",
    "            set values w/= i <- val1*2.0;\n",
    "        }\n",
    "        else {\n",
    "            set values w/= i <- val1;\n",
    "        }\n",
    "    }\n",
    "\n",
    "    return values;\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "### Quantum operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "There are three kinds of operations applied to neighboring pairs of qubits on the lattice: (i) $e^{-i(X \\otimes X)\\theta}$; (ii) $e^{-i(Z \\otimes Z)\\theta}$; and (iii) $e^{-i(Y \\otimes Y)\\theta}$.\n",
    "\n",
    "The following operation applies $e^{i(P \\otimes P)\\theta}$ for Pauli $P \\in \\{X, Y, Z\\}$ on pairs of horizontally (resp. vertically) neighboring qubits in a single row of the lattice if `dir` is `True` (resp. `False`). Additionally, the exponential is applied only to pairs where the first qubit id is even (resp. odd) if `grp` is `True` (resp. `False`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "microsoft": {
     "language": "qsharp"
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    },
    "vscode": {
     "languageId": "qsharp"
    }
   },
   "outputs": [],
   "source": [
    "%%qsharp\n",
    "operation ApplyPoPRot(n : Int, m : Int, qArr : Qubit[][], P : Pauli, theta : Double, dir : Bool, grp : Bool) : Unit {\n",
    "\n",
    "    let start = (grp ? 0 | 1);    // Choose either odd or even indices based on group number\n",
    "    let P_op = [P, P];\n",
    "    let jmp = 2;\n",
    "    let end = dir ? m-2 | m-1;\n",
    "    let v_end = dir? n-1 | n-2;\n",
    "\n",
    "    for row in 0..v_end {\n",
    "        for col in start..jmp..end {    // Either odd or even columns are picked based on group number\n",
    "\n",
    "            let row2 = dir? row | row+1;    // Choose the next row if for vertical direction\n",
    "            let col2 = dir? col+1 | col;    // Choose the next column for horizontal direction\n",
    "\n",
    "            Exp(P_op, theta, [qArr[row][col], qArr[row2][col2]]);\n",
    "        }\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "The next operation puts everything together and calls the operations needed to\n",
    "simulate the Heisenberg model Hamiltonian using a fourth order product formula.\n",
    "Observe that the `ApplyPoPRot` operation is called four times for each row for different\n",
    "choices of direction and starting index to ensure all possible pairs of qubits\n",
    "are appropriately considered.\n",
    "\n",
    "The various parameters taken in by the operation correspond to:\n",
    "\n",
    "- `J`: parameter by which the Hamiltonian terms are scaled.\n",
    "- `N1`, `N2`: row and columns size of the lattice.\n",
    "- `totTime`: the total time for the Trotter simulation.\n",
    "- `dt` : the step size for the simulation, sometimes denoted as $\\Delta$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "microsoft": {
     "language": "qsharp"
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    },
    "vscode": {
     "languageId": "qsharp"
    }
   },
   "outputs": [],
   "source": [
    "%%qsharp\n",
    "import Std.Math.*;\n",
    "import Std.Arrays.*;\n",
    "\n",
    "operation HeisModel2DSim(N1 : Int, N2 : Int, J : Double, totTime : Double, dt : Double) : Unit {\n",
    "\n",
    "    use qs = Qubit[N1*N2];\n",
    "\n",
    "    let qubitArray = Chunks(N2, qs); // Create a 2D array of qubits\n",
    "\n",
    "    let len = Length(qs);\n",
    "\n",
    "    let p = 1.0/(4.0 - 4.0^(1.0/3.0));\n",
    "\n",
    "    let nsteps = Ceiling(totTime/dt);\n",
    "\n",
    "    let seqLen = 20*nsteps+1;\n",
    "\n",
    "    let angSeq = SetAngleSequence(p, dt, J);\n",
    "\n",
    "    let paulis = [PauliX, PauliZ, PauliY];\n",
    "\n",
    "    for i in 0..seqLen-1 {\n",
    "\n",
    "        let j = (i%4==0) ? 0 | ( (i%2==1) ? 1 | 2); // Choose which pauli rotation will be applied\n",
    "        let theta = (i==0 or i==seqLen-1) ? -J*p*dt/2.0 | angSeq[i%20];\n",
    "\n",
    "        for (dir, grp) in [(true, true), (true, false), (false, true), (false, false)] {\n",
    "            ApplyPoPRot(N1, N2, qubitArray, paulis[j], theta, dir, grp);\n",
    "        }\n",
    "    }\n",
    "\n",
    "    //Resetting qubits before release;\n",
    "    ResetAll(qs);\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that for typical choices of `totTime` and `dt` used in real-world applications, the `seqLen` can range in the hundreds or thousands of terms each of whose resources will be individually estimated. This results in a lot of time consumed (> 15 min) to compute the resource estimates. However, from the expansion of $U_4(\\Delta)$, it is clear that there are a few unique terms that repeat `nsteps` number of times. Hence, to efficiently estimate resources, we compute the resources for running one step of $U_4(\\Delta)$ and multiply the cost by `nsteps` repetitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "qsharp"
    }
   },
   "outputs": [],
   "source": [
    "%%qsharp\n",
    "import Std.ResourceEstimation.*;\n",
    "\n",
    "operation EstimateHeisModel2DSim(N1 : Int, N2 : Int, J : Double, totTime : Double, dt : Double) : Unit {\n",
    "\n",
    "    use qs = Qubit[N1*N2];\n",
    "\n",
    "    let qubitArray = Chunks(N2, qs); // Create a 2D array of qubits\n",
    "\n",
    "    let len = Length(qs);\n",
    "\n",
    "    let p = 1.0/(4.0 - 4.0^(1.0/3.0));\n",
    "\n",
    "    let t = Ceiling(totTime/dt);\n",
    "\n",
    "    let angSeq = SetAngleSequence(p, dt, J);\n",
    "\n",
    "    let paulis = [PauliX, PauliZ, PauliY];\n",
    "\n",
    "    within {\n",
    "        RepeatEstimates(t);\n",
    "    }\n",
    "    apply {\n",
    "        for i in 0..19 {\n",
    "            let j = (i%4==0) ? 0 | ( (i%2==1) ? 1 | 2); // Choose which pauli rotation will be applied\n",
    "            let theta = angSeq[i%20];\n",
    "\n",
    "            for (dir, grp) in [(true, true), (true, false), (false, true), (false, false)] {\n",
    "                ApplyPoPRot(N1, N2, qubitArray, paulis[j], theta, dir, grp);\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    let angle = -J*p*dt/2.0;\n",
    "    for (dir, grp) in [(true, true), (true, false), (false, true), (false, false)] {\n",
    "        ApplyPoPRot(N1, N2, qubitArray, paulis[0], angle, dir, grp);\n",
    "    }\n",
    "\n",
    "    //Resetting qubits before release;\n",
    "    ResetAll(qs);\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Getting logical resource counts\n",
    "\n",
    "For the purpose of generating the rQOPS for some target runtime, it suffices that we obtain the logical resource estimates to simulate the Heisenberg model Hamiltonian. We consider three problem instances with lattice sizes $\\{20 \\times 20, 30 \\times 30, 40 \\times 40\\}$ with $J = 1.0$. These instances are simulated for a total time of $1000$ s, with step size `dt`$ = 0.1$, and overall probability of failure $\\varepsilon = 0.01$. Any one of the six pre-defined qubit parameters will do to obtain the logical coounts and in this notebook we choose a Majorana based qubit with the `floquet code`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up problem parameters\n",
    "N1 = [20, 30, 40]\n",
    "N2 = [20, 30, 40]\n",
    "J = 1.0\n",
    "totTime = 1000.0\n",
    "dt = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "We submit a resource estimation job with all the problem instances sequentially and collect the estimates in `results`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit jobs\n",
    "results = []\n",
    "\n",
    "for i in range(3):\n",
    "    qsharp_string = f\"EstimateHeisModel2DSim({N1[i]}, {N2[i]}, {J}, {totTime}, {dt})\"\n",
    "\n",
    "    result = qsharp.estimate(qsharp_string, params={\"errorBudget\": 0.01, \"qubitParams\": {\"name\": \"qubit_maj_ns_e6\"}, \"qecScheme\": {\"name\": \"floquet_code\"}, \"constraints\": {\"logicalDepthFactor\": 4}})\n",
    "    results.append(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see the complete information provided when invoking the resource estimator, we output the result for the $20 \\times 20$ lattice by displaying `results[0]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Displaying estimates for 20x20 lattice size.\n",
    "results[0]\n",
    "# Change index to 1 (resp. 2) for 30x30 (resp. 40x40) lattice size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Visualizing and understanding the results\n",
    "\n",
    "### Result summary table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "# Define function to display information in summary format\n",
    "def get_summary_table(results, labels):\n",
    "    logical_qubits = []\n",
    "    logical_depth = []\n",
    "    t_states = []\n",
    "    code_distance = []\n",
    "    t_factories = []\n",
    "    t_factory_fraction = []\n",
    "    physical_qubits = []\n",
    "    rqops = []\n",
    "    runtime = []\n",
    "    logical_error = []\n",
    "\n",
    "    for i in range(3):\n",
    "        logical_qubits.append(results[i]['physicalCounts']['breakdown']['algorithmicLogicalQubits'])\n",
    "        logical_depth.append(results[i]['physicalCountsFormatted']['logicalDepth'])\n",
    "        t_states.append(results[i]['physicalCountsFormatted']['numTstates'])\n",
    "        t_factories.append(results[i]['physicalCounts']['breakdown']['numTfactories'])\n",
    "        logical_error.append(results[i]['physicalCountsFormatted']['requiredLogicalQubitErrorRate'])\n",
    "        physical_qubits.append(results[i]['physicalCountsFormatted']['physicalQubits'])\n",
    "        rqops.append(results[i]['physicalCountsFormatted']['rqops'])\n",
    "        runtime.append(results[i]['physicalCountsFormatted']['runtime'])\n",
    "        code_distance.append(results[i]['logicalQubit']['codeDistance'])\n",
    "        t_factory_fraction.append(results[i]['physicalCountsFormatted']['physicalQubitsForTfactoriesPercentage'])\n",
    "\n",
    "    data = pd.DataFrame()\n",
    "    pd.options.display.float_format = '{:.2E}'.format\n",
    "    data['Logical qubits'] = logical_qubits\n",
    "    data['Logical depth'] = logical_depth\n",
    "    data['Logical error'] = logical_error\n",
    "    data['T states'] = t_states\n",
    "    # data['T states'] = data['T states'].astype('float64')\n",
    "    data['Code Distance'] = code_distance\n",
    "    data['T factories'] = t_factories\n",
    "    data['T factory fraction'] = t_factory_fraction\n",
    "    data['Physical qubits'] = physical_qubits\n",
    "    data['rQOPS'] = rqops\n",
    "    data['Physical runtime'] = runtime\n",
    "    data.index = labels\n",
    "\n",
    "    return data\n",
    "\n",
    "# Display summarized information for all problem instances\n",
    "labels = [\"Hei20\", \"Hei30\", \"Hei40\"]\n",
    "table = get_summary_table(results, labels)\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Note that in general, there is a trade-off between the logical depth and number of T factories used. \n",
    "\n",
    "> To ensure that T factories do not dominate the resource requirements, we set the `logical_depth_factor`${}=4$ adding some number of `noops` to increase the logical depth."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting the target rQOPS\n",
    "\n",
    "While the resource estimator generates a runtime for the given hardware profile, we want to set a target runtime of 1 week i.e., 604800 seconds to obtain a practical quantum advantage. We collect the previous results to compute the corresponding target rQOPS as \n",
    "$$ \\text{Target rQOPS} = \\frac{\\text{Logical qubits}\\cdot\\text{Logical Depth}}{\\text{Target runtime}}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_target_rqops(results, labels):\n",
    "\n",
    "    target_runtime = 604800\n",
    "    logical_qubits = []\n",
    "    logical_depth = []\n",
    "    target_rqops = []\n",
    "    logical_error = []\n",
    "\n",
    "    for i in range(3):\n",
    "        logical_qubits.append(results[i]['physicalCounts']['breakdown']['algorithmicLogicalQubits'])\n",
    "        logical_depth.append(results[i]['physicalCountsFormatted']['logicalDepth'])\n",
    "        logical_error.append(results[i]['physicalCountsFormatted']['requiredLogicalQubitErrorRate'])\n",
    "        target_rqops.append(round(results[i]['physicalCounts']['breakdown']['algorithmicLogicalQubits'] * results[i]['physicalCounts']['breakdown']['logicalDepth'] / target_runtime))\n",
    "\n",
    "    data = pd.DataFrame()\n",
    "    pd.options.display.float_format = '{:.2e}'.format\n",
    "    data['Logical qubits'] = logical_qubits\n",
    "    data['Logical depth'] = logical_depth\n",
    "    data['Logical error'] = logical_error\n",
    "    data['Target rQOPS'] = target_rqops\n",
    "    data['Target rQOPS'] = data['Target rQOPS'].astype('float64')\n",
    "    data.index = labels\n",
    "\n",
    "    return data\n",
    "\n",
    "rQOPS_table = get_target_rqops(results, labels)\n",
    "rQOPS_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Next steps\n",
    "\n",
    "Feel free to use this notebook as a starting point for your own experiments.  For\n",
    "example, you can\n",
    "\n",
    "* explore how the results change considering other problem instances of the Heisenberg model\n",
    "* explore space- and time-trade-offs by changing the value for\n",
    "  `logical_depth_factor` or `max_t_factories`\n",
    "* visualize these trade-offs with the space and time diagrams\n",
    "* use other or customized qubit parameters"
   ]
  }
 ],
 "metadata": {
  "kernel_info": {
   "name": "python3"
  },
  "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.7"
  },
  "nteract": {
   "version": "nteract-front-end@1.0.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
