From 45e7150886b1dfa6808dedad674610bd27b0344c Mon Sep 17 00:00:00 2001 From: liwt31 Date: Thu, 21 Dec 2023 19:24:56 +0800 Subject: [PATCH] enhance vha allow more direct parameter sharing --- example/vbe_ti.py | 23 +++++++------ tencirchem/dynamic/__init__.py | 2 +- tencirchem/dynamic/time_derivative.py | 49 ++++++++++++++++++--------- tencirchem/dynamic/transform.py | 40 ++++++++++++++-------- tencirchem/static/hea.py | 6 ++-- 5 files changed, 75 insertions(+), 45 deletions(-) diff --git a/example/vbe_ti.py b/example/vbe_ti.py index 5f3be8e..3f463ed 100644 --- a/example/vbe_ti.py +++ b/example/vbe_ti.py @@ -15,7 +15,7 @@ import tensorcircuit as tc from tencirchem import set_backend, Op, BasisSHO, BasisSimpleElectron, Mpo, Model -from tencirchem.dynamic import get_ansatz, qubit_encode_op, qubit_encode_basis +from tencirchem.dynamic import get_ansatz, qubit_encode_op_grouped, qubit_encode_basis from tencirchem.utils import scipy_opt_wrap from tencirchem.applications.vbe_lib import get_psi_indices, get_contracted_mpo, get_contract_args @@ -25,7 +25,6 @@ omega = 1 v = 1 # two qubit for each mode -# modify param_ids before modifying this n_qubit_per_mode = 2 nbas_v = 1 << n_qubit_per_mode @@ -51,8 +50,7 @@ def get_vha_terms(): ansatz_terms = [] for i in range(nsite): j = (i + 1) % nsite - ansatz_terms.append(Op(r"a^\dagger a", [i, j], v)) - ansatz_terms.append(Op(r"a^\dagger a", [j, i], -v)) + ansatz_terms.append(Op(r"a^\dagger a", [i, j], v) - Op(r"a^\dagger a", [j, i], v)) ansatz_terms.append(Op(r"a^\dagger a b^\dagger-b", [i, i, (i, 0)], g * omega)) basis = [] @@ -60,15 +58,18 @@ def get_vha_terms(): basis.append(BasisSimpleElectron(i)) basis.append(BasisSHO((i, 0), omega, nbas_v)) - ansatz_terms, _ = qubit_encode_op(ansatz_terms, basis, boson_encoding="gray") + ansatz_terms, _ = qubit_encode_op_grouped(ansatz_terms, basis, boson_encoding="gray") + + # More flexible ansatz by decoupling the parameters in the e-ph coupling term + ansatz_terms_extended = [] + for i in range(nsite): + ansatz_terms_extended.extend([ansatz_terms[2 * i]] + ansatz_terms[2 * i + 1]) spin_basis = qubit_encode_basis(basis, boson_encoding="gray") - # this is currently hard-coded for `n_qubit_per_mode==2` - param_ids = [1, -1, 0, 2, 3, 4, 5, 6, 7, 8] + [9, -9] + list(range(10, 18)) + [18, -18] + list(range(19, 27)) - return ansatz_terms, spin_basis, param_ids + return ansatz_terms_extended, spin_basis -ansatz_terms, spin_basis, param_ids = get_vha_terms() -ansatz = get_ansatz(ansatz_terms, spin_basis, n_layers, c, param_ids) +ansatz_terms, spin_basis = get_vha_terms() +ansatz = get_ansatz(ansatz_terms, spin_basis, n_layers, c) def cost_fn(params, h): @@ -141,7 +142,7 @@ def f(x): def main(): vqe_e = [] - thetas = np.zeros((max(param_ids) + 1) * n_layers) + thetas = np.zeros(len(ansatz_terms) * n_layers) for g in [1.5, 3]: for nbas in [4, 8, 12, 16, 20, 24, 28, 32]: diff --git a/tencirchem/dynamic/__init__.py b/tencirchem/dynamic/__init__.py index d441550..fc1c549 100644 --- a/tencirchem/dynamic/__init__.py +++ b/tencirchem/dynamic/__init__.py @@ -5,6 +5,6 @@ from tencirchem.dynamic.model import sbm, pyrazine -from tencirchem.dynamic.transform import qubit_encode_op, qubit_encode_basis +from tencirchem.dynamic.transform import qubit_encode_op, qubit_encode_op_grouped, qubit_encode_basis from tencirchem.dynamic.time_derivative import get_ansatz, get_jacobian_func, get_deriv, regularized_inversion from tencirchem.dynamic.time_evolution import TimeEvolution diff --git a/tencirchem/dynamic/time_derivative.py b/tencirchem/dynamic/time_derivative.py index 26b0a59..c111dd6 100644 --- a/tencirchem/dynamic/time_derivative.py +++ b/tencirchem/dynamic/time_derivative.py @@ -5,9 +5,12 @@ import logging +from typing import List, Union import numpy as np import scipy +from renormalizer.model.basis import BasisSet +from renormalizer import Op import tensorcircuit as tc from tencirchem.utils.backend import jit @@ -17,7 +20,7 @@ logger = logging.getLogger(__name__) -def construct_ansatz_op(ham_terms, spin_basis): +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 = [] @@ -44,13 +47,26 @@ def construct_ansatz_op(ham_terms, spin_basis): return ansatz_op_list -def get_circuit(ham_terms, spin_basis, n_layers, init_state, params, param_ids=None, compile_evolution=False): +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(ham_terms))) + param_ids = list(range(len(ansatz_terms))) params = tc.backend.reshape(params, [n_layers, max(param_ids) + 1]) - ansatz_op_list = construct_ansatz_op(ham_terms, spin_basis) + 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, tc.Circuit): c = tc.Circuit.from_qir(init_state.to_qir(), circuit_params=init_state.circuit_param) @@ -58,18 +74,19 @@ def get_circuit(ham_terms, spin_basis, n_layers, init_state, params, param_ids=N c = tc.Circuit(len(spin_basis), inputs=init_state) for i in range(0, n_layers): - for j, (ansatz_op, _, name, qubit_idx_list) in enumerate(ansatz_op_list): - param_id = np.abs(param_ids[j]) - # +0.1 is to avoid np.sign(0) problem - sign = np.sign(param_ids[j] + 0.1) - theta = sign * params[i, param_id] - if not compile_evolution: - np.testing.assert_allclose(ansatz_op @ ansatz_op, np.eye(len(ansatz_op))) - name = f"exp(-iθ{name})" - c.exp1(*qubit_idx_list, unitary=ansatz_op, theta=theta, name=name) - else: - pauli_string = tuple(zip(qubit_idx_list, name)) - c = evolve_pauli(c, pauli_string, theta=2 * theta) + 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 diff --git a/tencirchem/dynamic/transform.py b/tencirchem/dynamic/transform.py index 8572ebc..79868a4 100644 --- a/tencirchem/dynamic/transform.py +++ b/tencirchem/dynamic/transform.py @@ -4,7 +4,7 @@ # and WITHOUT ANY WARRANTY. See the LICENSE file for details. -from typing import Any, List +from typing import Any, List, Union, Tuple import numpy as np from renormalizer.model import Op, OpSum, Model @@ -19,7 +19,7 @@ DOF_SALT = "TCCQUBIT" -def check_basis_type(basis): +def check_basis_type(basis: List[BasisSet]): for b in basis: if isinstance(b, (BasisMultiElectronVac,)): raise TypeError(f"Unsupported basis: {type(b)}") @@ -27,21 +27,18 @@ def check_basis_type(basis): raise ValueError(f"For only two DOFs are allowed in BasisMultiElectron. Got {b}") -def qubit_encode_op(terms, basis, boson_encoding=None): +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, terms) + model = Model(basis, []) - old_ham_terms = [] - # `reversed` is for normal ordering with `pop` - for op in reversed(terms): - old_ham_terms.append(op.split_elementary(model.dof_to_siteidx)) - - new_ham_terms = [] - while old_ham_terms: - terms, factor = old_ham_terms.pop() + 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) @@ -52,13 +49,13 @@ def qubit_encode_op(terms, basis, boson_encoding=None): new_term = new_term * opsum new_term = new_term * factor - new_ham_terms.extend(new_term) + new_terms.extend(new_term) # post process # pick out constants - identity_terms = [] + identity_terms: List[Op] = [] non_identity_terms = OpSum() - for op in new_ham_terms: + for op in new_terms: if op.is_identity: identity_terms.append(op) else: @@ -69,6 +66,19 @@ def qubit_encode_op(terms, basis, boson_encoding=None): 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: diff --git a/tencirchem/static/hea.py b/tencirchem/static/hea.py index 322fcae..d93315c 100644 --- a/tencirchem/static/hea.py +++ b/tencirchem/static/hea.py @@ -216,7 +216,7 @@ def from_molecule(cls, m: Mole, active_space=None, n_layers=3, mapping="parity", """ from tencirchem import UCC - ucc = UCC(m, active_space=active_space) + ucc = UCC(m, active_space=active_space, run_ccsd=False, run_fci=False) return cls.ry(ucc.int1e, ucc.int2e, ucc.n_elec, ucc.e_core, n_layers=n_layers, mapping=mapping, **kwargs) @classmethod @@ -398,7 +398,9 @@ class FakeFCISolver: def __init__(self): self.instance: HEA = None self.config_function = config_function - self.instance_kwargs = kwargs + 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)