{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e18767c",
   "metadata": {},
   "source": [
    "# QSCI\n",
    "\n",
    "QSCI (see [arxiv:2302.11320](https://arxiv.org/pdf/2302.11320)) is a group of hybrid algorithms to approximate the ground state (or low energy eigenstates) of a system hamiltonian, usually for molecule structure problems.\n",
    "In principle it functions similar to other CI methods:\n",
    "- Identify a suitable truncated Fock Space through sampling from a quantum state\n",
    "- Diagonalize the subspace hamiltonian to solve the eigenvalue problem classically\n",
    "\n",
    "For more details we refer the reader to the aforementioned article.\n",
    "\n",
    "This example serves as a demonstration to show how existing algorithms can be ported to a tierkreis graph to leverage its automatic scaling features.\n",
    "For this were going to break down QSCI into 4 tasks:\n",
    "1. Construct the target circuits\n",
    "    1. Build a state preparation circuit ansatz (JW)\n",
    "    2. Generate circuits simulating the system hamiltonian (through time evolution)\n",
    "2. Run (simulate) the quantum circuits to sample the configurations\n",
    "3. Generate and diagonalize the approximate Hamiltonian to find the ground state energy\n",
    "The granularity of these tasks is kept very low for comprehensibility.\n",
    "For tierkreis worker tasks are considered atomic and will run in parallel if the data allows it.\n",
    "It is possible to break down the tasks even further if desired.\n",
    "\n",
    "We will bundle most of the functionality in a worker but repeat the here for brevity.\n",
    "Using the worker functions we will construct a tierkreis graph implementing QSCI.\n",
    "\n",
    "## The core tasks\n",
    "In the next section we will provide the implementations for 1.1, 1.2 and 3. which are part of the `qsci_worker`.\n",
    "First we will define some helper types:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c9b8dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install tierkreis pytket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9318ae35",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import NamedTuple\n",
    "\n",
    "\n",
    "class Molecule(NamedTuple):\n",
    "    geometry: list[tuple[str, list[float]]]\n",
    "    basis: str\n",
    "    charge: int\n",
    "\n",
    "\n",
    "class CompleteActiveSpace(NamedTuple):\n",
    "    n: int\n",
    "    n_ele: int\n",
    "\n",
    "\n",
    "class Hamiltonian(NamedTuple):\n",
    "    h0: float\n",
    "    h1: list[list[float]]\n",
    "    h2: list[list[list[list[float]]]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0f05e29",
   "metadata": {},
   "source": [
    "Then we define our core functionality.\n",
    "Here we reference our core library `qsci` which is part of the worker code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a676aa07",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from typing import cast\n",
    "\n",
    "import numpy as np\n",
    "from example_workers.qsci_worker.src.qsci.active_space import get_n_active, get_n_core\n",
    "from example_workers.qsci_worker.src.qsci.jordan_wigner import (\n",
    "    qubit_mapping_jordan_wigner,\n",
    ")\n",
    "from example_workers.qsci_worker.src.qsci.state_prep import perform_state_preparation\n",
    "from example_workers.qsci_worker.src.qsci.utils import rhf2ghf\n",
    "from pytket._tket.circuit import Circuit\n",
    "from pytket.pauli import QubitPauliString\n",
    "from pytket.utils.operators import CoeffTypeAccepted, QubitPauliOperator\n",
    "\n",
    "\n",
    "def state_prep(\n",
    "    ham_init: Hamiltonian,\n",
    "    reference_state: list[int],\n",
    "    max_iteration_prep: int,\n",
    "    atol: float,\n",
    "    mo_occ: list[int],\n",
    "    cas_init: CompleteActiveSpace,\n",
    "    cas_hsim: CompleteActiveSpace,\n",
    ") -> Circuit:\n",
    "    ham_init_operator = QubitPauliOperator(\n",
    "        cast(\n",
    "            \"dict[QubitPauliString, CoeffTypeAccepted]\",\n",
    "            qubit_mapping_jordan_wigner(\n",
    "                *rhf2ghf(\n",
    "                    ham_init.h0,\n",
    "                    np.array(ham_init.h1),\n",
    "                    np.array(ham_init.h2),\n",
    "                ),\n",
    "            ),\n",
    "        ),\n",
    "    )\n",
    "    # time-evolve CASCI ground state.\n",
    "    n_core_init = get_n_core(mo_occ, cas_init.n_ele)\n",
    "    n_core_hsim = get_n_core(mo_occ, cas_hsim.n_ele)\n",
    "    n_core = n_core_init - n_core_hsim\n",
    "    logging.info(\n",
    "        f\"mo_occ={mo_occ} n_cas_hsim={cas_hsim.n} n_elecas_hsim={cas_hsim.n_ele}\",\n",
    "    )\n",
    "    n_active_hsim = get_n_active(mo_occ, cas_hsim.n, cas_hsim.n_ele)\n",
    "    prepared_circ = Circuit(n_active_hsim * 2)\n",
    "    for i in range(n_core * 2):\n",
    "        prepared_circ.X(i)\n",
    "    return perform_state_preparation(\n",
    "        reference_state=reference_state,\n",
    "        ham_init=ham_init_operator,\n",
    "        n_cas_init=cas_init.n,\n",
    "        max_iteration=max_iteration_prep,\n",
    "        atol=atol,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f21464c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from example_workers.qsci_worker.src.qsci.utils import make_time_evolution_circuits\n",
    "from pytket.circuit import Qubit\n",
    "\n",
    "\n",
    "def circuits_from_hamiltonians(\n",
    "    ham_init: Hamiltonian,\n",
    "    ham_hsim: Hamiltonian,\n",
    "    adapt_circuit: Circuit,\n",
    "    t_step_list: list[float],\n",
    "    cas_init: CompleteActiveSpace,\n",
    "    cas_hsim: CompleteActiveSpace,\n",
    "    mo_occ: list[int],\n",
    "    max_cx_gates: int,\n",
    ") -> list[Circuit]:\n",
    "    ham_init_operator = QubitPauliOperator(\n",
    "        cast(\n",
    "            \"dict[QubitPauliString, CoeffTypeAccepted]\",\n",
    "            qubit_mapping_jordan_wigner(\n",
    "                *rhf2ghf(\n",
    "                    ham_init.h0,\n",
    "                    np.array(ham_init.h1),\n",
    "                    np.array(ham_init.h2),\n",
    "                ),\n",
    "            ),\n",
    "        ),\n",
    "    )\n",
    "    ham_hsim_operator = QubitPauliOperator(\n",
    "        cast(\n",
    "            \"dict[QubitPauliString, CoeffTypeAccepted]\",\n",
    "            qubit_mapping_jordan_wigner(\n",
    "                *rhf2ghf(\n",
    "                    ham_hsim.h0,\n",
    "                    np.array(ham_hsim.h1),\n",
    "                    np.array(ham_hsim.h2),\n",
    "                ),\n",
    "            ),\n",
    "        ),\n",
    "    )\n",
    "    # Load the input data.\n",
    "    n_core_init = get_n_core(mo_occ, cas_init.n_ele)\n",
    "    n_core_hsim = get_n_core(mo_occ, cas_hsim.n_ele)\n",
    "    n_core = n_core_init - n_core_hsim\n",
    "    n_active_hsim = get_n_active(mo_occ, cas_hsim.n, cas_hsim.n_ele)\n",
    "    prepared_circ = Circuit(n_active_hsim * 2)\n",
    "    for i in range(n_core * 2):\n",
    "        prepared_circ.X(i)\n",
    "    prepared_circ.add_circuit(\n",
    "        adapt_circuit,\n",
    "        [2 * n_core + i for i in range(adapt_circuit.n_qubits)],\n",
    "    )\n",
    "    ham_init_shifted = QubitPauliOperator(\n",
    "        {\n",
    "            QubitPauliString(\n",
    "                {\n",
    "                    Qubit(qubit.index[0] + 2 * n_core): pauli\n",
    "                    for qubit, pauli in qps.map.items()\n",
    "                },\n",
    "            ): coeff\n",
    "            for qps, coeff in ham_init_operator._dict.items()\n",
    "        },\n",
    "    )\n",
    "    return make_time_evolution_circuits(\n",
    "        t_step_list,\n",
    "        prepared_circ,\n",
    "        h_hsim=ham_hsim_operator,\n",
    "        h_init=ham_init_shifted,\n",
    "        max_cx_gates=max_cx_gates,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8de28c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "from example_workers.qsci_worker.src.qsci.postprocess import (\n",
    "    get_ci_matrix,\n",
    "    postprocess_configs,\n",
    ")\n",
    "from example_workers.qsci_worker.src.qsci.utils import get_config_from_cas_init\n",
    "from pytket.backends.backendresult import BackendResult\n",
    "\n",
    "\n",
    "def energy_from_results(\n",
    "    ham_hsim: Hamiltonian,\n",
    "    backend_results: list[BackendResult],\n",
    "    mo_occ: list[int],\n",
    "    cas_init: CompleteActiveSpace,\n",
    "    cas_hsim: CompleteActiveSpace,\n",
    ") -> float:\n",
    "    counts = Counter()\n",
    "    for r in backend_results:\n",
    "        for k, v in r.get_counts().items():\n",
    "            counts[k] += v\n",
    "    phis = list(counts.keys())\n",
    "    phis_init_orig = get_config_from_cas_init(\n",
    "        mo_occ,\n",
    "        cas_init.n,\n",
    "        cas_init.n_ele,\n",
    "        cas_hsim.n,\n",
    "        cas_hsim.n_ele,\n",
    "    )\n",
    "    for p in phis_init_orig:\n",
    "        if p not in phis:\n",
    "            phis.append(p)\n",
    "    # phis = get_config(backend_results)\n",
    "    logging.info(f\"CONFIG (before): {len(phis)}\")\n",
    "    phis = postprocess_configs(phis_init_orig[0], phis)\n",
    "    logging.info(f\"CONFIG (after):  {len(phis)}\")\n",
    "\n",
    "    hsd = get_ci_matrix(\n",
    "        phis,\n",
    "        h1=np.array(ham_hsim.h1),\n",
    "        h2=np.array(ham_hsim.h2),\n",
    "        enuc=ham_hsim.h0,\n",
    "    )\n",
    "    energy = np.linalg.eigh(hsd.todense())[0][0]\n",
    "    logging.info(f\"ENERGY: {energy}\")\n",
    "\n",
    "    return energy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb6ee165",
   "metadata": {},
   "source": [
    "Additionally were going to use a helper function to generate the hamiltonian from the molecule definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d4d4ae7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from example_workers.qsci_worker.src.chemistry.active_space import get_frozen\n",
    "from example_workers.qsci_worker.src.chemistry.molecule import extract_hamiltonian_rhf\n",
    "\n",
    "\n",
    "def make_ham(\n",
    "    molecule: Molecule,\n",
    "    mo_occ: list[int],\n",
    "    cas: CompleteActiveSpace,\n",
    ") -> Hamiltonian:\n",
    "    # Construct the frozen orbital lists.\n",
    "    frozen = get_frozen(mo_occ, cas.n, cas.n_ele)\n",
    "    h0, h1, h2 = extract_hamiltonian_rhf(\n",
    "        molecule.geometry,\n",
    "        molecule.basis,\n",
    "        charge=molecule.charge,\n",
    "        frozen=frozen,\n",
    "    )\n",
    "    return Hamiltonian(h0=h0, h1=h1.tolist(), h2=h2.tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44509762",
   "metadata": {},
   "source": [
    "As before we can annotate all these functions with the `@worker.task()` decorator to make the functionality available in tierkreis.\n",
    "\n",
    "## Constructing the QSCI graph\n",
    "\n",
    "Using the above defined functions (now as a worker) were going to use preexisting workest to compile and run a circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96dbf6b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tierkreis.aer_worker import submit_single\n",
    "from tierkreis.builder import GraphBuilder\n",
    "from tierkreis.controller.data.models import TKR, OpaqueType\n",
    "from tierkreis.quantinuum_worker import compile_circuit_quantinuum\n",
    "\n",
    "\n",
    "def _compile_and_run() -> GraphBuilder[\n",
    "    TKR[OpaqueType[\"pytket._tket.circuit.Circuit\"]],  # noqa: F821\n",
    "    TKR[OpaqueType[\"pytket.backends.backendresult.BackendResult\"]],  # noqa: F821\n",
    "]:\n",
    "    g = GraphBuilder(\n",
    "        TKR[OpaqueType[\"pytket._tket.circuit.Circuit\"]],\n",
    "        TKR[OpaqueType[\"pytket.backends.backendresult.BackendResult\"]],\n",
    "    )\n",
    "\n",
    "    n_shots = g.const(500)\n",
    "    compiled_circuit = g.task(compile_circuit_quantinuum(g.inputs))\n",
    "    backend_result = g.task(submit_single(compiled_circuit, n_shots))\n",
    "\n",
    "    g.outputs(backend_result)\n",
    "    return g"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "099b4b14",
   "metadata": {},
   "source": [
    "Defining the inputs and outputs of the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50fc9fd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class QSCIInputs(NamedTuple):\n",
    "    molecule: TKR[Molecule]\n",
    "    mo_occ: TKR[list[int]]\n",
    "    reference_state: TKR[list[int]]\n",
    "    cas_init: TKR[CompleteActiveSpace]\n",
    "    cas_hsim: TKR[CompleteActiveSpace]\n",
    "    t_step_list: TKR[list[float]]\n",
    "    max_iteration_prep: TKR[int]\n",
    "    max_cx_gates_hsim: TKR[int]\n",
    "    atol: TKR[float]\n",
    "\n",
    "\n",
    "class QSCIOutputs(NamedTuple):\n",
    "    energy: TKR[float]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f10a1e1",
   "metadata": {},
   "source": [
    "leads to the definition of the graph.\n",
    "We are redefining the above functions to point at the correct stubs (as we would normally use them)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d502187",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qsci_worker import (  # noqa: F811\n",
    "    make_ham,\n",
    "    circuits_from_hamiltonians,\n",
    "    energy_from_results,\n",
    "    make_ham,\n",
    "    state_prep,\n",
    ")\n",
    "\n",
    "qsci_graph = GraphBuilder(QSCIInputs, QSCIOutputs)\n",
    "# Separate tasks 'make_h_init'+'state_pre' and 'make_h_hsim' run in parallel\n",
    "ham_init = qsci_graph.task(\n",
    "    make_ham(\n",
    "        qsci_graph.inputs.molecule,\n",
    "        qsci_graph.inputs.mo_occ,\n",
    "        qsci_graph.inputs.cas_init,\n",
    "    ),\n",
    ")\n",
    "ham_hsim = qsci_graph.task(\n",
    "    make_ham(\n",
    "        qsci_graph.inputs.molecule,\n",
    "        qsci_graph.inputs.mo_occ,\n",
    "        qsci_graph.inputs.cas_hsim,\n",
    "    ),\n",
    ")\n",
    "\n",
    "adapt_circuit = qsci_graph.task(\n",
    "    state_prep(\n",
    "        ham_init,\n",
    "        qsci_graph.inputs.reference_state,\n",
    "        qsci_graph.inputs.max_iteration_prep,\n",
    "        qsci_graph.inputs.atol,\n",
    "        qsci_graph.inputs.mo_occ,\n",
    "        qsci_graph.inputs.cas_init,\n",
    "        qsci_graph.inputs.cas_hsim,\n",
    "    ),\n",
    ")\n",
    "circuits = qsci_graph.task(\n",
    "    circuits_from_hamiltonians(\n",
    "        ham_init,\n",
    "        ham_hsim,\n",
    "        adapt_circuit,\n",
    "        qsci_graph.inputs.t_step_list,\n",
    "        qsci_graph.inputs.cas_init,\n",
    "        qsci_graph.inputs.cas_hsim,\n",
    "        qsci_graph.inputs.mo_occ,\n",
    "        qsci_graph.inputs.max_cx_gates_hsim,\n",
    "    ),\n",
    ")\n",
    "backend_results = qsci_graph.map(_compile_and_run(), circuits)\n",
    "energy = qsci_graph.task(\n",
    "    energy_from_results(\n",
    "        ham_hsim,\n",
    "        backend_results,\n",
    "        qsci_graph.inputs.mo_occ,\n",
    "        qsci_graph.inputs.cas_init,\n",
    "        qsci_graph.inputs.cas_hsim,\n",
    "    ),\n",
    ")\n",
    "\n",
    "qsci_graph.outputs(QSCIOutputs(energy))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd0423ee",
   "metadata": {},
   "source": [
    "Finally we can run the graph as before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d309a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from pathlib import Path\n",
    "from uuid import UUID\n",
    "\n",
    "from tierkreis.consts import PACKAGE_PATH\n",
    "from tierkreis.controller import run_graph\n",
    "from tierkreis.controller.executor.multiple import MultipleExecutor\n",
    "from tierkreis.controller.executor.uv_executor import UvExecutor\n",
    "from tierkreis.controller.storage.filestorage import ControllerFileStorage\n",
    "from tierkreis.storage import read_outputs\n",
    "\n",
    "workflow_id = UUID(int=104)\n",
    "storage = ControllerFileStorage(workflow_id, name=\"qsci\", do_cleanup=True)\n",
    "\n",
    "registry_path = Path().parent / \"example_workers\"\n",
    "custom_executor = UvExecutor(registry_path=registry_path, logs_path=storage.logs_path)\n",
    "common_registry_path = PACKAGE_PATH.parent / \"tierkreis_workers\"\n",
    "common_executor = UvExecutor(\n",
    "    registry_path=common_registry_path,\n",
    "    logs_path=storage.logs_path,\n",
    ")\n",
    "multi_executor = MultipleExecutor(\n",
    "    common_executor,\n",
    "    executors={\"custom\": custom_executor},\n",
    "    assignments={\"qsci_worker\": \"custom\"},\n",
    ")\n",
    "\n",
    "run_graph(\n",
    "    storage,\n",
    "    multi_executor,\n",
    "    qsci_graph,\n",
    "    {\n",
    "        k: json.dumps(v).encode()\n",
    "        for k, v in (\n",
    "            {\n",
    "                \"molecule\": {\n",
    "                    \"geometry\": [\n",
    "                        [\"H\", [0, 0, 0]],\n",
    "                        [\"H\", [0, 0, 1.0]],\n",
    "                        [\"H\", [0, 0, 2.0]],\n",
    "                        [\"H\", [0, 0, 3.0]],\n",
    "                    ],\n",
    "                    \"basis\": \"sto-3g\",\n",
    "                    \"charge\": 0,\n",
    "                },\n",
    "                \"reference_state\": [1, 1, 0, 0],\n",
    "                \"cas_init\": {\n",
    "                    \"n\": 2,\n",
    "                    \"n_ele\": 2,\n",
    "                },\n",
    "                \"cas_hsim\": {\n",
    "                    \"n\": 4,\n",
    "                    \"n_ele\": 4,\n",
    "                },\n",
    "                \"max_cx_gates_hsim\": 2000,\n",
    "                \"mo_occ\": [2, 2, 0, 0],\n",
    "                \"t_step_list\": [0.5, 1.0, 1.5],\n",
    "                \"max_iteration_prep\": 5,\n",
    "                \"atol\": 0.03,\n",
    "            }\n",
    "        ).items()\n",
    "    },\n",
    "    polling_interval_seconds=0.01,\n",
    ")\n",
    "output = read_outputs(qsci_graph, storage)"
   ]
  }
 ],
 "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
}