enhance vha allow more direct parameter sharing

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,24 +50,26 @@ 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 = []
     for i in range(nsite):
         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]:

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

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,32 +47,46 @@ 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)
     else:
         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
 
 

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,29 +19,26 @@
 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)}")
         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, 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:

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)