Merge pull request #1 from jjren/trotter

add trotter algorithm for dynamics

Author: liwt31

tencirchem/dynamic/time_derivative.py

@@ -10,10 +10,8 @@
 
 logger = logging.getLogger(__name__)
 
+def construct_ansatz_op(ham_terms, spin_basis):
 
-def get_circuit(ham_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)))
  dof_idx_dict = {b.dof: i for i, b in enumerate(spin_basis)}
  ansatz_op_list = []
 
@@ -35,17 +33,25 @@ def get_circuit(ham_terms, spin_basis, n_layers, init_state, params, param_ids=N
  op_mat = np.kron(op_mat, tc.gates._z_matrix)
  name += "Z"
  qubit_idx_list = [dof_idx_dict[dof] for dof in op.dofs]
- ansatz_op_list.append((op_mat, name, qubit_idx_list))
+ ansatz_op_list.append((op_mat, op.factor, name, qubit_idx_list))
+ 
+ return ansatz_op_list
+
+def get_circuit(ham_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)))
 
  params = tc.backend.reshape(params, [n_layers, max(param_ids) + 1])
 
+ ansatz_op_list = construct_ansatz_op(ham_terms, spin_basis)
+
  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):
+ for j, (ansatz_op, op_factor, 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)
@@ -59,6 +65,20 @@ def get_circuit(ham_terms, spin_basis, n_layers, init_state, params, param_ids=N
  c = evolve_pauli(c, pauli_string, theta=theta)
  return c
 
+def one_trotter_step(ham_terms, spin_basis, init_state, dt):
+ """
+ one step first order trotter decompostion 
+ """
+ ansatz_op_list = construct_ansatz_op(ham_terms, spin_basis)
+
+ 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 (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

tencirchem/dynamic/time_evolution.py

@@ -1,4 +1,5 @@
 import time
+import logging
 from functools import partial
 from collections import defaultdict
 from typing import Dict, List, Optional
@@ -12,13 +13,13 @@
 from renormalizer.model.basis import BasisSet
 
 from tencirchem.dynamic.transform import qubit_encode_op, qubit_encode_basis, get_init_circuit
-from tencirchem.dynamic.time_derivative import get_circuit, get_ansatz, get_jacobian_func, get_deriv, get_pvqd_loss_func
+from tencirchem.dynamic.time_derivative import get_circuit, get_ansatz, get_jacobian_func, get_deriv, get_pvqd_loss_func, one_trotter_step
 
+logger = logging.getLogger(__name__)
 
 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,
@@ -30,6 +31,7 @@ def __init__(
  eps: float = 1e-5,
  property_op_dict: Dict = None,
  ref_only: bool = False,
+ ivp_config = None,
  ):
  # handling defaults
  if init_condition is None:
@@ -61,9 +63,14 @@ def __init__(
  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)
 
+ self.current_circuit = self.init_circuit
  # setup runtime components
  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)
@@ -102,26 +109,36 @@ def solve_pvqd(_params, delta_t):
 
  self.wall_time_list = []
 
- def kernel(self, tau, pvqd=False):
+ 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 not pvqd:
- scipy_sol = solve_ivp(self.scipy_deriv, [self.t, self.t + tau], self.params)
- new_params = scipy_sol.y[:, -1]
+ 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:
- scipy_sol = self.solve_pvqd(self.params, tau)
- new_params = self.params + scipy_sol.x
+ assert algo == "trotter"
+ self.current_circuit = one_trotter_step(self.model.ham_terms, self.model.basis,
+ self.current_circuit, tau)
+ state = self.current_circuit.state()
+
  time1 = time.time()
  self.t_list.append(self.t + tau)
- self.params_list.append(new_params)
  # t and params already updated
- state = self.ansatz(self.params)
  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)
- self.scipy_sol_list.append(scipy_sol)
  # calculate properties
  self.update_property_dict(state, state_ref)
 
@@ -203,3 +220,13 @@ def property_dict(self):
  @property
  def wall_time(self):
  return self.wall_time_list[-1]
+
+class Ivp_Config:
+ def __init__(self,
+ method="RK45",
+ rtol=1e-3,
+ atol=1e-6):
+ self.method = method
+ self.rtol = rtol
+ self.atol = atol
+

tencirchem/dynamic/transform.py

@@ -248,20 +248,32 @@ def get_encoding(m, boson_encoding):
 
 
 def get_init_circuit(model_ref, model, boson_encoding, init_condition):
- for v in init_condition.values():
+ for k, v in init_condition.items():
+ basis = model_ref.dof_to_basis[k]
  if not isinstance(v, int):
- return get_init_circuit_general(model_ref, model, boson_encoding, init_condition)
+ if isinstance(basis, BasisHalfSpin) and v.shape == (2,2) and \
+ np.allclose(np.eye(2), v @ v.T.conj) and np.allclose(np.eye(2), v.T.conj @ v):
+ continue
+ else:
+ return get_init_circuit_general(model_ref, model, boson_encoding, init_condition)
+ 
  # replace the dof_name key to site_index key
  circuit = tc.Circuit(len(model.basis))
  for k, v in init_condition.items():
  basis = model_ref.dof_to_basis[k]
  if isinstance(basis, BasisHalfSpin):
  if v == 1:
  circuit.X(model.dof_to_siteidx[k])
+ elif v.shape == (2,2):
+ circuit.ANY(idx, unitary=v)
+ else:
+ assert v == 0
  elif isinstance(basis, BasisMultiElectron):
  if v == 1:
  idx = model.dof_to_siteidx[basis.dofs]
  circuit.X(idx)
+ else:
+ assert v == 0
  else:
  assert basis.is_phonon
  if boson_encoding is None:

tests/dynamics/test_dynamics.py

@@ -6,7 +6,7 @@
 from tencirchem.dynamic import sbm, qubit_encode_op, qubit_encode_basis, TimeEvolution
 
 
-@pytest.mark.parametrize("algorithm", ["vanilla", "include_phase", "p-VQD"])
+@pytest.mark.parametrize("algorithm", ["vanilla", "include_phase", "p-VQD", "trotter"])
 def test_sbm(reset_backend, algorithm):
  set_backend("jax")
  epsilon = 0
@@ -30,11 +30,19 @@ def test_sbm(reset_backend, algorithm):
  eps=1e-5,
  )
  te.include_phase = algorithm == "include_phase"
-
- pvqd = algorithm == "p-VQD"
- tau = 0.1
+ 
+ if algorithm in ["vanilla", "include_phase"]:
+ algo = "vqd"
+ tau = 0.1
+ elif algorithm == "p-VQD":
+ algo = "pvqd"
+ tau = 0.1
+ else:
+ algo = "trotter"
+ tau = 0.02
+ 
  for _ in range(50):
- te.kernel(tau, pvqd=pvqd)
+ te.kernel(tau, algo=algo)
  z = te.property_dict["Z"]
  np.testing.assert_allclose(z[:, 0], z[:, 1], atol=1e-2)