{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1a4c7d2-9b3f-4e18-8a52-6cd07f1b94ae",
   "metadata": {},
   "source": [
    "# Scale SQD chemistry workflows with the SBD eigensolver\n",
    "\n",
    "This guide shows how to use the [SBD eigensolver](https://github.com/Qiskit/sbd-eigensolver-python) as an alternative SCI solver to diagonalize larger Fermionic problems beyond what the default PySCF solver can support. SBD (Selected Basis Diagonalization) runs in-process, is parallelized with MPI, and offers both CPU and CUDA backends. For information on how to install and use it, [visit the documentation](https://github.com/Qiskit/sbd-eigensolver-python).\n",
    "\n",
    "SBD compiles a C++ extension against the MPI and BLAS libraries available on the machine, so both must be installed before you install the ``sbd`` extra:\n",
    "\n",
    "```bash\n",
    "pip install qiskit-addon-sqd[sbd]\n",
    "```\n",
    "\n",
    "Because ``solve_sci_batch`` already matches the solver interface expected by `diagonalize_fermionic_hamiltonian`, it can be passed directly to `sci_solver` once its configuration options have been bound with `functools.partial`. To run on a CUDA-capable GPU instead of the CPU, pass `device_config=DeviceConfig.gpu()` from `sbd.device_config`.\n",
    "\n",
    "For more details on the SQD code used in this example, see the [chemistry Hamiltonian tutorial](https://quantum.cloud.ibm.com/docs/tutorials/sample-based-quantum-diagonalization)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f2c091-5d84-4a37-9e6b-1f3a8c25d70b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "import numpy as np\n",
    "import pyscf\n",
    "import pyscf.cc\n",
    "import pyscf.mcscf\n",
    "from qiskit_addon_sqd.counts import generate_bit_array_uniform\n",
    "from qiskit_addon_sqd.fermion import diagonalize_fermionic_hamiltonian\n",
    "from sbd.sbd_solver import solve_sci_batch\n",
    "\n",
    "# Specify molecule properties\n",
    "num_orbitals = 16\n",
    "num_elec_a = num_elec_b = 5\n",
    "spin_sq = 0\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)\n",
    "eri = pyscf.ao2mo.restore(1, cas.get_h2cas(mo), num_orbitals)\n",
    "\n",
    "# Compute exact energy\n",
    "exact_energy = cas.run().e_tot\n",
    "\n",
    "# Create a seed to control randomness throughout this workflow\n",
    "rng = np.random.default_rng(24)\n",
    "\n",
    "\n",
    "# Generate random samples\n",
    "bit_array = generate_bit_array_uniform(10_000, num_orbitals * 2, rand_seed=rng)\n",
    "\n",
    "# Bind the SBD solver options. method=0 selects the Davidson eigensolver.\n",
    "sbd_solver = partial(\n",
    "    solve_sci_batch,\n",
    "    sbd_config={\"method\": 0, \"eps\": 1e-5, \"max_it\": 10, \"max_nb\": 10},\n",
    ")\n",
    "\n",
    "# Run SQD\n",
    "result = diagonalize_fermionic_hamiltonian(\n",
    "    hcore,\n",
    "    eri,\n",
    "    bit_array,\n",
    "    samples_per_batch=300,\n",
    "    norb=num_orbitals,\n",
    "    nelec=(num_elec_a, num_elec_b),\n",
    "    num_batches=5,\n",
    "    max_iterations=5,\n",
    "    sci_solver=sbd_solver,\n",
    "    symmetrize_spin=True,\n",
    "    seed=rng,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d95e46-8a12-4f70-b6e9-52d418c0fa37",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Exact energy: {exact_energy}\")\n",
    "print(f\"Estimated energy: {result.energy + nuclear_repulsion_energy}\")"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}