{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d53b3dba",
   "metadata": {},
   "source": [
    "# Sample-based quantum diagonalization with Fulqrum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0237276",
   "metadata": {},
   "source": [
    "In this tutorial, we will use Fulqrum to find an approximation to the ground state of a chemistry Hamiltonian following a sampled-based quantum diagonalization (SQD) approach.\n",
    "\n",
    "We will use different blocks from Fulqrum to construct a full end-to-end workflow to estimate molecular Hamiltonian ground state energy using configuration recovery.\n",
    "\n",
    "- Convert a molecular Hamiltonian from one- and two-body integrals to a Fulqrum compatible operator.\n",
    "- Construct a subspace from bitstrings (configurations or Slater Determinants) and project the operator on to the subsapce as a sparse matrix in the CSR format.\n",
    "- Eigensolve the sparse matrix with a choice of an eigensolver.\n",
    "- Compute electron orbital occupancies from the eigenvector and subspace bitstrings and use them to recover configurations.\n",
    "- Use recovered configurations along with _carryover bitstrings_ in the next iteration to refine ground state estimation.\n",
    "\n",
    "This tutorial requires following two packages in addition to Fulqrum:\n",
    "- `PySCF`: To define the molecule and compute the one- and two-body integrals.\n",
    "- `Openfermion`: To convert the one- and two-body integrals into a Fulqrum compatible operator internally. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec3f6fa9",
   "metadata": {},
   "source": [
    "### Define Molecule"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b92fd13",
   "metadata": {},
   "source": [
    "In this tutorial, we will approximate the ground state energy of an $N_2$ molecule in the `6-31g` basis. First, we will specify the molecule and its properties. Then, we will compute one- and two-body integrals using the `PySCF` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "25856b6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "converged SCF energy = -108.835236570774\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "\n",
    "import pyscf\n",
    "import pyscf.mcscf\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# Build N2 molecule\n",
    "mol = pyscf.gto.Mole()\n",
    "mol.build(\n",
    "    atom=[[\"N\", (0, 0, 0)], [\"N\", (1.0, 0, 0)]],\n",
    "    basis=\"6-31g\",\n",
    "    symmetry=\"Dooh\",\n",
    ")\n",
    "\n",
    "# Define active space\n",
    "n_frozen = 2\n",
    "active_space = range(n_frozen, mol.nao_nr())\n",
    "\n",
    "# Get molecular integrals\n",
    "scf = pyscf.scf.RHF(mol).run()\n",
    "num_orbitals = len(active_space)\n",
    "n_electrons = int(sum(scf.mo_occ[active_space]))\n",
    "num_elec_a = (n_electrons + mol.spin) // 2\n",
    "num_elec_b = (n_electrons - mol.spin) // 2\n",
    "cas = pyscf.mcscf.CASCI(scf, num_orbitals, (num_elec_a, num_elec_b))\n",
    "mo = cas.sort_mo(active_space, base=0)\n",
    "hcore, nuclear_repulsion_energy = cas.get_h1cas(mo)  # hcore: one-body integrals\n",
    "eri = pyscf.ao2mo.restore(1, cas.get_h2cas(mo), num_orbitals)  # eri: two-body integrals"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a82f7ab",
   "metadata": {},
   "source": [
    "### Convert to the Fulqrum operator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fde3d3b0",
   "metadata": {},
   "source": [
    "Next, we will convert the one- and two-body integrals into a Fulqrum `FermionicOperator` using the `integrals_to_fq_fermionic_op()` converter available from Fulqrum. Next, we will perform an _extended_ Jordan-Wigner transformation using `extended_jw_transform()` to convert it to a Fulqrum `QubitOperator`. The `QubitOperator` will be used in the subsequent projection step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d3b1a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num terms operator: 36960\n",
      "Num terms operator [after JW transform]: 13280\n"
     ]
    }
   ],
   "source": [
    "from fulqrum.convert.integrals import integrals_to_fq_fermionic_op\n",
    "\n",
    "fermionic_op = integrals_to_fq_fermionic_op(\n",
    "    one_body_integrals=hcore, two_body_integrals=eri\n",
    ")\n",
    "\n",
    "print(f\"Num terms operator: {fermionic_op.size()}\")\n",
    "\n",
    "fulqrum_operator = fermionic_op.extended_jw_transformation()\n",
    "print(f\"Num terms operator [after JW transform]: {fulqrum_operator.size()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6cfafd1",
   "metadata": {},
   "source": [
    "### Load random bitstrings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f9fbc4a",
   "metadata": {},
   "source": [
    "Now, we load some pre-generated random bitstrings for the subspace construction and projection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7a99546a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num bitstrings in raw counts: 10000\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "with open(\"./data/random_counts_n2_10000_samples.json\", \"r\") as jf:\n",
    "    counts = json.load(jf)\n",
    "print(f\"Num bitstrings in raw counts: {len(counts)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7ac04d9",
   "metadata": {},
   "source": [
    "Next, we will separate bitstrings and their respective probabilities into two separate lists so that it becomes easier to use them later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dcb8b317",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "raw_bitstrings = []\n",
    "raw_probs = []\n",
    "for bs, cnt in counts.items():\n",
    "    raw_bitstrings.append(bs)\n",
    "    raw_probs.append(cnt)\n",
    "raw_probs = np.array(raw_probs, dtype=float)\n",
    "raw_probs /= np.linalg.norm(raw_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4950f33",
   "metadata": {},
   "source": [
    "### Configuration recovery loop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15a8b683",
   "metadata": {},
   "source": [
    "Now, we are ready to run the full sample-based quantum diagonalization (SQD) routine with iterative configuration recovery to refine our estimation of the ground state energy of $N_2$ molecule. Fulqrum has an `sqd` module that uses the [qiskit-addon-sqd-hpc](https://qiskit.github.io/qiskit-addon-sqd-hpc/) as a submodule, which provides postselection, subsampling, and configuration recovery capabilities. Each SQD loop has the following steps:\n",
    "\n",
    "1. Construct alpha (beta) half strings by combining Hartree-Fock state, carryover configurations from previous round, and recovered configurations from raw counts - in that order.\n",
    "2. Construct a Fulqrum `Subspace` using the spin-up (alpha) and spin-down (beta) half strings. For this $N_2$ example, both will be same. Fulqrum `Subspace` internally sorts the half strings and takes a Cartesian product to construct full strings.\n",
    "3. Project the `fulqrum_operator` on to this subspace and construct a CSR-format sparse matrix representation of the projected Hamiltonian. The CSR matrix will be wrapped into a [LinearOperator](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.LinearOperator.html) for faster eigensolving.\n",
    "4. Prepare a simple initial guess vector $v0$ for the eigensolver. While it is optional, it can speed up eigensolving.\n",
    "5. Eigensolve the projected Hamiltonian matrix to get the minimum eigenvalue (ground state energy estimate) and the corresponding eigenvector.\n",
    "6. Use the eigenvector to compute electron orbital occupancies that will be used to recover configurations in the next round and also to select important configurations to carryover to the next round."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "94f4060b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "\n",
    "\n",
    "def unique_alpha_beta_combined(bitstrings):\n",
    "    \"\"\"A utility function for getting unique alpha and beta halves from full bitstrings.\n",
    "\n",
    "    Args:\n",
    "        bitstrings (list[str]): A list of full bitstrings consisting of both alpha and beta parts.\n",
    "\n",
    "    Returns:\n",
    "        A dictionary with combined unique ``alpha`` and ``beta`` half strings.\n",
    "    \"\"\"\n",
    "    if not bitstrings:\n",
    "        return OrderedDict()\n",
    "\n",
    "    unique_ab = OrderedDict()\n",
    "    num_spatial_orb = len(bitstrings[0]) // 2\n",
    "\n",
    "    for bs in bitstrings:\n",
    "        a = bs[num_spatial_orb:]\n",
    "        b = bs[:num_spatial_orb]\n",
    "\n",
    "        unique_ab[a] = 1\n",
    "        unique_ab[b] = 1\n",
    "\n",
    "    return unique_ab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "82a03e2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ITERATION 0\n",
      "  num total half strs:  75\n",
      "  num selected half strs:  75\n",
      "  Half strs construction took: 0.0001 seconds\n",
      "  Subspace construction took: 0.001658 seconds\n",
      "  Subspace dimension: 75 x 75 = 5_625\n",
      "  Operator projection took: 0.008842 seconds\n",
      "  CSR matrix memory: 0.287434 MBs\n",
      "  Initial guess vector v0 construction took: 0.0000 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.0060 seconds\n",
      "  Electronic Energy: [-32.60417101]\n",
      "  Total Energy: [-108.83527355]\n",
      "  num carryover full strs: 5\n",
      "Iter 0 took: 0.0234 seconds\n",
      "\n",
      "ITERATION 1\n",
      "  num total half strs:  3867\n",
      "  num selected half strs:  500\n",
      "  Half strs construction took: 0.0026 seconds\n",
      "  Subspace construction took: 0.099804 seconds\n",
      "  Subspace dimension: 500 x 500 = 250_000\n",
      "  Operator projection took: 0.587633 seconds\n",
      "  CSR matrix memory: 111.487309 MBs\n",
      "  Initial guess vector v0 construction took: 0.0007 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.1056 seconds\n",
      "  Electronic Energy: [-32.70163608]\n",
      "  Total Energy: [-108.93273862]\n",
      "  num carryover full strs: 2221\n",
      "Iter 1 took: 0.9739 seconds\n",
      "\n",
      "ITERATION 2\n",
      "  num total half strs:  3867\n",
      "  num selected half strs:  500\n",
      "  Half strs construction took: 0.0031 seconds\n",
      "  Subspace construction took: 0.092869 seconds\n",
      "  Subspace dimension: 500 x 500 = 250_000\n",
      "  Operator projection took: 0.671684 seconds\n",
      "  CSR matrix memory: 164.557041 MBs\n",
      "  Initial guess vector v0 construction took: 0.0008 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.1171 seconds\n",
      "  Electronic Energy: [-32.76803522]\n",
      "  Total Energy: [-108.99913776]\n",
      "  num carryover full strs: 4517\n",
      "Iter 2 took: 1.0560 seconds\n",
      "\n",
      "ITERATION 3\n",
      "  num total half strs:  3887\n",
      "  num selected half strs:  500\n",
      "  Half strs construction took: 0.0039 seconds\n",
      "  Subspace construction took: 0.096041 seconds\n",
      "  Subspace dimension: 500 x 500 = 250_000\n",
      "  Operator projection took: 0.697432 seconds\n",
      "  CSR matrix memory: 199.318607 MBs\n",
      "  Initial guess vector v0 construction took: 0.0009 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.1243 seconds\n",
      "  Electronic Energy: [-32.7892532]\n",
      "  Total Energy: [-109.02035574]\n",
      "  num carryover full strs: 5693\n",
      "Iter 3 took: 1.1003 seconds\n",
      "\n",
      "ITERATION 4\n",
      "  num total half strs:  3865\n",
      "  num selected half strs:  500\n",
      "  Half strs construction took: 0.0041 seconds\n",
      "  Subspace construction took: 0.103407 seconds\n",
      "  Subspace dimension: 500 x 500 = 250_000\n",
      "  Operator projection took: 0.703007 seconds\n",
      "  CSR matrix memory: 215.561298 MBs\n",
      "  Initial guess vector v0 construction took: 0.0007 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.1297 seconds\n",
      "  Electronic Energy: [-32.79159265]\n",
      "  Total Energy: [-109.02269518]\n",
      "  num carryover full strs: 6342\n",
      "Iter 4 took: 1.1150 seconds\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "import scipy.sparse.linalg as spla\n",
    "\n",
    "import fulqrum as fq\n",
    "from fulqrum.core.sqd import (\n",
    "    postselect_by_hamming_right_and_left,\n",
    "    subsample,\n",
    "    recover_configurations,\n",
    "    get_carryover_full_strs,\n",
    ")\n",
    "\n",
    "# options\n",
    "tol = 1e-5\n",
    "num_batches = 1\n",
    "carryover_threshold = 1e-4\n",
    "seed = 0\n",
    "nelec = (num_elec_a, num_elec_b)\n",
    "samples_per_batch = 500\n",
    "current_occupancies = None\n",
    "max_iterations = 5\n",
    "\n",
    "carryover_full_strs = []\n",
    "S = None  # Subspace\n",
    "\n",
    "Hsub = fq.SubspaceHamiltonian(fulqrum_operator, S)\n",
    "\n",
    "seed_rng = np.random.default_rng(seed=seed)\n",
    "for ni in range(max_iterations):\n",
    "    iter_start = time.perf_counter()\n",
    "    print(f\"ITERATION {ni}\")\n",
    "    if current_occupancies is None:\n",
    "        # If we don't have average orbital occupancy information, simply postselect\n",
    "        # bitstrings with the correct numbers of spin-up and spin-down electrons\n",
    "        bitstrings, probs = postselect_by_hamming_right_and_left(\n",
    "            raw_bitstrings, raw_probs, num_elec_a, num_elec_b\n",
    "        )\n",
    "    else:\n",
    "        # If we do have average orbital occupancy information, use it to refine the\n",
    "        # full set of noisy configurations\n",
    "        bitstrings, probs = recover_configurations(\n",
    "            raw_bitstrings,\n",
    "            raw_probs,\n",
    "            current_occupancies[0],  # occs_a\n",
    "            current_occupancies[1],  # occs_b\n",
    "            num_elec_a,\n",
    "            num_elec_b,\n",
    "            seed,\n",
    "        )\n",
    "\n",
    "    subsamples = [\n",
    "        subsample(\n",
    "            bitstrings,\n",
    "            probs,\n",
    "            min(len(bitstrings), len(bitstrings)),\n",
    "            seed_rng.integers(0, np.iinfo(np.int32).max),\n",
    "        )\n",
    "        for _ in range(num_batches)\n",
    "    ]\n",
    "\n",
    "    for batch in subsamples:\n",
    "        hs_start = time.perf_counter()\n",
    "\n",
    "        # In this tutorial, we include half strings with the following precedence\n",
    "        # Hartree-Fock state (always include) > carryover (from previous round) > recovered samples.\n",
    "        # We add bitstrings in a dictionary as it preserves order and also filters out duplicates.\n",
    "        half_strs_dict = OrderedDict()\n",
    "\n",
    "        # Add hartree-fock state\n",
    "        hartree_fock_state = \"0\" * (num_orbitals - num_elec_a) + \"1\" * num_elec_a\n",
    "        half_strs_dict[hartree_fock_state] = 1\n",
    "\n",
    "        # Add carry-overs\n",
    "        # First, we split carryover full strs into unique alpha and beta halves\n",
    "        unique_ab_co = unique_alpha_beta_combined(carryover_full_strs)\n",
    "        for bs in unique_ab_co.keys():\n",
    "            half_strs_dict[bs] = 1\n",
    "\n",
    "        # Add recovered bitstrings\n",
    "        unique_ab_recovered = unique_alpha_beta_combined(batch)\n",
    "        for bs in unique_ab_recovered.keys():\n",
    "            half_strs_dict[bs] = 1\n",
    "\n",
    "        half_strs = list(half_strs_dict.keys())\n",
    "        print(\"  num total half strs: \", len(half_strs))\n",
    "\n",
    "        # We will use a subset of accumulated half strings in the Subspace construction\n",
    "        # defined by the `samples_per_batch` parameter. For this tutorial, the subspace\n",
    "        # dimension will be samples_per_batch x samples_per_batch\n",
    "        half_strs = half_strs[:samples_per_batch]\n",
    "        print(\"  num selected half strs: \", len(half_strs))\n",
    "        half_strs.sort()  # fq.Subspace() also sorts bitstrings, so can be skipped.\n",
    "\n",
    "        hs_end = time.perf_counter()\n",
    "        print(f\"  Half strs construction took: {hs_end - hs_start:.4f} seconds\")\n",
    "\n",
    "        s_start = time.perf_counter()\n",
    "        S = fq.Subspace([half_strs, half_strs])\n",
    "        s_end = time.perf_counter()\n",
    "        print(f\"  Subspace construction took: {s_end - s_start:4f} seconds\")\n",
    "        print(\n",
    "            f\"  Subspace dimension: {len(half_strs)} x {len(half_strs)} = {len(half_strs) ** 2:_}\"\n",
    "        )\n",
    "\n",
    "        proj_start = time.perf_counter()\n",
    "        Hsub.update_subspace(S)\n",
    "        # Convert the SubspaceHamiltonian into a CSR format sparse matrix,\n",
    "        # which will also be wrapped in a LinearOperator for faster eigensolving.\n",
    "        Hsub_csr_linop = Hsub.to_csr_linearoperator_fast()\n",
    "        proj_end = time.perf_counter()\n",
    "        print(f\"  Operator projection took: {proj_end - proj_start:4f} seconds\")\n",
    "\n",
    "        total_bytes = Hsub_csr_linop.memory_size\n",
    "        total_mega_bytes = total_bytes / (1024 * 1024)\n",
    "        print(f\"  CSR matrix memory: {total_mega_bytes:.6f} MBs\")\n",
    "\n",
    "        # Constructing a simple initial guess vector\n",
    "        v0_start = time.perf_counter()\n",
    "        diag_vec = Hsub.diagonal_vector()\n",
    "        min_idx = np.where(diag_vec == diag_vec.min())[0]\n",
    "        v0 = np.zeros(len(S), dtype=Hsub.dtype)\n",
    "        v0[min_idx] = 1\n",
    "        v0_end = time.perf_counter()\n",
    "        print(\n",
    "            f\"  Initial guess vector v0 construction took: {v0_end - v0_start:.4f} seconds\"\n",
    "        )\n",
    "\n",
    "        print(\"  Starting eigensolving ...\")\n",
    "        start = time.perf_counter()\n",
    "        eigvals, eigvecs = spla.eigsh(\n",
    "            Hsub_csr_linop,  # use `Hsub` for the matrix-free mode. Memory-efficient but slower.\n",
    "            k=1,\n",
    "            which=\"SA\",\n",
    "            tol=tol,\n",
    "            v0=v0,\n",
    "        )\n",
    "        end = time.perf_counter()\n",
    "        print(f\"  Eigensolving took: {end - start:.4f} seconds\")\n",
    "        print(f\"  Electronic Energy: {eigvals}\")\n",
    "        print(f\"  Total Energy: {eigvals + nuclear_repulsion_energy}\")\n",
    "\n",
    "        eigvecs = eigvecs.ravel()\n",
    "\n",
    "        # Subspace class has a method named `get_orbital_occupancies()` to compute\n",
    "        # electron orbital occupancies from the eigenvector and subspace bitstrings.\n",
    "        # The method accepts probabilities, and thus, we must do |amplitude| ^ 2.\n",
    "        # Returns alpha and beta occupanices as a length-2 tuple in the order\n",
    "        # ([a0 ... aN], [b0 ... bN]).\n",
    "        current_occupancies = S.get_orbital_occupancies(\n",
    "            probs=np.abs(eigvecs) ** 2, norb=num_orbitals\n",
    "        )\n",
    "\n",
    "        # Next, we get important bitstrings to include (carryover to) in the next round.\n",
    "        # Important bitstrings are the ones with |amplitude| > threshold.\n",
    "        # It is a must to provide |amplitude|, i.e., np.abs(eigenvector) to the function.\n",
    "        # The function returns a list of length-2 tuples, where the first element of the tuple\n",
    "        # is the full bitstrings, and the second element is its weight (|amplitude|).\n",
    "        # The list is also sorted in the descending order of |amplitude|.\n",
    "        carryover_full_strs_and_weights = get_carryover_full_strs(\n",
    "            S, np.abs(eigvecs), carryover_threshold\n",
    "        )\n",
    "\n",
    "        carryover_full_strs = [item[0] for item in carryover_full_strs_and_weights]\n",
    "\n",
    "        print(f\"  num carryover full strs: {len(carryover_full_strs)}\")\n",
    "\n",
    "    iter_end = time.perf_counter()\n",
    "    print(f\"Iter {ni} took: {iter_end - iter_start:.4f} seconds\\n\")"
   ]
  }
 ],
 "metadata": {
  "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.14.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}