{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "65adf860-4443-494f-ad90-5d8ea444ac4e",
   "metadata": {},
   "source": [
    "# Scale SQD chemistry workflows with Dice solver\n",
    "\n",
    "For information on how to install and use ``qiskit-addon-dice-solver``, [visit the docs](https://qiskit.github.io/qiskit-addon-dice-solver/).\n",
    "\n",
    "For more details on the SQD code used in this example, check out [tutorial 1](https://qiskit.github.io/qiskit-addon-sqd/tutorials/01_chemistry_hamiltonian.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "baf7d074-4443-4695-875e-65f7fb8cef25",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "from pyscf import ao2mo, tools\n",
    "from qiskit_addon_dice_solver import solve_fermion\n",
    "from qiskit_addon_sqd.configuration_recovery import recover_configurations\n",
    "from qiskit_addon_sqd.counts import counts_to_arrays, generate_counts_uniform\n",
    "from qiskit_addon_sqd.fermion import (\n",
    "    flip_orbital_occupancies,\n",
    ")\n",
    "from qiskit_addon_sqd.subsampling import postselect_and_subsample\n",
    "\n",
    "# Specify molecule properties\n",
    "num_orbitals = 16\n",
    "num_elec_a = num_elec_b = 5\n",
    "open_shell = False\n",
    "spin_sq = 0\n",
    "\n",
    "# Read in molecule from disk\n",
    "active_space_path = os.path.abspath(os.path.join(\"..\", \"molecules\", \"n2_fci.txt\"))\n",
    "mf_as = tools.fcidump.to_scf(active_space_path)\n",
    "hcore = mf_as.get_hcore()\n",
    "eri = ao2mo.restore(1, mf_as._eri, num_orbitals)\n",
    "nuclear_repulsion_energy = mf_as.mol.energy_nuc()\n",
    "\n",
    "# Create a seed to control randomness throughout this workflow\n",
    "rand_seed = 42\n",
    "\n",
    "# Generate random samples\n",
    "counts_dict = generate_counts_uniform(10_000, num_orbitals * 2, rand_seed=rand_seed)\n",
    "\n",
    "# Convert counts into bitstring and probability arrays\n",
    "bitstring_matrix_full, probs_arr_full = counts_to_arrays(counts_dict)\n",
    "\n",
    "# SQD options\n",
    "iterations = 5\n",
    "\n",
    "# Eigenstate solver options\n",
    "n_batches = 5\n",
    "samples_per_batch = 300\n",
    "\n",
    "# Self-consistent configuration recovery loop\n",
    "occupancies_bitwise = None  # orbital i corresponds to column i in bitstring matrix\n",
    "e_hist = np.zeros((iterations, n_batches))\n",
    "for i in range(iterations):\n",
    "    # On the first iteration, we have no orbital occupancy information from the\n",
    "    # solver, so we just post-select from the full bitstring set based on hamming weight.\n",
    "    if occupancies_bitwise is None:\n",
    "        bs_mat_tmp = bitstring_matrix_full\n",
    "        probs_arr_tmp = probs_arr_full\n",
    "\n",
    "    # In following iterations, we use both the occupancy info and the target hamming\n",
    "    # weight to refine bitstrings.\n",
    "    else:\n",
    "        bs_mat_tmp, probs_arr_tmp = recover_configurations(\n",
    "            bitstring_matrix_full,\n",
    "            probs_arr_full,\n",
    "            occupancies_bitwise,\n",
    "            num_elec_a,\n",
    "            num_elec_b,\n",
    "            rand_seed=rand_seed,\n",
    "        )\n",
    "\n",
    "    # Throw out samples with incorrect hamming weight and create batches of subsamples.\n",
    "    batches = postselect_and_subsample(\n",
    "        bs_mat_tmp,\n",
    "        probs_arr_tmp,\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",
    "    # Run eigenstate solvers in a loop. This loop should be parallelized for larger problems.\n",
    "    int_e = np.zeros(n_batches)\n",
    "    int_occs = np.zeros((n_batches, 2 * num_orbitals))\n",
    "    for j, batch in enumerate(batches):\n",
    "        energy_sci, wf_mags, avg_occs = solve_fermion(\n",
    "            batch,\n",
    "            hcore,\n",
    "            eri,\n",
    "            mpirun_options=[\"-n\", \"8\"],\n",
    "        )\n",
    "        int_e[j] = energy_sci + nuclear_repulsion_energy\n",
    "        int_occs[j, :num_orbitals] = avg_occs[0]\n",
    "        int_occs[j, num_orbitals:] = avg_occs[1]\n",
    "\n",
    "    # Combine batch results\n",
    "    avg_occupancy = np.mean(int_occs, axis=0)\n",
    "    # The occupancies from the solver should be flipped to match the bits in the bitstring matrix.\n",
    "    occupancies_bitwise = flip_orbital_occupancies(avg_occupancy)\n",
    "\n",
    "    # Track optimization history\n",
    "    e_hist[i, :] = int_e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "99709a88-f9ec-4dbc-a561-e682f90aa4c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exact energy: -109.10288938\n",
      "Estimated energy: -109.03013363477585\n"
     ]
    }
   ],
   "source": [
    "print(\"Exact energy: -109.10288938\")\n",
    "print(f\"Estimated energy: {np.min(e_hist[-1])}\")"
   ]
  }
 ],
 "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
}