From 95570a773562b9d665b047fec236ec0b6f277da6 Mon Sep 17 00:00:00 2001 From: William Galvin Date: Sat, 23 Dec 2023 01:00:11 -0800 Subject: [PATCH] Add main function --- src/hamiltonian.py | 19 ------ src/main.py | 152 +++++++++++++++++++++++++++++++++++++++++++++ src/utils.py | 3 - 3 files changed, 152 insertions(+), 22 deletions(-) create mode 100644 src/main.py diff --git a/src/hamiltonian.py b/src/hamiltonian.py index 697bc07..264db32 100644 --- a/src/hamiltonian.py +++ b/src/hamiltonian.py @@ -101,17 +101,6 @@ def annihilation(i: int, N: int) -> np.ndarray: return A -def target_energy(path: str, basis: str = "sto-3g") -> float: - """ - Uses pyscf to calculate the target electronic (e1) energy. - When we minimize theta, the expectation of H w.r.t. psi - should be close to this - """ - mol_pyscf = pyscf.gto.M(atom=path, basis=basis) - hartree, coulomb = mol_pyscf.energy_elec() - return (hartree - coulomb).real - - def expectation(H: np.ndarray, psi: np.ndarray) -> float: """ Calculates / = the expectation @@ -119,11 +108,3 @@ def expectation(H: np.ndarray, psi: np.ndarray) -> float: """ return ((psi.conj() @ H @ psi) / (psi.conj() @ psi)).real - -if __name__ == "__main__": - file = "data/h4.xyz" - H = hamiltonian(file) - print("H shape:", H.shape) - print("target energy:", target_energy(file)) - np.save("h4_H.npy", H) - diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..fe68340 --- /dev/null +++ b/src/main.py @@ -0,0 +1,152 @@ +import tempfile +import sys +import argparse +import io + +import jax +jax.config.update("jax_platform_name", "cpu") + +import numpy as np + +from jax import numpy as jnp +from jax import value_and_grad + +from givens import givens +from givens import hf_ground +from givens import full_wavefunction +from givens import excitations +from hamiltonian import expectation +from hamiltonian import hamiltonian +from utils import get_electrons_qubits +from utils import random_xyz + + +def energy(H: np.ndarray, electrons: int, qubits: int, thetas: np.ndarray) -> float: + """ + Returns the expected value of H for the system specified by (electrons, qubits, thetas) + + Params: + - H: hamiltonian + - electrons: n_electrons + - qubits: n_qubits + - thetas: vector of angles + """ + G = givens(electrons, qubits, thetas) + + psi_valid = hf_ground(electrons, qubits) @ G + psi = full_wavefunction(psi_valid, electrons, qubits) + + return expectation(H, psi) + + +def get_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="This script implements a very classical VQE to find the lowest enegery state of a hydrogen molecule", + epilog="For more information or context, see the paper accompanying this repo, or any VQE for quantum chemistry tutorial (PennyLane, Azure, etc)." + ) + + parser.add_argument( + "--random", + action=argparse.BooleanOptionalAction, + default=False, + help="Specify whether to use a random molecule or supply your own. If --random is set, --num-H must be too. If not, --file must be set" + ) + + parser.add_argument( + "--num-H", + type=int, + default=4, + help="The number of hydrogen atoms to randomly place. This argument is only used if --random is used. If PySCF doesn't support this molecule, program crashes." + ) + + parser.add_argument( + "--file", + help="Path to .xyz file of input molecule if --random is not set" + ) + + parser.add_argument( + "--max-iter", + type=int, + default=25, + help="Maximum number of optimization cycles" + ) + + parser.add_argument( + "--min-iter", + type=int, + default=5, + help="Minimum number of optimization cycles" + ) + + parser.add_argument( + "--threshold", + type=float, + default=0.001, + help="Convergence threshdold" + ) + + parser.add_argument( + "--step-size", + type=float, + default=1, + help="Optimization step size aka learning rate" + ) + + return parser.parse_args() + + +if __name__ == "__main__": + linebreak = 80 * "=" + + args = get_args() + + with tempfile.NamedTemporaryFile() as f: + if args.random: + if args.num_H is None: + raise ValueError("Must specify --num-H if --random is used") + + random_xyz({"H": args.num_H}, f.name) + print("Random H4 molecule generated:\n") + print(open(f.name).read()) + print(linebreak) + + file = f.name + else: + file = args.file + + try: + electrons, qubits = get_electrons_qubits(file) + except RuntimeError: + raise RuntimeError("Molecule not valid for PySCF") + + H = hamiltonian(file) + print(f"Hamiltonian created for {electrons} electrons and {qubits} spin orbitals") + print(linebreak) + + singles, doubles = excitations(electrons, qubits) + print(linebreak) + print(f"{len(singles)} single excitations and {len(doubles)} double excitations") + print(linebreak) + + theta = np.random.random(len(singles) + len(doubles)) * jnp.pi + + dE = value_and_grad(energy, argnums=3) + + print("Beginning optimization:") + + sys.stderr = io.StringIO() + prev = float("inf") + for i in range(args.max_iter): + E, theta_prime = dE(H, electrons, qubits, theta) + theta -= theta_prime * args.step_size + + print(f"\tstep: {i: 2}\t\tE: {E: .3f}\t\tDifference: {E - prev: .3f}") + + if abs(E - prev) < args.threshold and i >= args.min_iter: + break + + prev = E + + print(linebreak) + print(f"Final E converged to {E: 5f} in {i} steps") + \ No newline at end of file diff --git a/src/utils.py b/src/utils.py index 798ce09..bee02dd 100644 --- a/src/utils.py +++ b/src/utils.py @@ -15,9 +15,6 @@ def random_xyz(mol: Dict[str, int], path: str): - path: place to write file """ with open(path, "w+") as f: - f.write(str(sum(mol.values()))) - f.write("\n\n") - for sym, count in mol.items(): for _ in range(count): f.write(sym)