Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Add QCQMC code and notebooks. #341

Closed
wants to merge 6 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
314 changes: 314 additions & 0 deletions docs/full_workflow.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,314 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"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 but now 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."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"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": {},
"outputs": [],
"source": [
"from recirq.qcqmc.utilities import make_output_directories\n",
"make_output_directories()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from recirq.qcqmc.hamiltonian import build_hamiltonian_from_pyscf, PyscfHamiltonianParams\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 = build_hamiltonian_from_pyscf(pyscf_params)\n",
"chk_path = pyscf_params.base_path.with_suffix(\".chk\")"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Build Perfect Pairing Wavefunction\n",
"\n",
"Next we build a perfect pairing Hamiltonian based upon our `pyscf_params`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"from recirq.qcqmc.trial_wf import (\n",
" build_pp_plus_trial_wavefunction,\n",
" PerfectPairingPlusTrialWavefunctionParams,\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=False,\n",
")\n",
"\n",
"trial_wf = build_pp_plus_trial_wavefunction(\n",
" pp_params, dependencies={pyscf_params: pyscf_hamiltonian}, do_print=False\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from recirq.qcqmc.blueprint import BlueprintParamsTrialWf, build_blueprint\n",
"from recirq.qcqmc.experiment import build_experiment, SimulatedExperimentParams\n",
"\n",
"blueprint_params = BlueprintParamsTrialWf(\n",
" name='blueprint_test_medium',\n",
" trial_wf_params=pp_params,\n",
" n_cliffords=3,\n",
" qubit_partition=(tuple(qubit for qubit in pp_params.qubits_jordan_wigner_ordered),),\n",
" seed=1,\n",
")\n",
"\n",
"blueprint = build_blueprint(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 = build_experiment(\n",
" params=simulated_experiment_params, dependencies={blueprint.params: blueprint}\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import h5py\n",
"from typing import Dict\n",
"from recirq.qcqmc.analysis import build_analysis, OverlapAnalysisParams\n",
"from recirq.qcqmc.newtilities import save_wavefunction_for_ipie\n",
"from recirq.qcqmc.utilities import Data, Params\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 = build_analysis(analysis_params, dependencies=all_dependencies)\n",
"xd = save_wavefunction_for_ipie(\n",
" hamiltonian_data=pyscf_hamiltonian, trial_wf_data=trial_wf, overlap_analysis_data=analysis\n",
")\n",
"with h5py.File(xd.path) as fh5:\n",
" coeffs = fh5[\"coeffs_rotated\"][:]\n",
" occa = fh5[\"occa_rotated\"][:]\n",
" occb = fh5[\"occb_rotated\"][:]"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"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": {},
"outputs": [],
"source": [
"from ipie.utils.from_pyscf import generate_hamiltonian_from_chk\n",
"from ipie.systems.generic import Generic\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(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": {},
"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": {},
"outputs": [],
"source": [
"from ipie.trial_wavefunction.particle_hole import ParticleHole\n",
"\n",
"wfn = ParticleHole((coeffs, occa, occb), num_elec, pyscf_params.n_orb, verbose=True)\n",
"wfn.half_rotate(ham)\n",
"wfn.calculate_energy(system, ham)\n",
"print(f\"Trial wavefunction varitaional energy = {wfn.energy.real}\")"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Note that the variational energy differs slightly from the result in a previous cell. 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": {},
"outputs": [],
"source": [
"from ipie.qmc.afqmc import AFQMC\n",
"\n",
"qmc_driver = AFQMC.build(num_elec, ham, wfn, num_blocks=100, num_walkers=50, seed=7)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"qmc_driver.run(estimator_filename='./estimates.h5')"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Next we can visualize the data which is by default saved to `estimates.0.h5`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from ipie.analysis.extraction import extract_observable\n",
"import matplotlib.pyplot as plt\n",
"data = extract_observable(\"estimates.h5\")\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": {},
"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": {},
"outputs": [],
"source": [
"from ipie.analysis.blocking import reblock_minimal \n",
"reblock_minimal(qmc_driver.estimators.filename)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"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 would can use MPI to distribute the work among many MPI processes. See [ipie](https://github.com/JoonhoLee-Group/ipie/) for further details."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "qcqmc",
"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.10.12"
},
"orig_nbformat": 4
},
"nbformat": 4,
"nbformat_minor": 2
}
Loading
Loading