{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "65adf860-4443-494f-ad90-5d8ea444ac4e",
   "metadata": {},
   "source": [
    "# Scale SQD chemistry workflows with Fulqrum\n",
    "\n",
    "[Fulqrum](https://github.com/qiskit-community/fulqrum) is a software package for enabling the computation of eigenvalues and eigenvectors of Hamiltonians projected into a subspace. It achieves this by constructing a portable representation of the projected Hamiltonian that can then be passed to an eigenvalue solver of your choice, including standard off-the-shelf solvers. These solvers can provide a much faster alternative to the default PySCF solver used for diagonalization in the SQD addon. For information on how to install and use Fulqrum, [visit the docs](https://qiskit-community.github.io/fulqrum/dev/).\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": null,
   "id": "baf7d074-4443-4695-875e-65f7fb8cef25",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "converged SCF energy = -108.835236570775\n",
      "CASCI E = -109.046671778080  E(CI) = -32.8155692383188  S^2 = 0.0000000\n"
     ]
    }
   ],
   "source": [
    "import fulqrum as fq\n",
    "import numpy as np\n",
    "import pyscf\n",
    "import pyscf.cc\n",
    "import pyscf.mcscf\n",
    "import scipy.sparse.linalg\n",
    "from fulqrum.convert.integrals import integrals_to_fq_fermionic_op\n",
    "from qiskit_addon_sqd.counts import generate_bit_array_uniform\n",
    "from qiskit_addon_sqd.fermion import SCIResult, SCIState, diagonalize_fermionic_hamiltonian\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",
    "# Convert Hamiltonian to Fulqrum format\n",
    "fermionic_op = integrals_to_fq_fermionic_op(one_body_integrals=hcore, two_body_integrals=eri)\n",
    "fulqrum_operator = fermionic_op.extended_jw_transformation()\n",
    "\n",
    "# Initialize a Hamiltonian with a trivial subspace to be updated later\n",
    "Hsub = fq.SubspaceHamiltonian(\n",
    "    fulqrum_operator, fq.Subspace([[\"0\" * num_orbitals], [\"0\" * num_orbitals]])\n",
    ")\n",
    "\n",
    "\n",
    "# Define fulqrum solver\n",
    "def solve_sci_batch(ci_batches, _h1, _h2, norb, nelec, tol=1e-12):\n",
    "    # Ignore the _h1 and _h2 inputs because we store the Hamiltonian separately\n",
    "    results = []\n",
    "    for ints_a, ints_b in ci_batches:\n",
    "        strs_a = [format(x, f\"0{norb}b\") for x in ints_a]\n",
    "        strs_b = [format(x, f\"0{norb}b\") for x in ints_b]\n",
    "        subspace = fq.Subspace([strs_a, strs_b])\n",
    "        Hsub.update_subspace(subspace)\n",
    "        Hsub_csr_linop = Hsub.to_csr_linearoperator_fast()\n",
    "        diag = Hsub.diagonal_vector()\n",
    "        v0 = np.zeros(len(subspace), dtype=Hsub.dtype)\n",
    "        v0[np.argmin(diag)] = 1.0\n",
    "        eigs, vecs = scipy.sparse.linalg.eigsh(Hsub_csr_linop, k=1, which=\"SA\", tol=tol, v0=v0)\n",
    "        vec = vecs[:, 0]\n",
    "        occs = subspace.get_orbital_occupancies(probs=np.abs(vec) ** 2, norb=norb)\n",
    "        # SQD ordering and Fulqrum ordering are transpose of each other.\n",
    "        # ascontiguousarray() is needed for matching C contiguous ordering inside PySCF\n",
    "        state = SCIState(\n",
    "            amplitudes=np.ascontiguousarray(vec.reshape(len(strs_b), len(strs_a)).T),\n",
    "            ci_strs_a=ints_a,\n",
    "            ci_strs_b=ints_b,\n",
    "            norb=norb,\n",
    "            nelec=nelec,\n",
    "        )\n",
    "        results.append(SCIResult(energy=eigs[0], sci_state=state, orbital_occupancies=occs))\n",
    "    return results\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=solve_sci_batch,\n",
    "    symmetrize_spin=True,\n",
    "    seed=rng,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "99709a88-f9ec-4dbc-a561-e682f90aa4c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exact energy: -109.04667177808028\n",
      "Estimated energy: -109.0340266755937\n"
     ]
    }
   ],
   "source": [
    "print(f\"Exact energy: {exact_energy}\")\n",
    "print(f\"Estimated energy: {result.energy + nuclear_repulsion_energy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab599762",
   "metadata": {},
   "source": [
    "## Use a different eigensolver\n",
    "\n",
    "Because Fulqrum exposes the projected Hamiltonian as a standard [`scipy.sparse.linalg.LinearOperator`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.LinearOperator.html), you are free to diagonalize it with any eigensolver of your choosing rather than being tied to `scipy.sparse.linalg.eigsh`. This makes it easy to swap in a solver that is better suited to your problem size or hardware.\n",
    "\n",
    "As an example, we demonstrate using [PRIMME](https://github.com/primme/primme), a high-performance eigensolver library. Its Python interface, `primme.eigsh`, follows a similar calling convention as `scipy.sparse.linalg.eigsh`, so it can be used as a drop-in replacement.\n",
    "\n",
    "The only change from the previous solver is that we call `primme.eigsh` instead of `scipy.sparse.linalg.eigsh`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48850fb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import primme\n",
    "\n",
    "\n",
    "# Define a fulqrum solver that uses PRIMME to diagonalize the projected Hamiltonian\n",
    "def solve_sci_batch_primme(ci_batches, _h1, _h2, norb, nelec, tol=1e-12):\n",
    "    # Ignore the _h1 and _h2 inputs because we store the Hamiltonian separately\n",
    "    results = []\n",
    "    for ints_a, ints_b in ci_batches:\n",
    "        strs_a = [format(x, f\"0{norb}b\") for x in ints_a]\n",
    "        strs_b = [format(x, f\"0{norb}b\") for x in ints_b]\n",
    "        subspace = fq.Subspace([strs_a, strs_b])\n",
    "        Hsub.update_subspace(subspace)\n",
    "        Hsub_csr_linop = Hsub.to_csr_linearoperator_fast()\n",
    "        diag = Hsub.diagonal_vector()\n",
    "        v0 = np.zeros(len(subspace), dtype=Hsub.dtype)\n",
    "        v0[np.argmin(diag)] = 1.0\n",
    "        # primme.eigsh follows the same calling convention as scipy.sparse.linalg.eigsh,\n",
    "        # except that it expects the initial guess v0 as a 2D array of column vectors.\n",
    "        eigs, vecs = primme.eigsh(Hsub_csr_linop, k=1, which=\"SA\", tol=tol, v0=v0.reshape(-1, 1))\n",
    "        vec = vecs[:, 0]\n",
    "        occs = subspace.get_orbital_occupancies(probs=np.abs(vec) ** 2, norb=norb)\n",
    "        # SQD ordering and Fulqrum ordering are transpose of each other.\n",
    "        # ascontiguousarray() is needed for matching C contiguous ordering inside PySCF\n",
    "        state = SCIState(\n",
    "            amplitudes=np.ascontiguousarray(vec.reshape(len(strs_b), len(strs_a)).T),\n",
    "            ci_strs_a=ints_a,\n",
    "            ci_strs_b=ints_b,\n",
    "            norb=norb,\n",
    "            nelec=nelec,\n",
    "        )\n",
    "        results.append(SCIResult(energy=eigs[0], sci_state=state, orbital_occupancies=occs))\n",
    "    return results\n",
    "\n",
    "\n",
    "# Run SQD using the PRIMME-based solver\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=solve_sci_batch_primme,\n",
    "    symmetrize_spin=True,\n",
    "    seed=rng,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "39b960fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exact energy: -109.04667177808028\n",
      "Estimated energy: -109.03937429196165\n"
     ]
    }
   ],
   "source": [
    "print(f\"Exact energy: {exact_energy}\")\n",
    "print(f\"Estimated energy: {result.energy + nuclear_repulsion_energy}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.13.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}