{
 "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": 7,
   "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": 8,
   "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.num_terms}\")\n",
    "\n",
    "fulqrum_operator = fermionic_op.extended_jw_transformation()\n",
    "print(f\"Num terms operator [after JW transform]: {fulqrum_operator.num_terms}\")"
   ]
  },
  {
   "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": 9,
   "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": 10,
   "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": 11,
   "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": 12,
   "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.0005 seconds\n",
      "  Subspace construction took: 0.002349 seconds\n",
      "  Subspace dimension: 75 x 75 = 5_625\n",
      "  Operator projection took: 0.055319 seconds\n",
      "  CSR matrix memory: 0.287434 MBs\n",
      "  Initial guess vector v0 construction took: 0.0002 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.0072 seconds\n",
      "  Electronic Energy: [-32.60417101]\n",
      "  Total Energy: [-108.83527355]\n",
      "  num carryover full strs: 5\n",
      "Iter 0 took: 0.0773 seconds\n",
      "\n",
      "ITERATION 1\n",
      "  num total half strs:  3867\n",
      "  num selected half strs:  500\n",
      "  Half strs construction took: 0.0029 seconds\n",
      "  Subspace construction took: 0.111522 seconds\n",
      "  Subspace dimension: 500 x 500 = 250_000\n",
      "  Operator projection took: 1.050048 seconds\n",
      "  CSR matrix memory: 112.871311 MBs\n",
      "  Initial guess vector v0 construction took: 0.0011 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.1030 seconds\n",
      "  Electronic Energy: [-32.72992823]\n",
      "  Total Energy: [-108.96103077]\n",
      "  num carryover full strs: 2508\n",
      "Iter 1 took: 1.4591 seconds\n",
      "\n",
      "ITERATION 2\n",
      "  num total half strs:  3866\n",
      "  num selected half strs:  500\n",
      "  Half strs construction took: 0.0036 seconds\n",
      "  Subspace construction took: 0.104520 seconds\n",
      "  Subspace dimension: 500 x 500 = 250_000\n",
      "  Operator projection took: 1.153417 seconds\n",
      "  CSR matrix memory: 157.851444 MBs\n",
      "  Initial guess vector v0 construction took: 0.0007 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.1192 seconds\n",
      "  Electronic Energy: [-32.78091965]\n",
      "  Total Energy: [-109.01202219]\n",
      "  num carryover full strs: 4171\n",
      "Iter 2 took: 1.6060 seconds\n",
      "\n",
      "ITERATION 3\n",
      "  num total half strs:  3860\n",
      "  num selected half strs:  500\n",
      "  Half strs construction took: 0.0037 seconds\n",
      "  Subspace construction took: 0.096912 seconds\n",
      "  Subspace dimension: 500 x 500 = 250_000\n",
      "  Operator projection took: 1.597543 seconds\n",
      "  CSR matrix memory: 183.260990 MBs\n",
      "  Initial guess vector v0 construction took: 0.0004 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.1196 seconds\n",
      "  Electronic Energy: [-32.78314636]\n",
      "  Total Energy: [-109.0142489]\n",
      "  num carryover full strs: 5323\n",
      "Iter 3 took: 1.9858 seconds\n",
      "\n",
      "ITERATION 4\n",
      "  num total half strs:  3861\n",
      "  num selected half strs:  500\n",
      "  Half strs construction took: 0.0041 seconds\n",
      "  Subspace construction took: 0.099916 seconds\n",
      "  Subspace dimension: 500 x 500 = 250_000\n",
      "  Operator projection took: 1.330199 seconds\n",
      "  CSR matrix memory: 206.659351 MBs\n",
      "  Initial guess vector v0 construction took: 0.0008 seconds\n",
      "  Starting eigensolving ...\n",
      "  Eigensolving took: 0.2047 seconds\n",
      "  Electronic Energy: [-32.789638]\n",
      "  Total Energy: [-109.02074054]\n",
      "  num carryover full strs: 6196\n",
      "Iter 4 took: 1.8276 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",
    "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(bitstrings, probs, min(len(bitstrings), len(bitstrings)), seed)\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 = fq.SubspaceHamiltonian(fulqrum_operator, 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(verbose=False)\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": "fq-public",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}