{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "64a64d79",
      "metadata": {},
      "source": [
        "# Sample-Based Krylov Quantum Diagonalization (SKQD)\n",
        "\n",
        "**Sample-Based Krylov Quantum Diagonalization (SKQD)** is a hybrid quantum-classical algorithm that combines the theoretical guarantees of Krylov Quantum Diagonalization (KQD) with the practical efficiency of sample-based methods. Instead of expensive quantum measurements to compute Hamiltonian matrix elements, SKQD samples from quantum states to construct a computational subspace, then diagonalizes the Hamiltonian within that subspace classically.\n",
        "\n",
        "## Why SKQD?\n",
        "\n",
        "Traditional quantum algorithms like VQE face several fundamental challenges:\n",
        "\n",
        "- **Optimization complexity**: VQE requires optimizing many variational parameters in a high-dimensional, non-convex landscape\n",
        "- **Measurement overhead**: Computing expectation values $\\langle\\psi(\\theta)|H|\\psi(\\theta)\\rangle$ requires many measurements for each Pauli term\n",
        "- **Barren plateaus**: Optimization landscapes can become exponentially flat, making gradient-based optimization ineffective\n",
        "- **Parameter initialization**: Poor initial parameters can lead to local minima far from the global optimum\n",
        "\n",
        "SKQD addresses these fundamental limitations:\n",
        "\n",
        "- **No optimization required**: Uses deterministic quantum time evolution instead of variational circuits\n",
        "- **Provable convergence**: Theoretical guarantees based on the Rayleigh-Ritz variational principle\n",
        "- **Measurement efficient**: Only requires computational basis measurements (Z-basis), the most natural measurement on quantum hardware  \n",
        "- **Noise resilient**: Can filter out problematic measurement outcomes and handle finite sampling\n",
        "- **Systematic improvement**: Increasing Krylov dimension monotonically improves ground state estimates\n",
        "- **Hardware friendly**: Time evolution circuits are more amenable to near-term quantum devices than deep variational ansätze\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "79bc6748",
      "metadata": {},
      "source": [
        "## Understanding Krylov Subspaces\n",
        "\n",
        "### What is a Krylov Subspace?\n",
        "\n",
        "A **Krylov subspace** $\\mathcal{K}^r$ of dimension $r$ is the space spanned by vectors obtained by repeatedly applying an operator $A$ to a reference vector $|\\psi\\rangle$:\n",
        "\n",
        "$$\\mathcal{K}^r = \\text{span}\\{|\\psi\\rangle, A|\\psi\\rangle, A^2|\\psi\\rangle, \\ldots, A^{r-1}|\\psi\\rangle\\}$$\n",
        "\n",
        "\n",
        "\n",
        "### The SKQD Algorithm\n",
        "\n",
        "The key insight of SKQD is that we can:\n",
        "1. **Generate** Krylov states $U^k|\\psi\\rangle$ using quantum time evolution\n",
        "2. **Sample** from these states to get computational basis measurements\n",
        "3. **Combine** all samples to form a computational subspace\n",
        "4. **Diagonalize** the Hamiltonian within this subspace classically\n",
        "\n",
        "This approach is much more efficient than computing matrix elements via quantum measurements!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "568f214d",
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.12/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using: CUDA-Q Version 0.14.0 (https://github.com/NVIDIA/cuda-quantum d84568366c3f9e9a33b3829ff43fb794b3a703ab)\n"
          ]
        }
      ],
      "source": [
        "import cudaq\n",
        "import numpy as np\n",
        "import cupy as cp\n",
        "import pickle\n",
        "\n",
        "from skqd_src.pre_and_postprocessing import *\n",
        "\n",
        "from scipy.sparse import csr_matrix as sp_csr_matrix\n",
        "from scipy.sparse.linalg import eigsh as sp_eigsh\n",
        "\n",
        "from cupyx.scipy.sparse import csr_matrix as cp_csr_matrix\n",
        "from cupyx.scipy.sparse.linalg import eigsh as cp_eigsh\n",
        "\n",
        "if cudaq.num_available_gpus() > 0 and cudaq.has_target(\"nvidia\"):\n",
        "    cudaq.set_target(\"nvidia\")\n",
        "else:\n",
        "    cudaq.set_target(\"qpp-cpu\")\n",
        "\n",
        "cudaq.set_random_seed(42)\n",
        "np.random.seed(43)\n",
        "cp.random.seed(44)\n",
        "\n",
        "print(\"Using:\", cudaq.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "28cb9f2d",
      "metadata": {},
      "source": [
        "## Problem Setup: 22-Qubit Heisenberg Model\n",
        "\n",
        "We demonstrate SKQD on a 1D Heisenberg spin chain with 22 qubits:\n",
        "\n",
        "$$H = \\sum_{i} \\left(J_x \\sigma_i^x \\sigma_{i+1}^x + J_y \\sigma_i^y \\sigma_{i+1}^y + J_z \\sigma_i^z \\sigma_{i+1}^z\\right)$$\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "85c715f1",
      "metadata": {},
      "outputs": [],
      "source": [
        "num_spins = 22 \n",
        "\n",
        "if num_spins >= 63:\n",
        "    raise ValueError(f\"Vectorized implementation of postprocessing supports max 62 qubits due to int64 packing. Requested: {num_spins}\")\n",
        "\n",
        "shots = 100_000  \n",
        "\n",
        "total_time_evolution = np.pi  \n",
        "num_trotter_steps = 8  \n",
        "dt = total_time_evolution / num_trotter_steps  \n",
        "\n",
        "max_k = 12  # largest k for U^k\n",
        "\n",
        "Jx, Jy, Jz = 1.0, 1.0, 1.0 \n",
        "H = create_heisenberg_hamiltonian(num_spins, Jx, Jy, Jz)\n",
        "exact_ground_state_energy = -38.272304 # Computed via exact diagonalization\n",
        "hamiltonian_coefficients, pauli_words = extract_coeffs_and_paulis(H)\n",
        "hamiltonian_coefficients_numpy = np.array(hamiltonian_coefficients)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9cd4c837",
      "metadata": {},
      "source": [
        "## Krylov State Generation via Repeated Evolution\n",
        "\n",
        "For SKQD, we generate the Krylov sequence:\n",
        "$$|\\psi_0\\rangle,  U|\\psi_0\\rangle,  U^2|\\psi_0\\rangle, \\ldots,  U^{k-1}|\\psi_0\\rangle$$\n",
        "\n",
        "where $U = e^{-iHT}$ is approximated via Trotter decomposition.\n",
        "\n",
        "**Implementation Strategy**:\n",
        "1. Start with reference state $|\\psi_0\\rangle$\n",
        "2. Apply Trotter-decomposed time evolution $k$ times for $U^k|\\psi_0\\rangle$\n",
        "3. Measure each Krylov state in computational basis\n",
        "4. Accumulate measurement statistics across all Krylov powers\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "44e3422c",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2025-09-08T14:32:51.773472Z",
          "iopub.status.busy": "2025-09-08T14:32:51.773267Z",
          "iopub.status.idle": "2025-09-08T14:32:51.783809Z",
          "shell.execute_reply": "2025-09-08T14:32:51.783445Z"
        }
      },
      "outputs": [],
      "source": [
        "@cudaq.kernel\n",
        "def quantum_krylov_evolution_circuit(\n",
        "        num_qubits: int, \n",
        "        krylov_power: int, \n",
        "        trotter_steps: int,\n",
        "        dt: float,\n",
        "        H_pauli_words: list[cudaq.pauli_word],\n",
        "        H_coeffs: list[float]):\n",
        "    \n",
        "    \"\"\"\n",
        "    Generate Krylov states via repeated time evolution.\n",
        "    \n",
        "    Args:\n",
        "        num_qubits: Number of qubits in the system\n",
        "        krylov_power: Power k for computing U^k|ψ⟩\n",
        "        trotter_steps: Number of Trotter steps for each U application\n",
        "        dt: Time step for Trotter decomposition\n",
        "        H_pauli_words: Pauli decomposition of Hamiltonian\n",
        "        H_coeffs: Coefficients for each Pauli term\n",
        "    \n",
        "    Returns:\n",
        "        Measurement results in computational basis\n",
        "    \"\"\"\n",
        "    \n",
        "    qubits = cudaq.qvector(num_qubits)\n",
        "\n",
        "    #Prepare Néel state as reference |010101...⟩\n",
        "    for qubit_index in range(num_qubits):\n",
        "        if qubit_index % 2 == 0:\n",
        "            x(qubits[qubit_index])\n",
        "\n",
        "\n",
        "    for _ in range(krylov_power):                                            #applies U^k where U = exp(-iHT)\n",
        "        \n",
        "        for _ in range(trotter_steps):                                       #applies exp(-iHT)\n",
        "            \n",
        "            for i in range(len(H_coeffs)):                                   #applies exp(-iHdt)\n",
        "                exp_pauli( -1 * H_coeffs[i] * dt, qubits, H_pauli_words[i])  #applies exp(-ihdt)\n",
        "\n",
        "\n",
        "    mz(qubits)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "981f54b4",
      "metadata": {},
      "source": [
        "## Quantum Measurements and Sampling\n",
        "\n",
        "\n",
        "### The Sampling Process\n",
        "\n",
        "For each Krylov power $k = 0, 1, 2, \\ldots, K-1$:\n",
        "1. **Prepare** the state $U^k|\\psi_0\\rangle$ using our quantum circuit\n",
        "2. **Measure** in the computational basis many times\n",
        "3. **Collect** the resulting bitstring counts\n",
        "\n",
        "The key insight: these measurement outcomes give us a statistical representation of each Krylov state, which we can then use to construct our computational subspace classically.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "ccca2170",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generating Krylov state U^0...\n",
            "Generating Krylov state U^1...\n",
            "Generating Krylov state U^2...\n",
            "Generating Krylov state U^3...\n",
            "Generating Krylov state U^4...\n",
            "Generating Krylov state U^5...\n",
            "Generating Krylov state U^6...\n",
            "Generating Krylov state U^7...\n",
            "Generating Krylov state U^8...\n",
            "Generating Krylov state U^9...\n",
            "Generating Krylov state U^10...\n",
            "Generating Krylov state U^11...\n"
          ]
        }
      ],
      "source": [
        "all_measurement_results = []\n",
        "\n",
        "for krylov_power in range(max_k):\n",
        "    print(f\"Generating Krylov state U^{krylov_power}...\")\n",
        "    \n",
        "    sampling_result = cudaq.sample(\n",
        "        quantum_krylov_evolution_circuit,\n",
        "        num_spins,  \n",
        "        krylov_power, \n",
        "        num_trotter_steps,  \n",
        "        dt,  \n",
        "        pauli_words,  \n",
        "        hamiltonian_coefficients, \n",
        "        shots_count=shots) \n",
        "        \n",
        "    all_measurement_results.append(dict(sampling_result.items()))\n",
        "    \n",
        "cumulative_results = calculate_cumulative_results(all_measurement_results)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "36a00631",
      "metadata": {},
      "source": [
        "## Classical Post-Processing and Diagonalization\n",
        "\n",
        "Now comes the classical part of SKQD: we use our quantum measurement data to construct and diagonalize the Hamiltonian within each Krylov subspace.\n",
        "\n",
        "1. **Extract basis states** from measurement counts\n",
        "2. **Project Hamiltonian** onto the computational subspace spanned by these states  \n",
        "3. **Diagonalize** the projected Hamiltonian classically\n",
        "4. **Extract ground state energy** estimate\n",
        "\n",
        "### Matrix Construction Details\n",
        "\n",
        "The core of SKQD is constructing the effective Hamiltonian matrix within the computational subspace:\n",
        "\n",
        "1. **Computational Subspace Formation**: From quantum measurements, we obtain a set of computational basis states $\\{|s_1\\rangle, |s_2\\rangle, \\ldots, |s_d\\rangle\\}$ that spans our approximation to the Krylov subspace.\n",
        "\n",
        "2. **Matrix Element Computation**: For each Pauli term $P_k$ in the Hamiltonian with coefficient $h_k$:\n",
        "   $$H = \\sum_k h_k P_k$$\n",
        "   \n",
        "   We compute matrix elements $\\langle s_i | P_k | s_j \\rangle$ by applying the Pauli string $P_k$ to each basis state $|s_j\\rangle$.\n",
        "\n",
        "3. **Effective Hamiltonian**: The projected Hamiltonian becomes:\n",
        "   $$H_{\\text{eff}}[i,j] = \\sum_k h_k \\langle s_i | P_k | s_j \\rangle$$\n",
        "\n",
        "We demonstrate two eigenvalue solver approaches below: an explicit CSR sparse matrix and a matrix-free Lanczos iteration.\n",
        "\n",
        "### Approach 1: GPU-Vectorized CSR Sparse Matrix\n",
        "\n",
        "We build the projected Hamiltonian as a CuPy CSR sparse matrix on the GPU and solve with `cupyx.scipy.sparse.linalg.eigsh`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "6bed2897",
      "metadata": {},
      "outputs": [],
      "source": [
        "csr_energies = []\n",
        "\n",
        "for k in range(1, max_k):\n",
        "\n",
        "    cumulative_subspace_results = cumulative_results[k]\n",
        "    basis_states = get_basis_states_as_array(cumulative_subspace_results, num_spins)\n",
        "    subspace_dimension = len(cumulative_subspace_results)\n",
        "    assert len(cumulative_subspace_results) == basis_states.shape[0]\n",
        "    \n",
        "    # #slower non-vectorized  CPU implementation with CSR matrix input \n",
        "    # matrix_rows, matrix_cols, matrix_elements = projected_hamiltonian_cpu(basis_states, pauli_words, hamiltonian_coefficients_numpy)\n",
        "    # projected_hamiltonian = sp_csr_matrix((matrix_elements, (matrix_rows, matrix_cols)), shape=(subspace_dimension, subspace_dimension))\n",
        "    # eigenvalue = sp_eigsh(projected_hamiltonian, return_eigenvectors=False, **eigenvalue_solver_options)\n",
        "\n",
        "    v0_np = np.random.random((subspace_dimension,)).astype(np.float64)\n",
        "    v0_np = v0_np / np.linalg.norm(v0_np)\n",
        "    v0 = cp.asarray(v0_np, dtype=cp.complex128)\n",
        "\n",
        "    vectorized_pj_on_gpu = True\n",
        "    matrix_rows, matrix_cols, matrix_elements = vectorized_projected_hamiltonian(\n",
        "        basis_states, pauli_words, hamiltonian_coefficients_numpy, vectorized_pj_on_gpu)\n",
        "    \n",
        "    projected_hamiltonian = cp_csr_matrix(\n",
        "        (matrix_elements, (matrix_rows, matrix_cols)),\n",
        "        shape=(subspace_dimension, subspace_dimension))\n",
        "    \n",
        "    eigenvalue = cp_eigsh(projected_hamiltonian, return_eigenvectors=False,\n",
        "                          k=1, which='SA', v0=v0)\n",
        "\n",
        "    csr_energies.append(np.min(eigenvalue).item())\n",
        " "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fcd96a42",
      "metadata": {},
      "source": [
        "### Approach 2: Matrix-Free Lanczos via `distributed_eigsh`\n",
        "\n",
        "Instead of explicitly constructing the projected Hamiltonian as a sparse matrix, we can apply the Hamiltonian as a linear operator during Lanczos iteration. This matrix-free approach avoids storing the full sparse matrix and naturally extends to multiple GPUs (demonstrated in the multi-GPU section below)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "695b53e1",
      "metadata": {},
      "outputs": [],
      "source": [
        "lo_energies = []\n",
        "\n",
        "for k in range(1, max_k):\n",
        "\n",
        "    cumulative_subspace_results = cumulative_results[k]\n",
        "    basis_states = get_basis_states_as_array(cumulative_subspace_results, num_spins)\n",
        "    subspace_dimension = len(cumulative_subspace_results)\n",
        "\n",
        "    v0_np = np.random.random((subspace_dimension,)).astype(np.float64)\n",
        "    v0_np = v0_np / np.linalg.norm(v0_np)\n",
        "    v0 = cp.asarray(v0_np, dtype=cp.complex128)\n",
        "\n",
        "    ham_data = prepare_hamiltonian_data(basis_states, pauli_words, hamiltonian_coefficients_numpy)\n",
        "    cp.get_default_memory_pool().free_all_blocks()\n",
        "    ncv = min(20, subspace_dimension - 1)\n",
        "\n",
        "    eigenvalues, = distributed_eigsh(\n",
        "        ham_data, rank=0, size=1, nccl_comm_obj=None,\n",
        "        k=1, which='SA', ncv=ncv, v0_full=v0,\n",
        "        tile_size=32, return_eigenvectors=False)\n",
        "\n",
        "    eigenvalue = eigenvalues[0]\n",
        "    eigenvalue_real = float(cp.asnumpy(eigenvalue).real)\n",
        "    lo_energies.append(eigenvalue_real)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f36b27e7",
      "metadata": {},
      "source": [
        "## Results Analysis and Convergence\n",
        "\n",
        "Let's visualize our results and analyze how SKQD converges to the true ground state energy. \n",
        "\n",
        "### What to Expect:\n",
        "\n",
        "- **Monotonic improvement**: Each additional Krylov dimension should give a better (lower) energy estimate\n",
        "- **Convergence**: The estimates should approach the exact ground state energy  \n",
        "- **Diminishing returns**: Later Krylov dimensions provide smaller improvements\n",
        "\n",
        "The exact ground state energy for our selected Hamiltonian was computed earlier via classical exact diagonalization and will be used as the reference for comparison with SKQD results.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "b82f10f2",
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 59776 characters]",
            "text/plain": [
              "<Figure size 500x400 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "plot_skqd_convergence(csr_energies, lo_energies, range(1, max_k), exact_ground_state_energy)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "042aa6aa",
      "metadata": {},
      "source": [
        "## Postprocessing Acceleration: CSR matrix approach, single GPU vs CPU \n",
        "\n",
        "The data in the plot below was gathered by averaging over 5 runs on a NVIDIA H100 GPU and an Intel Xeon Platinum 8480CL CPU with 224 threads. The only thing that changed between the 2 data points was `vectorized_pj_on_gpu = True` and `vectorized_pj_on_gpu = False`\n",
        "\n",
        "This substantial acceleration comes from:\n",
        "- **Parallel Pauli string evaluation** across thousands of basis states simultaneously\n",
        "- **Vectorized matrix element computation** leveraging CUDA cores\n",
        "- **GPU-optimized sparse linear algebra** via cuSPARSE library for eigenvalue solving\n",
        "\n",
        "The speedup becomes increasingly critical as the Krylov dimension grows, since the computational subspace dimension (and thus the matrix size) scales exponentially with k. For higher k values, GPU acceleration transforms previously intractable postprocessing into feasible computation times.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "230e700d",
      "metadata": {},
      "source": [
        "<img src=\"images/speedup.png\" width=\"500\">\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "14f8f02f",
      "metadata": {},
      "source": [
        "## Postprocessing Scale-Up and Scale-Out: Linear Operator Approach, Multi-GPU Multi-Node\n",
        "\n",
        "The plots below were generated using `distributed_eigsh()` with the Hamiltonian of an active site of a GPCR (G protein-coupled receptor) molecule. This demonstrates scaling behavior on a problem with a much larger computational subspace than the 22-qubit demo above.\n",
        "\n",
        "<div style=\"display: flex; gap: 10px;\">\n",
        "  <img src=\"images/strong_scaling_timing.jpeg\" width=\"450\">\n",
        "  <img src=\"images/efficiency.jpeg\" width=\"450\">\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "916ebaeb",
      "metadata": {},
      "source": [
        "### Saving Hamiltonian Data\n",
        "\n",
        "To run the multi-GPU solver, we first save the Hamiltonian data (pauli words, coefficients) and the sampled basis states so they can be loaded by each MPI rank."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "987a33ce",
      "metadata": {},
      "outputs": [],
      "source": [
        "# # uncomment this to save the data\n",
        "\n",
        "# cumulative_subspace_results = cumulative_results[max_k-1]\n",
        "# basis_states = get_basis_states_as_array(cumulative_subspace_results, num_spins)\n",
        "\n",
        "# with open('skqd_data.pkl', 'wb') as f:\n",
        "#     pickle.dump({\n",
        "#         'basis_states': basis_states,\n",
        "#         'pauli_words': pauli_words,\n",
        "#         'hamiltonian_coefficients': hamiltonian_coefficients_numpy,\n",
        "#     }, f)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "157113ef",
      "metadata": {},
      "source": [
        "### Running the Distributed Solver\n",
        "\n",
        "Save the code below as `distributed_postprocessing.py` and launch it with MPI. For example, on a system with 8 GPUs:\n",
        "\n",
        "```bash\n",
        "mpirun -np 8 --bind-to none --allow-run-as-root python distributed_postprocessing.py\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "bf58d019",
      "metadata": {},
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "Multi-GPU SKQD Eigenvalue Solver\n",
        "\n",
        "Usage:\n",
        "    Save the code below in a distributed_postprocessing.py file\n",
        "    Execute with the command below, for example on a system with 8 GPUs: \n",
        "    mpirun -np 8 --bind-to none --allow-run-as-root python distributed_postprocessing.py \n",
        "\"\"\"\n",
        "\n",
        "import numpy as np\n",
        "import cupy as cp\n",
        "from skqd_src.pre_and_postprocessing import *\n",
        "from mpi4py import MPI\n",
        "import pickle\n",
        "comm = MPI.COMM_WORLD\n",
        "rank = comm.Get_rank()\n",
        "size = comm.Get_size()\n",
        "\n",
        "n_devices = cp.cuda.runtime.getDeviceCount()\n",
        "my_gpu = rank % n_devices\n",
        "cp.cuda.Device(my_gpu).use()\n",
        "\n",
        "tile_size = 32\n",
        "\n",
        "#uncomment this to load the data\n",
        "# with open('skqd_data.pkl', 'rb') as f:\n",
        "#     data = pickle.load(f)\n",
        "#     basis_states = data['basis_states']\n",
        "#     pauli_words = data['pauli_words']\n",
        "#     hamiltonian_coefficients_numpy = data['hamiltonian_coefficients']\n",
        "\n",
        "\n",
        "\n",
        "subspace_dimension = basis_states.shape[0]\n",
        "\n",
        "ham_data = prepare_hamiltonian_data(basis_states, pauli_words, hamiltonian_coefficients_numpy)\n",
        "cp.cuda.Device(my_gpu).synchronize()\n",
        "\n",
        "nccl_comm_obj = None\n",
        "if size > 1:\n",
        "    from cupy.cuda import nccl\n",
        "    if rank == 0:\n",
        "        nccl_id = nccl.get_unique_id()\n",
        "    else:\n",
        "        nccl_id = None\n",
        "    nccl_id = comm.bcast(nccl_id, root=0)\n",
        "    nccl_comm_obj = nccl.NcclCommunicator(size, nccl_id, rank)\n",
        "\n",
        "np.random.seed(42)\n",
        "v0_np = np.random.random((subspace_dimension,)).astype(np.float64)\n",
        "v0_np = v0_np / np.linalg.norm(v0_np)\n",
        "v0 = cp.asarray(v0_np, dtype=cp.complex128)\n",
        "\n",
        "cp.get_default_memory_pool().free_all_blocks()\n",
        "\n",
        "ncv = min(20, subspace_dimension - 1)\n",
        "\n",
        "eigenvalues, = distributed_eigsh(\n",
        "    ham_data, rank, size, nccl_comm_obj,\n",
        "    k=1, which='SA', ncv=ncv, v0_full=v0,\n",
        "    tile_size=tile_size, return_eigenvectors=False)\n",
        "\n",
        "eigenvalue = eigenvalues[0]\n",
        "cp.cuda.Device(my_gpu).synchronize()\n",
        "eigenvalue_real = float(cp.asnumpy(eigenvalue).real)\n",
        "\n",
        "MPI.Finalize()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b01835f6",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "In this notebook we demonstrated the **Sample-Based Krylov Quantum Diagonalization (SKQD)** algorithm using CUDA-Q:\n",
        "\n",
        "1. **Quantum phase**: Generated a sequence of Krylov states $U^k|\\psi_0\\rangle$ via Trotterized time evolution on a 22-qubit Heisenberg model, and sampled each state in the computational basis.\n",
        "\n",
        "2. **Classical post-processing**: Projected the Hamiltonian onto the computational subspace defined by the measurement outcomes, and extracted the ground state energy via eigenvalue solvers.\n",
        "\n",
        "3. **Two solver approaches**: Compared an explicit GPU-accelerated CSR sparse matrix solver with a matrix-free Lanczos iteration (`distributed_eigsh`), showing both converge to the same result.\n",
        "\n",
        "4. **GPU acceleration**: Demonstrated significant speedup of the classical post-processing when using GPU-vectorized operations compared to CPU.\n",
        "\n",
        "5. **Multi-GPU scaling**: Showed how the matrix-free Lanczos solver scales across multiple GPUs using NCCL for distributed computation.\n"
      ]
    }
  ],
  "metadata": {
    "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.12.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
