{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "d56e6c6d",
      "metadata": {},
      "source": [
        "# Hamiltonian Simulation\n",
        "\n",
        "In this example we're going to apply the previously learned concept to run a hamiltonian simulation.\n",
        "As before we are going to define a symbolic circuit as an ansatz."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3563e447",
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install tierkreis pytket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bfbaa4df",
      "metadata": {},
      "outputs": [],
      "source": [
        "from pytket._tket.circuit import Circuit, fresh_symbol\n",
        "\n",
        "\n",
        "def build_ansatz() -> Circuit:\n",
        "    a = fresh_symbol(\"a\")\n",
        "    b = fresh_symbol(\"b\")\n",
        "    c = fresh_symbol(\"c\")\n",
        "    circ = Circuit(4)\n",
        "    circ.CX(0, 1)\n",
        "    circ.CX(1, 2)\n",
        "    circ.CX(2, 3)\n",
        "    circ.Rz(a, 3)\n",
        "    circ.CX(2, 3)\n",
        "    circ.CX(1, 2)\n",
        "    circ.CX(0, 1)\n",
        "    circ.Rz(b, 0)\n",
        "    circ.CX(0, 1)\n",
        "    circ.CX(1, 2)\n",
        "    circ.CX(2, 3)\n",
        "    circ.Rz(c, 3)\n",
        "    circ.CX(2, 3)\n",
        "    circ.CX(1, 2)\n",
        "    circ.CX(0, 1)\n",
        "    return circ"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d1993864",
      "metadata": {},
      "source": [
        "We are going to simulate a hamiltonian given as a list of Pauli strings and their weights.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f5d2e7be",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.builder import GraphBuilder\n",
        "from tierkreis.controller.data.models import TKR\n",
        "from typing import NamedTuple, Literal\n",
        "\n",
        "from substitution_worker import substitute\n",
        "\n",
        "\n",
        "class SymbolicExecutionInputs(NamedTuple):\n",
        "    a: TKR[float]\n",
        "    b: TKR[float]\n",
        "    c: TKR[float]\n",
        "    ham: TKR[list[tuple[Literal[\"pytket._tket.pauli.QubitPauliString\"], float]]]\n",
        "    ansatz: TKR[Literal[\"pytket._tket.circuit.Circuit\"]]\n",
        "\n",
        "\n",
        "simulation_graph = GraphBuilder(SymbolicExecutionInputs, TKR[float])\n",
        "substituted_circuit = simulation_graph.task(\n",
        "    substitute(\n",
        "        simulation_graph.inputs.ansatz,\n",
        "        simulation_graph.inputs.a,\n",
        "        simulation_graph.inputs.b,\n",
        "        simulation_graph.inputs.c,\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "161a5e02",
      "metadata": {},
      "source": [
        "We will evaluate this circuit with an observable based on a Pauli string"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "39b21062",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.pytket_worker import (\n",
        "    append_pauli_measurement_impl,\n",
        "    optimise_phase_gadgets,\n",
        "    expectation,\n",
        ")\n",
        "from tierkreis.aer_worker import submit_single\n",
        "\n",
        "\n",
        "class SubmitInputs(NamedTuple):\n",
        "    circuit: TKR[Literal[\"pytket._tket.circuit.Circuit\"]]\n",
        "    pauli_string: TKR[Literal[\"pytket._tket.pauli.QubitPauliString\"]]\n",
        "    n_shots: TKR[int]\n",
        "\n",
        "\n",
        "def exp_val():\n",
        "    g = GraphBuilder(SubmitInputs, TKR[float])\n",
        "\n",
        "    circuit = g.inputs.circuit\n",
        "    pauli_string = g.inputs.pauli_string\n",
        "    n_shots = g.inputs.n_shots\n",
        "\n",
        "    measurement_circuit = g.task(append_pauli_measurement_impl(circuit, pauli_string))\n",
        "\n",
        "    compiled_circuit = g.task(optimise_phase_gadgets(measurement_circuit))\n",
        "\n",
        "    backend_result = g.task(submit_single(compiled_circuit, n_shots))\n",
        "    av = g.task(expectation(backend_result))\n",
        "    g.outputs(av)\n",
        "    return g"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c7f5f9a1",
      "metadata": {},
      "source": [
        "Since `exp_val` runs independently for each Pauli string we can implement this using a map, \n",
        "but we need to prepare the inputs first."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "021e0e5b",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.builtins import unzip\n",
        "from tierkreis.controller.data.models import TKR\n",
        "\n",
        "pauli_strings_list, parameters_list = simulation_graph.task(\n",
        "    unzip(simulation_graph.inputs.ham)\n",
        ")\n",
        "input_circuits = simulation_graph.map(\n",
        "    lambda x: SubmitInputs(substituted_circuit, x, simulation_graph.const(100)),\n",
        "    pauli_strings_list,\n",
        ")\n",
        "exp_values = simulation_graph.map(exp_val(), input_circuits)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7788777e",
      "metadata": {},
      "source": [
        "To estimate the energy, we can take a weighted sum of the expectation values.\n",
        "For this we want to compute a reduction (\\(x,y) \\z --> x*y+z) which we implement as a fold function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cd24e08c",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.builtins import add, times, untuple\n",
        "from tierkreis.graphs.fold import FoldFunctionInput\n",
        "\n",
        "ComputeTermsInputs = FoldFunctionInput[\n",
        "    tuple[float, float], float\n",
        "]  # (value, accum) -> new_accum\n",
        "\n",
        "\n",
        "def compute_terms():\n",
        "    g = GraphBuilder(ComputeTermsInputs, TKR[float])\n",
        "\n",
        "    res_0, res_1 = g.task(untuple(g.inputs.value))\n",
        "    prod = g.task(times(res_0, res_1))\n",
        "    sum = g.task(add(g.inputs.accum, prod))\n",
        "\n",
        "    g.outputs(sum)\n",
        "    return g"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "af57c8ac",
      "metadata": {},
      "source": [
        "Preparing the inputs for the fold graph we are going to use tuples (x=exp_val, y=weight)  and defining the start value z=0."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9bab0de3",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.builtins import tkr_zip\n",
        "from tierkreis.graphs.fold import FoldGraphInputs, fold_graph\n",
        "\n",
        "tuple_values = simulation_graph.task(tkr_zip(exp_values, parameters_list))\n",
        "fold_inputs = FoldGraphInputs(simulation_graph.const(0.0), tuple_values)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "da997ea2",
      "metadata": {},
      "source": [
        "applying the fold operation yields the final output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d7c6fa18",
      "metadata": {},
      "outputs": [],
      "source": [
        "computed = simulation_graph.eval(fold_graph(compute_terms()), fold_inputs)\n",
        "simulation_graph.outputs(computed)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5999c0a5",
      "metadata": {},
      "source": [
        "As before we now have to set up tierkreis storage and executors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "aec32e19",
      "metadata": {},
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "from uuid import UUID\n",
        "\n",
        "from tierkreis.consts import PACKAGE_PATH\n",
        "from tierkreis.controller.executor.multiple import MultipleExecutor\n",
        "from tierkreis.storage import FileStorage\n",
        "from tierkreis.controller.executor.uv_executor import UvExecutor\n",
        "\n",
        "storage = FileStorage(UUID(int=102), name=\"hamiltonian\")\n",
        "example_executor = UvExecutor(\n",
        "    registry_path=Path().parent / \"example_workers\", logs_path=storage.logs_path\n",
        ")\n",
        "common_executor = UvExecutor(\n",
        "    registry_path=PACKAGE_PATH.parent / \"tierkreis_workers\", logs_path=storage.logs_path\n",
        ")\n",
        "multi_executor = MultipleExecutor(\n",
        "    common_executor,\n",
        "    executors={\"custom\": example_executor},\n",
        "    assignments={\"substitution_worker\": \"custom\"},\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "20848215",
      "metadata": {},
      "source": [
        "and provide the inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2d3b86c0",
      "metadata": {},
      "outputs": [],
      "source": [
        "from pytket._tket.unit_id import Qubit\n",
        "from pytket.pauli import Pauli, QubitPauliString\n",
        "\n",
        "qubits = [Qubit(0), Qubit(1), Qubit(2), Qubit(3)]\n",
        "hamiltonian = [\n",
        "    (QubitPauliString(qubits, [Pauli.X, Pauli.Y, Pauli.X, Pauli.I]).to_list(), 0.1),\n",
        "    (QubitPauliString(qubits, [Pauli.Y, Pauli.Z, Pauli.X, Pauli.Z]).to_list(), 0.5),\n",
        "    (QubitPauliString(qubits, [Pauli.X, Pauli.Y, Pauli.Z, Pauli.I]).to_list(), 0.3),\n",
        "    (QubitPauliString(qubits, [Pauli.Z, Pauli.Y, Pauli.X, Pauli.Y]).to_list(), 0.6),\n",
        "]\n",
        "inputs = {\n",
        "    \"ansatz\": build_ansatz().to_dict(),\n",
        "    \"a\": 0.2,\n",
        "    \"b\": 0.55,\n",
        "    \"c\": 0.75,\n",
        "    \"ham\": hamiltonian,\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "037d98b1",
      "metadata": {},
      "source": [
        "before we can run the simulation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "53446615",
      "metadata": {},
      "outputs": [],
      "source": [
        "from tierkreis.controller import run_graph\n",
        "from tierkreis.storage import read_outputs\n",
        "\n",
        "\n",
        "storage.clean_graph_files()\n",
        "run_graph(\n",
        "    storage,\n",
        "    multi_executor,\n",
        "    simulation_graph,\n",
        "    inputs,\n",
        "    polling_interval_seconds=0.2,\n",
        ")\n",
        "output = read_outputs(simulation_graph, storage)\n",
        "print(output)"
      ]
    }
  ],
  "metadata": {
    "execution": {
      "timeout": -1
    },
    "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
}