diff --git a/pyproject.toml b/pyproject.toml index 1b86e2d..99c3701 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,12 @@ dependencies = [ "tensornetwork>=0.4.6", "torch>=2.7.1", "matplotlib", + "pandas", + "pyscf", + "openfermion", + "pylatexenc", + "noisyopt", + "renormalizer==0.0.11", ] [project.optional-dependencies] diff --git a/src/tyxonq/chem/__init__.py b/src/tyxonq/chem/__init__.py new file mode 100644 index 0000000..1e520e2 --- /dev/null +++ b/src/tyxonq/chem/__init__.py @@ -0,0 +1,65 @@ +__version__ = "0.1.0" +__author__ = "TyxonQ" + +# Forked from https://github.com/tencent-quantum-lab/TenCirChem +# Reconstruction IS IN PROGRESS + + +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + +import os +import logging + +os.environ["JAX_ENABLE_X64"] = "True" +# for debugging +# os.environ["CUDA_VISIBLE_DEVICES"] = "-1" + +# disable CUDA 11.1 warning +os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" + +logger = logging.getLogger("tyxonq") +logger.setLevel(logging.FATAL) + +os.environ["RENO_LOG_LEVEL"] = "100" + +logger = logging.getLogger("tencirchem") +logger.setLevel(logging.WARNING) + +# finish logger stuff +del logger + +from .utils.backend import set_backend, set_dtype + +# by default use float64 rather than float32 +set_dtype("complex128") + +# static module +# as an external interface +from pyscf import M + +from .static.ucc import UCC +from .static.uccsd import UCCSD, ROUCCSD +from .static.kupccgsd import KUPCCGSD +from .static.puccd import PUCCD +from .static.hea import HEA, parity, binary, get_noise_conf + +# dynamic module +# as an external interface +from renormalizer import Op, BasisSHO, BasisHalfSpin, BasisSimpleElectron, BasisMultiElectron, Model, Mpo +from renormalizer.model import OpSum + +from .utils.misc import get_dense_operator +from .dynamic.time_evolution import TimeEvolution + + +def clear_cache(): + from .utils.backend import ALL_JIT_LIBS + from .static.evolve_civector import CI_OPERATOR_CACHE, CI_OPERATOR_BATCH_CACHE + + for l in ALL_JIT_LIBS: + l.clear() + CI_OPERATOR_CACHE.clear() + CI_OPERATOR_BATCH_CACHE.clear() diff --git a/src/tyxonq/chem/constants.py b/src/tyxonq/chem/constants.py new file mode 100644 index 0000000..0421d9c --- /dev/null +++ b/src/tyxonq/chem/constants.py @@ -0,0 +1,20 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + +import numpy as np + +a = np.array([[0, 1], [0, 0]]) +# a^\dagger +ad = a.T +# a^\dagger_i a_j +ad_a = np.kron(ad, a) +# a^\dagger_i a_j - a_j a^\dagger_i +ad_a_hc = ad_a - ad_a.T +adad_aa = np.kron(np.kron(np.kron(ad, ad), a), a) +adad_aa_hc = adad_aa - adad_aa.T +ad_a_hc2 = ad_a_hc @ ad_a_hc +adad_aa_hc2 = adad_aa_hc @ adad_aa_hc + +DISCARD_EPS = 1e-12 diff --git a/src/tyxonq/chem/dynamic/__init__.py b/src/tyxonq/chem/dynamic/__init__.py new file mode 100644 index 0000000..1303e89 --- /dev/null +++ b/src/tyxonq/chem/dynamic/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from ..dynamic.model import sbm, pyrazine +from ..dynamic.transform import qubit_encode_op, qubit_encode_op_grouped, qubit_encode_basis +from ..dynamic.time_derivative import get_ansatz, get_jacobian_func, get_deriv, regularized_inversion +from ..dynamic.time_evolution import TimeEvolution diff --git a/src/tyxonq/chem/dynamic/model/__init__.py b/src/tyxonq/chem/dynamic/model/__init__.py new file mode 100644 index 0000000..0cf883e --- /dev/null +++ b/src/tyxonq/chem/dynamic/model/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. diff --git a/src/tyxonq/chem/dynamic/model/pyrazine.py b/src/tyxonq/chem/dynamic/model/pyrazine.py new file mode 100644 index 0000000..447e1cb --- /dev/null +++ b/src/tyxonq/chem/dynamic/model/pyrazine.py @@ -0,0 +1,140 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +import logging +from itertools import permutations as permut +from itertools import product + + +import numpy as np +from renormalizer.model import Op +from renormalizer.model.basis import BasisSHO, BasisMultiElectron +from renormalizer.utils.constant import ev2au + + +logger = logging.getLogger(__name__) + + +r""" +Bi-linear vibronic coupling model for Pyrazine, 4 modes +See: Raab, Worth, Meyer, Cederbaum. J.Chem.Phys. 110 (1999) 936 +The parameters are from heidelberg mctdh package pyr4+.op +""" + + +# frequencies +w10a = 0.1139 * ev2au +w6a = 0.0739 * ev2au +w1 = 0.1258 * ev2au +w9a = 0.1525 * ev2au + +# energy-gap +delta = 0.42300 * ev2au + +# linear, on-diagonal coupling coefficients +# H(1,1) +_6a_s1_s1 = 0.09806 * ev2au +_1_s1_s1 = 0.05033 * ev2au +_9a_s1_s1 = 0.14521 * ev2au +# H(2,2) +_6a_s2_s2 = -0.13545 * ev2au +_1_s2_s2 = 0.17100 * ev2au +_9a_s2_s2 = 0.03746 * ev2au + +# quadratic, on-diagonal coupling coefficients +# H(1,1) +_10a_10a_s1_s1 = -0.01159 * ev2au +_6a_6a_s1_s1 = 0.00000 * ev2au +_1_1_s1_s1 = 0.00000 * ev2au +_9a_9a_s1_s1 = 0.00000 * ev2au +# H(2,2) +_10a_10a_s2_s2 = -0.01159 * ev2au +_6a_6a_s2_s2 = 0.00000 * ev2au +_1_1_s2_s2 = 0.00000 * ev2au +_9a_9a_s2_s2 = 0.00000 * ev2au + +# bilinear, on-diagonal coupling coefficients +# H(1,1) +_6a_1_s1_s1 = 0.00108 * ev2au +_1_9a_s1_s1 = -0.00474 * ev2au +_6a_9a_s1_s1 = 0.00204 * ev2au +# H(2,2) +_6a_1_s2_s2 = -0.00298 * ev2au +_1_9a_s2_s2 = -0.00155 * ev2au +_6a_9a_s2_s2 = 0.00189 * ev2au + +# linear, off-diagonal coupling coefficients +_10a_s1_s2 = 0.20804 * ev2au + +# bilinear, off-diagonal coupling coefficients +# H(1,2) and H(2,1) +_1_10a_s1_s2 = 0.00553 * ev2au +_6a_10a_s1_s2 = 0.01000 * ev2au +_9a_10a_s1_s2 = 0.00126 * ev2au + +e_list = ["s1", "s2"] +v_list = ["10a", "6a", "9a", "1"] + + +def get_ham_terms(): + ham_terms = [] + for e_isymbol, e_jsymbol in product(e_list, repeat=2): + e_op = Op(r"a^\dagger a", [e_isymbol, e_jsymbol]) + for v_isymbol, v_jsymbol in product(v_list, repeat=2): + # linear + if v_isymbol == v_jsymbol: + # if one of the permutation is defined, then the `e_idx_tuple` term should + # be defined as required by Hermitian Hamiltonian + for eterm1, eterm2 in permut([e_isymbol, e_jsymbol], 2): + factor = globals().get(f"_{v_isymbol}_{eterm1}_{eterm2}") + if factor is not None: + factor *= np.sqrt(eval(f"w{v_isymbol}")) + ham_terms.append(e_op * Op("x", v_isymbol) * factor) + logger.debug(f"term: {v_isymbol}_{e_isymbol}_{e_jsymbol}") + break + else: + logger.debug(f"no term: {v_isymbol}_{e_isymbol}_{e_jsymbol}") + + # quadratic + # use product to guarantee `break` breaks the whole loop + it = product(permut([v_isymbol, v_jsymbol], 2), permut([e_isymbol, e_jsymbol], 2)) + for (vterm1, vterm2), (eterm1, eterm2) in it: + factor = globals().get(f"_{vterm1}_{vterm2}_{eterm1}_{eterm2}") + + if factor is not None: + factor *= np.sqrt(eval(f"w{v_isymbol}") * eval(f"w{v_jsymbol}")) + if vterm1 == vterm2: + v_op = Op("x^2", vterm1) + else: + v_op = Op("x", vterm1) * Op("x", vterm2) + ham_terms.append(e_op * v_op * factor) + logger.debug(f"term: {v_isymbol}_{v_jsymbol}_{e_isymbol}_{e_jsymbol}") + break + else: + logger.debug(f"no term: {v_isymbol}_{v_jsymbol}_{e_isymbol}_{e_jsymbol}") + + # electronic coupling + ham_terms.append(Op(r"a^\dagger a", "s1", -delta, [0, 0])) + ham_terms.append(Op(r"a^\dagger a", "s2", delta, [0, 0])) + + # vibrational kinetic and potential + for v_isymbol in v_list: + ham_terms.extend([Op("p^2", v_isymbol, 0.5), Op("x^2", v_isymbol, 0.5 * eval("w" + v_isymbol) ** 2)]) + + ham_terms = [term for term in ham_terms if term.factor != 0] + + return ham_terms + + +def get_basis(nlevels): + if isinstance(nlevels, int): + nlevels = [nlevels] * len(v_list) + if not isinstance(nlevels, list): + raise TypeError(f"`nlevels` must be int or list. Got {type(nlevels)}") + basis = [BasisMultiElectron(e_list, [0, 0])] + for v_isymbol, n in zip(v_list, nlevels): + basis.append(BasisSHO(v_isymbol, globals()[f"w{v_isymbol}"], n)) + return basis diff --git a/src/tyxonq/chem/dynamic/model/sbm.py b/src/tyxonq/chem/dynamic/model/sbm.py new file mode 100644 index 0000000..3acef0d --- /dev/null +++ b/src/tyxonq/chem/dynamic/model/sbm.py @@ -0,0 +1,27 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from renormalizer import BasisHalfSpin, BasisSHO, Op + + +def get_ham_terms(epsilon, delta, nmode, omega_list, g_list): + terms = [Op("sigma_z", "spin", epsilon), Op("sigma_x", "spin", delta)] + for i in range(nmode): + g = g_list[i] + omega = omega_list[i] + terms.extend([Op(r"b^\dagger b", f"v{i}", omega), Op("sigma_z", "spin", g) * Op(r"b^\dagger+b", f"v{i}")]) + return terms + + +def get_basis(omega_list, nlevels): + if isinstance(nlevels, int): + nlevels = [nlevels] * len(omega_list) + if not isinstance(nlevels, list): + raise TypeError(f"`nlevels` must be int or list. Got {type(nlevels)}") + basis = [BasisHalfSpin("spin")] + for i in range(len(omega_list)): + basis.append(BasisSHO(f"v{i}", omega=omega_list[i], nbas=nlevels[i])) + return basis diff --git a/src/tyxonq/chem/dynamic/time_derivative.py b/src/tyxonq/chem/dynamic/time_derivative.py new file mode 100644 index 0000000..98f0460 --- /dev/null +++ b/src/tyxonq/chem/dynamic/time_derivative.py @@ -0,0 +1,177 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +import logging +from typing import List + +import numpy as np +import scipy +from renormalizer.model.basis import BasisSet +from renormalizer import Op +import tyxonq as tq + +from ..utils.backend import jit +from ..utils.circuit import evolve_pauli + + +logger = logging.getLogger(__name__) + + +def construct_ansatz_op(ham_terms: List[Op], spin_basis: List[BasisSet]): + dof_idx_dict = {b.dof: i for i, b in enumerate(spin_basis)} + ansatz_op_list = [] + + for i, op in enumerate(ham_terms): + logger.info(f"Ansatz operator: {i}, {op}") + + op_mat = 1 + name = "" + for symbol in op.split_symbol: + if symbol in ["X", "x", "sigma_x"]: + op_mat = np.kron(op_mat, tq.gates._x_matrix) + name += "X" + elif symbol in ["Y", "y", "sigma_y"]: + op_mat = np.kron(op_mat, tq.gates._y_matrix) + name += "Y" + else: + if symbol not in ["Z", "z", "sigma_z"]: + raise ValueError(f"Hamiltonian must be sum of Pauli strings, got term {op}") + op_mat = np.kron(op_mat, tq.gates._z_matrix) + name += "Z" + qubit_idx_list = [dof_idx_dict[dof] for dof in op.dofs] + ansatz_op_list.append((op_mat, op.factor, name, qubit_idx_list)) + + return ansatz_op_list + + +def get_circuit(ansatz_terms, spin_basis, n_layers, init_state, params, param_ids=None, compile_evolution=False): + if param_ids is None: + param_ids = list(range(len(ansatz_terms))) + + params = tq.backend.reshape(params, [n_layers, max(param_ids) + 1]) + + ansatz_terms_grouped = [] + for term in ansatz_terms: + if isinstance(term, Op): + ansatz_terms_grouped.append([term]) + else: + ansatz_terms_grouped.append(term) + assert isinstance(ansatz_terms_grouped[0], list) and isinstance(ansatz_terms_grouped[0][0], Op) + + ansatz_op_list_grouped = [] + for ansatz_terms in ansatz_terms_grouped: + ansatz_op_list = construct_ansatz_op(ansatz_terms, spin_basis) + factors = np.array([factor for _, factor, _, _ in ansatz_op_list]) + assert np.allclose(factors.real, 0) or np.allclose(factors.imag, 0) + ansatz_op_list_grouped.append(ansatz_op_list) + + if isinstance(init_state, tq.Circuit): + c = tq.Circuit.from_qir(init_state.to_qir(), circuit_params=init_state.circuit_param) + else: + c = tq.Circuit(len(spin_basis), inputs=init_state) + + for i in range(0, n_layers): + for j, ansatz_op_list in enumerate(ansatz_op_list_grouped): + param_id = param_ids[j] + theta = params[i, param_id] + for ansatz_op, coeff, name, qubit_idx_list in ansatz_op_list: + if coeff.real == 0: + coeff = coeff.imag + if not compile_evolution: + np.testing.assert_allclose(ansatz_op.conj().T @ ansatz_op, np.eye(len(ansatz_op))) + name = f"exp(-iθ{name})" + c.exp1(*qubit_idx_list, unitary=ansatz_op, theta=coeff * theta, name=name) + else: + pauli_string = tuple(zip(qubit_idx_list, name)) + c = evolve_pauli(c, pauli_string, theta=2 * coeff * theta) + return c + + +def one_trotter_step(ham_terms, spin_basis, init_state, dt, inplace=False): + """ + one step first order trotter decompostion + """ + ansatz_op_list = construct_ansatz_op(ham_terms, spin_basis) + + if isinstance(init_state, tq.Circuit): + if inplace: + c = init_state + else: + c = tq.Circuit.from_qir(init_state.to_qir(), circuit_params=init_state.circuit_param) + else: + c = tq.Circuit(len(spin_basis), inputs=init_state) + + for ansatz_op, op_factor, name, qubit_idx_list in ansatz_op_list: + c.exp1(*qubit_idx_list, unitary=ansatz_op, theta=dt * op_factor, name=name) + return c + + +def get_ansatz(ham_terms, spin_basis, n_layers, init_state, param_ids=None): + @jit + def ansatz(theta): + c = get_circuit(ham_terms, spin_basis, n_layers, init_state, theta, param_ids) + return c.state() + + return ansatz + + +def get_jacobian_func(ansatz): + return jit(tq.backend.jacfwd(ansatz, argnums=0)) + + +def regularized_inversion(m, eps): + evals, evecs = scipy.linalg.eigh(m) + evals += eps * np.exp(-evals / eps) + new_evals = 1 / evals + return evecs @ np.diag(new_evals) @ evecs.T + + +def regularized_inversion2(m, eps): + evals, evecs = scipy.linalg.eigh(m) + mask = np.abs(evals) > 1e-10 + evals = evals[mask] + assert np.all(evals > 0) + new_evals = np.zeros(len(evecs)) + new_evals[mask] = 1 / evals + return evecs @ np.diag(new_evals) @ evecs.T + + +def get_deriv(ansatz, jacobian_func, params, hamiltonian: np.ndarray, eps: float = 1e-5, include_phase: bool = False): + params = tq.array_to_tensor(params) + + jacobian = tq.backend.numpy(jacobian_func(params)).astype(np.complex128) + lhs = jacobian.conj().T @ jacobian + + psi = tq.backend.numpy(ansatz(params)).astype(np.complex128) + hpsi = hamiltonian @ psi + rhs = jacobian.conj().T @ hpsi + + lhs = lhs.real + rhs = rhs.imag + + # global phase term in https://arxiv.org/pdf/1812.08767.pdf + if include_phase: + ovlp = jacobian.conj().T @ psi + lhs += (ovlp.reshape(-1, 1) * ovlp.reshape(1, -1)).real + e = psi.conj() @ hpsi + rhs -= (e * ovlp).imag + + lhs_inverse = regularized_inversion(lhs, eps) + theta_deriv = lhs_inverse @ rhs + np.testing.assert_allclose(lhs @ lhs_inverse @ rhs, rhs, atol=2e2 * eps) + np.testing.assert_allclose(theta_deriv.imag, 0) + return theta_deriv.real.astype(np.float64) + + +def get_pvqd_loss_func(ansatz): + def loss(delta_params, params, hamiltonian: np.ndarray, delta_t: float): + ket = ansatz(params + delta_params) + bra = ansatz(params) + evolution = tq.backend.convert_to_tensor(tq.backend.expm(1j * delta_t * hamiltonian)) + return 1 - tq.backend.norm(bra.conj() @ (evolution @ ket)) ** 2 + + loss = tq.interfaces.scipy.scipy_optimize_interface(loss) + return loss diff --git a/src/tyxonq/chem/dynamic/time_evolution.py b/src/tyxonq/chem/dynamic/time_evolution.py new file mode 100644 index 0000000..df7edca --- /dev/null +++ b/src/tyxonq/chem/dynamic/time_evolution.py @@ -0,0 +1,247 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +import time +import logging +from functools import partial +from collections import defaultdict +from typing import Dict, List, Optional + +import numpy as np +import scipy +from scipy.integrate import solve_ivp +from scipy.optimize import minimize +import tyxonq as tq +from renormalizer import Model, Mps, Mpo, BasisHalfSpin +from renormalizer.model.basis import BasisSet + +from ..dynamic.transform import qubit_encode_op, qubit_encode_basis, get_init_circuit +from ..dynamic.time_derivative import ( + get_circuit, + get_ansatz, + get_jacobian_func, + get_deriv, + get_pvqd_loss_func, + one_trotter_step, +) + +logger = logging.getLogger(__name__) + + +class Ivp_Config: + def __init__(self, method="RK45", rtol=1e-3, atol=1e-6): + self.method = method + self.rtol = rtol + self.atol = atol + + +def evolve_exact(evals: np.ndarray, evecs: np.ndarray, init: np.ndarray, t: float): + return evecs @ (np.diag(np.exp(-1j * t * evals)) @ (evecs.T @ init)) + + +class TimeEvolution: + def __init__( + self, + ham_terms, + basis: List[BasisSet], + boson_encoding: str = "gray", + init_condition: Optional[Dict] = None, + n_layers: int = 3, + eps: float = 1e-5, + property_op_dict: Dict = None, + ref_only: bool = False, + ivp_config: Ivp_Config = None, + ): + # handling defaults + if init_condition is None: + init_condition = {} + if property_op_dict is None: + property_op_dict = {} + + # setup refs + self.model_ref = Model(basis, ham_terms) + self.h_mpo_ref = Mpo(self.model_ref) + self.h_ref = self.h_mpo_ref.todense() + self.evals_ref, self.evecs_ref = scipy.linalg.eigh(self.h_ref) + self.init_ref = Mps.hartree_product_state(self.model_ref, init_condition).todense() + # only do ref in kernel. Could be useful for debugging + self.ref_only = ref_only + + # setup transformed Hamiltonian + ham_terms_spin, self.constant = qubit_encode_op(ham_terms, basis, boson_encoding) + basis_spin: List[BasisHalfSpin] = qubit_encode_basis(basis, boson_encoding) + # help perform some sanity checks + self.model = Model(basis_spin, ham_terms_spin) + self.h_mpo = Mpo(self.model) + self.h = self.h_mpo.todense() + + # setup the ansatz + self.init_circuit = get_init_circuit(self.model_ref, self.model, boson_encoding, init_condition) + self.n_layers = n_layers + self.n_params = self.n_layers * len(self.model.ham_terms) + self.ansatz = get_ansatz(self.model.ham_terms, self.model.basis, self.n_layers, self.init_circuit) + self.jacobian_func = get_jacobian_func(self.ansatz) + + # setup runtime components + self.current_circuit = self.init_circuit + self.eps = eps + self.include_phase = False + if ivp_config is None: + self.ivp_config = Ivp_Config() + else: + self.ivp_config = ivp_config + + def scipy_deriv(t, _params): + return get_deriv(self.ansatz, self.jacobian_func, _params, self.h, self.eps, self.include_phase) + + self.scipy_deriv = scipy_deriv + + self.pvqd_loss = get_pvqd_loss_func(self.ansatz) + + def solve_pvqd(_params, delta_t): + hamiltonian = tq.backend.convert_to_tensor(self.h) + loss = partial(self.pvqd_loss, params=_params, hamiltonian=hamiltonian, delta_t=delta_t) + opt_res = minimize(loss, np.zeros_like(_params), jac=True, method="L-BFGS-B") + return opt_res + + self.solve_pvqd = solve_pvqd + + self.property_op_dict = property_op_dict + self.property_mat_dict = {} + for k, op in self.property_op_dict.items(): + transformed_op, constant = qubit_encode_op(op, basis, boson_encoding) + mat = Mpo(self.model, transformed_op).todense() + mat += constant * np.eye(len(mat)) + mat_ref = Mpo(self.model_ref, op).todense() + self.property_mat_dict[k] = mat, mat_ref + + # time evolution result + self.t_list = [0] + self.params_list = [np.zeros(self.n_params, dtype=np.float64)] + state = self.init_circuit.state() + self.state_list = [state] + state_ref = self.init_ref + self.state_ref_list = [state_ref] + self.scipy_sol_list = [] + self._property_dict = defaultdict(list) + self.update_property_dict(state, state_ref) + + self.wall_time_list = [] + + def kernel(self, tau, algo="vqd"): + # one step of time evolution + if self.ref_only: + return self.kernel_ref_only(tau) + time0 = time.time() + if algo == "vqd" or algo == "pvqd": + if algo == "vqd": + method, rtol, atol = self.ivp_config.method, self.ivp_config.rtol, self.ivp_config.atol + scipy_sol = solve_ivp( + self.scipy_deriv, [self.t, self.t + tau], self.params, method=method, rtol=rtol, atol=atol + ) + new_params = scipy_sol.y[:, -1] + else: + scipy_sol = self.solve_pvqd(self.params, tau) + new_params = self.params + scipy_sol.x + + self.params_list.append(new_params) + state = self.ansatz(self.params) + self.scipy_sol_list.append(scipy_sol) + else: + assert algo == "trotter" + self.current_circuit = one_trotter_step( + self.model.ham_terms, self.model.basis, self.current_circuit, tau, inplace=True + ) + shortcut = one_trotter_step(self.model.ham_terms, self.model.basis, self.state, tau) + state = shortcut.state() + + time1 = time.time() + self.t_list.append(self.t + tau) + # t and params already updated + self.state_list.append(state) + state_ref = evolve_exact(self.evals_ref, self.evecs_ref, self.init_ref, self.t) + self.state_ref_list.append(state_ref) + # calculate properties + self.update_property_dict(state, state_ref) + + self.wall_time_list.append(time1 - time0) + + def kernel_ref_only(self, tau): + # Let's do code duplication to keep the source code simple + self.t_list.append(self.t + tau) + state_ref = evolve_exact(self.evals_ref, self.evecs_ref, self.init_ref, self.t) + self.state_ref_list.append(state_ref) + # calculate properties + self.update_property_dict(None, state_ref) + + def update_property_dict(self, state, state_ref): + for k, (mat, mat_ref) in self.property_mat_dict.items(): + if not self.ref_only: + res = state.T.conj() @ (mat @ state), state_ref.T.conj() @ (mat_ref @ state_ref) + else: + res = state_ref.T.conj() @ (mat_ref @ state_ref) + self._property_dict[k].append(res) + + def get_circuit(self, params, param_ids=None, compile_evolution=False): + return get_circuit( + self.ham_terms_spin, + self.basis_spin, + self.n_layers, + self.init_circuit, + params, + param_ids, + compile_evolution=compile_evolution, + ) + + @property + def ham_terms_spin(self): + return self.model.ham_terms + + @property + def basis_spin(self): + return self.model.basis + + @property + def t(self): + return self.t_list[-1] + + @property + def t_array(self): + return np.array(self.t_list) + + @property + def params(self): + return self.params_list[-1] + + @property + def params_array(self): + return np.array(self.params_list) + + @property + def state(self): + return self.state_list[-1] + + @property + def state_array(self): + return np.array(self.state_list) + + @property + def state_ref(self): + return self.state_ref_list[-1] + + @property + def state_ref_array(self): + return np.array(self.state_ref_list) + + @property + def property_dict(self): + return {k: np.array(v) for k, v in self._property_dict.items()} + + properties = property_dict + + @property + def wall_time(self): + return self.wall_time_list[-1] diff --git a/src/tyxonq/chem/dynamic/transform.py b/src/tyxonq/chem/dynamic/transform.py new file mode 100644 index 0000000..fb861de --- /dev/null +++ b/src/tyxonq/chem/dynamic/transform.py @@ -0,0 +1,339 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from typing import Any, List, Union, Tuple +import numpy as np + +from renormalizer.model import Op, OpSum, Model +from renormalizer import BasisHalfSpin, BasisSimpleElectron, BasisMultiElectron, BasisMultiElectronVac, Mps +from renormalizer.model.basis import BasisSet +import tyxonq as tq + +from ..constants import DISCARD_EPS + + +# to be added to new DOFs +DOF_SALT = "TCCQUBIT" + + +def check_basis_type(basis: List[BasisSet]): + for b in basis: + if isinstance(b, (BasisMultiElectronVac,)): + raise TypeError(f"Unsupported basis: {type(b)}") + if isinstance(b, BasisMultiElectron) and len(b.dofs) != 2: + raise ValueError(f"For only two DOFs are allowed in BasisMultiElectron. Got {b}") + + +def qubit_encode_op( + terms: Union[List[Op], Op], basis: List[BasisSet], boson_encoding: str = None +) -> Tuple[OpSum, float]: + check_basis_type(basis) + if isinstance(terms, Op): + terms = [terms] + + model = Model(basis, []) + + new_terms = [] + for op in terms: + terms, factor = op.split_elementary(model.dof_to_siteidx) + opsum_list = [] + for op in terms: + opsum = transform_op(op, model.dof_to_basis[op.dofs[0]], boson_encoding) + opsum_list.append(opsum) + + new_term = 1 + for opsum in opsum_list: + new_term = new_term * opsum + new_term = new_term * factor + + new_terms.extend(new_term) + + # post process + # pick out constants + identity_terms: List[Op] = [] + non_identity_terms = OpSum() + for op in new_terms: + if op.is_identity: + identity_terms.append(op) + else: + non_identity_terms.append(op.squeeze_identity()) + + constant = sum([op.factor for op in identity_terms]) + + return non_identity_terms.simplify(atol=DISCARD_EPS), constant + + +def qubit_encode_op_grouped( + terms: List[Union[List[Op], Op]], basis: List[BasisSet], boson_encoding: str = None +) -> Tuple[List[OpSum], float]: + new_terms = [] + constant_sum = 0 + for op in terms: + opsum, constant = qubit_encode_op(op, basis, boson_encoding) + new_terms.append(opsum) + constant_sum += constant + + return new_terms, constant_sum + + +def qubit_encode_basis(basis: List[BasisSet], boson_encoding=None): + spin_basis = [] + for b in basis: + if isinstance(b, BasisMultiElectron): + assert b.nbas == 2 + spin_basis.append(BasisHalfSpin(b.dofs)) + elif b.is_phonon: + if boson_encoding is None: + new_dofs = [b.dof] + elif boson_encoding == "unary": + new_dofs = [(b.dof, f"{DOF_SALT}-{i}") for i in range(b.nbas)] + else: + assert boson_encoding.lower() in ["binary", "gray"] + n_qubits = int(np.ceil(np.log2(b.nbas))) + new_dofs = [(b.dof, f"{DOF_SALT}-{i}") for i in range(n_qubits)] + new_basis = [BasisHalfSpin(dof) for dof in new_dofs] + spin_basis.extend(new_basis) + else: + spin_basis.append(BasisHalfSpin(b.dof)) + + return spin_basis + + +def transform_op(op: Op, basis: BasisSet, boson_encoding: str = None) -> OpSum: + # `op` is "elementary": all terms are in the `basis` + assert op.factor == 1 + + if set(op.split_symbol) == {"I"}: + return OpSum([op]) + + if isinstance(basis, (BasisHalfSpin, BasisSimpleElectron, BasisMultiElectron)): + if isinstance(basis, BasisMultiElectron): + assert len(basis.dof) == 2 + new_dof = basis.dofs + else: + new_dof = op.dofs[0] + return transform_op_direct(op, new_dof, basis) + + # in principle can encode basis sets such as BasisMultiElectron, + # but I think it's unnatural and not necessary + assert basis.is_phonon + return transform_op_boson(op, basis, boson_encoding) + + +def get_elem_qubit_op_direct(row_idx: int, col_idx: int, dof: Any): + if (row_idx, col_idx) == (0, 0): + return 1 / 2 * (Op("I", dof) + Op("Z", dof)) + elif (row_idx, col_idx) == (0, 1): + return 1 / 2 * (Op("X", dof) + 1j * Op("Y", dof)) + elif (row_idx, col_idx) == (1, 0): + return 1 / 2 * (Op("X", dof) - 1j * Op("Y", dof)) + else: + assert (row_idx, col_idx) == (1, 1) + return 1 / 2 * (Op("I", dof) - Op("Z", dof)) + + +def transform_op_direct(op: Op, dof: Any, basis: BasisSet): + if basis.nbas != 2: + raise ValueError("Direct encoding only support two level basis") + mat = basis.op_mat(op) + ret = OpSum() + for row_idx, col_idx in zip(*np.nonzero(mat)): + ret += mat[row_idx, col_idx] * get_elem_qubit_op_direct(row_idx, col_idx, dof) + return ret.simplify(atol=DISCARD_EPS) + + +def get_elem_qubit_op_unary(row_idx: int, col_idx: int, new_dofs: List[Any]): + if row_idx == col_idx: + dof_list = [new_dofs[row_idx]] + return 1 / 2 * (Op("I", dof_list) - Op("Z", dof_list)) + else: + des = 1 / 2 * (Op("X", new_dofs[col_idx]) + 1j * Op("Y", new_dofs[col_idx])) + cre = 1 / 2 * (Op("X", new_dofs[row_idx]) - 1j * Op("Y", new_dofs[row_idx])) + # exploit the commutation property for simplification + if new_dofs[row_idx] < new_dofs[col_idx]: + return cre * des + else: + return des * cre + + +def transform_op_boson_unary(op: Op, dof: Any, basis: BasisSet): + new_dofs = [(dof, f"{DOF_SALT}-{i}") for i in range(basis.nbas - 1, -1, -1)] + mat = basis.op_mat(op) + ret = OpSum() + for row_idx, col_idx in zip(*np.nonzero(mat)): + ret += mat[row_idx, col_idx] * get_elem_qubit_op_unary(row_idx, col_idx, new_dofs) + return ret.simplify(atol=DISCARD_EPS) + + +def get_elem_qubit_op_binary(row_idx: int, col_idx: int, new_dofs: List[Any], code_strs: List[str]): + # gray_code_ints = [int(s, base=2) for s in gray_code_strs] + n_qubits = len(new_dofs) + if row_idx == col_idx: + op_list = [] + for i in range(n_qubits): + dof = new_dofs[i] + if code_strs[row_idx][i] == "0": + new_op = 1 / 2 * (Op("I", dof) + Op("Z", dof)) + else: + new_op = 1 / 2 * (Op("I", dof) - Op("Z", dof)) + op_list.append(new_op) + return OpSum.product(op_list) + else: + # |code1> np.iinfo(uint_type).max: + raise ValueError(f"Too many qubits: {n_qubits}, try using complex128 datatype") + na, nb = unpack_nelec(n_elec_s) + if mode in ["fermion", "qubit"]: + beta = cistring.make_strings(range(n_qubits // 2), nb) + beta = xp.array(beta, dtype=uint_type) + if na == nb: + alpha = beta + else: + alpha = cistring.make_strings(range(n_qubits // 2), na) + alpha = xp.array(alpha, dtype=uint_type) + ci_strings = ((alpha << (n_qubits // 2)).reshape(-1, 1) + beta.reshape(1, -1)).ravel() + if strs2addr: + if na == nb: + strs2addr = xp.zeros(2 ** (n_qubits // 2), dtype=uint_type) + strs2addr[beta] = xp.arange(len(beta)) + else: + strs2addr = xp.zeros((2, 2 ** (n_qubits // 2)), dtype=uint_type) + strs2addr[0][alpha] = xp.arange(len(alpha)) + strs2addr[1][beta] = xp.arange(len(beta)) + return ci_strings, strs2addr + else: + assert mode == "hcb" + assert na == nb + ci_strings = cistring.make_strings(range(n_qubits), na).astype(uint_type) + if strs2addr: + strs2addr = xp.zeros(2**n_qubits, dtype=uint_type) + strs2addr[ci_strings] = xp.arange(len(ci_strings)) + return ci_strings, strs2addr + + return ci_strings + + +def get_addr(excitation, n_qubits, n_elec_s, strs2addr, mode, num_strings=None): + if mode == "hcb": + return strs2addr[excitation] + assert mode in ["fermion", "qubit"] + alpha = excitation >> (n_qubits // 2) + beta = excitation & (2 ** (n_qubits // 2) - 1) + na, nb = n_elec_s + if na == nb: + alpha_addr = strs2addr[alpha] + beta_addr = strs2addr[beta] + else: + alpha_addr = strs2addr[0][alpha] + beta_addr = strs2addr[1][beta] + if num_strings is None: + num_strings = cistring.num_strings(n_qubits // 2, nb) + return alpha_addr * num_strings + beta_addr + + +def get_ex_bitstring(n_qubits, n_elec_s, ex_op, mode): + na, nb = n_elec_s + if mode in ["fermion", "qubit"]: + bitstring_basea = ["0"] * (n_qubits // 2 - na) + ["1"] * na + bitstring_baseb = ["0"] * (n_qubits // 2 - nb) + ["1"] * nb + bitstring_base = bitstring_basea + bitstring_baseb + else: + assert mode == "hcb" + assert na == nb + bitstring_base = ["0"] * (n_qubits - na) + ["1"] * na + + bitstring = bitstring_base.copy()[::-1] + # first annihilation then creation + if len(ex_op) == 2: + bitstring[ex_op[1]] = "0" + bitstring[ex_op[0]] = "1" + else: + assert len(ex_op) == 4 + bitstring[ex_op[3]] = "0" + bitstring[ex_op[2]] = "0" + bitstring[ex_op[1]] = "1" + bitstring[ex_op[0]] = "1" + + return "".join(reversed(bitstring)) + + +def civector_to_statevector(civector, n_qubits, ci_strings): + statevector = tq.backend.zeros(2**n_qubits, dtype=tq.rdtypestr) + return tensor_set_elem(statevector, ci_strings, civector) + + +def statevector_to_civector(statevector, ci_strings): + return statevector[ci_strings] + + +@partial(jit, static_argnums=[0]) +def get_init_civector(len_ci): + civector = tq.backend.zeros(len_ci, dtype=tq.rdtypestr) + civector = tensor_set_elem(civector, 0, 1) + return civector diff --git a/src/tyxonq/chem/static/engine_hea.py b/src/tyxonq/chem/static/engine_hea.py new file mode 100644 index 0000000..72fd3fd --- /dev/null +++ b/src/tyxonq/chem/static/engine_hea.py @@ -0,0 +1,169 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from functools import partial + +import numpy as np +import tyxonq as tq +from tyxonq import Circuit, DMCircuit +from tyxonq.noisemodel import circuit_with_noise +from tyxonq.experimental import parameter_shift_grad +from tyxonq.cloud.wrapper import batch_expectation_ps + +from ..utils.backend import jit + + +class QpuConf: + def __init__(self, device=None, provider=None, initial_mapping=None): + if device is None: + device = "tianji_s2" + self.device = device + self.privider = provider + self.initial_mapping = initial_mapping + + +@partial(jit, static_argnums=[1]) +def get_statevector(params, get_circuit): + return get_circuit(params).state() + + +@partial(jit, static_argnums=[1, 2]) +def get_densitymatrix(params, get_dmcircuit, noise_conf): + if noise_conf is None: + raise ValueError("NoiseConf not provided") + circuit = get_dmcircuit(params) + circuit = circuit_with_noise(circuit, noise_conf) + return circuit.densitymatrix() + + +@partial(jit, static_argnums=[2]) +def get_energy_tensornetwork(params, h_array, get_circuit): + s = get_statevector(params, get_circuit) + return (s.conj() @ (h_array @ s)).real + + +@partial(jit, static_argnums=[2, 3]) +def get_energy_tensornetwork_noise(params, h_array, get_dmcircuit, noise_conf): + dm = get_densitymatrix(params, get_dmcircuit, noise_conf) + return tq.backend.trace(dm @ h_array).real + + +def sample_expectation_pauli(c, paulis, coeffs, shots, noise_conf): + expectations = [] + for pauli, coef in zip(paulis, coeffs): + x = [] + y = [] + z = [] + m = {"X": x, "Y": y, "Z": z} + for idx, symbol in pauli: + m[symbol].append(idx) + expectation = coef + if len(x + y + z) != 0: + expectation *= c.sample_expectation_ps(x=x, y=y, z=z, shots=shots, noise_conf=noise_conf) + expectations.append(expectation) + # already real + return sum(expectations) + + +@partial(jit, static_argnums=[1, 3, 4]) +def get_energy_tensornetwork_shot(params, paulis, coeffs, get_circuit, shots): + c = get_circuit(params) + return sample_expectation_pauli(c, paulis, coeffs, shots, None) + + +@partial(jit, static_argnums=[1, 3, 4, 5]) +def get_energy_tensornetwork_noise_shot(params, paulis, coeffs, get_dmcircuit, noise_conf, shots: int): + dm = get_densitymatrix(params, get_dmcircuit, noise_conf) + n_qubits = round(np.log2(dm.shape[0])) + c = DMCircuit(n_qubits, dminputs=dm) + return sample_expectation_pauli(c, paulis, coeffs, shots, noise_conf) + + +def get_energy_qpu(params, paulis, coeffs, get_circuit, qpu_conf: QpuConf, shots: int): + c: Circuit = get_circuit(params) + pss = [] + symbol_mapping = {"X": 1, "Y": 2, "Z": 3} + ps_identity_idx = [] + for i, pauli in enumerate(paulis): + ps = [0] * c.circuit_param["nqubits"] + if len(pauli) == 0: + ps_identity_idx.append(i) + continue + for idx, symbol in pauli: + ps[idx] = symbol_mapping[symbol] + pss.append(ps) + assert len(pss) + len(ps_identity_idx) == len(paulis) + coeffs_non_identity = [coeffs[i] for i in range(len(coeffs)) if i not in ps_identity_idx] + assert len(pss) == len(coeffs_non_identity) + es = [] + for _ in range((shots - 1) // 8192 + 1): + e = batch_expectation_ps(c, pss, device=qpu_conf.device, ws=coeffs_non_identity, shots=8192) + es.append(e) + print(paulis) + print(coeffs) + print(es) + return np.mean(es) + sum([coeffs[i] for i in ps_identity_idx]) + + +def _get_energy_and_grad(partial_get_energy, params, grad): + if grad == "param-shift": + e = partial_get_energy(params) + grad_f = parameter_shift_grad(partial_get_energy) + g = grad_f(params) + else: + assert grad == "autodiff" + e, g = tq.backend.value_and_grad(partial_get_energy)(params) + return e, g + + +@partial(jit, static_argnums=[2, 3]) +def get_energy_and_grad_tensornetwork(params, h_array, get_circuit, grad): + partial_get_energy = partial(get_energy_tensornetwork, h_array=h_array, get_circuit=get_circuit) + return _get_energy_and_grad(partial_get_energy, params, grad) + + +@partial(jit, static_argnums=[2, 3, 4]) +def get_energy_and_grad_tensornetwork_noise(params, h_array, get_dmcircuit, noise_conf, grad): + partial_get_energy = partial( + get_energy_tensornetwork_noise, h_array=h_array, get_dmcircuit=get_dmcircuit, noise_conf=noise_conf + ) + return _get_energy_and_grad(partial_get_energy, params, grad) + + +@partial(jit, static_argnums=[1, 3, 4, 5]) +def get_energy_and_grad_tensornetwork_shot(params, paulis, coeffs, get_circuit, shots, grad): + partial_get_energy = partial( + get_energy_tensornetwork_shot, + paulis=paulis, + coeffs=coeffs, + get_circuit=get_circuit, + shots=shots, + ) + return _get_energy_and_grad(partial_get_energy, params, grad) + + +@partial(jit, static_argnums=[1, 3, 4, 5, 6]) +def get_energy_and_grad_tensornetwork_noise_shot(params, paulis, coeffs, get_dmcircuit, noise_conf, shots, grad): + partial_get_energy = partial( + get_energy_tensornetwork_noise_shot, + paulis=paulis, + coeffs=coeffs, + get_dmcircuit=get_dmcircuit, + noise_conf=noise_conf, + shots=shots, + ) + return _get_energy_and_grad(partial_get_energy, params, grad) + + +def get_energy_and_grad_qpu(params, paulis, coeffs, get_circuit, shots: int, grad): + partial_get_energy = partial( + get_energy_qpu, + paulis=paulis, + coeffs=coeffs, + get_circuit=get_circuit, + shots=shots, + ) + return _get_energy_and_grad(partial_get_energy, params, grad) diff --git a/src/tyxonq/chem/static/engine_ucc.py b/src/tyxonq/chem/static/engine_ucc.py new file mode 100644 index 0000000..710cf7b --- /dev/null +++ b/src/tyxonq/chem/static/engine_ucc.py @@ -0,0 +1,177 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from functools import partial +from typing import Tuple +import logging + +import tyxonq as tq + +from ..utils.backend import jit, value_and_grad +from ..utils.misc import unpack_nelec +from ..static.hamiltonian import apply_op +from ..static.ci_utils import get_ci_strings, civector_to_statevector, statevector_to_civector +from ..static.evolve_civector import ( + get_civector_nocache, + get_energy_and_grad_civector, + get_energy_and_grad_civector_nocache, + apply_excitation_civector, + apply_excitation_civector_nocache, +) +from ..static.evolve_civector import get_civector as get_civector_ +from ..static.evolve_statevector import apply_excitation_statevector +from ..static.evolve_statevector import get_statevector as get_statevector_ +from ..static.evolve_tensornetwork import get_statevector_tensornetwork +from ..static.evolve_pyscf import get_civector_pyscf, get_energy_and_grad_pyscf, apply_excitation_pyscf + + +logger = logging.getLogger(__name__) + + +GETVECTOR_MAP = { + "tensornetwork": get_statevector_tensornetwork, + "statevector": get_statevector_, + "civector": get_civector_, + "civector-large": get_civector_nocache, + "pyscf": get_civector_pyscf, +} + + +def get_ket(params, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state, ci_strings, engine): + if param_ids is None: + param_ids = range(len(ex_ops)) + func = GETVECTOR_MAP[engine] + init_state = translate_init_state(init_state, n_qubits, ci_strings) + ket = func( + params, + n_qubits, + n_elec_s, + ex_ops=tuple(ex_ops), + param_ids=tuple(param_ids), + mode=mode, + init_state=init_state, + ) + return ket + + +def get_civector(params, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state, engine): + ci_strings = get_ci_strings(n_qubits, n_elec_s, mode) + ket = get_ket(params, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state, ci_strings, engine) + if engine.startswith("civector") or engine == "pyscf": + civector = ket + else: + civector = ket[ci_strings] + return civector + + +def get_statevector(params, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state, engine): + ci_strings = get_ci_strings(n_qubits, n_elec_s, mode) + ket = get_ket(params, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state, ci_strings, engine) + if engine.startswith("civector") or engine == "pyscf": + statevector = civector_to_statevector(ket, n_qubits, ci_strings) + else: + statevector = ket + return statevector + + +def get_energy(params, hamiltonian, n_qubits, n_elec_s, ex_ops: Tuple, param_ids: list, mode: str, init_state, engine): + if param_ids is None: + param_ids = range(len(ex_ops)) + logger.info(f"Entering `get_energy`") + ci_strings = get_ci_strings(n_qubits, n_elec_s, mode) + init_state = translate_init_state(init_state, n_qubits, ci_strings) + ket = GETVECTOR_MAP[engine]( + params, n_qubits, n_elec_s, tuple(ex_ops), tuple(param_ids), mode=mode, init_state=init_state + ) + hket = apply_op(hamiltonian, ket) + return ket @ hket + + +get_energy_statevector = partial(get_energy, engine="statevector") +try: + get_energy_and_grad_statevector = jit(value_and_grad(get_energy_statevector), static_argnums=[2, 3, 4, 5, 6]) +except NotImplementedError: + + def get_energy_and_grad_statevector(*args, **kwargs): + raise NotImplementedError("Non JAX-backend for statevector engine") + + +get_energy_tensornetwork = partial(get_energy, engine="tensornetwork") +try: + get_energy_and_grad_tensornetwork = jit(value_and_grad(get_energy_tensornetwork), static_argnums=[2, 3, 4, 5, 6]) +except NotImplementedError: + + def get_energy_and_grad_tensornetwork(*args, **kwargs): + raise NotImplementedError("Non JAX-backend for tensornetwork engine") + + +ENERGY_AND_GRAD_MAP = { + "tensornetwork": get_energy_and_grad_tensornetwork, + "statevector": get_energy_and_grad_statevector, + "civector": get_energy_and_grad_civector, + "civector-large": get_energy_and_grad_civector_nocache, + "pyscf": get_energy_and_grad_pyscf, +} + + +def get_energy_and_grad(params, hamiltonian, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state, engine): + if engine not in ENERGY_AND_GRAD_MAP: + raise ValueError(f"Engine '{engine}' not supported") + + func = ENERGY_AND_GRAD_MAP[engine] + ci_strings = get_ci_strings(n_qubits, n_elec_s, mode) + init_state = translate_init_state(init_state, n_qubits, ci_strings) + return func(params, hamiltonian, n_qubits, n_elec_s, tuple(ex_ops), tuple(param_ids), mode, init_state) + + +APPLY_EXCITATION_MAP = { + # share the same function with statevector engine + "tensornetwork": apply_excitation_statevector, + "statevector": apply_excitation_statevector, + "civector": apply_excitation_civector, + "civector-large": apply_excitation_civector_nocache, + "pyscf": apply_excitation_pyscf, +} + + +def apply_excitation(state, n_qubits, n_elec_s, ex_op, mode, engine): + if engine not in APPLY_EXCITATION_MAP: + raise ValueError(f"Engine '{engine}' not supported") + + state = tq.backend.convert_to_tensor(state) + + is_statevector_input = len(state) == (1 << n_qubits) + is_statevector_engine = engine in ["tensornetwork", "statevector"] + + n_elec_s = unpack_nelec(n_elec_s) + + if is_statevector_input and not is_statevector_engine: + ci_strings = get_ci_strings(n_qubits, n_elec_s, mode) + state = statevector_to_civector(state, ci_strings) + if not is_statevector_input and is_statevector_engine: + ci_strings = get_ci_strings(n_qubits, n_elec_s, mode) + state = civector_to_statevector(state, n_qubits, ci_strings) + func = APPLY_EXCITATION_MAP[engine] + res_state = func(state, n_qubits, n_elec_s, ex_op, mode) + if is_statevector_input and not is_statevector_engine: + return civector_to_statevector(res_state, n_qubits, ci_strings) + if not is_statevector_input and is_statevector_engine: + return statevector_to_civector(res_state, ci_strings) + return res_state + + +def translate_init_state(init_state, n_qubits, ci_strings): + if init_state is None: + return None + # translate to civector first for all engines to be JAX-compatible + if isinstance(init_state, tq.Circuit): + # note no cupy backend for tq + init_state = statevector_to_civector(tq.backend.convert_to_tensor(init_state.state().real), ci_strings) + else: + is_statevector_input = len(init_state) == (1 << n_qubits) + if is_statevector_input: + init_state = statevector_to_civector(init_state, ci_strings) + return tq.backend.convert_to_tensor(init_state) diff --git a/src/tyxonq/chem/static/evolve_civector.py b/src/tyxonq/chem/static/evolve_civector.py new file mode 100644 index 0000000..6de8173 --- /dev/null +++ b/src/tyxonq/chem/static/evolve_civector.py @@ -0,0 +1,323 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from functools import partial +import logging +from typing import Tuple + +import numpy as np +from openfermion import jordan_wigner +import tyxonq as tq + +from ..utils.backend import jit, fori_loop, scan, get_xp, get_uint_type +from ..utils.misc import ex_op_to_fop +from ..static.hamiltonian import apply_op +from ..static.ci_utils import get_ci_strings, get_addr, get_init_civector + + +logger = logging.getLogger(__name__) + + +def get_fket_permutation(f_idx, n_qubits, n_elec_s, ci_strings, strs2addr, mode): + mask = 0 + for i in f_idx: + mask += 1 << i + excitation = ci_strings ^ mask + return get_addr(excitation, n_qubits, n_elec_s, strs2addr, mode) + + +def get_fket_phase(f_idx, ci_strings): + if len(f_idx) == 2: + mask1 = 1 << f_idx[0] + mask2 = 1 << f_idx[1] + else: + assert len(f_idx) == 4 + mask1 = (1 << f_idx[0]) + (1 << f_idx[1]) + mask2 = (1 << f_idx[2]) + (1 << f_idx[3]) + flip = ci_strings ^ mask1 + mask = mask1 | mask2 + masked = flip & mask + positive = masked == mask + negative = masked == 0 + return positive, negative + + +FERMION_PHASE_MASK_CACHE = {} + + +def get_fermion_phase(f_idx, n_qubits, ci_strings): + if f_idx in FERMION_PHASE_MASK_CACHE: + mask, sign = FERMION_PHASE_MASK_CACHE[(f_idx, n_qubits)] + else: + # fermion operator index, not sorted + fop = ex_op_to_fop(f_idx) + + # pauli string index, already sorted + qop = jordan_wigner(fop) + mask_str = ["0"] * n_qubits + for idx, term in next(iter(qop.terms.keys())): + if term != "Z": + assert idx in f_idx + continue + mask_str[n_qubits - 1 - idx] = "1" + mask = get_uint_type()(int("".join(mask_str), base=2)) + + if sorted(qop.terms.items())[0][1].real > 0: + sign = -1 + else: + sign = 1 + + FERMION_PHASE_MASK_CACHE[(f_idx, n_qubits)] = mask, sign + + parity = ci_strings & mask + assert parity.dtype in [np.uint32, np.uint64] + if parity.dtype == np.uint32: + mask = 0x11111111 + shift = 28 + else: + mask = 0x1111111111111111 + shift = 60 + parity ^= parity >> 1 + parity ^= parity >> 2 + parity = (parity & mask) * mask + parity = (parity >> shift) & 1 + + return sign * np.sign(parity - 0.5).astype(np.int8) + + +def get_operators(n_qubits, n_elec_s, strs2addr, f_idx, ci_strings, mode): + if len(set(f_idx)) != len(f_idx): + raise ValueError(f"Excitation {f_idx} not supported") + xp = get_xp(tq.backend) + fket_permutation = get_fket_permutation(f_idx, n_qubits, n_elec_s, ci_strings, strs2addr, mode) + fket_phase = xp.zeros(len(ci_strings)) + positive, negative = get_fket_phase(f_idx, ci_strings) + fket_phase -= positive + fket_phase += negative + if mode == "fermion": + fket_phase *= get_fermion_phase(f_idx, n_qubits, ci_strings) + f2ket_phase = xp.zeros(len(ci_strings)) + f2ket_phase -= positive + f2ket_phase -= negative + + return fket_permutation, fket_phase, f2ket_phase + + +CI_OPERATOR_BATCH_CACHE = {} +CI_OPERATOR_CACHE = {} + + +@partial(jit, static_argnums=[0, 1, 2, 3]) +def get_operator_tensors(n_qubits, n_elec_s, ex_ops, mode="fermion"): + xp = get_xp(tq.backend) + batch_key = (xp, tq.rdtypestr, n_qubits, n_elec_s, ex_ops, mode) + is_jax_backend = tq.backend.name == "jax" + if not is_jax_backend and batch_key in CI_OPERATOR_BATCH_CACHE: + return CI_OPERATOR_BATCH_CACHE[batch_key] + + ci_strings, strs2addr = get_ci_strings(n_qubits, n_elec_s, mode, strs2addr=True) + + xp = get_xp(tq.backend) + fket_permutation_tensor = xp.zeros((len(ex_ops), len(ci_strings)), dtype=get_uint_type()) + fket_phase_tensor = xp.zeros((len(ex_ops), len(ci_strings)), dtype=np.int8) + f2ket_phase_tensor = xp.zeros((len(ex_ops), len(ci_strings)), dtype=np.int8) + for i, f_idx in enumerate(ex_ops): + if 64 < len(ex_ops): + logger.info((i, f_idx)) + op_key = (xp, tq.rdtypestr, n_qubits, n_elec_s, f_idx, mode) + if not is_jax_backend and op_key in CI_OPERATOR_CACHE: + fket_permutation, fket_phase, f2ket_phase = CI_OPERATOR_CACHE[op_key] + else: + fket_permutation, fket_phase, f2ket_phase = get_operators( + n_qubits, n_elec_s, strs2addr, f_idx, ci_strings, mode + ) + if not is_jax_backend: + CI_OPERATOR_CACHE[op_key] = fket_permutation, fket_phase, f2ket_phase + fket_permutation_tensor[i] = fket_permutation + fket_phase_tensor[i] = fket_phase + f2ket_phase_tensor[i] = f2ket_phase + + fket_permutation_tensor = tq.backend.convert_to_tensor(fket_permutation_tensor) + fket_phase_tensor = tq.backend.convert_to_tensor(fket_phase_tensor) + f2ket_phase_tensor = tq.backend.convert_to_tensor(f2ket_phase_tensor) + + ret = ci_strings, fket_permutation_tensor, fket_phase_tensor, f2ket_phase_tensor + if not is_jax_backend: + CI_OPERATOR_BATCH_CACHE[batch_key] = ret + return ret + + +@partial(jit, static_argnums=[1]) +def get_theta_tensors(params, param_ids): + theta_list = [] + for param_id in param_ids: + theta_list.append(params[param_id]) + + theta_tensor = tq.backend.convert_to_tensor(theta_list) + theta_sin_tensor = tq.backend.sin(theta_tensor) + theta_1mcos_tensor = 1 - tq.backend.cos(theta_tensor) + return theta_sin_tensor, theta_1mcos_tensor + + +@jit +def evolve_civector_by_tensor( + civector, fket_permutation_tensor, fket_phase_tensor, f2ket_phase_tensor, theta_sin, theta_1mcos +): + def _evolve_excitation(j, _civector): + _fket_phase = fket_phase_tensor[j] + _fket_permutation = fket_permutation_tensor[j] + fket = _civector[_fket_permutation] * _fket_phase + f2ket = f2ket_phase_tensor[j] * _civector + _civector += theta_1mcos[j] * f2ket + theta_sin[j] * fket + return _civector + + return fori_loop(0, len(fket_permutation_tensor), _evolve_excitation, civector) + + +@partial(jit, static_argnums=[1, 2, 3, 4, 5]) +def get_civector(params, n_qubits, n_elec_s, ex_ops, param_ids, mode="fermion", init_state=None): + if tq.backend.name == "jax": + logger.info(f"Entering `get_civector`. n_qubit: {n_qubits}") + + ci_strings, fket_permutation_tensor, fket_phase_tensor, f2ket_phase_tensor = get_operator_tensors( + n_qubits, n_elec_s, ex_ops, mode + ) + theta_sin, theta_1mcos = get_theta_tensors(params, param_ids) + + if init_state is None: + civector = get_init_civector(len(ci_strings)) + else: + civector = tq.backend.convert_to_tensor(init_state) + civector = evolve_civector_by_tensor( + civector, fket_permutation_tensor, fket_phase_tensor, f2ket_phase_tensor, theta_sin, theta_1mcos + ) + + return civector.reshape(-1) + + +def get_energy_and_grad_civector( + params, hamiltonian, n_qubits, n_elec_s, ex_ops: Tuple, param_ids: Tuple, mode: str = "fermion", init_state=None +): + ket = get_civector(params, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state) + bra = apply_op(hamiltonian, ket) + energy = bra @ ket + # already cached + op_tensors = get_operator_tensors(n_qubits, n_elec_s, ex_ops, mode=mode) + theta_tensors = get_theta_tensors(params, param_ids) + op_tensors = list(op_tensors) + list(theta_tensors) + gradients_beforesum = _get_gradients_civector(bra, ket, *op_tensors[1:]) + + gradients_beforesum = tq.backend.numpy(gradients_beforesum) + gradients = np.zeros(params.shape) + for grad, param_id in zip(gradients_beforesum, param_ids): + gradients[param_id] += grad + + return energy, 2 * gradients + + +@jit +def _get_gradients_civector( + bra, ket, fket_permutation_tensor, fket_phase_tensor, f2ket_phase_tensor, theta_sin, theta_1mcos +): + scan_xs = fket_permutation_tensor, fket_phase_tensor, f2ket_phase_tensor, -theta_sin, theta_1mcos + + def _evolve_excitation(_civector, _fket_permutation, _fket_phase, _f2ket_phase): + _civector += _f2ket_phase * _civector + _civector[_fket_permutation] * _fket_phase + return _civector + + def get_grad(braket, scan_x): + _bra, _ket = braket + _fket_permutation, _fket_phase, _f2ket_phase, _theta_msin, _theta_1mcos = scan_x + _ket = _evolve_excitation(_ket, _fket_permutation, _fket_phase * _theta_msin, _f2ket_phase * _theta_1mcos) + _bra = _evolve_excitation(_bra, _fket_permutation, _fket_phase * _theta_msin, _f2ket_phase * _theta_1mcos) + _fket = _ket[_fket_permutation] * _fket_phase + grad = _bra @ _fket + + return (_bra, _ket), grad + + _, gradients = scan(get_grad, (bra, ket), scan_xs, len(fket_permutation_tensor), reverse=True) + + return gradients + + +def evolve_excitation_nocache(civector, fket_permutation, fket_phase, f2ket_phase, theta_1mcos, theta_sin): + fket = civector[fket_permutation] * fket_phase + f2ket = civector * f2ket_phase + civector += theta_1mcos * f2ket + theta_sin * fket + return civector + + +@partial(jit, static_argnums=[1, 2, 3, 4, 5]) +def get_civector_nocache(params, n_qubits, n_elec_s, ex_ops, param_ids, mode="fermion", init_state=None): + if tq.backend.name == "jax": + logger.info(f"Entering `get_civector_nocache`. n_qubit: {n_qubits}") + + theta_sin_tensor, theta_1mcos_tensor = get_theta_tensors(params, param_ids) + ci_strings, strs2addr = get_ci_strings(n_qubits, n_elec_s, mode, strs2addr=True) + + if init_state is None: + civector = get_init_civector(len(ci_strings)) + else: + civector = tq.backend.convert_to_tensor(init_state) + + for theta_sin, theta_1mcos, f_idx in zip(theta_sin_tensor, theta_1mcos_tensor, ex_ops): + fket_permutation, fket_phase, f2ket_phase = get_operators( + n_qubits, n_elec_s, strs2addr, f_idx, ci_strings, mode + ) + civector = evolve_excitation_nocache( + civector, fket_permutation, fket_phase, f2ket_phase, theta_1mcos, theta_sin + ) + + return civector.reshape(-1) + + +def get_energy_and_grad_civector_nocache( + params, hamiltonian, n_qubits, n_elec_s, ex_ops: Tuple, param_ids: Tuple, mode: str = "fermion", init_state=None +): + ket = get_civector_nocache(params, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state) + bra = apply_op(hamiltonian, ket) + energy = bra @ ket + + gradients_beforesum = _get_gradients_civector_nocache(bra, ket, params, n_qubits, n_elec_s, ex_ops, param_ids, mode) + gradients_beforesum = tq.backend.numpy(gradients_beforesum) + + gradients = np.zeros(params.shape) + for grad, param_id in zip(gradients_beforesum, param_ids): + gradients[param_id] += grad + + return energy, 2 * gradients + + +@partial(jit, static_argnums=[3, 4, 5, 6, 7]) +def _get_gradients_civector_nocache(bra, ket, params, n_qubits, n_elec_s, ex_ops, param_ids, mode): + ci_strings, strs2addr = get_ci_strings(n_qubits, n_elec_s, mode, True) + theta_sin_tensor, theta_1mcos_tensor = get_theta_tensors(params, param_ids) + + gradients_beforesum = [] + for theta_sin, theta_1mcos, f_idx in reversed(list(zip(theta_sin_tensor, theta_1mcos_tensor, ex_ops))): + fket_permutation, fket_phase, f2ket_phase = get_operators( + n_qubits, n_elec_s, strs2addr, f_idx, ci_strings, mode + ) + bra = evolve_excitation_nocache(bra, fket_permutation, fket_phase, f2ket_phase, theta_1mcos, -theta_sin) + ket = evolve_excitation_nocache(ket, fket_permutation, fket_phase, f2ket_phase, theta_1mcos, -theta_sin) + fket = ket[fket_permutation] * fket_phase + grad = bra @ fket + gradients_beforesum.append(grad) + gradients_beforesum = list(reversed(gradients_beforesum)) + gradients_beforesum = tq.backend.convert_to_tensor(gradients_beforesum) + + return gradients_beforesum + + +def apply_excitation_civector(civector, n_qubits, n_elec_s, f_idx, mode): + _, fket_permutation, fket_phase, _ = get_operator_tensors(n_qubits, n_elec_s, ex_ops=(f_idx,), mode=mode) + return civector[fket_permutation[0]] * fket_phase[0] + + +def apply_excitation_civector_nocache(civector, n_qubits, n_elec_s, f_idx, mode): + ci_strings, strs2addr = get_ci_strings(n_qubits, n_elec_s, mode, strs2addr=True) + fket_permutation, fket_phase, _ = get_operators(n_qubits, n_elec_s, strs2addr, f_idx, ci_strings, mode) + return civector[fket_permutation] * fket_phase diff --git a/src/tyxonq/chem/static/evolve_pyscf.py b/src/tyxonq/chem/static/evolve_pyscf.py new file mode 100644 index 0000000..2e16dc4 --- /dev/null +++ b/src/tyxonq/chem/static/evolve_pyscf.py @@ -0,0 +1,161 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from typing import Tuple, Any +import logging + +import numpy as np +from pyscf.fci import cistring +from pyscf.fci.addons import des_a, cre_a, des_b, cre_b +import tyxonq as tq + +from ..static.hamiltonian import apply_op +from ..static.ci_utils import get_init_civector +from ..utils.misc import unpack_nelec + +Tensor = Any + +logger = logging.getLogger(__name__) + + +class CIvectorPySCF: + def __init__(self, civector, n_orb, n_elec_a, n_elec_b): + assert isinstance(civector, np.ndarray) + self.civector = civector.reshape(-1) + self.n_orb = n_orb + self.n_elec_a = n_elec_a + self.n_elec_b = n_elec_b + + def cre(self, i): + n_elec_a, n_elec_b = self.n_elec_a, self.n_elec_b + if i >= self.n_orb: + cre = cre_a + n_elec_a += 1 + else: + cre = cre_b + n_elec_b += 1 + + new_civector = cre(self.civector, self.n_orb, (self.n_elec_a, self.n_elec_b), i % self.n_orb) + return CIvectorPySCF(new_civector, self.n_orb, n_elec_a, n_elec_b) + + def des(self, i): + n_elec_a, n_elec_b = self.n_elec_a, self.n_elec_b + if i >= self.n_orb: + des = des_a + n_elec_a -= 1 + else: + des = des_b + n_elec_b -= 1 + + new_civector = des(self.civector, self.n_orb, (self.n_elec_a, self.n_elec_b), i % self.n_orb) + return CIvectorPySCF(new_civector, self.n_orb, n_elec_a, n_elec_b) + + def pq(self, p, q): + return self.des(q).cre(p) + + def pqqp(self, p, q): + return self.des(p).cre(q).des(q).cre(p) + + def pqrs(self, p, q, r, s): + return self.des(s).des(r).cre(q).cre(p) + + def pqrssrqp(self, p, q, r, s): + return self.des(p).des(q).cre(r).cre(s).des(s).des(r).cre(q).cre(p) + + +def apply_a2_pyscf(civector: CIvectorPySCF, ex_op) -> Tensor: + if len(ex_op) == 2: + apply_f = civector.pqqp + else: + assert len(ex_op) == 4 + apply_f = civector.pqrssrqp + civector1 = apply_f(*ex_op) + civector2 = apply_f(*reversed(ex_op)) + return -civector1.civector - civector2.civector + + +def apply_a_pyscf(civector: CIvectorPySCF, ex_op) -> Tensor: + if len(ex_op) == 2: + apply_func = civector.pq + else: + assert len(ex_op) == 4 + apply_func = civector.pqrs + civector1 = apply_func(*ex_op) + civector2 = apply_func(*reversed(ex_op)) + return civector1.civector - civector2.civector + + +def evolve_excitation_pyscf(civector: Tensor, ex_op, n_orb, n_elec_s, theta) -> Tensor: + na, nb = unpack_nelec(n_elec_s) + ket = CIvectorPySCF(civector, n_orb, na, nb) + aket = apply_a_pyscf(ket, ex_op) + a2ket = apply_a2_pyscf(ket, ex_op) + return civector + (1 - np.cos(theta)) * a2ket + np.sin(theta) * aket + + +def get_civector_pyscf(params, n_qubits, n_elec_s, ex_ops, param_ids, mode="fermion", init_state=None): + assert mode == "fermion" + n_orb = n_qubits // 2 + na, nb = unpack_nelec(n_elec_s) + num_strings = cistring.num_strings(n_orb, na) * cistring.num_strings(n_orb, nb) + + if init_state is None: + civector = get_init_civector(num_strings) + else: + civector = tq.backend.convert_to_tensor(init_state) + + civector = tq.backend.numpy(civector) + + for ex_op, param_id in zip(ex_ops, param_ids): + theta = params[param_id] + civector = evolve_excitation_pyscf(civector, ex_op, n_orb, n_elec_s, theta) + + return civector.reshape(-1) + + +def get_energy_and_grad_pyscf( + params, hamiltonian, n_qubits, n_elec_s, ex_ops: Tuple, param_ids: Tuple, mode: str = "fermion", init_state=None +): + params = tq.backend.numpy(params) + ket = get_civector_pyscf(params, n_qubits, n_elec_s, ex_ops, param_ids, mode, init_state) + bra = tq.backend.numpy(apply_op(hamiltonian, ket)) + energy = bra @ ket + + gradients_beforesum = _get_gradients_pyscf(bra, ket, params, n_qubits, n_elec_s, ex_ops, param_ids, mode) + + gradients = np.zeros(params.shape) + for grad, param_id in zip(gradients_beforesum, param_ids): + gradients[param_id] += grad + + return energy, 2 * gradients + + +def _get_gradients_pyscf(bra, ket, params, n_qubits, n_elec_s, ex_ops, param_ids, mode): + assert mode == "fermion" + + n_orb = n_qubits // 2 + na, nb = unpack_nelec(n_elec_s) + + gradients_beforesum = [] + for param_id, ex_op in reversed(list(zip(param_ids, ex_ops))): + theta = params[param_id] + bra = evolve_excitation_pyscf(bra, ex_op, n_orb, n_elec_s, -theta) + ket = evolve_excitation_pyscf(ket, ex_op, n_orb, n_elec_s, -theta) + ket_pyscf = CIvectorPySCF(ket, n_orb, na, nb) + fket = apply_a_pyscf(ket_pyscf, ex_op) + grad = bra @ fket + gradients_beforesum.append(grad) + gradients_beforesum = list(reversed(gradients_beforesum)) + gradients_beforesum = np.array(gradients_beforesum) + + return gradients_beforesum + + +def apply_excitation_pyscf(civector, n_qubits, n_elec_s, f_idx, mode): + assert mode == "fermion" + na, nb = unpack_nelec(n_elec_s) + civector_pyscf = CIvectorPySCF(civector, n_qubits // 2, na, nb) + return apply_a_pyscf(civector_pyscf, f_idx) diff --git a/src/tyxonq/chem/static/evolve_statevector.py b/src/tyxonq/chem/static/evolve_statevector.py new file mode 100644 index 0000000..d48aa57 --- /dev/null +++ b/src/tyxonq/chem/static/evolve_statevector.py @@ -0,0 +1,102 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +import logging +from functools import partial + +import numpy as np +from openfermion import jordan_wigner +import tyxonq as tq + +from ..utils.backend import jit +from ..constants import ad_a_hc2, adad_aa_hc2, ad_a_hc, adad_aa_hc +from ..utils.misc import ex_op_to_fop +from ..static.evolve_tensornetwork import get_init_circuit + + +logger = logging.getLogger(__name__) + + +@partial(jit, static_argnums=[1, 2, 3, 4, 5]) +def get_statevector(params, n_qubits, n_elec, ex_ops, param_ids, mode="fermion", init_state=None): + if tq.backend.name == "jax": + logger.info(f"Entering `get_statevector`. n_qubit: {n_qubits}") + if param_ids is None: + assert len(params) == len(ex_ops) + param_ids = list(range(len(params))) + + circuit = get_init_circuit(n_qubits, n_elec, mode, init_state) + + # note only the real part is taken + statevector = tq.backend.convert_to_tensor(circuit.state()) + + for param_id, f_idx in zip(param_ids, ex_ops): + theta = params[param_id] + statevector = evolve_excitation(statevector, f_idx, theta, mode) + + return statevector.real.reshape(-1) + + +@partial(jit, static_argnums=[1, 3]) +def evolve_excitation(statevector, f_idx, theta, mode): + n_qubits = round(np.log2(statevector.shape[0])) + qubit_idx = [n_qubits - 1 - idx for idx in f_idx] + # fermion operator operated on ket, twice + f2ket = tq.Circuit(n_qubits, inputs=statevector) + # fermion operator index, not sorted + if len(qubit_idx) == 2: + f2ket.any(*qubit_idx, unitary=ad_a_hc2) + else: + assert len(qubit_idx) == 4 + f2ket.any(*qubit_idx, unitary=adad_aa_hc2) + + # fermion operator operated on ket + fket = apply_excitation(statevector, f_idx, mode) + statevector += (1 - tq.backend.cos(theta)) * f2ket.state() + tq.backend.sin(theta) * fket + return statevector + + +@partial(jit, static_argnums=[1, 2]) +def apply_excitation(statevector, f_idx, mode): + n_qubits = round(np.log2(statevector.shape[0])) + qubit_idx = [n_qubits - 1 - idx for idx in f_idx] + circuit = tq.Circuit(n_qubits, inputs=statevector) + # fermion operator index, not sorted + if len(qubit_idx) == 2: + circuit.any(*qubit_idx, unitary=ad_a_hc) + else: + assert len(qubit_idx) == 4 + circuit.any(*qubit_idx, unitary=adad_aa_hc) + + if mode != "fermion": + return circuit.state() + + # apply all Z operators + # pauli string index, already sorted + fop = ex_op_to_fop(f_idx) + qop = jordan_wigner(fop) + z_indices = [] + for idx, term in next(iter(qop.terms.keys())): + if term != "Z": + assert idx in f_idx + continue + z_indices.append(n_qubits - 1 - idx) + if sorted(qop.terms.items())[0][1].real > 0: + sign = 1 + else: + sign = -1 + phase_vector = [sign] + for i in range(n_qubits): + if i in z_indices: + phase_vector = np.kron(phase_vector, [1, -1]) + else: + phase_vector = np.kron(phase_vector, [1, 1]) + return phase_vector * circuit.state() + + +# external interface +def apply_excitation_statevector(statevector, n_qubits, n_elec, f_idx, mode): + return apply_excitation(statevector, f_idx, mode).real diff --git a/src/tyxonq/chem/static/evolve_tensornetwork.py b/src/tyxonq/chem/static/evolve_tensornetwork.py new file mode 100644 index 0000000..bd4fcd0 --- /dev/null +++ b/src/tyxonq/chem/static/evolve_tensornetwork.py @@ -0,0 +1,204 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +import logging +from typing import Any, List, Tuple +from functools import partial + +import numpy as np +from openfermion import jordan_wigner +import tyxonq as tq + +from ..static.ci_utils import get_ci_strings, civector_to_statevector +from ..utils.backend import jit +from ..utils.misc import ex_op_to_fop, reverse_qop_idx, unpack_nelec +from ..utils.circuit import evolve_pauli, multicontrol_ry + + +logger = logging.getLogger(__name__) + + +Tensor = Any + + +def evolve_excitation(circuit: tq.Circuit, f_idx: tuple, qop, theta, mode: str, decompose_multicontrol: bool): + n_qubits = circuit.circuit_param["nqubits"] + f_idx = [n_qubits - 1 - idx for idx in f_idx] + + z_indices = [] + if mode == "fermion": + for idx, term in next(iter(qop.terms.keys())): + if term != "Z": + assert idx in f_idx + continue + z_indices.append(idx) + + def parity_gates(target, reverse=False): + if len(z_indices) == 0: + return + if not reverse: + for ii in range(len(z_indices) - 1): + circuit.CNOT(z_indices[ii], z_indices[ii + 1]) + circuit.cz(z_indices[-1], target) + else: + circuit.cz(z_indices[-1], target) + for ii in reversed(range(len(z_indices) - 1)): + circuit.CNOT(z_indices[ii], z_indices[ii + 1]) + + if len(f_idx) == 2: + k, l = f_idx + circuit.CNOT(k, l) + parity_gates(k) + circuit.cry(l, k, theta=theta) + parity_gates(k, reverse=True) + circuit.CNOT(k, l) + else: + assert len(f_idx) == 4 + k, l, i, j = f_idx + circuit.CNOT(l, k) + circuit.CNOT(j, i) + circuit.CNOT(l, j) + parity_gates(l) + if not decompose_multicontrol: + try: + name = f"Ry({theta:.4f})" + except TypeError: + # jax tracer can't be formatted + name = "Ry" + circuit.multicontrol(i, j, k, l, ctrl=[0, 1, 0], unitary=tq.gates.ry_gate(theta).tensor, name=name) + else: + circuit.append(multicontrol_ry(theta), indices=[i, j, k, l]) + parity_gates(l, reverse=True) + circuit.CNOT(l, j) + circuit.CNOT(j, i) + circuit.CNOT(l, k) + + return circuit + + +def get_init_circuit(n_qubits, n_elec_s, mode: str, init_state=None, givens_swap=False): + na, nb = unpack_nelec(n_elec_s) + if init_state is None: + # prepare HF state + circuit = tq.Circuit(n_qubits) + if mode in ["fermion", "qubit"]: + for i in range(nb): + circuit.X(n_qubits - 1 - i) + for i in range(na): + circuit.X(n_qubits // 2 - 1 - i) + else: + assert mode == "hcb" + if not givens_swap: + for i in range(na): + circuit.X(n_qubits - 1 - i) + else: + for i in range(na): + circuit.X(i) + elif isinstance(init_state, tq.Circuit): + return init_state + else: + ci_strings = get_ci_strings(n_qubits, n_elec_s, mode == "hcb") + statevector = civector_to_statevector(init_state, n_qubits, ci_strings) + if givens_swap: + statevector = statevector.reshape([2] * n_qubits) + new_idx = list(range(n_qubits - na, n_qubits)) + list(range(n_qubits - na)) + statevector = statevector.transpose(new_idx).ravel() + circuit = tq.Circuit(n_qubits, inputs=statevector) + return circuit + + +def get_circuit( + params, + n_qubits, + n_elec, + ex_ops, + param_ids, + mode: str = "fermion", + init_state: tq.Circuit = None, + decompose_multicontrol: bool = False, + trotter: bool = False, +): + if param_ids is None: + assert len(params) == len(ex_ops) + param_ids = list(range(len(params))) + + circuit = get_init_circuit(n_qubits, n_elec, mode, init_state) + + for param_id, f_idx in zip(param_ids, ex_ops): + theta = params[param_id] + fop = ex_op_to_fop(f_idx, with_conjugation=True) + qop = reverse_qop_idx(jordan_wigner(fop), n_qubits) + if trotter: + for pauli_string, v in qop.terms.items(): + if mode in ["qubit", "hcb"]: + pauli_string = [(idx, symbol) for idx, symbol in pauli_string if symbol != "Z"] + circuit = evolve_pauli(circuit, pauli_string, -2 * v.imag * theta) + else: + # https://arxiv.org/pdf/2005.14475.pdf + circuit = evolve_excitation(circuit, f_idx, qop, 2 * theta, mode, decompose_multicontrol) + + return circuit + + +def get_gs_unitary(theta): + # SWAP @ Givens Rotation + a = [ + [1, 0, 0, 0], + [0, -tq.backend.sin(theta), tq.backend.cos(theta), 0], + [0, tq.backend.cos(theta), tq.backend.sin(theta), 0], + [0, 0, 0, 1], + ] + return tq.backend.convert_to_tensor(a) + + +def get_gs_indices(no: int, nv: int) -> List[Tuple[int, int]]: + """Givens-Swap indices""" + layer1 = [np.array([no - 1, no])] + for _ in range(nv - 1): + layer1.append(layer1[-1] + 1) + layer1 = np.array(layer1) + ret = [layer1] + for _ in range(no - 1): + ret.append(ret[-1] - 1) + ret = np.array(ret).reshape(-1, 2) + assert len(ret) == no * nv + return ret.tolist() + + +def get_circuit_givens_swap(params, n_qubits, n_elec, init_state=None): + # https://arxiv.org/pdf/2002.00035.pdf + # the swapped qubit index represents molecule orbitals 0 to n_qubits - 1 + # so we need to take negative of theta + # ┌───┐ ┌──────┐ + # q_0(MO 1) : ┤ X ├────────┤ ├───────── MO 3 + # ├───┤┌──────┐│ GS │┌──────┐ + # q_1(MO 0) : ┤ X ├┤ ├┤ 3, 1 ├┤ ├─ MO 2 + # └───┘│ GS │├──────┤│ GS │ + # q_2(MO 3) : ─────┤ 3, 0 ├┤ ├┤ 2, 1 ├─ MO 1 + # └──────┘│ GS │└──────┘ + # q_3(MO 2) : ─────────────┤ 2, 0 ├───────── MO 0 + # └──────┘ + circuit = get_init_circuit(n_qubits, n_elec, mode="hcb", init_state=init_state, givens_swap=True) + gs_indices = get_gs_indices(n_elec // 2, n_qubits - n_elec // 2) + for i, (j, k) in enumerate(gs_indices): + theta = params[i] + unitary = get_gs_unitary(theta) + try: + name = f"Givens-SWAP({theta:.4f})" + except TypeError: + # jax tracer can't be formatted + name = "Givens-SWAP" + circuit.any(j, k, unitary=unitary, name=name) + return circuit + + +@partial(jit, static_argnums=[1, 2, 3, 4, 5]) +def get_statevector_tensornetwork(params, n_qubits, n_elec, ex_ops, param_ids, mode="fermion", init_state=None): + if tq.backend.name == "jax": + logger.info(f"Entering `get_statevector_tensornetwork`. n_qubit: {n_qubits}") + + circuit = get_circuit(params, n_qubits, n_elec, ex_ops, param_ids, mode, init_state) + return circuit.state().real.reshape(-1) diff --git a/src/tyxonq/chem/static/hamiltonian.py b/src/tyxonq/chem/static/hamiltonian.py new file mode 100644 index 0000000..10e1585 --- /dev/null +++ b/src/tyxonq/chem/static/hamiltonian.py @@ -0,0 +1,280 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +import logging +from inspect import isfunction +from itertools import product +from typing import Tuple, List + +import numpy as np +import tensornetwork as tn +from renormalizer import Mpo +from openfermion import FermionOperator, QubitOperator +from pyscf.scf.hf import RHF +from pyscf.mcscf import CASCI +from pyscf.fci import direct_nosym, cistring +from pyscf import ao2mo +import tyxonq as tq +from tyxonq import QuOperator + +from ..utils.misc import hcb_to_coo, fop_to_coo, reverse_qop_idx, canonical_mo_coeff, get_n_qubits +from ..constants import DISCARD_EPS + + +logger = logging.getLogger(__name__) + + +def get_integral_from_hf(hf: RHF, active_space: Tuple = None, aslst: List[int] = None): + if not isinstance(hf, RHF): + raise TypeError(f"hf object must be RHF class, got {type(hf)}") + m = hf.mol + assert hf.mo_coeff is not None + # todo(weitangli): don't modify inplace + hf.mo_coeff = canonical_mo_coeff(hf.mo_coeff) + + if active_space is None: + nelecas = m.nelectron + ncas = m.nao + else: + nelecas, ncas = active_space + + casci = CASCI(hf, ncas, nelecas) + if aslst is None: + int1e, e_core = casci.get_h1eff() + int2e = ao2mo.restore("s1", casci.get_h2eff(), ncas) + else: + mo = casci.sort_mo(aslst, base=0) + int1e, e_core = casci.get_h1eff(mo) + int2e = ao2mo.restore("s1", casci.get_h2eff(mo), ncas) + + return int1e, int2e, e_core + + +def get_hop_from_integral(int1e, int2e): + n_orb = int1e.shape[0] + if int1e.shape != (n_orb, n_orb): + raise ValueError(f"Invalid one-boby integral array shape: {int1e.shape}") + int2e = ao2mo.restore(1, int2e, n_orb) + assert int2e.shape == (n_orb, n_orb, n_orb, n_orb) + n_sorb = n_orb * 2 + + logger.info("Creating Hamiltonian operators") + + h1e = np.zeros((n_sorb, n_sorb)) + h2e = np.zeros((n_sorb, n_sorb, n_sorb, n_sorb)) + + h1e[:n_orb, :n_orb] = h1e[n_orb:, n_orb:] = int1e + + for p, q, r, s in product(range(n_sorb), repeat=4): + # a_p^\dagger a_q^\dagger a_r a_s + if ((p < n_orb) == (s < n_orb)) and ((q < n_orb) == (r < n_orb)): + # note the different orders of the indices + h2e[p, q, r, s] = int2e[p % n_orb, s % n_orb, q % n_orb, r % n_orb] + + op1e = [] + for p, q in product(range(n_sorb), repeat=2): + # a_p^\dagger a_q + v = h1e[p, q] + if np.abs(v) < DISCARD_EPS: + continue + op = FermionOperator(f"{p}^ {q}", v) + op1e.append(op) + + op2e = [] + for q, s in product(range(n_sorb), repeat=2): + for p, r in product(range(q), range(s)): + # a_p^\dagger a_q^\dagger a_r a_s + v = h2e[p, q, r, s] - h2e[q, p, r, s] + if np.abs(v) < DISCARD_EPS: + continue + op = FermionOperator(f"{p}^ {q}^ {r} {s}", v) + op2e.append(op) + + logger.info("Summing Hamiltonian operators") + ops = FermionOperator() + for op in op1e + op2e: + ops += op + + return ops + + +def qubit_operator(string: str, coeff: float) -> QubitOperator: + ret = coeff + terms = string.split(" ") + for term in terms: + if term[-1] == "^": + sign = -1 + term = term[:-1] + else: + sign = 1 + idx = int(term) + ret *= (QubitOperator(f"X{idx}") + sign * 1j * QubitOperator(f"Y{idx}")) / 2 + return ret + + +def get_hop_hcb_from_integral(int1e, int2e): + # Hard core boson Hamiltonian + # https://arxiv.org/pdf/2002.00035.pdf + n_orb = int1e.shape[0] + qop = QubitOperator() + for p in range(n_orb): + for q in range(p + 1): + if p == q: + qop += qubit_operator(f"{p}^ {p}", 2 * int1e[p, p] + int2e[p, p, p, p]) + else: + qop += qubit_operator(f"{p}^ {q}", int2e[p, q, p, q]) + qop += qubit_operator(f"{q}^ {p}", int2e[q, p, q, p]) + qop += qubit_operator(f"{p}^ {p} {q}^ {q}", 4 * int2e[p, p, q, q] - 2 * int2e[p, q, p, q]) + qop = reverse_qop_idx(qop, n_orb) + return qop + + +def get_h_sparse_from_integral(int1e, int2e, mode="fermion", do_log=False): + if mode in ["fermion", "qubit"]: + ops = get_hop_from_integral(int1e, int2e) + else: + assert mode == "hcb" + ops = get_hop_hcb_from_integral(int1e, int2e) + if do_log: + logger.info("Constructing sparse Hamiltonian") + if mode in ["fermion", "qubit"]: + h_sparse = fop_to_coo(ops, n_qubits=2 * len(int1e)) + else: + assert mode == "hcb" + h_sparse = hcb_to_coo(ops, n_qubits=len(int1e)) + if do_log: + logger.info("Sparse Hamiltonian constructed") + return h_sparse + + +def get_h_fcifunc_from_integral(int1e, int2e, n_elec): + n_orb = len(int1e) + h2e = direct_nosym.absorb_h1e(int1e, int2e, n_orb, n_elec, 0.5) + + def fci_func(civector): + civector = tq.backend.numpy(civector).astype(np.float64) + civector = direct_nosym.contract_2e(h2e, civector, norb=n_orb, nelec=n_elec) + return tq.backend.convert_to_tensor(civector).astype(tq.rdtypestr) + + return fci_func + + +def get_h_fcifunc_hcb_from_integral(int1e, int2e, n_elec): + # todo: how about using https://github.com/pyscf/doci + n_orb = len(int1e) + ci_strings = cistring.make_strings(range(n_orb), n_elec // 2) + + def fci_func(civector): + res = tq.backend.zeros(len(civector), dtype=tq.rdtypestr) + for p in range(n_orb): + for q in range(p + 1): + if p == q: + bitmask = 1 << p + arraymask = (ci_strings & bitmask) == bitmask + res += (civector * arraymask) * (2 * int1e[p, p] + int2e[p, p, p, p]) + else: + bitmask = (1 << p) + (1 << q) + excitation = ci_strings ^ bitmask + addr = cistring.strs2addr(n_orb, n_elec // 2, excitation) + flip = ci_strings ^ (1 << p) + masked_flip = flip & bitmask + arraymask = (masked_flip == bitmask) | (masked_flip == 0) + res += civector[addr] * arraymask * int2e[p, q, p, q] + arraymask = (ci_strings & bitmask) == bitmask + res += (civector * arraymask) * (4 * int2e[p, p, q, q] - 2 * int2e[p, q, p, q]) + return res + + return fci_func + + +def get_h_from_integral(int1e, int2e, n_elec_s, mode: str, htype: str): + if htype == "sparse": + hamiltonian = get_h_sparse_from_integral(int1e, int2e, mode=mode, do_log=True) + else: + assert htype.lower() == "fcifunc" + if mode in ["fermion", "qubit"]: + hamiltonian = get_h_fcifunc_from_integral(int1e, int2e, n_elec_s) + else: + assert mode == "hcb" + n_elec = sum(n_elec_s) + hamiltonian = get_h_fcifunc_hcb_from_integral(int1e, int2e, n_elec) + return hamiltonian + + +def get_h_from_hf(hf: RHF, active_space: Tuple = None, mode: str = "fermion", htype="sparse"): + if not isinstance(hf, RHF): + raise TypeError(f"hf object must RHF class, got {type(hf)}") + htype = htype.lower() + if not htype in ["sparse", "mpo", "fcifunc"]: + raise ValueError(f"htype must be 'sparse' or 'mpo', got '{htype}'") + int1e, int2e, e_core = get_integral_from_hf(hf, active_space) + if active_space is None: + n_elec = hf.mol.nelectron + else: + n_elec = active_space[0] + assert n_elec % 2 == 0 + n_elec_s = [n_elec // 2, n_elec // 2] + + hamiltonian = get_h_from_integral(int1e, int2e, n_elec_s, mode, htype) + + if active_space is None: + return hamiltonian + else: + return hamiltonian, e_core + + +def mpo_to_quoperator(mpo: Mpo): + array_list = [m.array for m in mpo] + # squeeze dim 1 out + assert len(array_list) >= 2 + a, b, c, d = array_list[0].shape + array_list[0] = array_list[0].reshape(b, c, d) + a, b, c, d = array_list[-1].shape + array_list[-1] = array_list[-1].reshape(a, b, c) + # convert MPO to tensor-network + node_list = [tn.Node(array) for array in array_list] + + node_list[0][2] ^ node_list[1][0] + for i in range(1, len(node_list) - 1): + node_list[i][3] ^ node_list[i + 1][0] + + in_edges = [node_list[0][0]] + [node[1] for node in node_list[1:]] + out_edges = [node_list[0][1]] + [node[2] for node in node_list[1:]] + + qop = QuOperator(out_edges, in_edges) + return qop + + +def apply_op(op, state): + if isinstance(op, list): + # in MPO form + n_qubit = get_n_qubits(state) + n_qubit_op = get_n_qubits(op) + assert n_qubit == n_qubit_op + mps = tq.QuVector.from_tensor(state.reshape([2] * n_qubit)) + h_qop = mpo_to_quoperator(op) + return (h_qop @ mps).eval().reshape(-1) + if isfunction(op): + return op(state) + else: + return op @ state + + +def random_integral(nao: int, seed: int = 2077): + np.random.seed(seed) + int1e = np.random.uniform(-1, 1, size=(nao, nao)) + int2e = np.random.uniform(-1, 1, size=(nao, nao, nao, nao)) + int1e = 0.5 * (int1e + int1e.T) + int2e = symmetrize_int2e(int2e) + return int1e, int2e + + +def symmetrize_int2e(int2e): + int2e = 0.25 * ( + int2e + int2e.transpose((0, 1, 3, 2)) + int2e.transpose((1, 0, 2, 3)) + int2e.transpose((2, 3, 0, 1)) + ) + int2e = 0.5 * (int2e + int2e.transpose(3, 2, 1, 0)) + return int2e diff --git a/src/tyxonq/chem/static/hea.py b/src/tyxonq/chem/static/hea.py new file mode 100644 index 0000000..c8f6446 --- /dev/null +++ b/src/tyxonq/chem/static/hea.py @@ -0,0 +1,1108 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +import logging +from functools import partial +from itertools import product +from time import time +from typing import Callable, Union, Any, List, Tuple + +import numpy as np +from scipy.optimize import minimize +from noisyopt import minimizeSPSA +from pyscf.gto.mole import Mole +from openfermion import ( + hermitian_conjugated, + jordan_wigner, + bravyi_kitaev, + binary_code_transform, + checksum_code, + parity_code, + get_sparse_operator, + QubitOperator, + FermionOperator, +) +from qiskit import QuantumCircuit +import tyxonq as tq +from tyxonq import Circuit, DMCircuit +from tyxonq.noisemodel import NoiseConf, circuit_with_noise + +from ..static.engine_hea import ( + QpuConf, + get_statevector, + get_densitymatrix, + get_energy_tensornetwork, + get_energy_tensornetwork_noise, + get_energy_tensornetwork_shot, + get_energy_tensornetwork_noise_shot, + get_energy_qpu, + get_energy_and_grad_tensornetwork, + get_energy_and_grad_tensornetwork_noise, + get_energy_and_grad_tensornetwork_shot, + get_energy_and_grad_tensornetwork_noise_shot, + get_energy_and_grad_qpu, +) +from ..static.hamiltonian import get_hop_from_integral +from ..utils.misc import reverse_fop_idx, scipy_opt_wrap, reverse_qop_idx +from ..utils.circuit import get_circuit_dataframe + +logger = logging.getLogger(__file__) + + +Tensor = Any + + +def get_noise_conf(probability=0.02): + noise_conf = NoiseConf() + # note TQ definition of error probability + channel = tq.channels.isotropicdepolarizingchannel(probability, 2) + + def condition(qir): + return len(qir["index"]) == 2 + + noise_conf.add_noise_by_condition(condition, channel) + return noise_conf + + +def binary(fermion_operator: FermionOperator, n_modes: int, n_elec: int) -> QubitOperator: + """ + Performs binary transformation. + + Parameters + ---------- + fermion_operator: FermionOperator + The fermion operator. + n_modes: int + The number of modes. + n_elec: int + The number of electrons. + + Returns + ------- + qubit_operator: QubitOperator + """ + return binary_code_transform(fermion_operator, 2 * checksum_code(n_modes // 2, (n_elec // 2) % 2)) + + +def _parity(fermion_operator, n_modes): + return binary_code_transform(fermion_operator, parity_code(n_modes)) + + +def parity(fermion_operator: FermionOperator, n_modes: int, n_elec: Union[int, Tuple[int, int]]) -> QubitOperator: + """ + Performs parity transformation and saves 2 qubits assuming electron number conservation in each spin sector. + For example, ``011011`` is first transformed to ``010010`` and then ``0101``. + ``110100`` is transformed to ``100111`` and then `1011`. + + Parameters + ---------- + fermion_operator: FermionOperator + The fermion operator. + n_modes: int + The number of modes (spin-orbitals). + n_elec: int or tuple + The number of electrons, or numbers of alpha/beta electrons + + Returns + ------- + qubit_operator: QubitOperator + """ + qubit_operator = _parity(reverse_fop_idx(fermion_operator, n_modes), n_modes) + res = 0 + assert n_modes % 2 == 0 + reduction_indices = [n_modes // 2 - 1, n_modes - 1] + if isinstance(n_elec, int): + if n_elec % 2 != 0: + raise ValueError("Specify alpha and beta spin as a tuple when the total number or electron is not even") + n_elec_s = [n_elec // 2, n_elec // 2] + else: + n_elec_s = n_elec + phase_alpha = (-1) ** n_elec_s[0] + phase_beta = (-1) ** sum(n_elec_s) + for qop in qubit_operator: + # qop example: 0.5 [Z1 X2 X3] + pauli_string, coeff = next(iter(qop.terms.items())) + # pauli_string example: ((1, 'Z'), (2, 'X'), (3, 'X')) + # coeff example: 0.5 + new_pauli_string = [] + for idx, symbol in pauli_string: + is_alpha = idx <= reduction_indices[0] + if idx in reduction_indices: + if symbol in ["X", "Y"]: + # discard this term because the bit will never change + continue + else: + assert symbol == "Z" + if is_alpha: + coeff *= phase_alpha + else: + coeff *= phase_beta + continue + if not is_alpha: + idx -= 1 + new_pauli_string.append((idx, symbol)) + qop.terms = {tuple(new_pauli_string): coeff} + res += qop + return res + + +def fop_to_qop(fop: FermionOperator, mapping: str, n_sorb: int, n_elec: Union[int, Tuple[int, int]]) -> QubitOperator: + if mapping == "parity": + qop = parity(fop, n_sorb, n_elec) + elif mapping in ["jordan-wigner", "jordan_wigner"]: + qop = reverse_qop_idx(jordan_wigner(fop), n_sorb) + elif mapping in ["bravyi-kitaev", "bravyi_kitaev"]: + qop = reverse_qop_idx(bravyi_kitaev(fop, n_sorb), n_sorb) + else: + raise ValueError(f"Unknown mapping: {mapping}") + return qop + + +def get_ry_circuit(params: Tensor, n_qubits: int, n_layers: int) -> Circuit: + """ + Get the parameterized :math:`R_y` circuit. + + Parameters + ---------- + params: Tensor + The circuit parameters. + n_qubits: int + The number of qubits. + n_layers: int + The number of layers in the ansatz. + + Returns + ------- + c: Circuit + """ + c = Circuit(n_qubits) + params = params.reshape(n_layers + 1, n_qubits) + for i in range(n_qubits): + c.ry(i, theta=params[0, i]) + c.barrier_instruction(*range(n_qubits)) + for l in range(n_layers): + for i in range(n_qubits - 1): + c.cnot(i, (i + 1)) + for i in range(n_qubits): + c.ry(i, theta=params[l + 1, i]) + c.barrier_instruction(*range(n_qubits)) + return c + + +class HEA: + """ + Run hardware-efficient ansatz calculation. + For a comprehensive tutorial see :doc:`/tutorial_jupyter/noisy_simulation`. + """ + + @classmethod + def from_molecule(cls, m: Mole, active_space=None, n_layers=3, mapping="parity", **kwargs): + """ + Construct the HEA class from the given molecule. + The :math:`R_y` ansatz is employed. By default, the number of layers is set to 3. + + + Parameters + ---------- + m: Mole + The molecule object. + active_space: Tuple[int, int], optional + Active space approximation. The first integer is the number of electrons and the second integer is + the number or spatial-orbitals. Defaults to None. + n_layers: int + The number of layers in the :math:`R_y` ansatz. Defaults to 3. + mapping: str + The fermion to qubit mapping. Supported mappings are ``"parity"``, + and ``"bravyi-kitaev"``. + + kwargs: + Other arguments to be passed to the :func:`__init__` function such as ``engine``. + + + Returns + ------- + hea: :class:`HEA` + An HEA instance + """ + from tencirchem import UCC + + if m.spin == 0: + init_method = "mp2" + else: + init_method = "zeros" + ucc = UCC(m, init_method=init_method, active_space=active_space, run_ccsd=False, run_fci=False) + return cls.ry(ucc.int1e, ucc.int2e, ucc.n_elec_s, ucc.e_core, n_layers=n_layers, mapping=mapping, **kwargs) + + @classmethod + def ry( + cls, + int1e: np.ndarray, + int2e: np.ndarray, + n_elec: Union[int, Tuple[int, int]], + e_core: float, + n_layers: int, + init_circuit: Circuit = None, + mapping: str = "parity", + **kwargs, + ): + r""" + Construct the HEA class from electron integrals and the :math:`R_y` ansatz. + The circuit consists of interleaved layers of $R_y$ and CNOT gates + + .. math:: + + |\Psi(\theta)\rangle=\prod_{l=k}^1\left [ L_{R_y}^{(l)}(\theta) L_{CNOT}^{(l)} \right ] L_{R_y}^{(0)}(\theta) |{\phi}\rangle + + where $k$ is the total number of layers, and the layers are defined as + + .. math:: + L_{CNOT}^{(l)}=\prod_{j=N/2-1}^1 CNOT_{2j, 2j+1} \prod_{j=N/2}^{1} CNOT_{2j-1, 2j} + + .. math:: + L_{R_y}^{(l)}(\theta)=\prod_{j=N}^{1} RY_{j}(\theta_{lj}) + + Overlap integral is assumed to be identity. + Parity transformation is used to transform from fermion operators to qubit operators by default. + + Parameters + ---------- + int1e: np.ndarray + One-body integral in spatial orbital. + int2e: np.ndarray + Two-body integral, in spatial orbital, chemists' notation, and without considering symmetry. + n_elec: int or tuple + The number of electrons, or numbers of alpha/beta electrons + e_core: float + The nuclear energy or core energy if active space approximation is involved. + n_layers: int + The number of layers in the ansatz. + init_circuit: Circuit + The initial circuit before the :math:`R_y` ansatz. Defaults to None. + mapping: str + The fermion to qubit mapping. Supported mappings are ``"parity"``, + and ``"bravyi-kitaev"``. + + kwargs: + Other arguments to be passed to the :func:`__init__` function such as ``engine``. + + Returns + ------- + hea: :class:`HEA` + An HEA instance + """ + n_sorb = 2 * len(int1e) + if mapping == "parity": + n_qubits = n_sorb - 2 + elif mapping in ["jordan-wigner", "jordan_wigner", "bravyi-kitaev", "bravyi_kitaev"]: + n_qubits = n_sorb + else: + raise ValueError(f"Unknown mapping: {mapping}") + + init_guess = np.random.random((n_layers + 1, n_qubits)).ravel() + + def get_circuit(params): + if init_circuit is None: + c = Circuit(n_qubits) + else: + c = Circuit.from_qir(init_circuit.to_qir(), init_circuit.circuit_param) + return c.append(get_ry_circuit(params, n_qubits, n_layers)) + + instance = cls.from_integral(int1e, int2e, n_elec, e_core, mapping, get_circuit, init_guess, **kwargs) + instance.mapping = mapping + return instance + + @classmethod + def from_integral( + cls, + int1e: np.ndarray, + int2e: np.ndarray, + n_elec: Union[int, Tuple[int, int]], + e_core: float, + mapping: str, + circuit: Union[Callable, QuantumCircuit], + init_guess: np.ndarray, + **kwargs, + ): + """ + Construct the HEA class from electron integrals and custom quantum circuit. + Overlap integral is assumed to be identity. + + Parameters + ---------- + int1e: np.ndarray + One-body integral in spatial orbital. + int2e: np.ndarray + Two-body integral, in spatial orbital, chemists' notation, and without considering symmetry. + n_elec: int or tuple + The number of electrons, or numbers of alpha/beta electrons + e_core: float + The nuclear energy or core energy if active space approximation is involved. + mapping: str + The fermion to qubit mapping. Supported mappings are ``"parity"``, + and ``"bravyi-kitaev"``. + circuit: Callable or QuantumCircuit + The ansatz as a function or Qiskit :class:`QuantumCircuit` + init_guess: list of float or :class:`np.ndarray` + The parameter initial guess. + + kwargs: + Other arguments to be passed to the :func:`__init__` function such as ``engine``. + + Returns + ------- + hea: :class:`HEA` + An HEA instance + """ + + if isinstance(n_elec, tuple): + n_elec_s = n_elec + n_elec = sum(n_elec) + spin = abs(n_elec_s[0] - n_elec_s[1]) + else: + if n_elec % 2 != 0: + raise ValueError("Specify alpha and beta spin as a tuple when the total number or electron is not even") + n_elec_s = (n_elec // 2, n_elec // 2) + spin = 0 + + hop = get_hop_from_integral(int1e, int2e) + e_core + n_sorb = 2 * len(int1e) + h_qubit_op = fop_to_qop(hop, mapping, n_sorb, n_elec_s) + + instance = cls(h_qubit_op, circuit, init_guess, **kwargs) + + instance.mapping = mapping + instance.int1e = int1e + instance.int2e = int2e + instance.n_elec = n_elec + instance.spin = spin + instance.e_core = e_core + instance.hop = hop + return instance + + @classmethod + def as_pyscf_solver(cls, config_function: Callable = None, opt_engine: str = None, **kwargs): + """ + Converts the ``HEA`` class to a PySCF FCI solver using :math:`R_y` ansatz. + + Parameters + ---------- + config_function: callable + A function to configure the ``HEA`` instance. + Accepts the ``HEA`` instance and modifies it inplace before :func:`kernel` is called. + opt_engine: str + The engine to use when performing the circuit parameter optimization. + kwargs + Other arguments to be passed to the :func:`__init__` function such as ``engine``. + + Returns + ------- + FCISolver + + Examples + -------- + >>> from pyscf.mcscf import CASSCF + >>> from tencirchem import HEA + >>> from tencirchem.molecule import h8 + >>> # normal PySCF workflow + >>> hf = h8.HF() + >>> round(hf.kernel(), 6) + -4.149619 + >>> casscf = CASSCF(hf, 2, 2) + >>> # set the FCI solver for CASSCF to be HEA + >>> casscf.fcisolver = HEA.as_pyscf_solver(n_layers=1) + >>> round(casscf.kernel()[0], 6) + -4.166473 + """ + + class FakeFCISolver: + def __init__(self): + self.instance: HEA = None + self.config_function = config_function + self.instance_kwargs = kwargs.copy() + if "n_layers" not in self.instance_kwargs: + self.instance_kwargs["n_layers"] = 1 + + def kernel(self, h1, h2, norb, nelec, ci0=None, ecore=0, **kwargs): + self.instance = cls.ry(h1, h2, nelec, e_core=ecore, **self.instance_kwargs) + if self.config_function is not None: + self.config_function(self.instance) + if opt_engine is None: + e = self.instance.kernel() + else: + engine_bak = self.instance.engine + self.instance.engine = opt_engine + self.instance.kernel() + self.instance.engine = engine_bak + e = self.instance.energy() + return e, self.instance.params + + def make_rdm1(self, params, norb, nelec): + rdm1 = self.instance.make_rdm1(params) + return rdm1 + + def make_rdm12(self, params, norb, nelec): + rdm1 = self.instance.make_rdm1(params) + rdm2 = self.instance.make_rdm2(params) + return rdm1, rdm2 + + def spin_square(self, params, norb, nelec): + return 0, 1 + + return FakeFCISolver() + + def __init__( + self, + h_qubit_op: QubitOperator, + circuit: Union[Callable, QuantumCircuit], + init_guess: Union[List[float], np.ndarray], + engine: str = None, + engine_conf: [NoiseConf, QpuConf] = None, + ): + """ + Construct the HEA class from Hamiltonian in :class:`QubitOperator` form and the ansatz. + + Parameters + ---------- + h_qubit_op: QubitOperator + Hamiltonian as openfermion :class:`QubitOperator`. + circuit: Callable or QuantumCircuit + The ansatz as a function or Qiskit :class:`QuantumCircuit` + init_guess: list of float or :class:`np.ndarray` + The parameter initial guess. + engine: str, optional + The engine to run the calculation. See :ref:`advanced:Engines` for details. + Defaults to ``"tensornetwork"``. + engine_conf: NoiseConf, optional + The noise configuration for the circuit. Defaults to None, in which case + if a noisy engine is used, an isotropic depolarizing error channel for all 2-qubit gates + with :math:`p=0.02` is added to the circuit. + """ + self._check_engine(engine) + + if engine is None: + engine = "tensornetwork" + if engine == "tensornetwork" and engine_conf is not None: + raise ValueError("Tensornetwork engine does not have engine configuration") + if engine_conf is None: + if engine.startswith("tensornetwork-noise"): + engine_conf = get_noise_conf() + if engine.startswith("qpu"): + engine_conf = QpuConf() + + init_guess = np.array(init_guess) + + # convert to real coefficients and prevent complex energy expectation outcome + self.h_qubit_op = QubitOperator() + for k, v in h_qubit_op.terms.items(): + if not np.iscomplex(v): + v = v.real + self.h_qubit_op.terms[k] = v + + if isinstance(circuit, Callable): + self.get_circuit = circuit + # sanity check + c: Circuit = self.get_circuit(init_guess) + if isinstance(c, DMCircuit): + raise TypeError("`circuit` function should return Circuit instead of DMCircuit") + self.n_qubits = c.circuit_param["nqubits"] + elif isinstance(circuit, QuantumCircuit): + + def get_circuit(params): + return Circuit.from_qiskit(circuit, binding_params=params) + + self.get_circuit = get_circuit + assert circuit.num_parameters == len(init_guess) + self.n_qubits = circuit.num_qubits + else: + raise TypeError("circuit must be callable or qiskit QuantumCircuit") + + self.h_array = np.array(get_sparse_operator(self.h_qubit_op, self.n_qubits).todense()) + + if init_guess.ndim != 1: + raise ValueError(f"Init guess should be one-dimensional. Got shape {init_guess}") + self.init_guess = init_guess + self.engine = engine + self.engine_conf = engine_conf + self.shots = 4096 + self._grad = "param-shift" + + self.scipy_minimize_options = None + self._params = None + self.opt_res = None + + # allow setting these attributes for features such as calculating RDM + # could make it a function for customized mapping + self.mapping: str = None # fermion-to-qubit mapping + self.int1e = None + self.int2e = None + self.n_elec = None + self.spin = None + self.e_core = None + self.hop = None + + def get_dmcircuit(self, params: Tensor = None, noise_conf: NoiseConf = None) -> DMCircuit: + """ + Get the :class:`DMCircuit` with noise. + Only valid for ``"tensornetwork-noise"`` and ``"tensornetwork-noise&shot"`` engines. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + noise_conf: NoiseConf, optional + The noise configuration for the circuit. Defaults to None, in which case + ``self.engine_conf`` is used. + + Returns + ------- + dmcircuit: DMCircuit + """ + params = self._check_params_argument(params) + dmcircuit = self.get_dmcircuit_no_noise(params) + if noise_conf is None: + assert self.engine.startswith("tensornetwork-noise") + noise_conf = self.engine_conf + dmcircuit = circuit_with_noise(dmcircuit, noise_conf) + return dmcircuit + + def get_dmcircuit_no_noise(self, params: Tensor = None) -> DMCircuit: + """ + Get the :class:`DMCircuit` without noise. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + + Returns + ------- + dmcircuit: DMCircuit + """ + qir = self.get_circuit(params).to_qir() + dmcircuit = DMCircuit.from_qir(qir) + return dmcircuit + + def _check_engine(self, engine): + supported_engine = [ + None, + "tensornetwork", + "tensornetwork-noise", + "tensornetwork-shot", + "tensornetwork-noise&shot", + "qpu", + ] + if not engine in supported_engine: + raise ValueError(f"Engine '{engine}' not supported") + + def _check_params_argument(self, params, strict=True): + if params is None: + if self.params is not None: + params = self.params + else: + if strict: + raise ValueError("Run the `.kernel` method to determine the parameters first") + else: + if self.init_guess is not None: + params = self.init_guess + else: + params = np.zeros(self.n_params) + + if len(params) != self.n_params: + raise ValueError(f"Incompatible parameter shape. {self.n_params} is desired. Got {len(params)}") + return tq.backend.convert_to_tensor(params).astype(tq.rdtypestr) + + def statevector(self, params: Tensor = None) -> Tensor: + """ + Evaluate the circuit state vector without considering noise. + Valid for ``"tensornetwork"`` and ``"tensornetwork-shot"`` engine. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + + Returns + ------- + statevector: Tensor + Corresponding state vector + + See Also + -------- + densitymatrix: Evaluate the circuit density matrix in the presence of circuit noise. + energy: Evaluate the total energy. + energy_and_grad: Evaluate the total energy and parameter gradients. + + Examples + -------- + >>> import numpy as np + >>> from tencirchem import UCC, HEA + >>> from tencirchem.molecule import h2 + >>> ucc = UCC(h2) + >>> hea = HEA.ry(ucc.int1e, ucc.int2e, ucc.n_elec, ucc.e_core, n_layers=1) + >>> np.round(hea.statevector([0, np.pi, 0, 0]), 8) # HF state + array([0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j]) + """ + # Evaluate the circuit state vector without noise. + # only engine is "tensornetwork" + params = self._check_params_argument(params) + statevector = get_statevector(params, self.get_circuit) + return statevector + + def densitymatrix(self, params: Tensor = None) -> Tensor: + """ + Evaluate the circuit density matrix in the presence of circuit noise. + Only valid for ``"tensornetwork-noise"`` and ``"tensornetwork-noise&shot"`` engines. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + + Returns + ------- + densitymatrix: Tensor + + See Also + -------- + statevector: Evaluate the circuit state vector. + energy: Evaluate the total energy. + energy_and_grad: Evaluate the total energy and parameter gradients. + + Examples + -------- + >>> import numpy as np + >>> from tencirchem import UCC, HEA + >>> from tencirchem.molecule import h2 + >>> ucc = UCC(h2) + >>> hea = HEA.ry(ucc.int1e, ucc.int2e, ucc.n_elec, ucc.e_core, + ... n_layers=1, engine="tensornetwork-noise") + >>> np.round(hea.densitymatrix([0, np.pi, 0, 0]), 8) # HF state with noise + array([[0.00533333+0.j, 0. +0.j, 0. +0.j, 0. +0.j], + [0. +0.j, 0.984 +0.j, 0. +0.j, 0. +0.j], + [0. +0.j, 0. +0.j, 0.00533333+0.j, 0. +0.j], + [0. +0.j, 0. +0.j, 0. +0.j, 0.00533333+0.j]]) + """ + # engines are "tensornetwork-noise" and "tensornetwork-noise&shot" + params = self._check_params_argument(params) + # the last two arguments should be identical and not garbage collected for each call for `jit` to work + densitymatrix = get_densitymatrix(params, self.get_dmcircuit_no_noise, self.engine_conf) + return densitymatrix + + def energy(self, params: Tensor = None, engine: str = None) -> float: + """ + Evaluate the total energy. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + engine: str, optional + The engine to use. Defaults to ``None``, which uses ``self.engine``. + + Returns + ------- + energy: float + Total energy + + See Also + -------- + statevector: Evaluate the circuit state vector. + densitymatrix: Evaluate the circuit density matrix in the presence of circuit noise. + energy_and_grad: Evaluate the total energy and parameter gradients. + + Examples + -------- + >>> import numpy as np + >>> from tencirchem import UCC, HEA + >>> from tencirchem.molecule import h2 + >>> ucc = UCC(h2) + >>> hea = HEA.ry(ucc.int1e, ucc.int2e, ucc.n_elec, ucc.e_core, + ... n_layers=1, engine="tensornetwork-noise") + >>> # HF state, no noise + >>> round(hea.energy([0, np.pi, 0, 0], "tensornetwork"), 8) + -1.11670614 + >>> # HF state, gate noise + >>> round(hea.energy([0, np.pi, 0, 0], "tensornetwork-noise"), 4) + -1.1001 + >>> # HF state, measurement noise. Set the number of shots by `hea.shots` + >>> round(hea.energy([0, np.pi, 0, 0], "tensornetwork-shot"), 1) + -1.1 + >>> # HF state, gate+measurement noise + >>> hea.energy([0, np.pi, 0, 0], "tensornetwork-noise&shot") # doctest:+ELLIPSIS + -1... + """ + params = self._check_params_argument(params) + if engine is None: + engine = self.engine + if engine == "tensornetwork": + e = get_energy_tensornetwork(params, self.h_array, self.get_circuit) + elif engine == "tensornetwork-noise": + e = get_energy_tensornetwork_noise(params, self.h_array, self.get_dmcircuit_no_noise, self.engine_conf) + elif engine == "tensornetwork-shot": + e = get_energy_tensornetwork_shot( + params, + tuple(self.h_qubit_op.terms.keys()), + list(self.h_qubit_op.terms.values()), + self.get_circuit, + self.shots, + ) + elif engine == "tensornetwork-noise&shot": + e = get_energy_tensornetwork_noise_shot( + params, + tuple(self.h_qubit_op.terms.keys()), + list(self.h_qubit_op.terms.values()), + self.get_dmcircuit_no_noise, + self.engine_conf, + self.shots, + ) + else: + assert engine == "qpu" + e = get_energy_qpu( + params, + tuple(self.h_qubit_op.terms.keys()), + list(self.h_qubit_op.terms.values()), + self.get_circuit, + self.engine_conf, + self.shots, + ) + return e + + def energy_and_grad(self, params: Tensor = None, engine: str = None, grad: str = None) -> Tuple[float, Tensor]: + """ + Evaluate the total energy and parameter gradients using parameter-shift rule. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + engine: str, optional + The engine to use. Defaults to ``None``, which uses ``self.engine``. + grad: str, optional + The algorithm to use for the gradient. Defaults to ``None``, which means ``self.grad`` will be used. + Possible options are ``"param-shift"`` for parameter-shift rule and + ``"autodiff"`` for auto-differentiation. + Note that ``"autodiff"`` is not compatible with ``"tensornetwork-shot"`` + and ``"tensornetwork-noise&shot"`` engine. + + Returns + ------- + energy: float + Total energy + grad: Tensor + The parameter gradients + + See Also + -------- + statevector: Evaluate the circuit state vector. + densitymatrix: Evaluate the circuit density matrix in the presence of circuit noise. + energy: Evaluate the total energy. + """ + + params = self._check_params_argument(params) + + if engine is None: + engine = self.engine + + if grad is None: + grad = self.grad + + if grad == "free": + raise ValueError("Must provide a gradient algorithm") + + if engine == "tensornetwork": + e, grad_array = get_energy_and_grad_tensornetwork(params, self.h_array, self.get_circuit, grad) + elif engine == "tensornetwork-noise": + e, grad_array = get_energy_and_grad_tensornetwork_noise( + params, self.h_array, self.get_dmcircuit_no_noise, self.engine_conf, grad + ) + elif engine == "tensornetwork-shot": + if grad == "autodiff": + raise ValueError(f"Engine {engine} is incompatible with grad method {grad}") + e, grad_array = get_energy_and_grad_tensornetwork_shot( + params, + tuple(self.h_qubit_op.terms.keys()), + list(self.h_qubit_op.terms.values()), + self.get_circuit, + self.shots, + grad, + ) + elif engine == "tensornetwork-noise&shot": + if grad == "autodiff": + raise ValueError(f"Engine {engine} is incompatible with grad method {grad}") + e, grad_array = get_energy_and_grad_tensornetwork_noise_shot( + params, + tuple(self.h_qubit_op.terms.keys()), + list(self.h_qubit_op.terms.values()), + self.get_dmcircuit_no_noise, + self.engine_conf, + self.shots, + grad, + ) + else: + assert engine == "qpu" + if grad == "autodiff": + raise ValueError(f"Engine {engine} is incompatible with grad method {grad}") + e, grad_array = get_energy_and_grad_qpu( + params, + tuple(self.h_qubit_op.terms.keys()), + list(self.h_qubit_op.terms.values()), + self.get_circuit, + self.shots, + grad, + ) + return e, grad_array + + def kernel(self): + logger.info("Begin optimization") + + func, stating_time = self.get_opt_function(with_time=True) + + time1 = time() + if self.grad == "free": + if self.engine in ["tensornetwork", "tensornetwork-noise", "qpu"]: + opt_res = minimize(func, x0=self.init_guess, method="COBYLA", options=self.scipy_minimize_options) + else: + assert self.engine in ["tensornetwork-shot", "tensornetwork-noise&shot"] + opt_res = minimizeSPSA(func, x0=self.init_guess, paired=False, niter=100, disp=True) + else: + opt_res = minimize( + func, x0=self.init_guess, jac=True, method="L-BFGS-B", options=self.scipy_minimize_options + ) + + time2 = time() + + if not opt_res.success: + logger.warning("Optimization failed. See `.opt_res` for details.") + + opt_res["staging_time"] = stating_time + opt_res["opt_time"] = time2 - time1 + opt_res["init_guess"] = self.init_guess + opt_res["e"] = float(opt_res.fun) + self.opt_res = opt_res + # prepare for future modification + self.params = opt_res.x.copy() + return opt_res.e + + def get_opt_function(self, grad: str = None, with_time: bool = False) -> Union[Callable, Tuple[Callable, float]]: + """ + Returns the cost function in SciPy format for optimization. + Basically a wrapper to :func:`energy_and_grad` or :func:`energy`, + + + Parameters + ---------- + with_time: bool, optional + Whether return staging time. Defaults to False. + grad: str, optional + The algorithm to use for the gradient. Defaults to ``None``, which means ``self.grad`` will be used. + Possible options are ``"param-shift"`` for parameter-shift rule and + ``"autodiff"`` for auto-differentiation. + Note that ``"autodiff"`` is not compatible with ``"tensornetwork-noise&shot"`` engine. + Returns + ------- + opt_function: Callable + The optimization cost function in SciPy format. + time: float + Staging time. Returned when ``with_time`` is set to ``True``. + """ + if grad is None: + grad = self.grad + + if grad != "free": + func = scipy_opt_wrap(partial(self.energy_and_grad, engine=self.engine)) + else: + func = scipy_opt_wrap(partial(self.energy, engine=self.engine), gradient=False) + + time1 = time() + if tq.backend.name == "jax": + logger.info("JIT compiling the circuit") + _ = func(np.zeros(self.n_params)) + logger.info("Circuit JIT compiled") + time2 = time() + if with_time: + return func, time2 - time1 + return func + + def make_rdm1(self, params: Tensor = None) -> np.ndarray: + r""" + Evaluate the spin-traced one-body reduced density matrix (1RDM). + + .. math:: + + \textrm{1RDM}[p,q] = \langle p_{\alpha}^\dagger q_{\alpha} \rangle + + \langle p_{\beta}^\dagger q_{\beta} \rangle + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + + Returns + ------- + rdm1: np.ndarray + The spin-traced one-body RDM. + + See Also + -------- + make_rdm2: Evaluate the spin-traced two-body reduced density matrix (2RDM). + """ + + if params is None: + params = self._check_params_argument(params) + if self.mapping is None: + raise ValueError("Must first set the fermion-to-qubit mapping") + + if self.mapping == "parity": + n_sorb = self.n_qubits + 2 + else: + n_sorb = self.n_qubits + n_orb = n_sorb // 2 + rdm1 = np.zeros([n_orb] * 2) + + # could optimize for tn engine by caching the statevector or dm + for i in range(n_orb): + for j in range(i + 1): + if self.spin == 0: + fop = 2 * FermionOperator(f"{i}^ {j}") + else: + fop = FermionOperator(f"{i}^ {j}") + FermionOperator(f"{i+n_orb}^ {j+n_orb}") + fop = fop + hermitian_conjugated(fop) + qop = fop_to_qop(fop, self.mapping, n_sorb, self.n_elec_s) + hea = HEA(qop, self.get_circuit, params, self.engine, self.engine_conf) + # divide by two since we have added the hermitian conjugation + rdm1[i, j] = rdm1[j, i] = hea.energy(params) / 2 + + return rdm1 + + def make_rdm2(self, params: Tensor = None) -> np.ndarray: + r""" + Evaluate the spin-traced two-body reduced density matrix (2RDM). + + .. math:: + + \begin{aligned} + \textrm{2RDM}[p,q,r,s] & = \langle p_{\alpha}^\dagger r_{\alpha}^\dagger + s_{\alpha} q_{\alpha} \rangle + + \langle p_{\beta}^\dagger r_{\alpha}^\dagger s_{\alpha} q_{\beta} \rangle \\ + & \quad + \langle p_{\alpha}^\dagger r_{\beta}^\dagger s_{\beta} q_{\alpha} \rangle + + \langle p_{\beta}^\dagger r_{\beta}^\dagger s_{\beta} q_{\beta} \rangle + \end{aligned} + + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + + Returns + ------- + rdm2: np.ndarray + The spin-traced two-body RDM. + + See Also + -------- + make_rdm1: Evaluate the spin-traced one-body reduced density matrix (1RDM). + """ + if params is None: + params = self._check_params_argument(params) + if self.mapping is None: + raise ValueError("Must first set the fermion-to-qubit mapping") + + if self.mapping == "parity": + n_sorb = self.n_qubits + 2 + else: + n_sorb = self.n_qubits + n_orb = n_sorb // 2 + rdm2 = np.zeros([n_orb] * 4) + + calculated_indices = set() + # a^\dagger_p a^\dagger_q a_r a_s + # possible spins: aaaa, abba, baab, bbbb + for p, q, r, s in product(range(n_orb), repeat=4): + if (p, q, r, s) in calculated_indices: + continue + + fop_aaaa = FermionOperator(f"{p}^ {q}^ {r} {s}") + fop_abba = FermionOperator(f"{p}^ {q+n_orb}^ {r+n_orb} {s}") + if self.spin == 0: + # aaaa is the same as bbbb, abba is the same as baab + fop = 2 * (fop_aaaa + fop_abba) + else: + fop_bbbb = FermionOperator(f"{p+n_orb}^ {q+n_orb}^ {r+n_orb} {s+n_orb}") + fop_baab = FermionOperator(f"{p+n_orb}^ {q}^ {r} {s+n_orb}") + fop = fop_aaaa + fop_abba + fop_bbbb + fop_baab + fop = fop + hermitian_conjugated(fop) + qop = fop_to_qop(fop, self.mapping, n_sorb, self.n_elec_s) + hea = HEA(qop, self.get_circuit, params, self.engine, self.engine_conf) + # divide by two since we have added the hermitian conjugation + v = hea.energy(params) / 2 + indices = [(p, q, r, s), (s, r, q, p), (q, p, s, r), (r, s, p, q)] + for idx in indices: + rdm2[idx] = v + calculated_indices.add(idx) + # transpose to PySCF notation: rdm2[p,q,r,s] = + rdm2 = rdm2.transpose(0, 3, 1, 2) + return rdm2 + + def print_circuit(self): + c = self.get_circuit(self.init_guess) + df = get_circuit_dataframe(c) + + def format_flop(f): + return f"{f:.3e}" + + formatters = {"flop": format_flop} + print(df.to_string(index=True, formatters=formatters)) + + def print_summary(self): + print("############################### Circuit ###############################") + self.print_circuit() + print("######################### Optimization Result #########################") + if self.opt_res is None: + print("Optimization not run yet") + else: + print(self.opt_res) + + @property + def grad(self): + return self._grad + + @grad.setter + def grad(self, v): + if not v in ["param-shift", "autodiff", "free"]: + raise ValueError(f"Invalid gradient method {v}") + self._grad = v + + @property + def n_elec_s(self): + """The number of electrons for alpha and beta spin""" + return (self.n_elec + self.spin) // 2, (self.n_elec - self.spin) // 2 + + @property + def n_params(self): + """The number of parameter in the ansatz/circuit.""" + return len(self.init_guess) + + @property + def params(self): + """The circuit parameters.""" + if self._params is not None: + return self._params + if self.opt_res is not None: + return self.opt_res.x + return None + + @params.setter + def params(self, params): + self._params = params diff --git a/src/tyxonq/chem/static/kupccgsd.py b/src/tyxonq/chem/static/kupccgsd.py new file mode 100644 index 0000000..347c726 --- /dev/null +++ b/src/tyxonq/chem/static/kupccgsd.py @@ -0,0 +1,279 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from time import time +import logging +from typing import Tuple, List, Union + +import numpy as np +from pyscf.gto.mole import Mole +from pyscf.scf import RHF + +from ..static.ucc import UCC + +logger = logging.getLogger(__name__) + + +class KUPCCGSD(UCC): + """ + Run :math:`k`-UpCCGSD calculation. + The interfaces are similar to :class:`UCCSD `. + """ + + def __init__( + self, + mol: Union[Mole, RHF], + active_space: Tuple[int, int] = None, + aslst: List[int] = None, + mo_coeff: np.ndarray = None, + k: int = 3, + n_tries: int = 1, + engine: str = None, + run_hf: bool = True, + run_mp2: bool = True, + run_ccsd: bool = True, + run_fci: bool = True, + ): + r""" + Initialize the class with molecular input. + + Parameters + ---------- + mol: Mole or RHF + The molecule as PySCF ``Mole`` object or the PySCF ``RHF`` object + active_space: Tuple[int, int], optional + Active space approximation. The first integer is the number of electrons and the second integer is + the number or spatial-orbitals. Defaults to None. + aslst: List[int], optional + Pick orbitals for the active space. Defaults to None which means the orbitals are sorted by energy. + The orbital index is 0-based. + + .. note:: + See `PySCF document `_ + for choosing the active space orbitals. Here orbital index is 0-based, whereas in PySCF by default it + is 1-based. + + mo_coeff: np.ndarray, optional + Molecule coefficients. If provided then RHF is skipped. + Can be used in combination with the ``init_state`` attribute. + Defaults to None which means RHF orbitals are used. + k: int, optional + The number of layers in the ansatz. Defaults to 3 + n_tries: int, optional + The number of different initial points used for VQE calculation. + For large circuits usually a lot of runs are required for good accuracy. + Defaults to 1. + engine: str, optional + The engine to run the calculation. See :ref:`advanced:Engines` for details. + run_hf: bool, optional + Whether run HF for molecule orbitals. Defaults to ``True``. + run_mp2: bool, optional + Whether run MP2 for energy reference. Defaults to ``True``. + run_ccsd: bool, optional + Whether run CCSD for energy reference. Defaults to ``True``. + run_fci: bool, optional + Whether run FCI for energy reference. Defaults to ``True``. + + See Also + -------- + tencirchem.UCCSD + tencirchem.PUCCD + tencirchem.UCC + """ + super().__init__( + mol, + init_method=None, + active_space=active_space, + aslst=aslst, + mo_coeff=mo_coeff, + engine=engine, + run_hf=run_hf, + run_mp2=run_mp2, + run_ccsd=run_ccsd, + run_fci=run_fci, + ) + # the number of layers + self.k = k + # the number of different initialization + self.n_tries = n_tries + self.ex_ops, self.param_ids, self.init_guess = self.get_ex_ops(self.t1, self.t2) + self.init_guess_list = [self.init_guess] + for _ in range(self.n_tries - 1): + self.init_guess_list.append(np.random.rand(self.n_params) - 0.5) + self.e_tries_list = [] + self.opt_res_list = [] + self.staging_time = self.opt_time = None + + def kernel(self): + _, stating_time = self.get_opt_function(with_time=True) + + time1 = time() + for i in range(self.n_tries): + logger.info(f"k-UpCCGSD try {i}") + if self.n_tries == 1: + if not np.allclose(self.init_guess, self.init_guess_list[0]): + logger.info("Inconsistent `self.init_guess` and `self.init_guess_list`. Use `self.init_guess`.") + else: + self.init_guess = self.init_guess_list[i] + super().kernel() + logger.info(f"k-UpCCGSD try {i} energy {self.opt_res.fun}") + self.opt_res_list.append(self.opt_res) + self.opt_res_list.sort(key=lambda x: x.fun) + self.e_tries_list = [float(res.fun) for res in self.opt_res_list] + time2 = time() + + self.staging_time = stating_time + self.opt_time = time2 - time1 + self.opt_res = self.opt_res_list[0] + self.opt_res.e_tries = self.e_tries_list + + if not self.opt_res.success: + logger.warning("Optimization failed. See `.opt_res` for details.") + + self.init_guess = self.opt_res.init_guess + return float(self.opt_res.e) + + def get_ex_ops(self, t1: np.ndarray = None, t2: np.ndarray = None) -> Tuple[List[Tuple], List[int], np.ndarray]: + """ + Get one-body and two-body excitation operators for :math:`k`-UpCCGSD ansatz. + The excitations are generalized and two-body excitations are restricted to paired ones. + Initial guesses are generated randomly. + + Parameters + ---------- + t1: np.ndarray, optional + Not used. Kept for consistency with the parent method. + t2: np.ndarray, optional + Not used. Kept for consistency with the parent method. + + Returns + ------- + ex_op: List[Tuple] + The excitation operators. Each operator is represented by a tuple of ints. + param_ids: List[int] + The mapping from excitations to parameters. + init_guess: np.ndarray + The initial guess for the parameters. + + See Also + -------- + get_ex1_ops: Get generalized one-body excitation operators. + get_ex2_ops: Get generalized paired two-body excitation operators. + + Examples + -------- + >>> from tencirchem import KUPCCGSD + >>> from tencirchem.molecule import h2 + >>> kupccgsd = KUPCCGSD(h2) + >>> ex_op, param_ids, init_guess = kupccgsd.get_ex_ops() + >>> ex_op + [(1, 3, 2, 0), (3, 2), (1, 0), (1, 3, 2, 0), (3, 2), (1, 0), (1, 3, 2, 0), (3, 2), (1, 0)] + >>> param_ids + [0, 1, 1, 2, 3, 3, 4, 5, 5] + >>> init_guess # doctest:+ELLIPSIS + array([...]) + """ + ex1_ops, ex1_param_id, _ = self.get_ex1_ops() + ex2_ops, ex2_param_id, _ = self.get_ex2_ops() + + ex_op = [] + param_ids = [-1] + for _ in range(self.k): + ex_op.extend(ex2_ops + ex1_ops) + param_ids.extend([i + param_ids[-1] + 1 for i in ex2_param_id]) + param_ids.extend([i + param_ids[-1] + 1 for i in ex1_param_id]) + init_guess = np.random.rand(max(param_ids) + 1) - 0.5 + return ex_op, param_ids[1:], init_guess + + def get_ex1_ops(self, t1: np.ndarray = None) -> Tuple[List[Tuple], List[int], np.ndarray]: + """ + Get generalized one-body excitation operators. + + Parameters + ---------- + t1: np.ndarray, optional + Not used. Kept for consistency with the parent method. + + Returns + ------- + ex_op: List[Tuple] + The excitation operators. Each operator is represented by a tuple of ints. + param_ids: List[int] + The mapping from excitations to parameters. + init_guess: np.ndarray + The initial guess for the parameters. + + See Also + -------- + get_ex2_ops: Get generalized paired two-body excitation operators. + get_ex_ops: Get one-body and two-body excitation operators for :math:`k`-UpCCGSD ansatz. + """ + assert t1 is None + no, nv = self.no, self.nv + + ex1_ops = [] + ex1_param_id = [-1] + + for a in range(no + nv): + for i in range(a): + # alpha to alpha + ex_op_a = (no + nv + a, no + nv + i) + # beta to beta + ex_op_b = (a, i) + ex1_ops.extend([ex_op_a, ex_op_b]) + ex1_param_id.extend([ex1_param_id[-1] + 1] * 2) + + ex1_init_guess = np.zeros(max(ex1_param_id) + 1) + return ex1_ops, ex1_param_id[1:], ex1_init_guess + + def get_ex2_ops(self, t2: np.ndarray = None) -> Tuple[List[Tuple], List[int], np.ndarray]: + """ + Get generalized paired two-body excitation operators. + + Parameters + ---------- + t2: np.ndarray, optional + Not used. Kept for consistency with the parent method. + + Returns + ------- + ex_op: List[Tuple] + The excitation operators. Each operator is represented by a tuple of ints. + param_ids: List[int] + The mapping from excitations to parameters. + init_guess: np.ndarray + The initial guess for the parameters. + + See Also + -------- + get_ex1_ops: Get one-body excitation operators. + get_ex_ops: Get one-body and two-body excitation operators for :math:`k`-UpCCGSD ansatz. + """ + + assert t2 is None + no, nv = self.no, self.nv + + ex2_ops = [] + ex2_param_id = [-1] + + for a in range(no + nv): + for i in range(a): + # i correspond to a and j correspond to b, as in PySCF convention + # otherwise the t2 amplitude has incorrect phase + # paired + ex_op_ab = (a, no + nv + a, no + nv + i, i) + ex2_ops.append(ex_op_ab) + ex2_param_id.append(ex2_param_id[-1] + 1) + + ex2_init_guess = np.zeros(max(ex2_param_id) + 1) + return ex2_ops, ex2_param_id[1:], ex2_init_guess + + @property + def e_kupccgsd(self): + """ + Returns :math:`k`-UpCCGSD energy + """ + return self.energy() diff --git a/src/tyxonq/chem/static/puccd.py b/src/tyxonq/chem/static/puccd.py new file mode 100644 index 0000000..faf79a1 --- /dev/null +++ b/src/tyxonq/chem/static/puccd.py @@ -0,0 +1,248 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from typing import Tuple, List, Union + +import numpy as np +from pyscf.gto.mole import Mole +from pyscf.scf import RHF +from pyscf.cc.addons import spatial2spin +from pyscf.fci import cistring +from tyxonq import Circuit + +from ..utils.misc import rdm_mo2ao +from ..static.ucc import UCC +from ..static.ci_utils import get_ci_strings +from ..static.evolve_tensornetwork import get_circuit_givens_swap + + +class PUCCD(UCC): + """ + Run paired UCC calculation. + The interfaces are similar to :class:`UCCSD `. + """ + + # todo: more documentation here and make the references right. + # separate docstring examples for a variety of methods, such as energy() + # also need to add a few comment on make_rdm1/2 + # https://arxiv.org/pdf/2002.00035.pdf + # https://arxiv.org/pdf/1503.04878.pdf + + def __init__( + self, + mol: Union[Mole, RHF], + init_method: str = "mp2", + active_space: Tuple[int, int] = None, + aslst: List[int] = None, + mo_coeff: np.ndarray = None, + engine: str = None, + run_hf: bool = True, + run_mp2: bool = True, + run_ccsd: bool = True, + run_fci: bool = True, + ): + r""" + Initialize the class with molecular input. + + Parameters + ---------- + mol: Mole or RHF + The molecule as PySCF ``Mole`` object or the PySCF ``RHF`` object + init_method: str, optional + How to determine the initial amplitude guess. Accepts ``"mp2"`` (default), ``"ccsd"``,``"fe"`` + and ``"zeros"``. + active_space: Tuple[int, int], optional + Active space approximation. The first integer is the number of electrons and the second integer is + the number or spatial-orbitals. Defaults to None. + aslst: List[int], optional + Pick orbitals for the active space. Defaults to None which means the orbitals are sorted by energy. + The orbital index is 0-based. + + .. note:: + See `PySCF document `_ + for choosing the active space orbitals. Here orbital index is 0-based, whereas in PySCF by default it + is 1-based. + + mo_coeff: np.ndarray, optional + Molecule coefficients. If provided then RHF is skipped. + Can be used in combination with the ``init_state`` attribute. + Defaults to None which means RHF orbitals are used. + engine: str, optional + The engine to run the calculation. See :ref:`advanced:Engines` for details. + run_hf: bool, optional + Whether run HF for molecule orbitals. Defaults to ``True``. + run_mp2: bool, optional + Whether run MP2 for initial guess and energy reference. Defaults to ``True``. + run_ccsd: bool, optional + Whether run CCSD for initial guess and energy reference. Defaults to ``True``. + run_fci: bool, optional + Whether run FCI for energy reference. Defaults to ``True``. + + See Also + -------- + tencirchem.UCCSD + tencirchem.KUPCCGSD + tencirchem.UCC + """ + super().__init__( + mol, + init_method, + active_space, + aslst, + mo_coeff, + mode="hcb", + engine=engine, + run_hf=run_hf, + run_mp2=run_mp2, + run_ccsd=run_ccsd, + run_fci=run_fci, + ) + self.ex_ops, self.param_ids, self.init_guess = self.get_ex_ops(self.t1, self.t2) + + def get_ex1_ops(self, t1: np.ndarray = None): + """Not applicable. Use :func:`get_ex_ops`.""" + raise NotImplementedError + + def get_ex2_ops(self, t2: np.ndarray = None): + """Not applicable. Use :func:`get_ex_ops`.""" + raise NotImplementedError + + def get_ex_ops(self, t1: np.ndarray = None, t2: np.ndarray = None) -> Tuple[List[Tuple], List[int], List[float]]: + """ + Get paired excitation operators for pUCCD ansatz. + + Parameters + ---------- + t1: np.ndarray, optional + Not used. Kept for consistency with the parent method. + t2: np.ndarray, optional + Not used. Kept for consistency with the parent method. + + Returns + ------- + ex_op: List[Tuple] + The excitation operators. Each operator is represented by a tuple of ints. + param_ids: List[int] + The mapping from excitations to parameters. + init_guess: List[float] + The initial guess for the parameters. + + Examples + -------- + >>> from tencirchem import PUCCD + >>> from tencirchem.molecule import h2 + >>> puccd = PUCCD(h2) + >>> ex_op, param_ids, init_guess = puccd.get_ex_ops() + >>> ex_op + [(1, 0)] + >>> param_ids + [0] + >>> init_guess # doctest:+ELLIPSIS + [...] + """ + no, nv = self.no, self.nv + if t2 is None: + t2 = np.zeros((no, no, nv, nv)) + t2 = spatial2spin(t2) + + ex_ops = [] + ex_init_guess = [] + # to be consistent with givens rotation circuit + for i in range(no): + for a in range(nv - 1, -1, -1): + ex_ops.append((no + a, i)) + ex_init_guess.append(t2[2 * i, 2 * i + 1, 2 * a, 2 * a + 1]) + return ex_ops, list(range(len(ex_ops))), ex_init_guess + + def make_rdm1(self, statevector=None, basis="AO"): + __doc__ = super().make_rdm1.__doc__ + + civector = self._statevector_to_civector(statevector) + ci_strings = get_ci_strings(self.n_qubits, self.n_elec_s, self.mode) + + n_active = self.n_qubits + rdm1_cas = np.zeros([n_active] * 2) + for i in range(n_active): + bitmask = 1 << i + arraymask = (ci_strings & bitmask) == bitmask + value = float(civector @ (arraymask * civector)) + rdm1_cas[i, i] = 2 * value + rdm1 = self.embed_rdm_cas(rdm1_cas) + if basis == "MO": + return rdm1 + else: + return rdm_mo2ao(rdm1, self.hf.mo_coeff) + + def make_rdm2(self, statevector=None, basis="AO"): + __doc__ = super().make_rdm2.__doc__ + + civector = self._statevector_to_civector(statevector) + ci_strings = get_ci_strings(self.n_qubits, self.n_elec_s, self.mode) + + n_active = self.n_qubits + rdm2_cas = np.zeros([n_active] * 4) + for p in range(n_active): + for q in range(p + 1): + maskq = 1 << q + maskp = 1 << p + maskpq = maskp + maskq + arraymask = (ci_strings & maskq) == maskq + if p == q: + value = float(civector @ (arraymask * civector)) + else: + arraymask &= ((~ci_strings) & maskp) == maskp + excitation = ci_strings ^ maskpq + addr = cistring.strs2addr(n_active, self.n_elec // 2, excitation) + value = float(civector @ (arraymask * civector[addr])) + + rdm2_cas[p, q, p, q] = rdm2_cas[q, p, q, p] = value + if p == q: + continue + arraymask = (ci_strings & maskpq) == maskpq + value = float(civector @ (arraymask * civector)) + + rdm2_cas[p, p, q, q] = rdm2_cas[q, q, p, p] = 2 * value + rdm2_cas[p, q, q, p] = rdm2_cas[q, p, p, q] = -value + rdm2_cas *= 2 + rdm2 = self.embed_rdm_cas(rdm2_cas) + # no need to transpose + if basis == "MO": + return rdm2 + else: + return rdm_mo2ao(rdm2, self.hf.mo_coeff) + + def get_circuit(self, params=None, trotter=False, givens_swap=False) -> Circuit: + """ + Get the circuit as TyxonQ ``Circuit`` object. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter. + If :func:`kernel` is not called before, the initial guess is used. + trotter: bool, optional + Whether Trotterize the UCC factor into Pauli strings. + Defaults to False. + givens_swap: bool, optional + Whether return the circuit with Givens-Swap gates. + + Returns + ------- + circuit: :class:`tc.Circuit` + The quantum circuit. + """ + if not givens_swap: + return super().get_circuit(params, trotter=trotter) + else: + params = self._check_params_argument(params, strict=False) + return get_circuit_givens_swap(params, self.n_qubits, self.n_elec, self.init_state) + + @property + def e_puccd(self): + """ + Returns pUCCD energy + """ + return self.energy() diff --git a/src/tyxonq/chem/static/ucc.py b/src/tyxonq/chem/static/ucc.py new file mode 100644 index 0000000..2babbdd --- /dev/null +++ b/src/tyxonq/chem/static/ucc.py @@ -0,0 +1,1523 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from functools import partial +from itertools import product +from collections import defaultdict +from time import time +import logging +from typing import Any, Tuple, Callable, List, Union + +import numpy as np +from scipy.optimize import minimize +from scipy.special import comb +import pandas as pd +from openfermion import jordan_wigner, FermionOperator, QubitOperator +from pyscf.gto.mole import Mole +from pyscf.scf import RHF +from pyscf.scf import ROHF +from pyscf.scf.hf import RHF as RHF_TYPE +from pyscf.scf.rohf import ROHF as ROHF_TYPE +from pyscf.cc.addons import spatial2spin +from pyscf.mcscf import CASCI +from pyscf import fci +import tyxonq as tq + +from ..constants import DISCARD_EPS +from ..molecule import _Molecule +from ..utils.misc import reverse_qop_idx, scipy_opt_wrap, rdm_mo2ao, canonical_mo_coeff +from ..utils.circuit import get_circuit_dataframe +from ..static.engine_ucc import ( + get_civector, + get_statevector, + get_energy, + get_energy_and_grad, + apply_excitation, + translate_init_state, +) +from ..static.hamiltonian import ( + get_integral_from_hf, + get_h_from_integral, + get_hop_from_integral, + get_hop_hcb_from_integral, +) +from ..static.ci_utils import get_ci_strings, get_ex_bitstring, get_addr, get_init_civector +from ..static.evolve_civector import get_fermion_phase +from ..static.evolve_tensornetwork import get_circuit + + +logger = logging.getLogger(__name__) + + +Tensor = Any + + +class UCC: + """ + Base class for :class:`UCCSD`. + """ + + @classmethod + def from_integral( + cls, + int1e: np.ndarray, + int2e: np.ndarray, + n_elec: Union[int, Tuple[int, int]], + e_core: float = 0, + ovlp: np.ndarray = None, + **kwargs, + ): + """ + Construct UCC classes from electron integrals. + + Parameters + ---------- + int1e: np.ndarray + One-body integral in spatial orbital. + int2e: np.ndarray + Two-body integral, in spatial orbital, chemists' notation, and without considering symmetry. + n_elec: int or tuple + The number of electrons, or numbers of alpha/beta electrons + e_core: float, optional + The nuclear energy or core energy if active space approximation is involved. + Defaults to 0. + ovlp: np.ndarray + The overlap integral. Defaults to ``None`` and identity matrix is used. + kwargs: + Other arguments to be passed to the :func:`__init__` function such as ``engine``. + + Returns + ------- + ucc: :class:`UCC` + A UCC instance + """ + if isinstance(n_elec, tuple): + spin = abs(n_elec[0] - n_elec[1]) + n_elec = n_elec[0] + n_elec[1] + else: + assert n_elec % 2 == 0 + spin = 0 + m = _Molecule(int1e, int2e, n_elec, spin, e_core, ovlp) + return cls(m, **kwargs) + + @classmethod + def as_pyscf_solver(cls, config_function: Callable = None, **kwargs): + """ + Converts the ``UCC`` class to a PySCF FCI solver. + + Parameters + ---------- + config_function: callable + A function to configure the ``UCC`` instance. + Accepts the ``UCC`` instance and modifies it inplace before :func:`kernel` is called. + kwargs + Other arguments to be passed to the :func:`__init__` function such as ``engine``. + + Returns + ------- + FCISolver + + Examples + -------- + >>> from pyscf.mcscf import CASSCF + >>> from tencirchem import UCCSD + >>> from tencirchem.molecule import h8 + >>> # normal PySCF workflow + >>> hf = h8.HF() + >>> round(hf.kernel(), 8) + -4.14961853 + >>> casscf = CASSCF(hf, 2, 2) + >>> # set the FCI solver for CASSCF to be UCCSD + >>> casscf.fcisolver = UCCSD.as_pyscf_solver() + >>> round(casscf.kernel()[0], 8) + -4.16647335 + """ + + class FakeFCISolver: + def __init__(self): + self.instance: UCC = None + self.config_function = config_function + self.instance_kwargs = kwargs + for arg in ["run_ccsd", "run_fci"]: + # keep MP2 for initial guess + self.instance_kwargs[arg] = False + + def kernel(self, h1, h2, norb, nelec, ci0=None, ecore=0, **kwargs): + self.instance = cls.from_integral(h1, h2, nelec, **self.instance_kwargs) + if self.config_function is not None: + self.config_function(self.instance) + e = self.instance.kernel() + return e + ecore, self.instance.params + + def make_rdm1(self, params, norb, nelec): + civector = self.instance.civector(params) + return self.instance.make_rdm1(civector) + + def make_rdm12(self, params, norb, nelec): + civector = self.instance.civector(params) + rdm1 = self.instance.make_rdm1(civector) + rdm2 = self.instance.make_rdm2(civector) + return rdm1, rdm2 + + def spin_square(self, params, norb, nelec): + return 0, 1 + + return FakeFCISolver() + + def __init__( + self, + mol: Union[Mole, RHF], + init_method="mp2", + active_space=None, + aslst=None, + mo_coeff=None, + mode="fermion", + engine=None, + run_hf=True, + run_mp2=True, + run_ccsd=True, + run_fci=True, + ): + r""" + Initialize the class with molecular input. + + Parameters + ---------- + mol: Mole or RHF + The molecule as PySCF ``Mole`` object or the PySCF ``RHF`` object + init_method: str, optional + How to determine the initial amplitude guess. Accepts ``"mp2"`` (default), ``"ccsd"``, ``"fe"`` + and ``"zeros"``. + active_space: Tuple[int, int], optional + Active space approximation. The first integer is the number of electrons and the second integer is + the number or spatial-orbitals. Defaults to None. + aslst: List[int], optional + Pick orbitals for the active space. Defaults to None which means the orbitals are sorted by energy. + The orbital index is 0-based. + + .. note:: + See `PySCF document `_ + for choosing the active space orbitals. Here orbital index is 0-based, whereas in PySCF by default it + is 1-based. + + mo_coeff: np.ndarray, optional + Molecule coefficients. If provided then RHF is skipped. + Can be used in combination with the ``init_state`` attribute. + Defaults to None which means RHF orbitals are used. + mode: str, optional + How to deal with particle symmetry, such as whether force electrons to pair as hard-core boson (HCB). + Possible values are ``"fermion"``, ``"qubit"`` and ``"hcb"``. + Default to ``"fermion"``. + engine: str, optional + The engine to run the calculation. See :ref:`advanced:Engines` for details. + run_hf: bool, optional + Whether run HF for molecule orbitals. Defaults to ``True``. + The argument has no effect if ``mol`` is a ``RHF`` object. + run_mp2: bool, optional + Whether run MP2 for initial guess and energy reference. Defaults to ``True``. + run_ccsd: bool, optional + Whether run CCSD for initial guess and energy reference. Defaults to ``True``. + run_fci: bool, optional + Whether run FCI for energy reference. Defaults to ``True``. + + See Also + -------- + tencirchem.UCCSD + tencirchem.KUPCCGSD + tencirchem.PUCCD + """ + # process mol + if isinstance(mol, _Molecule): + self.mol = mol + self.mol.verbose = 0 + self.hf: RHF = None + elif isinstance(mol, Mole): + # to set verbose = 0 + self.mol = mol.copy() + # be cautious when modifying mol. Custom mols are common in practice + self.mol.verbose = 0 + self.hf: RHF = None + elif isinstance(mol, RHF_TYPE): + self.hf: RHF = mol + self.mol = self.hf.mol + mol = self.mol + else: + raise TypeError( + f"Unknown input type {type(mol)}. If you're performing open shell calculations, " + "please use ROHF instead." + ) + + if active_space is None: + active_space = (mol.nelectron, int(mol.nao)) + + self.mode = mode + self.spin = self.mol.spin + if mode == "hcb": + assert self.spin == 0 + self.n_qubits = 2 * active_space[1] + if mode == "hcb": + self.n_qubits //= 2 + + # process activate space + self.active_space = active_space + self.n_elec = active_space[0] + self.active = active_space[1] + self.inactive_occ = (mol.nelectron - active_space[0]) // 2 + assert (mol.nelectron - active_space[0]) % 2 == 0 + self.inactive_vir = mol.nao - active_space[1] - self.inactive_occ + if aslst is None: + aslst = list(range(self.inactive_occ, mol.nao - self.inactive_vir)) + if len(aslst) != active_space[1]: + raise ValueError("sort_mo should have the same length as the number of active orbitals.") + frozen_idx = [i for i in range(mol.nao) if i not in aslst] + self.aslst = aslst + + # process backend + self._check_engine(engine) + + if engine is None: + # no need to be too precise + if self.n_qubits <= 16: + engine = "civector" + else: + engine = "civector-large" + self.engine = engine + + # classical quantum chemistry + # hf + if self.hf is not None: + self.e_hf = self.hf.e_tot + self.hf.mo_coeff = canonical_mo_coeff(self.hf.mo_coeff) + elif run_hf: + if self.spin == 0: + self.hf = RHF(self.mol) + else: + self.hf = ROHF(self.mol) + # avoid serialization warnings for `_Molecule` + self.hf.chkfile = None + # run this even when ``mo_coeff is not None`` because MP2 and CCSD + # reference energy might be desired + self.e_hf = self.hf.kernel(dump_chk=False) + self.hf.mo_coeff = canonical_mo_coeff(self.hf.mo_coeff) + else: + self.e_hf = None + # otherwise, can't run casci.get_h2eff() based on HF + self.hf = RHF(self.mol) + self.hf._eri = mol.intor("int2e", aosym="s8") + if mo_coeff is None: + raise ValueError("Must provide MO coefficient if HF is skipped") + + # mp2 + if run_mp2 and not isinstance(self.hf, ROHF_TYPE): + mp2 = self.hf.MP2() + if frozen_idx: + mp2.frozen = frozen_idx + e_corr_mp2, mp2_t2 = mp2.kernel() + self.e_mp2 = self.e_hf + e_corr_mp2 + else: + self.e_mp2 = None + mp2_t2 = None + if init_method is not None and init_method.lower() == "mp2": + raise ValueError("Must run RHF and MP2 to use MP2 as the initial guess method") + + # ccsd + if run_ccsd and not isinstance(self.hf, ROHF_TYPE): + ccsd = self.hf.CCSD() + if frozen_idx: + ccsd.frozen = frozen_idx + e_corr_ccsd, ccsd_t1, ccsd_t2 = ccsd.kernel() + self.e_ccsd = self.e_hf + e_corr_ccsd + else: + self.e_ccsd = None + ccsd_t1 = ccsd_t2 = None + if init_method is not None and init_method.lower() == "ccsd": + raise ValueError("Must run CCSD to use CCSD as the initial guess method") + + # MP2 and CCSD rely on canonical HF orbitals but FCI doesn't + # so set custom mo_coeff after MP2 and CCSD and before FCI + if mo_coeff is not None: + # use user defined coefficient + self.hf.mo_coeff = canonical_mo_coeff(mo_coeff) + + # fci + if run_fci: + fci = CASCI(self.hf, self.active_space[1], self.active_space[0]) + fci.max_memory = 32000 + mo = fci.sort_mo(aslst, base=0) + res = fci.kernel(mo) + self.e_fci = res[0] + self.civector_fci = res[2].ravel() + else: + self.e_fci = None + self.civector_fci = None + + self.e_nuc = mol.energy_nuc() + + # Hamiltonian related + self.hamiltonian_lib = {} + self.int1e = self.int2e = None + # e_core includes nuclear repulsion energy + self.hamiltonian, self.e_core, _ = self._get_hamiltonian_and_core(self.engine) + + # initial guess + self.t1 = np.zeros([self.no, self.nv]) + self.t2 = np.zeros([self.no, self.no, self.nv, self.nv]) + self.init_method = init_method + if init_method is None or init_method in ["zeros", "zero"]: + pass + elif init_method.lower() == "ccsd": + self.t1, self.t2 = ccsd_t1, ccsd_t2 + elif init_method.lower() == "fe": + self.t2 = compute_fe_t2(self.no, self.nv, self.int1e, self.int2e) + elif init_method.lower() == "mp2": + self.t2 = mp2_t2 + else: + raise ValueError(f"Unknown initialization method: {init_method}") + + # circuit related + self._init_state = None + self.ex_ops = None + self._param_ids = None + self.init_guess = None + + # optimization related + self.scipy_minimize_options = None + # optimization result + self.opt_res = None + # for manually set + self._params = None + + def kernel(self) -> float: + """ + The kernel to perform the VQE algorithm. + The L-BFGS-B method in SciPy is used for optimization + and configuration is possible by setting the ``self.scipy_minimize_options`` attribute. + + Returns + ------- + e: float + The optimized energy + """ + assert len(self.param_ids) == len(self.ex_ops) + + energy_and_grad, stating_time = self.get_opt_function(with_time=True) + + if self.init_guess is None: + self.init_guess = np.zeros(self.n_params) + + # optimization options + if self.scipy_minimize_options is None: + # quite strict + options = {"ftol": 1e1 * np.finfo(tq.rdtypestr).eps, "gtol": 1e2 * np.finfo(tq.rdtypestr).eps} + else: + options = self.scipy_minimize_options + + logger.info("Begin optimization") + + time1 = time() + opt_res = minimize(energy_and_grad, x0=self.init_guess, jac=True, method="L-BFGS-B", options=options) + time2 = time() + + if not opt_res.success: + logger.warning("Optimization failed. See `.opt_res` for details.") + + opt_res["staging_time"] = stating_time + opt_res["opt_time"] = time2 - time1 + opt_res["init_guess"] = self.init_guess + opt_res["e"] = float(opt_res.fun) + self.opt_res = opt_res + # prepare for future modification + self.params = opt_res.x.copy() + return opt_res.e + + def get_opt_function(self, with_time: bool = False) -> Union[Callable, Tuple[Callable, float]]: + """ + Returns the cost function in SciPy format for optimization. + The gradient is included. + Basically a wrapper to :func:`energy_and_grad`. + + Parameters + ---------- + with_time: bool, optional + Whether return staging time. Defaults to False. + + Returns + ------- + opt_function: Callable + The optimization cost function in SciPy format. + time: float + Staging time. Returned when ``with_time`` is set to ``True``. + """ + energy_and_grad = scipy_opt_wrap(partial(self.energy_and_grad, engine=self.engine)) + + time1 = time() + if tq.backend.name == "jax": + logger.info("JIT compiling the circuit") + _ = energy_and_grad(np.zeros(self.n_params)) + logger.info("Circuit JIT compiled") + time2 = time() + if with_time: + return energy_and_grad, time2 - time1 + return energy_and_grad + + def _check_params_argument(self, params, strict=True): + if params is None: + if self.params is not None: + params = self.params + else: + if strict: + raise ValueError("Run the `.kernel` method to determine the parameters first") + else: + if self.init_guess is not None: + params = self.init_guess + else: + params = np.zeros(self.n_params) + + if len(params) != self.n_params: + raise ValueError(f"Incompatible parameter shape. {self.n_params} is desired. Got {len(params)}") + return tq.backend.convert_to_tensor(params).astype(tq.rdtypestr) + + def _check_engine(self, engine): + supported_engine = [None, "tensornetwork", "statevector", "civector", "civector-large", "pyscf"] + if not engine in supported_engine: + raise ValueError(f"Engine '{engine}' not supported") + + def _sanity_check(self): + if self.ex_ops is None or self.param_ids is None: + raise ValueError("`ex_ops` or `param_ids` not defined") + if self.param_ids is not None and (len(self.ex_ops) != len(self.param_ids)): + raise ValueError( + f"Excitation operator size {len(self.ex_ops)} and parameter size {len(self.param_ids)} do not match" + ) + + def civector(self, params: Tensor = None, engine: str = None) -> Tensor: + """ + Evaluate the configuration interaction (CI) vector. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + engine: str, optional + The engine to use. Defaults to ``None``, which uses ``self.engine``. + + Returns + ------- + civector: Tensor + Corresponding CI vector + + See Also + -------- + statevector: Evaluate the circuit state vector. + energy: Evaluate the total energy. + energy_and_grad: Evaluate the total energy and parameter gradients. + + Examples + -------- + >>> from tencirchem import UCCSD + >>> from tencirchem.molecule import h2 + >>> uccsd = UCCSD(h2) + >>> uccsd.civector([0, 0]) # HF state + array([1., 0., 0., 0.]) + """ + self._sanity_check() + params = self._check_params_argument(params) + self._check_engine(engine) + if engine is None: + engine = self.engine + civector = get_civector( + params, self.n_qubits, self.n_elec_s, self.ex_ops, self.param_ids, self.mode, self.init_state, engine + ) + return civector + + def get_ci_strings(self, strs2addr: bool = False) -> np.ndarray: + """ + Get the CI bitstrings for all configurations in the CI vector. + + Parameters + ---------- + strs2addr: bool, optional. + Whether return the reversed mapping for one spin sector. Defaults to ``False``. + + Returns + ------- + cistrings: np.ndarray + The CI bitstrings. + string_addr: np.ndarray + The address of the string in one spin sector. Returned when ``strs2addr`` is set to ``True``. + + Examples + -------- + >>> from tencirchem import UCCSD, PUCCD + >>> from tencirchem.molecule import h2 + >>> uccsd = UCCSD(h2) + >>> uccsd.get_ci_strings() + array([ 5, 6, 9, 10], dtype=uint64) + >>> [f"{bin(i)[2:]:0>4}" for i in uccsd.get_ci_strings()] + ['0101', '0110', '1001', '1010'] + >>> uccsd.get_ci_strings(True)[1] # only one spin sector + array([0, 0, 1, 0], dtype=uint64) + """ + return get_ci_strings(self.n_qubits, self.n_elec_s, self.mode, strs2addr=strs2addr) + + # since there's ci_vector method + ci_strings = get_ci_strings + + def get_addr(self, bitstring: str) -> int: + """ + Get the address (index) of a CI bitstring in the CI vector. + + Parameters + ---------- + bitstring: str + The bitstring such as ``"0101"``. + + Returns + ------- + address: int + The bitstring address. + + Examples + -------- + >>> from tencirchem import UCCSD, PUCCD + >>> from tencirchem.molecule import h2 + >>> uccsd = UCCSD(h2) + >>> uccsd.get_addr("0101") # the HF state + 0 + >>> uccsd.get_addr("1010") + 3 + >>> puccd = PUCCD(h2) + >>> puccd.get_addr("01") # the HF state + 0 + >>> puccd.get_addr("10") + 1 + """ + _, strs2addr = self.get_ci_strings(strs2addr=True) + return int(get_addr(int(bitstring, base=2), self.n_qubits, self.n_elec_s, strs2addr, self.mode)) + + def statevector(self, params: Tensor = None, engine: str = None) -> Tensor: + """ + Evaluate the circuit state vector. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + engine: str, optional + The engine to use. Defaults to ``None``, which uses ``self.engine``. + + Returns + ------- + statevector: Tensor + Corresponding state vector + + See Also + -------- + civector: Evaluate the configuration interaction (CI) vector. + energy: Evaluate the total energy. + energy_and_grad: Evaluate the total energy and parameter gradients. + + Examples + -------- + >>> from tencirchem import UCCSD + >>> from tencirchem.molecule import h2 + >>> uccsd = UCCSD(h2) + >>> uccsd.statevector([0, 0]) # HF state + array([0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]) + """ + self._sanity_check() + params = self._check_params_argument(params) + self._check_engine(engine) + if engine is None: + engine = self.engine + statevector = get_statevector( + params, self.n_qubits, self.n_elec_s, self.ex_ops, self.param_ids, self.mode, self.init_state, engine + ) + return statevector + + def _get_hamiltonian_and_core(self, engine): + self._check_engine(engine) + if engine is None: + engine = self.engine + hamiltonian = self.hamiltonian + e_core = self.e_core + else: + if engine.startswith("civector") or engine == "pyscf": + htype = "fcifunc" + else: + assert engine in ["tensornetwork", "statevector"] + htype = "sparse" + hamiltonian = self.hamiltonian_lib.get(htype) + if hamiltonian is None: + if self.int1e is None: + self.int1e, self.int2e, e_core = get_integral_from_hf(self.hf, self.active_space, self.aslst) + else: + e_core = self.e_core + hamiltonian = get_h_from_integral(self.int1e, self.int2e, self.n_elec_s, self.mode, htype) + self.hamiltonian_lib[htype] = hamiltonian + else: + e_core = self.e_core + return hamiltonian, e_core, engine + + def energy(self, params: Tensor = None, engine: str = None) -> float: + """ + Evaluate the total energy. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + engine: str, optional + The engine to use. Defaults to ``None``, which uses ``self.engine``. + + Returns + ------- + energy: float + Total energy + + See Also + -------- + civector: Get the configuration interaction (CI) vector. + statevector: Evaluate the circuit state vector. + energy_and_grad: Evaluate the total energy and parameter gradients. + + Examples + -------- + >>> from tencirchem import UCCSD + >>> from tencirchem.molecule import h2 + >>> uccsd = UCCSD(h2) + >>> round(uccsd.energy([0, 0]), 8) # HF state + -1.11670614 + """ + self._sanity_check() + params = self._check_params_argument(params) + if params is self.params and self.opt_res is not None: + return self.opt_res.e + hamiltonian, _, engine = self._get_hamiltonian_and_core(engine) + e = get_energy( + params, + hamiltonian, + self.n_qubits, + self.n_elec_s, + self.ex_ops, + self.param_ids, + self.mode, + self.init_state, + engine, + ) + return float(e) + self.e_core + + def energy_and_grad(self, params: Tensor = None, engine: str = None) -> Tuple[float, Tensor]: + """ + Evaluate the total energy and parameter gradients. + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter + and :func:`kernel` must be called before. + engine: str, optional + The engine to use. Defaults to ``None``, which uses ``self.engine``. + + Returns + ------- + energy: float + Total energy + grad: Tensor + The parameter gradients + + See Also + -------- + civector: Get the configuration interaction (CI) vector. + statevector: Evaluate the circuit state vector. + energy: Evaluate the total energy. + + Examples + -------- + >>> from tencirchem import UCCSD + >>> from tencirchem.molecule import h2 + >>> uccsd = UCCSD(h2) + >>> e, g = uccsd.energy_and_grad([0, 0]) + >>> round(e, 8) + -1.11670614 + >>> g # doctest:+ELLIPSIS + array([..., ...]) + """ + self._sanity_check() + params = self._check_params_argument(params) + hamiltonian, _, engine = self._get_hamiltonian_and_core(engine) + e, g = get_energy_and_grad( + params, + hamiltonian, + self.n_qubits, + self.n_elec_s, + self.ex_ops, + self.param_ids, + self.mode, + self.init_state, + engine, + ) + return float(e + self.e_core), tq.backend.numpy(g) + + def apply_excitation(self, state: Tensor, ex_op: Tuple, engine: str = None) -> Tensor: + """ + Apply a given excitation operator to a given state. + + Parameters + ---------- + state: Tensor + The input state in statevector or CI vector form. + ex_op: tuple of ints + The excitation operator. + engine: str, optional + The engine to use. Defaults to ``None``, which uses ``self.engine``. + Returns + ------- + tensor: Tensor + The resulting tensor. + + Examples + -------- + >>> from tencirchem import UCC + >>> from tencirchem.molecule import h2 + >>> ucc = UCC(h2) + >>> ucc.apply_excitation([1, 0, 0, 0], (3, 1, 0, 2)) + array([0, 0, 0, 1]) + """ + self._check_engine(engine) + if engine is None: + engine = self.engine + return apply_excitation(state, self.n_qubits, self.n_elec_s, ex_op, mode=self.mode, engine=engine) + + def _statevector_to_civector(self, statevector=None): + if statevector is None: + civector = self.civector() + else: + if len(statevector) == self.statevector_size: + ci_strings = self.get_ci_strings() + civector = statevector[ci_strings] + else: + if len(statevector) == self.civector_size: + civector = statevector + else: + raise ValueError(f"Incompatible statevector size: {len(statevector)}") + + civector = tq.backend.numpy(tq.backend.convert_to_tensor(civector)) + return civector + + def make_rdm1(self, statevector: Tensor = None, basis: str = "AO") -> np.ndarray: + r""" + Evaluate the spin-traced one-body reduced density matrix (1RDM). + + .. math:: + + \textrm{1RDM}[p,q] = \langle p_{\alpha}^\dagger q_{\alpha} \rangle + + \langle p_{\beta}^\dagger q_{\beta} \rangle + + If active space approximation is employed, returns the full RDM of all orbitals. + + Parameters + ---------- + statevector: Tensor, optional + Custom system state. Could be CI vector or state vector. + Defaults to None, which uses the optimized state by :func:`civector`. + + basis: str, optional + One of ``"AO"`` or ``"MO"``. Defaults to ``"AO"``, which is for consistency with PySCF. + + Returns + ------- + rdm1: np.ndarray + The spin-traced one-body RDM. + + See Also + -------- + make_rdm2: Evaluate the spin-traced two-body reduced density matrix (2RDM). + """ + assert self.mode in ["fermion", "qubit"] + civector = self._statevector_to_civector(statevector).astype(np.float64) + + rdm1_cas = fci.direct_spin1.make_rdm1(civector, self.n_qubits // 2, self.n_elec_s) + + rdm1 = self.embed_rdm_cas(rdm1_cas) + + if basis == "MO": + return rdm1 + else: + return rdm_mo2ao(rdm1, self.hf.mo_coeff) + + def make_rdm2(self, statevector: Tensor = None, basis: str = "AO") -> np.ndarray: + r""" + Evaluate the spin-traced two-body reduced density matrix (2RDM). + + .. math:: + + \begin{aligned} + \textrm{2RDM}[p,q,r,s] & = \langle p_{\alpha}^\dagger r_{\alpha}^\dagger + s_{\alpha} q_{\alpha} \rangle + + \langle p_{\beta}^\dagger r_{\alpha}^\dagger s_{\alpha} q_{\beta} \rangle \\ + & \quad + \langle p_{\alpha}^\dagger r_{\beta}^\dagger s_{\beta} q_{\alpha} \rangle + + \langle p_{\beta}^\dagger r_{\beta}^\dagger s_{\beta} q_{\beta} \rangle + \end{aligned} + + If active space approximation is employed, returns the full RDM of all orbitals. + + Parameters + ---------- + statevector: Tensor, optional + Custom system state. Could be CI vector or state vector. + Defaults to None, which uses the optimized state by :func:`civector`. + + basis: str, optional + One of ``"AO"`` or ``"MO"``. Defaults to ``"AO"``, which is for consistency with PySCF. + + Returns + ------- + rdm2: np.ndarray + The spin-traced two-body RDM. + + See Also + -------- + make_rdm1: Evaluate the spin-traced one-body reduced density matrix (1RDM). + + Examples + -------- + >>> import numpy as np + >>> from tencirchem import UCC + >>> from tencirchem.molecule import h2 + >>> ucc = UCC(h2) + >>> state = [1, 0, 0, 0] ## HF state + >>> rdm1 = ucc.make_rdm1(state, basis="MO") + >>> rdm2 = ucc.make_rdm2(state, basis="MO") + >>> e_hf = ucc.int1e.ravel() @ rdm1.ravel() + 1/2 * ucc.int2e.ravel() @ rdm2.ravel() + >>> np.testing.assert_allclose(e_hf + ucc.e_nuc, ucc.e_hf, atol=1e-10) + """ + assert self.mode in ["fermion", "qubit"] + civector = self._statevector_to_civector(statevector).astype(np.float64) + + rdm2_cas = fci.direct_spin1.make_rdm12(civector.astype(np.float64), self.n_qubits // 2, self.n_elec_s)[1] + + rdm2 = self.embed_rdm_cas(rdm2_cas) + + if basis == "MO": + return rdm2 + else: + return rdm_mo2ao(rdm2, self.hf.mo_coeff) + + def embed_rdm_cas(self, rdm_cas): + """ + Embed CAS RDM into RDM of the whole system + """ + if self.inactive_occ == 0 and self.inactive_vir == 0: + # active space approximation not employed + return rdm_cas + # slice of indices in rdm corresponding to cas + slice_cas = slice(self.inactive_occ, self.inactive_occ + len(rdm_cas)) + if rdm_cas.ndim == 2: + rdm1_cas = rdm_cas + rdm1 = np.zeros((self.mol.nao, self.mol.nao)) + for i in range(self.inactive_occ): + rdm1[i, i] = 2 + rdm1[slice_cas, slice_cas] = rdm1_cas + return rdm1 + else: + rdm2_cas = rdm_cas + # active space approximation employed + rdm1 = self.make_rdm1(basis="MO") + rdm1_cas = rdm1[slice_cas, slice_cas] + rdm2 = np.zeros((self.mol.nao, self.mol.nao, self.mol.nao, self.mol.nao)) + rdm2[slice_cas, slice_cas, slice_cas, slice_cas] = rdm2_cas + for i in range(self.inactive_occ): + for j in range(self.inactive_occ): + rdm2[i, i, j, j] += 4 + rdm2[i, j, j, i] -= 2 + rdm2[i, i, slice_cas, slice_cas] = rdm2[slice_cas, slice_cas, i, i] = 2 * rdm1_cas + rdm2[i, slice_cas, slice_cas, i] = rdm2[slice_cas, i, i, slice_cas] = -rdm1_cas + return rdm2 + + def get_ex_ops(self, t1: np.ndarray = None, t2: np.ndarray = None): + """Virtual method to be implemented""" + raise NotImplementedError + + def get_ex1_ops(self, t1: np.ndarray = None) -> Tuple[List[Tuple], List[int], List[float]]: + """ + Get one-body excitation operators. + + Parameters + ---------- + t1: np.ndarray, optional + Initial one-body amplitudes based on e.g. CCSD + + Returns + ------- + ex_op: List[Tuple] + The excitation operators. Each operator is represented by a tuple of ints. + param_ids: List[int] + The mapping from excitations to parameters. + init_guess: List[float] + The initial guess for the parameters. + + See Also + -------- + get_ex2_ops: Get two-body excitation operators. + get_ex_ops: Get one-body and two-body excitation operators for UCCSD ansatz. + """ + # single excitations + no, nv = self.no, self.nv + if t1 is None: + t1 = self.t1 + + if t1.shape == (self.no, self.nv): + t1 = spatial2spin(t1) + else: + assert t1.shape == (2 * self.no, 2 * self.nv) + + ex1_ops = [] + # unique parameters. -1 is a place holder + ex1_param_ids = [-1] + ex1_init_guess = [] + for i in range(no): + for a in range(nv): + # alpha to alpha + ex_op_a = (2 * no + nv + a, no + nv + i) + # beta to beta + ex_op_b = (no + a, i) + ex1_ops.extend([ex_op_a, ex_op_b]) + ex1_param_ids.extend([ex1_param_ids[-1] + 1] * 2) + ex1_init_guess.append(t1[i, a]) + ex1_param_ids = ex1_param_ids[1:] + + # deal with qubit symmetry + if self.mode == "qubit": + ex1_ops, ex1_param_ids, ex1_init_guess = self._qubit_phase(ex1_ops, ex1_param_ids, ex1_init_guess) + + return ex1_ops, ex1_param_ids, ex1_init_guess + + def get_ex2_ops(self, t2: np.ndarray = None) -> Tuple[List[Tuple], List[int], List[float]]: + """ + Get two-body excitation operators. + + Parameters + ---------- + t2: np.ndarray, optional + Initial two-body amplitudes based on e.g. MP2 + + Returns + ------- + ex_op: List[Tuple] + The excitation operators. Each operator is represented by a tuple of ints. + param_ids: List[int] + The mapping from excitations to parameters. + init_guess: List[float] + The initial guess for the parameters. + + See Also + -------- + get_ex1_ops: Get one-body excitation operators. + get_ex_ops: Get one-body and two-body excitation operators for UCCSD ansatz. + """ + + # t2 in oovv 1212 format + no, nv = self.no, self.nv + if t2 is None: + t2 = self.t2 + + if t2.shape == (self.no, self.no, self.nv, self.nv): + t2 = spatial2spin(t2) + else: + assert t2.shape == (2 * self.no, 2 * self.no, 2 * self.nv, 2 * self.nv) + + def alpha_o(_i): + return no + nv + _i + + def alpha_v(_i): + return 2 * no + nv + _i + + def beta_o(_i): + return _i + + def beta_v(_i): + return no + _i + + # double excitations + ex_ops = [] + ex2_param_ids = [-1] + ex2_init_guess = [] + # 2 alphas or 2 betas + for i in range(no): + for j in range(i): + for a in range(nv): + for b in range(a): + # i correspond to a and j correspond to b, as in PySCF convention + # otherwise the t2 amplitude has incorrect phase + # 2 alphas + ex_op_aa = (alpha_v(b), alpha_v(a), alpha_o(i), alpha_o(j)) + # 2 betas + ex_op_bb = (beta_v(b), beta_v(a), beta_o(i), beta_o(j)) + ex_ops.extend([ex_op_aa, ex_op_bb]) + ex2_param_ids.extend([ex2_param_ids[-1] + 1] * 2) + ex2_init_guess.append(t2[2 * i, 2 * j, 2 * a, 2 * b]) + assert len(ex_ops) == 2 * (no * (no - 1) / 2) * (nv * (nv - 1) / 2) + # 1 alpha + 1 beta + for i in range(no): + for j in range(i + 1): + for a in range(nv): + for b in range(a + 1): + # i correspond to a and j correspond to b, as in PySCF convention + # otherwise the t2 amplitude has incorrect phase + if i == j and a == b: + # paired + ex_op_ab = (beta_v(a), alpha_v(a), alpha_o(i), beta_o(i)) + ex_ops.append(ex_op_ab) + ex2_param_ids.append(ex2_param_ids[-1] + 1) + ex2_init_guess.append(t2[2 * i, 2 * i + 1, 2 * a, 2 * a + 1]) + continue + # simple reflection + ex_op_ab1 = (beta_v(b), alpha_v(a), alpha_o(i), beta_o(j)) + ex_op_ab2 = (alpha_v(b), beta_v(a), beta_o(i), alpha_o(j)) + ex_ops.extend([ex_op_ab1, ex_op_ab2]) + ex2_param_ids.extend([ex2_param_ids[-1] + 1] * 2) + ex2_init_guess.append(t2[2 * i, 2 * j + 1, 2 * a, 2 * b + 1]) + if (i != j) and (a != b): + # exchange alpha and beta + ex_op_ab3 = (beta_v(a), alpha_v(b), alpha_o(i), beta_o(j)) + ex_op_ab4 = (alpha_v(a), beta_v(b), beta_o(i), alpha_o(j)) + ex_ops.extend([ex_op_ab3, ex_op_ab4]) + ex2_param_ids.extend([ex2_param_ids[-1] + 1] * 2) + ex2_init_guess.append(t2[2 * i, 2 * j + 1, 2 * b, 2 * a + 1]) + ex2_param_ids = ex2_param_ids[1:] + + # deal with qubit symmetry + if self.mode == "qubit": + ex_ops, ex2_param_ids, ex2_init_guess = self._qubit_phase(ex_ops, ex2_param_ids, ex2_init_guess) + + return ex_ops, ex2_param_ids, ex2_init_guess + + def _qubit_phase(self, ex_ops, ex_param_ids, ex_init_guess): + + hf_str = np.array(self.get_ci_strings()[:1], dtype=np.uint64) + iterated_ids = set() + for i, ex_op in enumerate(ex_ops): + if ex_param_ids[i] in iterated_ids: + continue + phase = get_fermion_phase(ex_op, self.n_qubits, hf_str)[0] + ex_init_guess[ex_param_ids[i]] *= phase + iterated_ids.add(ex_param_ids[i]) + return ex_ops, ex_param_ids, ex_init_guess + + @property + def e_ucc(self) -> float: + """ + Returns UCC energy + """ + return self.energy() + + def print_ansatz(self): + df_dict = { + "#qubits": [self.n_qubits], + "#params": [self.n_params], + "#excitations": [len(self.ex_ops)], + } + if self.init_state is None: + df_dict["initial condition"] = "RHF" + else: + df_dict["initial condition"] = "custom" + print(pd.DataFrame(df_dict).to_string(index=False)) + + def get_circuit( + self, params: Tensor = None, decompose_multicontrol: bool = False, trotter: bool = False + ) -> tq.Circuit: + """ + Get the circuit as TyxonQ ``Circuit`` object + + Parameters + ---------- + params: Tensor, optional + The circuit parameters. Defaults to None, which uses the optimized parameter. + If :func:`kernel` is not called before, the initial guess is used. + decompose_multicontrol: bool, optional + Whether decompose the Multicontrol gate in the circuit into CNOT gates. + Defaults to False. + trotter: bool, optional + Whether Trotterize the UCC factor into Pauli strings. + Defaults to False. + + Returns + ------- + circuit: :class:`tq.Circuit` + The quantum circuit. + """ + if self.ex_ops is None: + raise ValueError("Excitation operators not defined") + params = self._check_params_argument(params, strict=False) + return get_circuit( + params, + self.n_qubits, + self.n_elec_s, + self.ex_ops, + self.param_ids, + self.mode, + self.init_state, + decompose_multicontrol=decompose_multicontrol, + trotter=trotter, + ) + + def print_circuit(self): + """ + Prints the circuit information. If you wish to print the circuit diagram, + use :func:`get_circuit` and then call ``draw()`` such as ``print(ucc.get_circuit().draw())``. + """ + c = self.get_circuit() + df = get_circuit_dataframe(c) + + def format_flop(f): + return f"{f:.3e}" + + formatters = {"flop": format_flop} + print(df.to_string(index=False, formatters=formatters)) + + def get_init_state_dataframe(self, coeff_epsilon: float = DISCARD_EPS) -> pd.DataFrame: + """ + Returns initial state information dataframe. + + Parameters + ---------- + coeff_epsilon: float, optional + The threshold to screen out states with small coefficients. + Defaults to 1e-12. + + Returns + ------- + pd.DataFrame + + See Also + -------- + init_state: The circuit initial state before applying the excitation operators. + + Examples + -------- + >>> from tencirchem import UCC + >>> from tencirchem.molecule import h2 + >>> ucc = UCC(h2) + >>> ucc.init_state = [0.707, 0, 0, 0.707] + >>> ucc.get_init_state_dataframe() # doctest: +NORMALIZE_WHITESPACE + configuration coefficient + 0 0101 0.707 + 1 1010 0.707 + """ + columns = ["configuration", "coefficient"] + if self.init_state is None: + init_state = get_init_civector(self.civector_size) + else: + init_state = self.init_state + ci_strings = self.get_ci_strings() + ci_coeffs = translate_init_state(init_state, self.n_qubits, ci_strings) + data_list = [] + for ci_string, coeff in zip(ci_strings, ci_coeffs): + if np.abs(coeff) < coeff_epsilon: + continue + ci_string = bin(ci_string)[2:] + ci_string = "0" * (self.n_qubits - len(ci_string)) + ci_string + data_list.append((ci_string, coeff)) + return pd.DataFrame(data_list, columns=columns) + + def print_init_state(self): + print(self.get_init_state_dataframe().to_string()) + + def get_excitation_dataframe(self) -> pd.DataFrame: + columns = ["excitation", "configuration", "parameter", "initial guess"] + if self.ex_ops is None: + return pd.DataFrame(columns=columns) + + if self.params is None: + # optimization not done + params = [None] * len(self.init_guess) + else: + params = self.params + + if self.param_ids is None: + # see self.n_params + param_ids = range(len(self.ex_ops)) + else: + param_ids = self.param_ids + + data_list = [] + + for i, ex_op in zip(param_ids, self.ex_ops): + bitstring = get_ex_bitstring(self.n_qubits, self.n_elec_s, ex_op, self.mode) + data_list.append((ex_op, bitstring, params[i], self.init_guess[i])) + return pd.DataFrame(data_list, columns=columns) + + def print_excitations(self): + print(self.get_excitation_dataframe().to_string()) + + def get_energy_dataframe(self) -> pd.DataFrame: + """ + Returns energy information dataframe + """ + if self.params is None: + series_dict = {"HF": self.e_hf, "MP2": self.e_mp2, "CCSD": self.e_ccsd, "FCI": self.e_fci} + else: + ucc_name = self.__class__.__name__ + series_dict = { + "HF": self.e_hf, + "MP2": self.e_mp2, + "CCSD": self.e_ccsd, + ucc_name: self.energy(), + "FCI": self.e_fci, + } + df = pd.DataFrame() + energy = pd.Series(series_dict) + df["energy (Hartree)"] = energy + if self.e_fci is not None: + df["error (mH)"] = 1e3 * (energy - self.e_fci) + df["correlation energy (%)"] = 100 * (energy - self.e_hf) / (self.e_fci - self.e_hf) + return df + + def print_energy(self): + df = self.get_energy_dataframe() + + def format_ce(f): + return f"{f:.3f}" + + formatters = {"correlation energy (%)": format_ce} + print(df.to_string(index=True, formatters=formatters)) + + def print_summary(self, include_circuit: bool = False): + """ + Print a summary of the class. + + Parameters + ---------- + include_circuit: bool + Whether include the circuit section. + + """ + print("################################ Ansatz ###############################") + self.print_ansatz() + if self.init_state is not None: + print("############################ Initial Condition ########################") + self.print_init_state() + if include_circuit: + print("############################### Circuit ###############################") + self.print_circuit() + print("############################### Energy ################################") + self.print_energy() + print("############################# Excitations #############################") + self.print_excitations() + print("######################### Optimization Result #########################") + if self.opt_res is None: + print("Optimization not run (.opt_res is None)") + else: + print(self.opt_res) + + @property + def n_elec_s(self): + """The number of electrons for alpha and beta spin""" + return (self.n_elec + self.spin) // 2, (self.n_elec - self.spin) // 2 + + @property + def no(self) -> int: + """The number of occupied orbitals.""" + return self.n_elec // 2 + + @property + def nv(self) -> int: + """The number of virtual (unoccupied orbitals).""" + return self.active - self.no + + @property + def h_fermion_op(self) -> FermionOperator: + """ + Hamiltonian as openfermion.FermionOperator + """ + if self.mode == "hcb": + raise ValueError("No FermionOperator available for hard-core boson Hamiltonian") + return get_hop_from_integral(self.int1e, self.int2e) + self.e_core + + @property + def h_qubit_op(self) -> QubitOperator: + """ + Hamiltonian as openfermion.QubitOperator, mapped by + Jordan-Wigner transformation. + """ + if self.mode in ["fermion", "qubit"]: + return reverse_qop_idx(jordan_wigner(self.h_fermion_op), self.n_qubits) + else: + assert self.mode == "hcb" + return get_hop_hcb_from_integral(self.int1e, self.int2e) + self.e_core + + @property + def n_params(self) -> int: + """The number of parameter in the ansatz/circuit.""" + # this definition ensures that `param[param_id]` is always valid + if not self.param_ids: + return 0 + return max(self.param_ids) + 1 + + @property + def statevector_size(self) -> int: + """The size of the statevector.""" + return 1 << self.n_qubits + + @property + def civector_size(self) -> int: + """ + The size of the CI vector. + """ + if self.mode in ["fermion", "qubit"]: + na, nb = self.n_elec_s + return round(comb(self.n_qubits // 2, na)) * round(comb(self.n_qubits // 2, nb)) + else: + assert self.mode == "hcb" + return round(comb(self.n_qubits, self.n_elec // 2)) + + @property + def init_state(self) -> Tensor: + """ + The circuit initial state before applying the excitation operators. Usually RHF. + + See Also + -------- + get_init_state_dataframe: Returns initial state information dataframe. + """ + return self._init_state + + @init_state.setter + def init_state(self, init_state): + self._init_state = init_state + + @property + def params(self) -> Tensor: + """The circuit parameters.""" + if self._params is not None: + return self._params + if self.opt_res is not None: + return self.opt_res.x + return None + + @params.setter + def params(self, params): + self._params = params + + @property + def param_ids(self) -> List[int]: + """The mapping from excitations operators to parameters.""" + if self._param_ids is None: + if self.ex_ops is None: + raise ValueError("Excitation operators not defined") + else: + return tuple(range(len(self.ex_ops))) + return self._param_ids + + @param_ids.setter + def param_ids(self, v): + self._param_ids = v + + @property + def param_to_ex_ops(self): + d = defaultdict(list) + for i, j in enumerate(self.param_ids): + d[j].append(self.ex_ops[i]) + return d + + +def compute_fe_t2(no, nv, int1e, int2e): + n_orb = no + nv + + def translate_o(n): + if n % 2 == 0: + return n // 2 + n_orb + else: + return n // 2 + + def translate_v(n): + if n % 2 == 0: + return n // 2 + no + n_orb + else: + return n // 2 + no + + t2 = np.zeros((2 * no, 2 * no, 2 * nv, 2 * nv)) + for i, j, k, l in product(range(2 * no), range(2 * no), range(2 * nv), range(2 * nv)): + # spin not conserved + if i % 2 != k % 2 or j % 2 != l % 2: + continue + a = translate_o(i) + b = translate_o(j) + s = translate_v(l) + r = translate_v(k) + if len(set([a, b, s, r])) != 4: + continue + # r^ s^ b a + rr, ss, bb, aa = [i % n_orb for i in [r, s, b, a]] + if (r < n_orb and s < n_orb) or (r >= n_orb and s >= n_orb): + e_inter = int2e[aa, rr, bb, ss] - int2e[aa, ss, bb, rr] + else: + e_inter = int2e[aa, rr, bb, ss] + if np.allclose(e_inter, 0): + continue + e_diff = _compute_e_diff(r, s, b, a, int1e, int2e, n_orb, no) + if np.allclose(e_diff, 0): + raise RuntimeError("RHF degenerate ground state") + theta = np.arctan(-2 * e_inter / e_diff) / 2 + t2[i, j, k, l] = theta + return t2 + + +def _compute_e_diff(r, s, b, a, int1e, int2e, n_orb, no): + inert_a = list(range(no)) + inert_b = list(range(no)) + old_a = [] + old_b = [] + for i in [b, a]: + if i < n_orb: + inert_b.remove(i) + old_b.append(i) + else: + inert_a.remove(i % n_orb) + old_a.append(i % n_orb) + + new_a = [] + new_b = [] + for i in [r, s]: + if i < n_orb: + new_b.append(i) + else: + new_a.append(i % n_orb) + + diag1e = np.diag(int1e) + diagj = np.einsum("iijj->ij", int2e) + diagk = np.einsum("ijji->ij", int2e) + + e_diff_1e = diag1e[new_a].sum() + diag1e[new_b].sum() - diag1e[old_a].sum() - diag1e[old_b].sum() + # fmt: off + e_diff_j = _compute_j_outer(diagj, inert_a, inert_b, new_a, new_b) \ + - _compute_j_outer(diagj, inert_a, inert_b, old_a, old_b) + e_diff_k = _compute_k_outer(diagk, inert_a, inert_b, new_a, new_b) \ + - _compute_k_outer(diagk, inert_a, inert_b, old_a, old_b) + # fmt: on + return e_diff_1e + 1 / 2 * (e_diff_j - e_diff_k) + + +def _compute_j_outer(diagj, inert_a, inert_b, outer_a, outer_b): + # fmt: off + v = diagj[inert_a][:, outer_a].sum() + diagj[outer_a][:, inert_a].sum() + diagj[outer_a][:, outer_a].sum() \ + + diagj[inert_a][:, outer_b].sum() + diagj[outer_a][:, inert_b].sum() + diagj[outer_a][:, outer_b].sum() \ + + diagj[inert_b][:, outer_a].sum() + diagj[outer_b][:, inert_a].sum() + diagj[outer_b][:, outer_a].sum() \ + + diagj[inert_b][:, outer_b].sum() + diagj[outer_b][:, inert_b].sum() + diagj[outer_b][:, outer_b].sum() + # fmt: on + return v + + +def _compute_k_outer(diagk, inert_a, inert_b, outer_a, outer_b): + # fmt: off + v = diagk[inert_a][:, outer_a].sum() + diagk[outer_a][:, inert_a].sum() + diagk[outer_a][:, outer_a].sum() \ + + diagk[inert_b][:, outer_b].sum() + diagk[outer_b][:, inert_b].sum() + diagk[outer_b][:, outer_b].sum() + # fmt: on + return v diff --git a/src/tyxonq/chem/static/uccsd.py b/src/tyxonq/chem/static/uccsd.py new file mode 100644 index 0000000..dc92c8c --- /dev/null +++ b/src/tyxonq/chem/static/uccsd.py @@ -0,0 +1,303 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from typing import Tuple, List, Union + +import numpy as np +from pyscf.gto.mole import Mole +from pyscf.scf import RHF +from pyscf.scf import ROHF + +from ..static.ucc import UCC +from ..constants import DISCARD_EPS + + +class UCCSD(UCC): + """ + Run UCCSD calculation. For a comprehensive tutorial see :doc:`/tutorial_jupyter/ucc_functions`. + + Examples + -------- + >>> import numpy as np + >>> from tencirchem import UCCSD + >>> from tencirchem.molecule import h2 + >>> uccsd = UCCSD(h2) + >>> e_ucc = uccsd.kernel() + >>> np.testing.assert_allclose(e_ucc, uccsd.e_fci, atol=1e-10) + >>> e_hf = uccsd.energy(np.zeros(uccsd.n_params)) + >>> np.testing.assert_allclose(e_hf, uccsd.e_hf, atol=1e-10) + """ + + def __init__( + self, + mol: Union[Mole, RHF], + init_method: str = "mp2", + active_space: Tuple[int, int] = None, + aslst: List[int] = None, + mo_coeff: np.ndarray = None, + pick_ex2: bool = True, + epsilon: float = DISCARD_EPS, + sort_ex2: bool = True, + mode: str = "fermion", + engine: str = None, + run_hf: bool = True, + run_mp2: bool = True, + run_ccsd: bool = True, + run_fci: bool = True, + ): + r""" + Initialize the class with molecular input. + + Parameters + ---------- + mol: Mole or RHF + The molecule as PySCF ``Mole`` object or the PySCF ``RHF`` object + init_method: str, optional + How to determine the initial amplitude guess. Accepts ``"mp2"`` (default), ``"ccsd"``,``"fe"`` + and ``"zeros"``. + active_space: Tuple[int, int], optional + Active space approximation. The first integer is the number of electrons and the second integer is + the number or spatial-orbitals. Defaults to None. + aslst: List[int], optional + Pick orbitals for the active space. Defaults to None which means the orbitals are sorted by energy. + The orbital index is 0-based. + + .. note:: + See `PySCF document `_ + for choosing the active space orbitals. Here orbital index is 0-based, whereas in PySCF by default it + is 1-based. + + mo_coeff: np.ndarray, optional + Molecule coefficients. If provided then RHF is skipped. + Can be used in combination with the ``init_state`` attribute. + Defaults to None which means RHF orbitals are used. + pick_ex2: bool, optional + Whether screen out two body excitations based on the inital guess amplitude. + Defaults to True, which means excitations with amplitude less than ``epsilon`` (see below) are discarded. + The argument will be set to ``False`` if initial guesses are set to zero. + mode: str, optional + How to deal with particle symmetry. Possible values are ``"fermion"``, ``"qubit"``. + Default to ``"fermion"``. + epsilon: float, optional + The threshold to discard two body excitations. Defaults to 1e-12. + sort_ex2: bool, optional + Whether sort two-body excitations in the ansatz based on the initial guess amplitude. + Large excitations come first. Defaults to True. + Note this could lead to different ansatz for the same molecule at different geometry. + The argument will be set to ``False`` if initial guesses are set to zero. + engine: str, optional + The engine to run the calculation. See :ref:`advanced:Engines` for details. + run_hf: bool, optional + Whether run HF for molecule orbitals. Defaults to ``True``. + run_mp2: bool, optional + Whether run MP2 for initial guess and energy reference. Defaults to ``True``. + run_ccsd: bool, optional + Whether run CCSD for initial guess and energy reference. Defaults to ``True``. + run_fci: bool, optional + Whether run FCI for energy reference. Defaults to ``True``. + + See Also + -------- + tencirchem.KUPCCGSD + tencirchem.PUCCD + tencirchem.UCC + """ + super().__init__( + mol, + init_method, + active_space, + aslst, + mo_coeff, + mode=mode, + engine=engine, + run_hf=run_hf, + run_mp2=run_mp2, + run_ccsd=run_ccsd, + run_fci=run_fci, + ) + if self.init_method == "zeros": + self.pick_ex2 = self.sort_ex2 = False + else: + self.pick_ex2 = pick_ex2 + self.sort_ex2 = sort_ex2 + # screen out excitation operators based on t2 amplitude + self.t2_discard_eps = epsilon + self.ex_ops, self.param_ids, self.init_guess = self.get_ex_ops(self.t1, self.t2) + + def get_ex_ops(self, t1: np.ndarray = None, t2: np.ndarray = None) -> Tuple[List[Tuple], List[int], List[float]]: + """ + Get one-body and two-body excitation operators for UCCSD ansatz. + Pick and sort two-body operators if ``self.pick_ex2`` and ``self.sort_ex2`` are set to ``True``. + + Parameters + ---------- + t1: np.ndarray, optional + Initial one-body amplitudes based on e.g. CCSD + t2: np.ndarray, optional + Initial two-body amplitudes based on e.g. MP2 + + Returns + ------- + ex_op: List[Tuple] + The excitation operators. Each operator is represented by a tuple of ints. + param_ids: List[int] + The mapping from excitations to parameters. + init_guess: List[float] + The initial guess for the parameters. + + See Also + -------- + get_ex1_ops: Get one-body excitation operators. + get_ex2_ops: Get two-body excitation operators. + + Examples + -------- + >>> from tencirchem import UCCSD + >>> from tencirchem.molecule import h2 + >>> uccsd = UCCSD(h2) + >>> ex_op, param_ids, init_guess = uccsd.get_ex_ops() + >>> ex_op + [(3, 2), (1, 0), (1, 3, 2, 0)] + >>> param_ids + [0, 0, 1] + >>> init_guess # doctest:+ELLIPSIS + [0.0, ...] + """ + ex1_ops, ex1_param_ids, ex1_init_guess = self.get_ex1_ops(t1) + ex2_ops, ex2_param_ids, ex2_init_guess = self.get_ex2_ops(t2) + + # screen out symmetrically not allowed excitation + ex2_ops, ex2_param_ids, ex2_init_guess = self.pick_and_sort( + ex2_ops, ex2_param_ids, ex2_init_guess, self.pick_ex2, self.sort_ex2 + ) + + ex_op = ex1_ops + ex2_ops + param_ids = ex1_param_ids + [i + max(ex1_param_ids) + 1 for i in ex2_param_ids] + init_guess = ex1_init_guess + ex2_init_guess + return ex_op, param_ids, init_guess + + def pick_and_sort(self, ex_ops, param_ids, init_guess, do_pick=True, do_sort=True): + # sort operators according to amplitude + if do_sort: + sorted_ex_ops = sorted(zip(ex_ops, param_ids), key=lambda x: -np.abs(init_guess[x[1]])) + else: + sorted_ex_ops = list(zip(ex_ops, param_ids)) + ret_ex_ops = [] + ret_param_ids = [] + for ex_op, param_id in sorted_ex_ops: + # discard operators with tiny amplitude. + # The default eps is so small that the screened out excitations are probably not allowed + if do_pick and np.abs(init_guess[param_id]) < self.t2_discard_eps: + continue + ret_ex_ops.append(ex_op) + ret_param_ids.append(param_id) + assert len(ret_ex_ops) != 0 + unique_ids = np.unique(ret_param_ids) + ret_init_guess = np.array(init_guess)[unique_ids] + id_mapping = {old: new for new, old in enumerate(unique_ids)} + ret_param_ids = [id_mapping[i] for i in ret_param_ids] + return ret_ex_ops, ret_param_ids, list(ret_init_guess) + + @property + def e_uccsd(self) -> float: + """ + Returns UCCSD energy + """ + return self.energy() + + +class ROUCCSD(UCC): + def __init__( + self, + mol: Union[Mole, ROHF], + active_space: Tuple[int, int] = None, + aslst: List[int] = None, + mo_coeff: np.ndarray = None, + engine: str = "civector", + run_hf: bool = True, + # for API consistency with UCC + run_mp2: bool = False, + run_ccsd: bool = False, + run_fci: bool = True, + ): + init_method: str = "zeros" + # ROHF does not support mp2 and ccsd + run_mp2: bool = False + run_ccsd: bool = False + + super().__init__( + mol, + init_method, + active_space, + aslst, + mo_coeff, + engine=engine, + run_hf=run_hf, + run_mp2=run_mp2, + run_ccsd=run_ccsd, + run_fci=run_fci, + ) + no = int(np.sum(self.hf.mo_occ == 2)) - self.inactive_occ + ns = int(np.sum(self.hf.mo_occ == 1)) + nv = int(np.sum(self.hf.mo_occ == 0)) - self.inactive_vir + assert no + ns + nv == self.active_space[1] + # assuming single electrons in alpha + noa = no + ns + nva = nv + nob = no + nvb = ns + nv + + def alpha_o(_i): + return self.active_space[1] + _i + + def alpha_v(_i): + return self.active_space[1] + noa + _i + + def beta_o(_i): + return _i + + def beta_v(_i): + return nob + _i + + # single excitations + self.ex_ops = [] + for i in range(noa): + for a in range(nva): + # alpha to alpha + ex_op_a = (alpha_v(a), alpha_o(i)) + self.ex_ops.append(ex_op_a) + for i in range(nob): + for a in range(nvb): + # beta to beta + ex_op_b = (beta_v(a), beta_o(i)) + self.ex_ops.append(ex_op_b) + + # double excitations + # 2 alphas + for i in range(noa): + for j in range(i): + for a in range(nva): + for b in range(a): + ex_op_aa = (alpha_v(b), alpha_v(a), alpha_o(i), alpha_o(j)) + self.ex_ops.append(ex_op_aa) + # 2 betas + for i in range(nob): + for j in range(i): + for a in range(nvb): + for b in range(a): + ex_op_bb = (beta_v(b), beta_v(a), beta_o(i), beta_o(j)) + self.ex_ops.append(ex_op_bb) + + # 1 alpha + 1 beta + for i in range(noa): + for j in range(nob): + for a in range(nva): + for b in range(nvb): + ex_op_ab = (beta_v(b), alpha_v(a), alpha_o(i), beta_o(j)) + self.ex_ops.append(ex_op_ab) + + self.param_ids = list(range(len(self.ex_ops))) + self.init_guess = np.zeros_like(self.param_ids) diff --git a/src/tyxonq/chem/utils/__init__.py b/src/tyxonq/chem/utils/__init__.py new file mode 100644 index 0000000..15a6721 --- /dev/null +++ b/src/tyxonq/chem/utils/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from ..utils.misc import scipy_opt_wrap diff --git a/src/tyxonq/chem/utils/backend.py b/src/tyxonq/chem/utils/backend.py new file mode 100644 index 0000000..0769892 --- /dev/null +++ b/src/tyxonq/chem/utils/backend.py @@ -0,0 +1,130 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from typing import Any, List, Dict +from functools import wraps +import logging + +import numpy as np +import tyxonq as tq + +try: + import jax + + JAXIMPORTERROR = None +except ImportError as e: + jax = None + JAXIMPORTERROR = e + +try: + import cupy as cp +except ImportError: + cp = None + +Tensor = Any + +logger = logging.getLogger(__name__) + + +ALL_JIT_LIBS: List[Dict] = [] + + +def jit(f, static_argnums=None): + backend_jit_lib = {} + + @wraps(f) + def _wrap_jit(*args, **kwargs): + if tq.backend.name not in backend_jit_lib: + backend_jit_lib[tq.backend.name] = tq.backend.jit(f, static_argnums=static_argnums) + return backend_jit_lib[tq.backend.name](*args, **kwargs) + + ALL_JIT_LIBS.append(backend_jit_lib) + return _wrap_jit + + +def value_and_grad(f, argnums=0, has_aux=False): + backend_vg_lib = {} + + @wraps(f) + def _wrap_value_and_grad(*args, **kwargs): + if tq.backend.name not in backend_vg_lib: + try: + vg_func = tq.backend.value_and_grad(f, argnums=argnums, has_aux=has_aux) + except NotImplementedError: + + def vg_func(*args, **kwargs): + msg = ( + f"The engine requires auto-differentiation, " + f"which is not available for backend {tq.backend.name}" + ) + raise NotImplementedError(msg) + + backend_vg_lib[tq.backend.name] = vg_func + return backend_vg_lib[tq.backend.name](*args, **kwargs) + + ALL_JIT_LIBS.append(backend_vg_lib) + return _wrap_value_and_grad + + +set_backend = tq.set_backend + + +set_dtype = tq.set_dtype + + +def tensor_set_elem(tensor, idx, elem): + if tq.backend.name == "jax": + return tensor.at[idx].set(elem) + else: + assert tq.backend.name == "numpy" or tq.backend.name == "cupy" + tensor[idx] = elem + return tensor + + +def fori_loop(lower, upper, body_fun, init_val): + if tq.backend.name == "jax": + if jax is None: + raise JAXIMPORTERROR + return jax.lax.fori_loop(lower, upper, body_fun, init_val) + + assert tq.backend.name == "numpy" or tq.backend.name == "cupy" + val = init_val + for i in range(lower, upper): + val = body_fun(i, val) + return val + + +def scan(f, init, xs, length=None, reverse=False, unroll=1): + if tq.backend.name == "jax": + if jax is None: + raise JAXIMPORTERROR + return jax.lax.scan(f, init, xs, length, reverse, unroll) + + assert tq.backend.name == "numpy" or tq.backend.name == "cupy" + carry = init + ys = [] + xs = zip(*xs) + if reverse: + xs = reversed(list(xs)) + for x in xs: + carry, y = f(carry, x) + ys.append(y) + return carry, np.stack(list(reversed(ys))) + + +def get_xp(_backend): + if _backend.name == "cupy": + return cp + else: + return np + + +def get_uint_type(): + if tq.rdtypestr == "float64": + return np.uint64 + else: + assert tq.rdtypestr == "float32" + return np.uint32 diff --git a/src/tyxonq/chem/utils/circuit.py b/src/tyxonq/chem/utils/circuit.py new file mode 100644 index 0000000..21a1ef8 --- /dev/null +++ b/src/tyxonq/chem/utils/circuit.py @@ -0,0 +1,122 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from typing import Tuple + +import pandas as pd +from opt_einsum import contract_path +import tensornetwork as tn +import tyxonq as tq + + +def get_circuit_dataframe(circuit: tq.Circuit): + n_gates = circuit.gate_count() + n_cnot = circuit.gate_count(["cnot"]) + n_mc = circuit.gate_count(["multicontrol"]) + # avoid int overflow and for better formatting + n_qubits = circuit.circuit_param["nqubits"] + flop = count_circuit_flop(circuit) + + df_dict = { + "#qubits": n_qubits, + "#gates": [n_gates], + "#CNOT": [n_cnot], + "#multicontrol": [n_mc], + "depth": [circuit.to_qiskit().depth()], + "#FLOP": [flop], + } + return pd.DataFrame(df_dict) + + +def count_circuit_flop(circuit: tq.Circuit): + # tensor network contraction flops by greedy algorithm + nodes = circuit._copy()[0] + input_set_list = [set([id(e) for e in node.edges]) for node in nodes] + array_list = [node.tensor for node in nodes] + output_set = set([id(e) for e in tn.get_subgraph_dangling(nodes)]) + args = [] + for i in range(len(nodes)): + args.extend([array_list[i], input_set_list[i]]) + args.append(output_set) + _, desc = contract_path(*args) + return desc.opt_cost + + +def evolve_pauli(circuit: tq.Circuit, pauli_string: Tuple, theta: float): + # pauli_string in openfermion.QubitOperator.terms format + for idx, symbol in pauli_string: + if symbol == "X": + circuit.H(idx) + elif symbol == "Y": + circuit.SD(idx) + circuit.H(idx) + elif symbol == "Z": + continue + else: + raise ValueError(f"Invalid Pauli String: {pauli_string}") + + for i in range(len(pauli_string) - 1): + circuit.CNOT(pauli_string[i][0], pauli_string[i + 1][0]) + circuit.rz(pauli_string[-1][0], theta=theta) + + for i in reversed(range(len(pauli_string) - 1)): + circuit.CNOT(pauli_string[i][0], pauli_string[i + 1][0]) + + for idx, symbol in pauli_string: + if symbol == "X": + circuit.H(idx) + elif symbol == "Y": + circuit.H(idx) + circuit.S(idx) + elif symbol == "Z": + continue + else: + raise ValueError(f"Invalid Pauli String: {pauli_string}") + return circuit + + +def multicontrol_ry(theta): + # https://arxiv.org/pdf/2005.14475.pdf + c = tq.Circuit(4) + i, j, k, l = 0, 1, 2, 3 + + c.x(i) + c.x(k) + + c.ry(l, theta=theta / 8) + c.h(k) + c.cnot(l, k) + + c.ry(l, theta=-theta / 8) + c.h(i) + c.cnot(l, i) + + c.ry(l, theta=theta / 8) + c.cnot(l, k) + + c.ry(l, theta=-theta / 8) + c.h(j) + c.cnot(l, j) + + c.ry(l, theta=theta / 8) + c.cnot(l, k) + + c.ry(l, theta=-theta / 8) + c.cnot(l, i) + + c.ry(l, theta=theta / 8) + c.h(i) + c.cnot(l, k) + + # there's a typo in the paper + c.ry(l, theta=-theta / 8) + c.h(k) + c.cnot(l, j) + c.h(j) + + c.x(i) + c.x(k) + return c diff --git a/src/tyxonq/chem/utils/misc.py b/src/tyxonq/chem/utils/misc.py new file mode 100644 index 0000000..3889149 --- /dev/null +++ b/src/tyxonq/chem/utils/misc.py @@ -0,0 +1,143 @@ +# Copyright (c) 2023. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + + +from functools import wraps +from inspect import isfunction +from typing import List, Tuple + +import numpy as np +from scipy.sparse import coo_matrix +from renormalizer import Model, Mpo, Op +from renormalizer.model.basis import BasisSet +from openfermion import FermionOperator, QubitOperator, jordan_wigner, get_sparse_operator +from openfermion.utils import hermitian_conjugated +from qiskit.quantum_info import SparsePauliOp +import tyxonq as tq + +from ..constants import DISCARD_EPS + + +def csc_to_coo(csc): + coo = coo_matrix(csc) + mask = DISCARD_EPS < np.abs(coo.data.real) + indices = np.array([coo.row[mask], coo.col[mask]]).T + values = coo.data.real[mask].astype(tq.rdtypestr) + return tq.backend.coo_sparse_matrix(indices=indices, values=values, shape=coo.shape) + + +def fop_to_coo(fop: FermionOperator, n_qubits: int, real: bool = True): + op = get_sparse_operator(jordan_wigner(reverse_fop_idx(fop, n_qubits)), n_qubits=n_qubits) + if real: + op = op.real + return csc_to_coo(op) + + +def hcb_to_coo(qop: QubitOperator, n_qubits: int, real: bool = True): + op = get_sparse_operator(qop, n_qubits) + if real: + op = op.real + return csc_to_coo(op) + + +def qop_to_qiskit(qop: QubitOperator, n_qubits: int) -> SparsePauliOp: + sparse_list = [] + for k, v in qop.terms.items(): + s = "".join(kk[1] for kk in k) + idx = [kk[0] for kk in k] + sparse_list.append([s, idx, v]) + return SparsePauliOp.from_sparse_list(sparse_list, num_qubits=n_qubits) + + +def reverse_qop_idx(op: QubitOperator, n_qubits: int) -> QubitOperator: + ret = QubitOperator() + for pauli_string, v in op.terms.items(): + # internally QubitOperator assumes ascending index + pauli_string = tuple(reversed([(n_qubits - 1 - idx, symbol) for idx, symbol in pauli_string])) + ret.terms[pauli_string] = v + return ret + + +def reverse_fop_idx(op: FermionOperator, n_qubits: int) -> FermionOperator: + ret = FermionOperator() + for word, v in op.terms.items(): + word = tuple([(n_qubits - 1 - idx, symbol) for idx, symbol in word]) + ret.terms[word] = v + return ret + + +def format_ex_op(ex_op: Tuple) -> str: + if len(ex_op) == 2: + return f"{ex_op[0]}^ {ex_op[1]}" + else: + assert len(ex_op) == 4 + return f"{ex_op[0]}^ {ex_op[1]}^ {ex_op[2]} {ex_op[3]}" + + +def scipy_opt_wrap(f, gradient=True): + @wraps(f) + def _wrap_scipy_opt(_params, *args): + # scipy assumes 64bit https://github.com/scipy/scipy/issues/5832 + res = f(tq.backend.convert_to_tensor(_params), *args) + if gradient: + return [np.asarray(tq.backend.numpy(v), dtype=np.float64) for v in res] + else: + return np.asarray(tq.backend.numpy(res), dtype=np.float64) + + return _wrap_scipy_opt + + +def rdm_mo2ao(rdm: np.ndarray, mo_coeff: np.ndarray): + c = mo_coeff + if rdm.ndim == 2: + return c @ rdm @ c.T + else: + assert rdm.ndim == 4 + for _ in range(4): + rdm = np.tensordot(rdm, c.T, axes=1).transpose(3, 0, 1, 2) + return rdm + + +def canonical_mo_coeff(mo_coeff: np.ndarray): + # make the first large element positive + # all elements smaller than 1e-5 is highly unlikely (at least 1e10 basis) + largest_elem_idx = np.argmax(1e-5 < np.abs(mo_coeff), axis=0) + largest_elem = mo_coeff[(largest_elem_idx, np.arange(len(largest_elem_idx)))] + return mo_coeff * np.sign(largest_elem).reshape(1, -1) + + +def get_n_qubits(vector_or_matrix_or_mpo_func): + if isinstance(vector_or_matrix_or_mpo_func, list): + return len(vector_or_matrix_or_mpo_func) + if isfunction(vector_or_matrix_or_mpo_func): + return vector_or_matrix_or_mpo_func.n_qubit + # if hasattr(vector_or_matrix_or_mpo_func, "n_qubits"): + # return getattr(vector_or_matrix_or_mpo_func, "n_qubits") + return round(np.log2(vector_or_matrix_or_mpo_func.shape[0])) + + +def ex_op_to_fop(ex_op, with_conjugation=False): + if len(ex_op) == 2: + fop = FermionOperator(f"{ex_op[0]}^ {ex_op[1]}") + else: + assert len(ex_op) == 4 + fop = FermionOperator(f"{ex_op[0]}^ {ex_op[1]}^ {ex_op[2]} {ex_op[3]}") + if with_conjugation: + fop = fop - hermitian_conjugated(fop) + return fop + + +def get_dense_operator(basis: List[BasisSet], terms: List[Op]): + return Mpo(Model(basis, []), terms).todense() + + +def unpack_nelec(n_elec_s): + if isinstance(n_elec_s, tuple): + na, nb = n_elec_s + elif isinstance(n_elec_s, int): + na, nb = n_elec_s // 2, n_elec_s // 2 + else: + raise TypeError(f"Unknown electron number specification: {n_elec_s}") + return na, nb diff --git a/src/tyxonq/chem/utils/optimizer.py b/src/tyxonq/chem/utils/optimizer.py new file mode 100644 index 0000000..e08addd --- /dev/null +++ b/src/tyxonq/chem/utils/optimizer.py @@ -0,0 +1,155 @@ +# Copyright (c) 2024. The TenCirChem Developers. All Rights Reserved. +# +# This file is distributed under ACADEMIC PUBLIC LICENSE +# and WITHOUT ANY WARRANTY. See the LICENSE file for details. + +import numpy as np + +from scipy.optimize import OptimizeResult + + +def soap(fun, x0, args=(), u=0.1, maxfev=2000, atol=1e-3, callback=None, ret_traj=False, **kwargs): + """ + Scipy Optimizer interface for sequantial optimization with + approximate parabola (SOAP) + + Parameters + ---------- + fun : callable ``f(x, *args)`` + Function to be optimized. + x0 : ndarray, shape (n,) + Initial guess. Array of real elements of size (n,), + where 'n' is the number of independent variables. + args : tuple, optional + Extra arguments passed to the objective function. + u : float, optional + Step size for approximating the parabola. Defaults to 0.1. + maxfev : int + Maximum number of function evaluations to perform. + Empirically, ``5*len(x0)`` evaluations are sufficient for convergence. + Default: 2000. + atol : float + The tolerance for convergence. Defaults to 0.001. + callback : callable, optional + Called after each iteration. + ret_traj: bool, optional + Whether return trajectory for x or not. + + Returns + ------- + res : OptimizeResult + The optimization result represented as a SciPy ``OptimizeResult`` object. + Important attributes are: ``x`` the solution array. + """ + + nfev = 0 + nit = 0 + + def _fun(_x): + nonlocal nfev + nfev += 1 + return fun(_x, *args) + + trajectory = [x0.copy()] + vec_list = [] + # direction order + metric = np.abs(x0) + for i in np.argsort(metric)[::-1]: + vec = np.zeros_like(x0) + vec[i] = 1 + vec_list.append(vec) + vec_list_copy = vec_list.copy() + + e_list = [_fun(trajectory[-1])] + nfev_list = [nfev] + offset_list = [] + diff_list = [] + + scale = u + + while nfev < maxfev: + if len(vec_list) != 0: + vec = vec_list[0] + vec_list = vec_list[1:] + else: + vec_list = vec_list_copy.copy() + # continue + if len(trajectory) < len(vec_list_copy): + continue + p0 = trajectory[-1 - len(vec_list_copy)] + f0 = e_list[-1 - len(vec_list_copy)] + pn = trajectory[-1] + fn = e_list[-1] + fe = _fun(2 * pn - p0) + if fe > f0: + continue + average_direction = pn - p0 + if np.allclose(average_direction, 0): + continue + average_direction /= np.linalg.norm(average_direction) + replace_idx = np.argmax(np.abs(diff_list[-len(vec_list_copy) :])) + df = np.abs(diff_list[-len(vec_list_copy) :][replace_idx]) + if 2 * (f0 - 2 * fn + fe) * (f0 - fn - df) ** 2 > (f0 - fe) ** 2 * df: + continue + del vec_list[replace_idx] + vec_list = [average_direction] + vec_list + vec_list_copy = vec_list.copy() + continue + + vec_normed = vec / np.linalg.norm(vec) + x = [-scale, 0, scale] + es = [None, e_list[-1], None] + for j in [0, -1]: + es[j] = _fun(trajectory[-1] + x[j] * vec_normed) + if np.argmin(es) == 0: + x = [-scale * 4, -scale, 0, scale] + es = [None, es[0], es[1], es[2]] + for j in [0]: + es[j] = _fun(trajectory[-1] + x[j] * vec_normed) + elif np.argmin(es) == 2: + x = [-scale, 0, scale, scale * 4] + es = [es[0], es[1], es[2], None] + for j in [-1]: + es[j] = _fun(trajectory[-1] + x[j] * vec_normed) + else: + assert np.argmin(es) == 1 + a, b, c = np.polyfit(x, es, 2) + if np.argmin(es) not in [0, 3]: + offset = b / 2 / a + e = c - b**2 / 4 / a + else: + # print(a, b) + offset = -x[np.argmin(es)] + e = np.argmin(es) + offset_list.append(offset) + trajectory.append(trajectory[-1] - offset * vec_normed) + if len(es) == 3: + e_list.append(e) + else: + e_list.append(_fun(trajectory[-1])) + diff_list.append(e_list[-1] - e_list[-2]) + nfev_list.append(nfev) + + if callback is not None: + callback(np.copy(x0)) + + nit += 1 + # check convergence + if len(e_list) > 2 * len(x0): + if np.mean(e_list[-2 * len(x0) : -len(x0)]) - np.mean(e_list[-len(x0) :]) < atol: + break + + # the last measurement is probably from fitting + y = _fun(trajectory[-1]) + res = OptimizeResult( + fun=y, + x=trajectory[-1], + nit=nit, + nfev=nfev, + fun_list=np.array(e_list), + nfev_list=np.array(nfev_list), + success=True, + ) + if ret_traj: + res["trajectory"] = np.array(trajectory) + return res