{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "72420c62-2716-4f64-ad5b-73394b481cc1",
   "metadata": {},
   "source": [
    "# Understand open-shell vs closed-shell options and its effect in the subspace construction\n",
    "\n",
    "In this \"how-to\", we will show how to choose subpace dimensions in the `qiskit_addon_sqd` package to post-process quantum samples using the [self-consistent configuration recovery technique](https://arxiv.org/abs/2405.05068). \n",
    "\n",
    "More importantly, this \"how-to\" also highlights some differences in the behaviour in the susbapce construction when run in `open_shell = False` or `open_shell = True` modes:\n",
    "\n",
    "- `open_shell = False` only works when the number of spin-up and spin-down electrons is the same. \n",
    "\n",
    "- `open_shell = True` must be used when the number of spin-up and spin-down electrons is different. It can also be used when the number of spin-up and spin-down electrons is the same. However, in this last case, there is a difference in the sizes of the subspaces generated between `open_shell = False` and `open_shell = True`, as discussed in this notebook.\n",
    "\n",
    "**NOTE:** Some of the electronic-configuration (bitstring) manipulations in this package have as a goal to preserve the total spin symmetry $S^2$. Standard Selected Counfiguration Interaction (SCI) solvers cannot impose $S^2$ conservation exactly. Consequently, they do so approximately via a Lagrange multiplier. \n",
    "\n",
    "The choice of electronic configurations entering the eigenstate solver can also have a strong effect in the conservation of spin. For example, in a (2-electron,2-orbital) system, one may sample the configuration $|1001\\rangle$ (having a single spin-up excitation over the RHF state $|0101\\rangle$) which is a linear combination of the open-shell singlet and triplet states, respectively $(|1001\\rangle ± |0110\\rangle) /\\sqrt{2}$. If the configuration |0110⟩ is not sampled, one can construct neither eigenfunction of total spin, leading to spin contamination or redundancy (i.e. the configuration |1001⟩ is involved in a CI calculation, but has coefficient 0 in the CI vector). Consider that a single sample $|1001\\rangle$ is generated in the quantum computer, this is how the `sqd` package handles this situation:\n",
    "\n",
    "- `open_shell = False`: \n",
    "\n",
    "    1. The $1001$ bitstring is split in half, representing spin-up and spin-down configurations: $10$ (up) and $01$ (down).\n",
    "    \n",
    "    2. The list of unique spin-polarized configurations is constructed: $\\mathcal{U} = [01, 10]$.\n",
    "\n",
    "    3. We then consider all possible combinations of $\\mathcal{U}$ elements to form the basis: $\\left \\{ |0101\\rangle, |0110\\rangle , |1001\\rangle , |1010\\rangle \\right \\}$, which contains the singlet and triplet states.\n",
    "\n",
    "\n",
    "- `open_shell = True`: \n",
    "\n",
    "    1. Contrary to the `open_shell = False` case, we do not combine the halves of the bitstring to form the basis."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6755afb-ca1e-4473-974b-ba89acc8abce",
   "metadata": {},
   "source": [
    "## Closed-Shell\n",
    "\n",
    "This example shows how the bitstrings are manipulated in a (2-electron, 4-orbital) system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "677f54ac-b4ed-47e3-b5ba-5366d3a520f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify molecule properties\n",
    "num_orbitals = 4\n",
    "num_elec_a = num_elec_b = 1\n",
    "open_shell = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c58e988c-a109-44cd-a975-9df43250c318",
   "metadata": {},
   "source": [
    "### Specify by hand a dictionary of measurement outcomes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e9506e0b-ed64-48bb-a97a-ef851b604af1",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_dict = {\"00010010\": 1 / 2.0 - 0.01, \"01001000\": 1 / 2.0 - 0.01, \"00010001\": 0.02}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "851bc98e-9c08-4e78-9472-36301abc11d8",
   "metadata": {},
   "source": [
    "### Transform the counts dict into a bitstring matrix and probability array for post-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7a102a7f-aae6-4583-ab82-ae40fcb5496a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[False False False  True False False  True False]\n",
      " [False  True False False  True False False False]\n",
      " [False False False  True False False False  True]]\n",
      "[0.49 0.49 0.02]\n"
     ]
    }
   ],
   "source": [
    "from qiskit_addon_sqd.counts import counts_to_arrays\n",
    "\n",
    "# Convert counts into bitstring and probability arrays\n",
    "bitstring_matrix_full, probs_arr_full = counts_to_arrays(counts_dict)\n",
    "print(bitstring_matrix_full)\n",
    "print(probs_arr_full)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb704101-0fe8-4d12-b572-b1d844e35a90",
   "metadata": {},
   "source": [
    "### Subsample a single batch of size two:\n",
    "\n",
    "- ``n_batches = 1``: Number of batches of configurations used by the different calls to the eigenstate solver\n",
    "- ``samples_per_batch = 2``: Number of unique configurations to include in each batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fe60aee2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[False False False  True False False  True False]\n",
      " [False  True False False  True False False False]]\n"
     ]
    }
   ],
   "source": [
    "from qiskit_addon_sqd.subsampling import postselect_and_subsample\n",
    "\n",
    "n_batches = 1\n",
    "samples_per_batch = 2\n",
    "\n",
    "# seed for random number generator\n",
    "rand_seed = 48\n",
    "\n",
    "# Generate the batches\n",
    "batches = postselect_and_subsample(\n",
    "    bitstring_matrix_full,\n",
    "    probs_arr_full,\n",
    "    hamming_right=num_elec_a,\n",
    "    hamming_left=num_elec_b,\n",
    "    samples_per_batch=samples_per_batch,\n",
    "    num_batches=n_batches,\n",
    "    rand_seed=rand_seed,\n",
    ")\n",
    "\n",
    "print(batches[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93a6d05a",
   "metadata": {},
   "source": [
    "### Obtain decimal representation of the spin-up and spin-down bitstrings used by the eigenstate solver\n",
    "\n",
    "The fist element in the tuple corresponds to the decimal representation of the spin-up configurations, while the second element in the tuple corresponds to the decimal representation of the spin-down configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ef90e039",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(array([1, 2, 4, 8], dtype=int64), array([1, 2, 4, 8], dtype=int64))\n"
     ]
    }
   ],
   "source": [
    "from qiskit_addon_sqd.fermion import bitstring_matrix_to_ci_strs\n",
    "\n",
    "ci_strs = bitstring_matrix_to_ci_strs(batches[0], open_shell=open_shell)\n",
    "print(ci_strs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d7883c",
   "metadata": {},
   "source": [
    "Note that while the number of samples per batch is 2, and the sampled bitstrings are: $00010010$ and $01001000$, four electronic configurations are generated per spin-species. In this case, the set of unique spin-polarized configurations is given by:\n",
    "$$\n",
    "\\mathcal{U} = \\{ 0001, 0010, 0100, 1000 \\}\n",
    "$$\n",
    "whose base-10 decimal representation is \n",
    "$$\n",
    "\\mathcal{U}_{10} = \\{ 1, 2, 4, 8 \\}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03b32ed5",
   "metadata": {},
   "source": [
    "### Basis of the subspace:\n",
    "\n",
    "The eigenstate solver takes all possible pairs of spin-up and spin-down bitstrings to construnct the basis $\\mathcal{B}$ of the subspace:\n",
    " \n",
    "- Element 1: $|00010001\\rangle$ \n",
    "\n",
    "- Element 2: $|00010010\\rangle$ \n",
    "\n",
    "- Element 3: $|00010100\\rangle$ \n",
    "\n",
    "- Element 4: $|00011000\\rangle$ \n",
    "\n",
    "- Element 5: $|00100001\\rangle$ \n",
    "\n",
    "- Element 6: $|00100010\\rangle$ \n",
    "\n",
    "- Element 7: $|00100100\\rangle$ \n",
    "\n",
    "- Element 8: $|00101000\\rangle$ \n",
    " \n",
    "- Element 9: $|01000001\\rangle$ \n",
    "\n",
    "- Element 10: $|01000010\\rangle$ \n",
    "\n",
    "- Element 11: $|01000100\\rangle$ \n",
    "\n",
    "- Element 12: $|01001000\\rangle$ \n",
    "\n",
    "- Element 13: $|10000001\\rangle$ \n",
    "\n",
    "- Element 14: $|10000010\\rangle$ \n",
    "\n",
    "- Element 15: $|10000100\\rangle$ \n",
    "\n",
    "- Element 16: $|10001000\\rangle$ \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "11c924ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Basis elements of the subspace:\n",
      "|00010001>\n",
      "|00010010>\n",
      "|00010100>\n",
      "|00011000>\n",
      "|00100001>\n",
      "|00100010>\n",
      "|00100100>\n",
      "|00101000>\n",
      "|01000001>\n",
      "|01000010>\n",
      "|01000100>\n",
      "|01001000>\n",
      "|10000001>\n",
      "|10000010>\n",
      "|10000100>\n",
      "|10001000>\n"
     ]
    }
   ],
   "source": [
    "ci_strs_up, ci_strs_dn = ci_strs\n",
    "\n",
    "print(\"Basis elements of the subspace:\")\n",
    "\n",
    "for ci_str_up in ci_strs_up:\n",
    "    for ci_str_dn in ci_strs_dn:\n",
    "        format_name = \"{0:0\" + str(num_orbitals) + \"b}\"\n",
    "        print(\"|\" + format_name.format(ci_str_up) + format_name.format(ci_str_dn) + \">\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa43a4fc",
   "metadata": {},
   "source": [
    "**The subspace dimension is upper-bounded by**:  $2 \\cdot$ (`samples_per_batch`)$^2$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9412e52b",
   "metadata": {},
   "source": [
    "## Open-Shell\n",
    "\n",
    "This example shows how the bitstrings are manipulated in a (2-electron, 4-orbital) system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9b515b8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify molecule properties\n",
    "num_orbitals = 4\n",
    "num_elec_a = num_elec_b = 1\n",
    "open_shell = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef78560",
   "metadata": {},
   "source": [
    "### Specify by hand a dictionary of measurement outcomes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "06b2185a",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_dict = {\"00010010\": 1 / 2.0 - 0.01, \"01001000\": 1 / 2.0 - 0.01, \"00010001\": 0.02}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08b32957",
   "metadata": {},
   "source": [
    "### Transform the counts dict into a bitstring matrix and probability array for post-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "90561893",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[False False False  True False False  True False]\n",
      " [False  True False False  True False False False]\n",
      " [False False False  True False False False  True]]\n",
      "[0.49 0.49 0.02]\n"
     ]
    }
   ],
   "source": [
    "# Convert counts into bitstring and probability arrays\n",
    "bitstring_matrix_full, probs_arr_full = counts_to_arrays(counts_dict)\n",
    "print(bitstring_matrix_full)\n",
    "print(probs_arr_full)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "416bfb6c",
   "metadata": {},
   "source": [
    "### Subsample a single batch of size two:\n",
    "\n",
    "- ``n_batches = 1``: Number of batches of configurations used by the different calls to the eigenstate solver\n",
    "- ``samples_per_batch = 2``: Number of unique configurations to include in each batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cf4fe11d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[False False False  True False False  True False]\n",
      " [False  True False False  True False False False]]\n"
     ]
    }
   ],
   "source": [
    "n_batches = 1\n",
    "samples_per_batch = 2\n",
    "\n",
    "# seed for random number generator\n",
    "rand_seed = 48\n",
    "\n",
    "# Generate the batches\n",
    "batches = postselect_and_subsample(\n",
    "    bitstring_matrix_full,\n",
    "    probs_arr_full,\n",
    "    hamming_right=num_elec_a,\n",
    "    hamming_left=num_elec_b,\n",
    "    samples_per_batch=samples_per_batch,\n",
    "    num_batches=n_batches,\n",
    "    rand_seed=rand_seed,\n",
    ")\n",
    "\n",
    "print(batches[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54d699ca",
   "metadata": {},
   "source": [
    "### Obtain decimal representation of the spin-up and spin-down bitstrings used by the eigenstate solver\n",
    "\n",
    "The fist element in the tuple corresponds to the decimal representation of the spin-up configurations, while the second element in the tuple corresponds to the decimal representation of the spin-down configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b40b049b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(array([2, 8], dtype=int64), array([1, 4], dtype=int64))\n"
     ]
    }
   ],
   "source": [
    "ci_strs = bitstring_matrix_to_ci_strs(batches[0], open_shell=open_shell)\n",
    "print(ci_strs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28921e56",
   "metadata": {},
   "source": [
    "If we specify that `open_shell = True`, now we do not include all unique half-bitstrings as spin-up and spin-down configurations, thus yielding a smaller basis as when specifying `open_shell = False`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1959b72",
   "metadata": {},
   "source": [
    "### Basis of the subspace:\n",
    "\n",
    "The eigenstate solver takes all possible pairs of spin-up and spin-down bitstrings to construnct the basis $\\mathcal{B}$ of the subspace:\n",
    " \n",
    "- Element 1: $|00010010\\rangle$ \n",
    "\n",
    "- Element 2: $|00011000\\rangle$ \n",
    "\n",
    "- Element 3: $|01000010\\rangle$ \n",
    "\n",
    "- Element 4: $|01001000\\rangle$ \n",
    "\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a550aba2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Basis elements of the subspace:\n",
      "|00100001>\n",
      "|00100100>\n",
      "|10000001>\n",
      "|10000100>\n"
     ]
    }
   ],
   "source": [
    "ci_strs_up, ci_strs_dn = ci_strs\n",
    "\n",
    "print(\"Basis elements of the subspace:\")\n",
    "\n",
    "for ci_str_up in ci_strs_up:\n",
    "    for ci_str_dn in ci_strs_dn:\n",
    "        format_name = \"{0:0\" + str(num_orbitals) + \"b}\"\n",
    "        print(\"|\" + format_name.format(ci_str_up) + format_name.format(ci_str_dn) + \">\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7317bc49",
   "metadata": {},
   "source": [
    "**The subspace dimension is upper-bounded by**:  (`samples_per_batch`)$^2$"
   ]
  }
 ],
 "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}