From 837de385d76bf7192a59ab52ae742c00229ee8bc Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Thu, 13 Jun 2024 22:46:12 +0000 Subject: [PATCH 01/21] Add trial wavefunction. --- recirq/qcqmc/__init__.py | 44 + recirq/qcqmc/conftest.py | 92 ++ recirq/qcqmc/trial_wf.py | 1584 +++++++++++++++++++++++++++++++++ recirq/qcqmc/trial_wf_test.py | 387 ++++++++ 4 files changed, 2107 insertions(+) create mode 100644 recirq/qcqmc/conftest.py create mode 100644 recirq/qcqmc/trial_wf.py create mode 100644 recirq/qcqmc/trial_wf_test.py diff --git a/recirq/qcqmc/__init__.py b/recirq/qcqmc/__init__.py index 0f6532ad..2d86c742 100644 --- a/recirq/qcqmc/__init__.py +++ b/recirq/qcqmc/__init__.py @@ -11,3 +11,47 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from functools import lru_cache +from typing import Optional + +from cirq.protocols.json_serialization import DEFAULT_RESOLVERS, ObjectFactory + +from .hamiltonian import ( + HamiltonianData, + LoadFromFileHamiltonianParams, + PyscfHamiltonianParams, +) +from .trial_wf import ( + FermionicMode, + LayerSpec, + PerfectPairingPlusTrialWavefunctionParams, + TrialWavefunctionData, +) + + +@lru_cache() +def _resolve_json(cirq_type: str) -> Optional[ObjectFactory]: + """Resolve the types of `recirq.qcqmc.` json objects. + + This is a Cirq JSON resolver suitable for appending to + `cirq.protocols.json_serialization.DEFAULT_RESOLVERS`. + """ + if not cirq_type.startswith("recirq.qcqmc."): + return None + + cirq_type = cirq_type[len("recirq.qcqmc.") :] + return { + k.__name__: k + for k in [ + LoadFromFileHamiltonianParams, + HamiltonianData, + PyscfHamiltonianParams, + FermionicMode, + LayerSpec, + PerfectPairingPlusTrialWavefunctionParams, + TrialWavefunctionData, + ] + }.get(cirq_type, None) + + +DEFAULT_RESOLVERS.append(_resolve_json) diff --git a/recirq/qcqmc/conftest.py b/recirq/qcqmc/conftest.py new file mode 100644 index 00000000..7d6c5d08 --- /dev/null +++ b/recirq/qcqmc/conftest.py @@ -0,0 +1,92 @@ +from typing import Tuple + +import numpy as np +import pytest + +from recirq.qcqmc.hamiltonian import ( + HamiltonianData, + LoadFromFileHamiltonianParams, + build_hamiltonian_from_file, +) +from recirq.qcqmc.trial_wf import ( + PerfectPairingPlusTrialWavefunctionParams, + TrialWavefunctionData, + _get_qubits_a_b_reversed, + build_pp_plus_trial_wavefunction, +) + + +@pytest.fixture(scope="package") +def fixture_4_qubit_ham() -> HamiltonianData: + params = LoadFromFileHamiltonianParams( + name="test hamiltonian 4 qubits", integral_key="fh_sto3g", n_orb=2, n_elec=2 + ) + + hamiltonian_data = build_hamiltonian_from_file(params) + + return hamiltonian_data + + +@pytest.fixture(scope="package") +def fixture_8_qubit_ham() -> HamiltonianData: + params = LoadFromFileHamiltonianParams( + name="test hamiltonian 8 qubits", integral_key="h4_sto3g", n_orb=4, n_elec=4 + ) + + hamiltonian_data = build_hamiltonian_from_file(params) + + return hamiltonian_data + + +@pytest.fixture(scope="package") +def fixture_12_qubit_ham() -> HamiltonianData: + params = LoadFromFileHamiltonianParams( + name="test hamiltonian 12 qubits", + integral_key="diamond_dzvp/cas66", + n_orb=6, + n_elec=6, + do_eri_restore=True, + ) + + hamiltonian_data = build_hamiltonian_from_file(params) + + return hamiltonian_data + + +@pytest.fixture(scope="package") +def fixture_4_qubit_ham_and_trial_wf( + fixture_4_qubit_ham: HamiltonianData, +) -> Tuple[HamiltonianData, TrialWavefunctionData]: + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_1", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=tuple(), + do_pp=True, + restricted=True, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham} + ) + + return fixture_4_qubit_ham, trial_wf + + +@pytest.fixture(scope="package") +def fixture_8_qubit_ham_and_trial_wf( + fixture_8_qubit_ham: HamiltonianData, +) -> Tuple[HamiltonianData, TrialWavefunctionData]: + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_qchem", + hamiltonian_params=fixture_8_qubit_ham.params, + heuristic_layers=tuple(), + initial_orbital_rotation=None, + initial_two_body_qchem_amplitudes=np.asarray([0.3, 0.4]), + do_optimization=False, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, dependencies={fixture_8_qubit_ham.params: fixture_8_qubit_ham} + ) + + return fixture_8_qubit_ham, trial_wf diff --git a/recirq/qcqmc/trial_wf.py b/recirq/qcqmc/trial_wf.py new file mode 100644 index 00000000..73c3e9ec --- /dev/null +++ b/recirq/qcqmc/trial_wf.py @@ -0,0 +1,1584 @@ +import abc +import copy +import itertools +from typing import ( + Callable, + Dict, + Iterable, + Iterator, + List, + Mapping, + Optional, + Sequence, + Tuple, + Union, +) + +import attrs +import cirq +import fqe +import fqe.algorithm.low_rank +import fqe.hamiltonians.hamiltonian +import fqe.openfermion_utils +import numpy as np +import numpy.typing as npt +import openfermion as of +import scipy.optimize +from fqe.hamiltonians.restricted_hamiltonian import RestrictedHamiltonian +from fqe.wavefunction import Wavefunction as FqeWavefunction +from scipy.linalg import expm +from scipy.optimize import minimize +from scipy.sparse import csc_matrix + +from recirq.qcqmc.afqmc_circuits import GeminalStatePreparationGate +from recirq.qcqmc.config import OUTDIRS +from recirq.qcqmc.for_refactor import Data, Params +from recirq.qcqmc.hamiltonian import HamiltonianData, HamiltonianParams + + +@attrs.frozen +class FermionicMode: + orb_ind: int + spin: str # Should be "a" or "b" only + + def __attrs_post_init__(self): + if self.spin not in ["a", "b"]: + raise ValueError( + f'spin is set to {self.spin}, it should be either "a" or "b".' + ) + + @classmethod + def _json_namespace_(cls): + return "recirq.qcqmc" + + def _json_dict_(self): + # return cirq.dataclass_json_dict(self) + return attrs.asdict(self) + + @property + def openfermion_standard_index(self) -> int: + return 2 * self.orb_ind + (self.spin == "b") + + +@attrs.frozen +class LayerSpec: + """A specification of a hardware-efficient layer of gates. + + Args: + base_gate: 'charge_charge' for the e^{i n_i n_j} gate and 'givens' for a givens rotation. + layout: Should be 'in_pair', 'cross_pair', or 'cross_spin' only. + """ + + base_gate: str + layout: str + + def __post_init__(self): + if self.base_gate not in ["charge_charge", "givens"]: + raise ValueError( + f'base_gate is set to {self.base_gate}, it should be either "charge_charge or "givens".' + ) + if self.layout not in ["in_pair", "cross_pair", "cross_spin"]: + raise ValueError( + f'layout is set to {self.layout}, it should be either "cross_pair", "in_pair", or "cross_spin".' + ) + + @classmethod + def _json_namespace_(self): + return "recirq.qcqmc" + + def _json_dict_(self): + # return cirq.dataclass_json_dict(self) + return attrs.asdict(self) + + +@attrs.frozen +class TrialWavefunctionParams(Params, metaclass=abc.ABCMeta): + name: str + hamiltonian_params: HamiltonianParams + + @property + def bitstrings(self) -> Iterable[Tuple[bool, ...]]: + raise NotImplementedError( + "TrialWavefunctionParams should be subclassed and this method should be overwritten." + ) + + @property + def qubits_jordan_wigner_ordered(self) -> Tuple[cirq.GridQubit, ...]: + raise NotImplementedError( + "TrialWavefunctionParams should be subclassed and this method should be overwritten." + ) + + @property + def qubits_linearly_connected(self) -> Tuple[cirq.GridQubit, ...]: + raise NotImplementedError( + "TrialWavefunctionParams should be subclassed and this method should be overwritten." + ) + + +def _to_numpy(x: Optional[Iterable] = None) -> Optional[np.ndarray]: + return np.asarray(x) + + +def _to_tuple(x: Iterable[LayerSpec]) -> Sequence[LayerSpec]: + return tuple(x) + + +@attrs.frozen(repr=False) +class PerfectPairingPlusTrialWavefunctionParams(TrialWavefunctionParams): + """Class for storing the parameters that specify a TrialWavefunctionData. + + This class specifically stores the parameters for a trial wavefunction that + is a combination of a perfect pairing wavefunction with some number of + hardware-efficient layers appended. + """ + + name: str + hamiltonian_params: HamiltonianParams + heuristic_layers: Tuple[LayerSpec, ...] = attrs.field(converter=_to_tuple) + do_pp: bool = True + restricted: bool = False + random_parameter_scale: float = 1.0 + n_optimization_restarts: int = 1 + seed: int = 0 + initial_orbital_rotation: Optional[np.ndarray] = attrs.field( + default=None, converter=lambda v: _to_numpy(v) if v is not None else None + ) + initial_two_body_qchem_amplitudes: Optional[np.ndarray] = attrs.field( + default=None, converter=lambda v: _to_numpy(v) if v is not None else None + ) + do_optimization: bool = True + use_fast_gradients: bool = False + + @property + def n_orb(self) -> int: + return self.hamiltonian_params.n_orb + + @property + def n_elec(self) -> int: + return self.hamiltonian_params.n_elec + + @property + def n_qubits(self) -> int: + return 2 * self.n_orb + + @property + def n_pairs(self) -> int: + return self.n_elec // 2 + + @property + def path_string(self) -> str: + return OUTDIRS.DEFAULT_TRIAL_WAVEFUNCTION_DIRECTORY + self.name + + @property + def bitstrings(self) -> Iterable[Tuple[bool, ...]]: + return _get_bitstrings_a_b(n_orb=self.n_orb, n_elec=self.n_elec) + + def _json_dict_(self): + # return cirq.dataclass_json_dict(self) + return attrs.asdict(self) + + @property + def qubits_jordan_wigner_ordered(self) -> Tuple[cirq.GridQubit, ...]: + return _get_qubits_a_b(n_orb=self.n_orb) + + @property + def qubits_linearly_connected(self) -> Tuple[cirq.GridQubit, ...]: + return _get_qubits_a_b_reversed(n_orb=self.n_orb) + + @property + def mode_qubit_map(self) -> Dict[FermionicMode, cirq.GridQubit]: + return _get_mode_qubit_map_pp_plus(n_qubits=self.n_qubits) + + +def _get_qubits_a_b(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: + # This ordering creates qubits to facilitate a Jordan Wigner string + # threading through a row of alpha orbitals in ascending order followed by a + # row of beta orbitals in ascending order. + return tuple( + [cirq.GridQubit(0, i) for i in range(n_orb)] + + [cirq.GridQubit(1, i) for i in range(n_orb)] + ) + + +def _get_qubits_a_b_reversed(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: + # This ordering creates qubits to facilitate operations that need a linearly + # connected array of qubits with the order threading through a row of alpha + # orbitals in ascending order followed by a row of beta orbitals in + # descending order. + return tuple( + [cirq.GridQubit(0, i) for i in range(n_orb)] + + [cirq.GridQubit(1, i) for i in reversed(range(n_orb))] + ) + + +def _get_bitstrings_a_b(*, n_orb: int, n_elec: int) -> Iterable[Tuple[bool, ...]]: + """Iterates over bitstrings with the right symmetry assuming a_b ordering. + + This function assumes that the first n_orb qubits correspond to the alpha + orbitals and the second n_orb qubits correspond to the beta orbitals. The + ordering within the alpha and beta sectors doesn't matter (because we + iterate over all bitstrings with Hamming weight n_elec//2 in each sector. + """ + + if n_orb != n_elec: + raise NotImplementedError("n_orb must equal n_elec.") + + initial_bitstring = tuple(False for _ in range(n_orb - n_elec // 2)) + tuple( + True for _ in range(n_elec // 2) + ) + + spin_sector_bitstrings = set() + for perm in itertools.permutations(initial_bitstring): + spin_sector_bitstrings.add(perm) + + for bitstring_a, bitstring_b in itertools.product(spin_sector_bitstrings, repeat=2): + yield bitstring_a + bitstring_b + + +def _get_fermion_qubit_map_pp_plus(*, n_qubits: int) -> Dict[int, cirq.GridQubit]: + if n_qubits == 4: + return get_4_qubit_fermion_qubit_map() + elif n_qubits == 8: + return get_8_qubit_fermion_qubit_map() + elif n_qubits == 12: + return get_12_qubit_fermion_qubit_map() + elif n_qubits == 16: + return get_16_qubit_fermion_qubit_map() + else: + raise NotImplementedError() + + +def _get_mode_qubit_map_pp_plus( + *, n_qubits: int +) -> Dict[FermionicMode, cirq.GridQubit]: + """A map from Fermionic modes to qubits for our particular circuits. + + This function dispatches to _get_fermion_qubit_map_pp_plus, and ultimately + to get_X_qubit_fermion_qubit_map for specific values of X but it translates + this logic to a new system that uses the FermionicMode class rather than + opaque combinations of integers and strings. + """ + old_fermion_qubit_map = _get_fermion_qubit_map_pp_plus(n_qubits=n_qubits) + + n_orb = n_qubits // 2 + + mode_qubit_map = {} + + for i in range(n_orb): + mode_qubit_map[FermionicMode(i, "a")] = old_fermion_qubit_map[2 * i] + mode_qubit_map[FermionicMode(i, "b")] = old_fermion_qubit_map[2 * i + 1] + + return mode_qubit_map + + +def _get_reorder_func( + *, + mode_qubit_map: Mapping[FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> Callable[[int, int], int]: + """This is a helper function that allows us to reorder fermionic modes. + + Under the Jordan-Wigner transform, each fermionic mode is assigned to a + qubit. If we are provided an openfermion FermionOperator with the modes + assigned to qubits as described by mode_qubit_map this function gives us a + reorder_func that we can use to reorder the modes (with + openfermion.reorder(...)) so that they match the order of the qubits in + ordered_qubits. This is necessary to make a correspondence between + fermionic operators / wavefunctions and their qubit counterparts. + + Args: + mode_qubit_map: A dict that shows how each FermionicMode is mapped to a qubit. + + ordered_qubits: An ordered sequence of qubits. + """ + qubits = list(mode_qubit_map.values()) + assert len(qubits) == len(ordered_qubits) + + # We sort the key: value pairs by the order of the values (qubits) in + # ordered_qubits. + sorted_mapping = list(mode_qubit_map.items()) + sorted_mapping.sort(key=lambda x: ordered_qubits.index(x[1])) + + remapping_map = {} + for i, (mode, _) in enumerate(sorted_mapping): + openfermion_index = 2 * mode.orb_ind + (0 if mode.spin == "a" else 1) + remapping_map[openfermion_index] = i + + print("remapping_map:") + print(remapping_map) + + def remapper(index: int, _: int) -> int: + """A function that maps from the old index to the new one. + + The _ argument is because it's expected by openfermion.reorder""" + return remapping_map[index] + + return remapper + + +def _get_pp_plus_gate_generators( + *, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...], do_pp: bool = True +) -> List[of.FermionOperator]: + heuristic_gate_generators = get_heuristic_gate_generators(n_elec, heuristic_layers) + if not do_pp: + return heuristic_gate_generators + + n_pairs = n_elec // 2 + pair_gate_generators = get_pair_hopping_gate_generators(n_pairs, n_elec) + return pair_gate_generators + heuristic_gate_generators + + +def _get_ansatz_qubit_wf( + *, ansatz_circuit: cirq.Circuit, ordered_qubits: Sequence[cirq.Qid] +): + return cirq.final_state_vector( + ansatz_circuit, qubit_order=list(ordered_qubits), dtype=np.complex128 + ) + + +def _get_superposition_wf( + *, superposition_circuit: cirq.Circuit, ordered_qubits: Sequence[cirq.Qid] +) -> np.ndarray: + return cirq.final_state_vector( + superposition_circuit, qubit_order=list(ordered_qubits), dtype=np.complex128 + ) + + +def _get_fqe_wavefunctions( + *, + one_body_params: np.ndarray, + two_body_params: np.ndarray, + n_orb: int, + n_elec: int, + heuristic_layers: Tuple[LayerSpec, ...], + do_pp: bool = True, + restricted: Optional[bool] = False, + initial_orbital_rotation: Optional[np.ndarray] = None, +) -> Tuple[FqeWavefunction, FqeWavefunction]: + initial_wf = fqe.Wavefunction([[n_elec, 0, n_orb]]) + initial_wf.set_wfn(strategy="hartree-fock") + + wf, unrotated_wf = get_evolved_wf( + one_body_params=one_body_params, + two_body_params=two_body_params, + wf=initial_wf, + gate_generators=_get_pp_plus_gate_generators( + n_elec=n_elec, heuristic_layers=heuristic_layers, do_pp=do_pp + ), + n_orb=n_orb, + restricted=restricted, + initial_orbital_rotation=initial_orbital_rotation, + ) + + return wf, unrotated_wf + + +@attrs.frozen +class TrialWavefunctionData(Data): + """Class for storing a trial wavefunction's data.""" + + params: PerfectPairingPlusTrialWavefunctionParams + ansatz_circuit: cirq.Circuit + superposition_circuit: cirq.Circuit + hf_energy: float + ansatz_energy: float + fci_energy: float + one_body_basis_change_mat: np.ndarray = attrs.field(converter=_to_numpy) + one_body_params: np.ndarray = attrs.field(converter=_to_numpy) + two_body_params: np.ndarray = attrs.field(converter=_to_numpy) + + def _json_dict_(self): + # return cirq.dataclass_json_dict(self) + return attrs.asdict(self) + + +def get_and_check_energy( + *, + hamiltonian_data: HamiltonianData, + ansatz_circuit: cirq.Circuit, + one_body_params: np.ndarray, + two_body_params: np.ndarray, + one_body_basis_change_mat: np.ndarray, + params: PerfectPairingPlusTrialWavefunctionParams, +) -> Tuple[float, float]: + ansatz_qubit_wf = _get_ansatz_qubit_wf( + ansatz_circuit=ansatz_circuit, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + fqe_ham, e_core, sparse_ham = get_rotated_hamiltonians( + hamiltonian_data=hamiltonian_data, + one_body_basis_change_mat=one_body_basis_change_mat, + mode_qubit_map=params.mode_qubit_map, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + initial_wf = fqe.Wavefunction([[params.n_elec, 0, params.n_orb]]) + initial_wf.set_wfn(strategy="hartree-fock") + + hf_energy = initial_wf.expectationValue(fqe_ham) + e_core + + fqe_wf, unrotated_fqe_wf = get_evolved_wf( + one_body_params=one_body_params, + two_body_params=two_body_params, + wf=initial_wf, + gate_generators=_get_pp_plus_gate_generators( + n_elec=params.n_elec, + heuristic_layers=params.heuristic_layers, + do_pp=params.do_pp, + ), + n_orb=params.n_orb, + restricted=params.restricted, + initial_orbital_rotation=params.initial_orbital_rotation, + ) + + ansatz_energy = get_energy_and_check_sanity( + circuit_wf=ansatz_qubit_wf, + fqe_wf=fqe_wf, + unrotated_fqe_wf=unrotated_fqe_wf, + fqe_ham=fqe_ham, + sparse_ham=sparse_ham, + e_core=e_core, + mode_qubit_map=params.mode_qubit_map, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + return ansatz_energy, hf_energy + + +def build_pp_plus_trial_wavefunction( + params: PerfectPairingPlusTrialWavefunctionParams, + *, + dependencies: Dict[Params, Data], + do_print: bool = False, +) -> TrialWavefunctionData: + """Builds a TrialWavefunctionData from a TrialWavefunctionParams""" + + if do_print: + print("Building Trial Wavefunction") + np.random.seed(params.seed) + hamiltonian_data = dependencies[params.hamiltonian_params] + assert isinstance(hamiltonian_data, HamiltonianData) + + assert ( + params.n_orb == params.n_elec + ) ## Necessary for perfect pairing wavefunction to make sense. + + if params.do_optimization: + ( + one_body_params, + two_body_params, + one_body_basis_change_mat, + ) = get_pp_plus_params( + hamiltonian_data=hamiltonian_data, + restricted=params.restricted, + random_parameter_scale=params.random_parameter_scale, + initial_orbital_rotation=params.initial_orbital_rotation, + heuristic_layers=params.heuristic_layers, + do_pp=params.do_pp, + n_optimization_restarts=params.n_optimization_restarts, + do_print=do_print, + use_fast_gradients=params.use_fast_gradients, + ) + else: + if ( + params.initial_two_body_qchem_amplitudes is None + or params.initial_orbital_rotation is not None + ): + raise NotImplementedError("TODO: Implement whatever isn't finished here.") + + n_one_body_params = params.n_orb * (params.n_orb - 1) + one_body_params = np.zeros(n_one_body_params) + one_body_basis_change_mat = np.diag(np.ones(params.n_orb * 2)) + two_body_params = get_two_body_params_from_qchem_amplitudes( + params.initial_two_body_qchem_amplitudes + ) + + (superposition_circuit, ansatz_circuit) = get_circuits( + two_body_params=two_body_params, + n_orb=params.n_orb, + n_elec=params.n_elec, + heuristic_layers=params.heuristic_layers, + ) + + ansatz_energy, hf_energy = get_and_check_energy( + hamiltonian_data=hamiltonian_data, + ansatz_circuit=ansatz_circuit, + params=params, + one_body_params=one_body_params, + two_body_params=two_body_params, + one_body_basis_change_mat=one_body_basis_change_mat, + ) + + return TrialWavefunctionData( + params=params, + ansatz_circuit=ansatz_circuit, + superposition_circuit=superposition_circuit, + hf_energy=hf_energy, + ansatz_energy=ansatz_energy, + fci_energy=hamiltonian_data.e_fci, + one_body_basis_change_mat=one_body_basis_change_mat, + one_body_params=one_body_params, + two_body_params=two_body_params, + ) + + +def get_rotated_hamiltonians( + *, + hamiltonian_data: HamiltonianData, + one_body_basis_change_mat: np.ndarray, + mode_qubit_map: Mapping[FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> Tuple[fqe.hamiltonians.hamiltonian.Hamiltonian, float, csc_matrix]: + """A helper method that gets the hamiltonians in the basis of the trial_wf. + + Returns: + The hamiltonian in FQE form, minus a constant energy shift. + The constant part of the Hamiltonian missing from the FQE Hamiltonian. + The qubit Hamiltonian as a sparse matrix. + """ + n_qubits = len(mode_qubit_map) + + fqe_ham = hamiltonian_data.get_restricted_fqe_hamiltonian() + e_core = hamiltonian_data.e_core + + mol_ham = hamiltonian_data.get_molecular_hamiltonian() + mol_ham.rotate_basis(one_body_basis_change_mat) + fermion_operator_ham = of.get_fermion_operator(mol_ham) + + reorder_func = _get_reorder_func( + mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits + ) + fermion_operator_ham_qubit_ordered = of.reorder( + fermion_operator_ham, reorder_func, num_modes=n_qubits + ) + + sparse_qubit_ham = of.get_sparse_operator(fermion_operator_ham_qubit_ordered) + + return fqe_ham, e_core, sparse_qubit_ham + + +def get_energy_and_check_sanity( + *, + circuit_wf: np.ndarray, + fqe_wf: FqeWavefunction, + unrotated_fqe_wf: FqeWavefunction, + fqe_ham: fqe.hamiltonians.hamiltonian.Hamiltonian, + sparse_ham: csc_matrix, + e_core: float, + mode_qubit_map: Mapping[FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> float: + """A method that checks for consistency and returns the ansatz energy.""" + + unrotated_fqe_wf_as_cirq = convert_fqe_wf_to_cirq( + fqe_wf=unrotated_fqe_wf, + mode_qubit_map=mode_qubit_map, + ordered_qubits=ordered_qubits, + ) + ansatz_energy = np.real_if_close( + (np.conj(circuit_wf) @ sparse_ham @ circuit_wf) + ).item() + assert isinstance(ansatz_energy, float) + + fqe_energy = np.real(fqe_wf.expectationValue(fqe_ham) + e_core) + print(csc_matrix(circuit_wf)) + print(csc_matrix(unrotated_fqe_wf_as_cirq)) + np.testing.assert_array_almost_equal(ansatz_energy, fqe_energy) + np.testing.assert_array_almost_equal( + circuit_wf, unrotated_fqe_wf_as_cirq, decimal=5 + ) + return ansatz_energy + + +def get_4_qubit_fermion_qubit_map() -> Dict[int, cirq.GridQubit]: + """A helper function that provides the fermion qubit map for 4 qubits. + + We map the fermionic orbitals to grid qubits like so: + 3 1 + 2 0 + """ + fermion_index_to_qubit_map = { + 2: cirq.GridQubit(0, 0), + 3: cirq.GridQubit(1, 0), + 0: cirq.GridQubit(0, 1), + 1: cirq.GridQubit(1, 1), + } + + return fermion_index_to_qubit_map + + +def get_4_qubit_pp_circuits( + *, two_body_params: np.ndarray, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...] +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 3 1 + 2 0 + """ + assert n_elec == 2 + + fermion_index_to_qubit_map = get_4_qubit_fermion_qubit_map() + geminal_gate = GeminalStatePreparationGate(two_body_params[0], indicator=True) + + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + ) + ) + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[1:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + indicator = fermion_index_to_qubit_map[2] + superposition_circuit = cirq.Circuit([cirq.H(indicator) + ansatz_circuit]) + ansatz_circuit = cirq.Circuit([cirq.X(indicator) + ansatz_circuit]) + + return superposition_circuit, ansatz_circuit + + +def get_8_qubit_fermion_qubit_map(): + """A helper function that provides the fermion qubit map for 8 qubits. + + We map the fermionic orbitals to grid qubits like so: + 3 5 1 7 + 2 4 0 6 + """ + + # Linear connectivity is fine. + # This ordering is dictated by the way we specify perfect pairing (i) + # Elsewhere we generate the perfect pairing parameters using a specific + # convention for how we index the FermionOperators (which is itself) + # partly dictated by the OpenFermion conventions. Here we choose a + # mapping between the indices of these FermionOperators and the qubits in our + # grid that allows for the perfect pairing pairs to be in squares of four. + fermion_index_to_qubit_map = { + 2: cirq.GridQubit(0, 0), + 3: cirq.GridQubit(1, 0), + 4: cirq.GridQubit(0, 1), + 5: cirq.GridQubit(1, 1), + 0: cirq.GridQubit(0, 2), + 1: cirq.GridQubit(1, 2), + 6: cirq.GridQubit(0, 3), + 7: cirq.GridQubit(1, 3), + } + + return fermion_index_to_qubit_map + + +def get_8_qubit_circuits( + *, two_body_params: np.ndarray, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...] +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 3 5 1 7 + 2 4 0 6 + """ + fermion_index_to_qubit_map = get_8_qubit_fermion_qubit_map() + + geminal_gate_1 = GeminalStatePreparationGate(two_body_params[0], indicator=True) + geminal_gate_2 = GeminalStatePreparationGate(two_body_params[1], indicator=True) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[2:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[0]), + cirq.CNOT(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[6]), + cirq.SWAP(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[4]), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[4]), + cirq.X(fermion_index_to_qubit_map[6]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_12_qubit_fermion_qubit_map(): + """A helper function that provides the fermion qubit map for 12 qubits. + + We map the fermionic orbitals to grid qubits like so: + 5 7 3 9 1 11 + 4 6 2 8 0 10 + """ + + # Linear connectivity is fine. + # This ordering is dictated by the way we specify perfect pairing (i) + # Elsewhere we generate the perfect pairing parameters using a specific + # convention for how we index the FermionOperators (which is itself) + # partly dictated by the OpenFermion conventions. Here we choose a + # mapping between the indices of these FermionOperators and the qubits in our + # grid that allows for the perfect pairing pairs to be in squares of four. + fermion_index_to_qubit_map = { + 4: cirq.GridQubit(0, 0), + 5: cirq.GridQubit(1, 0), + 6: cirq.GridQubit(0, 1), + 7: cirq.GridQubit(1, 1), + 2: cirq.GridQubit(0, 2), + 3: cirq.GridQubit(1, 2), + 8: cirq.GridQubit(0, 3), + 9: cirq.GridQubit(1, 3), + 0: cirq.GridQubit(0, 4), + 1: cirq.GridQubit(1, 4), + 10: cirq.GridQubit(0, 5), + 11: cirq.GridQubit(1, 5), + } + + return fermion_index_to_qubit_map + + +def get_12_qubit_circuits( + *, two_body_params: np.ndarray, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...] +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 5 7 3 9 1 11 + 4 6 2 8 0 10 + """ + + fermion_index_to_qubit_map = get_12_qubit_fermion_qubit_map() + + geminal_gate_1 = GeminalStatePreparationGate(two_body_params[0], indicator=True) + geminal_gate_2 = GeminalStatePreparationGate(two_body_params[1], indicator=True) + geminal_gate_3 = GeminalStatePreparationGate(two_body_params[2], indicator=True) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[8], + fermion_index_to_qubit_map[9], + ) + ), + cirq.decompose( + geminal_gate_3.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[10], + fermion_index_to_qubit_map[11], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[3:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[8]), + cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[0]), + cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[2]), + cirq.SWAP( + fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[10] + ), + cirq.SWAP(fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[6]), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[6]), + cirq.X(fermion_index_to_qubit_map[8]), + cirq.X(fermion_index_to_qubit_map[10]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_16_qubit_fermion_qubit_map(): + """A helper function that provides the fermion qubit map for 16 qubits. + + We map the fermionic orbitals to grid qubits like so: + 7 9 5 11 3 13 1 15 + 6 8 4 10 2 12 0 14 + """ + + # Linear connectivity is fine. + # This ordering is dictated by the way we specify perfect pairing (i) + # Elsewhere we generate the perfect pairing parameters using a specific + # convention for how we index the FermionOperators (which is itself) + # partly dictated by the OpenFermion conventions. Here we choose a + # mapping between the indices of these FermionOperators and the qubits in our + # grid that allows for the perfect pairing pairs to be in squares of four. + fermion_index_to_qubit_map = { + 6: cirq.GridQubit(0, 0), + 7: cirq.GridQubit(1, 0), + 8: cirq.GridQubit(0, 1), + 9: cirq.GridQubit(1, 1), + 4: cirq.GridQubit(0, 2), + 5: cirq.GridQubit(1, 2), + 10: cirq.GridQubit(0, 3), + 11: cirq.GridQubit(1, 3), + 2: cirq.GridQubit(0, 4), + 3: cirq.GridQubit(1, 4), + 12: cirq.GridQubit(0, 5), + 13: cirq.GridQubit(1, 5), + 0: cirq.GridQubit(0, 6), + 1: cirq.GridQubit(1, 6), + 14: cirq.GridQubit(0, 7), + 15: cirq.GridQubit(1, 7), + } + + return fermion_index_to_qubit_map + + +def get_16_qubit_circuits( + *, two_body_params: np.ndarray, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...] +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 7 9 5 11 3 13 1 15 + 6 8 4 10 2 12 0 14 + """ + fermion_index_to_qubit_map = get_16_qubit_fermion_qubit_map() + + geminal_gate_1 = GeminalStatePreparationGate(two_body_params[0], indicator=True) + geminal_gate_2 = GeminalStatePreparationGate(two_body_params[1], indicator=True) + geminal_gate_3 = GeminalStatePreparationGate(two_body_params[2], indicator=True) + geminal_gate_4 = GeminalStatePreparationGate(two_body_params[3], indicator=True) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + fermion_index_to_qubit_map[8], + fermion_index_to_qubit_map[9], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + fermion_index_to_qubit_map[10], + fermion_index_to_qubit_map[11], + ) + ), + cirq.decompose( + geminal_gate_3.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[12], + fermion_index_to_qubit_map[13], + ) + ), + cirq.decompose( + geminal_gate_4.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[14], + fermion_index_to_qubit_map[15], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[4:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[10]), + cirq.CNOT( + fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[2] + ), + cirq.SWAP( + fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[12] + ), + cirq.CNOT( + fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[4] + ), + cirq.CNOT( + fermion_index_to_qubit_map[12], fermion_index_to_qubit_map[0] + ), + cirq.SWAP(fermion_index_to_qubit_map[4], fermion_index_to_qubit_map[8]), + cirq.SWAP( + fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[14] + ), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[8]), + cirq.X(fermion_index_to_qubit_map[10]), + cirq.X(fermion_index_to_qubit_map[12]), + cirq.X(fermion_index_to_qubit_map[14]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_circuits( + *, + two_body_params: np.ndarray, + # from wf_params: + n_orb: int, + n_elec: int, + heuristic_layers: Tuple[LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A function that runs a specialized method to get the ansatz circuits.""" + + # TODO(?): Just input one of these quantities. + if n_orb != n_elec: + raise ValueError("n_orb must equal n_elec.") + + circ_funcs = { + 2: get_4_qubit_pp_circuits, + 4: get_8_qubit_circuits, + 6: get_12_qubit_circuits, + 8: get_16_qubit_circuits, + } + try: + circ_func = circ_funcs[n_orb] + except KeyError: + raise NotImplementedError(f"No circuits for n_orb = {n_orb}") + + return circ_func( + two_body_params=two_body_params, + n_elec=n_elec, + heuristic_layers=heuristic_layers, + ) + + +def get_two_body_params_from_qchem_amplitudes( + qchem_amplitudes: np.ndarray, +) -> np.ndarray: + """Translates perfect pairing amplitudes from qchem to rotation angles. + + qchem style: 1 |1100> + t_i |0011> + our style: cos(\theta_i) |1100> + sin(\theta_i) |0011> + """ + + two_body_params = np.arccos(1 / np.sqrt(1 + qchem_amplitudes**2)) * np.sign( + qchem_amplitudes + ) + + # Numpy casts the array improperly to a float when we only have one parameter. + two_body_params = np.atleast_1d(two_body_params) + + return two_body_params + + +#################### Here be dragons.########################################### + + +def convert_fqe_wf_to_cirq( + fqe_wf: FqeWavefunction, + mode_qubit_map: Mapping[FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> np.ndarray: + """Converts an FQE wavefunction to one on qubits with a particular ordering.""" + n_qubits = len(mode_qubit_map) + fermion_op = fqe.openfermion_utils.fqe_to_fermion_operator(fqe_wf) + + reorder_func = _get_reorder_func( + mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits + ) + fermion_op = of.reorder(fermion_op, reorder_func, num_modes=n_qubits) + + qubit_op = of.jordan_wigner(fermion_op) + + return fqe.qubit_wavefunction_from_vacuum( + qubit_op, list(cirq.LineQubit.range(n_qubits)) + ) + + +def get_one_body_cluster_coef( + params: np.ndarray, n_orb: int, restricted: bool +) -> npt.NDArray[np.complex128]: + if restricted: + one_body_cluster_op = np.zeros((n_orb, n_orb), dtype=np.complex128) + else: + one_body_cluster_op = np.zeros((2 * n_orb, 2 * n_orb), dtype=np.complex128) + param_num = 0 + + for i in range(n_orb): + for j in range(i): + one_body_cluster_op[i, j] = params[param_num] + one_body_cluster_op[j, i] = -params[param_num] + param_num += 1 + + if not restricted: + for i in range(n_orb, 2 * n_orb): + for j in range(n_orb, i): + one_body_cluster_op[i, j] = params[param_num] + one_body_cluster_op[j, i] = -params[param_num] + param_num += 1 + + return one_body_cluster_op + + +def get_evolved_wf( + one_body_params, + two_body_params, + wf, + gate_generators, + n_orb, + restricted, + initial_orbital_rotation: Optional[np.ndarray] = None, +): + param_num = 0 + for gate_generator in gate_generators: + wf = wf.time_evolve(two_body_params[param_num], gate_generator) + param_num += 1 + + one_body_cluster_op = get_one_body_cluster_coef( + one_body_params, n_orb, restricted=restricted + ) + + if restricted: + one_body_ham = fqe.get_restricted_hamiltonian((1j * one_body_cluster_op,)) + else: + one_body_ham = fqe.get_sso_hamiltonian((1j * one_body_cluster_op,)) + + rotated_wf = wf.time_evolve(1.0, one_body_ham) + + if initial_orbital_rotation is not None: + rotated_wf = fqe.algorithm.low_rank.evolve_fqe_givens( + rotated_wf, initial_orbital_rotation + ) + + return rotated_wf, wf + + +def get_pair_hopping_gate_generators(n_pairs, n_elec): + gate_generators = [] + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + + fop_string = "{:d} {:d} {:d}^ {:d}^".format(to_b, to_a, from_b, from_a) + + gate_generator = of.FermionOperator(fop_string, 1.0) + gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) + gate_generators.append(gate_generator) + + return gate_generators + + +def get_indices_heuristic_layer_in_pair(n_elec) -> Iterator[Tuple[int, int]]: + # Indices that couple within a pair + n_pairs = n_elec // 2 + + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + yield (from_a, to_a) + yield (from_b, to_b) + + +def get_indices_heuristic_layer_cross_pair(n_elec) -> Iterator[Tuple[int, int]]: + # Indices that couple adjacent pairs + n_pairs = n_elec // 2 + + for pair in range(n_pairs - 1): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_next_a = n_elec - 2 * (pair + 1) - 2 + from_next_b = n_elec - 2 * (pair + 1) - 1 + yield (to_a, from_next_a) + yield (to_b, from_next_b) + + +def get_indices_heuristic_layer_cross_spin(n_elec) -> Iterator[Tuple[int, int]]: + # Indices that couple the two spin sectors + n_pairs = n_elec // 2 + + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + yield (to_a, to_b) + yield (from_a, from_b) + + +def get_charge_charge_generator(indices: Tuple[int, int]) -> of.FermionOperator: + # Returns the generator for density evolution between the indices + + fop_string = "{:d}^ {:d} {:d}^ {:d}".format( + indices[0], indices[0], indices[1], indices[1] + ) + gate_generator = of.FermionOperator(fop_string, 1.0) + + return gate_generator + + +def get_charge_charge_gate( + qubits: Tuple[cirq.Qid, ...], param: float +) -> cirq.Operation: + return cirq.CZ(qubits[0], qubits[1]) ** (-param / np.pi) + + +def get_givens_generator(indices): + # Returns the generator for density evolution between the indices + + fop_string = "{:d}^ {:d}".format(indices[0], indices[1]) + gate_generator = of.FermionOperator(fop_string, 1.0) + gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) + + return gate_generator + + +def get_givens_gate(qubits: Tuple[cirq.Qid, ...], param: float) -> cirq.Operation: + return cirq.givens(param).on(qubits[0], qubits[1]) + + +def get_layer_indices(layer_spec: LayerSpec, n_elec: int) -> List[Tuple[int, int]]: + indices_generators = { + "in_pair": get_indices_heuristic_layer_in_pair(n_elec), + "cross_pair": get_indices_heuristic_layer_cross_pair(n_elec), + "cross_spin": get_indices_heuristic_layer_cross_spin(n_elec), + } + indices_generator = indices_generators[layer_spec.layout] + + return [indices for indices in indices_generator] + + +def get_layer_gates( + layer_spec: LayerSpec, + n_elec: int, + params: np.ndarray, + fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], +) -> List[cirq.Operation]: + """Gets the gates for a hardware efficient layer of the ansatz.""" + + indices_list = get_layer_indices(layer_spec, n_elec) + + gate_funcs = {"givens": get_givens_gate, "charge_charge": get_charge_charge_gate} + gate_func = gate_funcs[layer_spec.base_gate] + + gates = [] + for indices, param in zip(indices_list, params): + qubits = tuple(fermion_index_to_qubit_map[ind] for ind in indices) + gates.append(gate_func(qubits, param)) + + return gates + + +def get_layer_generators(layer_spec: LayerSpec, n_elec: int): + """Gets the generators for rotations in a hardware efficient layer of the ansatz.""" + + indices_list = get_layer_indices(layer_spec, n_elec) + + gate_funcs = { + "givens": get_givens_generator, + "charge_charge": get_charge_charge_generator, + } + gate_func = gate_funcs[layer_spec.base_gate] + + return [gate_func(indices) for indices in indices_list] + + +def get_heuristic_gate_generators(n_elec: int, layer_specs: Sequence[LayerSpec]): + gate_generators = [] + + for layer_spec in layer_specs: + gate_generators += get_layer_generators(layer_spec, n_elec) + + return gate_generators + + +def get_heuristic_circuit( + layer_specs: Sequence[LayerSpec], + n_elec: int, + params: np.ndarray, + fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], +) -> cirq.Circuit: + gates: List[cirq.Operation] = [] + + for layer_spec in layer_specs: + params_slice = params[len(gates) :] + gates += get_layer_gates( + layer_spec, n_elec, params_slice, fermion_index_to_qubit_map + ) + + return cirq.Circuit(gates) + + +def orbital_rotation_gradient_matrix(generator_mat, a, b): + """The gradient of the orbital rotation unitary with respect to its parameters. + + Args: + generator_mat: The orbital rotation one-body generator matrix. + a, b: row and column indices corresponding to the location in the matrix + of the parameter we wish to find the gradient with respect to. + Returns: + The orbital rotation matrix gradient wrt theta_{a, b}. Corresponds to + expression in G15 of https://arxiv.org/abs/2004.04174. + """ + w_full, v_full = np.linalg.eigh(-1j * generator_mat) + eigs_diff = np.zeros((w_full.shape[0], w_full.shape[0]), dtype=np.complex128) + for i, j in itertools.product(range(w_full.shape[0]), repeat=2): + if np.isclose(abs(w_full[i] - w_full[j]), 0): + eigs_diff[i, j] = 1 + else: + eigs_diff[i, j] = (np.exp(1j * (w_full[i] - w_full[j])) - 1) / ( + 1j * (w_full[i] - w_full[j]) + ) + + Y_full = np.zeros_like(v_full, dtype=np.complex128) + if a == b: + Y_full[a, b] = 0 + else: + Y_full[a, b] = 1.0 + Y_full[b, a] = -1.0 + + Y_kl_full = v_full.conj().T @ Y_full @ v_full + # now rotate Y_{kl} * (exp(i(l_{k} - l_{l})) - 1) / (i(l_{k} - l_{l})) + # into the original basis + pre_matrix_full = v_full @ (eigs_diff * Y_kl_full) @ v_full.conj().T + + return pre_matrix_full + + +def evaluate_gradient_and_cost_function( + initial_wf: FqeWavefunction, + fqe_ham: RestrictedHamiltonian, + n_orb: int, + one_body_params: npt.NDArray[np.float64], + two_body_params: npt.NDArray[np.float64], + gate_generators: List[of.FermionOperator], + restricted: bool, + e_core: float, +) -> Tuple[float, npt.NDArray[np.float64]]: + """Evaluate gradient and cost function for optimization. + + Args: + initial_wf: Initial state (typically Hartree--Fock). + fqe_ham: The restricted Hamiltonian in FQE format. + n_orb: The number of spatial orbitals. + one_body_params: The parameters of the single-particle rotations. + two_body_params: The parameters for the two-particle terms. + gate_generators: The generators for the two-particle terms. + retricted: Whether the single-particle rotations are restricted (True) + or unrestricted (False). Unrestricted implies different parameters + for the alpha- and beta-spin rotations. + Returns: + cost_val: The cost function (total energy) evaluated for the input wavefunction parameters. + grad: An array of gradients with respect to the one- and two-body + parameters. The first n_orb * (n_orb + 1) // 2 parameters correspond to + the one-body gradients. + """ + phi = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + )[0] + lam = copy.deepcopy(phi) + lam = lam.apply(fqe_ham) + cost_val = fqe.vdot(lam, phi) + e_core + + # 1body + one_body_cluster_op = get_one_body_cluster_coef( + one_body_params, n_orb, restricted=restricted + ) + tril = np.tril_indices(n_orb, k=-1) + if restricted: + one_body_ham = fqe.get_restricted_hamiltonian((-1j * one_body_cluster_op,)) + else: + one_body_ham = fqe.get_sso_hamiltonian((-1j * one_body_cluster_op,)) + # Apply U1b^{dag} + phi.time_evolve(1, one_body_ham, inplace=True) + lam.time_evolve(1, one_body_ham, inplace=True) + one_body_grad = np.zeros_like(one_body_params) + n_one_body_params = len(one_body_params) + grad_position = n_one_body_params - 1 + for iparam in range(len(one_body_params)): + mu_state = copy.deepcopy(phi) + pidx = n_one_body_params - iparam - 1 + pidx_spin = 0 if restricted else pidx // (n_one_body_params // 2) + pidx_spat = pidx if restricted else pidx - (n_one_body_params // 2) * pidx_spin + p, q = (tril[0][pidx_spat], tril[1][pidx_spat]) + p += n_orb * pidx_spin + q += n_orb * pidx_spin + pre_matrix = orbital_rotation_gradient_matrix(-one_body_cluster_op, p, q) + assert of.is_hermitian(1j * pre_matrix) + if restricted: + fqe_quad_ham_pre = fqe.get_restricted_hamiltonian((pre_matrix,)) + else: + fqe_quad_ham_pre = fqe.get_sso_hamiltonian((pre_matrix,)) + mu_state = mu_state.apply(fqe_quad_ham_pre) + one_body_grad[grad_position] = 2 * fqe.vdot(lam, mu_state).real + grad_position -= 1 + # Get two-body contributions + two_body_grad = np.zeros(len(two_body_params)) + for pidx in reversed(range(len(gate_generators))): + mu = copy.deepcopy(phi) + mu = mu.apply(gate_generators[pidx]) + two_body_grad[pidx] = -np.real(2 * 1j * (fqe.vdot(lam, mu))) + phi = phi.time_evolve(-two_body_params[pidx], gate_generators[pidx]) + lam = lam.time_evolve(-two_body_params[pidx], gate_generators[pidx]) + + return cost_val, np.concatenate((two_body_grad, one_body_grad)) + + +def get_pp_plus_params( + *, + hamiltonian_data: HamiltonianData, + restricted: bool = False, + random_parameter_scale: float = 1.0, + initial_orbital_rotation: Optional[np.ndarray] = None, + heuristic_layers: Tuple[LayerSpec, ...], + do_pp: bool = True, + n_optimization_restarts: int = 1, + do_print: bool = True, + use_fast_gradients: bool = False, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + n_elec = hamiltonian_data.params.n_elec + n_orb = hamiltonian_data.params.n_orb + sz = 0 + + initial_wf = fqe.Wavefunction([[n_elec, sz, n_orb]]) + initial_wf.set_wfn(strategy="hartree-fock") + + fqe_ham = hamiltonian_data.get_restricted_fqe_hamiltonian() + e_core = hamiltonian_data.e_core + + hf_energy = initial_wf.expectationValue(fqe_ham) + e_core + + # We're only supporting closed shell stuff here. + assert n_elec % 2 == 0 + assert n_elec <= n_orb + if use_fast_gradients: + err_msg = "use_fast_gradients does not work with initial orbital rotation." + assert initial_orbital_rotation is None, err_msg + + gate_generators = _get_pp_plus_gate_generators( + n_elec=n_elec, heuristic_layers=heuristic_layers, do_pp=do_pp + ) + + n_two_body_params = len(gate_generators) + + if restricted: + n_one_body_params = n_orb * (n_orb - 1) // 2 + else: + n_one_body_params = n_orb * (n_orb - 1) + + best = np.inf + best_res: Union[None, scipy.optimize.OptimizeResult] = None + for i in range(n_optimization_restarts): + if do_print: + print(f"Optimization restart {i}", flush=True) + + def progress_cb(_): + print(".", end="", flush=True) + + else: + + def progress_cb(_): + pass + + params = random_parameter_scale * np.random.normal( + size=(n_two_body_params + n_one_body_params) + ) + + def objective(params): + one_body_params = params[-n_one_body_params:] + two_body_params = params[:n_two_body_params] + + wf, _ = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + initial_orbital_rotation=initial_orbital_rotation, + ) + + energy = wf.expectationValue(fqe_ham) + e_core + if do_print: + print(f"energy {energy}") + if np.abs(energy.imag) < 1e-6: + return energy.real + else: + return 1e6 + + def fast_obj_grad(params): + one_body_params = params[-n_one_body_params:] + two_body_params = params[:n_two_body_params] + energy, grad = evaluate_gradient_and_cost_function( + initial_wf, + fqe_ham, + n_orb, + one_body_params, + two_body_params, + gate_generators, + restricted, + e_core, + ) + if do_print: + print(f"energy {energy}, max|grad| {np.max(np.abs(grad))}") + if np.abs(energy.imag) < 1e-6: + return energy.real, grad + else: + return 1e6, 1e6 + + if use_fast_gradients: + res = minimize( + fast_obj_grad, params, jac=True, method="BFGS", callback=progress_cb + ) + else: + res = minimize(objective, params, callback=progress_cb) + if res.fun < best: + best = res.fun + best_res = res + + if do_print: + print(res, flush=True) + + assert best_res is not None + params = best_res.x + one_body_params = params[-n_one_body_params:] + two_body_params = params[:n_two_body_params] + + wf, _ = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + initial_orbital_rotation=initial_orbital_rotation, + ) + + one_body_cluster_mat = get_one_body_cluster_coef( + one_body_params, n_orb, restricted=restricted + ) + # We need to change the ordering to match OpenFermion's abababab ordering + if not restricted: + index_rearrangement = np.asarray( + [i // 2 % (n_orb) + (i % 2) * n_orb for i in range(2 * n_orb)] + ) + one_body_cluster_mat = one_body_cluster_mat[:, index_rearrangement] + one_body_cluster_mat = one_body_cluster_mat[index_rearrangement, :] + + one_body_basis_change_mat = expm(one_body_cluster_mat) + + if initial_orbital_rotation is not None: + if restricted: + one_body_basis_change_mat = ( + initial_orbital_rotation @ one_body_basis_change_mat + ) + else: + big_initial_orbital_rotation = np.zeros_like(one_body_basis_change_mat) + + for i in range(len(initial_orbital_rotation)): + for j in range(len(initial_orbital_rotation)): + big_initial_orbital_rotation[2 * i, 2 * j] = ( + initial_orbital_rotation[i, j] + ) + big_initial_orbital_rotation[2 * i + 1, 2 * j + 1] = ( + initial_orbital_rotation[i, j] + ) + + one_body_basis_change_mat = ( + big_initial_orbital_rotation @ one_body_basis_change_mat + ) + + if do_print: + print("Hartree-Fock Energy:") + print(hf_energy) + initial_wf.print_wfn() + print("-" * 80) + print("FCI Energy:") + print(hamiltonian_data.e_fci) + print("-" * 80) + print(best_res) + + print("-" * 80) + print("Ansatz Energy:") + print(np.real_if_close(wf.expectationValue(fqe_ham) + e_core)) + wf.print_wfn() + print("Basis Rotation Matrix:") + print(one_body_basis_change_mat) + print("Two Body Rotation Parameters:") + print(two_body_params) + + return one_body_params, two_body_params, one_body_basis_change_mat diff --git a/recirq/qcqmc/trial_wf_test.py b/recirq/qcqmc/trial_wf_test.py new file mode 100644 index 00000000..b2f64d72 --- /dev/null +++ b/recirq/qcqmc/trial_wf_test.py @@ -0,0 +1,387 @@ +import cirq +import fqe +import numpy as np +import pytest +import scipy.special +from fqe.hamiltonians.restricted_hamiltonian import RestrictedHamiltonian + +from recirq.qcqmc.hamiltonian import HamiltonianData +from recirq.qcqmc.trial_wf import (FermionicMode, LayerSpec, + PerfectPairingPlusTrialWavefunctionParams, + _get_ansatz_qubit_wf, _get_bitstrings_a_b, + _get_pp_plus_gate_generators, + build_pp_plus_trial_wavefunction, + evaluate_gradient_and_cost_function, + get_evolved_wf, + get_two_body_params_from_qchem_amplitudes) + + +def test_fermionic_mode(): + fm = FermionicMode(5, "a") + fm2 = cirq.read_json(json_text=cirq.to_json(fm)) + assert fm == fm2 + + with pytest.raises(ValueError, match="spin.*"): + _ = FermionicMode(10, "c") + + +def test_get_bitstrings_a_b(): + with pytest.raises(NotImplementedError): + list(_get_bitstrings_a_b(n_orb=4, n_elec=3)) + + bitstrings = np.array(list(_get_bitstrings_a_b(n_orb=4, n_elec=4))) + + assert bitstrings.shape[0] == scipy.special.binom(4, 2) ** 2 + assert bitstrings.shape[1] == 2 * 4 # n_qubits columns = 2 * n_orb. + hamming_weight_left = np.sum(bitstrings[:, 0:4], axis=1) + hamming_weight_right = np.sum(bitstrings[:, 4:8], axis=1) + + assert np.all(hamming_weight_left == 2) + assert np.all(hamming_weight_right == 2) + + +def test_pp_wf_energy(fixture_4_qubit_ham: HamiltonianData): + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_1", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=(), + do_pp=True, + restricted=True, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham} + ) + + assert np.isclose(trial_wf.ansatz_energy, fixture_4_qubit_ham.e_fci) + + +def test_pp_wf_energy_with_layer(fixture_4_qubit_ham: HamiltonianData): + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_2", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=(LayerSpec("charge_charge", "cross_spin"),), + do_pp=True, + restricted=True, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham} + ) + + assert np.isclose(trial_wf.ansatz_energy, fixture_4_qubit_ham.e_fci) + + +def test_qchem_pp_eight_qubit_wavefunctions_consistent( + fixture_8_qubit_ham: HamiltonianData, +): + """Tests (without optimization) that the eight qubit wavefunctions work. + + Specifically, that constructing the wavefunction with FQE and then + converting it to a cirq wavefunction yields the same result as constructing + the parameters with the circuit directly. + """ + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_qchem", + hamiltonian_params=fixture_8_qubit_ham.params, + heuristic_layers=tuple(), + initial_orbital_rotation=None, + initial_two_body_qchem_amplitudes=np.asarray([0.3, 0.4]), + do_optimization=False, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_8_qubit_ham.params: fixture_8_qubit_ham}, + do_print=False, + ) + + one_body_params = trial_wf.one_body_params + two_body_params = trial_wf.two_body_params + basis_change_mat = trial_wf.one_body_basis_change_mat + + np.testing.assert_array_almost_equal(one_body_params, np.zeros((12,))) + np.testing.assert_array_almost_equal(basis_change_mat, np.diag(np.ones(8))) + + np.testing.assert_equal(two_body_params.shape, (2,)) + + +def test_pp_plus_wf_energy_sloppy_1(fixture_8_qubit_ham: HamiltonianData): + params = PerfectPairingPlusTrialWavefunctionParams( + "pp_plus_test", + hamiltonian_params=fixture_8_qubit_ham.params, + heuristic_layers=tuple(), + do_pp=True, + restricted=False, + random_parameter_scale=1, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_8_qubit_ham.params: fixture_8_qubit_ham}, + do_print=True, + ) + + assert trial_wf.ansatz_energy < -1.947 + + +# TODO: Speed up this test and add a similar one with non-trivial heuristic layers. + + +def test_diamond_pp_wf_energy(fixture_12_qubit_ham: HamiltonianData): + params = PerfectPairingPlusTrialWavefunctionParams( + name="diamind_pp_test_wf_1", + hamiltonian_params=fixture_12_qubit_ham.params, + heuristic_layers=tuple(), + do_pp=True, + restricted=True, + random_parameter_scale=0.1, + n_optimization_restarts=1, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_12_qubit_ham.params: fixture_12_qubit_ham}, + do_print=True, + ) + + assert trial_wf.ansatz_energy < -10.4 + + +@pytest.mark.parametrize( + "initial_two_body_qchem_amplitudes, expected_ansatz_qubit_wf", + [ + ( + [1], + np.array( + [ + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.70710678 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.70710678 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + ] + ), + ), + ( + [0], + np.array( + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ] + ), + ), + ], +) +def test_qchem_pp_runs( + initial_two_body_qchem_amplitudes, + expected_ansatz_qubit_wf, + fixture_4_qubit_ham: HamiltonianData, +): + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_qchem", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=tuple(), + initial_orbital_rotation=None, + initial_two_body_qchem_amplitudes=initial_two_body_qchem_amplitudes, + do_optimization=False, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham}, + do_print=False, + ) + + ansatz_qubit_wf = _get_ansatz_qubit_wf( + ansatz_circuit=trial_wf.ansatz_circuit, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + np.testing.assert_array_almost_equal(ansatz_qubit_wf, expected_ansatz_qubit_wf) + + +def test_qchem_conversion_negative(fixture_4_qubit_ham: HamiltonianData): + qchem_amplitudes = np.asarray(-0.1) + + two_body_params = get_two_body_params_from_qchem_amplitudes(qchem_amplitudes) + + assert two_body_params.item() < 0 + + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_qchem_neg", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=tuple(), + initial_orbital_rotation=None, + initial_two_body_qchem_amplitudes=qchem_amplitudes, + do_optimization=False, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham}, + do_print=False, + ) + + ansatz_qubit_wf = _get_ansatz_qubit_wf( + ansatz_circuit=trial_wf.ansatz_circuit, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + assert any(ansatz_qubit_wf < 0) + + +def gen_random_restricted_ham(n_orb: int) -> RestrictedHamiltonian: + """8-fold symmetry restricted hamiltonian""" + h1e = np.random.random((n_orb,) * 2) + h1e = h1e + h1e.T + h2e = np.random.random((n_orb,) * 4) + h2e = h2e + h2e.transpose(2, 3, 0, 1) + h2e = h2e + h2e.transpose(3, 2, 1, 0) + h2e = h2e + h2e.transpose(1, 0, 2, 3) + h2e = np.asarray(h2e.transpose(0, 2, 3, 1), order="C") + fqe_ham = RestrictedHamiltonian((h1e, np.einsum("ijlk", -0.5 * h2e))) + return fqe_ham + + +def get_fd_grad( + n_orb, + n_elec, + one_body_params, + two_body_params, + ham, + initial_wf, + dtheta=1e-4, + restricted=False, +): + generators = _get_pp_plus_gate_generators( + n_elec=n_elec, heuristic_layers=tuple(), do_pp=True + ) + one_body_gradient = np.zeros_like(one_body_params) + for ig, _ in enumerate(one_body_gradient): + new_param = one_body_params.copy() + new_param[ig] = new_param[ig] + dtheta + phi = get_evolved_wf( + new_param, + two_body_params, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_plus = phi.expectationValue(ham) + new_param[ig] = new_param[ig] - 2 * dtheta + phi = get_evolved_wf( + new_param, + two_body_params, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_minu = phi.expectationValue(ham) + one_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) + two_body_gradient = np.zeros_like(two_body_params) + for ig, _ in enumerate(two_body_gradient): + new_param = two_body_params.copy() + new_param[ig] = new_param[ig] + dtheta + phi = get_evolved_wf( + one_body_params, + new_param, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_plus = phi.expectationValue(ham) + new_param[ig] = new_param[ig] - 2 * dtheta + phi = get_evolved_wf( + one_body_params, + new_param, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_minu = phi.expectationValue(ham) + two_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) + return one_body_gradient, two_body_gradient + + +@pytest.mark.parametrize("n_elec, n_orb", ((2, 2), (4, 4), (6, 6))) +@pytest.mark.parametrize("restricted", (True, False)) +def test_gradient(n_elec, n_orb, restricted): + sz = 0 + initial_wf = fqe.Wavefunction([[n_elec, sz, n_orb]]) + initial_wf.set_wfn(strategy="hartree-fock") + + fqe_ham = gen_random_restricted_ham(n_orb) + + if restricted: + n_one_body_params = n_orb * (n_orb - 1) // 2 + else: + n_one_body_params = n_orb * (n_orb - 1) + + gate_generators = _get_pp_plus_gate_generators( + n_elec=n_elec, heuristic_layers=tuple(), do_pp=True + ) + # reference implementation + one_body_params = np.random.random(n_one_body_params) + two_body_params = np.random.random(len(gate_generators)) + phi = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + )[0] + obj_val, grad = evaluate_gradient_and_cost_function( + initial_wf, + fqe_ham, + n_orb, + one_body_params, + two_body_params, + gate_generators, + restricted, + 0.0, + ) + ob_fd_grad, tb_fd_grad = get_fd_grad( + n_orb, + n_elec, + one_body_params, + two_body_params, + fqe_ham, + initial_wf, + restricted=restricted, + ) + assert np.isclose(obj_val, phi.expectationValue(fqe_ham)) + assert np.allclose(ob_fd_grad, grad[-n_one_body_params:]) + n_two_body_params = len(two_body_params) + assert np.allclose(tb_fd_grad, grad[:n_two_body_params]) From f1cc7fb7d87ee7e6cb7c66f17ad29d93480d72e9 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Fri, 21 Jun 2024 20:25:35 +0000 Subject: [PATCH 02/21] Addressing review comments. --- recirq/qcqmc/conftest.py | 13 ++++ recirq/qcqmc/trial_wf.py | 152 ++++++++++++++++++++++++--------------- 2 files changed, 109 insertions(+), 56 deletions(-) diff --git a/recirq/qcqmc/conftest.py b/recirq/qcqmc/conftest.py index 7d6c5d08..3d224fd0 100644 --- a/recirq/qcqmc/conftest.py +++ b/recirq/qcqmc/conftest.py @@ -1,3 +1,16 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. from typing import Tuple import numpy as np diff --git a/recirq/qcqmc/trial_wf.py b/recirq/qcqmc/trial_wf.py index 73c3e9ec..95602354 100644 --- a/recirq/qcqmc/trial_wf.py +++ b/recirq/qcqmc/trial_wf.py @@ -1,5 +1,20 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import abc import copy +import enum import itertools from typing import ( Callable, @@ -30,34 +45,36 @@ from scipy.optimize import minimize from scipy.sparse import csc_matrix -from recirq.qcqmc.afqmc_circuits import GeminalStatePreparationGate -from recirq.qcqmc.config import OUTDIRS -from recirq.qcqmc.for_refactor import Data, Params -from recirq.qcqmc.hamiltonian import HamiltonianData, HamiltonianParams +from recirq.qcqmc import afqmc_circuits, bitstrings, config, data, hamiltonian + + +class Spin(enum.Enum): + ALPHA = 0 + BETA = 1 @attrs.frozen class FermionicMode: - orb_ind: int - spin: str # Should be "a" or "b" only + """A specification of a fermionic mode. - def __attrs_post_init__(self): - if self.spin not in ["a", "b"]: - raise ValueError( - f'spin is set to {self.spin}, it should be either "a" or "b".' - ) + Args: + orb_ind: The spatial orbital index. + spin: The spin state of the fermion mode (up or down (alpha or beta)). + """ + + orb_ind: int + spin: Spin @classmethod def _json_namespace_(cls): return "recirq.qcqmc" def _json_dict_(self): - # return cirq.dataclass_json_dict(self) return attrs.asdict(self) @property def openfermion_standard_index(self) -> int: - return 2 * self.orb_ind + (self.spin == "b") + return 2 * self.orb_ind + self.spin.value @attrs.frozen @@ -83,18 +100,24 @@ def __post_init__(self): ) @classmethod - def _json_namespace_(self): + def _json_namespace_(cls): return "recirq.qcqmc" def _json_dict_(self): - # return cirq.dataclass_json_dict(self) return attrs.asdict(self) @attrs.frozen -class TrialWavefunctionParams(Params, metaclass=abc.ABCMeta): +class TrialWavefunctionParams(data.Params, metaclass=abc.ABCMeta): + """Parameters specifying a trial wavefunction. + + Args: + name: A descriptive name for the wavefunction parameters. + hamiltonian_params: Hamiltonian parameters specifying the molecule. + """ + name: str - hamiltonian_params: HamiltonianParams + hamiltonian_params: hamiltonian.HamiltonianParams @property def bitstrings(self) -> Iterable[Tuple[bool, ...]]: @@ -125,15 +148,35 @@ def _to_tuple(x: Iterable[LayerSpec]) -> Sequence[LayerSpec]: @attrs.frozen(repr=False) class PerfectPairingPlusTrialWavefunctionParams(TrialWavefunctionParams): - """Class for storing the parameters that specify a TrialWavefunctionData. + """Class for storing the parameters that specify the trial wavefunction. This class specifically stores the parameters for a trial wavefunction that is a combination of a perfect pairing wavefunction with some number of hardware-efficient layers appended. + + Args: + name: A name for the trial wavefunction. + hamiltonian_params: The hamiltonian parameters specifying the molecule. + heuristic_layers: A tuple of circuit layers to append to the perfect pairing circuit. + do_pp: Implement the perfect pairing circuit along with the heuristic + layers. Defaults to true. + restricted: Use a restricted perfect pairing ansatz. Defaults to false, + i.e. allow spin-symmetry breaking. + random_parameter_scale: A float to scale the random parameters by. + n_optimization_restarts: The number of times to restart the optimization + from a random guess in an attempt at global optimization. + seed: The random number seed to initialize the RNG with. + initial_orbital_rotation: An optional initial orbital rotation matrix, + which will be implmented as a givens circuit. + initial_two_body_qchem_amplitudes: Initial perfect pairing two-body + amplitudes using a qchem convention. + do_optimization: Optimize the ansatz using BFGS. + use_fast_gradients: Compute the parameter gradients using an anlytic + form. Default to false (use finite difference gradients). """ name: str - hamiltonian_params: HamiltonianParams + hamiltonian_params: hamiltonian.HamiltonianParams heuristic_layers: Tuple[LayerSpec, ...] = attrs.field(converter=_to_tuple) do_pp: bool = True restricted: bool = False @@ -167,11 +210,11 @@ def n_pairs(self) -> int: @property def path_string(self) -> str: - return OUTDIRS.DEFAULT_TRIAL_WAVEFUNCTION_DIRECTORY + self.name + return config.OUTDIRS.DEFAULT_TRIAL_WAVEFUNCTION_DIRECTORY + self.name @property def bitstrings(self) -> Iterable[Tuple[bool, ...]]: - return _get_bitstrings_a_b(n_orb=self.n_orb, n_elec=self.n_elec) + return bitstrings.get_bitstrings_a_b(n_orb=self.n_orb, n_elec=self.n_elec) def _json_dict_(self): # return cirq.dataclass_json_dict(self) @@ -211,30 +254,6 @@ def _get_qubits_a_b_reversed(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: ) -def _get_bitstrings_a_b(*, n_orb: int, n_elec: int) -> Iterable[Tuple[bool, ...]]: - """Iterates over bitstrings with the right symmetry assuming a_b ordering. - - This function assumes that the first n_orb qubits correspond to the alpha - orbitals and the second n_orb qubits correspond to the beta orbitals. The - ordering within the alpha and beta sectors doesn't matter (because we - iterate over all bitstrings with Hamming weight n_elec//2 in each sector. - """ - - if n_orb != n_elec: - raise NotImplementedError("n_orb must equal n_elec.") - - initial_bitstring = tuple(False for _ in range(n_orb - n_elec // 2)) + tuple( - True for _ in range(n_elec // 2) - ) - - spin_sector_bitstrings = set() - for perm in itertools.permutations(initial_bitstring): - spin_sector_bitstrings.add(perm) - - for bitstring_a, bitstring_b in itertools.product(spin_sector_bitstrings, repeat=2): - yield bitstring_a + bitstring_b - - def _get_fermion_qubit_map_pp_plus(*, n_qubits: int) -> Dict[int, cirq.GridQubit]: if n_qubits == 4: return get_4_qubit_fermion_qubit_map() @@ -449,7 +468,7 @@ def get_and_check_energy( def build_pp_plus_trial_wavefunction( params: PerfectPairingPlusTrialWavefunctionParams, *, - dependencies: Dict[Params, Data], + dependencies: Dict[data.Params, data.Data], do_print: bool = False, ) -> TrialWavefunctionData: """Builds a TrialWavefunctionData from a TrialWavefunctionParams""" @@ -620,7 +639,9 @@ def get_4_qubit_pp_circuits( assert n_elec == 2 fermion_index_to_qubit_map = get_4_qubit_fermion_qubit_map() - geminal_gate = GeminalStatePreparationGate(two_body_params[0], indicator=True) + geminal_gate = afqmc_circuits.afqmc_circuits.GeminalStatePreparation( + two_body_params[0], indicator=True + ) ansatz_circuit = cirq.Circuit( cirq.decompose( @@ -686,8 +707,12 @@ def get_8_qubit_circuits( """ fermion_index_to_qubit_map = get_8_qubit_fermion_qubit_map() - geminal_gate_1 = GeminalStatePreparationGate(two_body_params[0], indicator=True) - geminal_gate_2 = GeminalStatePreparationGate(two_body_params[1], indicator=True) + geminal_gate_1 = afqmc_circuits.GeminalStatePreparation( + two_body_params[0], indicator=True + ) + geminal_gate_2 = afqmc_circuits.GeminalStatePreparation( + two_body_params[1], indicator=True + ) # We'll add the initial bit flips later. ansatz_circuit = cirq.Circuit( @@ -784,9 +809,15 @@ def get_12_qubit_circuits( fermion_index_to_qubit_map = get_12_qubit_fermion_qubit_map() - geminal_gate_1 = GeminalStatePreparationGate(two_body_params[0], indicator=True) - geminal_gate_2 = GeminalStatePreparationGate(two_body_params[1], indicator=True) - geminal_gate_3 = GeminalStatePreparationGate(two_body_params[2], indicator=True) + geminal_gate_1 = afqmc_circuits.GeminalStatePreparation( + two_body_params[0], indicator=True + ) + geminal_gate_2 = afqmc_circuits.GeminalStatePreparation( + two_body_params[1], indicator=True + ) + geminal_gate_3 = afqmc_circuits.GeminalStatePreparation( + two_body_params[2], indicator=True + ) # We'll add the initial bit flips later. ansatz_circuit = cirq.Circuit( @@ -899,10 +930,18 @@ def get_16_qubit_circuits( """ fermion_index_to_qubit_map = get_16_qubit_fermion_qubit_map() - geminal_gate_1 = GeminalStatePreparationGate(two_body_params[0], indicator=True) - geminal_gate_2 = GeminalStatePreparationGate(two_body_params[1], indicator=True) - geminal_gate_3 = GeminalStatePreparationGate(two_body_params[2], indicator=True) - geminal_gate_4 = GeminalStatePreparationGate(two_body_params[3], indicator=True) + geminal_gate_1 = afqmc_circuits.GeminalStatePreparation( + two_body_params[0], indicator=True + ) + geminal_gate_2 = afqmc_circuits.GeminalStatePreparation( + two_body_params[1], indicator=True + ) + geminal_gate_3 = afqmc_circuits.GeminalStatePreparation( + two_body_params[2], indicator=True + ) + geminal_gate_4 = afqmc_circuits.GeminalStatePreparation( + two_body_params[3], indicator=True + ) # We'll add the initial bit flips later. ansatz_circuit = cirq.Circuit( @@ -1582,3 +1621,4 @@ def fast_obj_grad(params): print(two_body_params) return one_body_params, two_body_params, one_body_basis_change_mat + return one_body_params, two_body_params, one_body_basis_change_mat From 7d21ac9aaafcf588a863aa25a66a5061e4227f8a Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Sat, 22 Jun 2024 07:16:36 +0000 Subject: [PATCH 03/21] Addressing review comments. --- recirq/qcqmc/trial_wf.py | 271 ++++++++++++++++++++++++++++++++------- 1 file changed, 222 insertions(+), 49 deletions(-) diff --git a/recirq/qcqmc/trial_wf.py b/recirq/qcqmc/trial_wf.py index 95602354..f93db598 100644 --- a/recirq/qcqmc/trial_wf.py +++ b/recirq/qcqmc/trial_wf.py @@ -33,14 +33,13 @@ import cirq import fqe import fqe.algorithm.low_rank -import fqe.hamiltonians.hamiltonian +import fqe.hamiltonians.hamiltonian as fqe_ham import fqe.openfermion_utils +import fqe.wavefunction as fqe_wfn import numpy as np import numpy.typing as npt import openfermion as of import scipy.optimize -from fqe.hamiltonians.restricted_hamiltonian import RestrictedHamiltonian -from fqe.wavefunction import Wavefunction as FqeWavefunction from scipy.linalg import expm from scipy.optimize import minimize from scipy.sparse import csc_matrix @@ -49,6 +48,8 @@ class Spin(enum.Enum): + """A simple enum for distinguishing spin up (alpha) and spin down (beta) electrons.""" + ALPHA = 0 BETA = 1 @@ -217,7 +218,6 @@ def bitstrings(self) -> Iterable[Tuple[bool, ...]]: return bitstrings.get_bitstrings_a_b(n_orb=self.n_orb, n_elec=self.n_elec) def _json_dict_(self): - # return cirq.dataclass_json_dict(self) return attrs.asdict(self) @property @@ -234,9 +234,15 @@ def mode_qubit_map(self) -> Dict[FermionicMode, cirq.GridQubit]: def _get_qubits_a_b(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: - # This ordering creates qubits to facilitate a Jordan Wigner string - # threading through a row of alpha orbitals in ascending order followed by a - # row of beta orbitals in ascending order. + """Get grid alpha/beta grid qubits in ascending order. + + Args: + n_orb: The number of spatial orbitals. + + This ordering creates qubits to facilitate a Jordan Wigner string + threading through a row of alpha orbitals in ascending order followed by a + row of beta orbitals in ascending order. + """ return tuple( [cirq.GridQubit(0, i) for i in range(n_orb)] + [cirq.GridQubit(1, i) for i in range(n_orb)] @@ -244,10 +250,16 @@ def _get_qubits_a_b(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: def _get_qubits_a_b_reversed(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: - # This ordering creates qubits to facilitate operations that need a linearly - # connected array of qubits with the order threading through a row of alpha - # orbitals in ascending order followed by a row of beta orbitals in - # descending order. + """Get grid quibts with correct spin ordering. + + This ordering creates qubits to facilitate operations that need a linearly + connected array of qubits with the order threading through a row of alpha + orbitals in ascending order followed by a row of beta orbitals in + descending order. + + Args: + n_orb: The number of spatial orbitals. + """ return tuple( [cirq.GridQubit(0, i) for i in range(n_orb)] + [cirq.GridQubit(1, i) for i in reversed(range(n_orb))] @@ -255,6 +267,7 @@ def _get_qubits_a_b_reversed(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: def _get_fermion_qubit_map_pp_plus(*, n_qubits: int) -> Dict[int, cirq.GridQubit]: + """Dispatcher for qubit mappings.""" if n_qubits == 4: return get_4_qubit_fermion_qubit_map() elif n_qubits == 8: @@ -276,6 +289,9 @@ def _get_mode_qubit_map_pp_plus( to get_X_qubit_fermion_qubit_map for specific values of X but it translates this logic to a new system that uses the FermionicMode class rather than opaque combinations of integers and strings. + + Args: + n_qubits: The number of qubits. """ old_fermion_qubit_map = _get_fermion_qubit_map_pp_plus(n_qubits=n_qubits) @@ -284,8 +300,8 @@ def _get_mode_qubit_map_pp_plus( mode_qubit_map = {} for i in range(n_orb): - mode_qubit_map[FermionicMode(i, "a")] = old_fermion_qubit_map[2 * i] - mode_qubit_map[FermionicMode(i, "b")] = old_fermion_qubit_map[2 * i + 1] + mode_qubit_map[FermionicMode(i, Spin.ALPHA)] = old_fermion_qubit_map[2 * i] + mode_qubit_map[FermionicMode(i, Spin.BETA)] = old_fermion_qubit_map[2 * i + 1] return mode_qubit_map @@ -307,7 +323,6 @@ def _get_reorder_func( Args: mode_qubit_map: A dict that shows how each FermionicMode is mapped to a qubit. - ordered_qubits: An ordered sequence of qubits. """ qubits = list(mode_qubit_map.values()) @@ -323,9 +338,6 @@ def _get_reorder_func( openfermion_index = 2 * mode.orb_ind + (0 if mode.spin == "a" else 1) remapping_map[openfermion_index] = i - print("remapping_map:") - print(remapping_map) - def remapper(index: int, _: int) -> int: """A function that maps from the old index to the new one. @@ -373,7 +385,7 @@ def _get_fqe_wavefunctions( do_pp: bool = True, restricted: Optional[bool] = False, initial_orbital_rotation: Optional[np.ndarray] = None, -) -> Tuple[FqeWavefunction, FqeWavefunction]: +) -> Tuple[fqe_wfn.Wavefunction, fqe_wfn.Wavefunction]: initial_wf = fqe.Wavefunction([[n_elec, 0, n_orb]]) initial_wf.set_wfn(strategy="hartree-fock") @@ -477,7 +489,7 @@ def build_pp_plus_trial_wavefunction( print("Building Trial Wavefunction") np.random.seed(params.seed) hamiltonian_data = dependencies[params.hamiltonian_params] - assert isinstance(hamiltonian_data, HamiltonianData) + assert isinstance(hamiltonian_data, hamiltonian.HamiltonianData) assert ( params.n_orb == params.n_elec @@ -544,7 +556,7 @@ def build_pp_plus_trial_wavefunction( def get_rotated_hamiltonians( *, - hamiltonian_data: HamiltonianData, + hamiltonian_data: hamiltonian.HamiltonianData, one_body_basis_change_mat: np.ndarray, mode_qubit_map: Mapping[FermionicMode, cirq.Qid], ordered_qubits: Sequence[cirq.Qid], @@ -580,8 +592,8 @@ def get_rotated_hamiltonians( def get_energy_and_check_sanity( *, circuit_wf: np.ndarray, - fqe_wf: FqeWavefunction, - unrotated_fqe_wf: FqeWavefunction, + fqe_wf: fqe_wfn.Wavefunction, + unrotated_fqe_wf: fqe_wfn.Wavefunction, fqe_ham: fqe.hamiltonians.hamiltonian.Hamiltonian, sparse_ham: csc_matrix, e_core: float, @@ -1080,11 +1092,17 @@ def get_two_body_params_from_qchem_amplitudes( def convert_fqe_wf_to_cirq( - fqe_wf: FqeWavefunction, + fqe_wf: fqe_wfn.Wavefunction, mode_qubit_map: Mapping[FermionicMode, cirq.Qid], ordered_qubits: Sequence[cirq.Qid], ) -> np.ndarray: - """Converts an FQE wavefunction to one on qubits with a particular ordering.""" + """Converts an FQE wavefunction to one on qubits with a particular ordering. + + Args: + fqe_wf: The FQE wavefunction. + mode_qubit_map: A mapping from fermion modes to cirq qubits. + ordered_qubits: + """ n_qubits = len(mode_qubit_map) fermion_op = fqe.openfermion_utils.fqe_to_fermion_operator(fqe_wf) @@ -1126,14 +1144,29 @@ def get_one_body_cluster_coef( def get_evolved_wf( - one_body_params, - two_body_params, - wf, - gate_generators, - n_orb, - restricted, + one_body_params: np.ndarray, + two_body_params: np.ndarray, + wf: fqe.Wavefunction, + gate_generators: List[of.FermionOperator], + n_orb: int, + restricted: bool = True, initial_orbital_rotation: Optional[np.ndarray] = None, -): +) -> Tuple[fqe.Wavefunction, fqe.Wavefunction]: + """Get the wavefunction evaluated for this set of variational parameters. + + Args: + one_body_params: The variational parameters for the one-body terms in the ansatz. + two_body_params: The variational parameters for the two-body terms in the ansatz. + wf: The FQE wavefunction to evolve. + gate_generators: The generators of the two-body interaction terms. + n_orb: The number of spatial orbitals. + restricted: Whether the ansatz is restricted or not. + initial_orbital_rotation: Any initial orbital rotation to prepend to the circuit. + + Returs: + rotated_wf: the evolved wavefunction + wf: The original wavefunction + """ param_num = 0 for gate_generator in gate_generators: wf = wf.time_evolve(two_body_params[param_num], gate_generator) @@ -1158,7 +1191,18 @@ def get_evolved_wf( return rotated_wf, wf -def get_pair_hopping_gate_generators(n_pairs, n_elec): +def get_pair_hopping_gate_generators( + n_pairs: int, n_elec: int +) -> List[of.FermionOperator]: + """Get the generators of the pair-hopping unitaries. + + Args: + n_pairs: The number of pair coupling terms. + n_elec: The total number of electrons. + + Returns: + A list of gate generators + """ gate_generators = [] for pair in range(n_pairs): to_a = n_elec + 2 * pair @@ -1166,7 +1210,7 @@ def get_pair_hopping_gate_generators(n_pairs, n_elec): from_a = n_elec - 2 * pair - 2 from_b = n_elec - 2 * pair - 1 - fop_string = "{:d} {:d} {:d}^ {:d}^".format(to_b, to_a, from_b, from_a) + fop_string = f"{to_b} {to_a} {from_b}^ {from_a}^" gate_generator = of.FermionOperator(fop_string, 1.0) gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) @@ -1175,8 +1219,15 @@ def get_pair_hopping_gate_generators(n_pairs, n_elec): return gate_generators -def get_indices_heuristic_layer_in_pair(n_elec) -> Iterator[Tuple[int, int]]: - # Indices that couple within a pair +def get_indices_heuristic_layer_in_pair(n_elec: int) -> Iterator[Tuple[int, int]]: + """Get the indicies for the heuristic layers. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices + """ n_pairs = n_elec // 2 for pair in range(n_pairs): @@ -1189,7 +1240,14 @@ def get_indices_heuristic_layer_in_pair(n_elec) -> Iterator[Tuple[int, int]]: def get_indices_heuristic_layer_cross_pair(n_elec) -> Iterator[Tuple[int, int]]: - # Indices that couple adjacent pairs + """Indices that couple adjacent pairs. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices + """ n_pairs = n_elec // 2 for pair in range(n_pairs - 1): @@ -1202,7 +1260,14 @@ def get_indices_heuristic_layer_cross_pair(n_elec) -> Iterator[Tuple[int, int]]: def get_indices_heuristic_layer_cross_spin(n_elec) -> Iterator[Tuple[int, int]]: - # Indices that couple the two spin sectors + """Get indices that couple the two spin sectors. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices that couple spin sectors. + """ n_pairs = n_elec // 2 for pair in range(n_pairs): @@ -1215,7 +1280,14 @@ def get_indices_heuristic_layer_cross_spin(n_elec) -> Iterator[Tuple[int, int]]: def get_charge_charge_generator(indices: Tuple[int, int]) -> of.FermionOperator: - # Returns the generator for density evolution between the indices + """Returns the generator for density evolution between the indices + + Args: + indices: The indices to for charge-charge terms.:w + + Returns: + The generator for density evolution for this pair of electrons. + """ fop_string = "{:d}^ {:d} {:d}^ {:d}".format( indices[0], indices[0], indices[1], indices[1] @@ -1228,11 +1300,27 @@ def get_charge_charge_generator(indices: Tuple[int, int]) -> of.FermionOperator: def get_charge_charge_gate( qubits: Tuple[cirq.Qid, ...], param: float ) -> cirq.Operation: + """Get the cirq charge-charge gate. + + Args: + qubits: Two qubits you want to apply the gate to. + param: The parameter for the charge-charge interaction. + + Returns: + The charge-charge gate. + """ return cirq.CZ(qubits[0], qubits[1]) ** (-param / np.pi) -def get_givens_generator(indices): - # Returns the generator for density evolution between the indices +def get_givens_generator(indices: Tuple[int, int]) -> of.FermionOperator: + """Returns the generator for givens rotation between two orbitals. + + Args: + indices: The two indices for the givens rotation. + + Returns: + The givens generator for evolution for this pair of electrons. + """ fop_string = "{:d}^ {:d}".format(indices[0], indices[1]) gate_generator = of.FermionOperator(fop_string, 1.0) @@ -1242,10 +1330,28 @@ def get_givens_generator(indices): def get_givens_gate(qubits: Tuple[cirq.Qid, ...], param: float) -> cirq.Operation: + """Get a the givens rotation gate on two qubits. + + Args: + qubits: The two qubits to apply the gate to. + param: The parameter for the givens rotation. + + Returns: + The givens rotation gate. + """ return cirq.givens(param).on(qubits[0], qubits[1]) def get_layer_indices(layer_spec: LayerSpec, n_elec: int) -> List[Tuple[int, int]]: + """Get the indices for the heuristic layers. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + + Returns: + A list of indices for the layer. + """ indices_generators = { "in_pair": get_indices_heuristic_layer_in_pair(n_elec), "cross_pair": get_indices_heuristic_layer_cross_pair(n_elec), @@ -1262,7 +1368,17 @@ def get_layer_gates( params: np.ndarray, fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], ) -> List[cirq.Operation]: - """Gets the gates for a hardware efficient layer of the ansatz.""" + """Gets the gates for a hardware efficient layer of the ansatz. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + params: The variational parameters for the hardware efficient gate layer. + fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. + + Returns: + A list of gates for the layer. + """ indices_list = get_layer_indices(layer_spec, n_elec) @@ -1277,8 +1393,18 @@ def get_layer_gates( return gates -def get_layer_generators(layer_spec: LayerSpec, n_elec: int): - """Gets the generators for rotations in a hardware efficient layer of the ansatz.""" +def get_layer_generators( + layer_spec: LayerSpec, n_elec: int +) -> List[of.FermionOperator]: + """Gets the generators for rotations in a hardware efficient layer of the ansatz. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + + Returns: + A list of generators for the layers. + """ indices_list = get_layer_indices(layer_spec, n_elec) @@ -1291,7 +1417,18 @@ def get_layer_generators(layer_spec: LayerSpec, n_elec: int): return [gate_func(indices) for indices in indices_list] -def get_heuristic_gate_generators(n_elec: int, layer_specs: Sequence[LayerSpec]): +def get_heuristic_gate_generators( + n_elec: int, layer_specs: Sequence[LayerSpec] +) -> List[of.FermionOperator]: + """Get gate generators for the heuristic ansatz. + + Args: + n_elec: The number of electrons. + layer_specs: The layer specifications. + + Returns: + A list of generators for the layers. + """ gate_generators = [] for layer_spec in layer_specs: @@ -1306,6 +1443,17 @@ def get_heuristic_circuit( params: np.ndarray, fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], ) -> cirq.Circuit: + """Get a circuit for the heuristic ansatz. + + Args: + layer_specs: The layer specs for the heuristic layers. + n_elec: The number of electrons. + params: The variational parameters for the circuit. + fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. + + Returns: + A circuit for the heuristic ansatz. + """ gates: List[cirq.Operation] = [] for layer_spec in layer_specs: @@ -1317,13 +1465,16 @@ def get_heuristic_circuit( return cirq.Circuit(gates) -def orbital_rotation_gradient_matrix(generator_mat, a, b): +def orbital_rotation_gradient_matrix( + generator_mat: np.ndarray, a: int, b: int +) -> np.ndarray: """The gradient of the orbital rotation unitary with respect to its parameters. Args: generator_mat: The orbital rotation one-body generator matrix. a, b: row and column indices corresponding to the location in the matrix of the parameter we wish to find the gradient with respect to. + Returns: The orbital rotation matrix gradient wrt theta_{a, b}. Corresponds to expression in G15 of https://arxiv.org/abs/2004.04174. @@ -1354,8 +1505,8 @@ def orbital_rotation_gradient_matrix(generator_mat, a, b): def evaluate_gradient_and_cost_function( - initial_wf: FqeWavefunction, - fqe_ham: RestrictedHamiltonian, + initial_wf: fqe.Wavefunction, + fqe_ham: fqe_ham.RestrictedHamiltonian, n_orb: int, one_body_params: npt.NDArray[np.float64], two_body_params: npt.NDArray[np.float64], @@ -1375,6 +1526,7 @@ def evaluate_gradient_and_cost_function( retricted: Whether the single-particle rotations are restricted (True) or unrestricted (False). Unrestricted implies different parameters for the alpha- and beta-spin rotations. + Returns: cost_val: The cost function (total energy) evaluated for the input wavefunction parameters. grad: An array of gradients with respect to the one- and two-body @@ -1439,7 +1591,7 @@ def evaluate_gradient_and_cost_function( def get_pp_plus_params( *, - hamiltonian_data: HamiltonianData, + hamiltonian_data: hamiltonian.HamiltonianData, restricted: bool = False, random_parameter_scale: float = 1.0, initial_orbital_rotation: Optional[np.ndarray] = None, @@ -1449,6 +1601,28 @@ def get_pp_plus_params( do_print: bool = True, use_fast_gradients: bool = False, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Optimize the PP + Hardware layer ansatz. + + Args: + hamiltonian_data: Hamiltonian (molecular) specification. + restricted: Whether to use a spin-restricted ansatz or not. + random_parameter_scale: A float to scale the random parameters by. + initial_orbital_rotation: An optional initial orbital rotation matrix, + which will be implmented as a givens circuit. + heuristic_layers: A tuple of circuit layers to append to the perfect pairing circuit. + do_pp: Implement the perfect pairing circuit along with the heuristic + layers. Defaults to true. + n_optimization_restarts: The number of times to restart the optimization + from a random guess in an attempt at global optimization. + do_print: Whether to print optimization progress to stdout. + use_fast_gradients: Compute the parameter gradients anlytically using Wilcox formula. + Default to false (use finite difference gradients). + + Returns: + one_body_params: Optimized one-body parameters. + two_body_params: Optimized two-body parameters + one_body_basis_change_mat: The basis change matrix including any initial orbital rotation. + """ n_elec = hamiltonian_data.params.n_elec n_orb = hamiltonian_data.params.n_orb sz = 0 @@ -1621,4 +1795,3 @@ def fast_obj_grad(params): print(two_body_params) return one_body_params, two_body_params, one_body_basis_change_mat - return one_body_params, two_body_params, one_body_basis_change_mat From 158bc97e8b90b7188135972199fde161dfdc5177 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 17:42:06 +0000 Subject: [PATCH 04/21] Addressing comments. --- recirq/qcqmc/trial_wf_test.py | 52 ++++++++++++++++++++++++++--------- 1 file changed, 39 insertions(+), 13 deletions(-) diff --git a/recirq/qcqmc/trial_wf_test.py b/recirq/qcqmc/trial_wf_test.py index b2f64d72..801da266 100644 --- a/recirq/qcqmc/trial_wf_test.py +++ b/recirq/qcqmc/trial_wf_test.py @@ -1,9 +1,24 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import cirq import fqe +import fqe.hamiltonians as fqe_hams +import fqe.wavefunction as fqe_wfn import numpy as np import pytest import scipy.special -from fqe.hamiltonians.restricted_hamiltonian import RestrictedHamiltonian from recirq.qcqmc.hamiltonian import HamiltonianData from recirq.qcqmc.trial_wf import (FermionicMode, LayerSpec, @@ -257,7 +272,7 @@ def test_qchem_conversion_negative(fixture_4_qubit_ham: HamiltonianData): assert any(ansatz_qubit_wf < 0) -def gen_random_restricted_ham(n_orb: int) -> RestrictedHamiltonian: +def gen_random_restricted_ham(n_orb: int) -> fqe_hams.RestrictedHamiltonian: """8-fold symmetry restricted hamiltonian""" h1e = np.random.random((n_orb,) * 2) h1e = h1e + h1e.T @@ -266,20 +281,31 @@ def gen_random_restricted_ham(n_orb: int) -> RestrictedHamiltonian: h2e = h2e + h2e.transpose(3, 2, 1, 0) h2e = h2e + h2e.transpose(1, 0, 2, 3) h2e = np.asarray(h2e.transpose(0, 2, 3, 1), order="C") - fqe_ham = RestrictedHamiltonian((h1e, np.einsum("ijlk", -0.5 * h2e))) + fqe_ham = fqe_hams.RestrictedHamiltonian((h1e, np.einsum("ijlk", -0.5 * h2e))) return fqe_ham -def get_fd_grad( - n_orb, - n_elec, - one_body_params, - two_body_params, - ham, - initial_wf, - dtheta=1e-4, - restricted=False, +def compute_finite_difference_grad( + n_orb: int, + n_elec: int, + one_body_params: np.ndarray, + two_body_params: np.ndarray, + ham: fqe_hams.RestrictedHamiltonian, + initial_wf: fqe_wfn.Wavefunction, + dtheta: float = 1e-4, + restricted: bool = False, ): + """Compute the parameter gradient using finite differences. + + Args: + n_orb: the number of spatial orbitals. + n_elec: the number of electrons. + one_body_params: The variational parameters for the one-body terms in the ansatz. + two_body_params: The variational parameters for the two-body terms in the ansatz. + ham: The restricted FQE Hamiltonian. + initial_wf: The initial wavefunction (typically Hartree--Fock) + restricted: Whether we're using a restricted ansatz or not. + """ generators = _get_pp_plus_gate_generators( n_elec=n_elec, heuristic_layers=tuple(), do_pp=True ) @@ -372,7 +398,7 @@ def test_gradient(n_elec, n_orb, restricted): restricted, 0.0, ) - ob_fd_grad, tb_fd_grad = get_fd_grad( + ob_fd_grad, tb_fd_grad = compute_finite_difference_grad( n_orb, n_elec, one_body_params, From 40ebb8737637bd754a3331feb07974eaa2c48699 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 18:26:14 +0000 Subject: [PATCH 05/21] Revert enum. --- recirq/qcqmc/__init__.py | 4 +- recirq/qcqmc/conftest.py | 9 ++--- recirq/qcqmc/trial_wf.py | 69 +++++++++++++++++------------------ recirq/qcqmc/trial_wf_test.py | 37 ++++++------------- 4 files changed, 52 insertions(+), 67 deletions(-) diff --git a/recirq/qcqmc/__init__.py b/recirq/qcqmc/__init__.py index 2d86c742..1314592b 100644 --- a/recirq/qcqmc/__init__.py +++ b/recirq/qcqmc/__init__.py @@ -18,7 +18,7 @@ from .hamiltonian import ( HamiltonianData, - LoadFromFileHamiltonianParams, + HamiltonianFileParams, PyscfHamiltonianParams, ) from .trial_wf import ( @@ -43,7 +43,7 @@ def _resolve_json(cirq_type: str) -> Optional[ObjectFactory]: return { k.__name__: k for k in [ - LoadFromFileHamiltonianParams, + HamiltonianFileParams, HamiltonianData, PyscfHamiltonianParams, FermionicMode, diff --git a/recirq/qcqmc/conftest.py b/recirq/qcqmc/conftest.py index 3d224fd0..a61a6a22 100644 --- a/recirq/qcqmc/conftest.py +++ b/recirq/qcqmc/conftest.py @@ -18,20 +18,19 @@ from recirq.qcqmc.hamiltonian import ( HamiltonianData, - LoadFromFileHamiltonianParams, + HamiltonianFileParams, build_hamiltonian_from_file, ) from recirq.qcqmc.trial_wf import ( PerfectPairingPlusTrialWavefunctionParams, TrialWavefunctionData, - _get_qubits_a_b_reversed, build_pp_plus_trial_wavefunction, ) @pytest.fixture(scope="package") def fixture_4_qubit_ham() -> HamiltonianData: - params = LoadFromFileHamiltonianParams( + params = HamiltonianFileParams( name="test hamiltonian 4 qubits", integral_key="fh_sto3g", n_orb=2, n_elec=2 ) @@ -42,7 +41,7 @@ def fixture_4_qubit_ham() -> HamiltonianData: @pytest.fixture(scope="package") def fixture_8_qubit_ham() -> HamiltonianData: - params = LoadFromFileHamiltonianParams( + params = HamiltonianFileParams( name="test hamiltonian 8 qubits", integral_key="h4_sto3g", n_orb=4, n_elec=4 ) @@ -53,7 +52,7 @@ def fixture_8_qubit_ham() -> HamiltonianData: @pytest.fixture(scope="package") def fixture_12_qubit_ham() -> HamiltonianData: - params = LoadFromFileHamiltonianParams( + params = HamiltonianFileParams( name="test hamiltonian 12 qubits", integral_key="diamond_dzvp/cas66", n_orb=6, diff --git a/recirq/qcqmc/trial_wf.py b/recirq/qcqmc/trial_wf.py index f93db598..345978dd 100644 --- a/recirq/qcqmc/trial_wf.py +++ b/recirq/qcqmc/trial_wf.py @@ -33,7 +33,7 @@ import cirq import fqe import fqe.algorithm.low_rank -import fqe.hamiltonians.hamiltonian as fqe_ham +import fqe.hamiltonians.restricted_hamiltonian as fqe_hams import fqe.openfermion_utils import fqe.wavefunction as fqe_wfn import numpy as np @@ -47,13 +47,6 @@ from recirq.qcqmc import afqmc_circuits, bitstrings, config, data, hamiltonian -class Spin(enum.Enum): - """A simple enum for distinguishing spin up (alpha) and spin down (beta) electrons.""" - - ALPHA = 0 - BETA = 1 - - @attrs.frozen class FermionicMode: """A specification of a fermionic mode. @@ -64,7 +57,13 @@ class FermionicMode: """ orb_ind: int - spin: Spin + spin: str + + def __attrs_post_init__(self): + if self.spin not in ["a", "b"]: + raise ValueError( + "Spin must be either a or b for spin alpha(up) or beta(down) respectively." + ) @classmethod def _json_namespace_(cls): @@ -75,7 +74,7 @@ def _json_dict_(self): @property def openfermion_standard_index(self) -> int: - return 2 * self.orb_ind + self.spin.value + return 2 * self.orb_ind + (0 if self.spin.value == "a" else 1) @attrs.frozen @@ -300,8 +299,8 @@ def _get_mode_qubit_map_pp_plus( mode_qubit_map = {} for i in range(n_orb): - mode_qubit_map[FermionicMode(i, Spin.ALPHA)] = old_fermion_qubit_map[2 * i] - mode_qubit_map[FermionicMode(i, Spin.BETA)] = old_fermion_qubit_map[2 * i + 1] + mode_qubit_map[FermionicMode(i, "a")] = old_fermion_qubit_map[2 * i] + mode_qubit_map[FermionicMode(i, "b")] = old_fermion_qubit_map[2 * i + 1] return mode_qubit_map @@ -405,7 +404,7 @@ def _get_fqe_wavefunctions( @attrs.frozen -class TrialWavefunctionData(Data): +class TrialWavefunctionData(data.Data): """Class for storing a trial wavefunction's data.""" params: PerfectPairingPlusTrialWavefunctionParams @@ -425,7 +424,7 @@ def _json_dict_(self): def get_and_check_energy( *, - hamiltonian_data: HamiltonianData, + hamiltonian_data: hamiltonian.HamiltonianData, ansatz_circuit: cirq.Circuit, one_body_params: np.ndarray, two_body_params: np.ndarray, @@ -651,8 +650,8 @@ def get_4_qubit_pp_circuits( assert n_elec == 2 fermion_index_to_qubit_map = get_4_qubit_fermion_qubit_map() - geminal_gate = afqmc_circuits.afqmc_circuits.GeminalStatePreparation( - two_body_params[0], indicator=True + geminal_gate = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[0], inline_control=True ) ansatz_circuit = cirq.Circuit( @@ -719,11 +718,11 @@ def get_8_qubit_circuits( """ fermion_index_to_qubit_map = get_8_qubit_fermion_qubit_map() - geminal_gate_1 = afqmc_circuits.GeminalStatePreparation( - two_body_params[0], indicator=True + geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[0], inline_control=True ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparation( - two_body_params[1], indicator=True + geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[1], inline_control=True ) # We'll add the initial bit flips later. @@ -821,14 +820,14 @@ def get_12_qubit_circuits( fermion_index_to_qubit_map = get_12_qubit_fermion_qubit_map() - geminal_gate_1 = afqmc_circuits.GeminalStatePreparation( - two_body_params[0], indicator=True + geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[0], inline_control=True ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparation( - two_body_params[1], indicator=True + geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[1], inline_control=True ) - geminal_gate_3 = afqmc_circuits.GeminalStatePreparation( - two_body_params[2], indicator=True + geminal_gate_3 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[2], inline_control=True ) # We'll add the initial bit flips later. @@ -942,17 +941,17 @@ def get_16_qubit_circuits( """ fermion_index_to_qubit_map = get_16_qubit_fermion_qubit_map() - geminal_gate_1 = afqmc_circuits.GeminalStatePreparation( - two_body_params[0], indicator=True + geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[0], inline_control=True ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparation( - two_body_params[1], indicator=True + geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[1], inline_control=True ) - geminal_gate_3 = afqmc_circuits.GeminalStatePreparation( - two_body_params[2], indicator=True + geminal_gate_3 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[2], inline_control=True ) - geminal_gate_4 = afqmc_circuits.GeminalStatePreparation( - two_body_params[3], indicator=True + geminal_gate_4 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[3], inline_control=True ) # We'll add the initial bit flips later. @@ -1506,7 +1505,7 @@ def orbital_rotation_gradient_matrix( def evaluate_gradient_and_cost_function( initial_wf: fqe.Wavefunction, - fqe_ham: fqe_ham.RestrictedHamiltonian, + fqe_ham: fqe_hams.RestrictedHamiltonian, n_orb: int, one_body_params: npt.NDArray[np.float64], two_body_params: npt.NDArray[np.float64], diff --git a/recirq/qcqmc/trial_wf_test.py b/recirq/qcqmc/trial_wf_test.py index 801da266..6f34f3ac 100644 --- a/recirq/qcqmc/trial_wf_test.py +++ b/recirq/qcqmc/trial_wf_test.py @@ -14,21 +14,23 @@ import cirq import fqe -import fqe.hamiltonians as fqe_hams +import fqe.hamiltonians.restricted_hamiltonian as fqe_hams import fqe.wavefunction as fqe_wfn import numpy as np import pytest -import scipy.special from recirq.qcqmc.hamiltonian import HamiltonianData -from recirq.qcqmc.trial_wf import (FermionicMode, LayerSpec, - PerfectPairingPlusTrialWavefunctionParams, - _get_ansatz_qubit_wf, _get_bitstrings_a_b, - _get_pp_plus_gate_generators, - build_pp_plus_trial_wavefunction, - evaluate_gradient_and_cost_function, - get_evolved_wf, - get_two_body_params_from_qchem_amplitudes) +from recirq.qcqmc.trial_wf import ( + FermionicMode, + LayerSpec, + PerfectPairingPlusTrialWavefunctionParams, + _get_ansatz_qubit_wf, + _get_pp_plus_gate_generators, + build_pp_plus_trial_wavefunction, + evaluate_gradient_and_cost_function, + get_evolved_wf, + get_two_body_params_from_qchem_amplitudes, +) def test_fermionic_mode(): @@ -40,21 +42,6 @@ def test_fermionic_mode(): _ = FermionicMode(10, "c") -def test_get_bitstrings_a_b(): - with pytest.raises(NotImplementedError): - list(_get_bitstrings_a_b(n_orb=4, n_elec=3)) - - bitstrings = np.array(list(_get_bitstrings_a_b(n_orb=4, n_elec=4))) - - assert bitstrings.shape[0] == scipy.special.binom(4, 2) ** 2 - assert bitstrings.shape[1] == 2 * 4 # n_qubits columns = 2 * n_orb. - hamming_weight_left = np.sum(bitstrings[:, 0:4], axis=1) - hamming_weight_right = np.sum(bitstrings[:, 4:8], axis=1) - - assert np.all(hamming_weight_left == 2) - assert np.all(hamming_weight_right == 2) - - def test_pp_wf_energy(fixture_4_qubit_ham: HamiltonianData): params = PerfectPairingPlusTrialWavefunctionParams( name="pp_test_wf_1", From 671f88971c76ec23449916ca9b3dd8c622767699 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 18:28:22 +0000 Subject: [PATCH 06/21] Fix errors. --- recirq/qcqmc/trial_wf.py | 32 +------------------------------- 1 file changed, 1 insertion(+), 31 deletions(-) diff --git a/recirq/qcqmc/trial_wf.py b/recirq/qcqmc/trial_wf.py index 345978dd..dc4aab16 100644 --- a/recirq/qcqmc/trial_wf.py +++ b/recirq/qcqmc/trial_wf.py @@ -14,7 +14,6 @@ import abc import copy -import enum import itertools from typing import ( Callable, @@ -74,7 +73,7 @@ def _json_dict_(self): @property def openfermion_standard_index(self) -> int: - return 2 * self.orb_ind + (0 if self.spin.value == "a" else 1) + return 2 * self.orb_ind + (0 if self.spin == "a" else 1) @attrs.frozen @@ -374,35 +373,6 @@ def _get_superposition_wf( ) -def _get_fqe_wavefunctions( - *, - one_body_params: np.ndarray, - two_body_params: np.ndarray, - n_orb: int, - n_elec: int, - heuristic_layers: Tuple[LayerSpec, ...], - do_pp: bool = True, - restricted: Optional[bool] = False, - initial_orbital_rotation: Optional[np.ndarray] = None, -) -> Tuple[fqe_wfn.Wavefunction, fqe_wfn.Wavefunction]: - initial_wf = fqe.Wavefunction([[n_elec, 0, n_orb]]) - initial_wf.set_wfn(strategy="hartree-fock") - - wf, unrotated_wf = get_evolved_wf( - one_body_params=one_body_params, - two_body_params=two_body_params, - wf=initial_wf, - gate_generators=_get_pp_plus_gate_generators( - n_elec=n_elec, heuristic_layers=heuristic_layers, do_pp=do_pp - ), - n_orb=n_orb, - restricted=restricted, - initial_orbital_rotation=initial_orbital_rotation, - ) - - return wf, unrotated_wf - - @attrs.frozen class TrialWavefunctionData(data.Data): """Class for storing a trial wavefunction's data.""" From ccf839cff0d4efe182ef7e4f1edbb12ee7471189 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 18:55:17 +0000 Subject: [PATCH 07/21] Refactor trial_wf. --- recirq/qcqmc/bitstrings.py | 2 +- recirq/qcqmc/fermion_mode.py | 44 + recirq/qcqmc/fermion_mode_test.py | 25 + recirq/qcqmc/optimize_wf.py | 1369 ++++++++++++++++++++++++ recirq/qcqmc/optimize_wf_test.py | 387 +++++++ recirq/qcqmc/qubit_maps.py | 192 ++++ recirq/qcqmc/trial_wf.py | 1599 +---------------------------- recirq/qcqmc/trial_wf_test.py | 10 - 8 files changed, 2040 insertions(+), 1588 deletions(-) create mode 100644 recirq/qcqmc/fermion_mode.py create mode 100644 recirq/qcqmc/fermion_mode_test.py create mode 100644 recirq/qcqmc/optimize_wf.py create mode 100644 recirq/qcqmc/optimize_wf_test.py create mode 100644 recirq/qcqmc/qubit_maps.py diff --git a/recirq/qcqmc/bitstrings.py b/recirq/qcqmc/bitstrings.py index 50dbb8af..17d10ce1 100644 --- a/recirq/qcqmc/bitstrings.py +++ b/recirq/qcqmc/bitstrings.py @@ -37,4 +37,4 @@ def get_bitstrings_a_b(*, n_orb: int, n_elec: int) -> Iterable[Tuple[bool, ...]] spin_sector_bitstrings.add(perm) for bitstring_a, bitstring_b in itertools.product(spin_sector_bitstrings, repeat=2): - yield bitstring_a + bitstring_b + yield bitstring_a + bitstring_b \ No newline at end of file diff --git a/recirq/qcqmc/fermion_mode.py b/recirq/qcqmc/fermion_mode.py new file mode 100644 index 00000000..c37fe4b3 --- /dev/null +++ b/recirq/qcqmc/fermion_mode.py @@ -0,0 +1,44 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import attrs + + +@attrs.frozen +class FermionicMode: + """A specification of a fermionic mode. + + Args: + orb_ind: The spatial orbital index. + spin: The spin state of the fermion mode (up or down (alpha or beta)). + """ + + orb_ind: int + spin: str + + def __attrs_post_init__(self): + if self.spin not in ["a", "b"]: + raise ValueError( + "Spin must be either a or b for spin alpha(up) or beta(down) respectively." + ) + + @classmethod + def _json_namespace_(cls): + return "recirq.qcqmc" + + def _json_dict_(self): + return attrs.asdict(self) + + @property + def openfermion_standard_index(self) -> int: + return 2 * self.orb_ind + (0 if self.spin == "a" else 1) diff --git a/recirq/qcqmc/fermion_mode_test.py b/recirq/qcqmc/fermion_mode_test.py new file mode 100644 index 00000000..4d5ff965 --- /dev/null +++ b/recirq/qcqmc/fermion_mode_test.py @@ -0,0 +1,25 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import cirq + +from recirq.qcqmc.fermion_mode import FermionicMode + + +def test_fermionic_mode(): + fm = FermionicMode(5, "a") + fm2 = cirq.read_json(json_text=cirq.to_json(fm)) + assert fm == fm2 + + with pytest.raises(ValueError, match="spin.*"): + _ = FermionicMode(10, "c") diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py new file mode 100644 index 00000000..9b6da4d4 --- /dev/null +++ b/recirq/qcqmc/optimize_wf.py @@ -0,0 +1,1369 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import itertools +from typing import Callable, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple + +import cirq +import fqe +import fqe.hamiltonians.restricted_hamiltonian as fqe_hams +import fqe.wavefunction as fqe_wfn +import numpy as np +import openfermion as of +import scipy.linalg +import scipy.optimize +import scipy.sparse + +from recirq.qcqmc import ( + afqmc_circuits, + data, + fermion_mode, + hamiltonian, + qubit_maps, + trial_wf, +) + + +def _get_reorder_func( + *, + mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> Callable[[int, int], int]: + """This is a helper function that allows us to reorder fermionic modes. + + Under the Jordan-Wigner transform, each fermionic mode is assigned to a + qubit. If we are provided an openfermion FermionOperator with the modes + assigned to qubits as described by mode_qubit_map this function gives us a + reorder_func that we can use to reorder the modes (with + openfermion.reorder(...)) so that they match the order of the qubits in + ordered_qubits. This is necessary to make a correspondence between + fermionic operators / wavefunctions and their qubit counterparts. + + Args: + mode_qubit_map: A dict that shows how each FermionicMode is mapped to a qubit. + ordered_qubits: An ordered sequence of qubits. + """ + qubits = list(mode_qubit_map.values()) + assert len(qubits) == len(ordered_qubits) + + # We sort the key: value pairs by the order of the values (qubits) in + # ordered_qubits. + sorted_mapping = list(mode_qubit_map.items()) + sorted_mapping.sort(key=lambda x: ordered_qubits.index(x[1])) + + remapping_map = {} + for i, (mode, _) in enumerate(sorted_mapping): + openfermion_index = 2 * mode.orb_ind + (0 if mode.spin == "a" else 1) + remapping_map[openfermion_index] = i + + def remapper(index: int, _: int) -> int: + """A function that maps from the old index to the new one. + + The _ argument is because it's expected by openfermion.reorder""" + return remapping_map[index] + + return remapper + + +def _get_pp_plus_gate_generators( + *, n_elec: int, heuristic_layers: Tuple[trial_wf.LayerSpec, ...], do_pp: bool = True +) -> List[of.FermionOperator]: + heuristic_gate_generators = get_heuristic_gate_generators(n_elec, heuristic_layers) + if not do_pp: + return heuristic_gate_generators + + n_pairs = n_elec // 2 + pair_gate_generators = get_pair_hopping_gate_generators(n_pairs, n_elec) + return pair_gate_generators + heuristic_gate_generators + + +def _get_ansatz_qubit_wf( + *, ansatz_circuit: cirq.Circuit, ordered_qubits: Sequence[cirq.Qid] +): + return cirq.final_state_vector( + ansatz_circuit, qubit_order=list(ordered_qubits), dtype=np.complex128 + ) + + +def get_and_check_energy( + *, + hamiltonian_data: hamiltonian.HamiltonianData, + ansatz_circuit: cirq.Circuit, + one_body_params: np.ndarray, + two_body_params: np.ndarray, + one_body_basis_change_mat: np.ndarray, + params: trial_wf.PerfectPairingPlusTrialWavefunctionParams, +) -> Tuple[float, float]: + ansatz_qubit_wf = _get_ansatz_qubit_wf( + ansatz_circuit=ansatz_circuit, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + fqe_ham, e_core, sparse_ham = get_rotated_hamiltonians( + hamiltonian_data=hamiltonian_data, + one_body_basis_change_mat=one_body_basis_change_mat, + mode_qubit_map=params.mode_qubit_map, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + initial_wf = fqe_wfn.Wavefunction([[params.n_elec, 0, params.n_orb]]) + initial_wf.set_wfn(strategy="hartree-fock") + + hf_energy = initial_wf.expectationValue(fqe_ham) + e_core + + fqe_wf, unrotated_fqe_wf = get_evolved_wf( + one_body_params=one_body_params, + two_body_params=two_body_params, + wf=initial_wf, + gate_generators=_get_pp_plus_gate_generators( + n_elec=params.n_elec, + heuristic_layers=params.heuristic_layers, + do_pp=params.do_pp, + ), + n_orb=params.n_orb, + restricted=params.restricted, + initial_orbital_rotation=params.initial_orbital_rotation, + ) + + ansatz_energy = get_energy_and_check_sanity( + circuit_wf=ansatz_qubit_wf, + fqe_wf=fqe_wf, + unrotated_fqe_wf=unrotated_fqe_wf, + fqe_ham=fqe_ham, + sparse_ham=sparse_ham, + e_core=e_core, + mode_qubit_map=params.mode_qubit_map, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + return ansatz_energy, hf_energy + + +def build_pp_plus_trial_wavefunction( + params: trial_wf.PerfectPairingPlusTrialWavefunctionParams, + *, + dependencies: Dict[data.Params, data.Data], + do_print: bool = False, +) -> trial_wf.TrialWavefunctionData: + """Builds a TrialWavefunctionData from a TrialWavefunctionParams""" + + if do_print: + print("Building Trial Wavefunction") + np.random.seed(params.seed) + hamiltonian_data = dependencies[params.hamiltonian_params] + assert isinstance(hamiltonian_data, hamiltonian.HamiltonianData) + + assert ( + params.n_orb == params.n_elec + ) ## Necessary for perfect pairing wavefunction to make sense. + + if params.do_optimization: + ( + one_body_params, + two_body_params, + one_body_basis_change_mat, + ) = get_pp_plus_params( + hamiltonian_data=hamiltonian_data, + restricted=params.restricted, + random_parameter_scale=params.random_parameter_scale, + initial_orbital_rotation=params.initial_orbital_rotation, + heuristic_layers=params.heuristic_layers, + do_pp=params.do_pp, + n_optimization_restarts=params.n_optimization_restarts, + do_print=do_print, + use_fast_gradients=params.use_fast_gradients, + ) + else: + if ( + params.initial_two_body_qchem_amplitudes is None + or params.initial_orbital_rotation is not None + ): + raise NotImplementedError("TODO: Implement whatever isn't finished here.") + + n_one_body_params = params.n_orb * (params.n_orb - 1) + one_body_params = np.zeros(n_one_body_params) + one_body_basis_change_mat = np.diag(np.ones(params.n_orb * 2)) + two_body_params = get_two_body_params_from_qchem_amplitudes( + params.initial_two_body_qchem_amplitudes + ) + + (superposition_circuit, ansatz_circuit) = get_circuits( + two_body_params=two_body_params, + n_orb=params.n_orb, + n_elec=params.n_elec, + heuristic_layers=params.heuristic_layers, + ) + + ansatz_energy, hf_energy = get_and_check_energy( + hamiltonian_data=hamiltonian_data, + ansatz_circuit=ansatz_circuit, + params=params, + one_body_params=one_body_params, + two_body_params=two_body_params, + one_body_basis_change_mat=one_body_basis_change_mat, + ) + + return trial_wf.TrialWavefunctionData( + params=params, + ansatz_circuit=ansatz_circuit, + superposition_circuit=superposition_circuit, + hf_energy=hf_energy, + ansatz_energy=ansatz_energy, + fci_energy=hamiltonian_data.e_fci, + one_body_basis_change_mat=one_body_basis_change_mat, + one_body_params=one_body_params, + two_body_params=two_body_params, + ) + + +def get_rotated_hamiltonians( + *, + hamiltonian_data: hamiltonian.HamiltonianData, + one_body_basis_change_mat: np.ndarray, + mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> Tuple[fqe_hams.RestrictedHamiltonian, float, scipy.sparse.scipy.sparse.csc_matrix]: + """A helper method that gets the hamiltonians in the basis of the trial_wf. + + Returns: + The hamiltonian in FQE form, minus a constant energy shift. + The constant part of the Hamiltonian missing from the FQE Hamiltonian. + The qubit Hamiltonian as a sparse matrix. + """ + n_qubits = len(mode_qubit_map) + + fqe_ham = hamiltonian_data.get_restricted_fqe_hamiltonian() + e_core = hamiltonian_data.e_core + + mol_ham = hamiltonian_data.get_molecular_hamiltonian() + mol_ham.rotate_basis(one_body_basis_change_mat) + fermion_operator_ham = of.get_fermion_operator(mol_ham) + + reorder_func = _get_reorder_func( + mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits + ) + fermion_operator_ham_qubit_ordered = of.reorder( + fermion_operator_ham, reorder_func, num_modes=n_qubits + ) + + sparse_qubit_ham = of.get_sparse_operator(fermion_operator_ham_qubit_ordered) + + return fqe_ham, e_core, sparse_qubit_ham + + +def get_energy_and_check_sanity( + *, + circuit_wf: np.ndarray, + fqe_wf: fqe_wfn.Wavefunction, + unrotated_fqe_wf: fqe_wfn.Wavefunction, + fqe_ham: fqe_hams.RestrictedHamiltonian, + sparse_ham: scipy.sparse.csc_matrix, + e_core: float, + mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> float: + """A method that checks for consistency and returns the ansatz energy.""" + + unrotated_fqe_wf_as_cirq = convert_fqe_wf_to_cirq( + fqe_wf=unrotated_fqe_wf, + mode_qubit_map=mode_qubit_map, + ordered_qubits=ordered_qubits, + ) + ansatz_energy = np.real_if_close( + (np.conj(circuit_wf) @ sparse_ham @ circuit_wf) + ).item() + assert isinstance(ansatz_energy, float) + + fqe_energy = np.real(fqe_wf.expectationValue(fqe_ham) + e_core) + print(scipy.sparse.csc_matrix(circuit_wf)) + print(scipy.sparse.csc_matrix(unrotated_fqe_wf_as_cirq)) + np.testing.assert_array_almost_equal(ansatz_energy, fqe_energy) + np.testing.assert_array_almost_equal( + circuit_wf, unrotated_fqe_wf_as_cirq, decimal=5 + ) + return ansatz_energy + + +def get_4_qubit_pp_circuits( + *, + two_body_params: np.ndarray, + n_elec: int, + heuristic_layers: Tuple[trial_wf.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 3 1 + 2 0 + """ + assert n_elec == 2 + + fermion_index_to_qubit_map = qubit_maps.get_4_qubit_fermion_qubit_map() + geminal_gate = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[0], inline_control=True + ) + + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + ) + ) + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[1:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + indicator = fermion_index_to_qubit_map[2] + superposition_circuit = cirq.Circuit([cirq.H(indicator) + ansatz_circuit]) + ansatz_circuit = cirq.Circuit([cirq.X(indicator) + ansatz_circuit]) + + return superposition_circuit, ansatz_circuit + + +def get_8_qubit_circuits( + *, + two_body_params: np.ndarray, + n_elec: int, + heuristic_layers: Tuple[trial_wf.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 3 5 1 7 + 2 4 0 6 + """ + fermion_index_to_qubit_map = qubit_maps.get_8_qubit_fermion_qubit_map() + + geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[0], inline_control=True + ) + geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[1], inline_control=True + ) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[2:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[0]), + cirq.CNOT(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[6]), + cirq.SWAP(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[4]), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[4]), + cirq.X(fermion_index_to_qubit_map[6]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_12_qubit_circuits( + *, + two_body_params: np.ndarray, + n_elec: int, + heuristic_layers: Tuple[trial_wf.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 5 7 3 9 1 11 + 4 6 2 8 0 10 + """ + + fermion_index_to_qubit_map = qubit_maps.get_12_qubit_fermion_qubit_map() + + geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[0], inline_control=True + ) + geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[1], inline_control=True + ) + geminal_gate_3 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[2], inline_control=True + ) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[8], + fermion_index_to_qubit_map[9], + ) + ), + cirq.decompose( + geminal_gate_3.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[10], + fermion_index_to_qubit_map[11], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[3:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[8]), + cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[0]), + cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[2]), + cirq.SWAP( + fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[10] + ), + cirq.SWAP(fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[6]), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[6]), + cirq.X(fermion_index_to_qubit_map[8]), + cirq.X(fermion_index_to_qubit_map[10]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_16_qubit_circuits( + *, + two_body_params: np.ndarray, + n_elec: int, + heuristic_layers: Tuple[trial_wf.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 7 9 5 11 3 13 1 15 + 6 8 4 10 2 12 0 14 + """ + fermion_index_to_qubit_map = qubit_maps.get_16_qubit_fermion_qubit_map() + + geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[0], inline_control=True + ) + geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[1], inline_control=True + ) + geminal_gate_3 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[2], inline_control=True + ) + geminal_gate_4 = afqmc_circuits.GeminalStatePreparationGate( + two_body_params[3], inline_control=True + ) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + fermion_index_to_qubit_map[8], + fermion_index_to_qubit_map[9], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + fermion_index_to_qubit_map[10], + fermion_index_to_qubit_map[11], + ) + ), + cirq.decompose( + geminal_gate_3.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[12], + fermion_index_to_qubit_map[13], + ) + ), + cirq.decompose( + geminal_gate_4.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[14], + fermion_index_to_qubit_map[15], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[4:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[10]), + cirq.CNOT( + fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[2] + ), + cirq.SWAP( + fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[12] + ), + cirq.CNOT( + fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[4] + ), + cirq.CNOT( + fermion_index_to_qubit_map[12], fermion_index_to_qubit_map[0] + ), + cirq.SWAP(fermion_index_to_qubit_map[4], fermion_index_to_qubit_map[8]), + cirq.SWAP( + fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[14] + ), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[8]), + cirq.X(fermion_index_to_qubit_map[10]), + cirq.X(fermion_index_to_qubit_map[12]), + cirq.X(fermion_index_to_qubit_map[14]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_circuits( + *, + two_body_params: np.ndarray, + # from wf_params: + n_orb: int, + n_elec: int, + heuristic_layers: Tuple[trial_wf.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A function that runs a specialized method to get the ansatz circuits.""" + + # TODO(?): Just input one of these quantities. + if n_orb != n_elec: + raise ValueError("n_orb must equal n_elec.") + + circ_funcs = { + 2: get_4_qubit_pp_circuits, + 4: get_8_qubit_circuits, + 6: get_12_qubit_circuits, + 8: get_16_qubit_circuits, + } + try: + circ_func = circ_funcs[n_orb] + except KeyError: + raise NotImplementedError(f"No circuits for n_orb = {n_orb}") + + return circ_func( + two_body_params=two_body_params, + n_elec=n_elec, + heuristic_layers=heuristic_layers, + ) + + +def get_two_body_params_from_qchem_amplitudes( + qchem_amplitudes: np.ndarray, +) -> np.ndarray: + """Translates perfect pairing amplitudes from qchem to rotation angles. + + qchem style: 1 |1100> + t_i |0011> + our style: cos(\theta_i) |1100> + sin(\theta_i) |0011> + """ + + two_body_params = np.arccos(1 / np.sqrt(1 + qchem_amplitudes**2)) * np.sign( + qchem_amplitudes + ) + + # Numpy casts the array improperly to a float when we only have one parameter. + two_body_params = np.atleast_1d(two_body_params) + + return two_body_params + + +#################### Here be dragons.########################################### + + +def convert_fqe_wf_to_cirq( + fqe_wf: fqe_wfn.Wavefunction, + mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> np.ndarray: + """Converts an FQE wavefunction to one on qubits with a particular ordering. + + Args: + fqe_wf: The FQE wavefunction. + mode_qubit_map: A mapping from fermion modes to cirq qubits. + ordered_qubits: + """ + n_qubits = len(mode_qubit_map) + fermion_op = fqe.openfermion_utils.fqe_to_fermion_operator(fqe_wf) + + reorder_func = _get_reorder_func( + mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits + ) + fermion_op = of.reorder(fermion_op, reorder_func, num_modes=n_qubits) + + qubit_op = of.jordan_wigner(fermion_op) + + return fqe.qubit_wavefunction_from_vacuum( + qubit_op, list(cirq.LineQubit.range(n_qubits)) + ) + + +def get_one_body_cluster_coef( + params: np.ndarray, n_orb: int, restricted: bool +) -> np.ndarray: + if restricted: + one_body_cluster_op = np.zeros((n_orb, n_orb), dtype=np.complex128) + else: + one_body_cluster_op = np.zeros((2 * n_orb, 2 * n_orb), dtype=np.complex128) + param_num = 0 + + for i in range(n_orb): + for j in range(i): + one_body_cluster_op[i, j] = params[param_num] + one_body_cluster_op[j, i] = -params[param_num] + param_num += 1 + + if not restricted: + for i in range(n_orb, 2 * n_orb): + for j in range(n_orb, i): + one_body_cluster_op[i, j] = params[param_num] + one_body_cluster_op[j, i] = -params[param_num] + param_num += 1 + + return one_body_cluster_op + + +def get_evolved_wf( + one_body_params: np.ndarray, + two_body_params: np.ndarray, + wf: fqe.Wavefunction, + gate_generators: List[of.FermionOperator], + n_orb: int, + restricted: bool = True, + initial_orbital_rotation: Optional[np.ndarray] = None, +) -> Tuple[fqe.Wavefunction, fqe.Wavefunction]: + """Get the wavefunction evaluated for this set of variational parameters. + + Args: + one_body_params: The variational parameters for the one-body terms in the ansatz. + two_body_params: The variational parameters for the two-body terms in the ansatz. + wf: The FQE wavefunction to evolve. + gate_generators: The generators of the two-body interaction terms. + n_orb: The number of spatial orbitals. + restricted: Whether the ansatz is restricted or not. + initial_orbital_rotation: Any initial orbital rotation to prepend to the circuit. + + Returs: + rotated_wf: the evolved wavefunction + wf: The original wavefunction + """ + param_num = 0 + for gate_generator in gate_generators: + wf = wf.time_evolve(two_body_params[param_num], gate_generator) + param_num += 1 + + one_body_cluster_op = get_one_body_cluster_coef( + one_body_params, n_orb, restricted=restricted + ) + + if restricted: + one_body_ham = fqe.get_restricted_hamiltonian((1j * one_body_cluster_op,)) + else: + one_body_ham = fqe.get_sso_hamiltonian((1j * one_body_cluster_op,)) + + rotated_wf = wf.time_evolve(1.0, one_body_ham) + + if initial_orbital_rotation is not None: + rotated_wf = fqe.algorithm.low_rank.evolve_fqe_givens( + rotated_wf, initial_orbital_rotation + ) + + return rotated_wf, wf + + +def get_pair_hopping_gate_generators( + n_pairs: int, n_elec: int +) -> List[of.FermionOperator]: + """Get the generators of the pair-hopping unitaries. + + Args: + n_pairs: The number of pair coupling terms. + n_elec: The total number of electrons. + + Returns: + A list of gate generators + """ + gate_generators = [] + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + + fop_string = f"{to_b} {to_a} {from_b}^ {from_a}^" + + gate_generator = of.FermionOperator(fop_string, 1.0) + gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) + gate_generators.append(gate_generator) + + return gate_generators + + +def get_indices_heuristic_layer_in_pair(n_elec: int) -> Iterator[Tuple[int, int]]: + """Get the indicies for the heuristic layers. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices + """ + n_pairs = n_elec // 2 + + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + yield (from_a, to_a) + yield (from_b, to_b) + + +def get_indices_heuristic_layer_cross_pair(n_elec) -> Iterator[Tuple[int, int]]: + """Indices that couple adjacent pairs. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices + """ + n_pairs = n_elec // 2 + + for pair in range(n_pairs - 1): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_next_a = n_elec - 2 * (pair + 1) - 2 + from_next_b = n_elec - 2 * (pair + 1) - 1 + yield (to_a, from_next_a) + yield (to_b, from_next_b) + + +def get_indices_heuristic_layer_cross_spin(n_elec) -> Iterator[Tuple[int, int]]: + """Get indices that couple the two spin sectors. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices that couple spin sectors. + """ + n_pairs = n_elec // 2 + + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + yield (to_a, to_b) + yield (from_a, from_b) + + +def get_charge_charge_generator(indices: Tuple[int, int]) -> of.FermionOperator: + """Returns the generator for density evolution between the indices + + Args: + indices: The indices to for charge-charge terms.:w + + Returns: + The generator for density evolution for this pair of electrons. + """ + + fop_string = "{:d}^ {:d} {:d}^ {:d}".format( + indices[0], indices[0], indices[1], indices[1] + ) + gate_generator = of.FermionOperator(fop_string, 1.0) + + return gate_generator + + +def get_charge_charge_gate( + qubits: Tuple[cirq.Qid, ...], param: float +) -> cirq.Operation: + """Get the cirq charge-charge gate. + + Args: + qubits: Two qubits you want to apply the gate to. + param: The parameter for the charge-charge interaction. + + Returns: + The charge-charge gate. + """ + return cirq.CZ(qubits[0], qubits[1]) ** (-param / np.pi) + + +def get_givens_generator(indices: Tuple[int, int]) -> of.FermionOperator: + """Returns the generator for givens rotation between two orbitals. + + Args: + indices: The two indices for the givens rotation. + + Returns: + The givens generator for evolution for this pair of electrons. + """ + + fop_string = "{:d}^ {:d}".format(indices[0], indices[1]) + gate_generator = of.FermionOperator(fop_string, 1.0) + gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) + + return gate_generator + + +def get_givens_gate(qubits: Tuple[cirq.Qid, ...], param: float) -> cirq.Operation: + """Get a the givens rotation gate on two qubits. + + Args: + qubits: The two qubits to apply the gate to. + param: The parameter for the givens rotation. + + Returns: + The givens rotation gate. + """ + return cirq.givens(param).on(qubits[0], qubits[1]) + + +def get_layer_indices( + layer_spec: trial_wf.LayerSpec, n_elec: int +) -> List[Tuple[int, int]]: + """Get the indices for the heuristic layers. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + + Returns: + A list of indices for the layer. + """ + indices_generators = { + "in_pair": get_indices_heuristic_layer_in_pair(n_elec), + "cross_pair": get_indices_heuristic_layer_cross_pair(n_elec), + "cross_spin": get_indices_heuristic_layer_cross_spin(n_elec), + } + indices_generator = indices_generators[layer_spec.layout] + + return [indices for indices in indices_generator] + + +def get_layer_gates( + layer_spec: trial_wf.LayerSpec, + n_elec: int, + params: np.ndarray, + fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], +) -> List[cirq.Operation]: + """Gets the gates for a hardware efficient layer of the ansatz. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + params: The variational parameters for the hardware efficient gate layer. + fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. + + Returns: + A list of gates for the layer. + """ + + indices_list = get_layer_indices(layer_spec, n_elec) + + gate_funcs = {"givens": get_givens_gate, "charge_charge": get_charge_charge_gate} + gate_func = gate_funcs[layer_spec.base_gate] + + gates = [] + for indices, param in zip(indices_list, params): + qubits = tuple(fermion_index_to_qubit_map[ind] for ind in indices) + gates.append(gate_func(qubits, param)) + + return gates + + +def get_layer_generators( + layer_spec: trial_wf.LayerSpec, n_elec: int +) -> List[of.FermionOperator]: + """Gets the generators for rotations in a hardware efficient layer of the ansatz. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + + Returns: + A list of generators for the layers. + """ + + indices_list = get_layer_indices(layer_spec, n_elec) + + gate_funcs = { + "givens": get_givens_generator, + "charge_charge": get_charge_charge_generator, + } + gate_func = gate_funcs[layer_spec.base_gate] + + return [gate_func(indices) for indices in indices_list] + + +def get_heuristic_gate_generators( + n_elec: int, layer_specs: Sequence[trial_wf.LayerSpec] +) -> List[of.FermionOperator]: + """Get gate generators for the heuristic ansatz. + + Args: + n_elec: The number of electrons. + layer_specs: The layer specifications. + + Returns: + A list of generators for the layers. + """ + gate_generators = [] + + for layer_spec in layer_specs: + gate_generators += get_layer_generators(layer_spec, n_elec) + + return gate_generators + + +def get_heuristic_circuit( + layer_specs: Sequence[trial_wf.LayerSpec], + n_elec: int, + params: np.ndarray, + fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], +) -> cirq.Circuit: + """Get a circuit for the heuristic ansatz. + + Args: + layer_specs: The layer specs for the heuristic layers. + n_elec: The number of electrons. + params: The variational parameters for the circuit. + fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. + + Returns: + A circuit for the heuristic ansatz. + """ + gates: List[cirq.Operation] = [] + + for layer_spec in layer_specs: + params_slice = params[len(gates) :] + gates += get_layer_gates( + layer_spec, n_elec, params_slice, fermion_index_to_qubit_map + ) + + return cirq.Circuit(gates) + + +def orbital_rotation_gradient_matrix( + generator_mat: np.ndarray, a: int, b: int +) -> np.ndarray: + """The gradient of the orbital rotation unitary with respect to its parameters. + + Args: + generator_mat: The orbital rotation one-body generator matrix. + a, b: row and column indices corresponding to the location in the matrix + of the parameter we wish to find the gradient with respect to. + + Returns: + The orbital rotation matrix gradient wrt theta_{a, b}. Corresponds to + expression in G15 of https://arxiv.org/abs/2004.04174. + """ + w_full, v_full = np.linalg.eigh(-1j * generator_mat) + eigs_diff = np.zeros((w_full.shape[0], w_full.shape[0]), dtype=np.complex128) + for i, j in itertools.product(range(w_full.shape[0]), repeat=2): + if np.isclose(abs(w_full[i] - w_full[j]), 0): + eigs_diff[i, j] = 1 + else: + eigs_diff[i, j] = (np.exp(1j * (w_full[i] - w_full[j])) - 1) / ( + 1j * (w_full[i] - w_full[j]) + ) + + Y_full = np.zeros_like(v_full, dtype=np.complex128) + if a == b: + Y_full[a, b] = 0 + else: + Y_full[a, b] = 1.0 + Y_full[b, a] = -1.0 + + Y_kl_full = v_full.conj().T @ Y_full @ v_full + # now rotate Y_{kl} * (exp(i(l_{k} - l_{l})) - 1) / (i(l_{k} - l_{l})) + # into the original basis + pre_matrix_full = v_full @ (eigs_diff * Y_kl_full) @ v_full.conj().T + + return pre_matrix_full + + +def evaluate_gradient_and_cost_function( + initial_wf: fqe.Wavefunction, + fqe_ham: fqe_hams.RestrictedHamiltonian, + n_orb: int, + one_body_params: np.ndarray, + two_body_params: np.ndarray, + gate_generators: List[of.FermionOperator], + restricted: bool, + e_core: float, +) -> Tuple[float, np.ndarray]: + """Evaluate gradient and cost function for optimization. + + Args: + initial_wf: Initial state (typically Hartree--Fock). + fqe_ham: The restricted Hamiltonian in FQE format. + n_orb: The number of spatial orbitals. + one_body_params: The parameters of the single-particle rotations. + two_body_params: The parameters for the two-particle terms. + gate_generators: The generators for the two-particle terms. + retricted: Whether the single-particle rotations are restricted (True) + or unrestricted (False). Unrestricted implies different parameters + for the alpha- and beta-spin rotations. + + Returns: + cost_val: The cost function (total energy) evaluated for the input wavefunction parameters. + grad: An array of gradients with respect to the one- and two-body + parameters. The first n_orb * (n_orb + 1) // 2 parameters correspond to + the one-body gradients. + """ + phi = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + )[0] + lam = copy.deepcopy(phi) + lam = lam.apply(fqe_ham) + cost_val = fqe.vdot(lam, phi) + e_core + + # 1body + one_body_cluster_op = get_one_body_cluster_coef( + one_body_params, n_orb, restricted=restricted + ) + tril = np.tril_indices(n_orb, k=-1) + if restricted: + one_body_ham = fqe.get_restricted_hamiltonian((-1j * one_body_cluster_op,)) + else: + one_body_ham = fqe.get_sso_hamiltonian((-1j * one_body_cluster_op,)) + # Apply U1b^{dag} + phi.time_evolve(1, one_body_ham, inplace=True) + lam.time_evolve(1, one_body_ham, inplace=True) + one_body_grad = np.zeros_like(one_body_params) + n_one_body_params = len(one_body_params) + grad_position = n_one_body_params - 1 + for iparam in range(len(one_body_params)): + mu_state = copy.deepcopy(phi) + pidx = n_one_body_params - iparam - 1 + pidx_spin = 0 if restricted else pidx // (n_one_body_params // 2) + pidx_spat = pidx if restricted else pidx - (n_one_body_params // 2) * pidx_spin + p, q = (tril[0][pidx_spat], tril[1][pidx_spat]) + p += n_orb * pidx_spin + q += n_orb * pidx_spin + pre_matrix = orbital_rotation_gradient_matrix(-one_body_cluster_op, p, q) + assert of.is_hermitian(1j * pre_matrix) + if restricted: + fqe_quad_ham_pre = fqe.get_restricted_hamiltonian((pre_matrix,)) + else: + fqe_quad_ham_pre = fqe.get_sso_hamiltonian((pre_matrix,)) + mu_state = mu_state.apply(fqe_quad_ham_pre) + one_body_grad[grad_position] = 2 * fqe.vdot(lam, mu_state).real + grad_position -= 1 + # Get two-body contributions + two_body_grad = np.zeros(len(two_body_params)) + for pidx in reversed(range(len(gate_generators))): + mu = copy.deepcopy(phi) + mu = mu.apply(gate_generators[pidx]) + two_body_grad[pidx] = -np.real(2 * 1j * (fqe.vdot(lam, mu))) + phi = phi.time_evolve(-two_body_params[pidx], gate_generators[pidx]) + lam = lam.time_evolve(-two_body_params[pidx], gate_generators[pidx]) + + return cost_val, np.concatenate((two_body_grad, one_body_grad)) + + +def get_pp_plus_params( + *, + hamiltonian_data: hamiltonian.HamiltonianData, + restricted: bool = False, + random_parameter_scale: float = 1.0, + initial_orbital_rotation: Optional[np.ndarray] = None, + heuristic_layers: Tuple[trial_wf.LayerSpec, ...], + do_pp: bool = True, + n_optimization_restarts: int = 1, + do_print: bool = True, + use_fast_gradients: bool = False, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Optimize the PP + Hardware layer ansatz. + + Args: + hamiltonian_data: Hamiltonian (molecular) specification. + restricted: Whether to use a spin-restricted ansatz or not. + random_parameter_scale: A float to scale the random parameters by. + initial_orbital_rotation: An optional initial orbital rotation matrix, + which will be implmented as a givens circuit. + heuristic_layers: A tuple of circuit layers to append to the perfect pairing circuit. + do_pp: Implement the perfect pairing circuit along with the heuristic + layers. Defaults to true. + n_optimization_restarts: The number of times to restart the optimization + from a random guess in an attempt at global optimization. + do_print: Whether to print optimization progress to stdout. + use_fast_gradients: Compute the parameter gradients anlytically using Wilcox formula. + Default to false (use finite difference gradients). + + Returns: + one_body_params: Optimized one-body parameters. + two_body_params: Optimized two-body parameters + one_body_basis_change_mat: The basis change matrix including any initial orbital rotation. + """ + n_elec = hamiltonian_data.params.n_elec + n_orb = hamiltonian_data.params.n_orb + sz = 0 + + initial_wf = fqe.Wavefunction([[n_elec, sz, n_orb]]) + initial_wf.set_wfn(strategy="hartree-fock") + + fqe_ham = hamiltonian_data.get_restricted_fqe_hamiltonian() + e_core = hamiltonian_data.e_core + + hf_energy = initial_wf.expectationValue(fqe_ham) + e_core + + # We're only supporting closed shell stuff here. + assert n_elec % 2 == 0 + assert n_elec <= n_orb + if use_fast_gradients: + err_msg = "use_fast_gradients does not work with initial orbital rotation." + assert initial_orbital_rotation is None, err_msg + + gate_generators = _get_pp_plus_gate_generators( + n_elec=n_elec, heuristic_layers=heuristic_layers, do_pp=do_pp + ) + + n_two_body_params = len(gate_generators) + + if restricted: + n_one_body_params = n_orb * (n_orb - 1) // 2 + else: + n_one_body_params = n_orb * (n_orb - 1) + + best = np.inf + best_res: Union[None, scipy.optimize.OptimizeResult] = None + for i in range(n_optimization_restarts): + if do_print: + print(f"Optimization restart {i}", flush=True) + + def progress_cb(_): + print(".", end="", flush=True) + + else: + + def progress_cb(_): + pass + + params = random_parameter_scale * np.random.normal( + size=(n_two_body_params + n_one_body_params) + ) + + def objective(params): + one_body_params = params[-n_one_body_params:] + two_body_params = params[:n_two_body_params] + + wf, _ = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + initial_orbital_rotation=initial_orbital_rotation, + ) + + energy = wf.expectationValue(fqe_ham) + e_core + if do_print: + print(f"energy {energy}") + if np.abs(energy.imag) < 1e-6: + return energy.real + else: + return 1e6 + + def fast_obj_grad(params): + one_body_params = params[-n_one_body_params:] + two_body_params = params[:n_two_body_params] + energy, grad = evaluate_gradient_and_cost_function( + initial_wf, + fqe_ham, + n_orb, + one_body_params, + two_body_params, + gate_generators, + restricted, + e_core, + ) + if do_print: + print(f"energy {energy}, max|grad| {np.max(np.abs(grad))}") + if np.abs(energy.imag) < 1e-6: + return energy.real, grad + else: + return 1e6, 1e6 + + if use_fast_gradients: + res = scipy.optimize.minimize( + fast_obj_grad, params, jac=True, method="BFGS", callback=progress_cb + ) + else: + res = scipy.optimize.minimize(objective, params, callback=progress_cb) + if res.fun < best: + best = res.fun + best_res = res + + if do_print: + print(res, flush=True) + + assert best_res is not None + params = best_res.x + one_body_params = params[-n_one_body_params:] + two_body_params = params[:n_two_body_params] + + wf, _ = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + initial_orbital_rotation=initial_orbital_rotation, + ) + + one_body_cluster_mat = get_one_body_cluster_coef( + one_body_params, n_orb, restricted=restricted + ) + # We need to change the ordering to match OpenFermion's abababab ordering + if not restricted: + index_rearrangement = np.asarray( + [i // 2 % (n_orb) + (i % 2) * n_orb for i in range(2 * n_orb)] + ) + one_body_cluster_mat = one_body_cluster_mat[:, index_rearrangement] + one_body_cluster_mat = one_body_cluster_mat[index_rearrangement, :] + + one_body_basis_change_mat = scipy.linalg.expm(one_body_cluster_mat) + + if initial_orbital_rotation is not None: + if restricted: + one_body_basis_change_mat = ( + initial_orbital_rotation @ one_body_basis_change_mat + ) + else: + big_initial_orbital_rotation = np.zeros_like(one_body_basis_change_mat) + + for i in range(len(initial_orbital_rotation)): + for j in range(len(initial_orbital_rotation)): + big_initial_orbital_rotation[2 * i, 2 * j] = ( + initial_orbital_rotation[i, j] + ) + big_initial_orbital_rotation[2 * i + 1, 2 * j + 1] = ( + initial_orbital_rotation[i, j] + ) + + one_body_basis_change_mat = ( + big_initial_orbital_rotation @ one_body_basis_change_mat + ) + + if do_print: + print("Hartree-Fock Energy:") + print(hf_energy) + initial_wf.print_wfn() + print("-" * 80) + print("FCI Energy:") + print(hamiltonian_data.e_fci) + print("-" * 80) + print(best_res) + + print("-" * 80) + print("Ansatz Energy:") + print(np.real_if_close(wf.expectationValue(fqe_ham) + e_core)) + wf.print_wfn() + print("Basis Rotation Matrix:") + print(one_body_basis_change_mat) + print("Two Body Rotation Parameters:") + print(two_body_params) + + return one_body_params, two_body_params, one_body_basis_change_mat diff --git a/recirq/qcqmc/optimize_wf_test.py b/recirq/qcqmc/optimize_wf_test.py new file mode 100644 index 00000000..bea91d40 --- /dev/null +++ b/recirq/qcqmc/optimize_wf_test.py @@ -0,0 +1,387 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import fqe +import fqe.hamiltonians.restricted_hamiltonian as fqe_hams +import fqe.wavefunction as fqe_wfn +import numpy as np +import pytest + +from recirq.qcqmc.hamiltonian import HamiltonianData +from recirq.qcqmc.trial_wf import (LayerSpec, + PerfectPairingPlusTrialWavefunctionParams, + _get_ansatz_qubit_wf, + _get_pp_plus_gate_generators, + build_pp_plus_trial_wavefunction, + evaluate_gradient_and_cost_function, + get_evolved_wf, + get_two_body_params_from_qchem_amplitudes) + + +def test_pp_wf_energy(fixture_4_qubit_ham: HamiltonianData): + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_1", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=(), + do_pp=True, + restricted=True, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham} + ) + + assert np.isclose(trial_wf.ansatz_energy, fixture_4_qubit_ham.e_fci) + + +def test_pp_wf_energy_with_layer(fixture_4_qubit_ham: HamiltonianData): + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_2", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=(LayerSpec("charge_charge", "cross_spin"),), + do_pp=True, + restricted=True, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham} + ) + + assert np.isclose(trial_wf.ansatz_energy, fixture_4_qubit_ham.e_fci) + + +def test_qchem_pp_eight_qubit_wavefunctions_consistent( + fixture_8_qubit_ham: HamiltonianData, +): + """Tests (without optimization) that the eight qubit wavefunctions work. + + Specifically, that constructing the wavefunction with FQE and then + converting it to a cirq wavefunction yields the same result as constructing + the parameters with the circuit directly. + """ + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_qchem", + hamiltonian_params=fixture_8_qubit_ham.params, + heuristic_layers=tuple(), + initial_orbital_rotation=None, + initial_two_body_qchem_amplitudes=np.asarray([0.3, 0.4]), + do_optimization=False, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_8_qubit_ham.params: fixture_8_qubit_ham}, + do_print=False, + ) + + one_body_params = trial_wf.one_body_params + two_body_params = trial_wf.two_body_params + basis_change_mat = trial_wf.one_body_basis_change_mat + + np.testing.assert_array_almost_equal(one_body_params, np.zeros((12,))) + np.testing.assert_array_almost_equal(basis_change_mat, np.diag(np.ones(8))) + + np.testing.assert_equal(two_body_params.shape, (2,)) + + +def test_pp_plus_wf_energy_sloppy_1(fixture_8_qubit_ham: HamiltonianData): + params = PerfectPairingPlusTrialWavefunctionParams( + "pp_plus_test", + hamiltonian_params=fixture_8_qubit_ham.params, + heuristic_layers=tuple(), + do_pp=True, + restricted=False, + random_parameter_scale=1, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_8_qubit_ham.params: fixture_8_qubit_ham}, + do_print=True, + ) + + assert trial_wf.ansatz_energy < -1.947 + + +# TODO: Speed up this test and add a similar one with non-trivial heuristic layers. + + +def test_diamond_pp_wf_energy(fixture_12_qubit_ham: HamiltonianData): + params = PerfectPairingPlusTrialWavefunctionParams( + name="diamind_pp_test_wf_1", + hamiltonian_params=fixture_12_qubit_ham.params, + heuristic_layers=tuple(), + do_pp=True, + restricted=True, + random_parameter_scale=0.1, + n_optimization_restarts=1, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_12_qubit_ham.params: fixture_12_qubit_ham}, + do_print=True, + ) + + assert trial_wf.ansatz_energy < -10.4 + + +@pytest.mark.parametrize( + "initial_two_body_qchem_amplitudes, expected_ansatz_qubit_wf", + [ + ( + [1], + np.array( + [ + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.70710678 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.70710678 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + 0.0 + 0.0j, + ] + ), + ), + ( + [0], + np.array( + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ] + ), + ), + ], +) +def test_qchem_pp_runs( + initial_two_body_qchem_amplitudes, + expected_ansatz_qubit_wf, + fixture_4_qubit_ham: HamiltonianData, +): + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_qchem", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=tuple(), + initial_orbital_rotation=None, + initial_two_body_qchem_amplitudes=initial_two_body_qchem_amplitudes, + do_optimization=False, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham}, + do_print=False, + ) + + ansatz_qubit_wf = _get_ansatz_qubit_wf( + ansatz_circuit=trial_wf.ansatz_circuit, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + np.testing.assert_array_almost_equal(ansatz_qubit_wf, expected_ansatz_qubit_wf) + + +def test_qchem_conversion_negative(fixture_4_qubit_ham: HamiltonianData): + qchem_amplitudes = np.asarray(-0.1) + + two_body_params = get_two_body_params_from_qchem_amplitudes(qchem_amplitudes) + + assert two_body_params.item() < 0 + + params = PerfectPairingPlusTrialWavefunctionParams( + name="pp_test_wf_qchem_neg", + hamiltonian_params=fixture_4_qubit_ham.params, + heuristic_layers=tuple(), + initial_orbital_rotation=None, + initial_two_body_qchem_amplitudes=qchem_amplitudes, + do_optimization=False, + ) + + trial_wf = build_pp_plus_trial_wavefunction( + params, + dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham}, + do_print=False, + ) + + ansatz_qubit_wf = _get_ansatz_qubit_wf( + ansatz_circuit=trial_wf.ansatz_circuit, + ordered_qubits=params.qubits_jordan_wigner_ordered, + ) + + assert any(ansatz_qubit_wf < 0) + + +def gen_random_restricted_ham(n_orb: int) -> fqe_hams.RestrictedHamiltonian: + """8-fold symmetry restricted hamiltonian""" + h1e = np.random.random((n_orb,) * 2) + h1e = h1e + h1e.T + h2e = np.random.random((n_orb,) * 4) + h2e = h2e + h2e.transpose(2, 3, 0, 1) + h2e = h2e + h2e.transpose(3, 2, 1, 0) + h2e = h2e + h2e.transpose(1, 0, 2, 3) + h2e = np.asarray(h2e.transpose(0, 2, 3, 1), order="C") + fqe_ham = fqe_hams.RestrictedHamiltonian((h1e, np.einsum("ijlk", -0.5 * h2e))) + return fqe_ham + + +def compute_finite_difference_grad( + n_orb: int, + n_elec: int, + one_body_params: np.ndarray, + two_body_params: np.ndarray, + ham: fqe_hams.RestrictedHamiltonian, + initial_wf: fqe_wfn.Wavefunction, + dtheta: float = 1e-4, + restricted: bool = False, +): + """Compute the parameter gradient using finite differences. + + Args: + n_orb: the number of spatial orbitals. + n_elec: the number of electrons. + one_body_params: The variational parameters for the one-body terms in the ansatz. + two_body_params: The variational parameters for the two-body terms in the ansatz. + ham: The restricted FQE Hamiltonian. + initial_wf: The initial wavefunction (typically Hartree--Fock) + restricted: Whether we're using a restricted ansatz or not. + """ + generators = _get_pp_plus_gate_generators( + n_elec=n_elec, heuristic_layers=tuple(), do_pp=True + ) + one_body_gradient = np.zeros_like(one_body_params) + for ig, _ in enumerate(one_body_gradient): + new_param = one_body_params.copy() + new_param[ig] = new_param[ig] + dtheta + phi = get_evolved_wf( + new_param, + two_body_params, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_plus = phi.expectationValue(ham) + new_param[ig] = new_param[ig] - 2 * dtheta + phi = get_evolved_wf( + new_param, + two_body_params, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_minu = phi.expectationValue(ham) + one_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) + two_body_gradient = np.zeros_like(two_body_params) + for ig, _ in enumerate(two_body_gradient): + new_param = two_body_params.copy() + new_param[ig] = new_param[ig] + dtheta + phi = get_evolved_wf( + one_body_params, + new_param, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_plus = phi.expectationValue(ham) + new_param[ig] = new_param[ig] - 2 * dtheta + phi = get_evolved_wf( + one_body_params, + new_param, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_minu = phi.expectationValue(ham) + two_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) + return one_body_gradient, two_body_gradient + + +@pytest.mark.parametrize("n_elec, n_orb", ((2, 2), (4, 4), (6, 6))) +@pytest.mark.parametrize("restricted", (True, False)) +def test_gradient(n_elec, n_orb, restricted): + sz = 0 + initial_wf = fqe.Wavefunction([[n_elec, sz, n_orb]]) + initial_wf.set_wfn(strategy="hartree-fock") + + fqe_ham = gen_random_restricted_ham(n_orb) + + if restricted: + n_one_body_params = n_orb * (n_orb - 1) // 2 + else: + n_one_body_params = n_orb * (n_orb - 1) + + gate_generators = _get_pp_plus_gate_generators( + n_elec=n_elec, heuristic_layers=tuple(), do_pp=True + ) + # reference implementation + one_body_params = np.random.random(n_one_body_params) + two_body_params = np.random.random(len(gate_generators)) + phi = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + )[0] + obj_val, grad = evaluate_gradient_and_cost_function( + initial_wf, + fqe_ham, + n_orb, + one_body_params, + two_body_params, + gate_generators, + restricted, + 0.0, + ) + ob_fd_grad, tb_fd_grad = compute_finite_difference_grad( + n_orb, + n_elec, + one_body_params, + two_body_params, + fqe_ham, + initial_wf, + restricted=restricted, + ) + assert np.isclose(obj_val, phi.expectationValue(fqe_ham)) + assert np.allclose(ob_fd_grad, grad[-n_one_body_params:]) + n_two_body_params = len(two_body_params) + assert np.allclose(tb_fd_grad, grad[:n_two_body_params]) diff --git a/recirq/qcqmc/qubit_maps.py b/recirq/qcqmc/qubit_maps.py new file mode 100644 index 00000000..c2ea3f9d --- /dev/null +++ b/recirq/qcqmc/qubit_maps.py @@ -0,0 +1,192 @@ +from typing import Dict, Tuple + +import cirq + +from recirq.qcqmc.fermion_mode import FermionicMode + + +def get_qubits_a_b(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: + """Get grid alpha/beta grid qubits in ascending order. + + Args: + n_orb: The number of spatial orbitals. + + This ordering creates qubits to facilitate a Jordan Wigner string + threading through a row of alpha orbitals in ascending order followed by a + row of beta orbitals in ascending order. + """ + return tuple( + [cirq.GridQubit(0, i) for i in range(n_orb)] + + [cirq.GridQubit(1, i) for i in range(n_orb)] + ) + + +def get_qubits_a_b_reversed(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: + """Get grid quibts with correct spin ordering. + + This ordering creates qubits to facilitate operations that need a linearly + connected array of qubits with the order threading through a row of alpha + orbitals in ascending order followed by a row of beta orbitals in + descending order. + + Args: + n_orb: The number of spatial orbitals. + """ + return tuple( + [cirq.GridQubit(0, i) for i in range(n_orb)] + + [cirq.GridQubit(1, i) for i in reversed(range(n_orb))] + ) + + +def get_4_qubit_fermion_qubit_map() -> Dict[int, cirq.GridQubit]: + """A helper function that provides the fermion qubit map for 4 qubits. + + We map the fermionic orbitals to grid qubits like so: + 3 1 + 2 0 + """ + fermion_index_to_qubit_map = { + 2: cirq.GridQubit(0, 0), + 3: cirq.GridQubit(1, 0), + 0: cirq.GridQubit(0, 1), + 1: cirq.GridQubit(1, 1), + } + + return fermion_index_to_qubit_map + + +def get_8_qubit_fermion_qubit_map(): + """A helper function that provides the fermion qubit map for 8 qubits. + + We map the fermionic orbitals to grid qubits like so: + 3 5 1 7 + 2 4 0 6 + """ + + # Linear connectivity is fine. + # This ordering is dictated by the way we specify perfect pairing (i) + # Elsewhere we generate the perfect pairing parameters using a specific + # convention for how we index the FermionOperators (which is itself) + # partly dictated by the OpenFermion conventions. Here we choose a + # mapping between the indices of these FermionOperators and the qubits in our + # grid that allows for the perfect pairing pairs to be in squares of four. + fermion_index_to_qubit_map = { + 2: cirq.GridQubit(0, 0), + 3: cirq.GridQubit(1, 0), + 4: cirq.GridQubit(0, 1), + 5: cirq.GridQubit(1, 1), + 0: cirq.GridQubit(0, 2), + 1: cirq.GridQubit(1, 2), + 6: cirq.GridQubit(0, 3), + 7: cirq.GridQubit(1, 3), + } + + return fermion_index_to_qubit_map + + +def get_12_qubit_fermion_qubit_map(): + """A helper function that provides the fermion qubit map for 12 qubits. + + We map the fermionic orbitals to grid qubits like so: + 5 7 3 9 1 11 + 4 6 2 8 0 10 + """ + + # Linear connectivity is fine. + # This ordering is dictated by the way we specify perfect pairing (i) + # Elsewhere we generate the perfect pairing parameters using a specific + # convention for how we index the FermionOperators (which is itself) + # partly dictated by the OpenFermion conventions. Here we choose a + # mapping between the indices of these FermionOperators and the qubits in our + # grid that allows for the perfect pairing pairs to be in squares of four. + fermion_index_to_qubit_map = { + 4: cirq.GridQubit(0, 0), + 5: cirq.GridQubit(1, 0), + 6: cirq.GridQubit(0, 1), + 7: cirq.GridQubit(1, 1), + 2: cirq.GridQubit(0, 2), + 3: cirq.GridQubit(1, 2), + 8: cirq.GridQubit(0, 3), + 9: cirq.GridQubit(1, 3), + 0: cirq.GridQubit(0, 4), + 1: cirq.GridQubit(1, 4), + 10: cirq.GridQubit(0, 5), + 11: cirq.GridQubit(1, 5), + } + + return fermion_index_to_qubit_map + + +def get_16_qubit_fermion_qubit_map(): + """A helper function that provides the fermion qubit map for 16 qubits. + + We map the fermionic orbitals to grid qubits like so: + 7 9 5 11 3 13 1 15 + 6 8 4 10 2 12 0 14 + """ + + # Linear connectivity is fine. + # This ordering is dictated by the way we specify perfect pairing (i) + # Elsewhere we generate the perfect pairing parameters using a specific + # convention for how we index the FermionOperators (which is itself) + # partly dictated by the OpenFermion conventions. Here we choose a + # mapping between the indices of these FermionOperators and the qubits in our + # grid that allows for the perfect pairing pairs to be in squares of four. + fermion_index_to_qubit_map = { + 6: cirq.GridQubit(0, 0), + 7: cirq.GridQubit(1, 0), + 8: cirq.GridQubit(0, 1), + 9: cirq.GridQubit(1, 1), + 4: cirq.GridQubit(0, 2), + 5: cirq.GridQubit(1, 2), + 10: cirq.GridQubit(0, 3), + 11: cirq.GridQubit(1, 3), + 2: cirq.GridQubit(0, 4), + 3: cirq.GridQubit(1, 4), + 12: cirq.GridQubit(0, 5), + 13: cirq.GridQubit(1, 5), + 0: cirq.GridQubit(0, 6), + 1: cirq.GridQubit(1, 6), + 14: cirq.GridQubit(0, 7), + 15: cirq.GridQubit(1, 7), + } + + return fermion_index_to_qubit_map + + +def get_fermion_qubit_map_pp_plus(*, n_qubits: int) -> Dict[int, cirq.GridQubit]: + """Dispatcher for qubit mappings.""" + if n_qubits == 4: + return get_4_qubit_fermion_qubit_map() + elif n_qubits == 8: + return get_8_qubit_fermion_qubit_map() + elif n_qubits == 12: + return get_12_qubit_fermion_qubit_map() + elif n_qubits == 16: + return get_16_qubit_fermion_qubit_map() + else: + raise NotImplementedError() + + +def get_mode_qubit_map_pp_plus(*, n_qubits: int) -> Dict[FermionicMode, cirq.GridQubit]: + """A map from Fermionic modes to qubits for our particular circuits. + + This function dispatches to _get_fermion_qubit_map_pp_plus, and ultimately + to get_X_qubit_fermion_qubit_map for specific values of X but it translates + this logic to a new system that uses the FermionicMode class rather than + opaque combinations of integers and strings. + + Args: + n_qubits: The number of qubits. + """ + old_fermion_qubit_map = get_fermion_qubit_map_pp_plus(n_qubits=n_qubits) + + n_orb = n_qubits // 2 + + mode_qubit_map = {} + + for i in range(n_orb): + mode_qubit_map[FermionicMode(i, "a")] = old_fermion_qubit_map[2 * i] + mode_qubit_map[FermionicMode(i, "b")] = old_fermion_qubit_map[2 * i + 1] + + return mode_qubit_map diff --git a/recirq/qcqmc/trial_wf.py b/recirq/qcqmc/trial_wf.py index dc4aab16..f6ddd3ce 100644 --- a/recirq/qcqmc/trial_wf.py +++ b/recirq/qcqmc/trial_wf.py @@ -13,67 +13,13 @@ # limitations under the License. import abc -import copy -import itertools -from typing import ( - Callable, - Dict, - Iterable, - Iterator, - List, - Mapping, - Optional, - Sequence, - Tuple, - Union, -) +from typing import Dict, Iterable, Optional, Sequence, Tuple import attrs import cirq -import fqe -import fqe.algorithm.low_rank -import fqe.hamiltonians.restricted_hamiltonian as fqe_hams -import fqe.openfermion_utils -import fqe.wavefunction as fqe_wfn import numpy as np -import numpy.typing as npt -import openfermion as of -import scipy.optimize -from scipy.linalg import expm -from scipy.optimize import minimize -from scipy.sparse import csc_matrix -from recirq.qcqmc import afqmc_circuits, bitstrings, config, data, hamiltonian - - -@attrs.frozen -class FermionicMode: - """A specification of a fermionic mode. - - Args: - orb_ind: The spatial orbital index. - spin: The spin state of the fermion mode (up or down (alpha or beta)). - """ - - orb_ind: int - spin: str - - def __attrs_post_init__(self): - if self.spin not in ["a", "b"]: - raise ValueError( - "Spin must be either a or b for spin alpha(up) or beta(down) respectively." - ) - - @classmethod - def _json_namespace_(cls): - return "recirq.qcqmc" - - def _json_dict_(self): - return attrs.asdict(self) - - @property - def openfermion_standard_index(self) -> int: - return 2 * self.orb_ind + (0 if self.spin == "a" else 1) +from recirq.qcqmc import bitstrings, config, data, fermion_mode, hamiltonian, qubit_maps @attrs.frozen @@ -213,6 +159,7 @@ def path_string(self) -> str: @property def bitstrings(self) -> Iterable[Tuple[bool, ...]]: + """The full set of bitstrings (determinants) for this wavefunction.""" return bitstrings.get_bitstrings_a_b(n_orb=self.n_orb, n_elec=self.n_elec) def _json_dict_(self): @@ -220,162 +167,35 @@ def _json_dict_(self): @property def qubits_jordan_wigner_ordered(self) -> Tuple[cirq.GridQubit, ...]: - return _get_qubits_a_b(n_orb=self.n_orb) + """Get the cirq qubits assuming a Jordan-Wigner ordering.""" + return qubit_maps.get_qubits_a_b(n_orb=self.n_orb) @property def qubits_linearly_connected(self) -> Tuple[cirq.GridQubit, ...]: - return _get_qubits_a_b_reversed(n_orb=self.n_orb) + """Get the cirq qubits assuming a linear connected qubit ordering.""" + return qubit_maps.get_qubits_a_b_reversed(n_orb=self.n_orb) @property - def mode_qubit_map(self) -> Dict[FermionicMode, cirq.GridQubit]: - return _get_mode_qubit_map_pp_plus(n_qubits=self.n_qubits) - - -def _get_qubits_a_b(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: - """Get grid alpha/beta grid qubits in ascending order. - - Args: - n_orb: The number of spatial orbitals. - - This ordering creates qubits to facilitate a Jordan Wigner string - threading through a row of alpha orbitals in ascending order followed by a - row of beta orbitals in ascending order. - """ - return tuple( - [cirq.GridQubit(0, i) for i in range(n_orb)] - + [cirq.GridQubit(1, i) for i in range(n_orb)] - ) - - -def _get_qubits_a_b_reversed(*, n_orb: int) -> Tuple[cirq.GridQubit, ...]: - """Get grid quibts with correct spin ordering. + def mode_qubit_map(self) -> Dict[fermion_mode.FermionicMode, cirq.GridQubit]: + """Get the mapping between fermionic modes and cirq qubits.""" + return qubit_maps.get_mode_qubit_map_pp_plus(n_qubits=self.n_qubits) - This ordering creates qubits to facilitate operations that need a linearly - connected array of qubits with the order threading through a row of alpha - orbitals in ascending order followed by a row of beta orbitals in - descending order. - Args: - n_orb: The number of spatial orbitals. - """ - return tuple( - [cirq.GridQubit(0, i) for i in range(n_orb)] - + [cirq.GridQubit(1, i) for i in reversed(range(n_orb))] - ) - - -def _get_fermion_qubit_map_pp_plus(*, n_qubits: int) -> Dict[int, cirq.GridQubit]: - """Dispatcher for qubit mappings.""" - if n_qubits == 4: - return get_4_qubit_fermion_qubit_map() - elif n_qubits == 8: - return get_8_qubit_fermion_qubit_map() - elif n_qubits == 12: - return get_12_qubit_fermion_qubit_map() - elif n_qubits == 16: - return get_16_qubit_fermion_qubit_map() - else: - raise NotImplementedError() - - -def _get_mode_qubit_map_pp_plus( - *, n_qubits: int -) -> Dict[FermionicMode, cirq.GridQubit]: - """A map from Fermionic modes to qubits for our particular circuits. - - This function dispatches to _get_fermion_qubit_map_pp_plus, and ultimately - to get_X_qubit_fermion_qubit_map for specific values of X but it translates - this logic to a new system that uses the FermionicMode class rather than - opaque combinations of integers and strings. - - Args: - n_qubits: The number of qubits. - """ - old_fermion_qubit_map = _get_fermion_qubit_map_pp_plus(n_qubits=n_qubits) - - n_orb = n_qubits // 2 - - mode_qubit_map = {} - - for i in range(n_orb): - mode_qubit_map[FermionicMode(i, "a")] = old_fermion_qubit_map[2 * i] - mode_qubit_map[FermionicMode(i, "b")] = old_fermion_qubit_map[2 * i + 1] - - return mode_qubit_map - - -def _get_reorder_func( - *, - mode_qubit_map: Mapping[FermionicMode, cirq.Qid], - ordered_qubits: Sequence[cirq.Qid], -) -> Callable[[int, int], int]: - """This is a helper function that allows us to reorder fermionic modes. - - Under the Jordan-Wigner transform, each fermionic mode is assigned to a - qubit. If we are provided an openfermion FermionOperator with the modes - assigned to qubits as described by mode_qubit_map this function gives us a - reorder_func that we can use to reorder the modes (with - openfermion.reorder(...)) so that they match the order of the qubits in - ordered_qubits. This is necessary to make a correspondence between - fermionic operators / wavefunctions and their qubit counterparts. +@attrs.frozen +class TrialWavefunctionData(data.Data): + """Class for storing a trial wavefunction's data. Args: - mode_qubit_map: A dict that shows how each FermionicMode is mapped to a qubit. - ordered_qubits: An ordered sequence of qubits. + params: The trial wavefunction parameters. + ansatz_circuit: The circuit specifying the (pp) ansatz. + superposition_circuit: The superposition circuit. + hf_energy: The Hartree--Fock energy for the underlying molecule. + ansatze_energy: The expected energy optimized wavefunction. + fci_energy: The exact ground state energy of the underlying molecule. + one_body_basis_change_mat: The one-body basis change matrix. + one_body_params: The one-body variational parameters. + two_body_params: The two-body variational parameters. """ - qubits = list(mode_qubit_map.values()) - assert len(qubits) == len(ordered_qubits) - - # We sort the key: value pairs by the order of the values (qubits) in - # ordered_qubits. - sorted_mapping = list(mode_qubit_map.items()) - sorted_mapping.sort(key=lambda x: ordered_qubits.index(x[1])) - - remapping_map = {} - for i, (mode, _) in enumerate(sorted_mapping): - openfermion_index = 2 * mode.orb_ind + (0 if mode.spin == "a" else 1) - remapping_map[openfermion_index] = i - - def remapper(index: int, _: int) -> int: - """A function that maps from the old index to the new one. - - The _ argument is because it's expected by openfermion.reorder""" - return remapping_map[index] - - return remapper - - -def _get_pp_plus_gate_generators( - *, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...], do_pp: bool = True -) -> List[of.FermionOperator]: - heuristic_gate_generators = get_heuristic_gate_generators(n_elec, heuristic_layers) - if not do_pp: - return heuristic_gate_generators - - n_pairs = n_elec // 2 - pair_gate_generators = get_pair_hopping_gate_generators(n_pairs, n_elec) - return pair_gate_generators + heuristic_gate_generators - - -def _get_ansatz_qubit_wf( - *, ansatz_circuit: cirq.Circuit, ordered_qubits: Sequence[cirq.Qid] -): - return cirq.final_state_vector( - ansatz_circuit, qubit_order=list(ordered_qubits), dtype=np.complex128 - ) - - -def _get_superposition_wf( - *, superposition_circuit: cirq.Circuit, ordered_qubits: Sequence[cirq.Qid] -) -> np.ndarray: - return cirq.final_state_vector( - superposition_circuit, qubit_order=list(ordered_qubits), dtype=np.complex128 - ) - - -@attrs.frozen -class TrialWavefunctionData(data.Data): - """Class for storing a trial wavefunction's data.""" params: PerfectPairingPlusTrialWavefunctionParams ansatz_circuit: cirq.Circuit @@ -388,1379 +208,4 @@ class TrialWavefunctionData(data.Data): two_body_params: np.ndarray = attrs.field(converter=_to_numpy) def _json_dict_(self): - # return cirq.dataclass_json_dict(self) return attrs.asdict(self) - - -def get_and_check_energy( - *, - hamiltonian_data: hamiltonian.HamiltonianData, - ansatz_circuit: cirq.Circuit, - one_body_params: np.ndarray, - two_body_params: np.ndarray, - one_body_basis_change_mat: np.ndarray, - params: PerfectPairingPlusTrialWavefunctionParams, -) -> Tuple[float, float]: - ansatz_qubit_wf = _get_ansatz_qubit_wf( - ansatz_circuit=ansatz_circuit, - ordered_qubits=params.qubits_jordan_wigner_ordered, - ) - - fqe_ham, e_core, sparse_ham = get_rotated_hamiltonians( - hamiltonian_data=hamiltonian_data, - one_body_basis_change_mat=one_body_basis_change_mat, - mode_qubit_map=params.mode_qubit_map, - ordered_qubits=params.qubits_jordan_wigner_ordered, - ) - - initial_wf = fqe.Wavefunction([[params.n_elec, 0, params.n_orb]]) - initial_wf.set_wfn(strategy="hartree-fock") - - hf_energy = initial_wf.expectationValue(fqe_ham) + e_core - - fqe_wf, unrotated_fqe_wf = get_evolved_wf( - one_body_params=one_body_params, - two_body_params=two_body_params, - wf=initial_wf, - gate_generators=_get_pp_plus_gate_generators( - n_elec=params.n_elec, - heuristic_layers=params.heuristic_layers, - do_pp=params.do_pp, - ), - n_orb=params.n_orb, - restricted=params.restricted, - initial_orbital_rotation=params.initial_orbital_rotation, - ) - - ansatz_energy = get_energy_and_check_sanity( - circuit_wf=ansatz_qubit_wf, - fqe_wf=fqe_wf, - unrotated_fqe_wf=unrotated_fqe_wf, - fqe_ham=fqe_ham, - sparse_ham=sparse_ham, - e_core=e_core, - mode_qubit_map=params.mode_qubit_map, - ordered_qubits=params.qubits_jordan_wigner_ordered, - ) - - return ansatz_energy, hf_energy - - -def build_pp_plus_trial_wavefunction( - params: PerfectPairingPlusTrialWavefunctionParams, - *, - dependencies: Dict[data.Params, data.Data], - do_print: bool = False, -) -> TrialWavefunctionData: - """Builds a TrialWavefunctionData from a TrialWavefunctionParams""" - - if do_print: - print("Building Trial Wavefunction") - np.random.seed(params.seed) - hamiltonian_data = dependencies[params.hamiltonian_params] - assert isinstance(hamiltonian_data, hamiltonian.HamiltonianData) - - assert ( - params.n_orb == params.n_elec - ) ## Necessary for perfect pairing wavefunction to make sense. - - if params.do_optimization: - ( - one_body_params, - two_body_params, - one_body_basis_change_mat, - ) = get_pp_plus_params( - hamiltonian_data=hamiltonian_data, - restricted=params.restricted, - random_parameter_scale=params.random_parameter_scale, - initial_orbital_rotation=params.initial_orbital_rotation, - heuristic_layers=params.heuristic_layers, - do_pp=params.do_pp, - n_optimization_restarts=params.n_optimization_restarts, - do_print=do_print, - use_fast_gradients=params.use_fast_gradients, - ) - else: - if ( - params.initial_two_body_qchem_amplitudes is None - or params.initial_orbital_rotation is not None - ): - raise NotImplementedError("TODO: Implement whatever isn't finished here.") - - n_one_body_params = params.n_orb * (params.n_orb - 1) - one_body_params = np.zeros(n_one_body_params) - one_body_basis_change_mat = np.diag(np.ones(params.n_orb * 2)) - two_body_params = get_two_body_params_from_qchem_amplitudes( - params.initial_two_body_qchem_amplitudes - ) - - (superposition_circuit, ansatz_circuit) = get_circuits( - two_body_params=two_body_params, - n_orb=params.n_orb, - n_elec=params.n_elec, - heuristic_layers=params.heuristic_layers, - ) - - ansatz_energy, hf_energy = get_and_check_energy( - hamiltonian_data=hamiltonian_data, - ansatz_circuit=ansatz_circuit, - params=params, - one_body_params=one_body_params, - two_body_params=two_body_params, - one_body_basis_change_mat=one_body_basis_change_mat, - ) - - return TrialWavefunctionData( - params=params, - ansatz_circuit=ansatz_circuit, - superposition_circuit=superposition_circuit, - hf_energy=hf_energy, - ansatz_energy=ansatz_energy, - fci_energy=hamiltonian_data.e_fci, - one_body_basis_change_mat=one_body_basis_change_mat, - one_body_params=one_body_params, - two_body_params=two_body_params, - ) - - -def get_rotated_hamiltonians( - *, - hamiltonian_data: hamiltonian.HamiltonianData, - one_body_basis_change_mat: np.ndarray, - mode_qubit_map: Mapping[FermionicMode, cirq.Qid], - ordered_qubits: Sequence[cirq.Qid], -) -> Tuple[fqe.hamiltonians.hamiltonian.Hamiltonian, float, csc_matrix]: - """A helper method that gets the hamiltonians in the basis of the trial_wf. - - Returns: - The hamiltonian in FQE form, minus a constant energy shift. - The constant part of the Hamiltonian missing from the FQE Hamiltonian. - The qubit Hamiltonian as a sparse matrix. - """ - n_qubits = len(mode_qubit_map) - - fqe_ham = hamiltonian_data.get_restricted_fqe_hamiltonian() - e_core = hamiltonian_data.e_core - - mol_ham = hamiltonian_data.get_molecular_hamiltonian() - mol_ham.rotate_basis(one_body_basis_change_mat) - fermion_operator_ham = of.get_fermion_operator(mol_ham) - - reorder_func = _get_reorder_func( - mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits - ) - fermion_operator_ham_qubit_ordered = of.reorder( - fermion_operator_ham, reorder_func, num_modes=n_qubits - ) - - sparse_qubit_ham = of.get_sparse_operator(fermion_operator_ham_qubit_ordered) - - return fqe_ham, e_core, sparse_qubit_ham - - -def get_energy_and_check_sanity( - *, - circuit_wf: np.ndarray, - fqe_wf: fqe_wfn.Wavefunction, - unrotated_fqe_wf: fqe_wfn.Wavefunction, - fqe_ham: fqe.hamiltonians.hamiltonian.Hamiltonian, - sparse_ham: csc_matrix, - e_core: float, - mode_qubit_map: Mapping[FermionicMode, cirq.Qid], - ordered_qubits: Sequence[cirq.Qid], -) -> float: - """A method that checks for consistency and returns the ansatz energy.""" - - unrotated_fqe_wf_as_cirq = convert_fqe_wf_to_cirq( - fqe_wf=unrotated_fqe_wf, - mode_qubit_map=mode_qubit_map, - ordered_qubits=ordered_qubits, - ) - ansatz_energy = np.real_if_close( - (np.conj(circuit_wf) @ sparse_ham @ circuit_wf) - ).item() - assert isinstance(ansatz_energy, float) - - fqe_energy = np.real(fqe_wf.expectationValue(fqe_ham) + e_core) - print(csc_matrix(circuit_wf)) - print(csc_matrix(unrotated_fqe_wf_as_cirq)) - np.testing.assert_array_almost_equal(ansatz_energy, fqe_energy) - np.testing.assert_array_almost_equal( - circuit_wf, unrotated_fqe_wf_as_cirq, decimal=5 - ) - return ansatz_energy - - -def get_4_qubit_fermion_qubit_map() -> Dict[int, cirq.GridQubit]: - """A helper function that provides the fermion qubit map for 4 qubits. - - We map the fermionic orbitals to grid qubits like so: - 3 1 - 2 0 - """ - fermion_index_to_qubit_map = { - 2: cirq.GridQubit(0, 0), - 3: cirq.GridQubit(1, 0), - 0: cirq.GridQubit(0, 1), - 1: cirq.GridQubit(1, 1), - } - - return fermion_index_to_qubit_map - - -def get_4_qubit_pp_circuits( - *, two_body_params: np.ndarray, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...] -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. - - We map the fermionic orbitals to grid qubits like so: - 3 1 - 2 0 - """ - assert n_elec == 2 - - fermion_index_to_qubit_map = get_4_qubit_fermion_qubit_map() - geminal_gate = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[0], inline_control=True - ) - - ansatz_circuit = cirq.Circuit( - cirq.decompose( - geminal_gate.on( - fermion_index_to_qubit_map[0], - fermion_index_to_qubit_map[1], - fermion_index_to_qubit_map[2], - fermion_index_to_qubit_map[3], - ) - ) - ) - - heuristic_layer_circuit = get_heuristic_circuit( - heuristic_layers, n_elec, two_body_params[1:], fermion_index_to_qubit_map - ) - - ansatz_circuit += heuristic_layer_circuit - - indicator = fermion_index_to_qubit_map[2] - superposition_circuit = cirq.Circuit([cirq.H(indicator) + ansatz_circuit]) - ansatz_circuit = cirq.Circuit([cirq.X(indicator) + ansatz_circuit]) - - return superposition_circuit, ansatz_circuit - - -def get_8_qubit_fermion_qubit_map(): - """A helper function that provides the fermion qubit map for 8 qubits. - - We map the fermionic orbitals to grid qubits like so: - 3 5 1 7 - 2 4 0 6 - """ - - # Linear connectivity is fine. - # This ordering is dictated by the way we specify perfect pairing (i) - # Elsewhere we generate the perfect pairing parameters using a specific - # convention for how we index the FermionOperators (which is itself) - # partly dictated by the OpenFermion conventions. Here we choose a - # mapping between the indices of these FermionOperators and the qubits in our - # grid that allows for the perfect pairing pairs to be in squares of four. - fermion_index_to_qubit_map = { - 2: cirq.GridQubit(0, 0), - 3: cirq.GridQubit(1, 0), - 4: cirq.GridQubit(0, 1), - 5: cirq.GridQubit(1, 1), - 0: cirq.GridQubit(0, 2), - 1: cirq.GridQubit(1, 2), - 6: cirq.GridQubit(0, 3), - 7: cirq.GridQubit(1, 3), - } - - return fermion_index_to_qubit_map - - -def get_8_qubit_circuits( - *, two_body_params: np.ndarray, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...] -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. - - We map the fermionic orbitals to grid qubits like so: - 3 5 1 7 - 2 4 0 6 - """ - fermion_index_to_qubit_map = get_8_qubit_fermion_qubit_map() - - geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[0], inline_control=True - ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[1], inline_control=True - ) - - # We'll add the initial bit flips later. - ansatz_circuit = cirq.Circuit( - cirq.decompose( - geminal_gate_1.on( - fermion_index_to_qubit_map[2], - fermion_index_to_qubit_map[3], - fermion_index_to_qubit_map[4], - fermion_index_to_qubit_map[5], - ) - ), - cirq.decompose( - geminal_gate_2.on( - fermion_index_to_qubit_map[0], - fermion_index_to_qubit_map[1], - fermion_index_to_qubit_map[6], - fermion_index_to_qubit_map[7], - ) - ), - ) - - heuristic_layer_circuit = get_heuristic_circuit( - heuristic_layers, n_elec, two_body_params[2:], fermion_index_to_qubit_map - ) - - ansatz_circuit += heuristic_layer_circuit - - superposition_circuit = ( - cirq.Circuit( - [ - cirq.H(fermion_index_to_qubit_map[0]), - cirq.CNOT(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[6]), - cirq.SWAP(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[4]), - ] - ) - + ansatz_circuit - ) - - ansatz_circuit = ( - cirq.Circuit( - [ - cirq.X(fermion_index_to_qubit_map[4]), - cirq.X(fermion_index_to_qubit_map[6]), - ] - ) - + ansatz_circuit - ) - - return superposition_circuit, ansatz_circuit - - -def get_12_qubit_fermion_qubit_map(): - """A helper function that provides the fermion qubit map for 12 qubits. - - We map the fermionic orbitals to grid qubits like so: - 5 7 3 9 1 11 - 4 6 2 8 0 10 - """ - - # Linear connectivity is fine. - # This ordering is dictated by the way we specify perfect pairing (i) - # Elsewhere we generate the perfect pairing parameters using a specific - # convention for how we index the FermionOperators (which is itself) - # partly dictated by the OpenFermion conventions. Here we choose a - # mapping between the indices of these FermionOperators and the qubits in our - # grid that allows for the perfect pairing pairs to be in squares of four. - fermion_index_to_qubit_map = { - 4: cirq.GridQubit(0, 0), - 5: cirq.GridQubit(1, 0), - 6: cirq.GridQubit(0, 1), - 7: cirq.GridQubit(1, 1), - 2: cirq.GridQubit(0, 2), - 3: cirq.GridQubit(1, 2), - 8: cirq.GridQubit(0, 3), - 9: cirq.GridQubit(1, 3), - 0: cirq.GridQubit(0, 4), - 1: cirq.GridQubit(1, 4), - 10: cirq.GridQubit(0, 5), - 11: cirq.GridQubit(1, 5), - } - - return fermion_index_to_qubit_map - - -def get_12_qubit_circuits( - *, two_body_params: np.ndarray, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...] -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. - - We map the fermionic orbitals to grid qubits like so: - 5 7 3 9 1 11 - 4 6 2 8 0 10 - """ - - fermion_index_to_qubit_map = get_12_qubit_fermion_qubit_map() - - geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[0], inline_control=True - ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[1], inline_control=True - ) - geminal_gate_3 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[2], inline_control=True - ) - - # We'll add the initial bit flips later. - ansatz_circuit = cirq.Circuit( - cirq.decompose( - geminal_gate_1.on( - fermion_index_to_qubit_map[4], - fermion_index_to_qubit_map[5], - fermion_index_to_qubit_map[6], - fermion_index_to_qubit_map[7], - ) - ), - cirq.decompose( - geminal_gate_2.on( - fermion_index_to_qubit_map[2], - fermion_index_to_qubit_map[3], - fermion_index_to_qubit_map[8], - fermion_index_to_qubit_map[9], - ) - ), - cirq.decompose( - geminal_gate_3.on( - fermion_index_to_qubit_map[0], - fermion_index_to_qubit_map[1], - fermion_index_to_qubit_map[10], - fermion_index_to_qubit_map[11], - ) - ), - ) - - heuristic_layer_circuit = get_heuristic_circuit( - heuristic_layers, n_elec, two_body_params[3:], fermion_index_to_qubit_map - ) - - ansatz_circuit += heuristic_layer_circuit - - superposition_circuit = ( - cirq.Circuit( - [ - cirq.H(fermion_index_to_qubit_map[8]), - cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[0]), - cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[2]), - cirq.SWAP( - fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[10] - ), - cirq.SWAP(fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[6]), - ] - ) - + ansatz_circuit - ) - - ansatz_circuit = ( - cirq.Circuit( - [ - cirq.X(fermion_index_to_qubit_map[6]), - cirq.X(fermion_index_to_qubit_map[8]), - cirq.X(fermion_index_to_qubit_map[10]), - ] - ) - + ansatz_circuit - ) - - return superposition_circuit, ansatz_circuit - - -def get_16_qubit_fermion_qubit_map(): - """A helper function that provides the fermion qubit map for 16 qubits. - - We map the fermionic orbitals to grid qubits like so: - 7 9 5 11 3 13 1 15 - 6 8 4 10 2 12 0 14 - """ - - # Linear connectivity is fine. - # This ordering is dictated by the way we specify perfect pairing (i) - # Elsewhere we generate the perfect pairing parameters using a specific - # convention for how we index the FermionOperators (which is itself) - # partly dictated by the OpenFermion conventions. Here we choose a - # mapping between the indices of these FermionOperators and the qubits in our - # grid that allows for the perfect pairing pairs to be in squares of four. - fermion_index_to_qubit_map = { - 6: cirq.GridQubit(0, 0), - 7: cirq.GridQubit(1, 0), - 8: cirq.GridQubit(0, 1), - 9: cirq.GridQubit(1, 1), - 4: cirq.GridQubit(0, 2), - 5: cirq.GridQubit(1, 2), - 10: cirq.GridQubit(0, 3), - 11: cirq.GridQubit(1, 3), - 2: cirq.GridQubit(0, 4), - 3: cirq.GridQubit(1, 4), - 12: cirq.GridQubit(0, 5), - 13: cirq.GridQubit(1, 5), - 0: cirq.GridQubit(0, 6), - 1: cirq.GridQubit(1, 6), - 14: cirq.GridQubit(0, 7), - 15: cirq.GridQubit(1, 7), - } - - return fermion_index_to_qubit_map - - -def get_16_qubit_circuits( - *, two_body_params: np.ndarray, n_elec: int, heuristic_layers: Tuple[LayerSpec, ...] -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. - - We map the fermionic orbitals to grid qubits like so: - 7 9 5 11 3 13 1 15 - 6 8 4 10 2 12 0 14 - """ - fermion_index_to_qubit_map = get_16_qubit_fermion_qubit_map() - - geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[0], inline_control=True - ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[1], inline_control=True - ) - geminal_gate_3 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[2], inline_control=True - ) - geminal_gate_4 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[3], inline_control=True - ) - - # We'll add the initial bit flips later. - ansatz_circuit = cirq.Circuit( - cirq.decompose( - geminal_gate_1.on( - fermion_index_to_qubit_map[6], - fermion_index_to_qubit_map[7], - fermion_index_to_qubit_map[8], - fermion_index_to_qubit_map[9], - ) - ), - cirq.decompose( - geminal_gate_2.on( - fermion_index_to_qubit_map[4], - fermion_index_to_qubit_map[5], - fermion_index_to_qubit_map[10], - fermion_index_to_qubit_map[11], - ) - ), - cirq.decompose( - geminal_gate_3.on( - fermion_index_to_qubit_map[2], - fermion_index_to_qubit_map[3], - fermion_index_to_qubit_map[12], - fermion_index_to_qubit_map[13], - ) - ), - cirq.decompose( - geminal_gate_4.on( - fermion_index_to_qubit_map[0], - fermion_index_to_qubit_map[1], - fermion_index_to_qubit_map[14], - fermion_index_to_qubit_map[15], - ) - ), - ) - - heuristic_layer_circuit = get_heuristic_circuit( - heuristic_layers, n_elec, two_body_params[4:], fermion_index_to_qubit_map - ) - - ansatz_circuit += heuristic_layer_circuit - - superposition_circuit = ( - cirq.Circuit( - [ - cirq.H(fermion_index_to_qubit_map[10]), - cirq.CNOT( - fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[2] - ), - cirq.SWAP( - fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[12] - ), - cirq.CNOT( - fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[4] - ), - cirq.CNOT( - fermion_index_to_qubit_map[12], fermion_index_to_qubit_map[0] - ), - cirq.SWAP(fermion_index_to_qubit_map[4], fermion_index_to_qubit_map[8]), - cirq.SWAP( - fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[14] - ), - ] - ) - + ansatz_circuit - ) - - ansatz_circuit = ( - cirq.Circuit( - [ - cirq.X(fermion_index_to_qubit_map[8]), - cirq.X(fermion_index_to_qubit_map[10]), - cirq.X(fermion_index_to_qubit_map[12]), - cirq.X(fermion_index_to_qubit_map[14]), - ] - ) - + ansatz_circuit - ) - - return superposition_circuit, ansatz_circuit - - -def get_circuits( - *, - two_body_params: np.ndarray, - # from wf_params: - n_orb: int, - n_elec: int, - heuristic_layers: Tuple[LayerSpec, ...], -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A function that runs a specialized method to get the ansatz circuits.""" - - # TODO(?): Just input one of these quantities. - if n_orb != n_elec: - raise ValueError("n_orb must equal n_elec.") - - circ_funcs = { - 2: get_4_qubit_pp_circuits, - 4: get_8_qubit_circuits, - 6: get_12_qubit_circuits, - 8: get_16_qubit_circuits, - } - try: - circ_func = circ_funcs[n_orb] - except KeyError: - raise NotImplementedError(f"No circuits for n_orb = {n_orb}") - - return circ_func( - two_body_params=two_body_params, - n_elec=n_elec, - heuristic_layers=heuristic_layers, - ) - - -def get_two_body_params_from_qchem_amplitudes( - qchem_amplitudes: np.ndarray, -) -> np.ndarray: - """Translates perfect pairing amplitudes from qchem to rotation angles. - - qchem style: 1 |1100> + t_i |0011> - our style: cos(\theta_i) |1100> + sin(\theta_i) |0011> - """ - - two_body_params = np.arccos(1 / np.sqrt(1 + qchem_amplitudes**2)) * np.sign( - qchem_amplitudes - ) - - # Numpy casts the array improperly to a float when we only have one parameter. - two_body_params = np.atleast_1d(two_body_params) - - return two_body_params - - -#################### Here be dragons.########################################### - - -def convert_fqe_wf_to_cirq( - fqe_wf: fqe_wfn.Wavefunction, - mode_qubit_map: Mapping[FermionicMode, cirq.Qid], - ordered_qubits: Sequence[cirq.Qid], -) -> np.ndarray: - """Converts an FQE wavefunction to one on qubits with a particular ordering. - - Args: - fqe_wf: The FQE wavefunction. - mode_qubit_map: A mapping from fermion modes to cirq qubits. - ordered_qubits: - """ - n_qubits = len(mode_qubit_map) - fermion_op = fqe.openfermion_utils.fqe_to_fermion_operator(fqe_wf) - - reorder_func = _get_reorder_func( - mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits - ) - fermion_op = of.reorder(fermion_op, reorder_func, num_modes=n_qubits) - - qubit_op = of.jordan_wigner(fermion_op) - - return fqe.qubit_wavefunction_from_vacuum( - qubit_op, list(cirq.LineQubit.range(n_qubits)) - ) - - -def get_one_body_cluster_coef( - params: np.ndarray, n_orb: int, restricted: bool -) -> npt.NDArray[np.complex128]: - if restricted: - one_body_cluster_op = np.zeros((n_orb, n_orb), dtype=np.complex128) - else: - one_body_cluster_op = np.zeros((2 * n_orb, 2 * n_orb), dtype=np.complex128) - param_num = 0 - - for i in range(n_orb): - for j in range(i): - one_body_cluster_op[i, j] = params[param_num] - one_body_cluster_op[j, i] = -params[param_num] - param_num += 1 - - if not restricted: - for i in range(n_orb, 2 * n_orb): - for j in range(n_orb, i): - one_body_cluster_op[i, j] = params[param_num] - one_body_cluster_op[j, i] = -params[param_num] - param_num += 1 - - return one_body_cluster_op - - -def get_evolved_wf( - one_body_params: np.ndarray, - two_body_params: np.ndarray, - wf: fqe.Wavefunction, - gate_generators: List[of.FermionOperator], - n_orb: int, - restricted: bool = True, - initial_orbital_rotation: Optional[np.ndarray] = None, -) -> Tuple[fqe.Wavefunction, fqe.Wavefunction]: - """Get the wavefunction evaluated for this set of variational parameters. - - Args: - one_body_params: The variational parameters for the one-body terms in the ansatz. - two_body_params: The variational parameters for the two-body terms in the ansatz. - wf: The FQE wavefunction to evolve. - gate_generators: The generators of the two-body interaction terms. - n_orb: The number of spatial orbitals. - restricted: Whether the ansatz is restricted or not. - initial_orbital_rotation: Any initial orbital rotation to prepend to the circuit. - - Returs: - rotated_wf: the evolved wavefunction - wf: The original wavefunction - """ - param_num = 0 - for gate_generator in gate_generators: - wf = wf.time_evolve(two_body_params[param_num], gate_generator) - param_num += 1 - - one_body_cluster_op = get_one_body_cluster_coef( - one_body_params, n_orb, restricted=restricted - ) - - if restricted: - one_body_ham = fqe.get_restricted_hamiltonian((1j * one_body_cluster_op,)) - else: - one_body_ham = fqe.get_sso_hamiltonian((1j * one_body_cluster_op,)) - - rotated_wf = wf.time_evolve(1.0, one_body_ham) - - if initial_orbital_rotation is not None: - rotated_wf = fqe.algorithm.low_rank.evolve_fqe_givens( - rotated_wf, initial_orbital_rotation - ) - - return rotated_wf, wf - - -def get_pair_hopping_gate_generators( - n_pairs: int, n_elec: int -) -> List[of.FermionOperator]: - """Get the generators of the pair-hopping unitaries. - - Args: - n_pairs: The number of pair coupling terms. - n_elec: The total number of electrons. - - Returns: - A list of gate generators - """ - gate_generators = [] - for pair in range(n_pairs): - to_a = n_elec + 2 * pair - to_b = n_elec + 2 * pair + 1 - from_a = n_elec - 2 * pair - 2 - from_b = n_elec - 2 * pair - 1 - - fop_string = f"{to_b} {to_a} {from_b}^ {from_a}^" - - gate_generator = of.FermionOperator(fop_string, 1.0) - gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) - gate_generators.append(gate_generator) - - return gate_generators - - -def get_indices_heuristic_layer_in_pair(n_elec: int) -> Iterator[Tuple[int, int]]: - """Get the indicies for the heuristic layers. - - Args: - n_elec: The number of electrons - - Returns: - An iterator of the indices - """ - n_pairs = n_elec // 2 - - for pair in range(n_pairs): - to_a = n_elec + 2 * pair - to_b = n_elec + 2 * pair + 1 - from_a = n_elec - 2 * pair - 2 - from_b = n_elec - 2 * pair - 1 - yield (from_a, to_a) - yield (from_b, to_b) - - -def get_indices_heuristic_layer_cross_pair(n_elec) -> Iterator[Tuple[int, int]]: - """Indices that couple adjacent pairs. - - Args: - n_elec: The number of electrons - - Returns: - An iterator of the indices - """ - n_pairs = n_elec // 2 - - for pair in range(n_pairs - 1): - to_a = n_elec + 2 * pair - to_b = n_elec + 2 * pair + 1 - from_next_a = n_elec - 2 * (pair + 1) - 2 - from_next_b = n_elec - 2 * (pair + 1) - 1 - yield (to_a, from_next_a) - yield (to_b, from_next_b) - - -def get_indices_heuristic_layer_cross_spin(n_elec) -> Iterator[Tuple[int, int]]: - """Get indices that couple the two spin sectors. - - Args: - n_elec: The number of electrons - - Returns: - An iterator of the indices that couple spin sectors. - """ - n_pairs = n_elec // 2 - - for pair in range(n_pairs): - to_a = n_elec + 2 * pair - to_b = n_elec + 2 * pair + 1 - from_a = n_elec - 2 * pair - 2 - from_b = n_elec - 2 * pair - 1 - yield (to_a, to_b) - yield (from_a, from_b) - - -def get_charge_charge_generator(indices: Tuple[int, int]) -> of.FermionOperator: - """Returns the generator for density evolution between the indices - - Args: - indices: The indices to for charge-charge terms.:w - - Returns: - The generator for density evolution for this pair of electrons. - """ - - fop_string = "{:d}^ {:d} {:d}^ {:d}".format( - indices[0], indices[0], indices[1], indices[1] - ) - gate_generator = of.FermionOperator(fop_string, 1.0) - - return gate_generator - - -def get_charge_charge_gate( - qubits: Tuple[cirq.Qid, ...], param: float -) -> cirq.Operation: - """Get the cirq charge-charge gate. - - Args: - qubits: Two qubits you want to apply the gate to. - param: The parameter for the charge-charge interaction. - - Returns: - The charge-charge gate. - """ - return cirq.CZ(qubits[0], qubits[1]) ** (-param / np.pi) - - -def get_givens_generator(indices: Tuple[int, int]) -> of.FermionOperator: - """Returns the generator for givens rotation between two orbitals. - - Args: - indices: The two indices for the givens rotation. - - Returns: - The givens generator for evolution for this pair of electrons. - """ - - fop_string = "{:d}^ {:d}".format(indices[0], indices[1]) - gate_generator = of.FermionOperator(fop_string, 1.0) - gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) - - return gate_generator - - -def get_givens_gate(qubits: Tuple[cirq.Qid, ...], param: float) -> cirq.Operation: - """Get a the givens rotation gate on two qubits. - - Args: - qubits: The two qubits to apply the gate to. - param: The parameter for the givens rotation. - - Returns: - The givens rotation gate. - """ - return cirq.givens(param).on(qubits[0], qubits[1]) - - -def get_layer_indices(layer_spec: LayerSpec, n_elec: int) -> List[Tuple[int, int]]: - """Get the indices for the heuristic layers. - - Args: - layer_spec: The layer specification. - n_elec: The number of electrons. - - Returns: - A list of indices for the layer. - """ - indices_generators = { - "in_pair": get_indices_heuristic_layer_in_pair(n_elec), - "cross_pair": get_indices_heuristic_layer_cross_pair(n_elec), - "cross_spin": get_indices_heuristic_layer_cross_spin(n_elec), - } - indices_generator = indices_generators[layer_spec.layout] - - return [indices for indices in indices_generator] - - -def get_layer_gates( - layer_spec: LayerSpec, - n_elec: int, - params: np.ndarray, - fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], -) -> List[cirq.Operation]: - """Gets the gates for a hardware efficient layer of the ansatz. - - Args: - layer_spec: The layer specification. - n_elec: The number of electrons. - params: The variational parameters for the hardware efficient gate layer. - fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. - - Returns: - A list of gates for the layer. - """ - - indices_list = get_layer_indices(layer_spec, n_elec) - - gate_funcs = {"givens": get_givens_gate, "charge_charge": get_charge_charge_gate} - gate_func = gate_funcs[layer_spec.base_gate] - - gates = [] - for indices, param in zip(indices_list, params): - qubits = tuple(fermion_index_to_qubit_map[ind] for ind in indices) - gates.append(gate_func(qubits, param)) - - return gates - - -def get_layer_generators( - layer_spec: LayerSpec, n_elec: int -) -> List[of.FermionOperator]: - """Gets the generators for rotations in a hardware efficient layer of the ansatz. - - Args: - layer_spec: The layer specification. - n_elec: The number of electrons. - - Returns: - A list of generators for the layers. - """ - - indices_list = get_layer_indices(layer_spec, n_elec) - - gate_funcs = { - "givens": get_givens_generator, - "charge_charge": get_charge_charge_generator, - } - gate_func = gate_funcs[layer_spec.base_gate] - - return [gate_func(indices) for indices in indices_list] - - -def get_heuristic_gate_generators( - n_elec: int, layer_specs: Sequence[LayerSpec] -) -> List[of.FermionOperator]: - """Get gate generators for the heuristic ansatz. - - Args: - n_elec: The number of electrons. - layer_specs: The layer specifications. - - Returns: - A list of generators for the layers. - """ - gate_generators = [] - - for layer_spec in layer_specs: - gate_generators += get_layer_generators(layer_spec, n_elec) - - return gate_generators - - -def get_heuristic_circuit( - layer_specs: Sequence[LayerSpec], - n_elec: int, - params: np.ndarray, - fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], -) -> cirq.Circuit: - """Get a circuit for the heuristic ansatz. - - Args: - layer_specs: The layer specs for the heuristic layers. - n_elec: The number of electrons. - params: The variational parameters for the circuit. - fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. - - Returns: - A circuit for the heuristic ansatz. - """ - gates: List[cirq.Operation] = [] - - for layer_spec in layer_specs: - params_slice = params[len(gates) :] - gates += get_layer_gates( - layer_spec, n_elec, params_slice, fermion_index_to_qubit_map - ) - - return cirq.Circuit(gates) - - -def orbital_rotation_gradient_matrix( - generator_mat: np.ndarray, a: int, b: int -) -> np.ndarray: - """The gradient of the orbital rotation unitary with respect to its parameters. - - Args: - generator_mat: The orbital rotation one-body generator matrix. - a, b: row and column indices corresponding to the location in the matrix - of the parameter we wish to find the gradient with respect to. - - Returns: - The orbital rotation matrix gradient wrt theta_{a, b}. Corresponds to - expression in G15 of https://arxiv.org/abs/2004.04174. - """ - w_full, v_full = np.linalg.eigh(-1j * generator_mat) - eigs_diff = np.zeros((w_full.shape[0], w_full.shape[0]), dtype=np.complex128) - for i, j in itertools.product(range(w_full.shape[0]), repeat=2): - if np.isclose(abs(w_full[i] - w_full[j]), 0): - eigs_diff[i, j] = 1 - else: - eigs_diff[i, j] = (np.exp(1j * (w_full[i] - w_full[j])) - 1) / ( - 1j * (w_full[i] - w_full[j]) - ) - - Y_full = np.zeros_like(v_full, dtype=np.complex128) - if a == b: - Y_full[a, b] = 0 - else: - Y_full[a, b] = 1.0 - Y_full[b, a] = -1.0 - - Y_kl_full = v_full.conj().T @ Y_full @ v_full - # now rotate Y_{kl} * (exp(i(l_{k} - l_{l})) - 1) / (i(l_{k} - l_{l})) - # into the original basis - pre_matrix_full = v_full @ (eigs_diff * Y_kl_full) @ v_full.conj().T - - return pre_matrix_full - - -def evaluate_gradient_and_cost_function( - initial_wf: fqe.Wavefunction, - fqe_ham: fqe_hams.RestrictedHamiltonian, - n_orb: int, - one_body_params: npt.NDArray[np.float64], - two_body_params: npt.NDArray[np.float64], - gate_generators: List[of.FermionOperator], - restricted: bool, - e_core: float, -) -> Tuple[float, npt.NDArray[np.float64]]: - """Evaluate gradient and cost function for optimization. - - Args: - initial_wf: Initial state (typically Hartree--Fock). - fqe_ham: The restricted Hamiltonian in FQE format. - n_orb: The number of spatial orbitals. - one_body_params: The parameters of the single-particle rotations. - two_body_params: The parameters for the two-particle terms. - gate_generators: The generators for the two-particle terms. - retricted: Whether the single-particle rotations are restricted (True) - or unrestricted (False). Unrestricted implies different parameters - for the alpha- and beta-spin rotations. - - Returns: - cost_val: The cost function (total energy) evaluated for the input wavefunction parameters. - grad: An array of gradients with respect to the one- and two-body - parameters. The first n_orb * (n_orb + 1) // 2 parameters correspond to - the one-body gradients. - """ - phi = get_evolved_wf( - one_body_params, - two_body_params, - initial_wf, - gate_generators, - n_orb, - restricted=restricted, - )[0] - lam = copy.deepcopy(phi) - lam = lam.apply(fqe_ham) - cost_val = fqe.vdot(lam, phi) + e_core - - # 1body - one_body_cluster_op = get_one_body_cluster_coef( - one_body_params, n_orb, restricted=restricted - ) - tril = np.tril_indices(n_orb, k=-1) - if restricted: - one_body_ham = fqe.get_restricted_hamiltonian((-1j * one_body_cluster_op,)) - else: - one_body_ham = fqe.get_sso_hamiltonian((-1j * one_body_cluster_op,)) - # Apply U1b^{dag} - phi.time_evolve(1, one_body_ham, inplace=True) - lam.time_evolve(1, one_body_ham, inplace=True) - one_body_grad = np.zeros_like(one_body_params) - n_one_body_params = len(one_body_params) - grad_position = n_one_body_params - 1 - for iparam in range(len(one_body_params)): - mu_state = copy.deepcopy(phi) - pidx = n_one_body_params - iparam - 1 - pidx_spin = 0 if restricted else pidx // (n_one_body_params // 2) - pidx_spat = pidx if restricted else pidx - (n_one_body_params // 2) * pidx_spin - p, q = (tril[0][pidx_spat], tril[1][pidx_spat]) - p += n_orb * pidx_spin - q += n_orb * pidx_spin - pre_matrix = orbital_rotation_gradient_matrix(-one_body_cluster_op, p, q) - assert of.is_hermitian(1j * pre_matrix) - if restricted: - fqe_quad_ham_pre = fqe.get_restricted_hamiltonian((pre_matrix,)) - else: - fqe_quad_ham_pre = fqe.get_sso_hamiltonian((pre_matrix,)) - mu_state = mu_state.apply(fqe_quad_ham_pre) - one_body_grad[grad_position] = 2 * fqe.vdot(lam, mu_state).real - grad_position -= 1 - # Get two-body contributions - two_body_grad = np.zeros(len(two_body_params)) - for pidx in reversed(range(len(gate_generators))): - mu = copy.deepcopy(phi) - mu = mu.apply(gate_generators[pidx]) - two_body_grad[pidx] = -np.real(2 * 1j * (fqe.vdot(lam, mu))) - phi = phi.time_evolve(-two_body_params[pidx], gate_generators[pidx]) - lam = lam.time_evolve(-two_body_params[pidx], gate_generators[pidx]) - - return cost_val, np.concatenate((two_body_grad, one_body_grad)) - - -def get_pp_plus_params( - *, - hamiltonian_data: hamiltonian.HamiltonianData, - restricted: bool = False, - random_parameter_scale: float = 1.0, - initial_orbital_rotation: Optional[np.ndarray] = None, - heuristic_layers: Tuple[LayerSpec, ...], - do_pp: bool = True, - n_optimization_restarts: int = 1, - do_print: bool = True, - use_fast_gradients: bool = False, -) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """Optimize the PP + Hardware layer ansatz. - - Args: - hamiltonian_data: Hamiltonian (molecular) specification. - restricted: Whether to use a spin-restricted ansatz or not. - random_parameter_scale: A float to scale the random parameters by. - initial_orbital_rotation: An optional initial orbital rotation matrix, - which will be implmented as a givens circuit. - heuristic_layers: A tuple of circuit layers to append to the perfect pairing circuit. - do_pp: Implement the perfect pairing circuit along with the heuristic - layers. Defaults to true. - n_optimization_restarts: The number of times to restart the optimization - from a random guess in an attempt at global optimization. - do_print: Whether to print optimization progress to stdout. - use_fast_gradients: Compute the parameter gradients anlytically using Wilcox formula. - Default to false (use finite difference gradients). - - Returns: - one_body_params: Optimized one-body parameters. - two_body_params: Optimized two-body parameters - one_body_basis_change_mat: The basis change matrix including any initial orbital rotation. - """ - n_elec = hamiltonian_data.params.n_elec - n_orb = hamiltonian_data.params.n_orb - sz = 0 - - initial_wf = fqe.Wavefunction([[n_elec, sz, n_orb]]) - initial_wf.set_wfn(strategy="hartree-fock") - - fqe_ham = hamiltonian_data.get_restricted_fqe_hamiltonian() - e_core = hamiltonian_data.e_core - - hf_energy = initial_wf.expectationValue(fqe_ham) + e_core - - # We're only supporting closed shell stuff here. - assert n_elec % 2 == 0 - assert n_elec <= n_orb - if use_fast_gradients: - err_msg = "use_fast_gradients does not work with initial orbital rotation." - assert initial_orbital_rotation is None, err_msg - - gate_generators = _get_pp_plus_gate_generators( - n_elec=n_elec, heuristic_layers=heuristic_layers, do_pp=do_pp - ) - - n_two_body_params = len(gate_generators) - - if restricted: - n_one_body_params = n_orb * (n_orb - 1) // 2 - else: - n_one_body_params = n_orb * (n_orb - 1) - - best = np.inf - best_res: Union[None, scipy.optimize.OptimizeResult] = None - for i in range(n_optimization_restarts): - if do_print: - print(f"Optimization restart {i}", flush=True) - - def progress_cb(_): - print(".", end="", flush=True) - - else: - - def progress_cb(_): - pass - - params = random_parameter_scale * np.random.normal( - size=(n_two_body_params + n_one_body_params) - ) - - def objective(params): - one_body_params = params[-n_one_body_params:] - two_body_params = params[:n_two_body_params] - - wf, _ = get_evolved_wf( - one_body_params, - two_body_params, - initial_wf, - gate_generators, - n_orb, - restricted=restricted, - initial_orbital_rotation=initial_orbital_rotation, - ) - - energy = wf.expectationValue(fqe_ham) + e_core - if do_print: - print(f"energy {energy}") - if np.abs(energy.imag) < 1e-6: - return energy.real - else: - return 1e6 - - def fast_obj_grad(params): - one_body_params = params[-n_one_body_params:] - two_body_params = params[:n_two_body_params] - energy, grad = evaluate_gradient_and_cost_function( - initial_wf, - fqe_ham, - n_orb, - one_body_params, - two_body_params, - gate_generators, - restricted, - e_core, - ) - if do_print: - print(f"energy {energy}, max|grad| {np.max(np.abs(grad))}") - if np.abs(energy.imag) < 1e-6: - return energy.real, grad - else: - return 1e6, 1e6 - - if use_fast_gradients: - res = minimize( - fast_obj_grad, params, jac=True, method="BFGS", callback=progress_cb - ) - else: - res = minimize(objective, params, callback=progress_cb) - if res.fun < best: - best = res.fun - best_res = res - - if do_print: - print(res, flush=True) - - assert best_res is not None - params = best_res.x - one_body_params = params[-n_one_body_params:] - two_body_params = params[:n_two_body_params] - - wf, _ = get_evolved_wf( - one_body_params, - two_body_params, - initial_wf, - gate_generators, - n_orb, - restricted=restricted, - initial_orbital_rotation=initial_orbital_rotation, - ) - - one_body_cluster_mat = get_one_body_cluster_coef( - one_body_params, n_orb, restricted=restricted - ) - # We need to change the ordering to match OpenFermion's abababab ordering - if not restricted: - index_rearrangement = np.asarray( - [i // 2 % (n_orb) + (i % 2) * n_orb for i in range(2 * n_orb)] - ) - one_body_cluster_mat = one_body_cluster_mat[:, index_rearrangement] - one_body_cluster_mat = one_body_cluster_mat[index_rearrangement, :] - - one_body_basis_change_mat = expm(one_body_cluster_mat) - - if initial_orbital_rotation is not None: - if restricted: - one_body_basis_change_mat = ( - initial_orbital_rotation @ one_body_basis_change_mat - ) - else: - big_initial_orbital_rotation = np.zeros_like(one_body_basis_change_mat) - - for i in range(len(initial_orbital_rotation)): - for j in range(len(initial_orbital_rotation)): - big_initial_orbital_rotation[2 * i, 2 * j] = ( - initial_orbital_rotation[i, j] - ) - big_initial_orbital_rotation[2 * i + 1, 2 * j + 1] = ( - initial_orbital_rotation[i, j] - ) - - one_body_basis_change_mat = ( - big_initial_orbital_rotation @ one_body_basis_change_mat - ) - - if do_print: - print("Hartree-Fock Energy:") - print(hf_energy) - initial_wf.print_wfn() - print("-" * 80) - print("FCI Energy:") - print(hamiltonian_data.e_fci) - print("-" * 80) - print(best_res) - - print("-" * 80) - print("Ansatz Energy:") - print(np.real_if_close(wf.expectationValue(fqe_ham) + e_core)) - wf.print_wfn() - print("Basis Rotation Matrix:") - print(one_body_basis_change_mat) - print("Two Body Rotation Parameters:") - print(two_body_params) - - return one_body_params, two_body_params, one_body_basis_change_mat diff --git a/recirq/qcqmc/trial_wf_test.py b/recirq/qcqmc/trial_wf_test.py index 6f34f3ac..91bd187d 100644 --- a/recirq/qcqmc/trial_wf_test.py +++ b/recirq/qcqmc/trial_wf_test.py @@ -21,7 +21,6 @@ from recirq.qcqmc.hamiltonian import HamiltonianData from recirq.qcqmc.trial_wf import ( - FermionicMode, LayerSpec, PerfectPairingPlusTrialWavefunctionParams, _get_ansatz_qubit_wf, @@ -33,15 +32,6 @@ ) -def test_fermionic_mode(): - fm = FermionicMode(5, "a") - fm2 = cirq.read_json(json_text=cirq.to_json(fm)) - assert fm == fm2 - - with pytest.raises(ValueError, match="spin.*"): - _ = FermionicMode(10, "c") - - def test_pp_wf_energy(fixture_4_qubit_ham: HamiltonianData): params = PerfectPairingPlusTrialWavefunctionParams( name="pp_test_wf_1", From 48586f32d014b6f9ecb8e88d6f02089a376e5f5a Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 20:08:49 +0000 Subject: [PATCH 08/21] Fix tests. --- recirq/qcqmc/__init__.py | 16 +- recirq/qcqmc/conftest.py | 2 +- recirq/qcqmc/fermion_mode_test.py | 1 + recirq/qcqmc/optimize_wf.py | 2 +- recirq/qcqmc/optimize_wf_test.py | 17 +- recirq/qcqmc/trial_wf_test.py | 390 ------------------------------ 6 files changed, 17 insertions(+), 411 deletions(-) diff --git a/recirq/qcqmc/__init__.py b/recirq/qcqmc/__init__.py index 1314592b..6a6c1385 100644 --- a/recirq/qcqmc/__init__.py +++ b/recirq/qcqmc/__init__.py @@ -16,17 +16,11 @@ from cirq.protocols.json_serialization import DEFAULT_RESOLVERS, ObjectFactory -from .hamiltonian import ( - HamiltonianData, - HamiltonianFileParams, - PyscfHamiltonianParams, -) -from .trial_wf import ( - FermionicMode, - LayerSpec, - PerfectPairingPlusTrialWavefunctionParams, - TrialWavefunctionData, -) +from .fermion_mode import FermionicMode +from .hamiltonian import (HamiltonianData, HamiltonianFileParams, + PyscfHamiltonianParams) +from .trial_wf import (LayerSpec, PerfectPairingPlusTrialWavefunctionParams, + TrialWavefunctionData) @lru_cache() diff --git a/recirq/qcqmc/conftest.py b/recirq/qcqmc/conftest.py index a61a6a22..a216f21c 100644 --- a/recirq/qcqmc/conftest.py +++ b/recirq/qcqmc/conftest.py @@ -21,10 +21,10 @@ HamiltonianFileParams, build_hamiltonian_from_file, ) +from recirq.qcqmc.optimize_wf import build_pp_plus_trial_wavefunction from recirq.qcqmc.trial_wf import ( PerfectPairingPlusTrialWavefunctionParams, TrialWavefunctionData, - build_pp_plus_trial_wavefunction, ) diff --git a/recirq/qcqmc/fermion_mode_test.py b/recirq/qcqmc/fermion_mode_test.py index 4d5ff965..e57e73e5 100644 --- a/recirq/qcqmc/fermion_mode_test.py +++ b/recirq/qcqmc/fermion_mode_test.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import cirq +import pytest from recirq.qcqmc.fermion_mode import FermionicMode diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index 9b6da4d4..b760c6cf 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -234,7 +234,7 @@ def get_rotated_hamiltonians( one_body_basis_change_mat: np.ndarray, mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], ordered_qubits: Sequence[cirq.Qid], -) -> Tuple[fqe_hams.RestrictedHamiltonian, float, scipy.sparse.scipy.sparse.csc_matrix]: +) -> Tuple[fqe_hams.RestrictedHamiltonian, float, scipy.sparse.csc_matrix]: """A helper method that gets the hamiltonians in the basis of the trial_wf. Returns: diff --git a/recirq/qcqmc/optimize_wf_test.py b/recirq/qcqmc/optimize_wf_test.py index bea91d40..1b392d00 100644 --- a/recirq/qcqmc/optimize_wf_test.py +++ b/recirq/qcqmc/optimize_wf_test.py @@ -19,14 +19,15 @@ import pytest from recirq.qcqmc.hamiltonian import HamiltonianData -from recirq.qcqmc.trial_wf import (LayerSpec, - PerfectPairingPlusTrialWavefunctionParams, - _get_ansatz_qubit_wf, - _get_pp_plus_gate_generators, - build_pp_plus_trial_wavefunction, - evaluate_gradient_and_cost_function, - get_evolved_wf, - get_two_body_params_from_qchem_amplitudes) +from recirq.qcqmc.optimize_wf import ( + _get_ansatz_qubit_wf, + _get_pp_plus_gate_generators, + build_pp_plus_trial_wavefunction, + evaluate_gradient_and_cost_function, + get_evolved_wf, + get_two_body_params_from_qchem_amplitudes, +) +from recirq.qcqmc.trial_wf import LayerSpec, PerfectPairingPlusTrialWavefunctionParams def test_pp_wf_energy(fixture_4_qubit_ham: HamiltonianData): diff --git a/recirq/qcqmc/trial_wf_test.py b/recirq/qcqmc/trial_wf_test.py index 91bd187d..e69de29b 100644 --- a/recirq/qcqmc/trial_wf_test.py +++ b/recirq/qcqmc/trial_wf_test.py @@ -1,390 +0,0 @@ -# Copyright 2024 Google -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import cirq -import fqe -import fqe.hamiltonians.restricted_hamiltonian as fqe_hams -import fqe.wavefunction as fqe_wfn -import numpy as np -import pytest - -from recirq.qcqmc.hamiltonian import HamiltonianData -from recirq.qcqmc.trial_wf import ( - LayerSpec, - PerfectPairingPlusTrialWavefunctionParams, - _get_ansatz_qubit_wf, - _get_pp_plus_gate_generators, - build_pp_plus_trial_wavefunction, - evaluate_gradient_and_cost_function, - get_evolved_wf, - get_two_body_params_from_qchem_amplitudes, -) - - -def test_pp_wf_energy(fixture_4_qubit_ham: HamiltonianData): - params = PerfectPairingPlusTrialWavefunctionParams( - name="pp_test_wf_1", - hamiltonian_params=fixture_4_qubit_ham.params, - heuristic_layers=(), - do_pp=True, - restricted=True, - ) - - trial_wf = build_pp_plus_trial_wavefunction( - params, dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham} - ) - - assert np.isclose(trial_wf.ansatz_energy, fixture_4_qubit_ham.e_fci) - - -def test_pp_wf_energy_with_layer(fixture_4_qubit_ham: HamiltonianData): - params = PerfectPairingPlusTrialWavefunctionParams( - name="pp_test_wf_2", - hamiltonian_params=fixture_4_qubit_ham.params, - heuristic_layers=(LayerSpec("charge_charge", "cross_spin"),), - do_pp=True, - restricted=True, - ) - - trial_wf = build_pp_plus_trial_wavefunction( - params, dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham} - ) - - assert np.isclose(trial_wf.ansatz_energy, fixture_4_qubit_ham.e_fci) - - -def test_qchem_pp_eight_qubit_wavefunctions_consistent( - fixture_8_qubit_ham: HamiltonianData, -): - """Tests (without optimization) that the eight qubit wavefunctions work. - - Specifically, that constructing the wavefunction with FQE and then - converting it to a cirq wavefunction yields the same result as constructing - the parameters with the circuit directly. - """ - params = PerfectPairingPlusTrialWavefunctionParams( - name="pp_test_wf_qchem", - hamiltonian_params=fixture_8_qubit_ham.params, - heuristic_layers=tuple(), - initial_orbital_rotation=None, - initial_two_body_qchem_amplitudes=np.asarray([0.3, 0.4]), - do_optimization=False, - ) - - trial_wf = build_pp_plus_trial_wavefunction( - params, - dependencies={fixture_8_qubit_ham.params: fixture_8_qubit_ham}, - do_print=False, - ) - - one_body_params = trial_wf.one_body_params - two_body_params = trial_wf.two_body_params - basis_change_mat = trial_wf.one_body_basis_change_mat - - np.testing.assert_array_almost_equal(one_body_params, np.zeros((12,))) - np.testing.assert_array_almost_equal(basis_change_mat, np.diag(np.ones(8))) - - np.testing.assert_equal(two_body_params.shape, (2,)) - - -def test_pp_plus_wf_energy_sloppy_1(fixture_8_qubit_ham: HamiltonianData): - params = PerfectPairingPlusTrialWavefunctionParams( - "pp_plus_test", - hamiltonian_params=fixture_8_qubit_ham.params, - heuristic_layers=tuple(), - do_pp=True, - restricted=False, - random_parameter_scale=1, - ) - - trial_wf = build_pp_plus_trial_wavefunction( - params, - dependencies={fixture_8_qubit_ham.params: fixture_8_qubit_ham}, - do_print=True, - ) - - assert trial_wf.ansatz_energy < -1.947 - - -# TODO: Speed up this test and add a similar one with non-trivial heuristic layers. - - -def test_diamond_pp_wf_energy(fixture_12_qubit_ham: HamiltonianData): - params = PerfectPairingPlusTrialWavefunctionParams( - name="diamind_pp_test_wf_1", - hamiltonian_params=fixture_12_qubit_ham.params, - heuristic_layers=tuple(), - do_pp=True, - restricted=True, - random_parameter_scale=0.1, - n_optimization_restarts=1, - ) - - trial_wf = build_pp_plus_trial_wavefunction( - params, - dependencies={fixture_12_qubit_ham.params: fixture_12_qubit_ham}, - do_print=True, - ) - - assert trial_wf.ansatz_energy < -10.4 - - -@pytest.mark.parametrize( - "initial_two_body_qchem_amplitudes, expected_ansatz_qubit_wf", - [ - ( - [1], - np.array( - [ - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.70710678 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.70710678 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - 0.0 + 0.0j, - ] - ), - ), - ( - [0], - np.array( - [ - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - ] - ), - ), - ], -) -def test_qchem_pp_runs( - initial_two_body_qchem_amplitudes, - expected_ansatz_qubit_wf, - fixture_4_qubit_ham: HamiltonianData, -): - params = PerfectPairingPlusTrialWavefunctionParams( - name="pp_test_wf_qchem", - hamiltonian_params=fixture_4_qubit_ham.params, - heuristic_layers=tuple(), - initial_orbital_rotation=None, - initial_two_body_qchem_amplitudes=initial_two_body_qchem_amplitudes, - do_optimization=False, - ) - - trial_wf = build_pp_plus_trial_wavefunction( - params, - dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham}, - do_print=False, - ) - - ansatz_qubit_wf = _get_ansatz_qubit_wf( - ansatz_circuit=trial_wf.ansatz_circuit, - ordered_qubits=params.qubits_jordan_wigner_ordered, - ) - - np.testing.assert_array_almost_equal(ansatz_qubit_wf, expected_ansatz_qubit_wf) - - -def test_qchem_conversion_negative(fixture_4_qubit_ham: HamiltonianData): - qchem_amplitudes = np.asarray(-0.1) - - two_body_params = get_two_body_params_from_qchem_amplitudes(qchem_amplitudes) - - assert two_body_params.item() < 0 - - params = PerfectPairingPlusTrialWavefunctionParams( - name="pp_test_wf_qchem_neg", - hamiltonian_params=fixture_4_qubit_ham.params, - heuristic_layers=tuple(), - initial_orbital_rotation=None, - initial_two_body_qchem_amplitudes=qchem_amplitudes, - do_optimization=False, - ) - - trial_wf = build_pp_plus_trial_wavefunction( - params, - dependencies={fixture_4_qubit_ham.params: fixture_4_qubit_ham}, - do_print=False, - ) - - ansatz_qubit_wf = _get_ansatz_qubit_wf( - ansatz_circuit=trial_wf.ansatz_circuit, - ordered_qubits=params.qubits_jordan_wigner_ordered, - ) - - assert any(ansatz_qubit_wf < 0) - - -def gen_random_restricted_ham(n_orb: int) -> fqe_hams.RestrictedHamiltonian: - """8-fold symmetry restricted hamiltonian""" - h1e = np.random.random((n_orb,) * 2) - h1e = h1e + h1e.T - h2e = np.random.random((n_orb,) * 4) - h2e = h2e + h2e.transpose(2, 3, 0, 1) - h2e = h2e + h2e.transpose(3, 2, 1, 0) - h2e = h2e + h2e.transpose(1, 0, 2, 3) - h2e = np.asarray(h2e.transpose(0, 2, 3, 1), order="C") - fqe_ham = fqe_hams.RestrictedHamiltonian((h1e, np.einsum("ijlk", -0.5 * h2e))) - return fqe_ham - - -def compute_finite_difference_grad( - n_orb: int, - n_elec: int, - one_body_params: np.ndarray, - two_body_params: np.ndarray, - ham: fqe_hams.RestrictedHamiltonian, - initial_wf: fqe_wfn.Wavefunction, - dtheta: float = 1e-4, - restricted: bool = False, -): - """Compute the parameter gradient using finite differences. - - Args: - n_orb: the number of spatial orbitals. - n_elec: the number of electrons. - one_body_params: The variational parameters for the one-body terms in the ansatz. - two_body_params: The variational parameters for the two-body terms in the ansatz. - ham: The restricted FQE Hamiltonian. - initial_wf: The initial wavefunction (typically Hartree--Fock) - restricted: Whether we're using a restricted ansatz or not. - """ - generators = _get_pp_plus_gate_generators( - n_elec=n_elec, heuristic_layers=tuple(), do_pp=True - ) - one_body_gradient = np.zeros_like(one_body_params) - for ig, _ in enumerate(one_body_gradient): - new_param = one_body_params.copy() - new_param[ig] = new_param[ig] + dtheta - phi = get_evolved_wf( - new_param, - two_body_params, - initial_wf, - generators, - n_orb, - restricted=restricted, - )[0] - e_plus = phi.expectationValue(ham) - new_param[ig] = new_param[ig] - 2 * dtheta - phi = get_evolved_wf( - new_param, - two_body_params, - initial_wf, - generators, - n_orb, - restricted=restricted, - )[0] - e_minu = phi.expectationValue(ham) - one_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) - two_body_gradient = np.zeros_like(two_body_params) - for ig, _ in enumerate(two_body_gradient): - new_param = two_body_params.copy() - new_param[ig] = new_param[ig] + dtheta - phi = get_evolved_wf( - one_body_params, - new_param, - initial_wf, - generators, - n_orb, - restricted=restricted, - )[0] - e_plus = phi.expectationValue(ham) - new_param[ig] = new_param[ig] - 2 * dtheta - phi = get_evolved_wf( - one_body_params, - new_param, - initial_wf, - generators, - n_orb, - restricted=restricted, - )[0] - e_minu = phi.expectationValue(ham) - two_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) - return one_body_gradient, two_body_gradient - - -@pytest.mark.parametrize("n_elec, n_orb", ((2, 2), (4, 4), (6, 6))) -@pytest.mark.parametrize("restricted", (True, False)) -def test_gradient(n_elec, n_orb, restricted): - sz = 0 - initial_wf = fqe.Wavefunction([[n_elec, sz, n_orb]]) - initial_wf.set_wfn(strategy="hartree-fock") - - fqe_ham = gen_random_restricted_ham(n_orb) - - if restricted: - n_one_body_params = n_orb * (n_orb - 1) // 2 - else: - n_one_body_params = n_orb * (n_orb - 1) - - gate_generators = _get_pp_plus_gate_generators( - n_elec=n_elec, heuristic_layers=tuple(), do_pp=True - ) - # reference implementation - one_body_params = np.random.random(n_one_body_params) - two_body_params = np.random.random(len(gate_generators)) - phi = get_evolved_wf( - one_body_params, - two_body_params, - initial_wf, - gate_generators, - n_orb, - restricted=restricted, - )[0] - obj_val, grad = evaluate_gradient_and_cost_function( - initial_wf, - fqe_ham, - n_orb, - one_body_params, - two_body_params, - gate_generators, - restricted, - 0.0, - ) - ob_fd_grad, tb_fd_grad = compute_finite_difference_grad( - n_orb, - n_elec, - one_body_params, - two_body_params, - fqe_ham, - initial_wf, - restricted=restricted, - ) - assert np.isclose(obj_val, phi.expectationValue(fqe_ham)) - assert np.allclose(ob_fd_grad, grad[-n_one_body_params:]) - n_two_body_params = len(two_body_params) - assert np.allclose(tb_fd_grad, grad[:n_two_body_params]) From d3258093741bac23c09615168a92c2693e5902a7 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 21:29:44 +0000 Subject: [PATCH 09/21] More refactor. --- recirq/qcqmc/__init__.py | 7 +- recirq/qcqmc/afqmc_circuits.py | 432 ++++++++++++++++- recirq/qcqmc/afqmc_generators.py | 151 ++++++ recirq/qcqmc/converters.py | 105 ++++ recirq/qcqmc/layer_spec.py | 127 +++++ recirq/qcqmc/optimize_wf.py | 805 +++---------------------------- recirq/qcqmc/optimize_wf_test.py | 19 +- recirq/qcqmc/trial_wf.py | 46 +- 8 files changed, 900 insertions(+), 792 deletions(-) create mode 100644 recirq/qcqmc/afqmc_generators.py create mode 100644 recirq/qcqmc/converters.py create mode 100644 recirq/qcqmc/layer_spec.py diff --git a/recirq/qcqmc/__init__.py b/recirq/qcqmc/__init__.py index 6a6c1385..bbae0f6f 100644 --- a/recirq/qcqmc/__init__.py +++ b/recirq/qcqmc/__init__.py @@ -17,10 +17,9 @@ from cirq.protocols.json_serialization import DEFAULT_RESOLVERS, ObjectFactory from .fermion_mode import FermionicMode -from .hamiltonian import (HamiltonianData, HamiltonianFileParams, - PyscfHamiltonianParams) -from .trial_wf import (LayerSpec, PerfectPairingPlusTrialWavefunctionParams, - TrialWavefunctionData) +from .hamiltonian import HamiltonianData, HamiltonianFileParams, PyscfHamiltonianParams +from .layer_spec import LayerSpec +from .trial_wf import PerfectPairingPlusTrialWavefunctionParams, TrialWavefunctionData @lru_cache() diff --git a/recirq/qcqmc/afqmc_circuits.py b/recirq/qcqmc/afqmc_circuits.py index c76d572b..612bbc4e 100644 --- a/recirq/qcqmc/afqmc_circuits.py +++ b/recirq/qcqmc/afqmc_circuits.py @@ -15,8 +15,9 @@ """Trial wavefunction circuit ansatz primitives.""" import itertools -from typing import Iterable, List, Tuple +from typing import Dict, Iterable, Iterator, List, Sequence, Tuple +import attrs import cirq import numpy as np import openfermion @@ -26,6 +27,9 @@ from openfermion.linalg.givens_rotations import givens_decomposition from openfermion.linalg.sparse_tools import jw_sparse_givens_rotation +from recirq.qcqmc import layer_spec as lspec +from recirq.qcqmc import qubit_maps + class GeminalStatePreparationGate(cirq.Gate): def __init__(self, angle: float, inline_control: bool = False): @@ -434,3 +438,429 @@ def get_geminal_and_slater_det_overlap_via_simulation( else: measurement = cirq.Circuit(cirq.Z(ancilla)).unitary(qubit_order=qubits) return final_state.T.conj().dot(measurement.dot(final_state)) + + +def get_givens_gate(qubits: Tuple[cirq.Qid, ...], param: float) -> cirq.Operation: + """Get a the givens rotation gate on two qubits. + + Args: + qubits: The two qubits to apply the gate to. + param: The parameter for the givens rotation. + + Returns: + The givens rotation gate. + """ + return cirq.givens(param).on(qubits[0], qubits[1]) + + +def get_charge_charge_gate( + qubits: Tuple[cirq.Qid, ...], param: float +) -> cirq.Operation: + """Get the cirq charge-charge gate. + + Args: + qubits: Two qubits you want to apply the gate to. + param: The parameter for the charge-charge interaction. + + Returns: + The charge-charge gate. + """ + return cirq.CZ(qubits[0], qubits[1]) ** (-param / np.pi) + + +def get_layer_gates( + layer_spec: lspec.LayerSpec, + n_elec: int, + params: np.ndarray, + fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], +) -> List[cirq.Operation]: + """Gets the gates for a hardware efficient layer of the ansatz. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + params: The variational parameters for the hardware efficient gate layer. + fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. + + Returns: + A list of gates for the layer. + """ + + indices_list = lspec.get_layer_indices(layer_spec, n_elec) + + gate_funcs = {"givens": get_givens_gate, "charge_charge": get_charge_charge_gate} + gate_func = gate_funcs[layer_spec.base_gate] + + gates = [] + for indices, param in zip(indices_list, params): + qubits = tuple(fermion_index_to_qubit_map[ind] for ind in indices) + gates.append(gate_func(qubits, param)) + + return gates + + +def get_heuristic_circuit( + layer_specs: Sequence[lspec.LayerSpec], + n_elec: int, + params: np.ndarray, + fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], +) -> cirq.Circuit: + """Get a circuit for the heuristic ansatz. + + Args: + layer_specs: The layer specs for the heuristic layers. + n_elec: The number of electrons. + params: The variational parameters for the circuit. + fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. + + Returns: + A circuit for the heuristic ansatz. + """ + gates: List[cirq.Operation] = [] + + for layer_spec in layer_specs: + params_slice = params[len(gates) :] + gates += get_layer_gates( + layer_spec, n_elec, params_slice, fermion_index_to_qubit_map + ) + + return cirq.Circuit(gates) + + +def get_4_qubit_pp_circuits( + *, + two_body_params: np.ndarray, + n_elec: int, + heuristic_layers: Tuple[lspec.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 3 1 + 2 0 + """ + assert n_elec == 2 + + fermion_index_to_qubit_map = qubit_maps.get_4_qubit_fermion_qubit_map() + geminal_gate = GeminalStatePreparationGate(two_body_params[0], inline_control=True) + + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + ) + ) + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[1:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + indicator = fermion_index_to_qubit_map[2] + superposition_circuit = cirq.Circuit([cirq.H(indicator) + ansatz_circuit]) + ansatz_circuit = cirq.Circuit([cirq.X(indicator) + ansatz_circuit]) + + return superposition_circuit, ansatz_circuit + + +def get_8_qubit_circuits( + *, + two_body_params: np.ndarray, + n_elec: int, + heuristic_layers: Tuple[lspec.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 3 5 1 7 + 2 4 0 6 + """ + fermion_index_to_qubit_map = qubit_maps.get_8_qubit_fermion_qubit_map() + + geminal_gate_1 = GeminalStatePreparationGate( + two_body_params[0], inline_control=True + ) + geminal_gate_2 = GeminalStatePreparationGate( + two_body_params[1], inline_control=True + ) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[2:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[0]), + cirq.CNOT(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[6]), + cirq.SWAP(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[4]), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[4]), + cirq.X(fermion_index_to_qubit_map[6]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_12_qubit_circuits( + *, + two_body_params: np.ndarray, + n_elec: int, + heuristic_layers: Tuple[lspec.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 5 7 3 9 1 11 + 4 6 2 8 0 10 + """ + + fermion_index_to_qubit_map = qubit_maps.get_12_qubit_fermion_qubit_map() + + geminal_gate_1 = GeminalStatePreparationGate( + two_body_params[0], inline_control=True + ) + geminal_gate_2 = GeminalStatePreparationGate( + two_body_params[1], inline_control=True + ) + geminal_gate_3 = GeminalStatePreparationGate( + two_body_params[2], inline_control=True + ) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[8], + fermion_index_to_qubit_map[9], + ) + ), + cirq.decompose( + geminal_gate_3.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[10], + fermion_index_to_qubit_map[11], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[3:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[8]), + cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[0]), + cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[2]), + cirq.SWAP( + fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[10] + ), + cirq.SWAP(fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[6]), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[6]), + cirq.X(fermion_index_to_qubit_map[8]), + cirq.X(fermion_index_to_qubit_map[10]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_16_qubit_circuits( + *, + two_body_params: np.ndarray, + n_elec: int, + heuristic_layers: Tuple[lspec.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A helper function that builds the circuits for the four qubit ansatz. + + We map the fermionic orbitals to grid qubits like so: + 7 9 5 11 3 13 1 15 + 6 8 4 10 2 12 0 14 + """ + fermion_index_to_qubit_map = qubit_maps.get_16_qubit_fermion_qubit_map() + + geminal_gate_1 = GeminalStatePreparationGate( + two_body_params[0], inline_control=True + ) + geminal_gate_2 = GeminalStatePreparationGate( + two_body_params[1], inline_control=True + ) + geminal_gate_3 = GeminalStatePreparationGate( + two_body_params[2], inline_control=True + ) + geminal_gate_4 = GeminalStatePreparationGate( + two_body_params[3], inline_control=True + ) + + # We'll add the initial bit flips later. + ansatz_circuit = cirq.Circuit( + cirq.decompose( + geminal_gate_1.on( + fermion_index_to_qubit_map[6], + fermion_index_to_qubit_map[7], + fermion_index_to_qubit_map[8], + fermion_index_to_qubit_map[9], + ) + ), + cirq.decompose( + geminal_gate_2.on( + fermion_index_to_qubit_map[4], + fermion_index_to_qubit_map[5], + fermion_index_to_qubit_map[10], + fermion_index_to_qubit_map[11], + ) + ), + cirq.decompose( + geminal_gate_3.on( + fermion_index_to_qubit_map[2], + fermion_index_to_qubit_map[3], + fermion_index_to_qubit_map[12], + fermion_index_to_qubit_map[13], + ) + ), + cirq.decompose( + geminal_gate_4.on( + fermion_index_to_qubit_map[0], + fermion_index_to_qubit_map[1], + fermion_index_to_qubit_map[14], + fermion_index_to_qubit_map[15], + ) + ), + ) + + heuristic_layer_circuit = get_heuristic_circuit( + heuristic_layers, n_elec, two_body_params[4:], fermion_index_to_qubit_map + ) + + ansatz_circuit += heuristic_layer_circuit + + superposition_circuit = ( + cirq.Circuit( + [ + cirq.H(fermion_index_to_qubit_map[10]), + cirq.CNOT( + fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[2] + ), + cirq.SWAP( + fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[12] + ), + cirq.CNOT( + fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[4] + ), + cirq.CNOT( + fermion_index_to_qubit_map[12], fermion_index_to_qubit_map[0] + ), + cirq.SWAP(fermion_index_to_qubit_map[4], fermion_index_to_qubit_map[8]), + cirq.SWAP( + fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[14] + ), + ] + ) + + ansatz_circuit + ) + + ansatz_circuit = ( + cirq.Circuit( + [ + cirq.X(fermion_index_to_qubit_map[8]), + cirq.X(fermion_index_to_qubit_map[10]), + cirq.X(fermion_index_to_qubit_map[12]), + cirq.X(fermion_index_to_qubit_map[14]), + ] + ) + + ansatz_circuit + ) + + return superposition_circuit, ansatz_circuit + + +def get_circuits( + *, + two_body_params: np.ndarray, + n_orb: int, + n_elec: int, + heuristic_layers: Tuple[lspec.LayerSpec, ...], +) -> Tuple[cirq.Circuit, cirq.Circuit]: + """A function that runs a specialized method to get the ansatz circuits.""" + + if n_orb != n_elec: + raise ValueError("n_orb must equal n_elec.") + + circ_funcs = { + 2: get_4_qubit_pp_circuits, + 4: get_8_qubit_circuits, + 6: get_12_qubit_circuits, + 8: get_16_qubit_circuits, + } + try: + circ_func = circ_funcs[n_orb] + except KeyError: + raise NotImplementedError(f"No circuits for n_orb = {n_orb}") + + return circ_func( + two_body_params=two_body_params, + n_elec=n_elec, + heuristic_layers=heuristic_layers, + ) diff --git a/recirq/qcqmc/afqmc_generators.py b/recirq/qcqmc/afqmc_generators.py new file mode 100644 index 00000000..78987ce0 --- /dev/null +++ b/recirq/qcqmc/afqmc_generators.py @@ -0,0 +1,151 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Sequence, Tuple + +import openfermion as of + +from recirq.qcqmc import layer_spec as lspec + + +def get_pp_plus_gate_generators( + *, + n_elec: int, + heuristic_layers: Tuple[lspec.LayerSpec, ...], + do_pp: bool = True, +) -> List[of.FermionOperator]: + """Get PP+ gate generators for a given number of electrons. + + Args: + n_elec: The number of electrons + heuristic_layers: The heuristic HW layers + do_pp: Whether to add the PP gates to the ansatz too. + + Returns: + The list of generators necessary to construct the ansatz. + """ + heuristic_gate_generators = get_heuristic_gate_generators(n_elec, heuristic_layers) + if not do_pp: + return heuristic_gate_generators + + n_pairs = n_elec // 2 + pair_gate_generators = get_pair_hopping_gate_generators(n_pairs, n_elec) + return pair_gate_generators + heuristic_gate_generators + + +def get_pair_hopping_gate_generators( + n_pairs: int, n_elec: int +) -> List[of.FermionOperator]: + """Get the generators of the pair-hopping unitaries. + + Args: + n_pairs: The number of pair coupling terms. + n_elec: The total number of electrons. + + Returns: + A list of gate generators + """ + gate_generators = [] + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + + fop_string = f"{to_b} {to_a} {from_b}^ {from_a}^" + + gate_generator = of.FermionOperator(fop_string, 1.0) + gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) + gate_generators.append(gate_generator) + + return gate_generators + + +def get_charge_charge_generator(indices: Tuple[int, int]) -> of.FermionOperator: + """Returns the generator for density evolution between the indices + + Args: + indices: The indices to for charge-charge terms.:w + + Returns: + The generator for density evolution for this pair of electrons. + """ + + fop_string = "{:d}^ {:d} {:d}^ {:d}".format( + indices[0], indices[0], indices[1], indices[1] + ) + gate_generator = of.FermionOperator(fop_string, 1.0) + + return gate_generator + + +def get_givens_generator(indices: Tuple[int, int]) -> of.FermionOperator: + """Returns the generator for givens rotation between two orbitalspec. + + Args: + indices: The two indices for the givens rotation. + + Returns: + The givens generator for evolution for this pair of electrons. + """ + + fop_string = "{:d}^ {:d}".format(indices[0], indices[1]) + gate_generator = of.FermionOperator(fop_string, 1.0) + gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) + + return gate_generator + + +def get_layer_generators( + layer_spec: lspec.LayerSpec, n_elec: int +) -> List[of.FermionOperator]: + """Gets the generators for rotations in a hardware efficient layer of the ansatz. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + + Returns: + A list of generators for the layers. + """ + + indices_list = lspec.get_layer_indices(layer_spec, n_elec) + + gate_funcs = { + "givens": get_givens_generator, + "charge_charge": get_charge_charge_generator, + } + gate_func = gate_funcs[layer_spec.base_gate] + + return [gate_func(indices) for indices in indices_list] + + +def get_heuristic_gate_generators( + n_elec: int, layer_specs: Sequence[lspec.LayerSpec] +) -> List[of.FermionOperator]: + """Get gate generators for the heuristic ansatz. + + Args: + n_elec: The number of electrons. + layer_specs: The layer specifications. + + Returns: + A list of generators for the layers. + """ + gate_generators = [] + + for layer_spec in layer_specs: + gate_generators += get_layer_generators(layer_spec, n_elec) + + return gate_generators diff --git a/recirq/qcqmc/converters.py b/recirq/qcqmc/converters.py new file mode 100644 index 00000000..88365d80 --- /dev/null +++ b/recirq/qcqmc/converters.py @@ -0,0 +1,105 @@ +from typing import Callable, Mapping, Sequence + +import cirq +import fqe +import fqe.wavefunction as fqe_wfn +import numpy as np +import openfermion as of + +from recirq.qcqmc import fermion_mode + + +def get_reorder_func( + *, + mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> Callable[[int, int], int]: + """This is a helper function that allows us to reorder fermionic modes. + + Under the Jordan-Wigner transform, each fermionic mode is assigned to a + qubit. If we are provided an openfermion FermionOperator with the modes + assigned to qubits as described by mode_qubit_map this function gives us a + reorder_func that we can use to reorder the modes (with + openfermion.reorder(...)) so that they match the order of the qubits in + ordered_qubits. This is necessary to make a correspondence between + fermionic operators / wavefunctions and their qubit counterparts. + + Args: + mode_qubit_map: A dict that shows how each FermionicMode is mapped to a qubit. + ordered_qubits: An ordered sequence of qubits. + """ + qubits = list(mode_qubit_map.values()) + assert len(qubits) == len(ordered_qubits) + + # We sort the key: value pairs by the order of the values (qubits) in + # ordered_qubits. + sorted_mapping = list(mode_qubit_map.items()) + sorted_mapping.sort(key=lambda x: ordered_qubits.index(x[1])) + + remapping_map = {} + for i, (mode, _) in enumerate(sorted_mapping): + openfermion_index = 2 * mode.orb_ind + (0 if mode.spin == "a" else 1) + remapping_map[openfermion_index] = i + + def remapper(index: int, _: int) -> int: + """A function that maps from the old index to the new one. + + The _ argument is because it's expected by openfermion.reorder""" + return remapping_map[index] + + return remapper + + +def get_ansatz_qubit_wf( + *, ansatz_circuit: cirq.Circuit, ordered_qubits: Sequence[cirq.Qid] +): + """Get the cirq statevector from the ansatz circuit.""" + return cirq.final_state_vector( + ansatz_circuit, qubit_order=list(ordered_qubits), dtype=np.complex128 + ) + + +def get_two_body_params_from_qchem_amplitudes( + qchem_amplitudes: np.ndarray, +) -> np.ndarray: + """Translates perfect pairing amplitudes from qchem to rotation angles. + + qchem style: 1 |1100> + t_i |0011> + our style: cos(\theta_i) |1100> + sin(\theta_i) |0011> + """ + + two_body_params = np.arccos(1 / np.sqrt(1 + qchem_amplitudes**2)) * np.sign( + qchem_amplitudes + ) + + # Numpy casts the array improperly to a float when we only have one parameter. + two_body_params = np.atleast_1d(two_body_params) + + return two_body_params + + +def convert_fqe_wf_to_cirq( + fqe_wf: fqe_wfn.Wavefunction, + mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], + ordered_qubits: Sequence[cirq.Qid], +) -> np.ndarray: + """Converts an FQE wavefunction to one on qubits with a particular ordering. + + Args: + fqe_wf: The FQE wavefunction. + mode_qubit_map: A mapping from fermion modes to cirq qubits. + ordered_qubits: + """ + n_qubits = len(mode_qubit_map) + fermion_op = fqe.openfermion_utils.fqe_to_fermion_operator(fqe_wf) + + reorder_func = get_reorder_func( + mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits + ) + fermion_op = of.reorder(fermion_op, reorder_func, num_modes=n_qubits) + + qubit_op = of.jordan_wigner(fermion_op) + + return fqe.qubit_wavefunction_from_vacuum( + qubit_op, list(cirq.LineQubit.range(n_qubits)) + ) diff --git a/recirq/qcqmc/layer_spec.py b/recirq/qcqmc/layer_spec.py new file mode 100644 index 00000000..a1b046d2 --- /dev/null +++ b/recirq/qcqmc/layer_spec.py @@ -0,0 +1,127 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Iterator, List, Tuple + +import attrs + + +@attrs.frozen +class LayerSpec: + """A specification of a hardware-efficient layer of gates. + + Args: + base_gate: 'charge_charge' for the e^{i n_i n_j} gate and 'givens' for a givens rotation. + layout: Should be 'in_pair', 'cross_pair', or 'cross_spin' only. + """ + + base_gate: str + layout: str + + def __attrs_post_init__(self): + if self.base_gate not in ["charge_charge", "givens"]: + raise ValueError( + f'base_gate is set to {self.base_gate}, it should be either "charge_charge or "givens".' + ) + if self.layout not in ["in_pair", "cross_pair", "cross_spin"]: + raise ValueError( + f'layout is set to {self.layout}, it should be either "cross_pair", "in_pair", or "cross_spin".' + ) + + @classmethod + def _json_namespace_(cls): + return "recirq.qcqmc" + + def _json_dict_(self): + return attrs.asdict(self) + + +def get_indices_heuristic_layer_in_pair(n_elec: int) -> Iterator[Tuple[int, int]]: + """Get the indicies for the heuristic layers. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices + """ + n_pairs = n_elec // 2 + + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + yield (from_a, to_a) + yield (from_b, to_b) + + +def get_indices_heuristic_layer_cross_pair(n_elec) -> Iterator[Tuple[int, int]]: + """Indices that couple adjacent pairs. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices + """ + n_pairs = n_elec // 2 + + for pair in range(n_pairs - 1): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_next_a = n_elec - 2 * (pair + 1) - 2 + from_next_b = n_elec - 2 * (pair + 1) - 1 + yield (to_a, from_next_a) + yield (to_b, from_next_b) + + +def get_indices_heuristic_layer_cross_spin(n_elec) -> Iterator[Tuple[int, int]]: + """Get indices that couple the two spin sectors. + + Args: + n_elec: The number of electrons + + Returns: + An iterator of the indices that couple spin sectors. + """ + n_pairs = n_elec // 2 + + for pair in range(n_pairs): + to_a = n_elec + 2 * pair + to_b = n_elec + 2 * pair + 1 + from_a = n_elec - 2 * pair - 2 + from_b = n_elec - 2 * pair - 1 + yield (to_a, to_b) + yield (from_a, from_b) + + +def get_layer_indices(layer_spec: LayerSpec, n_elec: int) -> List[Tuple[int, int]]: + """Get the indices for the heuristic layers. + + Args: + layer_spec: The layer specification. + n_elec: The number of electrons. + + Returns: + A list of indices for the layer. + """ + indices_generators = { + "in_pair": get_indices_heuristic_layer_in_pair(n_elec), + "cross_pair": get_indices_heuristic_layer_cross_pair(n_elec), + "cross_spin": get_indices_heuristic_layer_cross_spin(n_elec), + } + indices_generator = indices_generators[layer_spec.layout] + + return [indices for indices in indices_generator] diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index b760c6cf..36244812 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -14,7 +14,7 @@ import copy import itertools -from typing import Callable, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple +from typing import Dict, List, Mapping, Optional, Sequence, Tuple, Union import cirq import fqe @@ -28,75 +28,16 @@ from recirq.qcqmc import ( afqmc_circuits, + afqmc_generators, + converters, data, fermion_mode, hamiltonian, - qubit_maps, + layer_spec, trial_wf, ) -def _get_reorder_func( - *, - mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], - ordered_qubits: Sequence[cirq.Qid], -) -> Callable[[int, int], int]: - """This is a helper function that allows us to reorder fermionic modes. - - Under the Jordan-Wigner transform, each fermionic mode is assigned to a - qubit. If we are provided an openfermion FermionOperator with the modes - assigned to qubits as described by mode_qubit_map this function gives us a - reorder_func that we can use to reorder the modes (with - openfermion.reorder(...)) so that they match the order of the qubits in - ordered_qubits. This is necessary to make a correspondence between - fermionic operators / wavefunctions and their qubit counterparts. - - Args: - mode_qubit_map: A dict that shows how each FermionicMode is mapped to a qubit. - ordered_qubits: An ordered sequence of qubits. - """ - qubits = list(mode_qubit_map.values()) - assert len(qubits) == len(ordered_qubits) - - # We sort the key: value pairs by the order of the values (qubits) in - # ordered_qubits. - sorted_mapping = list(mode_qubit_map.items()) - sorted_mapping.sort(key=lambda x: ordered_qubits.index(x[1])) - - remapping_map = {} - for i, (mode, _) in enumerate(sorted_mapping): - openfermion_index = 2 * mode.orb_ind + (0 if mode.spin == "a" else 1) - remapping_map[openfermion_index] = i - - def remapper(index: int, _: int) -> int: - """A function that maps from the old index to the new one. - - The _ argument is because it's expected by openfermion.reorder""" - return remapping_map[index] - - return remapper - - -def _get_pp_plus_gate_generators( - *, n_elec: int, heuristic_layers: Tuple[trial_wf.LayerSpec, ...], do_pp: bool = True -) -> List[of.FermionOperator]: - heuristic_gate_generators = get_heuristic_gate_generators(n_elec, heuristic_layers) - if not do_pp: - return heuristic_gate_generators - - n_pairs = n_elec // 2 - pair_gate_generators = get_pair_hopping_gate_generators(n_pairs, n_elec) - return pair_gate_generators + heuristic_gate_generators - - -def _get_ansatz_qubit_wf( - *, ansatz_circuit: cirq.Circuit, ordered_qubits: Sequence[cirq.Qid] -): - return cirq.final_state_vector( - ansatz_circuit, qubit_order=list(ordered_qubits), dtype=np.complex128 - ) - - def get_and_check_energy( *, hamiltonian_data: hamiltonian.HamiltonianData, @@ -106,7 +47,21 @@ def get_and_check_energy( one_body_basis_change_mat: np.ndarray, params: trial_wf.PerfectPairingPlusTrialWavefunctionParams, ) -> Tuple[float, float]: - ansatz_qubit_wf = _get_ansatz_qubit_wf( + """Compute the energy of the ansatz circuit and check against known values where possible. + + Args: + hamiltonian_data: The Hamiltonian data. + ansatz_circuit: The ansatz circuit. + one_body_params: The one-body variational parameters. + two_body_params: The two-body variational parameters. + one_body_basis_change_mat: The one-body basis change matrix. + params: The trial wavefunction parameters. + + Returns: + ansatz_energy: The total energy of the ansatz circuit. + hf_energy: The hartree-fock energy (initial energy of the ansatz circuit). + """ + ansatz_qubit_wf = converters.get_ansatz_qubit_wf( ansatz_circuit=ansatz_circuit, ordered_qubits=params.qubits_jordan_wigner_ordered, ) @@ -127,7 +82,7 @@ def get_and_check_energy( one_body_params=one_body_params, two_body_params=two_body_params, wf=initial_wf, - gate_generators=_get_pp_plus_gate_generators( + gate_generators=afqmc_generators.get_pp_plus_gate_generators( n_elec=params.n_elec, heuristic_layers=params.heuristic_layers, do_pp=params.do_pp, @@ -157,7 +112,16 @@ def build_pp_plus_trial_wavefunction( dependencies: Dict[data.Params, data.Data], do_print: bool = False, ) -> trial_wf.TrialWavefunctionData: - """Builds a TrialWavefunctionData from a TrialWavefunctionParams""" + """Builds a TrialWavefunctionData from a TrialWavefunctionParams + + Args: + params: The parameters specifying the PP+ trial wavefunction. + dependencies: Data dependencies + do_print: Print debugging information to stdout + + Returns: + The constructed TrialWavefunctionData object. + """ if do_print: print("Building Trial Wavefunction") @@ -165,9 +129,8 @@ def build_pp_plus_trial_wavefunction( hamiltonian_data = dependencies[params.hamiltonian_params] assert isinstance(hamiltonian_data, hamiltonian.HamiltonianData) - assert ( - params.n_orb == params.n_elec - ) ## Necessary for perfect pairing wavefunction to make sense. + if params.n_orb != params.n_elec: + raise ValueError("PP wavefunction must have n_orb = n_elec") if params.do_optimization: ( @@ -195,11 +158,11 @@ def build_pp_plus_trial_wavefunction( n_one_body_params = params.n_orb * (params.n_orb - 1) one_body_params = np.zeros(n_one_body_params) one_body_basis_change_mat = np.diag(np.ones(params.n_orb * 2)) - two_body_params = get_two_body_params_from_qchem_amplitudes( + two_body_params = converters.get_two_body_params_from_qchem_amplitudes( params.initial_two_body_qchem_amplitudes ) - (superposition_circuit, ansatz_circuit) = get_circuits( + (superposition_circuit, ansatz_circuit) = afqmc_circuits.get_circuits( two_body_params=two_body_params, n_orb=params.n_orb, n_elec=params.n_elec, @@ -237,6 +200,12 @@ def get_rotated_hamiltonians( ) -> Tuple[fqe_hams.RestrictedHamiltonian, float, scipy.sparse.csc_matrix]: """A helper method that gets the hamiltonians in the basis of the trial_wf. + Args: + hamiltonian_data: The specification of the Hamiltonian. + one_body_basis_change_mat: The basis change matrix for the one body hamiltonian. + mode_qubit_map: A mapping between fermionic modes and cirq qubits. + ordered_qubits: An ordered set of qubits. + Returns: The hamiltonian in FQE form, minus a constant energy shift. The constant part of the Hamiltonian missing from the FQE Hamiltonian. @@ -251,7 +220,7 @@ def get_rotated_hamiltonians( mol_ham.rotate_basis(one_body_basis_change_mat) fermion_operator_ham = of.get_fermion_operator(mol_ham) - reorder_func = _get_reorder_func( + reorder_func = converters.get_reorder_func( mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits ) fermion_operator_ham_qubit_ordered = of.reorder( @@ -274,9 +243,23 @@ def get_energy_and_check_sanity( mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], ordered_qubits: Sequence[cirq.Qid], ) -> float: - """A method that checks for consistency and returns the ansatz energy.""" + """A method that checks for consistency and returns the ansatz energy. - unrotated_fqe_wf_as_cirq = convert_fqe_wf_to_cirq( + Args: + circuit_wf: The cirq statevector constructed from the wavefunction ansatz circuit. + fqe_wf: The FQE wavefunction used for optimization. + unrotated_fqe_wf: The (unoptimized) unrotated FQE wavefunction. + fqe_ham: The restricted FQE hamiltonian. + sparse_ham: The qubit Hamiltonian as a sparse matrix. + e_core: The core energy (nuclear-repulision + any frozen core energy.) + mode_qubit_map: A mapping between fermionic modes and cirq qubits. + ordered_qubits: An ordered set of qubits. + + Returns: + The ansatz energy. + """ + + unrotated_fqe_wf_as_cirq = converters.convert_fqe_wf_to_cirq( fqe_wf=unrotated_fqe_wf, mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits, @@ -287,8 +270,6 @@ def get_energy_and_check_sanity( assert isinstance(ansatz_energy, float) fqe_energy = np.real(fqe_wf.expectationValue(fqe_ham) + e_core) - print(scipy.sparse.csc_matrix(circuit_wf)) - print(scipy.sparse.csc_matrix(unrotated_fqe_wf_as_cirq)) np.testing.assert_array_almost_equal(ansatz_energy, fqe_energy) np.testing.assert_array_almost_equal( circuit_wf, unrotated_fqe_wf_as_cirq, decimal=5 @@ -296,398 +277,6 @@ def get_energy_and_check_sanity( return ansatz_energy -def get_4_qubit_pp_circuits( - *, - two_body_params: np.ndarray, - n_elec: int, - heuristic_layers: Tuple[trial_wf.LayerSpec, ...], -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. - - We map the fermionic orbitals to grid qubits like so: - 3 1 - 2 0 - """ - assert n_elec == 2 - - fermion_index_to_qubit_map = qubit_maps.get_4_qubit_fermion_qubit_map() - geminal_gate = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[0], inline_control=True - ) - - ansatz_circuit = cirq.Circuit( - cirq.decompose( - geminal_gate.on( - fermion_index_to_qubit_map[0], - fermion_index_to_qubit_map[1], - fermion_index_to_qubit_map[2], - fermion_index_to_qubit_map[3], - ) - ) - ) - - heuristic_layer_circuit = get_heuristic_circuit( - heuristic_layers, n_elec, two_body_params[1:], fermion_index_to_qubit_map - ) - - ansatz_circuit += heuristic_layer_circuit - - indicator = fermion_index_to_qubit_map[2] - superposition_circuit = cirq.Circuit([cirq.H(indicator) + ansatz_circuit]) - ansatz_circuit = cirq.Circuit([cirq.X(indicator) + ansatz_circuit]) - - return superposition_circuit, ansatz_circuit - - -def get_8_qubit_circuits( - *, - two_body_params: np.ndarray, - n_elec: int, - heuristic_layers: Tuple[trial_wf.LayerSpec, ...], -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. - - We map the fermionic orbitals to grid qubits like so: - 3 5 1 7 - 2 4 0 6 - """ - fermion_index_to_qubit_map = qubit_maps.get_8_qubit_fermion_qubit_map() - - geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[0], inline_control=True - ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[1], inline_control=True - ) - - # We'll add the initial bit flips later. - ansatz_circuit = cirq.Circuit( - cirq.decompose( - geminal_gate_1.on( - fermion_index_to_qubit_map[2], - fermion_index_to_qubit_map[3], - fermion_index_to_qubit_map[4], - fermion_index_to_qubit_map[5], - ) - ), - cirq.decompose( - geminal_gate_2.on( - fermion_index_to_qubit_map[0], - fermion_index_to_qubit_map[1], - fermion_index_to_qubit_map[6], - fermion_index_to_qubit_map[7], - ) - ), - ) - - heuristic_layer_circuit = get_heuristic_circuit( - heuristic_layers, n_elec, two_body_params[2:], fermion_index_to_qubit_map - ) - - ansatz_circuit += heuristic_layer_circuit - - superposition_circuit = ( - cirq.Circuit( - [ - cirq.H(fermion_index_to_qubit_map[0]), - cirq.CNOT(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[6]), - cirq.SWAP(fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[4]), - ] - ) - + ansatz_circuit - ) - - ansatz_circuit = ( - cirq.Circuit( - [ - cirq.X(fermion_index_to_qubit_map[4]), - cirq.X(fermion_index_to_qubit_map[6]), - ] - ) - + ansatz_circuit - ) - - return superposition_circuit, ansatz_circuit - - -def get_12_qubit_circuits( - *, - two_body_params: np.ndarray, - n_elec: int, - heuristic_layers: Tuple[trial_wf.LayerSpec, ...], -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. - - We map the fermionic orbitals to grid qubits like so: - 5 7 3 9 1 11 - 4 6 2 8 0 10 - """ - - fermion_index_to_qubit_map = qubit_maps.get_12_qubit_fermion_qubit_map() - - geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[0], inline_control=True - ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[1], inline_control=True - ) - geminal_gate_3 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[2], inline_control=True - ) - - # We'll add the initial bit flips later. - ansatz_circuit = cirq.Circuit( - cirq.decompose( - geminal_gate_1.on( - fermion_index_to_qubit_map[4], - fermion_index_to_qubit_map[5], - fermion_index_to_qubit_map[6], - fermion_index_to_qubit_map[7], - ) - ), - cirq.decompose( - geminal_gate_2.on( - fermion_index_to_qubit_map[2], - fermion_index_to_qubit_map[3], - fermion_index_to_qubit_map[8], - fermion_index_to_qubit_map[9], - ) - ), - cirq.decompose( - geminal_gate_3.on( - fermion_index_to_qubit_map[0], - fermion_index_to_qubit_map[1], - fermion_index_to_qubit_map[10], - fermion_index_to_qubit_map[11], - ) - ), - ) - - heuristic_layer_circuit = get_heuristic_circuit( - heuristic_layers, n_elec, two_body_params[3:], fermion_index_to_qubit_map - ) - - ansatz_circuit += heuristic_layer_circuit - - superposition_circuit = ( - cirq.Circuit( - [ - cirq.H(fermion_index_to_qubit_map[8]), - cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[0]), - cirq.CNOT(fermion_index_to_qubit_map[8], fermion_index_to_qubit_map[2]), - cirq.SWAP( - fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[10] - ), - cirq.SWAP(fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[6]), - ] - ) - + ansatz_circuit - ) - - ansatz_circuit = ( - cirq.Circuit( - [ - cirq.X(fermion_index_to_qubit_map[6]), - cirq.X(fermion_index_to_qubit_map[8]), - cirq.X(fermion_index_to_qubit_map[10]), - ] - ) - + ansatz_circuit - ) - - return superposition_circuit, ansatz_circuit - - -def get_16_qubit_circuits( - *, - two_body_params: np.ndarray, - n_elec: int, - heuristic_layers: Tuple[trial_wf.LayerSpec, ...], -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. - - We map the fermionic orbitals to grid qubits like so: - 7 9 5 11 3 13 1 15 - 6 8 4 10 2 12 0 14 - """ - fermion_index_to_qubit_map = qubit_maps.get_16_qubit_fermion_qubit_map() - - geminal_gate_1 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[0], inline_control=True - ) - geminal_gate_2 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[1], inline_control=True - ) - geminal_gate_3 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[2], inline_control=True - ) - geminal_gate_4 = afqmc_circuits.GeminalStatePreparationGate( - two_body_params[3], inline_control=True - ) - - # We'll add the initial bit flips later. - ansatz_circuit = cirq.Circuit( - cirq.decompose( - geminal_gate_1.on( - fermion_index_to_qubit_map[6], - fermion_index_to_qubit_map[7], - fermion_index_to_qubit_map[8], - fermion_index_to_qubit_map[9], - ) - ), - cirq.decompose( - geminal_gate_2.on( - fermion_index_to_qubit_map[4], - fermion_index_to_qubit_map[5], - fermion_index_to_qubit_map[10], - fermion_index_to_qubit_map[11], - ) - ), - cirq.decompose( - geminal_gate_3.on( - fermion_index_to_qubit_map[2], - fermion_index_to_qubit_map[3], - fermion_index_to_qubit_map[12], - fermion_index_to_qubit_map[13], - ) - ), - cirq.decompose( - geminal_gate_4.on( - fermion_index_to_qubit_map[0], - fermion_index_to_qubit_map[1], - fermion_index_to_qubit_map[14], - fermion_index_to_qubit_map[15], - ) - ), - ) - - heuristic_layer_circuit = get_heuristic_circuit( - heuristic_layers, n_elec, two_body_params[4:], fermion_index_to_qubit_map - ) - - ansatz_circuit += heuristic_layer_circuit - - superposition_circuit = ( - cirq.Circuit( - [ - cirq.H(fermion_index_to_qubit_map[10]), - cirq.CNOT( - fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[2] - ), - cirq.SWAP( - fermion_index_to_qubit_map[2], fermion_index_to_qubit_map[12] - ), - cirq.CNOT( - fermion_index_to_qubit_map[10], fermion_index_to_qubit_map[4] - ), - cirq.CNOT( - fermion_index_to_qubit_map[12], fermion_index_to_qubit_map[0] - ), - cirq.SWAP(fermion_index_to_qubit_map[4], fermion_index_to_qubit_map[8]), - cirq.SWAP( - fermion_index_to_qubit_map[0], fermion_index_to_qubit_map[14] - ), - ] - ) - + ansatz_circuit - ) - - ansatz_circuit = ( - cirq.Circuit( - [ - cirq.X(fermion_index_to_qubit_map[8]), - cirq.X(fermion_index_to_qubit_map[10]), - cirq.X(fermion_index_to_qubit_map[12]), - cirq.X(fermion_index_to_qubit_map[14]), - ] - ) - + ansatz_circuit - ) - - return superposition_circuit, ansatz_circuit - - -def get_circuits( - *, - two_body_params: np.ndarray, - # from wf_params: - n_orb: int, - n_elec: int, - heuristic_layers: Tuple[trial_wf.LayerSpec, ...], -) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A function that runs a specialized method to get the ansatz circuits.""" - - # TODO(?): Just input one of these quantities. - if n_orb != n_elec: - raise ValueError("n_orb must equal n_elec.") - - circ_funcs = { - 2: get_4_qubit_pp_circuits, - 4: get_8_qubit_circuits, - 6: get_12_qubit_circuits, - 8: get_16_qubit_circuits, - } - try: - circ_func = circ_funcs[n_orb] - except KeyError: - raise NotImplementedError(f"No circuits for n_orb = {n_orb}") - - return circ_func( - two_body_params=two_body_params, - n_elec=n_elec, - heuristic_layers=heuristic_layers, - ) - - -def get_two_body_params_from_qchem_amplitudes( - qchem_amplitudes: np.ndarray, -) -> np.ndarray: - """Translates perfect pairing amplitudes from qchem to rotation angles. - - qchem style: 1 |1100> + t_i |0011> - our style: cos(\theta_i) |1100> + sin(\theta_i) |0011> - """ - - two_body_params = np.arccos(1 / np.sqrt(1 + qchem_amplitudes**2)) * np.sign( - qchem_amplitudes - ) - - # Numpy casts the array improperly to a float when we only have one parameter. - two_body_params = np.atleast_1d(two_body_params) - - return two_body_params - - -#################### Here be dragons.########################################### - - -def convert_fqe_wf_to_cirq( - fqe_wf: fqe_wfn.Wavefunction, - mode_qubit_map: Mapping[fermion_mode.FermionicMode, cirq.Qid], - ordered_qubits: Sequence[cirq.Qid], -) -> np.ndarray: - """Converts an FQE wavefunction to one on qubits with a particular ordering. - - Args: - fqe_wf: The FQE wavefunction. - mode_qubit_map: A mapping from fermion modes to cirq qubits. - ordered_qubits: - """ - n_qubits = len(mode_qubit_map) - fermion_op = fqe.openfermion_utils.fqe_to_fermion_operator(fqe_wf) - - reorder_func = _get_reorder_func( - mode_qubit_map=mode_qubit_map, ordered_qubits=ordered_qubits - ) - fermion_op = of.reorder(fermion_op, reorder_func, num_modes=n_qubits) - - qubit_op = of.jordan_wigner(fermion_op) - - return fqe.qubit_wavefunction_from_vacuum( - qubit_op, list(cirq.LineQubit.range(n_qubits)) - ) - - def get_one_body_cluster_coef( params: np.ndarray, n_orb: int, restricted: bool ) -> np.ndarray: @@ -761,282 +350,6 @@ def get_evolved_wf( return rotated_wf, wf -def get_pair_hopping_gate_generators( - n_pairs: int, n_elec: int -) -> List[of.FermionOperator]: - """Get the generators of the pair-hopping unitaries. - - Args: - n_pairs: The number of pair coupling terms. - n_elec: The total number of electrons. - - Returns: - A list of gate generators - """ - gate_generators = [] - for pair in range(n_pairs): - to_a = n_elec + 2 * pair - to_b = n_elec + 2 * pair + 1 - from_a = n_elec - 2 * pair - 2 - from_b = n_elec - 2 * pair - 1 - - fop_string = f"{to_b} {to_a} {from_b}^ {from_a}^" - - gate_generator = of.FermionOperator(fop_string, 1.0) - gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) - gate_generators.append(gate_generator) - - return gate_generators - - -def get_indices_heuristic_layer_in_pair(n_elec: int) -> Iterator[Tuple[int, int]]: - """Get the indicies for the heuristic layers. - - Args: - n_elec: The number of electrons - - Returns: - An iterator of the indices - """ - n_pairs = n_elec // 2 - - for pair in range(n_pairs): - to_a = n_elec + 2 * pair - to_b = n_elec + 2 * pair + 1 - from_a = n_elec - 2 * pair - 2 - from_b = n_elec - 2 * pair - 1 - yield (from_a, to_a) - yield (from_b, to_b) - - -def get_indices_heuristic_layer_cross_pair(n_elec) -> Iterator[Tuple[int, int]]: - """Indices that couple adjacent pairs. - - Args: - n_elec: The number of electrons - - Returns: - An iterator of the indices - """ - n_pairs = n_elec // 2 - - for pair in range(n_pairs - 1): - to_a = n_elec + 2 * pair - to_b = n_elec + 2 * pair + 1 - from_next_a = n_elec - 2 * (pair + 1) - 2 - from_next_b = n_elec - 2 * (pair + 1) - 1 - yield (to_a, from_next_a) - yield (to_b, from_next_b) - - -def get_indices_heuristic_layer_cross_spin(n_elec) -> Iterator[Tuple[int, int]]: - """Get indices that couple the two spin sectors. - - Args: - n_elec: The number of electrons - - Returns: - An iterator of the indices that couple spin sectors. - """ - n_pairs = n_elec // 2 - - for pair in range(n_pairs): - to_a = n_elec + 2 * pair - to_b = n_elec + 2 * pair + 1 - from_a = n_elec - 2 * pair - 2 - from_b = n_elec - 2 * pair - 1 - yield (to_a, to_b) - yield (from_a, from_b) - - -def get_charge_charge_generator(indices: Tuple[int, int]) -> of.FermionOperator: - """Returns the generator for density evolution between the indices - - Args: - indices: The indices to for charge-charge terms.:w - - Returns: - The generator for density evolution for this pair of electrons. - """ - - fop_string = "{:d}^ {:d} {:d}^ {:d}".format( - indices[0], indices[0], indices[1], indices[1] - ) - gate_generator = of.FermionOperator(fop_string, 1.0) - - return gate_generator - - -def get_charge_charge_gate( - qubits: Tuple[cirq.Qid, ...], param: float -) -> cirq.Operation: - """Get the cirq charge-charge gate. - - Args: - qubits: Two qubits you want to apply the gate to. - param: The parameter for the charge-charge interaction. - - Returns: - The charge-charge gate. - """ - return cirq.CZ(qubits[0], qubits[1]) ** (-param / np.pi) - - -def get_givens_generator(indices: Tuple[int, int]) -> of.FermionOperator: - """Returns the generator for givens rotation between two orbitals. - - Args: - indices: The two indices for the givens rotation. - - Returns: - The givens generator for evolution for this pair of electrons. - """ - - fop_string = "{:d}^ {:d}".format(indices[0], indices[1]) - gate_generator = of.FermionOperator(fop_string, 1.0) - gate_generator = 1j * (gate_generator - of.hermitian_conjugated(gate_generator)) - - return gate_generator - - -def get_givens_gate(qubits: Tuple[cirq.Qid, ...], param: float) -> cirq.Operation: - """Get a the givens rotation gate on two qubits. - - Args: - qubits: The two qubits to apply the gate to. - param: The parameter for the givens rotation. - - Returns: - The givens rotation gate. - """ - return cirq.givens(param).on(qubits[0], qubits[1]) - - -def get_layer_indices( - layer_spec: trial_wf.LayerSpec, n_elec: int -) -> List[Tuple[int, int]]: - """Get the indices for the heuristic layers. - - Args: - layer_spec: The layer specification. - n_elec: The number of electrons. - - Returns: - A list of indices for the layer. - """ - indices_generators = { - "in_pair": get_indices_heuristic_layer_in_pair(n_elec), - "cross_pair": get_indices_heuristic_layer_cross_pair(n_elec), - "cross_spin": get_indices_heuristic_layer_cross_spin(n_elec), - } - indices_generator = indices_generators[layer_spec.layout] - - return [indices for indices in indices_generator] - - -def get_layer_gates( - layer_spec: trial_wf.LayerSpec, - n_elec: int, - params: np.ndarray, - fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], -) -> List[cirq.Operation]: - """Gets the gates for a hardware efficient layer of the ansatz. - - Args: - layer_spec: The layer specification. - n_elec: The number of electrons. - params: The variational parameters for the hardware efficient gate layer. - fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. - - Returns: - A list of gates for the layer. - """ - - indices_list = get_layer_indices(layer_spec, n_elec) - - gate_funcs = {"givens": get_givens_gate, "charge_charge": get_charge_charge_gate} - gate_func = gate_funcs[layer_spec.base_gate] - - gates = [] - for indices, param in zip(indices_list, params): - qubits = tuple(fermion_index_to_qubit_map[ind] for ind in indices) - gates.append(gate_func(qubits, param)) - - return gates - - -def get_layer_generators( - layer_spec: trial_wf.LayerSpec, n_elec: int -) -> List[of.FermionOperator]: - """Gets the generators for rotations in a hardware efficient layer of the ansatz. - - Args: - layer_spec: The layer specification. - n_elec: The number of electrons. - - Returns: - A list of generators for the layers. - """ - - indices_list = get_layer_indices(layer_spec, n_elec) - - gate_funcs = { - "givens": get_givens_generator, - "charge_charge": get_charge_charge_generator, - } - gate_func = gate_funcs[layer_spec.base_gate] - - return [gate_func(indices) for indices in indices_list] - - -def get_heuristic_gate_generators( - n_elec: int, layer_specs: Sequence[trial_wf.LayerSpec] -) -> List[of.FermionOperator]: - """Get gate generators for the heuristic ansatz. - - Args: - n_elec: The number of electrons. - layer_specs: The layer specifications. - - Returns: - A list of generators for the layers. - """ - gate_generators = [] - - for layer_spec in layer_specs: - gate_generators += get_layer_generators(layer_spec, n_elec) - - return gate_generators - - -def get_heuristic_circuit( - layer_specs: Sequence[trial_wf.LayerSpec], - n_elec: int, - params: np.ndarray, - fermion_index_to_qubit_map: Dict[int, cirq.GridQubit], -) -> cirq.Circuit: - """Get a circuit for the heuristic ansatz. - - Args: - layer_specs: The layer specs for the heuristic layers. - n_elec: The number of electrons. - params: The variational parameters for the circuit. - fermion_index_to_qubit_map: A mapping between fermion mode indices and qubits. - - Returns: - A circuit for the heuristic ansatz. - """ - gates: List[cirq.Operation] = [] - - for layer_spec in layer_specs: - params_slice = params[len(gates) :] - gates += get_layer_gates( - layer_spec, n_elec, params_slice, fermion_index_to_qubit_map - ) - - return cirq.Circuit(gates) - - def orbital_rotation_gradient_matrix( generator_mat: np.ndarray, a: int, b: int ) -> np.ndarray: @@ -1167,7 +480,7 @@ def get_pp_plus_params( restricted: bool = False, random_parameter_scale: float = 1.0, initial_orbital_rotation: Optional[np.ndarray] = None, - heuristic_layers: Tuple[trial_wf.LayerSpec, ...], + heuristic_layers: Tuple[layer_spec.LayerSpec, ...], do_pp: bool = True, n_optimization_restarts: int = 1, do_print: bool = True, @@ -1214,7 +527,7 @@ def get_pp_plus_params( err_msg = "use_fast_gradients does not work with initial orbital rotation." assert initial_orbital_rotation is None, err_msg - gate_generators = _get_pp_plus_gate_generators( + gate_generators = afqmc_generators.get_pp_plus_gate_generators( n_elec=n_elec, heuristic_layers=heuristic_layers, do_pp=do_pp ) diff --git a/recirq/qcqmc/optimize_wf_test.py b/recirq/qcqmc/optimize_wf_test.py index 1b392d00..cfbcc4ab 100644 --- a/recirq/qcqmc/optimize_wf_test.py +++ b/recirq/qcqmc/optimize_wf_test.py @@ -18,16 +18,19 @@ import numpy as np import pytest +from recirq.qcqmc.afqmc_generators import get_pp_plus_gate_generators +from recirq.qcqmc.converters import ( + get_ansatz_qubit_wf, + get_two_body_params_from_qchem_amplitudes, +) from recirq.qcqmc.hamiltonian import HamiltonianData +from recirq.qcqmc.layer_spec import LayerSpec from recirq.qcqmc.optimize_wf import ( - _get_ansatz_qubit_wf, - _get_pp_plus_gate_generators, build_pp_plus_trial_wavefunction, evaluate_gradient_and_cost_function, get_evolved_wf, - get_two_body_params_from_qchem_amplitudes, ) -from recirq.qcqmc.trial_wf import LayerSpec, PerfectPairingPlusTrialWavefunctionParams +from recirq.qcqmc.trial_wf import PerfectPairingPlusTrialWavefunctionParams def test_pp_wf_energy(fixture_4_qubit_ham: HamiltonianData): @@ -209,7 +212,7 @@ def test_qchem_pp_runs( do_print=False, ) - ansatz_qubit_wf = _get_ansatz_qubit_wf( + ansatz_qubit_wf = get_ansatz_qubit_wf( ansatz_circuit=trial_wf.ansatz_circuit, ordered_qubits=params.qubits_jordan_wigner_ordered, ) @@ -239,7 +242,7 @@ def test_qchem_conversion_negative(fixture_4_qubit_ham: HamiltonianData): do_print=False, ) - ansatz_qubit_wf = _get_ansatz_qubit_wf( + ansatz_qubit_wf = get_ansatz_qubit_wf( ansatz_circuit=trial_wf.ansatz_circuit, ordered_qubits=params.qubits_jordan_wigner_ordered, ) @@ -281,7 +284,7 @@ def compute_finite_difference_grad( initial_wf: The initial wavefunction (typically Hartree--Fock) restricted: Whether we're using a restricted ansatz or not. """ - generators = _get_pp_plus_gate_generators( + generators = get_pp_plus_gate_generators( n_elec=n_elec, heuristic_layers=tuple(), do_pp=True ) one_body_gradient = np.zeros_like(one_body_params) @@ -349,7 +352,7 @@ def test_gradient(n_elec, n_orb, restricted): else: n_one_body_params = n_orb * (n_orb - 1) - gate_generators = _get_pp_plus_gate_generators( + gate_generators = get_pp_plus_gate_generators( n_elec=n_elec, heuristic_layers=tuple(), do_pp=True ) # reference implementation diff --git a/recirq/qcqmc/trial_wf.py b/recirq/qcqmc/trial_wf.py index f6ddd3ce..33475b11 100644 --- a/recirq/qcqmc/trial_wf.py +++ b/recirq/qcqmc/trial_wf.py @@ -19,37 +19,15 @@ import cirq import numpy as np -from recirq.qcqmc import bitstrings, config, data, fermion_mode, hamiltonian, qubit_maps - - -@attrs.frozen -class LayerSpec: - """A specification of a hardware-efficient layer of gates. - - Args: - base_gate: 'charge_charge' for the e^{i n_i n_j} gate and 'givens' for a givens rotation. - layout: Should be 'in_pair', 'cross_pair', or 'cross_spin' only. - """ - - base_gate: str - layout: str - - def __post_init__(self): - if self.base_gate not in ["charge_charge", "givens"]: - raise ValueError( - f'base_gate is set to {self.base_gate}, it should be either "charge_charge or "givens".' - ) - if self.layout not in ["in_pair", "cross_pair", "cross_spin"]: - raise ValueError( - f'layout is set to {self.layout}, it should be either "cross_pair", "in_pair", or "cross_spin".' - ) - - @classmethod - def _json_namespace_(cls): - return "recirq.qcqmc" - - def _json_dict_(self): - return attrs.asdict(self) +from recirq.qcqmc import ( + bitstrings, + config, + data, + fermion_mode, + hamiltonian, + layer_spec, + qubit_maps, +) @attrs.frozen @@ -87,7 +65,7 @@ def _to_numpy(x: Optional[Iterable] = None) -> Optional[np.ndarray]: return np.asarray(x) -def _to_tuple(x: Iterable[LayerSpec]) -> Sequence[LayerSpec]: +def _to_tuple(x: Iterable[layer_spec.LayerSpec]) -> Sequence[layer_spec.LayerSpec]: return tuple(x) @@ -122,7 +100,9 @@ class PerfectPairingPlusTrialWavefunctionParams(TrialWavefunctionParams): name: str hamiltonian_params: hamiltonian.HamiltonianParams - heuristic_layers: Tuple[LayerSpec, ...] = attrs.field(converter=_to_tuple) + heuristic_layers: Tuple[layer_spec.LayerSpec, ...] = attrs.field( + converter=_to_tuple + ) do_pp: bool = True restricted: bool = False random_parameter_scale: float = 1.0 From 063b769643c26649d2296a922b6a2b9158cddeda Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 21:37:20 +0000 Subject: [PATCH 10/21] Docstring. --- recirq/qcqmc/optimize_wf.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index 36244812..1bd257a1 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -26,16 +26,8 @@ import scipy.optimize import scipy.sparse -from recirq.qcqmc import ( - afqmc_circuits, - afqmc_generators, - converters, - data, - fermion_mode, - hamiltonian, - layer_spec, - trial_wf, -) +from recirq.qcqmc import (afqmc_circuits, afqmc_generators, converters, data, + fermion_mode, hamiltonian, layer_spec, trial_wf) def get_and_check_energy( @@ -280,6 +272,16 @@ def get_energy_and_check_sanity( def get_one_body_cluster_coef( params: np.ndarray, n_orb: int, restricted: bool ) -> np.ndarray: + """Get the matrix elements associated with the one-body cluster operator. + + Args: + params: The variational parameters of the one-body cluster operator. + n_orb: The number of spatial orbitals. + restricted: Whether a spin-restricted cluster operator is used. + + Returns: + The one-body cluster operator matrix. + """ if restricted: one_body_cluster_op = np.zeros((n_orb, n_orb), dtype=np.complex128) else: From dcd008735526b59fd5cc4ccedaabe8e38bdd7e41 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 22:00:11 +0000 Subject: [PATCH 11/21] Tidy up optimization.?: --- recirq/qcqmc/optimize_wf.py | 261 +++++++++++++++++++++++++----------- 1 file changed, 185 insertions(+), 76 deletions(-) diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index 1bd257a1..c06d9900 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -26,8 +26,16 @@ import scipy.optimize import scipy.sparse -from recirq.qcqmc import (afqmc_circuits, afqmc_generators, converters, data, - fermion_mode, hamiltonian, layer_spec, trial_wf) +from recirq.qcqmc import ( + afqmc_circuits, + afqmc_generators, + converters, + data, + fermion_mode, + hamiltonian, + layer_spec, + trial_wf, +) def get_and_check_energy( @@ -391,7 +399,7 @@ def orbital_rotation_gradient_matrix( return pre_matrix_full -def evaluate_gradient_and_cost_function( +def evaluate_energy_and_gradient( initial_wf: fqe.Wavefunction, fqe_ham: fqe_hams.RestrictedHamiltonian, n_orb: int, @@ -415,7 +423,7 @@ def evaluate_gradient_and_cost_function( for the alpha- and beta-spin rotations. Returns: - cost_val: The cost function (total energy) evaluated for the input wavefunction parameters. + energy: The cost function (total energy) evaluated for the input wavefunction parameters. grad: An array of gradients with respect to the one- and two-body parameters. The first n_orb * (n_orb + 1) // 2 parameters correspond to the one-body gradients. @@ -476,6 +484,178 @@ def evaluate_gradient_and_cost_function( return cost_val, np.concatenate((two_body_grad, one_body_grad)) +def objective( + params: np.ndarray, + n_one_body_params: int, + n_two_body_params: int, + initial_wf: fqe_wfn.Wavefunction, + fqe_ham: fqe_hams.RestrictedHamiltonian, + gate_generators: List[of.FermionOperator], + n_orb: int, + restricted: bool, + initial_orbital_rotation: np.ndarray, + e_core: float, + do_print: bool = False, +) -> float: + """Evaluate the objective function (total energy) for a set of variational pareters.""" + one_body_params = params[-n_one_body_params:] + two_body_params = params[:n_two_body_params] + + wf, _ = get_evolved_wf( + one_body_params, + two_body_params, + initial_wf, + gate_generators, + n_orb, + restricted=restricted, + initial_orbital_rotation=initial_orbital_rotation, + ) + + energy = wf.expectationValue(fqe_ham) + e_core + if do_print: + print(f"energy {energy}") + if np.abs(energy.imag) < 1e-6: + return energy.real + else: + return 1e6 + + +def fast_objective_and_gradient( + params: np.ndarray, + n_one_body_params: int, + n_two_body_params: int, + initial_wf: fqe_wfn.Wavefunction, + fqe_ham: fqe_hams.RestrictedHamiltonian, + gate_generators: List[of.FermionOperator], + n_orb: int, + restricted: bool, + e_core: float, + do_print: bool = False, +) -> Tuple[float, float]: + one_body_params = params[-n_one_body_params:] + two_body_params = params[:n_two_body_params] + energy, grad = evaluate_energy_and_cost_function( + initial_wf, + fqe_ham, + n_orb, + one_body_params, + two_body_params, + gate_generators, + restricted, + e_core, + ) + if do_print: + print(f"energy {energy}, max|grad| {np.max(np.abs(grad))}") + if np.abs(energy.imag) < 1e-6: + return energy.real, grad + else: + return 1e6, 1e6 + + +def optimize_parameters( + initial_wf: fqe_wfn.Wavefunction, + gate_generators: List[of.FermionOperator], + n_orb: int, + n_one_body_params: int, + n_two_body_params: int, + initial_orbital_rotation: np.ndarray, + fqe_ham: fqe_hams.RestrictedHamiltonian, + e_core: float, + restricted: bool = False, + use_fast_gradients: bool = False, + n_optimization_restarts: int = 1, + random_parameter_scale: float = 1.0, + do_print: bool = True, +) -> Optional[scipy.optimize.OptimizeResult]: + """Optimize the cost function (total energy) for the PP+ ansatz. + + Loops over n_optimization_restarts to try to find a good minimum value of the total energy. + + Args: + initial_wf: The initial wavefunction the circuit unitary is applied to. + gate_generators: The generators of the two-body interaction terms. + n_orb: The number of orbitals. + n_one_body_params: The number of variational parameters for the one-body terms. + n_two_body_params: The number of variational parameters for the two-body terms. + initial_orbital_rotation: An optional initial orbital rotation matrix, + which will be implmented as a givens circuit. + fqe_ham: The restricted FQE hamiltonian. + e_core: The Hamiltonian core (all the constants) energy. + use_fast_gradients: Compute the parameter gradients anlytically using Wilcox formula. + Default to false (use finite difference gradients). + n_optimization_restarts: The number of times to restart the optimization + from a random guess in an attempt at global optimization. + restricted: Whether to use a spin-restricted ansatz or not. + random_parameter_scale: A float to scale the random parameters by. + do_print: Whether to print optimization progress to stdout. + + Returns: + The optimization result. + """ + best = np.inf + best_res: Optional[scipy.optimize.OptimizeResult] = None + for i in range(n_optimization_restarts): + if do_print: + print(f"Optimization restart {i}", flush=True) + + def progress_cb(_): + print(".", end="", flush=True) + + else: + + def progress_cb(_): + pass + + params = random_parameter_scale * np.random.normal( + size=(n_two_body_params + n_one_body_params) + ) + + if use_fast_gradients: + res = scipy.optimize.minimize( + fast_objective_and_gradient, + params, + jac=True, + method="BFGS", + callback=progress_cb, + args=( + n_one_body_params, + n_two_body_params, + initial_wf, + fqe_ham, + gate_generators, + n_orb, + restricted, + e_core, + do_print, + ), + ) + else: + res = scipy.optimize.minimize( + objective, + params, + callback=progress_cb, + args=( + n_one_body_params, + n_two_body_params, + initial_wf, + fqe_ham, + gate_generators, + n_orb, + restricted, + initial_orbital_rotation, + e_core, + do_print, + ), + ) + if res.fun < best: + best = res.fun + best_res = res + + if do_print: + print(res, flush=True) + return best_res + + def get_pp_plus_params( *, hamiltonian_data: hamiltonian.HamiltonianData, @@ -540,78 +720,7 @@ def get_pp_plus_params( else: n_one_body_params = n_orb * (n_orb - 1) - best = np.inf - best_res: Union[None, scipy.optimize.OptimizeResult] = None - for i in range(n_optimization_restarts): - if do_print: - print(f"Optimization restart {i}", flush=True) - - def progress_cb(_): - print(".", end="", flush=True) - - else: - - def progress_cb(_): - pass - - params = random_parameter_scale * np.random.normal( - size=(n_two_body_params + n_one_body_params) - ) - - def objective(params): - one_body_params = params[-n_one_body_params:] - two_body_params = params[:n_two_body_params] - - wf, _ = get_evolved_wf( - one_body_params, - two_body_params, - initial_wf, - gate_generators, - n_orb, - restricted=restricted, - initial_orbital_rotation=initial_orbital_rotation, - ) - - energy = wf.expectationValue(fqe_ham) + e_core - if do_print: - print(f"energy {energy}") - if np.abs(energy.imag) < 1e-6: - return energy.real - else: - return 1e6 - - def fast_obj_grad(params): - one_body_params = params[-n_one_body_params:] - two_body_params = params[:n_two_body_params] - energy, grad = evaluate_gradient_and_cost_function( - initial_wf, - fqe_ham, - n_orb, - one_body_params, - two_body_params, - gate_generators, - restricted, - e_core, - ) - if do_print: - print(f"energy {energy}, max|grad| {np.max(np.abs(grad))}") - if np.abs(energy.imag) < 1e-6: - return energy.real, grad - else: - return 1e6, 1e6 - - if use_fast_gradients: - res = scipy.optimize.minimize( - fast_obj_grad, params, jac=True, method="BFGS", callback=progress_cb - ) - else: - res = scipy.optimize.minimize(objective, params, callback=progress_cb) - if res.fun < best: - best = res.fun - best_res = res - - if do_print: - print(res, flush=True) + best_res = optimize_parameters() assert best_res is not None params = best_res.x From d12dd9db16ecae2b0e80780fbda6d49b8aef6e46 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 23:11:41 +0000 Subject: [PATCH 12/21] Comments. --- recirq/qcqmc/optimize_wf.py | 99 +++++++++++++++++++++++++------- recirq/qcqmc/optimize_wf_test.py | 4 +- 2 files changed, 80 insertions(+), 23 deletions(-) diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index c06d9900..93a7aa2f 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -26,16 +26,8 @@ import scipy.optimize import scipy.sparse -from recirq.qcqmc import ( - afqmc_circuits, - afqmc_generators, - converters, - data, - fermion_mode, - hamiltonian, - layer_spec, - trial_wf, -) +from recirq.qcqmc import (afqmc_circuits, afqmc_generators, converters, data, + fermion_mode, hamiltonian, layer_spec, trial_wf) def get_and_check_energy( @@ -411,6 +403,8 @@ def evaluate_energy_and_gradient( ) -> Tuple[float, np.ndarray]: """Evaluate gradient and cost function for optimization. + Uses the linear scaling algorithm at the expense of three copies of the wavefunction. + Args: initial_wf: Initial state (typically Hartree--Fock). fqe_ham: The restricted Hamiltonian in FQE format. @@ -440,11 +434,11 @@ def evaluate_energy_and_gradient( lam = lam.apply(fqe_ham) cost_val = fqe.vdot(lam, phi) + e_core - # 1body + # First build the 1body cluster op as a matrix one_body_cluster_op = get_one_body_cluster_coef( one_body_params, n_orb, restricted=restricted ) - tril = np.tril_indices(n_orb, k=-1) + # Build the one-body FQE hamiltonian if restricted: one_body_ham = fqe.get_restricted_hamiltonian((-1j * one_body_cluster_op,)) else: @@ -455,14 +449,25 @@ def evaluate_energy_and_gradient( one_body_grad = np.zeros_like(one_body_params) n_one_body_params = len(one_body_params) grad_position = n_one_body_params - 1 + # The parameters correspond to the lower triangular part of the matrix. + # we need the row and column indices corresponding to each flattened lower triangular index. + tril = np.tril_indices(n_orb, k=-1) + # Now compute the gradient of the one-body orbital rotation operator for each parameter. + # The basic idea is to use the fact that |psi(theta)> = U_p .... U_1 # |psi_0>, + # and progressively move the gradient operator through the expression saving + # intermediate wavefunctions along the way. + # The gradient is then related to the overlap of a left and right wavefunction. for iparam in range(len(one_body_params)): mu_state = copy.deepcopy(phi) + # get the parameter index starting from the end and working backwards. pidx = n_one_body_params - iparam - 1 pidx_spin = 0 if restricted else pidx // (n_one_body_params // 2) pidx_spat = pidx if restricted else pidx - (n_one_body_params // 2) * pidx_spin + # Get the actual row and column indicies corresponding to this parameter index. p, q = (tril[0][pidx_spat], tril[1][pidx_spat]) p += n_orb * pidx_spin q += n_orb * pidx_spin + # Get the orbital rotation gradient "pre" matrix and apply it the |mu>. pre_matrix = orbital_rotation_gradient_matrix(-one_body_cluster_op, p, q) assert of.is_hermitian(1j * pre_matrix) if restricted: @@ -473,6 +478,9 @@ def evaluate_energy_and_gradient( one_body_grad[grad_position] = 2 * fqe.vdot(lam, mu_state).real grad_position -= 1 # Get two-body contributions + # Here we already have the generators so the gradient is simple to evaluate + # as the derivative just brings down a generator which we need to apply to + # the state before computing the overlap. two_body_grad = np.zeros(len(two_body_params)) for pidx in reversed(range(len(gate_generators))): mu = copy.deepcopy(phi) @@ -497,7 +505,23 @@ def objective( e_core: float, do_print: bool = False, ) -> float: - """Evaluate the objective function (total energy) for a set of variational pareters.""" + """Helper function to compute energy from the variational parameters. + + Args: + params: A packed array containing the one and two-body parameters. + n_one_body_params: The number of variational parameters for the one-body terms. + n_two_body_params: The number of variational parameters for the two-body terms. + initial_wf: The initial wavefunction the circuit unitary is applied to. + fqe_ham: The restricted FQE hamiltonian. + gate_generators: The list of gate generators. + n_orb: The number of spatial orbitals. + restricted: Whether to use a spin-restricted ansatz or not. + e_core: The Hamiltonian core (all the constants) energy. + do_print: Whether to print optimization progress to stdout. + + Returns: + The energy capped at 1e6 if the energy is imaginary. + """ one_body_params = params[-n_one_body_params:] two_body_params = params[:n_two_body_params] @@ -520,7 +544,7 @@ def objective( return 1e6 -def fast_objective_and_gradient( +def objective_and_gradient( params: np.ndarray, n_one_body_params: int, n_two_body_params: int, @@ -531,10 +555,28 @@ def fast_objective_and_gradient( restricted: bool, e_core: float, do_print: bool = False, -) -> Tuple[float, float]: +) -> Tuple[float, np.array]: + """Helper function to compute energy and gradient from the variational parameters + + Args: + params: A packed array containing the one and two-body parameters. + n_one_body_params: The number of variational parameters for the one-body terms. + n_two_body_params: The number of variational parameters for the two-body terms. + initial_wf: The initial wavefunction the circuit unitary is applied to. + fqe_ham: The restricted FQE hamiltonian. + gate_generators: The list of gate generators. + n_orb: The number of spatial orbitals. + restricted: Whether to use a spin-restricted ansatz or not. + e_core: The Hamiltonian core (all the constants) energy. + do_print: Whether to print optimization progress to stdout. + + Returns: + A tuple containing the energy and gradient. These are capped at 1e6 if + the energy is imaginary. + """ one_body_params = params[-n_one_body_params:] two_body_params = params[:n_two_body_params] - energy, grad = evaluate_energy_and_cost_function( + energy, grad = evaluate_energy_and_gradient( initial_wf, fqe_ham, n_orb, @@ -547,9 +589,9 @@ def fast_objective_and_gradient( if do_print: print(f"energy {energy}, max|grad| {np.max(np.abs(grad))}") if np.abs(energy.imag) < 1e-6: - return energy.real, grad + return energy.real, grad.real else: - return 1e6, 1e6 + return 1e6, np.array([1e6]) * len(grad) def optimize_parameters( @@ -558,9 +600,9 @@ def optimize_parameters( n_orb: int, n_one_body_params: int, n_two_body_params: int, - initial_orbital_rotation: np.ndarray, fqe_ham: fqe_hams.RestrictedHamiltonian, e_core: float, + initial_orbital_rotation: Optional[np.ndarray] = None, restricted: bool = False, use_fast_gradients: bool = False, n_optimization_restarts: int = 1, @@ -611,8 +653,9 @@ def progress_cb(_): ) if use_fast_gradients: + # Use analytic gradient rather than finite differences. res = scipy.optimize.minimize( - fast_objective_and_gradient, + objective_and_gradient, params, jac=True, method="BFGS", @@ -720,7 +763,21 @@ def get_pp_plus_params( else: n_one_body_params = n_orb * (n_orb - 1) - best_res = optimize_parameters() + best_res = optimize_parameters( + initial_wf, + gate_generators, + n_orb, + n_one_body_params, + n_two_body_params, + fqe_ham, + e_core, + restricted=restricted, + initial_orbital_rotation=initial_orbital_rotation, + use_fast_gradients=use_fast_gradients, + n_optimization_restarts=n_optimization_restarts, + random_parameter_scale=random_parameter_scale, + do_print=do_print, + ) assert best_res is not None params = best_res.x diff --git a/recirq/qcqmc/optimize_wf_test.py b/recirq/qcqmc/optimize_wf_test.py index cfbcc4ab..0db3a8a3 100644 --- a/recirq/qcqmc/optimize_wf_test.py +++ b/recirq/qcqmc/optimize_wf_test.py @@ -27,7 +27,7 @@ from recirq.qcqmc.layer_spec import LayerSpec from recirq.qcqmc.optimize_wf import ( build_pp_plus_trial_wavefunction, - evaluate_gradient_and_cost_function, + evaluate_energy_and_gradient, get_evolved_wf, ) from recirq.qcqmc.trial_wf import PerfectPairingPlusTrialWavefunctionParams @@ -366,7 +366,7 @@ def test_gradient(n_elec, n_orb, restricted): n_orb, restricted=restricted, )[0] - obj_val, grad = evaluate_gradient_and_cost_function( + obj_val, grad = evaluate_energy_and_gradient( initial_wf, fqe_ham, n_orb, From 80749761de90c9c61837487dffda3e92c4cb9ec4 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 23:30:12 +0000 Subject: [PATCH 13/21] Add note to reference. --- recirq/qcqmc/optimize_wf.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index 93a7aa2f..a58bba84 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -403,7 +403,9 @@ def evaluate_energy_and_gradient( ) -> Tuple[float, np.ndarray]: """Evaluate gradient and cost function for optimization. - Uses the linear scaling algorithm at the expense of three copies of the wavefunction. + Uses the linear scaling algorithm (see algo 2 from example: + https://arxiv.org/pdf/2009.02823) at the expense of three copies of the + wavefunction. Args: initial_wf: Initial state (typically Hartree--Fock). @@ -422,6 +424,7 @@ def evaluate_energy_and_gradient( parameters. The first n_orb * (n_orb + 1) // 2 parameters correspond to the one-body gradients. """ + # Build |phi> = U(theta)|phi_0> phi = get_evolved_wf( one_body_params, two_body_params, @@ -430,8 +433,11 @@ def evaluate_energy_and_gradient( n_orb, restricted=restricted, )[0] + # Set |lambda> = |phi> initially lam = copy.deepcopy(phi) + # H|lambda> lam = lam.apply(fqe_ham) + # E = cost_val = fqe.vdot(lam, phi) + e_core # First build the 1body cluster op as a matrix @@ -443,7 +449,8 @@ def evaluate_energy_and_gradient( one_body_ham = fqe.get_restricted_hamiltonian((-1j * one_body_cluster_op,)) else: one_body_ham = fqe.get_sso_hamiltonian((-1j * one_body_cluster_op,)) - # Apply U1b^{dag} + # |phi> = U1b U2b |phi_0> + # 1. Remove U1b from |phi> by U1b^dagger |phi> phi.time_evolve(1, one_body_ham, inplace=True) lam.time_evolve(1, one_body_ham, inplace=True) one_body_grad = np.zeros_like(one_body_params) @@ -453,10 +460,8 @@ def evaluate_energy_and_gradient( # we need the row and column indices corresponding to each flattened lower triangular index. tril = np.tril_indices(n_orb, k=-1) # Now compute the gradient of the one-body orbital rotation operator for each parameter. - # The basic idea is to use the fact that |psi(theta)> = U_p .... U_1 # |psi_0>, - # and progressively move the gradient operator through the expression saving - # intermediate wavefunctions along the way. - # The gradient is then related to the overlap of a left and right wavefunction. + # If we write E(theta) = + # Then d E(theta)/ d theta_p = -2 i Im for iparam in range(len(one_body_params)): mu_state = copy.deepcopy(phi) # get the parameter index starting from the end and working backwards. @@ -468,6 +473,8 @@ def evaluate_energy_and_gradient( p += n_orb * pidx_spin q += n_orb * pidx_spin # Get the orbital rotation gradient "pre" matrix and apply it the |mu>. + # For the orbital rotation part we compute dU(theta)/dtheta_p using the + # wilcox identity (see e.g.: https://arxiv.org/abs/2004.04174.) pre_matrix = orbital_rotation_gradient_matrix(-one_body_cluster_op, p, q) assert of.is_hermitian(1j * pre_matrix) if restricted: From 001c5ed70631240e17b4e81e85d8b8d605a36dc5 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 23:32:22 +0000 Subject: [PATCH 14/21] Tidy comment. --- recirq/qcqmc/optimize_wf.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index a58bba84..6507f62c 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -26,8 +26,16 @@ import scipy.optimize import scipy.sparse -from recirq.qcqmc import (afqmc_circuits, afqmc_generators, converters, data, - fermion_mode, hamiltonian, layer_spec, trial_wf) +from recirq.qcqmc import ( + afqmc_circuits, + afqmc_generators, + converters, + data, + fermion_mode, + hamiltonian, + layer_spec, + trial_wf, +) def get_and_check_energy( @@ -403,8 +411,8 @@ def evaluate_energy_and_gradient( ) -> Tuple[float, np.ndarray]: """Evaluate gradient and cost function for optimization. - Uses the linear scaling algorithm (see algo 2 from example: - https://arxiv.org/pdf/2009.02823) at the expense of three copies of the + Uses the linear scaling algorithm (see algo 1 from + https://arxiv.org/pdf/2009.02823 for example) at the expense of three copies of the wavefunction. Args: From 4387efc765fe267a8bcaa7e70c83869f135d4e90 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 17:07:12 -0700 Subject: [PATCH 15/21] Remove empty file. --- recirq/qcqmc/trial_wf_test.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 recirq/qcqmc/trial_wf_test.py diff --git a/recirq/qcqmc/trial_wf_test.py b/recirq/qcqmc/trial_wf_test.py deleted file mode 100644 index e69de29b..00000000 From 2a02fc3428016ba2f403ad566b2354f7bc65f386 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 17:14:26 -0700 Subject: [PATCH 16/21] Add missing docstring. --- recirq/qcqmc/afqmc_circuits.py | 59 ++++++++++++++++++++++++++++------ 1 file changed, 50 insertions(+), 9 deletions(-) diff --git a/recirq/qcqmc/afqmc_circuits.py b/recirq/qcqmc/afqmc_circuits.py index 612bbc4e..499e26e9 100644 --- a/recirq/qcqmc/afqmc_circuits.py +++ b/recirq/qcqmc/afqmc_circuits.py @@ -21,9 +21,8 @@ import cirq import numpy as np import openfermion -from openfermion.circuits.primitives.state_preparation import ( - _ops_from_givens_rotations_circuit_description, -) +from openfermion.circuits.primitives.state_preparation import \ + _ops_from_givens_rotations_circuit_description from openfermion.linalg.givens_rotations import givens_decomposition from openfermion.linalg.sparse_tools import jw_sparse_givens_rotation @@ -538,6 +537,14 @@ def get_4_qubit_pp_circuits( We map the fermionic orbitals to grid qubits like so: 3 1 2 0 + + Args: + two_body_params: The parameters of the two-body terms in the perfect pairing ansatz. + n_elec: The number of electrons. + heuristic_layers: A tuple of hardware efficient layers. + + Returns: + The superposition and ansatz circuits. """ assert n_elec == 2 @@ -574,11 +581,19 @@ def get_8_qubit_circuits( n_elec: int, heuristic_layers: Tuple[lspec.LayerSpec, ...], ) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. + """A helper function that builds the circuits for the 8 qubit ansatz. We map the fermionic orbitals to grid qubits like so: 3 5 1 7 2 4 0 6 + + Args: + two_body_params: The parameters of the two-body terms in the perfect pairing ansatz. + n_elec: The number of electrons. + heuristic_layers: A tuple of hardware efficient layers. + + Returns: + The superposition and ansatz circuits. """ fermion_index_to_qubit_map = qubit_maps.get_8_qubit_fermion_qubit_map() @@ -645,11 +660,19 @@ def get_12_qubit_circuits( n_elec: int, heuristic_layers: Tuple[lspec.LayerSpec, ...], ) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. + """A helper function that builds the circuits for the twelve qubit ansatz. We map the fermionic orbitals to grid qubits like so: 5 7 3 9 1 11 4 6 2 8 0 10 + + Args: + two_body_params: The parameters of the two-body terms in the perfect pairing ansatz. + n_elec: The number of electrons. + heuristic_layers: A tuple of hardware efficient layers. + + Returns: + The superposition and ansatz circuits. """ fermion_index_to_qubit_map = qubit_maps.get_12_qubit_fermion_qubit_map() @@ -733,11 +756,19 @@ def get_16_qubit_circuits( n_elec: int, heuristic_layers: Tuple[lspec.LayerSpec, ...], ) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A helper function that builds the circuits for the four qubit ansatz. + """A helper function that builds the circuits for the 16 qubit ansatz. We map the fermionic orbitals to grid qubits like so: 7 9 5 11 3 13 1 15 6 8 4 10 2 12 0 14 + + Args: + two_body_params: The parameters of the two-body terms in the perfect pairing ansatz. + n_elec: The number of electrons. + heuristic_layers: A tuple of hardware efficient layers. + + Returns: + The superposition and ansatz circuits. """ fermion_index_to_qubit_map = qubit_maps.get_16_qubit_fermion_qubit_map() @@ -843,7 +874,17 @@ def get_circuits( n_elec: int, heuristic_layers: Tuple[lspec.LayerSpec, ...], ) -> Tuple[cirq.Circuit, cirq.Circuit]: - """A function that runs a specialized method to get the ansatz circuits.""" + """A function that runs a specialized method to get the ansatz circuits. + + Args: + two_body_params: The parameters of the two-body terms in the perfect pairing ansatz. + n_orb: The number of spatial orbitals. + n_elec: The number of electrons. + heuristic_layers: A tuple of hardware efficient layers. + + Returns: + The superposition and ansatz circuits. + """ if n_orb != n_elec: raise ValueError("n_orb must equal n_elec.") @@ -856,8 +897,8 @@ def get_circuits( } try: circ_func = circ_funcs[n_orb] - except KeyError: - raise NotImplementedError(f"No circuits for n_orb = {n_orb}") + except KeyError as exc: + raise NotImplementedError(f"No circuits for n_orb = {n_orb}") from exc return circ_func( two_body_params=two_body_params, From 97553d2b7185b745698dfed4a72017458ef31f2f Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 17:15:00 -0700 Subject: [PATCH 17/21] Remove whitespace. --- recirq/qcqmc/afqmc_circuits.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/recirq/qcqmc/afqmc_circuits.py b/recirq/qcqmc/afqmc_circuits.py index 499e26e9..b9951f5d 100644 --- a/recirq/qcqmc/afqmc_circuits.py +++ b/recirq/qcqmc/afqmc_circuits.py @@ -21,8 +21,9 @@ import cirq import numpy as np import openfermion -from openfermion.circuits.primitives.state_preparation import \ - _ops_from_givens_rotations_circuit_description +from openfermion.circuits.primitives.state_preparation import ( + _ops_from_givens_rotations_circuit_description, +) from openfermion.linalg.givens_rotations import givens_decomposition from openfermion.linalg.sparse_tools import jw_sparse_givens_rotation @@ -875,7 +876,7 @@ def get_circuits( heuristic_layers: Tuple[lspec.LayerSpec, ...], ) -> Tuple[cirq.Circuit, cirq.Circuit]: """A function that runs a specialized method to get the ansatz circuits. - + Args: two_body_params: The parameters of the two-body terms in the perfect pairing ansatz. n_orb: The number of spatial orbitals. From 595a0b73b67dd0348b304a086bf4015ca99bf397 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 17:16:38 -0700 Subject: [PATCH 18/21] Add missing license. --- recirq/qcqmc/qubit_maps.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/recirq/qcqmc/qubit_maps.py b/recirq/qcqmc/qubit_maps.py index c2ea3f9d..0beb81e2 100644 --- a/recirq/qcqmc/qubit_maps.py +++ b/recirq/qcqmc/qubit_maps.py @@ -1,3 +1,16 @@ +# Copyright 2024 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. from typing import Dict, Tuple import cirq From e342b8b9daa8cd576f76e56304174863eeaf99a8 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Mon, 24 Jun 2024 17:29:49 -0700 Subject: [PATCH 19/21] Add module level description. --- recirq/qcqmc/afqmc_circuits.py | 1 - recirq/qcqmc/afqmc_generators.py | 10 ++++++++++ recirq/qcqmc/fermion_mode.py | 2 ++ recirq/qcqmc/layer_spec.py | 1 + recirq/qcqmc/optimize_wf.py | 13 +++---------- recirq/qcqmc/qubit_maps.py | 1 + recirq/qcqmc/trial_wf.py | 1 + 7 files changed, 18 insertions(+), 11 deletions(-) diff --git a/recirq/qcqmc/afqmc_circuits.py b/recirq/qcqmc/afqmc_circuits.py index b9951f5d..f7b01488 100644 --- a/recirq/qcqmc/afqmc_circuits.py +++ b/recirq/qcqmc/afqmc_circuits.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - """Trial wavefunction circuit ansatz primitives.""" import itertools diff --git a/recirq/qcqmc/afqmc_generators.py b/recirq/qcqmc/afqmc_generators.py index 78987ce0..f47614bc 100644 --- a/recirq/qcqmc/afqmc_generators.py +++ b/recirq/qcqmc/afqmc_generators.py @@ -11,6 +11,16 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +r"""Module for building generators for various fermionic wavefunction ansatzes for AFQMC. + +Given a unitary defined in terms of fermionic modes of the form +$$ +U = e^{\alpha a_p^\dagger a_q} +$$ +the generator of this unitary (gate) is $a_p^\dagger a_q$ where $a_p^{(\dagger)} +$ is the fermionic annihilation (creation) operator for a fermion in spin +orbital $p$. +""" from typing import List, Sequence, Tuple diff --git a/recirq/qcqmc/fermion_mode.py b/recirq/qcqmc/fermion_mode.py index c37fe4b3..8ffd5853 100644 --- a/recirq/qcqmc/fermion_mode.py +++ b/recirq/qcqmc/fermion_mode.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Specification of a fermion mode.""" import attrs @@ -21,6 +22,7 @@ class FermionicMode: Args: orb_ind: The spatial orbital index. spin: The spin state of the fermion mode (up or down (alpha or beta)). + Must be either 'a' or 'b'. """ orb_ind: int diff --git a/recirq/qcqmc/layer_spec.py b/recirq/qcqmc/layer_spec.py index a1b046d2..42efd550 100644 --- a/recirq/qcqmc/layer_spec.py +++ b/recirq/qcqmc/layer_spec.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Specification of hardware efficient layers and utility functions for indexing.""" from typing import Iterator, List, Tuple diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index 6507f62c..e1ffc6eb 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Functions for the variational optimizion of a trial wavefunction.""" import copy import itertools @@ -26,16 +27,8 @@ import scipy.optimize import scipy.sparse -from recirq.qcqmc import ( - afqmc_circuits, - afqmc_generators, - converters, - data, - fermion_mode, - hamiltonian, - layer_spec, - trial_wf, -) +from recirq.qcqmc import (afqmc_circuits, afqmc_generators, converters, data, + fermion_mode, hamiltonian, layer_spec, trial_wf) def get_and_check_energy( diff --git a/recirq/qcqmc/qubit_maps.py b/recirq/qcqmc/qubit_maps.py index 0beb81e2..b8151940 100644 --- a/recirq/qcqmc/qubit_maps.py +++ b/recirq/qcqmc/qubit_maps.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Various mappings between fermions and qubits.""" from typing import Dict, Tuple import cirq diff --git a/recirq/qcqmc/trial_wf.py b/recirq/qcqmc/trial_wf.py index 33475b11..0dfb43fb 100644 --- a/recirq/qcqmc/trial_wf.py +++ b/recirq/qcqmc/trial_wf.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Specification of a trial wavefunction.""" import abc from typing import Dict, Iterable, Optional, Sequence, Tuple From 1aad9b443a98d9290d8e977e8738c79082783c60 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Wed, 26 Jun 2024 07:57:56 -0700 Subject: [PATCH 20/21] Move finite difference grad out of the test file. --- recirq/qcqmc/optimize_wf.py | 75 ++++++++++++++++++++++++++++++ recirq/qcqmc/optimize_wf_test.py | 79 +------------------------------- 2 files changed, 76 insertions(+), 78 deletions(-) diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index e1ffc6eb..704a0001 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -499,6 +499,81 @@ def evaluate_energy_and_gradient( return cost_val, np.concatenate((two_body_grad, one_body_grad)) +def compute_finite_difference_grad( + n_orb: int, + n_elec: int, + one_body_params: np.ndarray, + two_body_params: np.ndarray, + ham: fqe_hams.RestrictedHamiltonian, + initial_wf: fqe_wfn.Wavefunction, + dtheta: float = 1e-4, + restricted: bool = False, +): + """Compute the parameter gradient using finite differences. + + Args: + n_orb: the number of spatial orbitals. + n_elec: the number of electrons. + one_body_params: The variational parameters for the one-body terms in the ansatz. + two_body_params: The variational parameters for the two-body terms in the ansatz. + ham: The restricted FQE Hamiltonian. + initial_wf: The initial wavefunction (typically Hartree--Fock) + restricted: Whether we're using a restricted ansatz or not. + """ + generators = get_pp_plus_gate_generators( + n_elec=n_elec, heuristic_layers=tuple(), do_pp=True + ) + one_body_gradient = np.zeros_like(one_body_params) + for ig, _ in enumerate(one_body_gradient): + new_param = one_body_params.copy() + new_param[ig] = new_param[ig] + dtheta + phi = get_evolved_wf( + new_param, + two_body_params, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_plus = phi.expectationValue(ham) + new_param[ig] = new_param[ig] - 2 * dtheta + phi = get_evolved_wf( + new_param, + two_body_params, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_minu = phi.expectationValue(ham) + one_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) + two_body_gradient = np.zeros_like(two_body_params) + for ig, _ in enumerate(two_body_gradient): + new_param = two_body_params.copy() + new_param[ig] = new_param[ig] + dtheta + phi = get_evolved_wf( + one_body_params, + new_param, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_plus = phi.expectationValue(ham) + new_param[ig] = new_param[ig] - 2 * dtheta + phi = get_evolved_wf( + one_body_params, + new_param, + initial_wf, + generators, + n_orb, + restricted=restricted, + )[0] + e_minu = phi.expectationValue(ham) + two_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) + return one_body_gradient, two_body_gradient + + def objective( params: np.ndarray, diff --git a/recirq/qcqmc/optimize_wf_test.py b/recirq/qcqmc/optimize_wf_test.py index 0db3a8a3..c3788ba6 100644 --- a/recirq/qcqmc/optimize_wf_test.py +++ b/recirq/qcqmc/optimize_wf_test.py @@ -27,6 +27,7 @@ from recirq.qcqmc.layer_spec import LayerSpec from recirq.qcqmc.optimize_wf import ( build_pp_plus_trial_wavefunction, + compute_finite_difference_grad, evaluate_energy_and_gradient, get_evolved_wf, ) @@ -118,9 +119,6 @@ def test_pp_plus_wf_energy_sloppy_1(fixture_8_qubit_ham: HamiltonianData): assert trial_wf.ansatz_energy < -1.947 -# TODO: Speed up this test and add a similar one with non-trivial heuristic layers. - - def test_diamond_pp_wf_energy(fixture_12_qubit_ham: HamiltonianData): params = PerfectPairingPlusTrialWavefunctionParams( name="diamind_pp_test_wf_1", @@ -263,81 +261,6 @@ def gen_random_restricted_ham(n_orb: int) -> fqe_hams.RestrictedHamiltonian: return fqe_ham -def compute_finite_difference_grad( - n_orb: int, - n_elec: int, - one_body_params: np.ndarray, - two_body_params: np.ndarray, - ham: fqe_hams.RestrictedHamiltonian, - initial_wf: fqe_wfn.Wavefunction, - dtheta: float = 1e-4, - restricted: bool = False, -): - """Compute the parameter gradient using finite differences. - - Args: - n_orb: the number of spatial orbitals. - n_elec: the number of electrons. - one_body_params: The variational parameters for the one-body terms in the ansatz. - two_body_params: The variational parameters for the two-body terms in the ansatz. - ham: The restricted FQE Hamiltonian. - initial_wf: The initial wavefunction (typically Hartree--Fock) - restricted: Whether we're using a restricted ansatz or not. - """ - generators = get_pp_plus_gate_generators( - n_elec=n_elec, heuristic_layers=tuple(), do_pp=True - ) - one_body_gradient = np.zeros_like(one_body_params) - for ig, _ in enumerate(one_body_gradient): - new_param = one_body_params.copy() - new_param[ig] = new_param[ig] + dtheta - phi = get_evolved_wf( - new_param, - two_body_params, - initial_wf, - generators, - n_orb, - restricted=restricted, - )[0] - e_plus = phi.expectationValue(ham) - new_param[ig] = new_param[ig] - 2 * dtheta - phi = get_evolved_wf( - new_param, - two_body_params, - initial_wf, - generators, - n_orb, - restricted=restricted, - )[0] - e_minu = phi.expectationValue(ham) - one_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) - two_body_gradient = np.zeros_like(two_body_params) - for ig, _ in enumerate(two_body_gradient): - new_param = two_body_params.copy() - new_param[ig] = new_param[ig] + dtheta - phi = get_evolved_wf( - one_body_params, - new_param, - initial_wf, - generators, - n_orb, - restricted=restricted, - )[0] - e_plus = phi.expectationValue(ham) - new_param[ig] = new_param[ig] - 2 * dtheta - phi = get_evolved_wf( - one_body_params, - new_param, - initial_wf, - generators, - n_orb, - restricted=restricted, - )[0] - e_minu = phi.expectationValue(ham) - two_body_gradient[ig] = (e_plus - e_minu).real / (2 * dtheta) - return one_body_gradient, two_body_gradient - - @pytest.mark.parametrize("n_elec, n_orb", ((2, 2), (4, 4), (6, 6))) @pytest.mark.parametrize("restricted", (True, False)) def test_gradient(n_elec, n_orb, restricted): From 0f297fb37457c44f7430cc04f96e6e2d936b9338 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Wed, 26 Jun 2024 08:03:22 -0700 Subject: [PATCH 21/21] Fix test failure. --- recirq/qcqmc/optimize_wf.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/recirq/qcqmc/optimize_wf.py b/recirq/qcqmc/optimize_wf.py index 704a0001..49e1c4ae 100644 --- a/recirq/qcqmc/optimize_wf.py +++ b/recirq/qcqmc/optimize_wf.py @@ -27,8 +27,16 @@ import scipy.optimize import scipy.sparse -from recirq.qcqmc import (afqmc_circuits, afqmc_generators, converters, data, - fermion_mode, hamiltonian, layer_spec, trial_wf) +from recirq.qcqmc import ( + afqmc_circuits, + afqmc_generators, + converters, + data, + fermion_mode, + hamiltonian, + layer_spec, + trial_wf, +) def get_and_check_energy( @@ -499,6 +507,7 @@ def evaluate_energy_and_gradient( return cost_val, np.concatenate((two_body_grad, one_body_grad)) + def compute_finite_difference_grad( n_orb: int, n_elec: int, @@ -520,7 +529,7 @@ def compute_finite_difference_grad( initial_wf: The initial wavefunction (typically Hartree--Fock) restricted: Whether we're using a restricted ansatz or not. """ - generators = get_pp_plus_gate_generators( + generators = afqmc_generators.get_pp_plus_gate_generators( n_elec=n_elec, heuristic_layers=tuple(), do_pp=True ) one_body_gradient = np.zeros_like(one_body_params) @@ -574,7 +583,6 @@ def compute_finite_difference_grad( return one_body_gradient, two_body_gradient - def objective( params: np.ndarray, n_one_body_params: int,