docs/qcqmc/full_workflow.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "8fcd86dea77a"
   },
   "source": [
    "# End-to-End Example\n",
    "\n",
    "We learned in the [high-level](high-level.ipynb) notebook how to setup the quantum side of the calculation, which involves specifying a Hamiltonian, defining a trial wavefunction ansatz before performing shadow tomography to extract the AFQMC trial wavefunction. In this example, we will repeat these steps again for the case of H4 which was studied in the [paper](https://arxiv.org/pdf/2106.16235.pdf). Then we will outline how to interface the output of the quantum half of the calculation with [ipie](https://github.com/JoonhoLee-Group/ipie/) to perform AFQMC with the quantum trial wavefunction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d6b7b1b4dab4"
   },
   "source": [
    "# Setup\n",
    "\n",
    "First install recirq:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0b0bebea3b3b"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import recirq\n",
    "except ImportError:\n",
    "    %pip install git+https://github.com/quantumlib/ReCirq"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "b3b7d7eca9a7"
   },
   "source": [
    "## Define an SCF Job\n",
    "\n",
    "We first setup a `PyscfHamiltonianParams` object which defines the SCF job we will run using [pyscf](https://github.com/pyscf/pyscf). Here we are simulation H4 in the (minimal) sto-3g basis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4a1faf7a8ab8"
   },
   "outputs": [],
   "source": [
    "from recirq.qcqmc.hamiltonian import PyscfHamiltonianParams, HamiltonianData\n",
    "\n",
    "pyscf_params = PyscfHamiltonianParams(\n",
    "    name=\"TEST_square_H4\",\n",
    "    n_orb=4,\n",
    "    n_elec=4,\n",
    "    geometry=((\"H\", (0, 0, 0)), (\"H\", (0, 0, 1.23)), (\"H\", (1.23, 0, 0)), (\"H\", (1.23, 0, 1.23))),\n",
    "    basis=\"sto3g\",\n",
    "    multiplicity=1,\n",
    "    charge=0,\n",
    "    save_chkfile=True,\n",
    "    overwrite_chk_file=True,\n",
    ")\n",
    "pyscf_hamiltonian = HamiltonianData.build_hamiltonian_from_pyscf(pyscf_params)\n",
    "chk_path = pyscf_params.base_path.with_suffix(\".chk\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "6b18f3db0743"
   },
   "source": [
    "## Build Perfect Pairing Wavefunction\n",
    "\n",
    "Next we build a perfect pairing wavefunction based upon our `pyscf_params`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "180a012e0579"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from recirq.qcqmc.trial_wf import (\n",
    "    PerfectPairingPlusTrialWavefunctionParams,\n",
    "    TrialWavefunctionData\n",
    ")\n",
    "\n",
    "pp_params = PerfectPairingPlusTrialWavefunctionParams(\n",
    "    'pp_plus_test',\n",
    "    hamiltonian_params=pyscf_params,\n",
    "    heuristic_layers=tuple(),\n",
    "    do_pp=True,\n",
    "    restricted=False,\n",
    "    initial_orbital_rotation=None,\n",
    "    initial_two_body_qchem_amplitudes=np.asarray([0.3, 0.4]),\n",
    "    do_optimization=True,\n",
    "    use_fast_gradients=True\n",
    ")\n",
    "\n",
    "trial_wf = TrialWavefunctionData.build_pp_plus_trial_from_dependencies(\n",
    "    pp_params, dependencies={pyscf_params: pyscf_hamiltonian}, do_print=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0fe3452d66a6"
   },
   "source": [
    "We can comare the wavefunction of the optimized ansatz to the exact ground state energy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f50a14acd6af"
   },
   "outputs": [],
   "source": [
    "print(f\"Trial wavefunction energy: {trial_wf.ansatz_energy}\")\n",
    "print(f\"Exact ground state energy: {trial_wf.fci_energy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "06ad291fd2e1"
   },
   "source": [
    "Next, we build an `experiment` which will combine the trial wavefunction circuit with that required for shadow tomography. In this case the experiment will be simulated using a statevector simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2301d967e68d"
   },
   "outputs": [],
   "source": [
    "from recirq.qcqmc.blueprint import BlueprintParamsTrialWf, BlueprintData\n",
    "from recirq.qcqmc.experiment import SimulatedExperimentParams, ExperimentData\n",
    "\n",
    "blueprint_params = BlueprintParamsTrialWf(\n",
    "    name=\"blueprint_test_medium\",\n",
    "    trial_wf_params=pp_params,\n",
    "    n_cliffords=100,\n",
    "    qubit_partition=(tuple(qubit for qubit in pp_params.qubits_jordan_wigner_ordered),),\n",
    "    seed=1,\n",
    ")\n",
    "\n",
    "blueprint = BlueprintData.build_blueprint_from_dependencies(blueprint_params, dependencies={pp_params: trial_wf})\n",
    "\n",
    "simulated_experiment_params = SimulatedExperimentParams(\n",
    "    name=\"test_1\",\n",
    "    blueprint_params=blueprint.params,\n",
    "    noise_model_name=\"None\",\n",
    "    noise_model_params=(0,),\n",
    "    n_samples_per_clifford=31,\n",
    "    seed=1,\n",
    ")\n",
    "experiment = ExperimentData.build_experiment_from_dependencies(\n",
    "    params=simulated_experiment_params, dependencies={blueprint.params: blueprint}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "55aa2e8b0050"
   },
   "source": [
    "The experimental output is post-processed to extract the reconstructed trial wavefunction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6d22cbe86339"
   },
   "outputs": [],
   "source": [
    "from typing import Dict\n",
    "from recirq.qcqmc.data import Params, Data\n",
    "from recirq.qcqmc.analysis import OverlapAnalysisData, OverlapAnalysisParams\n",
    "\n",
    "analysis_params = OverlapAnalysisParams(\n",
    "    \"TEST_analysis\", experiment_params=experiment.params, k_to_calculate=(1,)\n",
    ")\n",
    "all_dependencies: Dict[Params, Data] = {\n",
    "    pyscf_params: pyscf_hamiltonian,\n",
    "    pp_params: trial_wf,\n",
    "    blueprint_params: blueprint,\n",
    "    simulated_experiment_params: experiment,\n",
    "}\n",
    "analysis = OverlapAnalysisData.build_analysis_from_dependencies(\n",
    "    analysis_params, dependencies=all_dependencies\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "040bde1bfe54"
   },
   "source": [
    "Finally, we save the wavefunction in a format that is acceptable for ipie."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7cecbbc8052a"
   },
   "outputs": [],
   "source": [
    "from recirq.qcqmc.convert_to_ipie import save_wavefunction_for_ipie\n",
    "ipie_data = save_wavefunction_for_ipie(\n",
    "    hamiltonian_data=pyscf_hamiltonian, trial_wf_data=trial_wf, overlap_analysis_data=analysis, do_print=False\n",
    ")\n",
    "print(f\"Reconstructed shadow wavefunction energy: {ipie_data.variational_energy}\")\n",
    "print(f\"Ideal trial wavefunction energy: {trial_wf.ansatz_energy}\")\n",
    "print(f\"FCI energy: {trial_wf.fci_energy}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "bf3fed265e17"
   },
   "source": [
    "## Run AFQMC \n",
    "\n",
    "Now that we have built our quantum wavefunction, we can use it as a trial wavefunction in an AFQMC simulation. First, we need to build a factorized Hamiltonian which is required for AFQMC. \n",
    "\n",
    "In particular, we require the three-index Cholesky tensor (`ham.chol` below) which satisfies\n",
    "\n",
    "$$\n",
    "(pq|rs) = \\sum_X L_{pq}^X L_{rs}^X.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "927c01d9cf98"
   },
   "outputs": [],
   "source": [
    "from ipie.systems.generic import Generic\n",
    "from ipie.utils.from_pyscf import generate_hamiltonian_from_chk\n",
    "\n",
    "num_elec = (pyscf_params.n_elec // 2,) * 2\n",
    "system = Generic(num_elec)\n",
    "chk_path = pyscf_params.base_path.with_suffix(\".chk\")\n",
    "ham = generate_hamiltonian_from_chk(str(chk_path))\n",
    "assert ham.H1.shape == (2, pyscf_params.n_orb, pyscf_params.n_orb)\n",
    "assert ham.chol.shape[0] == pyscf_params.n_orb * pyscf_params.n_orb"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "261e70d57dfa"
   },
   "source": [
    "Next, we need to build a trial wavefunction from quantum wavefunction. In practice, the quantum trial is expanded as a linear combination of (orthogonal) Slater Determinants. Within ipie, this style of trial wavefunction is defined as a `ParticleHole` trial wavefunction. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fd9d66faea78"
   },
   "outputs": [],
   "source": [
    "import h5py\n",
    "from ipie.trial_wavefunction.particle_hole import ParticleHole\n",
    "# Read quantum wavefunction from file. \n",
    "with h5py.File(ipie_data.path, 'r') as fh5:\n",
    "    coeffs = fh5[\"coeffs_rotated\"][:]\n",
    "    occa = fh5[\"occa_rotated\"][:]\n",
    "    occb = fh5[\"occb_rotated\"][:]\n",
    "wfn = ParticleHole((coeffs, occa, occb), num_elec, pyscf_params.n_orb)\n",
    "wfn.half_rotate(ham)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2b9c637242f5"
   },
   "outputs": [],
   "source": [
    "wfn.calculate_energy(system, ham)\n",
    "print(f\"Trial wavefunction variational energy in ipie: {wfn.energy.real}\")\n",
    "assert np.isclose(wfn.energy.real, ipie_data.variational_energy)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "c16d33d36551"
   },
   "source": [
    "Note that the variational energy might differ slightly from the result compute from the previous section. This is because the cholesky factorization uses a threshold of $1\\times10^{-5}$ a stopping criteria for convergence. Reducing the parameter `chol_cut` in `generate_hamiltonian_from_chk` will yield better agreement.\n",
    "\n",
    "At this point, we are ready to run AFQMC. We need to build an `AFQMC` driver class which takes as input the factorized Hamiltonian and our `ParticleHole` trial wavefunction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5cce31ab8e71"
   },
   "outputs": [],
   "source": [
    "from ipie.qmc.afqmc import AFQMC\n",
    "\n",
    "qmc_driver = AFQMC.build(num_elec, ham, wfn, num_blocks=300, num_walkers=50, seed=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7cae84bd5b10"
   },
   "outputs": [],
   "source": [
    "from recirq.qcqmc.config import OUTDIRS\n",
    "import pathlib\n",
    "ipie_path = pathlib.Path(pyscf_hamiltonian.params.path_prefix + OUTDIRS.DEFAULT_QMC_DIRECTORY)\n",
    "if not ipie_path.is_dir():\n",
    "    ipie_path.mkdir(parents=True)\n",
    "qmc_driver.run(estimator_filename=f'{ipie_path}/estimates.h5')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "4eb32a627148"
   },
   "source": [
    "Next, we can visualize the data which is by default saved to `estimates.0.h5`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7f2dbce4d733"
   },
   "outputs": [],
   "source": [
    "from ipie.analysis.extraction import extract_observable\n",
    "import matplotlib.pyplot as plt\n",
    "data = extract_observable(qmc_driver.estimators.filename)\n",
    "plt.plot(data.ETotal, marker='o', label=\"AFQMC\", color=\"C0\")\n",
    "plt.axhline(-1.969512, label=\"Exact Result\", color=\"C1\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"Block number\")\n",
    "plt.ylabel(\"Total Energy (Ha)\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "e72a3c54bc56"
   },
   "source": [
    "Through visual inspection, we can determine when the calculation has equilibrated, after which we should be sampling the (approximate) ground state. We can use samples from this point on to estimate the AFQMC energy and error bar. The function `reblock_minimal` will perform the necessary error analysis taking into account the serial temporal correlation in the AFQMC data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "81fe96e4438d"
   },
   "outputs": [],
   "source": [
    "from ipie.analysis.blocking import reblock_minimal \n",
    "reblock_minimal(qmc_driver.estimators.filename, start_block=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "428c61015a05"
   },
   "source": [
    "Note that the number of walkers and the number of blocks is probably too low to obtain statistically significant results and we would advise increasing both of these. For larger scale simulations, we can use MPI to distribute the work among many MPI processes. See [ipie](https://github.com/JoonhoLee-Group/ipie/) for further details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "88925e9e87b9"
   },
   "source": [
    "# Next steps\n",
    "\n",
    "The results presented above are worse than reported in the paper. In practice one needs to:\n",
    "\n",
    "1. Carefully converge the trial wavefunction parameters by scanning over multiple random restarts.\n",
    "2. Converge the reconstructed shadow wavefunction with respect to: `n_cliffords`\n",
    "3. The AFQMC part of the simulations can be sped up using MPI parallelism. See [ipie](https://github.com/JoonhoLee-Group/ipie) for more information. The number of walkers should also be increased to obtain better statistics and reduce any population control biases. The number of samples (`num_blocks`) should also be increased (`ETotal_nsamp_ac` should be about 100 for statistically significant results.) "
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "full_workflow.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}