Repository: shuaigroup/Renormalizer (language: python)

Issue:
[WIP] Basic tree tensornetwork algorithm
This is more like a preview than an actual PR. Feedbacks are welcome.
Implemented features:
- Automatic tree tensornetwork operator (TTO) construction
- 2site optimization algorithm

Please see the test cases for how to use them.

Obviously, there are a lot of important features that remain to be implemented. The following ones have a higher priority
- QN support
- VMF (imaginary) time evolution

After that, we may have a proper comparison/benchmark with (ML-)MCTDH

On the longer term I'm going to look at
- GPU support
- PS time evolution
- Finite temperature (Could be *very* tedious)
- TTO/TTS (tree tensornetwork states) manipulations
- Compression
- Documentation

Of course, one may imagine a lot of other features such as excited states optimization algorithm and technical enhancements like dumping to or loading from hard disk. In the foreseeable future, I'm afraid I'll only work on ones that will enable finite temperature mobility calculation due to my limited time. 


Discussion / hints:

Fix patch:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index 9c40560e..a0d8ea36 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -40,7 +40,7 @@ jobs:
           command: |
               pip install pytest-xdist pytest-cov
               export RENO_NUM_THREADS=1
-              pytest -n 4 --durations=0 --cov=renormalizer
+              pytest -n 4 --durations=0 --cov=renormalizer renormalizer
               pip install primme==3.2.*
               pytest --durations=0 renormalizer/mps/tests/test_gs.py::test_multistate --cov=renormalizer --cov-append
       - run:
@@ -54,12 +54,6 @@ jobs:
       - run:
           name: Build docs
           command: cd doc; make html
-      - run:
-          name: Codecov report
-          command: |
-              curl -Os https://uploader.codecov.io/latest/linux/codecov
-              chmod +x codecov
-              ./codecov
       - persist_to_workspace:
           root: doc/
           paths: html
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index ba83d91f..d7e4c202 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -23,7 +23,7 @@ jobs:
         run: |
           pip install pytest-xdist pytest-cov
           export RENO_NUM_THREADS=1
-          pytest -n 4 --durations=0 --cov=renormalizer
+          pytest -n 4 --durations=0 --cov=renormalizer renormalizer
           pip install primme==3.2.*
           pytest --durations=0 renormalizer/mps/tests/test_gs.py::test_multistate --cov=renormalizer --cov-append
       - name: run examples
@@ -38,4 +38,6 @@ jobs:
           cd doc
           make html
       - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@v3
+        uses: codecov/codecov-action@v4
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
diff --git a/example/run.sh b/example/run.sh
index e2b72f6a..86eefd4f 100644
--- a/example/run.sh
+++ b/example/run.sh
@@ -1,6 +1,15 @@
 export PYTHONPATH=../:PYTHONPATH
 code=0
-for python_args in fmo.py sbm.py h2o_qc.py "dynamics.py std.yaml" "transport_kubo.py std.yaml"; do
+for python_args in fmo.py \
+                   sbm.py \
+                   h2o_qc.py \
+                   "dynamics.py std.yaml" \
+                   "transport_kubo.py std.yaml" \
+                   ./ttns/junction_zt.py \
+                   "./ttns/junction_ft.py 32 1 100" \
+                   "./ttns/sbm_zt.py 050 001 050"\
+                   ./ttns/sbm_ft.py
+do
     echo ============================$python_args=============================
     timeout 20s python $python_args
     exit_code=$?
diff --git a/example/ttns/junction_ft.py b/example/ttns/junction_ft.py
new file mode 100644
index 00000000..532fa18d
--- /dev/null
+++ b/example/ttns/junction_ft.py
@@ -0,0 +1,237 @@
+import sys
+
+from renormalizer import BasisHalfSpin, Op, Quantity, BasisSHO, BasisDummy
+from renormalizer.mps.mps import expand_bond_dimension_general
+from renormalizer.sbm import ColeDavidsonSDF
+from renormalizer.utils.configs import (
+    EvolveMethod,
+    EvolveConfig,
+    CompressConfig,
+    CompressCriteria,
+)
+from renormalizer.utils import constant, log
+from renormalizer.utils.log import package_logger as logger
+from renormalizer.tn import BasisTree, TreeNodeBasis, TTNS, TTNO
+
+import numpy as np
+
+
+Ms_str = sys.argv[1]  # 8 16 24 32
+initial_str = sys.argv[2]  # 1 0
+temperature_str = sys.argv[3]  # 000 100 200 300
+
+dump_dir = "./"
+job_name = f"Ms{Ms_str}_i{initial_str}_temperature{temperature_str}"  ####################
+log.register_file_output(dump_dir + job_name + ".log", mode="w")
+
+Ms = int(Ms_str)
+
+n_ph_mode = 1000
+omega_c = Quantity(500, "cm-1").as_au()
+ita = Quantity(2000, "cm-1").as_au() / 2
+beta = 0.5
+upper_limit = Quantity(1, "eV").as_au() * 10
+logger.info(("phonon parameters", omega_c, ita, beta, upper_limit))
+sdf = ColeDavidsonSDF(ita, omega_c, beta, upper_limit)
+w, c2 = sdf.Wang1(n_ph_mode)
+c = np.sqrt(c2)
+logger.info(w)
+logger.info(c)
+
+reno = sdf.reno(w[-1])
+logger.info(f"renormalization constant: {reno}")
+
+temperature = Quantity(int(temperature_str), "K").to_beta()
+logger.info(f"temperature beta: {temperature}")
+
+n_e_mode = 320
+
+beta_e = Quantity(1, "eV").as_au() * reno
+alpha_e = Quantity(0.2, "eV").as_au() * reno
+v = 0.1 * reno
+mu_l = Quantity(v / 2, "eV").as_au()
+mu_r = Quantity(-v / 2, "eV").as_au()
+
+
+e_k = np.arange(1, n_e_mode + 1) / (n_e_mode + 1) * 4 * beta_e - 2 * beta_e
+rho_e = 1 / (e_k[1] - e_k[0])
+e_k_l = e_k - mu_l
+e_k_r = e_k - mu_r
+
+mode_with_e = [(f"L{i}", e) for i, e in enumerate(e_k_l)] + [(f"R{i}", e) for i, e in enumerate(e_k_r)]
+mode_with_e.sort(key=lambda x: x[1])
+logger.info(mode_with_e)
+
+basis = []
+first_positive = True
+for mode, e in mode_with_e:
+    if e > 0 and first_positive:
+        first_positive = False
+        basis.append(BasisHalfSpin("s"))
+    basis.append(BasisHalfSpin((mode, "p")))
+    basis.append(BasisHalfSpin((mode, "q")))
+
+dofs = [b.dofs[0] for b in basis]
+logger.info(dofs)
+s_idx = dofs.index("s")
+basis_tree_l = BasisTree.binary_mctdh(basis[:s_idx], dummy_label="EL-dummy")
+basis_tree_r = BasisTree.binary_mctdh(basis[s_idx + 1 :], dummy_label="ER-dummy")
+
+
+ham_terms = []
+i_l_terms = []
+i_r_terms = []
+for mode, e in mode_with_e:
+    if mode[0] == "L":
+        mu = mu_l
+        i_terms = i_l_terms
+    else:
+        assert mode[0] == "R"
+        mu = mu_r
+        i_terms = i_r_terms
+
+    ham_terms.append(Op("+ -", (mode, "p"), e + mu))
+    ham_terms.append(Op("+ -", (mode, "q"), -(e + mu)))
+    v2 = alpha_e**2 / beta_e**2 * np.sqrt(4 * beta_e**2 - (e + mu) ** 2) / 2 / np.pi / rho_e
+    v = np.sqrt(v2)
+    theta = np.arctan(np.exp(-temperature * e / 2))
+    idx = dofs.index((mode, "p"))
+    if idx < s_idx:
+        z_idx = list(range(idx + 1, s_idx))
+    else:
+        assert s_idx < idx
+        z_idx = list(range(s_idx + 1, idx))
+    z_dofs = [dofs[i] for i in z_idx]
+    op1 = Op(
+        "+ " + "Z " * len(z_idx) + "-",
+        [(mode, "p")] + z_dofs + ["s"],
+        v * np.cos(theta),
+    )
+    op2 = Op(
+        "- " + "Z " * len(z_idx) + "+",
+        [(mode, "p")] + z_dofs + ["s"],
+        v * np.cos(theta),
+    )
+    idx = dofs.index((mode, "q"))
+    if idx < s_idx:
+        z_idx = list(range(idx + 1, s_idx))
+    else:
+        assert s_idx < idx
+        z_idx = list(range(s_idx + 1, idx))
+    z_dofs = [dofs[i] for i in z_idx]
+    op3 = Op(
+        "- " + "Z " * len(z_idx) + "-",
+        [(mode, "q")] + z_dofs + ["s"],
+        v * np.sin(theta),
+    )
+    op4 = Op(
+        "+ " + "Z " * len(z_idx) + "+",
+        [(mode, "q")] + z_dofs + ["s"],
+        v * np.sin(theta),
+    )
+    ham_terms.extend([op1, op2, op3, op4])
+    # move 1j to expectation
+    i_terms.extend(op2 - op1 + op4 - op3)
+
+initial_occupied = initial_str == "1"
+
+if initial_occupied:
+    ham_terms.append(Op("+ -", "s", qn=[0, 0], factor=-4 * (c**2 / w**2).sum()))
+
+# vibrations at last
+
+# boson energy
+for imode in range(n_ph_mode):
+    op1 = Op(r"p^2", f"v_{imode}_p", factor=0.5, qn=0)
+    op2 = Op(r"x^2", f"v_{imode}_p", factor=0.5 * w[imode] ** 2, qn=0)
+    ham_terms.extend([op1, op2])
+    op1 = Op(r"p^2", f"v_{imode}_q", factor=-0.5, qn=0)
+    op2 = Op(r"x^2", f"v_{imode}_q", factor=-0.5 * w[imode] ** 2, qn=0)
+    ham_terms.extend([op1, op2])
+
+theta_array = np.arctanh(np.exp(-w * temperature / 2))
+# system-boson coupling
+for imode in range(n_ph_mode):
+    sys_op = Op("+ -", "s", qn=[0, 0])
+    if initial_occupied:
+        sys_op = sys_op - Op.identity("s")
+
+    theta = theta_array[imode]
+    op1 = sys_op * Op(r"x", f"v_{imode}_p", factor=2 * c[imode] * np.cosh(theta), qn=[0])
+    op2 = sys_op * Op(r"x", f"v_{imode}_q", factor=2 * c[imode] * np.sinh(theta), qn=[0])
+    ham_terms.extend(op1 + op2)
+
+nbas = np.max([16 * c2 / w**3 * np.cosh(theta_array) ** 2, np.ones(n_ph_mode) * 4], axis=0)
+nbas = np.round(nbas).astype(int)
+nbas = np.min([nbas, np.ones(n_ph_mode) * 512], axis=0)
+logger.info(nbas)
+basis_list_phonon = []
+for imode in range(n_ph_mode):
+    basis_list_phonon.append(BasisSHO(f"v_{imode}_p", w[imode], int(nbas[imode])))
+    basis_list_phonon.append(BasisSHO(f"v_{imode}_q", w[imode], int(nbas[imode])))
+
+labels = np.array([[nbas > Ms], [nbas > Ms]]).T.ravel()
+basis_tree_phonon = BasisTree.binary_mctdh(
+    basis_list_phonon,
+    contract_primitive=True,
+    contract_label=labels,
+    dummy_label="phonon-dummy",
+)
+node1 = TreeNodeBasis([BasisDummy("dummy")])
+node1.add_child([basis_tree_l.root, basis_tree_r.root])
+node2 = TreeNodeBasis([basis[s_idx]])
+node2.add_child([node1, basis_tree_phonon.root])
+basis_tree = BasisTree(node2)
+basis_tree.print(logger.info)
+
+
+# model = Model(basis, ham_terms)
+ttno = TTNO(basis_tree, ham_terms)
+i_l_mpo = TTNO(basis_tree, i_l_terms)
+i_r_mpo = TTNO(basis_tree, i_r_terms)
+# n_l_mpo = TTNO(basis_tree, terms=[Op("+ -", f"L{i}") for i in range(n_e_mode)])
+# n_r_mpo = TTNO(basis_tree, terms=[Op("+ -", f"R{i}") for i in range(n_e_mode)])
+n_s_mpo = TTNO(basis_tree, terms=Op("+ -", "s"))
+ttno.print_shape(False, logger.info)
+i_l_mpo.print_shape(False, logger.info)
+i_r_mpo.print_shape(False, logger.info)
+# n_r_mpo.print_shape(False, logger.info)
+# n_s_mpo.print_shape(False, logger.info)
+# 0 - [1, 0] (spin up) means occupation, 1 - [0, 1] (spin down) means unoccupation
+# initial condition is taken care of in the Hamiltonian
+condition = {dofs[i]: 1 for i in range(len(dofs))}
+if initial_occupied:
+    condition["s"] = 0
+else:
+    condition["s"] = 1
+
+ttns = TTNS(basis_tree, condition=condition)
+ttns.compress_config = CompressConfig(CompressCriteria.fixed, max_bonddim=Ms)
+ttns = expand_bond_dimension_general(ttns, ttno, ex_mps=None)
+ttns.evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
+ttns.print_shape(print_function=logger.info, full=False)
+
+step = 0.5 * constant.fs2au
+# step = 5
+nsteps = 200
+au2muA = 6.623618237510e3
+i = 0
+current_list = []
+while True:
+    i_l = (1j * ttns.expectation(i_l_mpo)).real
+    i_r = (1j * ttns.expectation(i_r_mpo)).real
+    # n_l = mps.expectation(n_l_mpo)
+    # n_r = mps.expectation(n_r_mpo)
+    n_s = ttns.expectation(n_s_mpo)
+    logger.info((i, ttns.bond_dims))
+    current = (i_r - i_l) / 2 * au2muA
+    # logger.info((n_l, n_r, n_s, i_l*au2muA, i_r*au2muA, current))
+    logger.info((n_s, i_l * au2muA, i_r * au2muA, current))
+    current_list.append(current)
+    i += 1
+    if i == nsteps:
+        break
+    if i > 0:
+        ttns.evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
+    ttns = ttns.evolve(ttno, step)
+logger.info(current_list)
diff --git a/example/ttns/junction_zt.py b/example/ttns/junction_zt.py
new file mode 100644
index 00000000..7f3a283a
--- /dev/null
+++ b/example/ttns/junction_zt.py
@@ -0,0 +1,180 @@
+from renormalizer import BasisHalfSpin, Op, Quantity, BasisSHO, BasisDummy
+from renormalizer.mps.mps import expand_bond_dimension_general
+from renormalizer.sbm import ColeDavidsonSDF
+from renormalizer.utils.configs import (
+    EvolveMethod,
+    EvolveConfig,
+    CompressConfig,
+    CompressCriteria,
+)
+from renormalizer.utils import constant
+from renormalizer.utils.log import package_logger as logger
+from renormalizer.tn import BasisTree, TreeNodeBasis, TTNS, TTNO
+
+import numpy as np
+
+
+n_ph_mode = 500
+omega_c = Quantity(500, "cm-1").as_au()
+ita = Quantity(2000, "cm-1").as_au() / 2
+beta = 0.5
+upper_limit = Quantity(1, "eV").as_au() * 5
+logger.info(("phonon parameters", omega_c, ita, beta, upper_limit))
+sdf = ColeDavidsonSDF(ita, omega_c, beta, upper_limit)
+w, c2 = sdf.Wang1(n_ph_mode)
+c = np.sqrt(c2)
+logger.info(w)
+logger.info(c)
+
+reno = sdf.reno(w[-1])
+logger.info(f"renormalization constant: {reno}")
+
+
+n_e_mode = 160
+
+beta_e = Quantity(1, "eV").as_au() * reno
+alpha_e = Quantity(0.2, "eV").as_au() * reno
+v = 0.1 * reno
+mu_l = Quantity(v / 2, "eV").as_au()
+mu_r = Quantity(-v / 2, "eV").as_au()
+
+
+e_k = np.arange(1, n_e_mode + 1) / (n_e_mode + 1) * 4 * beta_e - 2 * beta_e
+rho_e = 1 / (e_k[1] - e_k[0])
+e_k_l = e_k - mu_l
+e_k_r = e_k - mu_r
+
+mode_with_e = [(f"L{i}", e) for i, e in enumerate(e_k_l)] + [(f"R{i}", e) for i, e in enumerate(e_k_r)]
+mode_with_e.sort(key=lambda x: x[1])
+logger.info(mode_with_e)
+
+basis = []
+first_positive = True
+for mode, e in mode_with_e:
+    if e > 0 and first_positive:
+        first_positive = False
+        basis.append(BasisHalfSpin("s"))
+    basis.append(BasisHalfSpin(mode))
+
+dofs = [b.dofs[0] for b in basis]
+logger.info(dofs)
+s_idx = dofs.index("s")
+basis_tree_l = BasisTree.binary_mctdh(basis[:s_idx], dummy_label="EL-dummy")
+basis_tree_r = BasisTree.binary_mctdh(basis[s_idx + 1 :], dummy_label="ER-dummy")
+
+
+ham_terms = []
+i_l_terms = []
+i_r_terms = []
+for mode, e in mode_with_e:
+    if mode[0] == "L":
+        mu = mu_l
+        i_terms = i_l_terms
+    else:
+        assert mode[0] == "R"
+        mu = mu_r
+        i_terms = i_r_terms
+
+    ham_terms.append(Op("+ -", mode, e + mu))
+    v2 = alpha_e**2 / beta_e**2 * np.sqrt(4 * beta_e**2 - (e + mu) ** 2) / 2 / np.pi / rho_e
+    v = np.sqrt(v2)
+    idx = dofs.index(mode)
+    if idx < s_idx:
+        z_idx = list(range(idx + 1, s_idx))
+    else:
+        assert s_idx < idx
+        z_idx = list(range(s_idx + 1, idx))
+    z_dofs = [dofs[i] for i in z_idx]
+    op1 = Op("+ " + "Z " * len(z_idx) + "-", [mode] + z_dofs + ["s"], v)
+    op2 = Op("- " + "Z " * len(z_idx) + "+", [mode] + z_dofs + ["s"], v)
+    ham_terms.extend([op1, op2])
+    # move 1j to expectation
+    i_terms.extend(op2 - op1)
+
+initial_occupied = True
+
+if initial_occupied:
+    ham_terms.append(Op("+ -", "s", qn=[0, 0], factor=-4 * (c**2 / w**2).sum()))
+
+# vibrations at last
+
+# boson energy
+for imode in range(n_ph_mode):
+    op1 = Op(r"p^2", f"v_{imode}", factor=0.5, qn=0)
+    op2 = Op(r"x^2", f"v_{imode}", factor=0.5 * w[imode] ** 2, qn=0)
+    ham_terms.extend([op1, op2])
+
+# system-boson coupling
+for imode in range(n_ph_mode):
+    sys_op = Op("+ -", "s", qn=[0, 0])
+    if initial_occupied:
+        sys_op = sys_op - Op.identity("s")
+
+    op = sys_op * Op(r"x", f"v_{imode}", factor=2 * c[imode], qn=[0])
+    ham_terms.extend(op)
+
+nbas = np.max([16 * c2 / w**3, np.ones(n_ph_mode) * 4], axis=0)
+nbas = np.round(nbas).astype(int)
+logger.info(nbas)
+basis_list_phonon = []
+for imode in range(n_ph_mode):
+    basis_list_phonon.append(BasisSHO(f"v_{imode}", w[imode], int(nbas[imode])))
+
+basis_tree_phonon = BasisTree.binary_mctdh(basis_list_phonon, dummy_label="phonon-dummy")
+node1 = TreeNodeBasis([BasisDummy("dummy")])
+node1.add_child([basis_tree_l.root, basis_tree_r.root])
+node2 = TreeNodeBasis([basis[s_idx]])
+node2.add_child([node1, basis_tree_phonon.root])
+basis_tree = BasisTree(node2)
+basis_tree.print(logger.info)
+
+
+# model = Model(basis, ham_terms)
+ttno = TTNO(basis_tree, ham_terms)
+i_l_mpo = TTNO(basis_tree, i_l_terms)
+i_r_mpo = TTNO(basis_tree, i_r_terms)
+n_l_mpo = TTNO(basis_tree, terms=[Op("+ -", f"L{i}") for i in range(n_e_mode)])
+n_r_mpo = TTNO(basis_tree, terms=[Op("+ -", f"R{i}") for i in range(n_e_mode)])
+n_s_mpo = TTNO(basis_tree, terms=Op("+ -", "s"))
+ttno.print_shape(False, logger.info)
+i_l_mpo.print_shape(False, logger.info)
+i_r_mpo.print_shape(False, logger.info)
+n_r_mpo.print_shape(False, logger.info)
+n_s_mpo.print_shape(False, logger.info)
+# 0 - [1, 0] (spin up) means occupation, 1 - [0, 1] (spin down) means unoccupation
+# initial condition is taken care of in the Hamiltonian
+condition = {dofs[i]: 1 for i in range(s_idx + 1, len(dofs))}
+if initial_occupied:
+    condition["s"] = 0
+else:
+    condition["s"] = 1
+
+ttns = TTNS(basis_tree, condition=condition)
+ttns.compress_config = CompressConfig(CompressCriteria.fixed, max_bonddim=32)
+ttns = expand_bond_dimension_general(ttns, ttno, ex_mps=None)
+ttns.evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
+ttns.print_shape(print_function=logger.info, full=False)
+
+step = 0.5 * constant.fs2au
+# step = 5
+nsteps = 100
+au2muA = 6.623618237510e3
+i = 0
+current_list = []
+while True:
+    i_l = (1j * ttns.expectation(i_l_mpo)).real
+    i_r = (1j * ttns.expectation(i_r_mpo)).real
+    n_l = ttns.expectation(n_l_mpo)
+    n_r = ttns.expectation(n_r_mpo)
+    n_s = ttns.expectation(n_s_mpo)
+    logger.info((i, ttns.bond_dims))
+    current = (i_r - i_l) / 2 * au2muA
+    logger.info((n_l, n_r, n_s, i_l * au2muA, i_r * au2muA, current))
+    current_list.append(current)
+    i += 1
+    if i == nsteps:
+        break
+    if i > 0:
+        ttns.evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
+    ttns = ttns.evolve(ttno, step)
+logger.info(current_list)
diff --git a/example/ttns/sbm_ft.py b/example/ttns/sbm_ft.py
new file mode 100644
index 00000000..a30a1784
--- /dev/null
+++ b/example/ttns/sbm_ft.py
@@ -0,0 +1,123 @@
+import sys
+
+from renormalizer.model import Op
+from renormalizer.model import basis as ba
+from renormalizer.mps.mps import expand_bond_dimension_general
+from renormalizer.sbm import ColeDavidsonSDF
+from renormalizer.utils import EvolveConfig, CompressConfig, CompressCriteria, EvolveMethod
+from renormalizer.utils import log
+from renormalizer.tn import BasisTree, TTNO, TTNS, TreeNodeBasis
+
+import numpy as np
+
+
+# ita_str = sys.argv[1]  # 010, 025, 100
+# omega_c_str = sys.argv[2]  # 001, 100
+# beta_str = sys.argv[3]  # 025, 100
+# temperature_str = sys.argv[4]   # 05, 10, 15
+ita_str = "010"
+omega_c_str = "010"
+beta_str = "025"
+temperature_str = "20"
+
+
+ita = int(ita_str) / 10 # 1, 2.5, 5, 10
+eps = 0
+Delta = 1
+omega_c = int(omega_c_str) / 10  # 0.1, 1, 10
+
+beta = int(beta_str) / 100  # 0.25, 0.5, 0.75, 1
+temperature = int(temperature_str) / 10 # 0.5, 1, 1.5
+
+from renormalizer.utils.log import package_logger
+logger = package_logger
+dump_dir = "./"
+job_name = f"ps1_binary_ita{ita_str}_omega{omega_c_str}_beta{beta_str}_temperature{temperature_str}"  ####################
+log.register_file_output(dump_dir+job_name+".log", mode="w")
+
+
+nmodes = 1000
+Ms = 20
+upper_limit = 30
+sdf = ColeDavidsonSDF(ita, omega_c, beta, upper_limit)
+
+w, c2 = sdf.Wang1(nmodes)
+c = np.sqrt(c2)
+logger.info(f"w:{w}")
+logger.info(f"c:{c}")
+
+reno = sdf.reno(w[-1])
+logger.info(f"renormalization constant: {reno}")
+Delta *= reno
+
+ham_terms = []
+
+# h_s
+ham_terms.extend([Op("sigma_z","spin",factor=eps, qn=0),
+        Op("sigma_x","spin",factor=Delta, qn=0)])
+
+
+# boson energy
+for imode in range(nmodes):
+    op1 = Op(r"p^2",f"v_{imode}_p",factor=0.5, qn=0)
+    op2 = Op(r"x^2",f"v_{imode}_p",factor=0.5*w[imode]**2, qn=0)
+    ham_terms.extend([op1,op2])
+    op1 = Op(r"p^2",f"v_{imode}_q",factor=-0.5, qn=0)
+    op2 = Op(r"x^2",f"v_{imode}_q",factor=-0.5*w[imode]**2, qn=0)
+    ham_terms.extend([op1,op2])
+
+theta_array = np.arctanh(np.exp(-w/ temperature / 2))
+# system-boson coupling
+for imode in range(nmodes):
+    theta = theta_array[imode]
+    op = Op(r"sigma_z x", ["spin", f"v_{imode}_p"], factor=np.cosh(theta) * c[imode], qn=[0,0])
+    ham_terms.append(op)
+    op = Op(r"sigma_z x", ["spin", f"v_{imode}_q"], factor=np.sinh(theta) * c[imode], qn=[0,0])
+    ham_terms.append(op)
+
+
+nbas_factor = 2
+nbas = np.max([16 * c2/w**3 * np.cosh(theta_array)**2, np.ones(nmodes)*4], axis=0)
+nbas = np.min([nbas, np.ones(nmodes)*512], axis=0)
+nbas = np.round(nbas).astype(int)
+nbas *= nbas_factor
+logger.info(nbas)
+basis = [ba.BasisHalfSpin("spin",[0,0])]
+for imode in range(nmodes):
+    basis.append(ba.BasisSHO(f"v_{imode}_p", w[imode], int(nbas[imode])))
+    basis.append(ba.BasisSHO(f"v_{imode}_q", w[imode], int(nbas[imode])))
+
+
+tree_order = 2
+basis_vib = basis[1:]
+elementary_nodes = []
+
+labels = np.array([[nbas>Ms], [nbas>Ms]]).T.ravel()
+root = BasisTree.binary_mctdh(basis_vib, contract_primitive=True, contract_label=labels, dummy_label="n").root
+
+root.add_child(TreeNodeBasis(basis[:1]))
+
+basis_tree = BasisTree(root)
+basis_tree.print(print_function=logger.info)
+
+# basis_tree = BasisTree.linear(basis)
+ttno = TTNO(basis_tree, ham_terms)
+exp_z = TTNO(basis_tree, Op("sigma_z", "spin"))
+exp_x = TTNO(basis_tree, Op("sigma_x", "spin"))
+ttns = TTNS(basis_tree)
+ttns.compress_config = CompressConfig(CompressCriteria.fixed, max_bonddim=Ms)
+ttns = expand_bond_dimension_general(ttns, ttno, ex_mps=None)
+logger.info(ttns.bond_dims)
+logger.info(ttno.bond_dims)
+logger.info(len(ttns))
+ttns.evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
+nsteps = 400
+dt = 0.1
+expectations = []
+for i in range(nsteps):
+    ttns = ttns.evolve(ttno, dt)
+    z = ttns.expectation(exp_z)
+    x = ttns.expectation(exp_x)
+    expectations.append((z, x))
+    logger.info((z, x))
+logger.info(expectations)
\ No newline at end of file
diff --git a/example/ttns/sbm_zt.py b/example/ttns/sbm_zt.py
new file mode 100644
index 00000000..0d9bd3b8
--- /dev/null
+++ b/example/ttns/sbm_zt.py
@@ -0,0 +1,107 @@
+import sys
+
+from renormalizer.model import Op
+from renormalizer.model import basis as ba
+from renormalizer.mps.mps import expand_bond_dimension_general
+from renormalizer.sbm import ColeDavidsonSDF
+from renormalizer.utils import EvolveConfig, CompressConfig, CompressCriteria, EvolveMethod
+from renormalizer.utils import log
+from renormalizer.tn import BasisTree, TTNO, TTNS, TreeNodeBasis
+
+import numpy as np
+
+ita_str = sys.argv[1]  # 010, 025, 100
+omega_c_str = sys.argv[2]  # 001, 100
+beta_str = sys.argv[3]  # 025, 100
+# ita_str = "050"
+# omega_c_str = "001"
+# beta_str = "050"
+
+ita = int(ita_str) / 10 # 1, 2.5, 5, 10
+eps = 0
+Delta = 1
+omega_c = int(omega_c_str) / 10  # 0.1, 1, 10
+
+beta = int(beta_str) / 100  # 0.25, 0.5, 0.75, 1
+
+from renormalizer.utils.log import package_logger
+logger = package_logger
+dump_dir = "./"
+job_name = f"ps1_binary_ita{ita_str}_omega{omega_c_str}_beta{beta_str}"  ####################
+log.register_file_output(dump_dir+job_name+".log", mode="w")
+
+
+nmodes = 1000
+Ms = 20
+upper_limit = 30
+sdf = ColeDavidsonSDF(ita, omega_c, beta, upper_limit)
+
+w, c2 = sdf.Wang1(nmodes)
+c = np.sqrt(c2)
+logger.info(f"w:{w}")
+logger.info(f"c:{c}")
+
+reno = sdf.reno(w[-1])
+logger.info(f"renormalization constant: {reno}")
+Delta *= reno
+
+ham_terms = []
+
+# h_s
+ham_terms.extend([Op("sigma_z","spin",factor=eps, qn=0),
+        Op("sigma_x","spin",factor=Delta, qn=0)])
+
+
+# boson energy
+for imode in range(nmodes):
+    op1 = Op(r"p^2",f"v_{imode}",factor=0.5, qn=0)
+    op2 = Op(r"x^2",f"v_{imode}",factor=0.5*w[imode]**2, qn=0)
+    ham_terms.extend([op1,op2])
+
+# system-boson coupling
+for imode in range(nmodes):
+    op = Op(r"sigma_z x", ["spin", f"v_{imode}"],
+            factor=c[imode], qn=[0,0])
+    ham_terms.append(op)
+
+nbas = np.max([16 * c2/w**3, np.ones(nmodes)*4], axis=0)
+nbas = np.round(nbas).astype(int)
+logger.info(nbas)
+basis = [ba.BasisHalfSpin("spin",[0,0])]
+for imode in range(nmodes):
+    basis.append(ba.BasisSHO(f"v_{imode}", w[imode], int(nbas[imode])))
+
+
+tree_order = 2
+basis_vib = basis[1:]
+elementary_nodes = []
+
+
+root = BasisTree.binary_mctdh(basis_vib, contract_primitive=True, contract_label=nbas>Ms, dummy_label="n").root
+
+root.add_child(TreeNodeBasis(basis[:1]))
+
+basis_tree = BasisTree(root)
+basis_tree.print(print_function=logger.info)
+
+# basis_tree = BasisTree.linear(basis)
+ttno = TTNO(basis_tree, ham_terms)
+exp_z = TTNO(basis_tree, Op("sigma_z", "spin"))
+exp_x = TTNO(basis_tree, Op("sigma_x", "spin"))
+ttns = TTNS(basis_tree)
+ttns.compress_config = CompressConfig(CompressCriteria.fixed, max_bonddim=Ms)
+ttns = expand_bond_dimension_general(ttns, ttno, ex_mps=None)
+logger.info(ttns.bond_dims)
+logger.info(ttno.bond_dims)
+logger.info(len(ttns))
+ttns.evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
+nsteps = 200
+dt = 0.2
+expectations = []
+for i in range(nsteps):
+    ttns = ttns.evolve(ttno, dt)
+    z = ttns.expectation(exp_z)
+    x = ttns.expectation(exp_x)
+    expectations.append((z, x))
+    logger.info((z, x))
+logger.info(expectations)
\ No newline at end of file
diff --git a/renormalizer/__init__.py b/renormalizer/__init__.py
index c15c9db4..768846a2 100644
--- a/renormalizer/__init__.py
+++ b/renormalizer/__init__.py
@@ -7,6 +7,9 @@
 
 
 reno_num_threads = os.environ.get("RENO_NUM_THREADS")
+# for debugging
+# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+# reno_num_threads = "1"
 if reno_num_threads is not None:
     # Set environment variables to limit NumPy CPU usage
     # Note that this should be done before NumPy is imported
@@ -32,8 +35,8 @@
 
 
 # user interfaces
-from renormalizer.model import Model, HolsteinModel, SpinBosonModel, TI1DModel, Op
+from renormalizer.model import Model, HolsteinModel, SpinBosonModel, TI1DModel, Op, OpSum
 from renormalizer.model.basis import BasisSHO, BasisHopsBoson, BasisSineDVR, BasisMultiElectron, \
-    BasisMultiElectronVac, BasisSimpleElectron, BasisHalfSpin
+    BasisMultiElectronVac, BasisSimpleElectron, BasisHalfSpin, BasisDummy
 from renormalizer.mps import Mpo, Mps, optimize_mps
 from renormalizer.utils import Quantity
diff --git a/renormalizer/lib/davidson/CHANGES b/renormalizer/lib/davidson/CHANGES
deleted file mode 100644
index 422d42b7..00000000
--- a/renormalizer/lib/davidson/CHANGES
+++ /dev/null
@@ -1,4 +0,0 @@
-CHANGES TO THE ORIGINAL CODE
-
-- Code Style has been changed a bit with `Black`.
-- Support Cupy array as `x0`.
\ No newline at end of file
diff --git a/renormalizer/lib/davidson/davidson.py b/renormalizer/lib/davidson/davidson.py
index 3e2954d9..c30607b9 100644
--- a/renormalizer/lib/davidson/davidson.py
+++ b/renormalizer/lib/davidson/davidson.py
@@ -1,4 +1,5 @@
-# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
+#!/usr/bin/env python
+# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,11 +16,26 @@
 # Author: Qiming Sun <osirpt.sun@gmail.com>
 #
 
+'''
+Extension to scipy.linalg module
+'''
+
+import sys
+import warnings
+import tempfile
+
 import numpy
+np = numpy
+import scipy.linalg
 import h5py
-import tempfile
-import scipy
 
+from renormalizer.lib.davidson import logger
+
+
+SAFE_EIGH_LINDEP = 1e-15
+DAVIDSON_LINDEP = 1e-14
+DSOLVE_LINDEP = 1e-15
+MAX_MEMORY = 2000
 
 # sort by similarity has problem which flips the ordering of eigenvalues when
 # the initial guess is closed to excited state.  In this situation, function
@@ -31,77 +47,214 @@
 # following iterations.  Projecting out the converged eigenvectors may lead to
 # large errors to the yet converged eigenvectors.
 PROJECT_OUT_CONV_EIGS = False
-# default max_memory 2000 MB
-def davidson(
-    aop,
-    x0,
-    precond,
-    tol=1e-12,
-    max_cycle=50,
-    max_space=12,
-    lindep=1e-14,
-    max_memory=2000,
-    dot=numpy.dot,
-    callback=None,
-    nroots=1,
-    lessio=False,
-    follow_state=False,
-):
-    e, x = davidson1(
-        lambda xs: [aop(x) for x in xs],
-        x0,
-        precond,
-        tol,
-        max_cycle,
-        max_space,
-        lindep,
-        max_memory,
-        dot,
-        callback,
-        nroots,
-        lessio,
-        follow_state,
-    )[1:]
+
+FOLLOW_STATE = False
+
+
+def _fill_heff_hermitian(heff, xs, ax, xt, axt, dot):
+    nrow = len(axt)
+    row1 = len(ax)
+    row0 = row1 - nrow
+    for ip, i in enumerate(range(row0, row1)):
+        for jp, j in enumerate(range(row0, i)):
+            heff[i,j] = dot(xt[ip].conj(), axt[jp])
+            heff[j,i] = heff[i,j].conj()
+        heff[i,i] = dot(xt[ip].conj(), axt[ip]).real
+
+    for i in range(row0):
+        axi = numpy.asarray(ax[i])
+        for jp, j in enumerate(range(row0, row1)):
+            heff[j,i] = dot(xt[jp].conj(), axi)
+            heff[i,j] = heff[j,i].conj()
+        axi = None
+    return heff
+
+
+def davidson(aop, x0, precond, tol=1e-12, max_cycle=50, max_space=12,
+             lindep=DAVIDSON_LINDEP, max_memory=MAX_MEMORY,
+             dot=numpy.dot, callback=None,
+             nroots=1, lessio=False, pick=None, verbose=logger.WARN,
+             follow_state=FOLLOW_STATE):
+    r'''Davidson diagonalization method to solve  a c = e c.  Ref
+    [1] E.R. Davidson, J. Comput. Phys. 17 (1), 87-94 (1975).
+    [2] http://people.inf.ethz.ch/arbenz/ewp/Lnotes/chapter11.pdf
+
+    Note: This function has an overhead of memory usage ~4*x0.size*nroots
+
+    Args:
+        aop : function(x) => array_like_x
+            aop(x) to mimic the matrix vector multiplication :math:`\sum_{j}a_{ij}*x_j`.
+            The argument is a 1D array.  The returned value is a 1D array.
+        x0 : 1D array or a list of 1D array
+            Initial guess.  The initial guess vector(s) are just used as the
+            initial subspace bases.  If the subspace is smaller than "nroots",
+            eg 10 roots and one initial guess, all eigenvectors are chosen as
+            the eigenvectors during the iterations.  The first iteration has
+            one eigenvector, the next iteration has two, the third iterastion
+            has 4, ..., until the subspace size > nroots.
+        precond : diagonal elements of the matrix or  function(dx, e, x0) => array_like_dx
+            Preconditioner to generate new trial vector.
+            The argument dx is a residual vector ``a*x0-e*x0``; e is the current
+            eigenvalue; x0 is the current eigenvector.
+
+    Kwargs:
+        tol : float
+            Convergence tolerance.
+        max_cycle : int
+            max number of iterations.
+        max_space : int
+            space size to hold trial vectors.
+        lindep : float
+            Linear dependency threshold.  The function is terminated when the
+            smallest eigenvalue of the metric of the trial vectors is lower
+            than this threshold.
+        max_memory : int or float
+            Allowed memory in MB.
+        dot : function(x, y) => scalar
+            Inner product
+        callback : function(envs_dict) => None
+            callback function takes one dict as the argument which is
+            generated by the builtin function :func:`locals`, so that the
+            callback function can access all local variables in the current
+            envrionment.
+        nroots : int
+            Number of eigenvalues to be computed.  When nroots > 1, it affects
+            the shape of the return value
+        lessio : bool
+            How to compute a*x0 for current eigenvector x0.  There are two
+            ways to compute a*x0.  One is to assemble the existed a*x.  The
+            other is to call aop(x0).  The default is the first method which
+            needs more IO and less computational cost.  When IO is slow, the
+            second method can be considered.
+        pick : function(w,v,nroots) => (e[idx], w[:,idx], idx)
+            Function to filter eigenvalues and eigenvectors.
+        follow_state : bool
+            If the solution dramatically changes in two iterations, clean the
+            subspace and restart the iteration with the old solution.  It can
+            help to improve numerical stability.  Default is False.
+
+    Returns:
+        e : float or list of floats
+            Eigenvalue.  By default it's one float number.  If :attr:`nroots` > 1, it
+            is a list of floats for the lowest :attr:`nroots` eigenvalues.
+        c : 1D array or list of 1D arrays
+            Eigenvector.  By default it's a 1D array.  If :attr:`nroots` > 1, it
+            is a list of arrays for the lowest :attr:`nroots` eigenvectors.
+    '''
+    e, x = davidson1(lambda xs: [aop(x) for x in xs],
+                     x0, precond, tol, max_cycle, max_space, lindep,
+                     max_memory, dot, callback, nroots, lessio, pick, verbose,
+                     follow_state)[1:]
     if nroots == 1:
         return e[0], x[0]
     else:
         return e, x
 
 
-def davidson1(
-    aop,
-    x0,
-    precond,
-    tol=1e-12,
-    max_cycle=50,
-    max_space=12,
-    lindep=1e-14,
-    max_memory=2000,
-    dot=numpy.dot,
-    callback=None,
-    nroots=1,
-    lessio=False,
-    follow_state=False,
-):
-
-    toloose = numpy.sqrt(tol)
-    # print('tol %g  toloose %g', tol, toloose)
-
-    if (not isinstance(x0, list)) and x0.ndim == 1:
+def davidson1(aop, x0, precond, tol=1e-12, max_cycle=50, max_space=12,
+              lindep=DAVIDSON_LINDEP, max_memory=MAX_MEMORY,
+              dot=numpy.dot, callback=None,
+              nroots=1, lessio=False, pick=None, verbose=logger.WARN,
+              follow_state=FOLLOW_STATE, tol_residual=None,
+              fill_heff=_fill_heff_hermitian):
+    r'''Davidson diagonalization method to solve  a c = e c.  Ref
+    [1] E.R. Davidson, J. Comput. Phys. 17 (1), 87-94 (1975).
+    [2] http://people.inf.ethz.ch/arbenz/ewp/Lnotes/chapter11.pdf
+
+    Note: This function has an overhead of memory usage ~4*x0.size*nroots
+
+    Args:
+        aop : function([x]) => [array_like_x]
+            Matrix vector multiplication :math:`y_{ki} = \sum_{j}a_{ij}*x_{jk}`.
+        x0 : 1D array or a list of 1D arrays or a function to generate x0 array(s)
+            Initial guess.  The initial guess vector(s) are just used as the
+            initial subspace bases.  If the subspace is smaller than "nroots",
+            eg 10 roots and one initial guess, all eigenvectors are chosen as
+            the eigenvectors during the iterations.  The first iteration has
+            one eigenvector, the next iteration has two, the third iterastion
+            has 4, ..., until the subspace size > nroots.
+        precond : diagonal elements of the matrix or  function(dx, e, x0) => array_like_dx
+            Preconditioner to generate new trial vector.
+            The argument dx is a residual vector ``a*x0-e*x0``; e is the current
+            eigenvalue; x0 is the current eigenvector.
+
+    Kwargs:
+        tol : float
+            Convergence tolerance.
+        max_cycle : int
+            max number of iterations.
+        max_space : int
+            space size to hold trial vectors.
+        lindep : float
+            Linear dependency threshold.  The function is terminated when the
+            smallest eigenvalue of the metric of the trial vectors is lower
+            than this threshold.
+        max_memory : int or float
+            Allowed memory in MB.
+        dot : function(x, y) => scalar
+            Inner product
+        callback : function(envs_dict) => None
+            callback function takes one dict as the argument which is
+            generated by the builtin function :func:`locals`, so that the
+            callback function can access all local variables in the current
+            envrionment.
+        nroots : int
+            Number of eigenvalues to be computed.  When nroots > 1, it affects
+            the shape of the return value
+        lessio : bool
+            How to compute a*x0 for current eigenvector x0.  There are two
+            ways to compute a*x0.  One is to assemble the existed a*x.  The
+            other is to call aop(x0).  The default is the first method which
+            needs more IO and less computational cost.  When IO is slow, the
+            second method can be considered.
+        pick : function(w,v,nroots) => (e[idx], w[:,idx], idx)
+            Function to filter eigenvalues and eigenvectors.
+        follow_state : bool
+            If the solution dramatically changes in two iterations, clean the
+            subspace and restart the iteration with the old solution.  It can
+            help to improve numerical stability.  Default is False.
+
+    Returns:
+        conv : bool
+            Converged or not
+        e : list of floats
+            The lowest :attr:`nroots` eigenvalues.
+        c : list of 1D arrays
+            The lowest :attr:`nroots` eigenvectors.
+    '''
+    if isinstance(verbose, logger.Logger):
+        log = verbose
+    else:
+        log = logger.Logger(sys.stdout, verbose)
+
+    if tol_residual is None:
+        toloose = numpy.sqrt(tol)
+    else:
+        toloose = tol_residual
+    log.debug1('tol %g  toloose %g', tol, toloose)
+
+    if not callable(precond):
+        precond = make_diag_precond(precond)
+
+    if callable(x0):  # lazy initialization to reduce memory footprint
+        x0 = x0()
+    if isinstance(x0, numpy.ndarray) and x0.ndim == 1:
         x0 = [x0]
-    # max_cycle = min(max_cycle, x0[0].size)
-    max_space = min(max_space + nroots * 3, max(1,x0[0].size//2))
+    #max_cycle = min(max_cycle, x0[0].size)
+    max_space = max_space + (nroots-1) * 3
     # max_space*2 for holding ax and xs, nroots*2 for holding axt and xt
-    _incore = max_memory * 1e6 / x0[0].nbytes > max_space * 2 + nroots * 3
+    _incore = max_memory*1e6/x0[0].nbytes > max_space*2+nroots*3
     lessio = lessio and not _incore
-    # print('max_cycle %d  max_space %d  max_memory %d  incore %s',
-    #           max_cycle, max_space, max_memory, _incore)
+    log.debug1('max_cycle %d  max_space %d  max_memory %d  incore %s',
+               max_cycle, max_space, max_memory, _incore)
+    dtype = None
     heff = None
     fresh_start = True
     e = 0
     v = None
     conv = [False] * nroots
     emin = None
+    norm_min = 1
 
     for icyc in range(max_cycle):
         if fresh_start:
@@ -112,15 +265,32 @@ def davidson1(
                 xs = _Xlist()
                 ax = _Xlist()
             space = 0
-            # Orthogonalize xt space because the basis of subspace xs must be orthogonal
-            # but the eigenvectors x0 might not be strictly orthogonal
+# Orthogonalize xt space because the basis of subspace xs must be orthogonal
+# but the eigenvectors x0 might not be strictly orthogonal
             xt = None
-            xt, x0 = _qr(x0, dot), None
+            x0len = len(x0)
+            xt = _qr(x0, dot, lindep)[0]
+            if len(xt) != x0len:
+                log.warn('QR decomposition removed %d vectors.  The davidson may fail.',
+                         x0len - len(xt))
+                if callable(pick):
+                    log.warn('Check to see if `pick` function %s is providing '
+                             'linear dependent vectors', pick.__name__)
+                if len(xt) == 0:
+                    if icyc == 0:
+                        msg = 'Initial guess is empty or zero'
+                    else:
+                        msg = ('No more linearly independent basis were found. '
+                               'Unless loosen the lindep tolerance (current value '
+                               f'{lindep}), the diagonalization solver is not able '
+                               'to find eigenvectors.')
+                    raise LinearDependenceError(msg)
+            x0 = None
             max_dx_last = 1e9
             if SORT_EIG_BY_SIMILARITY:
-                conv = numpy.array([False] * nroots)
+                conv = [False] * nroots
         elif len(xt) > 1:
-            xt = _qr(xt, dot)
+            xt = _qr(xt, dot, lindep)[0]
             xt = xt[:40]  # 40 trial vectors at most
 
         axt = aop(xt)
@@ -128,39 +298,35 @@ def davidson1(
             xs.append(xt[k])
             ax.append(axt[k])
         rnow = len(xt)
-        head, space = space, space + rnow
-
+        head, space = space, space+rnow
+
+        if dtype is None:
+            try:
+                dtype = numpy.result_type(axt[0], xt[0])
+            except IndexError:
+                raise LinearDependenceError('No linearly independent basis found '
+                                            'by the diagonalization solver.')
         if heff is None:  # Lazy initilize heff to determine the dtype
-            heff = numpy.empty(
-                (max_space + nroots, max_space + nroots), dtype=ax[0].dtype
-            )
+            heff = numpy.empty((max_space+nroots,max_space+nroots), dtype=dtype)
         else:
-            heff = numpy.asarray(heff, dtype=ax[0].dtype)
+            heff = numpy.asarray(heff, dtype=dtype)
 
         elast = e
         vlast = v
         conv_last = conv
-        for i in range(space):
-            if head <= i < head + rnow:
-                for k in range(i - head + 1):
-                    heff[head + k, i] = dot(xt[k].conj(), axt[i - head])
-                    heff[i, head + k] = heff[head + k, i].conj()
-            else:
-                for k in range(rnow):
-                    heff[head + k, i] = dot(xt[k].conj(), ax[i])
-                    heff[i, head + k] = heff[head + k, i].conj()
 
-        w, v = scipy.linalg.eigh(heff[:space, :space])
-        if SORT_EIG_BY_SIMILARITY:
-            e, v = _sort_by_similarity(w, v, nroots, conv, vlast, emin)
-            if elast.size != e.size:
-                de = e
-            else:
-                de = e - elast
-        else:
-            e = w[:nroots]
-            v = v[:, :nroots]
+        fill_heff(heff, xs, ax, xt, axt, dot)
+        xt = axt = None
+        w, v = scipy.linalg.eigh(heff[:space,:space])
+        if callable(pick):
+            w, v, idx = pick(w, v, nroots, locals())
+            if len(w) == 0:
+                raise RuntimeError(f'Not enough eigenvalues found by {pick}')
 
+        e = w[:nroots]
+        v = v[:,:nroots]
+
+        x0 = None
         x0 = _gen_x0(v, xs)
         if lessio:
             ax0 = aop(x0)
@@ -175,11 +341,12 @@ def davidson1(
                     xt[k] = ax0[k] - ek * x0[k]
                     dx_norm[k] = numpy.sqrt(dot(xt[k].conj(), xt[k]).real)
                     if abs(de[k]) < tol and dx_norm[k] < toloose:
-                        # print('root %d converged  |r|= %4.3g  e= %s  max|de|= %4.3g',
-                        #          k, dx_norm[k], ek, de[k])
+                        log.debug('root %d converged  |r|= %4.3g  e= %s  max|de|= %4.3g',
+                                  k, dx_norm[k], ek, de[k])
                         conv[k] = True
         else:
-            elast, conv_last = _sort_elast(elast, conv_last, vlast, v, fresh_start)
+            elast, conv_last = _sort_elast(elast, conv_last, vlast, v,
+                                           fresh_start, log)
             de = e - elast
             dx_norm = []
             xt = []
@@ -188,25 +355,21 @@ def davidson1(
                 xt.append(ax0[k] - ek * x0[k])
                 dx_norm.append(numpy.sqrt(dot(xt[k].conj(), xt[k]).real))
                 conv[k] = abs(de[k]) < tol and dx_norm[k] < toloose
-                # if conv[k] and not conv_last[k]:
-                #    print('root %d converged  |r|= %4.3g  e= %s  max|de|= %4.3g',
-                #              k, dx_norm[k], ek, de[k])
+                if conv[k] and not conv_last[k]:
+                    log.debug('root %d converged  |r|= %4.3g  e= %s  max|de|= %4.3g',
+                              k, dx_norm[k], ek, de[k])
         ax0 = None
         max_dx_norm = max(dx_norm)
         ide = numpy.argmax(abs(de))
         if all(conv):
-            # print('converge %d %d  |r|= %4.3g  e= %s  max|de|= %4.3g',
-            #          icyc, space, max_dx_norm, e, de[ide])
+            log.debug('converged %d %d  |r|= %4.3g  e= %s  max|de|= %4.3g',
+                      icyc, space, max_dx_norm, e, de[ide])
             break
-        elif (
-            follow_state
-            and max_dx_norm > 1
-            and max_dx_norm / max_dx_last > 3
-            and space > nroots * 1
-        ):
-            # print('davidson %d %d  |r|= %4.3g  e= %s  max|de|= %4.3g  lindep= %4.3g',
-            #          icyc, space, max_dx_norm, e, de[ide], norm_min)
-            # print('Large |r| detected, restore to previous x0')
+        elif (follow_state and max_dx_norm > 1 and
+              max_dx_norm/max_dx_last > 3 and space > nroots+2):
+            log.debug('davidson %d %d  |r|= %4.3g  e= %s  max|de|= %4.3g  lindep= %4.3g',
+                      icyc, space, max_dx_norm, e, de[ide], norm_min)
+            log.debug('Large |r| detected, restore to previous x0')
             x0 = _gen_x0(vlast, xs)
             fresh_start = True
             continue
@@ -216,115 +379,167 @@ def davidson1(
                 emin = min(e)
 
         # remove subspace linear dependency
-        if any(((not conv[k]) and n ** 2 > lindep) for k, n in enumerate(dx_norm)):
+        if any(((not conv[k]) and n**2>lindep) for k, n in enumerate(dx_norm)):
             for k, ek in enumerate(e):
-                if (not conv[k]) and dx_norm[k] ** 2 > lindep:
+                if (not conv[k]) and dx_norm[k]**2 > lindep:
                     xt[k] = precond(xt[k], e[0], x0[k])
-                    xt[k] *= 1 / numpy.sqrt(dot(xt[k].conj(), xt[k]).real)
+                    xt[k] *= 1/numpy.sqrt(dot(xt[k].conj(), xt[k]).real)
                 else:
                     xt[k] = None
         else:
             for k, ek in enumerate(e):
-                if dx_norm[k] ** 2 > lindep:
+                if dx_norm[k]**2 > lindep:
                     xt[k] = precond(xt[k], e[0], x0[k])
-                    xt[k] *= 1 / numpy.sqrt(dot(xt[k].conj(), xt[k]).real)
+                    xt[k] *= 1/numpy.sqrt(dot(xt[k].conj(), xt[k]).real)
                 else:
                     xt[k] = None
+                    log.debug1('Throwing out eigenvector %d with norm=%4.3g', k, dx_norm[k])
         xt = [xi for xi in xt if xi is not None]
 
         for i in range(space):
-            xsi = xs[i]
+            xsi = numpy.asarray(xs[i])
             for xi in xt:
                 xi -= xsi * dot(xsi.conj(), xi)
+            xsi = None
         norm_min = 1
-        for i, xi in enumerate(xt):
+        for i,xi in enumerate(xt):
             norm = numpy.sqrt(dot(xi.conj(), xi).real)
-            if norm ** 2 > lindep:
-                xt[i] *= 1 / norm
+            if norm**2 > lindep:
+                xt[i] *= 1/norm
                 norm_min = min(norm_min, norm)
             else:
                 xt[i] = None
         xt = [xi for xi in xt if xi is not None]
         xi = None
-        # print('davidson %d %d  |r|= %4.3g  e= %s  max|de|= %4.3g  lindep= %4.3g',
-        #          icyc, space, max_dx_norm, e, de[ide], norm_min)
+        log.debug('davidson %d %d  |r|= %4.3g  e= %s  max|de|= %4.3g  lindep= %4.3g',
+                  icyc, space, max_dx_norm, e, de[ide], norm_min)
         if len(xt) == 0:
-            # print('Linear dependency in trial subspace. |r| for each state %s', dx_norm)
-            conv = [conv[k] or (norm < toloose) for k, norm in enumerate(dx_norm)]
+            log.debug('Linear dependency in trial subspace. |r| for each state %s',
+                      dx_norm)
+            conv = [conv[k] or (norm < toloose) for k,norm in enumerate(dx_norm)]
             break
 
         max_dx_last = max_dx_norm
-        fresh_start = space + nroots > max_space
+        fresh_start = space+nroots > max_space
 
         if callable(callback):
             callback(locals())
 
-    return conv, e, x0
-
-
-def _qr(xs, dot):
-    norm = numpy.sqrt(dot(xs[0].conj(), xs[0]).real)
-    qs = [xs[0] / norm]
-    for i in range(1, len(xs)):
-        xi = xs[i].copy()
-        for j in range(len(qs)):
-            xi -= qs[j] * dot(qs[j].conj(), xi)
-        norm = numpy.sqrt(dot(xi.conj(), xi).real)
-        if norm > 1e-7:
-            qs.append(xi / norm)
-    return qs
-
+    x0 = [x for x in x0]  # nparray -> list
+
+    # Check whether the solver finds enough eigenvectors.
+    h_dim = x0[0].size
+    if len(x0) < min(h_dim, nroots):
+        # Two possible reasons:
+        # 1. All the initial guess are the eigenvectors. No more trial vectors
+        # can be generated.
+        # 2. The initial guess sits in the subspace which is smaller than the
+        # required number of roots.
+        msg = 'Not enough eigenvectors (len(x0)=%d, nroots=%d)' % (len(x0), nroots)
+        warnings.warn(msg)
+
+    return numpy.asarray(conv), e, x0
+
+
+def make_diag_precond(diag, level_shift=0):
+    '''Generate the preconditioner function with the diagonal function.'''
+    def precond(dx, e, *args):
+        diagd = diag - (e - level_shift)
+        diagd[abs(diagd)<1e-8] = 1e-8
+        return dx/diagd
+    return precond
+
+
+def _qr(xs, dot, lindep=1e-14):
+    '''QR decomposition for a list of vectors (for linearly independent vectors only).
+    xs = (r.T).dot(qs)
+    '''
+    nvec = len(xs)
+    dtype = xs[0].dtype
+    qs = numpy.empty((nvec,xs[0].size), dtype=dtype)
+    rmat = numpy.empty((nvec,nvec), order='F', dtype=dtype)
+
+    nv = 0
+    for i in range(nvec):
+        xi = numpy.array(xs[i], copy=True)
+        rmat[:,nv] = 0
+        rmat[nv,nv] = 1
+        for j in range(nv):
+            prod = dot(qs[j].conj(), xi)
+            xi -= qs[j] * prod
+            rmat[:,nv] -= rmat[:,j] * prod
+        innerprod = dot(xi.conj(), xi).real
+        norm = numpy.sqrt(innerprod)
+        if innerprod > lindep:
+            qs[nv] = xi/norm
+            rmat[:nv+1,nv] /= norm
+            nv += 1
+    return qs[:nv], numpy.linalg.inv(rmat[:nv,:nv])
 
 def _gen_x0(v, xs):
     space, nroots = v.shape
-    x0 = []
-    for k in range(nroots):
-        x0.append(xs[space - 1] * v[space - 1, k])
-    for i in reversed(range(space - 1)):
-        xsi = xs[i]
+    x0 = numpy.einsum('c,x->cx', v[space-1], numpy.asarray(xs[space-1]))
+    for i in reversed(range(space-1)):
+        xsi = numpy.asarray(xs[i])
         for k in range(nroots):
-            x0[k] += v[i, k] * xsi
+            x0[k] += v[i,k] * xsi
     return x0
 
 
-def _sort_by_similarity(w, v, nroots, conv, vlast, emin=None, heff=None):
-    if not any(conv) or vlast is None:
-        return w[:nroots], v[:, :nroots]
-
-    head, nroots = vlast.shape
-    conv = numpy.asarray(conv[:nroots])
-    ovlp = vlast[:, conv].T.conj().dot(v[:head])
-    ovlp = numpy.einsum("ij,ij->j", ovlp, ovlp)
-    nconv = numpy.count_nonzero(conv)
-    nleft = nroots - nconv
-    idx = ovlp.argsort()
-    sorted_idx = numpy.zeros(nroots, dtype=int)
-    sorted_idx[conv] = numpy.sort(idx[-nconv:])
-    sorted_idx[~conv] = numpy.sort(idx[:-nconv])[:nleft]
-
-    e = w[sorted_idx]
-    c = v[:, sorted_idx]
-    return e, c
-
-
-def _sort_elast(elast, conv_last, vlast, v, fresh_start):
+def _sort_elast(elast, conv_last, vlast, v, fresh_start, log):
+    '''
+    Eigenstates may be flipped during the Davidson iterations.  Reorder the
+    eigenvalues of last iteration to make them comparable to the eigenvalues
+    of the current iterations.
+    '''
     if fresh_start:
         return elast, conv_last
     head, nroots = vlast.shape
     ovlp = abs(numpy.dot(v[:head].conj().T, vlast))
     idx = numpy.argmax(ovlp, axis=1)
+
+    if log.verbose >= logger.DEBUG:
+        ordering_diff = (idx != numpy.arange(len(idx)))
+        if numpy.any(ordering_diff):
+            log.debug('Old state -> New state')
+            for i in numpy.where(ordering_diff)[0]:
+                log.debug('  %3d     ->   %3d ', idx[i], i)
+
     return [elast[i] for i in idx], [conv_last[i] for i in idx]
 
 
+class LinearDependenceError(RuntimeError):
+    pass
+
+
 class H5TmpFile(h5py.File):
-    def __init__(self, filename=None, *args, **kwargs):
+    '''Create and return an HDF5 temporary file.
+    Kwargs:
+        filename : str or None
+            If a string is given, an HDF5 file of the given filename will be
+            created. The temporary file will exist even if the H5TmpFile
+            object is released.  If nothing is specified, the HDF5 temporary
+            file will be deleted when the H5TmpFile object is released.
+    The return object is an h5py.File object. The file will be automatically
+    deleted when it is closed or the object is released (unless filename is
+    specified).
+    '''
+    def __init__(self, filename=None, mode='a', *args, **kwargs):
         if filename is None:
             tmpfile = tempfile.NamedTemporaryFile(dir=".")
             filename = tmpfile.name
-        h5py.File.__init__(self, filename, *args, **kwargs)
-
+        h5py.File.__init__(self, filename, mode, *args, **kwargs)
+#FIXME: Does GC flush/close the HDF5 file when releasing the resource?
+# To make HDF5 file reusable, file has to be closed or flushed
     def __del__(self):
-        self.close()
+        try:
+            self.close()
+        except AttributeError:  # close not defined in old h5py
+            pass
+        except ValueError:  # if close() is called twice
+            pass
+        except ImportError:  # exit program before de-referring the object
+            pass
 
 
 class _Xlist(list):
@@ -334,22 +549,26 @@ def __init__(self):
 
     def __getitem__(self, n):
         key = self.index[n]
-        return self.scr_h5[key].value
+        return self.scr_h5[str(key)]
 
     def append(self, x):
-        key = str(len(self.index) + 1)
-        if key in self.index:
-            for i in range(len(self.index) + 1):
-                if str(i) not in self.index:
-                    key = str(i)
-                    break
+        length = len(self.index)
+        key = length + 1
+        index_set = set(self.index)
+        if key in index_set:
+            key = set(range(length)).difference(index_set).pop()
         self.index.append(key)
-        self.scr_h5[key] = x
+
+        self.scr_h5[str(key)] = x
         self.scr_h5.flush()
 
+    def extend(self, x):
+        for xi in x:
+            self.append(xi)
+
     def __setitem__(self, n, x):
         key = self.index[n]
-        self.scr_h5[key][:] = x
+        self.scr_h5[str(key)][:] = x
         self.scr_h5.flush()
 
     def __len__(self):
@@ -357,4 +576,6 @@ def __len__(self):
 
     def pop(self, index):
         key = self.index.pop(index)
-        del (self.scr_h5[key])
+        del (self.scr_h5[str(key)])
+
+del (SAFE_EIGH_LINDEP, DAVIDSON_LINDEP, DSOLVE_LINDEP, MAX_MEMORY)
diff --git a/renormalizer/lib/davidson/logger.py b/renormalizer/lib/davidson/logger.py
new file mode 100644
index 00000000..ee689ac2
--- /dev/null
+++ b/renormalizer/lib/davidson/logger.py
@@ -0,0 +1,212 @@
+#!/usr/bin/env python
+# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Author: Qiming Sun <osirpt.sun@gmail.com>
+#
+
+
+'''
+Logging system
+Log level
+---------
+======= ======
+Level   number
+------- ------
+DEBUG4  9
+DEBUG3  8
+DEBUG2  7
+DEBUG1  6
+DEBUG   5
+INFO    4
+NOTE    3
+WARN    2
+ERROR   1
+QUIET   0
+======= ======
+Large value means more noise in the output file.
+.. note::
+    Error and warning messages are written to stderr.
+Each Logger object has its own output destination and verbose level.  So
+multiple Logger objects can be created to manage the message system without
+affecting each other.
+The methods provided by Logger class has the direct connection to the log level.
+E.g.  :func:`info` print messages if the verbose level >= 4 (INFO):
+'''
+
+import sys
+import time
+
+if sys.version_info < (3, 0):
+    process_clock = time.clock
+    perf_counter = time.time
+else:
+    process_clock = time.process_time
+    perf_counter = time.perf_counter
+
+VERBOSE_DEBUG  = 5
+VERBOSE_INFO   = 4
+VERBOSE_NOTICE = 3
+VERBOSE_WARN   = 2
+VERBOSE_ERR    = 1
+VERBOSE_QUIET  = 0
+VERBOSE_CRIT   = -1
+VERBOSE_ALERT  = -2
+VERBOSE_PANIC  = -3
+
+DEBUG4 = VERBOSE_DEBUG + 4
+DEBUG3 = VERBOSE_DEBUG + 3
+DEBUG2 = VERBOSE_DEBUG + 2
+DEBUG1 = VERBOSE_DEBUG + 1
+DEBUG  = VERBOSE_DEBUG
+INFO   = VERBOSE_INFO
+NOTE   = VERBOSE_NOTICE
+NOTICE = NOTE
+WARN   = VERBOSE_WARN
+WARNING = WARN
+ERR    = VERBOSE_ERR
+ERROR  = ERR
+QUIET  = VERBOSE_QUIET
+CRIT   = VERBOSE_CRIT
+ALERT  = VERBOSE_ALERT
+PANIC  = VERBOSE_PANIC
+
+TIMER_LEVEL  = DEBUG
+
+sys.verbose = NOTE
+
+def flush(rec, msg, *args):
+    rec.stdout.write(msg%args)
+    rec.stdout.write('\n')
+    rec.stdout.flush()
+
+def log(rec, msg, *args):
+    if rec.verbose > QUIET:
+        flush(rec, msg, *args)
+
+def error(rec, msg, *args):
+    if rec.verbose >= ERROR:
+        flush(rec, '\nERROR: '+msg+'\n', *args)
+    sys.stderr.write('ERROR: ' + (msg%args) + '\n')
+
+def warn(rec, msg, *args):
+    if rec.verbose >= WARN:
+        flush(rec, '\nWARN: '+msg+'\n', *args)
+        if rec.stdout is not sys.stdout:
+            sys.stderr.write('WARN: ' + (msg%args) + '\n')
+
+def info(rec, msg, *args):
+    if rec.verbose >= INFO:
+        flush(rec, msg, *args)
+
+def note(rec, msg, *args):
+    if rec.verbose >= NOTICE:
+        flush(rec, msg, *args)
+
+def debug(rec, msg, *args):
+    if rec.verbose >= DEBUG:
+        flush(rec, msg, *args)
+
+def debug1(rec, msg, *args):
+    if rec.verbose >= DEBUG1:
+        flush(rec, msg, *args)
+
+def debug2(rec, msg, *args):
+    if rec.verbose >= DEBUG2:
+        flush(rec, msg, *args)
+
+def debug3(rec, msg, *args):
+    if rec.verbose >= DEBUG3:
+        flush(rec, msg, *args)
+
+def debug4(rec, msg, *args):
+    if rec.verbose >= DEBUG4:
+        flush(rec, msg, *args)
+
+def stdout(rec, msg, *args):
+    if rec.verbose >= DEBUG:
+        flush(rec, msg, *args)
+    sys.stdout.write('>>> %s\n' % msg)
+
+def timer(rec, msg, cpu0=None, wall0=None):
+    if cpu0 is None:
+        cpu0 = rec._t0
+    if wall0:
+        rec._t0, rec._w0 = process_clock(), perf_counter()
+        if rec.verbose >= TIMER_LEVEL:
+            flush(rec, '    CPU time for %s %9.2f sec, wall time %9.2f sec'
+                  % (msg, rec._t0-cpu0, rec._w0-wall0))
+        return rec._t0, rec._w0
+    else:
+        rec._t0 = process_clock()
+        if rec.verbose >= TIMER_LEVEL:
+            flush(rec, '    CPU time for %s %9.2f sec' % (msg, rec._t0-cpu0))
+        return rec._t0
+
+def timer_debug1(rec, msg, cpu0=None, wall0=None):
+    if rec.verbose >= DEBUG1:
+        return timer(rec, msg, cpu0, wall0)
+    elif wall0:
+        rec._t0, rec._w0 = process_clock(), perf_counter()
+        return rec._t0, rec._w0
+    else:
+        rec._t0 = process_clock()
+        return rec._t0
+
+class Logger(object):
+    '''
+    Attributes:
+        stdout : file object or sys.stdout
+            The file to dump output message.
+        verbose : int
+            Large value means more noise in the output file.
+    '''
+    def __init__(self, stdout=sys.stdout, verbose=NOTE):
+        self.stdout = stdout
+        self.verbose = verbose
+        self._t0 = process_clock()
+        self._w0 = perf_counter()
+
+    log = log
+    error = error
+    warn = warn
+    note = note
+    info = info
+    debug  = debug
+    debug1 = debug1
+    debug2 = debug2
+    debug3 = debug3
+    debug4 = debug4
+    timer = timer
+    timer_debug1 = timer_debug1
+
+def new_logger(rec=None, verbose=None):
+    '''Create and return a :class:`Logger` object
+    Args:
+        rec : An object which carries the attributes stdout and verbose
+        verbose : a Logger object, or integer or None
+            The verbose level. If verbose is a Logger object, the Logger
+            object is returned. If verbose is not specified (None),
+            rec.verbose will be used in the new Logger object.
+    '''
+    if isinstance(verbose, Logger):
+        log = verbose
+    elif isinstance(verbose, int):
+        if getattr(rec, 'stdout', None):
+            log = Logger(rec.stdout, verbose)
+        else:
+            log = Logger(sys.stdout, verbose)
+    else:
+        log = Logger(rec.stdout, rec.verbose)
+    return log
diff --git a/renormalizer/lib/krylov/krylov.py b/renormalizer/lib/krylov/krylov.py
index 8862e32b..9b9557dd 100644
--- a/renormalizer/lib/krylov/krylov.py
+++ b/renormalizer/lib/krylov/krylov.py
@@ -17,7 +17,7 @@ def _expm_krylov(alpha, beta, V, v_norm, dt):
     try:
         w_hess, u_hess = eigh_tridiagonal(alpha, beta)
     except np.linalg.LinAlgError:
-        logger.warning("tridiagonal failed")
+        logger.warning(f"tridiagonal diagonalization failed, size:{len(alpha)}")
         h = np.diag(alpha) + np.diag(beta, k=-1) + np.diag(beta, k=1)
         w_hess, u_hess = np.linalg.eigh(h)
 
diff --git a/renormalizer/model/basis.py b/renormalizer/model/basis.py
index 9a5f1f83..08da0b9c 100644
--- a/renormalizer/model/basis.py
+++ b/renormalizer/model/basis.py
@@ -107,14 +107,6 @@ def copy(self, new_dof):
         raise NotImplementedError
 
 
-class BasisDummy(BasisSet):
-    """
-    Dummy basis set used in soft truncated TNPI
-    """
-    def  __init__(self, dof, nbas):
-        super().__init__(dof, nbas, [0] * nbas)
-
-
 class BasisSHO(BasisSet):
     """
     simple harmonic oscillator basis set
@@ -904,8 +896,10 @@ class BasisSimpleElectron(BasisSet):
     """
     is_electron = True
 
-    def __init__(self, dof):
-        super().__init__(dof, 2, [0, 1])
+    def __init__(self, dof, sigmaqn=None):
+        if sigmaqn is None:
+            sigmaqn = [0, 1]
+        super().__init__(dof, 2, sigmaqn)
 
     def op_mat(self, op):
         if not isinstance(op, Op):
@@ -998,6 +992,27 @@ def copy(self, new_dof):
         return self.__class__(new_dof, self.sigmaqn)
 
 
+class BasisDummy(BasisSet):
+    def __init__(self, dof, nbas=1, sigmaqn:List=None):
+        if sigmaqn is None:
+            sigmaqn = [0] * nbas
+        super().__init__(dof, nbas, sigmaqn)
+
+    def op_mat(self, op: Union[Op, str]):
+        if not isinstance(op, Op):
+            op = Op(op, None)
+        op_symbol, op_factor = op.split_symbol, op.factor
+
+        if len(op_symbol) == 1 and op_symbol[0] == "I":
+            mat = np.eye(1)
+        else:
+            raise ValueError(f"op_symbol:{op_symbol} is not supported")
+
+        return mat * op_factor
+
+    def copy(self, new_dof):
+        return self.__class__(new_dof, self.sigmaqn)
+
 def x_power_k(k, m, n):
 # <m|x^k|n>, origin is 0
 #\left\langle m\left|X^{k}\right| n\right\rangle=2^{-\frac{k}{2}} \sqrt{n ! m !}
diff --git a/renormalizer/model/model.py b/renormalizer/model/model.py
index 3a543204..17f1d5ad 100644
--- a/renormalizer/model/model.py
+++ b/renormalizer/model/model.py
@@ -515,3 +515,13 @@ def load_from_dict(param, scheme, lam: bool):
     j_constant = Quantity(*param["j constant"])
     model = HolsteinModel([Mol(Quantity(0), ph_list)] * param["mol num"], j_constant, scheme)
     return model, temperature
+
+
+def heisenberg_ops(nspin):
+    ham_terms = []
+    for ispin in range(nspin - 1):
+        op1 = Op("sigma_z sigma_z", [ispin, ispin + 1], 1.0 / 4)
+        op2 = Op("sigma_+ sigma_-", [ispin, ispin + 1], 1.0 / 2)
+        op3 = Op("sigma_- sigma_+", [ispin, ispin + 1], 1.0 / 2)
+        ham_terms.extend([op1, op2, op3])
+    return ham_terms
diff --git a/renormalizer/model/op.py b/renormalizer/model/op.py
index 0c665370..8e9d4b17 100644
--- a/renormalizer/model/op.py
+++ b/renormalizer/model/op.py
@@ -349,6 +349,7 @@ def __repr__(self):
         return str(self)
 
     def __add__(self, other):
+        # if want to add with a scalar, convert the scalar to identity operator first
         if isinstance(other, Op):
             return OpSum([self, other])
         elif isinstance(other, list):
diff --git a/renormalizer/mps/lib.py b/renormalizer/mps/lib.py
index f9561174..6cc04e74 100644
--- a/renormalizer/mps/lib.py
+++ b/renormalizer/mps/lib.py
@@ -433,7 +433,7 @@ def compressed_sum(mps_list, batchsize=5, temp_m_trunc=None):
 
 def _sum(mps_list, compress=True, temp_m_trunc=None):
     new_mps = reduce(lambda mps1, mps2: mps1.add(mps2), mps_list)
-    if compress and not mps_list[0].compress_add:
+    if compress:
         new_mps.canonicalise()
         new_mps.compress(temp_m_trunc=temp_m_trunc)
     return new_mps
diff --git a/renormalizer/mps/matrix.py b/renormalizer/mps/matrix.py
index 876a2b30..4d6fa5d7 100644
--- a/renormalizer/mps/matrix.py
+++ b/renormalizer/mps/matrix.py
@@ -291,11 +291,13 @@ def pair_tensor_contract(
     return tensordot(view_left, view_right, axes=(left_pos, right_pos))
 
 
-def asnumpy(array: Union[np.ndarray, xp.ndarray, Matrix]) -> np.ndarray:
+def asnumpy(array: Union[np.ndarray, xp.ndarray, Matrix, List]) -> np.ndarray:
     if array is None:
         return None
     if isinstance(array, Matrix):
         return array.array
+    if isinstance(array, List):
+        return np.array(array)
     if not USE_GPU:
         assert isinstance(array, np.ndarray)
         return array
@@ -314,3 +316,16 @@ def asxp(array: Union[np.ndarray, xp.ndarray, Matrix]) -> xp.ndarray:
         assert isinstance(array, np.ndarray)
         return array
     return xp.asarray(array)
+
+
+def asxp_oe_args(oe_args):
+    # opt_einsum.contract args in interleaved format
+    new_args = []
+    # the last one is the output index
+    for i in range(len(oe_args) - 1):
+        if i % 2 == 0:
+            new_args.append(asxp(oe_args[i]))
+        else:
+            new_args.append(oe_args[i])
+    new_args.append(oe_args[-1])
+    return new_args
diff --git a/renormalizer/mps/mp.py b/renormalizer/mps/mp.py
index 00d15e7e..259e3040 100644
--- a/renormalizer/mps/mp.py
+++ b/renormalizer/mps/mp.py
@@ -76,8 +76,6 @@ def __init__(self):
         # if sweeping to right: True else False
         self.to_right: bool = None
 
-        # compress after add?
-        self.compress_add: bool = False
 
     @property
     def site_num(self):
@@ -152,12 +150,12 @@ def move_qnidx(self, dstidx: int):
         and R system is tot.
         """
         # construct the L system qn
-        for idx in range(self.qnidx + 1, self.site_num):
-            self.qn[idx] = [self.qntot - i for i in self.qn[idx]]
+        for idx in range(self.qnidx + 1, self.site_num + 1):
+            self.qn[idx] = self.qntot - self.qn[idx]
 
         # set boundary to fsite:
-        for idx in range(self.site_num - 1, dstidx, -1):
-            self.qn[idx] = [self.qntot - i for i in self.qn[idx]]
+        for idx in range(self.site_num, dstidx, -1):
+            self.qn[idx] = self.qntot - self.qn[idx]
         self.qnidx = dstidx
 
     def check_left_canonical(self, atol=None):
@@ -415,21 +413,12 @@ def add(self, other: "MatrixProduct"):
         # qn at the boundary should have dimension 1
         new_mps.qn[0] = np.zeros((1, new_mps.qn[0].shape[1]), dtype=int)
         new_mps.qn[-1] = np.zeros((1, new_mps.qn[0].shape[1]), dtype=int)
-        if self.compress_add:
-            new_mps.canonicalise()
-            new_mps.compress()
         return new_mps
 
     def compress(self, temp_m_trunc=None, ret_s=False):
         """
         inp: canonicalise MPS (or MPO)
 
-        side='l': compress LEFT-canonicalised MPS
-                  by sweeping from RIGHT to LEFT
-                  output MPS is right canonicalised i.e. CRRR
-
-        side='r': reverse of 'l'
-
         Returns
         -------
              truncated MPS
@@ -1020,9 +1009,8 @@ def metacopy(self) -> "MatrixProduct":
         new.compress_config = self.compress_config.copy()
         new.qn = [qn.copy() for qn in self.qn]
         new.qnidx = self.qnidx
-        new.qntot = self.qntot
+        new.qntot = self.qntot.copy()
         new.to_right = self.to_right
-        new.compress_add = self.compress_add
         return new
 
     def _array2mt(self, array, idx, allow_dump=True):
diff --git a/renormalizer/mps/mpdm.py b/renormalizer/mps/mpdm.py
index 14427814..3956de79 100644
--- a/renormalizer/mps/mpdm.py
+++ b/renormalizer/mps/mpdm.py
@@ -38,7 +38,6 @@ def from_mps(cls, mps: Mps):
 
         mpo.optimize_config = mps.optimize_config
         mpo.evolve_config = mps.evolve_config
-        mpo.compress_add = mps.compress_add
 
         mpo.qn = [qn.copy() for qn in mps.qn]
         mpo.qntot = mps.qntot
diff --git a/renormalizer/mps/mps.py b/renormalizer/mps/mps.py
index 53c27890..1dadb782 100644
--- a/renormalizer/mps/mps.py
+++ b/renormalizer/mps/mps.py
@@ -177,7 +177,7 @@ def random(cls, model: Model, qntot, m_max, percent=1.0) -> "Mps":
         return mps
 
     @classmethod
-    def hartree_product_state(cls, model, condition: Dict):
+    def hartree_product_state(cls, model, condition: Dict, qn_idx:int=None):
         r"""
         Construct a Hartree product state
         
@@ -193,8 +193,9 @@ def hartree_product_state(cls, model, condition: Dict):
                     if a basis contains ``"e_1"``, ``"e_2"`` and ``"e_3"``,
                     ``{"e_1": 2}`` (``{"e_1": [0, 0, 1]}``) means ``"e_3"`` is occupied and
                     ``{"e_1": 1}`` (``{"e_1": [0, 1, 0]}``) means ``"e_2"`` is occupied.
-                    Be aware that in :class:`renormalizer.utils.basis.BasisMultiElectronVac` the vacuum state
+                    Be aware that in :class:`renormalizer.BasisMultiElectronVac` the vacuum state
                     is added to the ``0`` index.
+            qn_idx (int): the site index of the quantum number center.
 
         Returns:
             Constructed mps (:class:`Mps`)
@@ -235,8 +236,10 @@ def hartree_product_state(cls, model, condition: Dict):
         if len(condition) != 0:
             raise ValueError(f"Condition not completely used: {condition}")
         mps.qntot = mps.qn[-1][0]
-        mps.qn[-1] = np.zeros((1, qn_size), dtype=int)
-        mps.qnidx = model.nsite - 1
+        mps.qnidx = model.nsite
+        if qn_idx is None:
+            qn_idx = model.nsite - 1
+        mps.move_qnidx(qn_idx)
         mps.to_right = False
 
         return mps
@@ -248,7 +251,7 @@ def ground_state(cls, model: Model, max_entangled: bool,
         Obtain ground state at :math:`T = 0` or :math:`T = \infty` (maximum entangled).
         Electronic DOFs are always at ground state. and vibrational DOFs depend on ``max_entangled``.
         For Spin-Boson model the electronic DOF also depends on ``max_entangled``.
-        
+
 
         Parameters
         ----------
@@ -266,21 +269,21 @@ def ground_state(cls, model: Model, max_entangled: bool,
         Returns
         -------
             mps : renormalizer.mps.Mps
-        
+
         """
-        
+
         mps = cls()
         mps.model = model
         mps.qn = [np.zeros((1, model.qn_size), dtype=int)] * (model.nsite + 1)
         mps.qnidx = model.nsite - 1
         mps.to_right = False
         mps.qntot = np.zeros(model.qn_size, dtype=int)
-            
+
         mps.build_empty_mp(model.nsite)
-        
+
         if condition is not None:
             # check that the condition is not duplicated
-            # each site has at most 1 single key to assign the occupation  
+            # each site has at most 1 single key to assign the occupation
             index = [model.dof_to_siteidx[key] for key in condition.keys()]
             assert len(index) == len(set(index))
             # replace the dof_name key to site_index key
@@ -549,7 +552,7 @@ def e_occupations(self):
     def metacopy(self) -> "Mps":
         new: Mps = super().metacopy()
         new.coeff = self.coeff
-        new.optimize_config = self.optimize_config
+        new.optimize_config = self.optimize_config.copy()
         # evolve_config has its own data
         new.evolve_config = self.evolve_config.copy()
         return new
@@ -569,82 +572,15 @@ def normalize(self, kind):
         ``self`` is overwritten.
         '''
 
-        if kind == "mps_only":
-            new_coeff = self.coeff
-        elif kind == "mps_and_coeff":
-            new_coeff = self.coeff / np.linalg.norm(self.coeff)
-        elif kind == "mps_norm_to_coeff":
-            new_coeff = self.coeff * self.mp_norm
-        else:
-            raise ValueError(f"kind={kind} is not valid.")
-        new_mps = self.scale(1.0 / self.mp_norm, inplace=True)
-        new_mps.coeff = new_coeff
-
-        return new_mps
+        return normalize(self, kind)
 
 
     def expand_bond_dimension(self, hint_mpo=None, coef=1e-10, include_ex=True):
         """
         expand bond dimension as required in compress_config
         """
-        # expander m target
-        m_target = self.compress_config.bond_dim_max_value - self.bond_dims_mean
-        # will be restored at exit
-        self.compress_config.bond_dim_max_value = m_target
-        if self.compress_config.criteria is not CompressCriteria.fixed:
-            logger.warning("Setting compress criteria to fixed")
-            self.compress_config.criteria = CompressCriteria.fixed
-        logger.debug(f"target for expander: {m_target}")
-        if hint_mpo is None:
-            expander = self.__class__.random(self.model, 1, m_target)
-        else:
-            # fill states related to `hint_mpo`
-            logger.debug(
-                f"average bond dimension of hint mpo: {hint_mpo.bond_dims_mean}"
-            )
-            # in case of localized `self`
-            if include_ex:
-                if self.is_mps:
-                    ex_state: MatrixProduct = self.ground_state(self.model, False)
-                    # for self.qntot >= 1
-                    assert self.model.qn_size == 1 # otherwise not supported
-                    for i in range(self.qntot[0]):
-                        ex_state = Mpo.onsite(self.model, r"a^\dagger") @ ex_state
-                elif self.is_mpdm:
-                    assert self.qntot == 1
-                    ex_state: MatrixProduct = self.max_entangled_ex(self.model)
-                else:
-                    assert False
-                ex_state.compress_config = self.compress_config
-                ex_state.move_qnidx(self.qnidx)
-                ex_state.to_right = self.to_right
-                lastone = self + ex_state
+        return expand_bond_dimension(self, hint_mpo, coef, include_ex)
 
-            else:
-                lastone = self
-            expander_list: List["MatrixProduct"] = []
-            cumulated_m = 0
-            while True:
-                lastone.compress_config.criteria = CompressCriteria.fixed
-                expander_list.append(lastone)
-                expander = compressed_sum(expander_list)
-                if cumulated_m == expander.bond_dims_mean:
-                    # probably a small system, the required bond dimension can't be reached
-                    break
-                cumulated_m = expander.bond_dims_mean
-                logger.debug(
-                    f"cumulated bond dimension: {cumulated_m}. lastone bond dimension: {lastone.bond_dims}"
-                )
-                if m_target < cumulated_m:
-                    break
-                if m_target < 0.8 * (lastone.bond_dims_mean * hint_mpo.bond_dims_mean):
-                    lastone = lastone.canonicalise().compress(
-                        m_target // hint_mpo.bond_dims_mean + 1
-                    )
-                lastone = (hint_mpo @ lastone).normalize("mps_and_coeff")
-        logger.debug(f"expander bond dimension: {expander.bond_dims}")
-        self.compress_config.bond_dim_max_value += self.bond_dims_mean
-        return (self + expander.scale(coef*self.norm, inplace=True)).canonicalise().canonicalise().normalize("mps_norm_to_coeff")
 
     def evolve(self, mpo, evolve_dt, normalize=True) -> "Mps":
 
@@ -1831,6 +1767,7 @@ def projector(
     proj = Iden - proj
     return proj
 
+
 def integrand_func_factory(
     shape,
     hop,
@@ -1916,6 +1853,122 @@ def _mu_regularize(s, epsilon=1e-10):
     return s + epsilon * np.exp(-s / epsilon)
 
 
+def expand_bond_dimension(mps, hint_mpo=None, coef=1e-10, include_ex=True):
+    """
+    expand bond dimension as required in compress_config
+    """
+    if hint_mpo is not None and include_ex:
+        # fill states related to `hint_mpo`
+        logger.debug(
+            f"average bond dimension of hint mpo: {hint_mpo.bond_dims_mean}"
+        )
+        # in case of localized `self`
+        if mps.is_mps:
+            ex_state: MatrixProduct = mps.ground_state(mps.model, False)
+            # for self.qntot >= 1
+            assert mps.model.qn_size == 1 # otherwise not supported
+            for i in range(mps.qntot[0]):
+                ex_state = Mpo.onsite(mps.model, r"a^\dagger") @ ex_state
+        elif mps.is_mpdm:
+            assert mps.qntot == 1
+            ex_state: MatrixProduct = mps.max_entangled_ex(mps.model)
+        else:
+            assert False
+        ex_state.compress_config = mps.compress_config
+        ex_state.move_qnidx(mps.qnidx)
+        ex_state.to_right = mps.to_right
+    else:
+        ex_state = None
+    return expand_bond_dimension_general(mps, hint_mpo, coef, ex_state)
+
+
+def expand_bond_dimension_general(mps, hint_mpo=None, coef=1e-10, ex_mps=None):
+    """
+    expand bond dimension as required in compress_config. works for both mps and ttns
+    """
+    # expander m target
+    m_target = mps.compress_config.bond_dim_max_value - mps.bond_dims_mean
+    # will be restored at exit
+    mps.compress_config.bond_dim_max_value = m_target
+    if mps.compress_config.criteria is not CompressCriteria.fixed:
+        logger.warning("Setting compress criteria to fixed")
+        mps.compress_config.criteria = CompressCriteria.fixed
+    logger.debug(f"target for expander: {m_target}")
+    if hint_mpo is None:
+        expander = mps.__class__.random(mps.model, 1, m_target)
+    else:
+        # fill states related to `hint_mpo`
+        logger.debug(
+            f"average bond dimension of hint mpo: {hint_mpo.bond_dims_mean}"
+        )
+        # in case of localized `self`
+        if ex_mps:
+            lastone = mps + ex_mps
+
+        else:
+            lastone = mps
+        expander_list: List["MatrixProduct"] = []
+        cumulated_m = 0
+        while True:
+            lastone.compress_config.criteria = CompressCriteria.fixed
+            expander_list.append(lastone)
+            expander = compressed_sum(expander_list)
+            if cumulated_m == expander.bond_dims_mean:
+                # probably a small system, the required bond dimension can't be reached
+                break
+            cumulated_m = expander.bond_dims_mean
+            logger.debug(
+                f"cumulated bond dimension: {cumulated_m}. lastone bond dimension: {lastone.bond_dims}"
+            )
+            if m_target < cumulated_m:
+                break
+            if m_target < 0.8 * (lastone.bond_dims_mean * hint_mpo.bond_dims_mean):
+                lastone = lastone.canonicalise().compress(
+                    m_target // hint_mpo.bond_dims_mean + 1
+                )
+            lastone = (hint_mpo @ lastone).normalize("mps_and_coeff")
+    logger.debug(f"expander bond dimension: {expander.bond_dims}")
+    mps.compress_config.bond_dim_max_value += mps.bond_dims_mean
+    return (mps + expander.scale(coef * mps.norm, inplace=True)).canonicalise().canonicalise().normalize(
+        "mps_norm_to_coeff")
+
+
+def normalize(tn, kind):
+    r''' normalize the wavefunction
+
+    Parameters
+    ----------
+    kind: str
+        "mps_only": the mps part is normalized and coeff is not modified;
+        "mps_norm_to_coeff": the mps part is normalized and the norm is multiplied to coeff;
+        "mps_and_coeff": both mps and coeff is normalized
+
+    Returns
+    -------
+    ``self`` is overwritten.
+    '''
+
+    if hasattr(tn, "mp_norm"):
+        tn_norm = tn.mp_norm
+    elif hasattr(tn, "ttns_norm"):
+        tn_norm = tn.ttns_norm
+    else:
+        raise ValueError(f"{type(tn)} does not have norm attribute")
+
+    if kind in ["mps_only", "ttns_only"]:
+        new_coeff = tn.coeff
+    elif kind in ["mps_and_coeff", "ttns_and_coeff"]:
+        new_coeff = tn.coeff / np.linalg.norm(tn.coeff)
+    elif kind in ["mps_norm_to_coeff", "ttns_norm_to_coeff"]:
+        new_coeff = tn.coeff * tn_norm
+    else:
+        raise ValueError(f"kind={kind} is not valid.")
+
+    tn.scale(1.0 / tn_norm, inplace=True)
+    tn.coeff = new_coeff
+
+    return tn
+
 class BraKetPair:
     def __init__(self, bra_mps, ket_mps, mpo=None):
         self.bra_mps = bra_mps
diff --git a/renormalizer/mps/svd_qn.py b/renormalizer/mps/svd_qn.py
index 6eeb8342..3f82f91a 100644
--- a/renormalizer/mps/svd_qn.py
+++ b/renormalizer/mps/svd_qn.py
@@ -156,6 +156,7 @@ def svd_qn(
     SVD = not QR
     coef_matrix = coef_array.reshape((np.prod(qnbigl.shape[:-1]), np.prod(qnbigr.shape[:-1])))
 
+    assert qntot.ndim == 1
     qn_size = len(qntot)
     localqnl = qnbigl.reshape(-1, qn_size)
     localqnr = qnbigr.reshape(-1, qn_size)
diff --git a/renormalizer/mps/symbolic_mpo.py b/renormalizer/mps/symbolic_mpo.py
index 3948ba3f..714dcff0 100644
--- a/renormalizer/mps/symbolic_mpo.py
+++ b/renormalizer/mps/symbolic_mpo.py
@@ -9,6 +9,7 @@
 import scipy.sparse
 
 from renormalizer.model import Model
+from renormalizer.model.basis import BasisSet
 from renormalizer.model.op import Op
 from renormalizer.lib import bipartite_vertex_cover
 
@@ -105,15 +106,17 @@ def construct_symbolic_mpo(table, factor, algo="Hopcroft-Karp"):
     qn_size = len(table[0][0].qn)
     # Simplest case. Cut to the chase
     if len(table) == 1:
-        # The first layer: number of sites. The 2nd and 3rd layer: in and out virtual bond
+        # The first layer: number of sites. The middle array: in and out virtual bond
         # the 4th layer: operator sums
-        mpo: List[List[List[List[[Op]]]]] = []
+        mpo: List[np.ndarray[List[Op]]] = []
         mpoqn = [np.zeros((1, qn_size), dtype=int)]
         primary_ops = list(set(table[0]))
         op2idx = dict(zip(primary_ops, range(len(primary_ops))))
         out_ops_list: List[List[OpTuple]] = [[OpTuple([0], qn=0, factor=1)]]
         for op in table[0]:
-            mpo.append([[[op]]])
+            mo = np.full((1, 1), None)
+            mo[0][0] = [op]
+            mpo.append(mo)
             qn = mpoqn[-1][0] + op.qn
             mpoqn.append(np.array([qn]))
             out_ops_list.append([OpTuple([0, op2idx[op]], qn=qn, factor=1)])
@@ -127,60 +130,15 @@ def construct_symbolic_mpo(table, factor, algo="Hopcroft-Karp"):
         # the last two terms are not set for fast construction of the operators
         return mpo, mpoqn, qntot, qnidx, out_ops_list, primary_ops
 
-    # use np.uint32, np.uint16 to save memory
-    max_uint32 = np.iinfo(np.uint32).max
-    max_uint16 = np.iinfo(np.uint16).max
-
     logger.debug(f"symbolic mpo algorithm: {algo}")
     logger.debug(f"Input operator terms: {len(table)}")
 
-    # translate the symbolic operator table to an easy to manipulate numpy array
-    table = np.array(table)
-
-    # unique operators with DoF names taken into consideration
-    # The inclusion of DoF names is necessary for multi-dof basis.
-    unique_op = OrderedDict.fromkeys(table.ravel())
-    # Convert Set(table.ravel()) to List will change the Op order in list, OrderedDict made reproducible
-    unique_op = list(unique_op.keys())
-    # check the index of different operators could be represented with np.uint16
-    assert len(unique_op) < max_uint16
-
-    # Construct mapping from easy-to-manipulate integer to actual Op
-    primary_ops = list(unique_op)
-
-    op2idx = dict(zip(unique_op, range(len(unique_op))))
-    new_table = np.vectorize(op2idx.get)(table).astype(np.uint16)
-
-    del unique_op
-
-    if __debug__:
-        qn_table = np.array([[x.qn for x in ta] for ta in table])
-        factor_table = np.array([[x.factor for x in ta] for ta in table])
-        for idx in range(len(primary_ops)):
-            coord = np.nonzero(new_table == idx)
-            # check that op with the same symbol has the same factor and qn
-            for j in range(qn_size):
-                assert np.unique(qn_table[:, :, j][coord]).size == 1
-            assert np.all(factor_table[coord] == factor_table[coord][0])
-
-        del factor_table, qn_table
-
-    # combine the same terms but with different factors(add them together)
-    unique_term, unique_inverse = np.unique(new_table, axis=0, return_inverse=True)
-    # it is efficient to vectorize the operation that moves the rows and cols
-    # and sum them together
-    coord = np.array([[newidx, oldidx] for oldidx, newidx in enumerate(unique_inverse)])
-    mask = scipy.sparse.csr_matrix((np.ones(len(coord)), (coord[:, 0], coord[:, 1])))
-    factor = mask.dot(factor)
+    table, factor, primary_ops = _transform_table(table, factor)
 
     # add the first and last column for convenience
-    ta = np.zeros((unique_term.shape[0], 1), dtype=np.uint16)
-    table = np.concatenate((ta, unique_term, ta), axis=1)
+    ta = np.zeros((table.shape[0], 1), dtype=np.uint16)
+    table = np.concatenate((ta, table, ta), axis=1)
     logger.debug(f"After combination of the same terms: {table.shape[0]}")
-    # check the index of interaction could be represented with np.uint32
-    assert table.shape[0] < max_uint32
-
-    del unique_term, unique_inverse
 
     # 0 represents the identity symbol. Identity might not present
     # in `primary_ops` but the algorithm still works.
@@ -214,115 +172,14 @@ def _construct_symbolic_mpo(table, in_ops, factor, primary_ops, algo="Hopcroft-K
     out_ops_list = [in_ops]
 
     for isite in range(nsite):
-        # split table into the row and col part
-        term_row, row_unique_inverse = np.unique(table[:, :2], axis=0, return_inverse=True)
-
-        # faster version of the following code
-        # term_col, col_unique_inverse = np.unique(table[:, 2:], axis=0, return_inverse=True)
-        term_col = {}
-        col_unique_inverse = []
-        for row in table[:, 2:]:
-            row_bytes = row.tobytes()
-            i_and_row = term_col.get(row_bytes)
-            if i_and_row is None:
-                i_and_row = (len(term_col), row)
-                term_col[row_bytes] = i_and_row
-            col_unique_inverse.append(i_and_row[0])
-        term_col = [v[1] for v in term_col.values()]
-
-        # get the non_redudant ops
-        # the +1 trick is to use the csr sparse matrix format
-        non_red = scipy.sparse.diags(np.arange(1, table.shape[0] + 1), format="csr", dtype=np.uint32)
-        coord = np.array([[newidx, oldidx] for oldidx, newidx in enumerate(row_unique_inverse)])
-        mask = scipy.sparse.csr_matrix((np.ones(len(coord), dtype=np.uint32), (coord[:, 0], coord[:, 1])))
-        non_red = mask.dot(non_red)
-        coord = np.array([[oldidx, newidx] for oldidx, newidx in enumerate(col_unique_inverse)])
-        mask = scipy.sparse.csr_matrix((np.ones(len(coord), dtype=np.uint32), (coord[:, 0], coord[:, 1])))
-        non_red = non_red.dot(mask)
-        # use sparse matrix to represent non_red will be inefficient a little
-        # bit compared to dense matrix, but saves a lot of memory when the
-        # number of terms is huge
-        # logger.info(f"isite: {isite}, bipartite graph size: {non_red.shape}")
-
-        # the usual case
-
-        bigraph = []
-        if non_red.shape[0] < non_red.shape[1]:
-            for i in range(non_red.shape[0]):
-                bigraph.append(non_red.indices[non_red.indptr[i]:non_red.indptr[i + 1]])
-            rowbool, colbool = bipartite_vertex_cover(bigraph, algo=algo)
-        else:
-            non_red_csc = non_red.tocsc()
-            for i in range(non_red.shape[1]):
-                bigraph.append(non_red_csc.indices[non_red_csc.indptr[i]:non_red_csc.indptr[i + 1]])
-            colbool, rowbool = bipartite_vertex_cover(bigraph, algo=algo)
-
-        row_select = np.nonzero(rowbool)[0]
-        # largest cover first
-        row_select = sorted(row_select, key=lambda i: non_red.indptr[i + 1] - non_red.indptr[i], reverse=True)
-        col_select = np.nonzero(colbool)[0]
-
-
-        if len(row_select) > 0:
-            assert np.amax(row_select) < non_red.shape[0]
-        if len(col_select) > 0:
-            assert np.amax(col_select) < non_red.shape[1]
-
-        # select the reserved ops
-        out_ops: List[List[OpTuple]] = []
-        new_table = []
-        new_factor = []
-
-        for row_idx in row_select:
-            # construct out_op
-            # dealing with row (left side of the table). One row corresponds to multiple cols.
-            # Produce one out operator and multiple new_table entries
-            symbol = term_row[row_idx]
-            qn = in_ops[symbol[0]][0].qn + primary_ops[symbol[1]].qn
-            out_op = OpTuple(symbol, qn, factor=1.0)
-            out_ops.append([out_op])
-
-            col_link = non_red.indices[non_red.indptr[row_idx]:non_red.indptr[row_idx + 1]]
-            stack = np.array([len(out_ops) - 1] * len(col_link), dtype=np.uint16).reshape(-1, 1)
-            new_table.append(np.hstack((stack, [term_col[i] for i in col_link])))
-            new_factor.append(factor[non_red[row_idx, col_link].toarray() - 1])
-            non_red.data[non_red.indptr[row_idx]:non_red.indptr[row_idx + 1]] = 0
-
-        non_red.eliminate_zeros()
-
-        nonzero_row_idx, nonzero_col_idx = non_red.nonzero()
-        for col_idx in col_select:
-
-            out_ops.append([])
-            # complementary operator
-            # dealing with column (right side of the table). One col correspond to multiple rows.
-            # Produce multiple out operators and one new_table entry
-            non_red_one_col = non_red[:, col_idx].toarray().flatten()
-            for i in nonzero_row_idx[np.nonzero(nonzero_col_idx == col_idx)[0]]:
-                symbol = term_row[i]
-                qn = in_ops[symbol[0]][0].qn + primary_ops[symbol[1]].qn
-                out_op = OpTuple(symbol, qn, factor=factor[non_red_one_col[i] - 1])
-                out_ops[-1].append(out_op)
-
-            new_table.append(np.array([len(out_ops) - 1] + list(term_col[col_idx]), dtype=np.uint16).reshape(1, -1))
-            new_factor.append(1.0)
-
-            # it is not necessary to remove the column nonzero elements
-            # non_red[:, col_idx] = 0
-            # non_red.eliminate_zeros()
-
-        # reconstruct the table in new operator
-        table = np.concatenate(new_table)
-        # check the number of incoming operators could be represent as np.uint16
-        assert len(out_ops) <= np.iinfo(np.uint16).max
-        factor = np.concatenate(new_factor, axis=None)
-
-        assert len(table) == len(factor)
+        table_row = table[:, :2]
+        table_col = table[:, 2:]
+        out_ops, table, factor = _construct_symbolic_mpo_one_site(table_row, table_col, [in_ops], factor, primary_ops, algo)
 
         # debug
         # logger.debug(f"in_ops: {in_ops}")
         # logger.debug(f"out_ops: {out_ops}")
-        # logger.debug(f"new_factor: {new_factor}")
+        # logger.debug(f"new_factor: {factor}")
 
         in_ops = out_ops
         # in_ops_full_symbol = out_ops_full_symbol
@@ -331,12 +188,118 @@ def _construct_symbolic_mpo(table, in_ops, factor, primary_ops, algo="Hopcroft-K
     return out_ops_list
 
 
-def add_idx(symbol, idx):
-    symbols = symbol.split(" ")
-    for i in range(len(symbols)):
-        symbols[i] = symbols[i] + f"_{idx}"
-    return " ".join(symbols)
+def _construct_symbolic_mpo_one_site(table_row, table_col, in_ops_list, factor, primary_ops, algo, k=1):
+    # split table into the row and col part
+    term_row, row_unique_inverse = np.unique(table_row, axis=0, return_inverse=True)
+    assert len(in_ops_list) + k == term_row.shape[1]
+
+    # faster version of the following code
+    # term_col, col_unique_inverse = np.unique(table_col, axis=0, return_inverse=True)
+    term_col = {}
+    col_unique_inverse = []
+    for row in table_col:
+        row_bytes = row.tobytes()
+        i_and_row = term_col.get(row_bytes)
+        if i_and_row is None:
+            i_and_row = (len(term_col), row)
+            term_col[row_bytes] = i_and_row
+        col_unique_inverse.append(i_and_row[0])
+    term_col = [v[1] for v in term_col.values()]
+
+    # get the non_redudant ops
+    # the +1 trick is to use the csr sparse matrix format
+    non_red = scipy.sparse.diags(np.arange(1, table_row.shape[0] + 1), format="csr", dtype=np.uint32)
+    coord = np.array([[newidx, oldidx] for oldidx, newidx in enumerate(row_unique_inverse)])
+    mask = scipy.sparse.csr_matrix((np.ones(len(coord), dtype=np.uint32), (coord[:, 0], coord[:, 1])))
+    non_red = mask.dot(non_red)
+    coord = np.array([[oldidx, newidx] for oldidx, newidx in enumerate(col_unique_inverse)])
+    mask = scipy.sparse.csr_matrix((np.ones(len(coord), dtype=np.uint32), (coord[:, 0], coord[:, 1])))
+    non_red = non_red.dot(mask)
+    # use sparse matrix to represent non_red will be inefficient a little
+    # bit compared to dense matrix, but saves a lot of memory when the
+    # number of terms is huge
+    # logger.info(f"isite: {isite}, bipartite graph size: {non_red.shape}")
+
+    # the usual case
+
+    bigraph = []
+    if non_red.shape[0] < non_red.shape[1]:
+        for i in range(non_red.shape[0]):
+            bigraph.append(non_red.indices[non_red.indptr[i]:non_red.indptr[i + 1]])
+        rowbool, colbool = bipartite_vertex_cover(bigraph, algo=algo)
+    else:
+        non_red_csc = non_red.tocsc()
+        for i in range(non_red.shape[1]):
+            bigraph.append(non_red_csc.indices[non_red_csc.indptr[i]:non_red_csc.indptr[i + 1]])
+        colbool, rowbool = bipartite_vertex_cover(bigraph, algo=algo)
+
+    row_select = np.nonzero(rowbool)[0]
+    # largest cover first
+    row_select = sorted(row_select, key=lambda i: non_red.indptr[i + 1] - non_red.indptr[i], reverse=True)
+    col_select = np.nonzero(colbool)[0]
+
+    if len(row_select) > 0:
+        assert np.amax(row_select) < non_red.shape[0]
+    if len(col_select) > 0:
+        assert np.amax(col_select) < non_red.shape[1]
+
+    # select the reserved ops
+    out_ops: List[List[OpTuple]] = []
+    new_table = []
+    new_factor = []
 
+    for row_idx in row_select:
+        # construct out_op
+        # dealing with row (left side of the table). One row corresponds to multiple cols.
+        # Produce one out operator and multiple new_table entries
+        symbol = term_row[row_idx]
+        qn = _compute_qn(in_ops_list, symbol, primary_ops, k)
+        out_op = OpTuple(symbol, qn, factor=1.0)
+        out_ops.append([out_op])
+
+        col_link = non_red.indices[non_red.indptr[row_idx]:non_red.indptr[row_idx + 1]]
+        stack = np.array([len(out_ops) - 1] * len(col_link), dtype=np.uint16).reshape(-1, 1)
+        new_table.append(np.hstack((stack, [term_col[i] for i in col_link])))
+        new_factor.append(factor[non_red[row_idx, col_link].toarray() - 1])
+        non_red.data[non_red.indptr[row_idx]:non_red.indptr[row_idx + 1]] = 0
+
+    non_red.eliminate_zeros()
+
+    nonzero_row_idx, nonzero_col_idx = non_red.nonzero()
+    for col_idx in col_select:
+
+        out_ops.append([])
+        # complementary operator
+        # dealing with column (right side of the table). One col correspond to multiple rows.
+        # Produce multiple out operators and one new_table entry
+        non_red_one_col = non_red[:, col_idx].toarray().flatten()
+        for i in nonzero_row_idx[np.nonzero(nonzero_col_idx == col_idx)[0]]:
+            symbol = term_row[i]
+            qn = _compute_qn(in_ops_list, symbol, primary_ops, k)
+            out_op = OpTuple(symbol, qn, factor=factor[non_red_one_col[i] - 1])
+            out_ops[-1].append(out_op)
+
+        new_table.append(np.array([len(out_ops) - 1] + list(term_col[col_idx]), dtype=np.uint16).reshape(1, -1))
+        new_factor.append(1.0)
+
+        # it is not necessary to remove the column nonzero elements
+        # non_red[:, col_idx] = 0
+        # non_red.eliminate_zeros()
+
+    # reconstruct the table in new operator
+    table = np.concatenate(new_table)
+    # check the number of incoming operators could be represented as np.uint16
+    assert len(out_ops) <= np.iinfo(np.uint16).max
+    factor = np.concatenate(new_factor, axis=None)
+
+    assert len(table) == len(factor)
+    return out_ops, table, factor
+
+
+def _compute_qn(in_ops_list, symbol, primary_ops, k) -> int:
+    qn = sum(in_ops[i][0].qn for in_ops, i in zip(in_ops_list, symbol[:-k]))
+    qn += sum(primary_ops[i].qn for i in symbol[-k:])
+    return qn
 
 def _terms_to_table(model: Model, terms: List[Op], const: float):
     r"""
@@ -377,9 +340,65 @@ def _terms_to_table(model: Model, terms: List[Op], const: float):
     return table, factor_list
 
 
+def _transform_table(table, factor):
+    """Transforms the table to integer table and combine duplicate terms."""
+
+    # use np.uint32, np.uint16 to save memory
+    max_uint32 = np.iinfo(np.uint32).max
+    max_uint16 = np.iinfo(np.uint16).max
+
+    # translate the symbolic operator table to an easy to manipulate numpy array
+    table = np.array(table)
+    # unique operators with DoF names taken into consideration
+    # The inclusion of DoF names is necessary for multi-dof basis.
+    unique_op = OrderedDict.fromkeys(table.ravel())
+    # Convert Set(table.ravel()) to List will change the Op order in list, OrderedDict made reproducible
+    unique_op = list(unique_op.keys())
+
+    # check the index of different operators could be represented with np.uint16
+    assert len(unique_op) < max_uint16
+
+    # Construct mapping from easy-to-manipulate integer to actual Op
+    primary_ops = list(unique_op)
+
+    op2idx = dict(zip(unique_op, range(len(unique_op))))
+    new_table = np.vectorize(op2idx.get)(table).astype(np.uint16)
+
+    del unique_op
+
+    if __debug__:
+        qn_size = len(table[0][0].qn)
+        qn_table = np.array([[x.qn for x in ta] for ta in table])
+        factor_table = np.array([[x.factor for x in ta] for ta in table])
+        for idx in range(len(primary_ops)):
+            coord = np.nonzero(new_table == idx)
+            # check that op with the same symbol has the same factor and qn
+            for j in range(qn_size):
+                assert np.unique(qn_table[:, :, j][coord]).size == 1
+            assert np.all(factor_table[coord] == factor_table[coord][0])
+
+        del factor_table, qn_table
+
+    # combine the same terms but with different factors(add them together)
+    unique_term, unique_inverse = np.unique(new_table, axis=0, return_inverse=True)
+    # it is efficient to vectorize the operation that moves the rows and cols
+    # and sum them together
+    coord = np.array([[newidx, oldidx] for oldidx, newidx in enumerate(unique_inverse)])
+    mask = scipy.sparse.csr_matrix((np.ones(len(coord)), (coord[:, 0], coord[:, 1])))
+    factor = mask.dot(factor)
+
+    # check the index of interaction could be represented with np.uint32
+    assert table.shape[0] < max_uint32
+
+    return np.array(unique_term), factor, primary_ops
+
+
 # translate the numbers into symbolic Matrix Operator
 def compose_symbolic_mo(in_ops, out_ops, primary_ops):
-    mo = [[[] for o in range(len(out_ops))] for i in range(len(in_ops))]
+    shape = [len(in_ops), len(out_ops)]
+    mo = np.full(shape, None, dtype=object)
+    for i, _ in np.ndenumerate(mo):
+        mo[i] = []
     for iop, out_op in enumerate(out_ops):
         for composed_op in out_op:
             in_idx = composed_op.symbol[0]
@@ -389,15 +408,18 @@ def compose_symbolic_mo(in_ops, out_ops, primary_ops):
 
 
 # translate symbolic Matrix Operator to numerical matrix operator defined with certain basis
-def symbolic_mo_to_numeric_mo(basis, mo, dtype):
+def symbolic_mo_to_numeric_mo(basis: BasisSet, mo, dtype):
     pdim = basis.nbas
-    nrow, ncol = len(mo), len(mo[0])
-    mo_mat = np.zeros((nrow, pdim, pdim, ncol), dtype=dtype)
+    shape = list(mo.shape) + [pdim, pdim]
+    mo_mat = np.zeros(shape, dtype=dtype)
+
+    for i, terms in np.ndenumerate(mo):
+        for term in terms:
+            mo_mat[i] += basis.op_mat(term)
 
-    for irow, icol in itertools.product(range(nrow), range(ncol)):
-        for term in mo[irow][icol]:
-            mo_mat[irow, :, :, icol] += basis.op_mat(term)
-    return mo_mat
+    axes = list(range(mo.ndim + 2))
+    axes = axes[:-3] + axes[-2:] + [axes[-3]]
+    return mo_mat.transpose(axes)
 
 
 def _format_symbolic_mpo(symbolic_mpo):
diff --git a/renormalizer/sbm/__init__.py b/renormalizer/sbm/__init__.py
index d8eabfa6..e2dc93e1 100644
--- a/renormalizer/sbm/__init__.py
+++ b/renormalizer/sbm/__init__.py
@@ -1,4 +1,4 @@
 # -*- coding: utf-8 -*-
 
-from renormalizer.sbm.lib import SpectralDensityFunction, param2mollist
-from renormalizer.sbm.sbm import SpinBosonDynamics
\ No newline at end of file
+from renormalizer.sbm.lib import SpectralDensityFunction, param2mollist, DebyeSDF, OhmicSDF, ColeDavidsonSDF
+from renormalizer.sbm.sbm import SpinBosonDynamics
diff --git a/renormalizer/sbm/lib.py b/renormalizer/sbm/lib.py
index 59dd7f9f..2ed09975 100644
--- a/renormalizer/sbm/lib.py
+++ b/renormalizer/sbm/lib.py
@@ -1,14 +1,18 @@
 # -*- coding: utf-8 -*-
 
-import numpy as np
-import numpy.polynomial.laguerre as la
-import numpy.polynomial.legendre as le
+import logging
+from typing import Union, Tuple
+
 import scipy
 import scipy.special
 import scipy.optimize
 
-from renormalizer.model import Phonon, Mol, SpinBosonModel
+from renormalizer.model import Phonon, SpinBosonModel
 from renormalizer.utils import Quantity
+from renormalizer.mps.backend import np
+
+
+logger = logging.getLogger(__name__)
 
 
 class DebyeSpectralDensityFunction:
@@ -25,42 +29,73 @@ def func(self, omega_value):
         """
         the function of the Debye-type spectral density function
         """
-        return 2. * self.lamb * omega_value * self.omega_c / (omega_value ** 2 + self.omega_c ** 2)
+        return 2.0 * self.lamb * omega_value * self.omega_c / (omega_value**2 + self.omega_c**2)
+
+
+DebyeSDF = DebyeSpectralDensityFunction
 
 
 class SpectralDensityFunction:
     r"""
-    the ohmic spectral density function
+    The Ohmic spectral density function
     J(\omega) = \pi / 2 \alpha \omega e^{-\omega/\omega_c}
     """
 
-    def __init__(self, alpha, omega_c: Quantity):
+    def __init__(self, alpha: float, omega_c: Union[Quantity, float], s: float = 1):
         self.alpha = alpha
-        self.omega_c = omega_c.as_au()
+        if isinstance(omega_c, Quantity):
+            omega_c = omega_c.as_au()
+        self.omega_c = omega_c
+        self.s = s
+
+    def reno(self, omega_l) -> float:
+        def integrate_func(x):
+            return self.func(x) / x**2
+
+        res = scipy.integrate.quad(integrate_func, a=omega_l, b=self.omega_c * 30)
+        logger.info(f"integrate: {res[0]}, {res[1]}")
+        re = np.exp(-res[0] * 2 / np.pi)
 
-    def adiabatic_renormalization(self, delta_quan: Quantity, p: float):
-        delta = delta_quan.as_au()
+        return re
+
+    def adiabatic_renormalization(self, delta: Union[Quantity, float], p: float) -> Tuple[float, float]:
+        """
+        the cut-off omega_l is p*delta
+        """
+        if isinstance(delta, Quantity):
+            delta = delta.as_au()
         loop = 0
-        re = 1.
+        re = 1.0
         while loop < 50:
             re_old = re
             omega_l = delta * re * p
-            re = np.exp(-self.alpha * scipy.special.expn(1, omega_l / self.omega_c))
+
+            def integrate_func(x):
+                return self.func(x) / x**2
+
+            res = scipy.integrate.quad(integrate_func, a=omega_l, b=self.omega_c * 30)
+            logger.info(f"integrate: {res[0]}, {res[1]}")
+            re = np.exp(-res[0] * 2 / np.pi)
             loop += 1
+            logger.info(f"re, {re_old}, {re}")
             if np.allclose(re, re_old):
                 break
 
-        return delta_quan * re, Quantity(delta * re * p)
+        return delta * re, delta * re * p
 
     def func(self, omega_value):
         """
         the function of the ohmic spectral density function
         """
-        return np.pi / 2. * self.alpha * omega_value * np.exp(-omega_value / self.omega_c)
+        return (
+            np.pi / 2.0 * self.alpha
+            * omega_value**self.s * self.omega_c ** (1 - self.s)
+            * np.exp(-omega_value / self.omega_c)
+        )
 
     @staticmethod
-    def post_process(omega_value, c_j2, ifsort):
-        displacement_array = np.sqrt(c_j2) / omega_value ** 2
+    def post_process(omega_value, c_j2, ifsort=True):
+        displacement_array = np.sqrt(c_j2) / omega_value**2
         if ifsort:
             idx = np.argsort(c_j2 / omega_value)[::-1]
         else:
@@ -78,86 +113,105 @@ def _dos_Wang1(self, nb, omega_value):
         """
         return (nb + 1) / self.omega_c * np.exp(-omega_value / self.omega_c)
 
-    def Wang1(self, nb, ifsort=True):
-        """
+    def Wang1(self, nb):
+        r"""
         Wang's 1st scheme discretization
         """
-        omega_value = np.array([-np.log(-float(j) / (nb + 1) + 1.) * self.omega_c for j in range(1, nb + 1, 1)])
+        omega_value = np.array([-np.log(-float(j) / (nb + 1) + 1.0) * self.omega_c for j in range(1, nb + 1, 1)])
 
         # general form
-        # c_j2 = 2./np.pi * omega_value * self.func(omega_value) / self._dos_Wang1(nb, omega_value)
+        c_j2 = 2.0 / np.pi * omega_value * self.func(omega_value) / self._dos_Wang1(nb, omega_value)
 
-        # excat form
-        c_j2 = omega_value ** 2 * self.alpha * self.omega_c / (nb + 1)
+        return omega_value, c_j2
 
-        return self.post_process(omega_value, c_j2, ifsort)
-
-    def legendre(self, nb, x0, x1, ifsort=True):
-        """
-        Legendre polynomial fit [x0, x1] to [-1,1]
-        omega_m is the cutoff
-        """
-        omega_value, w = le.leggauss(nb)
-        omega_value = (omega_value + (x1 + x0) / (x1 - x0)) * (x1 - x0) / 2.
-        c_j2 = w * (x1 - x0) / 2. * self.alpha * omega_value ** 2 * np.exp(-omega_value / self.omega_c)
+    def trapz(self, nb, x0, x1):
+        dw = (x1 - x0) / float(nb)
+        xlist = [x0 + i * dw for i in range(nb + 1)]
+        omega_value = np.array([(xlist[i] + xlist[i + 1]) / 2.0 for i in range(nb)])
+        c_j2 = (
+            np.array([(self.func(xlist[i]) + self.func(xlist[i + 1])) / 2 for i in range(nb)])
+            * 2.0 / np.pi * omega_value * dw
+        )
 
-        return self.post_process(omega_value, c_j2, ifsort)
+        return omega_value, c_j2
 
-    def laguerre(self, nb, ifsort=True):
-        assert nb <= 100
 
-        omega_value, w = la.laggauss(nb)
-        omega_value *= self.omega_c
-        c_j2 = w * self.alpha * self.omega_c * omega_value ** 2
+OhmicSDF = SpectralDensityFunction
 
-        return self.post_process(omega_value, c_j2, ifsort)
 
-    def trapz(self, nb, x0, x1, ifsort=True):
-        dw = (x1 - x0) / float(nb)
-        xlist = [x0 + i * dw for i in range(nb + 1)]
-        omega_value = np.array([(xlist[i] + xlist[i + 1]) / 2. for i in range(nb)])
-        c_j2 = np.array([(self.func(xlist[i]) + self.func(xlist[i + 1])) / 2 for i in range(nb)]) * 2. / np.pi * omega_value * dw
+class ColeDavidsonSDF:
+    """
+    the ColeDavidson spectral density function
+    """
 
-        return self.post_process(omega_value, c_j2, ifsort)
+    def __init__(self, ita, omega_c, beta, omega_limit):
+        self.ita = ita
+        self.omega_c = omega_c
+        self.beta = beta
+        self.omega_limit = omega_limit
 
-    def _opt_cut(self, p):
-        """
-        p is the percent of the height of the SDF
-        """
-        assert 0 < p < 1
-        cut = np.exp(-1.) * p
+    def reno(self, omega_l):
+        def integrate_func(x):
+            return self.func(x) / x**2
 
-        def F(x):
-            return x * np.exp(-x) - cut
+        res = scipy.integrate.quad(integrate_func, a=omega_l, b=omega_l * 1000)
+        logger.info(f"integrate: {res[0]}, {res[1]}")
+        re = np.exp(-res[0] * 2 / np.pi)
 
-        def fprime(x):
-            return np.exp(-x) * (1. - x)
+        return re
 
-        def fprime2(x):
-            return (x - 2.) * np.exp(-x)
+    def func(self, omega_value):
+        """
+        the function of the spectral density function
+        """
+        theta = np.arctan(omega_value / self.omega_c)
+        return self.ita * np.sin(self.beta * theta) / (1 + omega_value**2 / self.omega_c**2) ** (self.beta / 2)
 
-        x1 = scipy.optimize.newton(F, 0.0, fprime=fprime, fprime2=fprime2)
-        x2 = scipy.optimize.newton(F, -np.log(cut), fprime=fprime, fprime2=fprime2)
+    def _dos_Wang1(self, A, omega_value):
+        """
+        Wang's 1st scheme DOS \rho(\omega)
+        """
 
-        return x1 * self.omega_c, x2 * self.omega_c
+        return A * self.func(omega_value) / omega_value
 
-    def plot_data(self, x0, x1, n, omega_value, c_j2, sigma=0.1):
+    def Wang1(self, nb):
         """
-        plot the spectral density function (continuous and discrete)
+        Wang's 1st scheme discretization
         """
-        x = np.linspace(x0, x1, n)
-        y_c = self.func(x)
 
-        y_d = np.einsum("i,ji -> j", c_j2 / omega_value * np.pi / 2. * 1 / np.sqrt(2 * np.pi * sigma ** 2),
-                        np.exp(-(np.subtract.outer(x, omega_value) / sigma) ** 2 / 2.))
+        def integrate_func(x):
+            return self.func(x) / x
+
+        A = (nb + 1) / scipy.integrate.quad(integrate_func, a=0, b=self.omega_limit)[0]
+        logger.info(scipy.integrate.quad(integrate_func, a=0, b=self.omega_limit)[0] * 4 / np.pi)
+        logger.info(2 * self.ita)
+        nsamples = int(1e7)
+        delta = self.omega_limit / nsamples
+        omega_value_big = np.linspace(delta, self.omega_limit, nsamples)
+        dos = self._dos_Wang1(A, omega_value_big)
+        rho_cumint = np.cumsum(dos) * delta
+        diff = (rho_cumint % 1)[1:] - (rho_cumint % 1)[:-1]
+        idx = np.where(diff < 0)[0]
+        omega_value = omega_value_big[idx]
+        logger.info(len(omega_value))
+        assert len(omega_value) == nb
+
+        # general form
+        c_j2 = 2.0 / np.pi * omega_value * self.func(omega_value) / self._dos_Wang1(A, omega_value)
 
-        return x, y_c, y_d
+        return omega_value, c_j2
 
 
-def param2mollist(alpha: float, raw_delta: Quantity, omega_c: Quantity, renormalization_p: float, n_phonons: int):
-    sdf = SpectralDensityFunction(alpha, omega_c)
+def param2mollist(
+    alpha: float,
+    raw_delta: Quantity,
+    omega_c: Quantity,
+    renormalization_p: float,
+    n_phonons: int,
+):
+    sdf = SpectralDensityFunction(alpha, omega_c, s=1)
     delta, max_omega = sdf.adiabatic_renormalization(raw_delta, renormalization_p)
-    omega_list, displacement_list = sdf.trapz(n_phonons, 0.0, max_omega.as_au())
-
-    ph_list = [Phonon.simplest_phonon(o, d) for o,d in zip(omega_list, displacement_list)]
-    return SpinBosonModel(Quantity(0), delta, ph_list)
+    omega_list, displacement_list = sdf.trapz(n_phonons, 0.0, max_omega)
+    omega_list, displacement_list = sdf.post_process(omega_list, displacement_list)
+    ph_list = [Phonon.simplest_phonon(o, d) for o, d in zip(omega_list, displacement_list)]
+    return SpinBosonModel(Quantity(0), Quantity(delta), ph_list)
diff --git a/renormalizer/sbm/sbm.py b/renormalizer/sbm/sbm.py
index a10e3521..79f783e9 100644
--- a/renormalizer/sbm/sbm.py
+++ b/renormalizer/sbm/sbm.py
@@ -1,12 +1,11 @@
 # -*- coding: utf-8 -*-
 
 import logging
-from functools import partial
 
 from renormalizer.model import Model
 from renormalizer.mps import Mpo, Mps
-from renormalizer.utils import TdMpsJob, Quantity, CompressConfig
-import numpy as np
+from renormalizer.utils import TdMpsJob, CompressConfig
+
 
 logger = logging.getLogger(__name__)
 
diff --git a/renormalizer/tn/__init__.py b/renormalizer/tn/__init__.py
new file mode 100644
index 00000000..6993573d
--- /dev/null
+++ b/renormalizer/tn/__init__.py
@@ -0,0 +1,5 @@
+from renormalizer.tn.node import TreeNodeBasis, TreeNodeTensor, copy_connection
+from renormalizer.tn.treebase import BasisTree
+from renormalizer.tn.tree import TTNO, TTNS
+# set the variable
+from renormalizer.tn.time_evolution import EVOLVE_METHODS
diff --git a/renormalizer/tn/gs.py b/renormalizer/tn/gs.py
new file mode 100644
index 00000000..b4814dda
--- /dev/null
+++ b/renormalizer/tn/gs.py
@@ -0,0 +1,131 @@
+import logging
+from typing import List
+
+
+import scipy
+
+from renormalizer.lib import davidson
+from renormalizer.mps.backend import primme, IMPORT_PRIMME_EXCEPTION, np
+from renormalizer.mps.matrix import asnumpy, asxp
+from renormalizer.tn.node import TreeNodeTensor
+from renormalizer.tn.tree import TTNS, TTNO, TTNEnviron
+from renormalizer.tn.hop_expr import hop_expr2
+
+
+logger = logging.getLogger(__name__)
+
+
+def optimize_ttns(ttns: TTNS, ttno: TTNO, procedure=None):
+    if procedure is None:
+        procedure = ttns.optimize_config.procedure
+    ttne = TTNEnviron(ttns, ttno)
+    e_list = []
+    for m, percent in procedure:
+        # todo: better converge condition
+        micro_e = optimize_recursion(ttns.root, ttns, ttno, ttne, m, percent)
+        logger.info(f"Micro e: {micro_e}")
+        e_list.append(micro_e[-1])
+    return e_list
+
+
+def optimize_recursion(
+    snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, m: int, percent: float = 0
+) -> List[float]:
+    """Optimize snode and all of its children"""
+    assert snode.children  # 2 site can't do only one node
+    micro_e = []
+    for ichild, child in enumerate(snode.children):
+        if child.children:
+            # optimize snode + child
+            e, c = optimize_2site(child, ttns, ttno, ttne)
+            micro_e.append(e)
+            # cano to child
+            ttns.update_2site(child, c, m, percent, cano_parent=False)
+            # update env
+            ttne.update_2site(child, ttns, ttno)
+            # recursive optimization
+            micro_e_child = optimize_recursion(child, ttns, ttno, ttne, m)
+            micro_e.extend(micro_e_child)
+
+        # optimize snode + child
+        e, c = optimize_2site(child, ttns, ttno, ttne)
+        micro_e.append(e)
+        # cano to snode
+        ttns.update_2site(child, c, m, percent, cano_parent=True)
+        # update env
+        ttne.update_2site(child, ttns, ttno)
+    return micro_e
+
+
+def optimize_2site(snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron):
+    cguess = ttns.merge_with_parent(snode)
+    qn_mask = ttns.get_qnmask(snode, include_parent=True)
+    cguess = cguess[qn_mask].ravel()
+    expr, hdiag = hop_expr2(snode, ttns, ttno, ttne)
+    hdiag = hdiag[qn_mask].ravel()
+
+    def hop(x):
+        cstruct = vec2tensor(x, qn_mask)
+        ret = expr(asxp(cstruct))[qn_mask].ravel()
+        return asnumpy(ret)
+
+    assert ttns.optimize_config.nroots == 1
+    algo: str = ttns.optimize_config.algo
+    e, c = eigh_iterative(hop, hdiag, cguess, algo)
+    c = vec2tensor(c, qn_mask)
+    return e, c
+
+
+def eigh_iterative(hop, hdiag, cguess, algo):
+    hdiag = asnumpy(hdiag)
+    cguess = asnumpy(cguess)
+    h_dim = len(hdiag)
+
+    if algo == "davidson":
+        precond = lambda x, e, *args: x / (hdiag - e + 1e-4)
+
+        e, c = davidson(hop, cguess, precond, max_cycle=100, nroots=1, max_memory=64000, verbose=0)
+    elif algo == "primme":
+        if primme is None:
+            logger.error("can not import primme")
+            raise IMPORT_PRIMME_EXCEPTION
+        precond = lambda x: scipy.sparse.diags(1 / (hdiag + 1e-4)) @ x
+        A = scipy.sparse.linalg.LinearOperator((h_dim, h_dim), matvec=hop, matmat=hop)
+        M = scipy.sparse.linalg.LinearOperator((h_dim, h_dim), matvec=precond, matmat=hop)
+        e, c = primme.eigsh(
+            A,
+            k=1,
+            which="SA",
+            v0=np.array(cguess).reshape(-1, 1),
+            OPinv=M,
+            method="PRIMME_DYNAMIC",
+            tol=1e-6,
+        )
+        c = c[:, 0]
+        e = e[0]
+    elif algo == "arpack":
+        A = scipy.sparse.linalg.LinearOperator((h_dim, h_dim), matvec=hop)
+        e, c = scipy.sparse.linalg.eigsh(A, k=1, which="SA", v0=cguess)
+        e = e[0]
+    elif algo == "direct":
+        # direct algorithm. Poor performance, debugging only.
+        a_list = []
+        for i in range(h_dim):
+            a = np.zeros(h_dim)
+            a[i] = 1
+            a_list.append(hop(a))
+        a = np.array(a_list)
+        assert np.allclose(a, a.conj().T)
+        evals, evecs = np.linalg.eigh(a)
+        e = evals[0]
+        c = evecs[:, 0]
+    else:
+        assert False
+
+    return e, c
+
+
+def vec2tensor(c, qn_mask):
+    cstruct = np.zeros(qn_mask.shape, dtype=c.dtype)
+    np.place(cstruct, qn_mask, c)
+    return cstruct
diff --git a/renormalizer/tn/hop_expr.py b/renormalizer/tn/hop_expr.py
new file mode 100644
index 00000000..2d5b8752
--- /dev/null
+++ b/renormalizer/tn/hop_expr.py
@@ -0,0 +1,153 @@
+import opt_einsum as oe
+
+from renormalizer.mps.backend import np
+from renormalizer.mps.matrix import asxp
+from renormalizer.tn.node import TreeNodeTensor
+from renormalizer.tn.tree import TTNS, TTNO, TTNEnviron
+
+
+def hop_expr0(snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron):
+    # zero site, used in tdvp time evolution
+    # assuming the first index connects child and the second index connects parent
+    # #--------o---------#
+    # child--coeff--parent
+    enode = ttne.node_list[ttns.node_idx[snode]]
+
+    args = []
+    input_indices = []
+    output_indices = []
+    shape = []
+
+    tensor = enode.parent.environ_children[enode.idx_as_child]
+    shape.append(tensor.shape[0])
+    args.append(tensor)
+    indices = ttne.get_child_indices(enode.parent, enode.idx_as_child, ttns, ttno)
+    output_indices.append(indices[0])
+    input_indices.append(indices[2])
+    args.append(indices)
+
+    tensor = enode.environ_parent
+    shape.append(tensor.shape[0])
+    args.append(tensor)
+    indices = ttne.get_parent_indices(enode, ttns, ttno)
+    assert len(indices) == 3
+    indices[0] = tuple(list(indices[0]) + ["hop0_conj"])
+    indices[2] = tuple(list(indices[2]) + ["hop0"])
+    output_indices.append(indices[0])
+    input_indices.append(indices[2])
+    args.append(indices)
+
+    expr = _contract_expression(args, shape, input_indices, output_indices)
+
+    return expr
+
+
+def hop_expr1(snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, return_hdiag=False):
+    # build one site effective hamiltonian operator as an opt_einsum expression
+    enode = ttne.node_list[ttns.node_idx[snode]]
+    onode = ttno.node_list[ttns.node_idx[snode]]
+
+    args = []
+    # enode children environments
+    for i, echild_environchild in enumerate(enode.environ_children):
+        args.append(echild_environchild)
+        args.append(ttne.get_child_indices(enode, i, ttns, ttno))
+    # parent environments
+    args.append(enode.environ_parent)
+    args.append(ttne.get_parent_indices(enode, ttns, ttno))
+    # operator
+    args.extend([onode.tensor, ttno.get_node_indices(onode)])
+
+    # input and output
+    input_indices = ttns.get_node_indices(snode, ttno=ttno)
+    output_indices = ttns.get_node_indices(snode, conj=True)
+
+    shape = snode.shape
+    # cache the contraction path
+    expr = _contract_expression(args, shape, input_indices, output_indices)
+    if not return_hdiag:
+        return expr
+    else:
+        hdiag = _get_hdiag(args, input_indices)
+        return expr, hdiag
+
+
+def hop_expr2(snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron):
+    # build two-site effective hamiltonian operator as an opt_einsum expression
+    sparent = snode.parent
+    enode = ttne.node_list[ttns.node_idx[snode]]
+    eparent = ttne.node_list[ttns.node_idx[sparent]]
+    onode = ttno.node_list[ttns.node_idx[snode]]
+    oparent = ttno.node_list[ttns.node_idx[sparent]]
+
+    args = []
+    # enode children environments
+    for i, echild_environchild in enumerate(enode.environ_children):
+        args.append(echild_environchild)
+        args.append(ttne.get_child_indices(enode, i, ttns, ttno))
+
+    # eparent children environments
+    for i, enode_environchild in enumerate(eparent.environ_children):
+        if eparent.children[i] is enode:
+            continue
+        args.append(enode_environchild)
+        args.append(ttne.get_child_indices(eparent, i, ttns, ttno))
+
+    # eparent parent environments
+    args.append(eparent.environ_parent)
+    args.append(ttne.get_parent_indices(eparent, ttns, ttno))
+
+    # operator
+    args.extend([oparent.tensor, ttno.get_node_indices(oparent)])
+    args.extend([onode.tensor, ttno.get_node_indices(onode)])
+
+    # input and output
+    input_indices = ttns.get_node_indices(snode, include_parent=True, ttno=ttno)
+    output_indices = ttns.get_node_indices(snode, conj=True, include_parent=True)
+
+    # shape
+    shape = list(snode.shape[:-1])
+    shape_parent = list(snode.parent.shape)
+    del shape_parent[snode.parent.children.index(snode)]
+    shape += shape_parent
+    # cache the contraction path
+    expr = _contract_expression(args, shape, input_indices, output_indices)
+    hdiag = _get_hdiag(args, input_indices)
+    return expr, hdiag
+
+
+def _contract_expression(args, x_shape, x_indices, y_indices):
+    # contract_expression in interleaved format
+    args_fake = args.copy()
+    args_fake.extend([np.empty(x_shape), x_indices])
+    args_fake.append(y_indices)
+    indices, tensors = oe.parser.convert_interleaved_input(args_fake)
+    args = [asxp(t) for t in tensors[:-1]] + [x_shape]
+    expr = oe.contract_expression(
+        indices,
+        *args,
+        constants=list(range(len(tensors)))[:-1],
+    )
+    return expr
+
+
+def _get_hdiag(args, input_indices):
+    new_args = []
+    for arg in args:
+        if not isinstance(arg, (tuple, list)):
+            # tensors
+            new_args.append(asxp(arg))
+            continue
+        # indices
+        arg = list(arg)
+        if arg[0][-5:] == "_conj":
+            # the environ
+            arg[0] = arg[0][:-5]
+        elif arg[1] == "up":
+            # mpo part
+            arg[1] = "down"
+        else:
+            pass
+        new_args.append(tuple(arg))
+    new_args.append(input_indices)
+    return oe.contract(*new_args)
diff --git a/renormalizer/tn/node.py b/renormalizer/tn/node.py
new file mode 100644
index 00000000..21d630b7
--- /dev/null
+++ b/renormalizer/tn/node.py
@@ -0,0 +1,135 @@
+from typing import List, Dict, Union, Sequence
+
+from renormalizer.mps.backend import np, backend
+from renormalizer.mps.matrix import asnumpy
+from renormalizer.model.basis import BasisSet
+
+
+class TreeNode:
+    def __init__(self):
+        self.children: List[__class__] = []
+        self.parent: __class__ = None
+
+    def add_child(self, node: Union["TreeNode", Sequence["TreeNode"]]) -> "TreeNode":
+        if isinstance(node, TreeNode):
+            nodes = [node]
+        else:
+            nodes = node
+
+        for node in nodes:
+            if node.parent is not None:
+                raise ValueError("Node already has parent")
+            self.children.append(node)
+            node.parent = self
+
+        return self
+
+    @property
+    def idx_as_child(self) -> int:
+        assert self.parent
+        return self.parent.children.index(self)
+
+
+class TreeNodeBasis(TreeNode):
+    # tree node whose data is basis sets
+    def __init__(self, basis_sets: List[BasisSet]):
+        super().__init__()
+        self.basis_sets: List[BasisSet] = basis_sets
+        self.n_sets = len(basis_sets)
+        qn_size_list = [b.sigmaqn.shape[1] for b in self.basis_sets]
+        if len(set(qn_size_list)) != 1:
+            raise ValueError(f"Inconsistent quantum number size: {set(qn_size_list)}")
+        self.qn_size: int = qn_size_list[0]
+        self.dofs = [b.dofs for b in basis_sets]
+        self.pbond_dims = [len(b.sigmaqn) for b in self.basis_sets]
+
+    def copy(self):
+        new = self.__class__(self.basis_sets)
+        if self.parent is not None:
+            new.parent = self.parent.copy()
+        new.children = self.children.copy()
+        return new
+
+
+class TreeNodeTensor(TreeNode):
+    # tree node whose data is numerical tensors for each TTN node/site
+    def __init__(self, tensor, qn=None):
+        super().__init__()
+        self.tensor: np.ndarray = tensor
+        self.qn: np.ndarray = qn
+
+    def check_canonical(self, atol=None, assertion=True):
+        if atol is None:
+            atol = backend.canonical_atol
+        tensor = self.tensor.reshape(-1, self.tensor.shape[-1])
+        s = tensor.conj().T @ tensor
+        res = np.allclose(s, np.eye(s.shape[0]), atol=atol)
+        if assertion:
+            assert res
+        return res
+
+    @property
+    def shape(self):
+        return self.tensor.shape
+
+    @property
+    def tensor(self):
+        return self._tensor
+
+    @tensor.setter
+    def tensor(self, tensor):
+        if np.iscomplexobj(tensor):
+            dtype = backend.complex_dtype
+        else:
+            dtype = backend.real_dtype
+        self._tensor = np.asarray(asnumpy(tensor), dtype=dtype)
+
+    @property
+    def qn(self):
+        return self._qn
+
+    @qn.setter
+    def qn(self, qn):
+        self._qn = np.array(qn)
+
+
+class TreeNodeEnviron(TreeNode):
+    # tree node whose data is environmental tensors
+    def __init__(self):
+        super().__init__()
+        self.parent: TreeNodeEnviron = None
+        # environ from parent
+        self.environ_parent: np.ndarray = None
+        # environ from children
+        self.environ_children: List[np.ndarray] = []
+
+
+NodeUnion = Union[TreeNodeTensor, TreeNodeBasis, TreeNodeEnviron]
+
+
+def copy_connection(source_node_list: List[NodeUnion], target_node_list: List[NodeUnion]) -> NodeUnion:
+    """
+    Create a new tree with the same connection structure as the source tree in the target tree.
+
+    Parameters
+    ----------
+    source_node_list : List[NodeUnion]
+        The list of nodes in the source tree.
+    target_node_list : List[NodeUnion]
+        The list of nodes in the target tree.
+
+    Returns
+    -------
+    NodeUnion
+        The root node of the target tree.
+    """
+    node2idx: Dict[NodeUnion, int] = {n: i for i, n in enumerate(source_node_list)}
+    root = None
+    for source_node, target_node in zip(source_node_list, target_node_list):
+        for child in source_node.children:
+            idx = node2idx[child]
+            target_node.add_child(target_node_list[idx])
+        if source_node.parent is None:
+            root = target_node
+    assert root is not None
+    return root
diff --git a/renormalizer/tn/symbolic_mpo.py b/renormalizer/tn/symbolic_mpo.py
new file mode 100644
index 00000000..0e9c3752
--- /dev/null
+++ b/renormalizer/tn/symbolic_mpo.py
@@ -0,0 +1,105 @@
+from itertools import chain
+import logging
+from typing import List
+
+from renormalizer.mps.backend import np
+from renormalizer import Op, Model
+from renormalizer.model.basis import BasisSet
+from renormalizer.tn.treebase import BasisTree
+from renormalizer.mps.symbolic_mpo import _terms_to_table, _transform_table, _construct_symbolic_mpo_one_site, OpTuple
+
+
+logger = logging.getLogger(__name__)
+
+
+# translate the numbers into symbolic Matrix Operator
+def compose_symbolic_mo_general(in_ops_list, out_ops, primary_ops, k):
+    shape = [len(in_ops) for in_ops in in_ops_list] + [len(out_ops)]
+    mo = np.full(shape, None, dtype=object)
+    for i, _ in np.ndenumerate(mo):
+        mo[i] = []
+    for iop, out_op in enumerate(out_ops):
+        for composed_op in out_op:
+            if in_ops_list:
+                in_idx = tuple(composed_op.symbol[:-k])
+                l = mo[in_idx][iop]
+            else:
+                l = mo[iop]
+            op = composed_op.factor
+            for s in composed_op.symbol[-k:]:
+                op = op * primary_ops[s]
+            l.append(op)
+    return mo
+
+
+# translate symbolic Matrix Operator to numerical matrix operator defined with certain basis
+def symbolic_mo_to_numeric_mo_general(basis_sets: List[BasisSet], mo, dtype):
+    model = Model(basis_sets, [])
+    pdims = [b.nbas for b in basis_sets]
+    shape = list(mo.shape) + list(chain(*[[pdim, pdim] for pdim in pdims]))
+    mo_tensor = np.zeros(shape, dtype=dtype)
+    terms: List[Op]
+    for i, terms in np.ndenumerate(mo):
+        for term in terms:
+            term_split, factor = term.split_elementary(model.dof_to_siteidx)
+            assert len(term_split) == len(basis_sets)
+            mo_elem = np.eye(1) * factor
+            for symbol, b in zip(term_split, basis_sets):
+                mo_elem = np.tensordot(mo_elem, b.op_mat(symbol)[None, :, :, None], axes=1)
+            assert not np.iscomplexobj(mo_elem), "complex operator not supported yet"
+            mo_tensor[i] += mo_elem[0, ..., 0]
+
+    return np.moveaxis(mo_tensor, mo.ndim - 1, -1)
+
+
+def construct_symbolic_mpo(tn: BasisTree, terms: List[Op], const: float = 0):
+    algo = "Hopcroft-Karp"
+    nodes = tn.postorder_list()
+    basis = list(chain(*[n.basis_sets for n in nodes]))
+    model = Model(basis, [])
+    qn_size = model.qn_size
+    table, factor = _terms_to_table(model, terms, const)
+    table, factor, primary_ops = _transform_table(table, factor)
+
+    dummy_in_ops = [[OpTuple([0], qn=np.zeros(qn_size, dtype=int), factor=1)]]
+    out_ops: List[List[OpTuple]]
+    out_ops_list = []
+
+    for i, node in enumerate(nodes):
+        k = node.n_sets
+        if not node.children:
+            ta = np.zeros((table.shape[0], 1), dtype=np.uint16)
+            table = np.concatenate((ta, table), axis=1)
+            table_row = table[:, : k + 1]
+            table_col = table[:, k + 1 :]
+            in_ops_list = [dummy_in_ops]
+        else:
+            # the children must have been visited
+            children_idx = [nodes.index(n) for n in node.children]
+            assert np.all(np.array(children_idx) < i)
+            in_ops_list = [out_ops_list[i] for i in children_idx]
+            m = len(node.children)
+            # roll relevant columns to the front
+            table = np.roll(table, m, axis=1)
+            table_row = table[:, : m + k]
+            table_col = table[:, m + k :]
+        out_ops, table, factor = _construct_symbolic_mpo_one_site(
+            table_row, table_col, in_ops_list, factor, primary_ops, algo, k
+        )
+        # move the new column at the first index to the last index
+        table = np.roll(table, -1, axis=1)
+        out_ops_list.append(out_ops)
+
+    mpo = []
+    for i, node in enumerate(nodes):
+        children_idx = [nodes.index(n) for n in node.children]
+        in_ops_list = [out_ops_list[i] for i in children_idx]
+        mo = compose_symbolic_mo_general(in_ops_list, out_ops_list[i], primary_ops, node.n_sets)
+        mpo.append(mo)
+
+    mpoqn = []
+    for out_ops in out_ops_list:
+        qn = np.array([out_op[0].qn for out_op in out_ops])
+        mpoqn.append(qn)
+
+    return mpo, mpoqn
diff --git a/renormalizer/tn/time_evolution.py b/renormalizer/tn/time_evolution.py
new file mode 100644
index 00000000..13daeb64
--- /dev/null
+++ b/renormalizer/tn/time_evolution.py
@@ -0,0 +1,302 @@
+from math import factorial
+from typing import Union, List, Tuple
+import logging
+
+import scipy
+from scipy import stats
+import opt_einsum as oe
+
+from renormalizer.mps.lib import compressed_sum
+from renormalizer.mps.backend import np, xp
+from renormalizer.mps.matrix import asxp
+from renormalizer.lib import solve_ivp, expm_krylov
+from renormalizer.utils.configs import EvolveMethod
+from renormalizer.tn.node import TreeNodeTensor
+from renormalizer.tn.tree import TTNO, TTNS, TTNEnviron, EVOLVE_METHODS
+from renormalizer.tn.hop_expr import hop_expr0, hop_expr1, hop_expr2
+
+
+logger = logging.getLogger(__name__)
+
+
+def time_derivative_vmf(ttns: TTNS, ttno: TTNO):
+    # todo: benchmark and optimize
+    # parallel over multiple processors?
+    environ_s = TTNEnviron(ttns, TTNO.identity(ttns.basis))
+    environ_h = TTNEnviron(ttns, ttno)
+
+    deriv_list = []
+    for inode, node in enumerate(ttns.node_list):
+        hop = hop_expr1(node, ttns, ttno, environ_h)
+        # idx1: children+physical, idx2: parent
+        dim_parent = node.shape[-1]
+        tensor = asxp(node.tensor)
+        shape_2d = (-1, dim_parent)
+        deriv = hop(tensor).reshape(shape_2d)
+        if node.parent is not None:
+            # apply projector and S^-1
+            tensor = tensor.reshape(shape_2d)
+            proj = tensor.conj() @ tensor.T
+            ovlp = environ_s.node_list[inode].environ_parent.reshape(dim_parent, dim_parent)
+            ovlp_inv = regularized_inversion(ovlp, ttns.evolve_config.reg_epsilon)
+            deriv = oe.contract("bf, bg, fh -> gh", deriv, xp.eye(proj.shape[0]) - proj, asxp(ovlp_inv.T))
+        qnmask = ttns.get_qnmask(node).reshape(deriv.shape)
+        deriv_list.append(deriv[qnmask].ravel())
+    return np.concatenate(deriv_list)
+
+
+def regularized_inversion(m, eps):
+    evals, evecs = scipy.linalg.eigh(m)
+    weight = np.exp(-evals / eps)
+    evals = evals + eps * weight
+    return evecs @ np.diag(1 / evals) @ evecs.T.conj()
+
+
+def evolve_tdvp_vmf(ttns: TTNS, ttno: TTNO, coeff: Union[complex, float], tau: float, first_step=None):
+    def ivp_func(t, params):
+        ttns_t = TTNS.from_tensors(ttns, params)
+        return coeff * time_derivative_vmf(ttns_t, ttno)
+
+    init_y = np.concatenate([node.tensor[ttns.get_qnmask(node)].ravel() for node in ttns.node_list])
+    atol = ttns.evolve_config.ivp_atol
+    rtol = ttns.evolve_config.ivp_rtol
+    sol = solve_ivp(ivp_func, (0, tau), init_y, first_step=first_step, atol=atol, rtol=rtol)
+    logger.info(f"VMF func called: {sol.nfev}. RKF steps: {len(sol.t)}")
+    new_ttns = TTNS.from_tensors(ttns, sol.y[:, -1])
+    new_ttns.canonicalise()
+    return new_ttns
+
+
+def evolve_prop_and_compress_tdrk4(ttns: TTNS, ttno: TTNO, coeff: Union[complex, float], tau: float):
+    termlist = [ttns]
+    for i in range(4):
+        termlist.append(ttno.contract(termlist[-1]))
+    for i, term in enumerate(termlist):
+        term.scale((coeff * tau) ** i / factorial(i), inplace=True)
+    return compressed_sum(termlist)
+
+
+def evolve_tdvp_ps(ttns: TTNS, ttno: TTNO, coeff: Union[complex, float], tau: float):
+    ttns.check_canonical()
+    # second order 1-site projector splitting
+    ttne = TTNEnviron(ttns, ttno)
+
+    # in MPS language: left to right sweep
+    local_steps1 = _tdvp_ps_forward(ttns, ttno, ttne, coeff, tau / 2)
+    # in MPS language: right to left sweep
+    local_steps2 = _tdvp_ps_backward(ttns, ttno, ttne, coeff, tau / 2)
+
+    # Used for consistency with MPS
+    # # in MPS language: right to left sweep
+    # local_steps1 = _tdvp_ps_backward(ttns.root, ttns, ttno, tte, coeff, tau / 2)
+    # # in MPS language: left to right sweep
+    # local_steps2 = _tdvp_ps_forward(ttns.root, ttns, ttno, tte, coeff, tau / 2)
+
+    steps_stat = stats.describe(local_steps1 + local_steps2)
+    # logger.debug(f"Local Krylov space forward: {local_steps1}")
+    # logger.debug(f"Local Krylov space backward: {local_steps2}")
+    logger.debug(f"TDVP-PS Krylov space: {steps_stat}")
+    return ttns
+
+
+def _tdvp_ps_forward(ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, coeff: Union[complex, float], tau: float) -> List[int]:
+    local_steps: List[int] = []
+    # current node and the child that has already been processed (once popped out)
+    stack: List[Tuple[TreeNodeTensor, int]] = [(ttns.root, -1)]
+    while stack:
+        snode, ichild = stack[-1]
+        # no children to evolve
+        if (not snode.children) or (ichild == len(snode.children) - 1):
+            ms, j = evolve_1site(snode, ttns, ttno, ttne, coeff, tau)
+            snode.tensor = ms.reshape(snode.shape)
+            local_steps.append(j)
+
+            if snode.parent is None:
+                assert len(stack) == 1
+                stack.pop()
+                continue
+            # decompose, the first index for parent, the second index for child
+            ms = ttns.decompose_to_parent(snode)
+            # update env
+            ttne.build_children_environ_node(snode, ttns, ttno)
+            # backward time evolution for snode
+            ms_t, j = evolve_0site(ms.T, snode, ttns, ttno, ttne, coeff, -tau)
+            ttns.merge_to_parent(snode, ms_t.reshape(ms.T.shape).T)
+            local_steps.append(j)
+
+            stack.pop()
+            continue
+
+        # there are children that has not been evolved
+        ichild += 1
+        child = snode.children[ichild]
+        # cano to child
+        ttns.push_cano_to_child(snode, ichild)
+        # update env
+        ttne.build_parent_environ_node(snode, ichild, ttns, ttno)
+        stack[-1] = (snode, ichild)
+        stack.append((child, -1))
+
+    return local_steps
+
+
+def _tdvp_ps_backward(ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, coeff: Union[complex, float], tau: float) -> List[int]:
+    local_steps: List[int] = []
+    # current node and the child that has already been processed (once popped out)
+    stack: List[Tuple[TreeNodeTensor, int]] = [(ttns.root, -1)]
+    while stack:
+        snode, ichild = stack[-1]
+        if ichild == -1:
+            ms, j = evolve_1site(snode, ttns, ttno, ttne, coeff, tau)
+            snode.tensor = ms.reshape(snode.shape)
+            local_steps.append(j)
+        if ichild == len(snode.children) - 1:
+            if snode is not ttns.root:
+                ttns.push_cano_to_parent(snode)
+                # update env
+                ttne.build_children_environ_node(snode, ttns, ttno)
+            stack.pop()
+            continue
+        ichild += 1
+        child = snode.children[ichild]
+        # decompose, the first index for child, the second index for parent
+        ms = ttns.decompose_to_child(snode, ichild)
+        # update env
+        ttne.build_parent_environ_node(snode, ichild, ttns, ttno)
+        # backward time evolution for snode
+        shape = ms.shape
+        ms, j = evolve_0site(ms, child, ttns, ttno, ttne, coeff, -tau)
+        ttns.merge_to_child(snode, ichild, ms.reshape(shape))
+        local_steps.append(j)
+        stack[-1] = snode, ichild
+        stack.append((child, -1))
+
+    return local_steps
+
+
+def evolve_tdvp_ps2(ttns: TTNS, ttno: TTNO, coeff: Union[complex, float], tau: float):
+    ttns.check_canonical()
+    # second order 2-site projector splitting
+    tte = TTNEnviron(ttns, ttno)
+    # in MPS language: left to right sweep
+    local_steps1 = _tdvp_ps2_recursion_forward(ttns.root, ttns, ttno, tte, coeff, tau / 2)
+    # in MPS language: right to left sweep
+    local_steps2 = _tdvp_ps2_recursion_backward(ttns.root, ttns, ttno, tte, coeff, tau / 2)
+    steps_stat = stats.describe(local_steps1 + local_steps2)
+    logger.debug(f"TDVP-PS Krylov space: {steps_stat}")
+    return ttns
+
+
+def _tdvp_ps2_recursion_forward(
+    snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, coeff: Union[complex, float], tau: float
+) -> List[int]:
+    """time evolution all of snode's children (without evolve snode!).
+    The exception is when snode == ttns.root, which is evolved.
+    Cano center at snode when entering and leaving"""
+    assert snode.children  # 2 site can't do only one node
+    # todo: update to more general cases like truncation based on singular values
+    m = ttns.compress_config.bond_dim_max_value
+    local_steps: List[int] = []
+    for ichild, child in enumerate(snode.children):
+        if child.children:
+            # cano to child
+            ttns.push_cano_to_child(snode, ichild)
+            # update env
+            ttne.update_1bond(child, ttns, ttno)
+            # recursive time evolution
+            local_steps_child = _tdvp_ps2_recursion_forward(child, ttns, ttno, ttne, coeff, tau)
+            local_steps.extend(local_steps_child)
+
+        # forward time evolution for snode + child
+        ms2, j = evolve_2site(child, ttns, ttno, ttne, coeff, tau)
+        local_steps.append(j)
+        # cano to snode
+        ttns.update_2site(child, ms2, m, cano_parent=True)
+        # update env
+        ttne.update_2site(child, ttns, ttno)
+        # backward time evolution for snode
+        if snode is ttns.root and ichild == len(snode.children) - 1:
+            continue
+        ms, j = evolve_1site(snode, ttns, ttno, ttne, coeff, -tau)
+        snode.tensor = ms.reshape(snode.shape)
+        local_steps.append(j)
+        # update env
+        ttne.update_1site(snode, ttns, ttno)
+    return local_steps
+
+
+def _tdvp_ps2_recursion_backward(
+    snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, coeff: Union[complex, float], tau: float
+) -> List[int]:
+    """time evolution all of snode's children (without evolve snode!).
+    The exception is when snode == ttns.root, which is evolved.
+    Cano center at snode when entering and leaving"""
+    assert snode.children  # 2 site can't do only one node
+    # todo: update to more general cases like truncation based on singular values
+    m = ttns.compress_config.bond_dim_max_value
+    local_steps: List[int] = []
+    for ichild, child in reversed(list(enumerate(snode.children))):
+        # backward time evolution for snode
+        if not (snode is ttns.root and ichild == len(snode.children) - 1):
+            ms, j = evolve_1site(snode, ttns, ttno, ttne, coeff, -tau)
+            snode.tensor = ms.reshape(snode.shape)
+            local_steps.append(j)
+            # update env
+            ttne.update_1site(snode, ttns, ttno)
+
+        # forward time evolution for snode + child
+        ms2, j = evolve_2site(child, ttns, ttno, ttne, coeff, tau)
+        local_steps.append(j)
+        # cano to snode
+        ttns.update_2site(child, ms2, m, cano_parent=not child.children)
+        # update env
+        ttne.update_2site(child, ttns, ttno)
+
+        if child.children:
+            # recursive time evolution
+            local_steps_child = _tdvp_ps2_recursion_backward(child, ttns, ttno, ttne, coeff, tau)
+            local_steps.extend(local_steps_child)
+            # cano to snode
+            ttns.push_cano_to_parent(child)
+            # update env
+            ttne.update_1bond(child, ttns, ttno)
+    return local_steps
+
+
+def evolve_2site(
+    snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, coeff: Union[complex, float], tau: float
+):
+    # evolve snode and parent
+    ms2 = ttns.merge_with_parent(snode)
+    hop, _ = hop_expr2(snode, ttns, ttno, ttne)
+    ms2_t, j = expm_krylov(lambda y: hop(y.reshape(ms2.shape)).ravel(), coeff * tau, ms2.ravel())
+    return ms2_t, j
+
+
+def evolve_1site(
+    snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, coeff: Union[complex, float], tau: float
+):
+    ms = snode.tensor
+    hop = hop_expr1(snode, ttns, ttno, ttne)
+    ms_t, j = expm_krylov(lambda y: hop(y.reshape(ms.shape)).ravel(), coeff * tau, ms.ravel())
+    return ms_t, j
+
+
+def evolve_0site(
+    ms: np.ndarray,
+    snode: TreeNodeTensor,
+    ttns: TTNS,
+    ttno: TTNO,
+    ttne: TTNEnviron,
+    coeff: Union[complex, float],
+    tau: float,
+):
+    hop = hop_expr0(snode, ttns, ttno, ttne)
+    ms_t, j = expm_krylov(lambda y: hop(y.reshape(ms.shape)).ravel(), coeff * tau, ms.ravel())
+    return ms_t, j
+
+
+EVOLVE_METHODS[EvolveMethod.tdvp_vmf] = evolve_tdvp_vmf
+EVOLVE_METHODS[EvolveMethod.prop_and_compress_tdrk4] = evolve_prop_and_compress_tdrk4
+EVOLVE_METHODS[EvolveMethod.tdvp_ps] = evolve_tdvp_ps
+EVOLVE_METHODS[EvolveMethod.tdvp_ps2] = evolve_tdvp_ps2
diff --git a/renormalizer/tn/tree.py b/renormalizer/tn/tree.py
new file mode 100644
index 00000000..466a9ba8
--- /dev/null
+++ b/renormalizer/tn/tree.py
@@ -0,0 +1,1313 @@
+from typing import List, Dict, Tuple, Union, Callable
+from numbers import Number
+import logging
+
+import scipy
+import opt_einsum as oe
+from print_tree import print_tree
+
+from renormalizer import Op, Mps, Model
+from renormalizer.model.basis import BasisSet, BasisDummy
+from renormalizer.mps.backend import np, backend, xp
+from renormalizer.mps.matrix import asnumpy, asxp_oe_args, tensordot
+from renormalizer.mps.svd_qn import add_outer, svd_qn, blockrecover, get_qn_mask
+from renormalizer.mps.lib import select_basis
+from renormalizer.mps.mps import normalize
+from renormalizer.utils.configs import CompressConfig, OptimizeConfig, EvolveConfig, EvolveMethod
+from renormalizer.utils import calc_vn_entropy
+from renormalizer.tn.node import TreeNodeTensor, TreeNodeBasis, copy_connection, TreeNodeEnviron
+from renormalizer.tn.treebase import Tree, BasisTree
+from renormalizer.tn.symbolic_mpo import construct_symbolic_mpo, symbolic_mo_to_numeric_mo_general
+
+
+logger = logging.getLogger(__name__)
+
+
+class TTNBase(Tree):
+    # A tree whose tree node is TreeNodeTensor
+
+    @classmethod
+    def load(cls, basis: BasisTree, fname: str, other_attrs=None):
+        npload = np.load(fname, allow_pickle=True)
+        assert npload["version"] == "0.1"
+
+        nsites = int(npload["nsites"])
+        nodes = []
+        for i in range(nsites):
+            tensor = npload[f"tensor_{i}"]
+            qn = npload[f"qn_{i}"]
+            nodes.append(TreeNodeTensor(tensor, qn))
+        copy_connection(basis.node_list, nodes)
+        instance = cls(basis, root=nodes[0])
+        for attr in other_attrs:
+            setattr(instance, attr, npload[attr])
+        return instance
+
+    def __init__(self, basis: BasisTree, root: TreeNodeTensor):
+        self.basis = basis
+        super().__init__(root)
+
+    def dump(self, fname: str, other_attrs=None):
+        if other_attrs is None:
+            other_attrs = []
+
+        data_dict = {
+            "version": "0.1",
+            "nsites": len(self),
+        }
+
+        for attr in other_attrs:
+            data_dict[attr] = getattr(self, attr)
+
+        for i, node in enumerate(self.node_list):
+            data_dict[f"tensor_{i}"] = node.tensor
+            data_dict[f"qn_{i}"] = node.qn
+
+        try:
+            np.savez(fname, **data_dict)
+        except Exception:
+            logger.exception(f"Dump MP failed.")
+
+    def print_shape(self, full: bool = False, print_function: Callable = None):
+        """
+        Print the shape of the TTN tree
+
+        Parameters
+        ----------
+        full: bool
+            whether print the full shape or only the shape of the last bond
+        print_function
+            the function used for printing. If None, use ``print``.
+            Could use ``logger.info`` instead.
+        """
+
+        class _print_shape(print_tree):
+            def get_children(self, node):
+                return node.children
+
+            def get_node_str(self, node: TreeNodeTensor):
+                if full:
+                    return str(node.tensor.shape)
+                else:
+                    return str(node.tensor.shape[-1])
+
+        tree = _print_shape(self.root)
+        if print_function is not None:
+            for row in tree.rows:
+                print_function(row)
+
+    def print_vn_entropy(self, print_function=None):
+        # todo: move ttns.compress to here
+        _, s_list = self.compress(ret_s=True)
+        vn_entropy = [calc_vn_entropy(s**2) for s in s_list]
+        nodes = [TreeNodeTensor([entropy]) for entropy in vn_entropy]
+        copy_connection(self.node_list, nodes)
+
+        class print_data(print_tree):
+            def get_children(self, node):
+                return node.children
+
+            def get_node_str(self, node: TreeNodeTensor):
+                return str(node.tensor.shape[-1])
+
+        tree = print_data(nodes[0])
+        if print_function is not None:
+            for row in tree.rows:
+                print_function(row)
+
+    @property
+    def bond_dims(self):
+        return [node.tensor.shape[-1] for node in self]
+
+    @property
+    def bond_dims_mean(self) -> int:
+        # duplicate with Mps
+        return int(round(np.mean(self.bond_dims)))
+
+    @property
+    def qntot(self):
+        return self.root.qn[0]
+
+
+class TTNO(TTNBase):
+    @classmethod
+    def identity(cls, basis: BasisTree):
+        if not basis.identity_ttno:
+            basis.identity_ttno = cls(basis, [basis.identity_op])
+        return basis.identity_ttno
+
+    def __init__(self, basis: BasisTree, terms: Union[List[Op], Op], root: TreeNodeTensor = None):
+        self.basis: BasisTree = basis
+        if isinstance(terms, Op):
+            terms = [terms]
+        self.terms: List[Op] = terms
+
+        if not root:
+            symbolic_mpo, mpoqn = construct_symbolic_mpo(basis, terms)
+            # from renormalizer.mps.symbolic_mpo import _format_symbolic_mpo
+            # print(_format_symbolic_mpo(symbolic_mpo))
+            node_list_basis = self.basis.postorder_list()
+            node_list_op = []
+            for impo, (mo, qn) in enumerate(zip(symbolic_mpo, mpoqn)):
+                node_basis: TreeNodeBasis = node_list_basis[impo]
+                mo_mat = symbolic_mo_to_numeric_mo_general(node_basis.basis_sets, mo, backend.real_dtype)
+                node_list_op.append(TreeNodeTensor(mo_mat, qn))
+            root: TreeNodeTensor = copy_connection(node_list_basis, node_list_op)
+        super().__init__(basis, root)
+        # tensor node to basis node
+        self.tn2bn = {tn: bn for tn, bn in zip(self.node_list, self.basis.node_list)}
+        self.tn2dofs = {tn: bn.dofs for tn, bn in self.tn2bn.items()}
+
+    def apply(self, ttns: "TTNS", canonicalise: bool = False) -> "TTNS":
+        """
+        Apply the operator to the TTNS.
+
+        Parameters
+        ----------
+        ttns: TTNS
+            the TTNS to apply the operator to
+        canonicalise: bool
+            whether canonicalise the new TTNS
+
+        Returns
+        -------
+        The new TTNS
+        """
+        new = ttns.metacopy()
+
+        for snode1, snode2, onode in zip(ttns, new, self):
+            assert len(snode1.children) == len(onode.children)
+
+            indices1 = ttns.get_node_indices(snode1, ttno=self)
+            indices2 = self.get_node_indices(onode)
+            args = [snode1.tensor, indices1, onode.tensor, indices2]
+            output_indices = []
+            output_shape = []
+            # children indices
+            for i in range(len(snode1.children)):
+                output_shape.append(snode1.shape[i] * onode.shape[i])
+                output_indices.extend([indices1[i], indices2[i]])
+            # physical indices
+            bnode = ttns.tn2bn[snode1]
+            for i in range(bnode.n_sets):
+                output_shape.append(snode1.shape[len(snode1.children) + i])
+                output_indices.append(("up", str(bnode.dofs[i])))
+            # parent indices
+            output_shape.append(snode1.shape[-1] * onode.shape[-1])
+            output_indices.extend([indices1[-1], indices2[-1]])
+            args.append(output_indices)
+            # do contraction
+            res = oe.contract(*(asxp_oe_args(args))).reshape(output_shape)
+            snode2.tensor = res
+            snode2.qn = add_outer(snode1.qn, onode.qn).reshape(output_shape[-1], ttns.basis.qn_size)
+
+        new.check_shape()
+        if canonicalise:
+            new.canonicalise()
+        return new
+
+    def contract(self, ttns: "TTNS", algo="svd") -> "TTNS":
+        """
+        Contract the operator to the TTNS.
+        Parameters
+        ----------
+        ttns: TTNS
+            the TTNS to contract the operator to
+        algo: str
+            the algorithm to use for the compression after applying the operator
+
+        Returns
+        -------
+        The new TTNS
+        """
+        assert algo == "svd", "variational compress not supported yet"
+        new_ttns = self.apply(ttns)
+        new_ttns.canonicalise()
+        new_ttns.compress()
+        return new_ttns
+
+    def todense(self, order: List[BasisSet] = None) -> np.ndarray:
+        """
+        Convert the TTNO operator to dense matrix.
+
+        Parameters
+        ----------
+        order: list
+            the order of the basis sets for the resulting tensor
+
+        Returns
+        -------
+        The dense matrix
+        """
+        _id = str(id(self))
+        args = self.to_contract_args("up", "down")
+        if order is None:
+            order = self.basis.basis_list
+        indices_up = []
+        indices_down = []
+        for basis in order:
+            if isinstance(basis, BasisDummy):
+                continue
+            indices_up.append(("up", str(basis.dofs)))
+            indices_down.append(("down", str(basis.dofs)))
+        output_indices = indices_up + indices_down
+        args.append(output_indices)
+        res = oe.contract(*asxp_oe_args(args))
+        # to be consistent with the behavior of MPS/MPO
+        res = asnumpy(res)
+        dim = round(np.sqrt(np.prod(res.shape)))
+        return res.reshape(dim, dim)
+
+    def to_contract_args(self, prefix_up, prefix_down) -> List:
+        """
+        Get the arguments for contraction based on ``opt_einsum`` of the whole TTNO.
+
+        Parameters
+        ----------
+        prefix_up
+            the prefix for the indices of the up physical indices
+        prefix_down
+            the prefix for the indices of the down physical indices
+        Returns
+        -------
+        The arguments for contraction
+        """
+        args = []
+        for node in self.node_list:
+            assert isinstance(node, TreeNodeTensor)
+            indices = self.get_node_indices(node, prefix_up, prefix_down)
+            indices = [indices[i] for i, s in enumerate(node.tensor.shape) if s != 1]
+            tensor = node.tensor.squeeze()
+            assert len(indices) == tensor.ndim
+            args.extend([tensor, indices])
+        return args
+
+    def get_node_indices(self, node: TreeNodeTensor, prefix_up="up", prefix_down="down") -> List:
+        """
+        Get the indices of the node for contraction based on ``opt_einsum``.
+
+        Parameters
+        ----------
+        node: TreeNodeTensor
+            The node to get the indices for contraction
+        prefix_up
+            the prefix for the indices of the up physical indices
+        prefix_down
+            the prefix for the indices of the down physical indices
+
+        Returns
+        -------
+        The indices for contraction
+        """
+        _id = str(id(self))
+        all_dofs = self.tn2dofs[node]
+        indices = []
+        for child in node.children:
+            indices.append((_id, str(all_dofs), str(self.tn2dofs[child])))
+        for dofs in all_dofs:
+            # interleaved up and down
+            indices.append((prefix_up, str(dofs)))
+            indices.append((prefix_down, str(dofs)))
+        if node.parent is None:
+            indices.append((_id, "root", str(all_dofs)))
+        else:
+            indices.append((_id, str(self.tn2dofs[node.parent]), str(all_dofs)))
+        assert len(indices) == node.tensor.ndim
+        return indices
+
+    def __matmul__(self, other):
+        # duplicate with Mpo
+        return self.apply(other)
+
+
+# values are set in time_evolution.py
+EVOLVE_METHODS = {}
+
+
+class TTNS(TTNBase):
+    @classmethod
+    def load(cls, basis: BasisTree, fname: str, other_attrs=None):
+        if other_attrs is None:
+            other_attrs = []
+        other_attrs = other_attrs + ["coeff"]
+        return super().load(basis, fname, other_attrs)
+
+    @classmethod
+    def random(cls, basis: BasisTree, qntot, m_max, percent=1.0):
+        """
+        Construct a random TTNS.
+
+        Parameters
+        ----------
+        basis: BasisTree
+            the basis for the TTNS
+        qntot: int or np.ndarray
+            the total quantum number
+        m_max: int
+            the maximum bond dimension
+        percent: float
+            the percentage of the states to be randomly chosen
+
+        Returns
+        -------
+        The random TTNS
+        """
+        ttns = cls(basis)
+        if isinstance(qntot, int):
+            qntot = np.array([qntot])
+        qn_size = len(qntot)
+        assert basis.qn_size == qn_size
+
+        for node in ttns.postorder_list()[:-1]:
+            qnbigl, _, _ = ttns.get_qnmat(node, include_parent=False)
+            qnbigl_shape = qnbigl.shape
+            qnbigl = qnbigl.reshape(-1, qn_size)
+            u_list = []
+            s_list = []
+            qn_list = []
+
+            for iblock in set([tuple(t) for t in qnbigl]):
+                if np.all(np.array(qntot) < np.array(iblock)):
+                    continue
+                # find the quantum number index
+                indices = [i for i, x in enumerate(qnbigl) if tuple(x) == iblock]
+                assert len(indices) != 0
+                a = np.random.random([len(indices), len(indices)]) - 0.5
+                a = a + a.T
+                s, u = scipy.linalg.eigh(a=a)
+                u_list.append(blockrecover(indices, u, len(qnbigl)))
+                s_list.append(s)
+                qn_list += [iblock] * len(indices)
+
+            u = np.concatenate(u_list, axis=1)
+            s = np.concatenate(s_list)
+            mt, mpsdim, mpsqn, nouse = select_basis(u, s, qn_list, u, m_max, percent=percent)
+            node.tensor = mt.reshape(list(qnbigl_shape)[:-1] + [mpsdim])
+            node.qn = mpsqn
+        # deal with root
+        ttns.root.qn = np.ones((1, qn_size), dtype=int) * qntot
+        qn_mask = ttns.get_qnmask(ttns.root, include_parent=False)
+        ttns.root.tensor = np.random.random(qn_mask.shape) - 0.5
+        ttns.root.tensor[~qn_mask] = 0
+        ttns.root.tensor /= np.linalg.norm(ttns.root.tensor.ravel())
+        ttns.check_shape()
+        ttns.check_canonical()
+        return ttns
+
+    @classmethod
+    def from_tensors(cls, template: "TTNS", tensors: xp.ndarray):
+        """
+        Construct a TTNS from a list of tensors.
+        The list of tensors is concatenated to one-dimension.
+        Note that QN is taken into account
+
+        Parameters
+        ----------
+        template: TTNS
+            TTNS template providing the topology, shape and QN of the new TTNS
+        tensors: xp.ndarray
+            list of tensors concatenated as one dimensional array
+        Returns
+        -------
+        The new TTNS
+        """
+        ttns = template.metacopy()
+        cursor = 0
+        for node, tnode in zip(ttns.node_list, template.node_list):
+            qnmask = template.get_qnmask(tnode)
+            length = np.sum(qnmask)
+            node.tensor = np.zeros(tnode.shape, dtype=tensors.dtype)
+            node.tensor[qnmask] = asnumpy(tensors[cursor : cursor + length])
+            node.qn = tnode.qn
+            cursor += length
+        assert len(tensors) == cursor
+        ttns.check_shape()
+        return ttns
+
+    def __init__(self, basis: BasisTree, condition: Dict = None, root: TreeNodeTensor = None):
+        """
+        Construct a TTNS.
+
+        Parameters
+        ----------
+        basis: BasisTree
+            The basis for the TTNS
+        condition: Dict
+            The condition for the TTNS. If provided, construct a direct product TTNS with bond dimension 1
+            and the condition specifies the tensors of the TTNS.
+            See :meth:`renormalizer.Mps.hartree_product_state` for more details
+        root: TreeNodeTensor
+            The root node of the TTNS. Only one of ``condition`` or ``root`` should be provided.
+
+        Returns
+        -------
+        The new TTNS
+        """
+        self.basis = basis
+        if not root:
+            if condition is None:
+                condition = {}
+            basis_list = basis.basis_list_postorder
+            mps = Mps.hartree_product_state(Model(basis_list, []), condition, len(basis_list))
+            # can't directly use MPS qn because the topology is different
+            site_qn = [mps.qn[i + 1] - mps.qn[i] for i in range(len(mps))]
+            node_list_state = []
+
+            for node_basis in basis.node_list:
+                mps_indices = [basis_list.index(b) for b in node_basis.basis_sets]
+                assert mps_indices
+                tensor = np.eye(1)
+                # here only the site qn (rather than bond qn) is set
+                qn = 0
+                for i in mps_indices:
+                    tensor = np.tensordot(tensor, mps[i].array, axes=1)
+                    qn += site_qn[i]
+                tensor = tensor.reshape([1] * len(node_basis.children) + list(tensor.shape)[1:-1] + [1])
+                node_list_state.append(TreeNodeTensor(tensor, qn))
+
+            root: TreeNodeTensor = copy_connection(basis.node_list, node_list_state)
+
+            super().__init__(basis, root)
+
+            # summing up the site qn
+            for node in self.postorder_list():
+                for child in node.children:
+                    node.qn += child.qn
+        else:
+            assert condition is None
+            super().__init__(basis, root)
+
+        self.coeff = 1
+        self.check_shape()
+        # tensor node to basis node. make a property?
+        self.tn2bn: Dict[TreeNodeTensor, TreeNodeBasis] = {
+            tn: bn for tn, bn in zip(self.node_list, self.basis.node_list)
+        }
+        self.tn2dofs = {tn: bn.dofs for tn, bn in self.tn2bn.items()}
+
+        self.compress_config = CompressConfig()
+        self.optimize_config = OptimizeConfig()
+        self.evolve_config = EvolveConfig(EvolveMethod.tdvp_vmf, force_ovlp=False)
+
+    def check_shape(self):
+        """
+        Assert the shape of the TTNS is consistent with the basis.
+        """
+        for snode, bnode in zip(self.node_list, self.basis.node_list):
+            assert snode.tensor.ndim == len(snode.children) + bnode.n_sets + 1
+            assert snode.qn.shape[0] == snode.tensor.shape[-1]
+            assert snode.qn.shape[1] == bnode.qn_size
+            for i, b in enumerate(bnode.basis_sets):
+                assert snode.shape[len(snode.children) + i] == b.nbas
+
+    def check_canonical(self, atol=None) -> bool:
+        for node in self.node_list[1:]:
+            node.check_canonical(atol)
+        return True
+
+    def is_canonical(self, atol=None) -> bool:
+        for node in self.node_list[1:]:
+            if not node.check_canonical(atol, assertion=False):
+                return False
+        return True
+
+    def canonicalise(self):
+        for node in self.postorder_list()[:-1]:
+            self.push_cano_to_parent(node)
+        return self
+
+    def compress(self, temp_m_trunc=None, ret_s=False):
+        if self.compress_config.bonddim_should_set:
+            self.compress_config.set_bonddim(len(self.node_list) + 1)
+        s_dict: Dict[TreeNodeTensor, np.ndarray] = {self.root: np.array([1])}
+        compress_recursion(self.root, self, s_dict, temp_m_trunc)
+        self.check_shape()
+        self.check_canonical()
+        if not ret_s:
+            return self
+        else:
+            s_list = [s_dict[n] for n in self.node_list]
+            return self, s_list
+
+    def expectation1(self, ttno: TTNO, bra: "TTNS" = None):
+        # old implementation. When the network is large, opt_einsum sometimes takes a long time to
+        # find the contraction path and the resulting path is not optimal
+        if bra is None:
+            bra = self
+        args = self.to_contract_args()
+        args.extend(bra.to_contract_args(conj=True))
+        args.extend(ttno.to_contract_args("up", "down"))
+        val = oe.contract(*asxp_oe_args(args), optimize="greedy").ravel()[0]
+
+        if np.isclose(float(val.imag), 0):
+            return float(val.real)
+        else:
+            return complex(val)
+
+    def expectation(self, ttno: TTNO, bra: "TTNS" = None) -> Number:
+        """
+        Calculate the expectation value of <bra|TTNO|self>.
+
+        Parameters
+        ----------
+        ttno: TTNO
+            The operator for expectation
+        bra: TTNS
+            The bra state in TTNS format. Defaults to None and the bra is the same as self.
+
+        Returns
+        -------
+        The expectation value
+        """
+        assert bra is None  # not implemented yet
+        basis_node = TreeNodeBasis([BasisDummy("expectation dummy")])
+        basis_node_ttns = basis_node
+        basis_node_ttno = basis_node.copy()
+        basis_node_ttns.add_child(self.basis.root.copy())
+        basis_node_ttno.add_child(ttno.basis.root.copy())
+        basis_tree_ttns = BasisTree(basis_node_ttns)
+        basis_tree_ttno = BasisTree(basis_node_ttno)
+        snode = TreeNodeTensor(np.ones((1, 1, 1)), qn=np.zeros((1, basis_tree_ttns.qn_size)))
+        snode.add_child(self.root)
+        onode = TreeNodeTensor(np.ones((1, 1, 1, 1)), qn=np.zeros((1, basis_tree_ttno.qn_size)))
+        onode.add_child(ttno.root)
+
+        ttns_extended = TTNS(basis_tree_ttns, root=snode)
+        ttno_extended = TTNO(basis_tree_ttno, [], root=onode)
+        environ = TTNEnviron(ttns_extended, ttno_extended, build_environ=False)
+        environ.build_children_environ(ttns_extended, ttno_extended)
+        val = environ.root.environ_children[0].ravel()[0]
+
+        for node in [self.basis.root, self.root, ttno.root]:
+            node.parent = None
+
+        if np.isclose(float(val.imag), 0):
+            return float(val.real)
+        else:
+            return complex(val)
+
+    def add(self, other: "TTNS") -> "TTNS":
+        """
+        Add two TTNSs.
+
+        Parameters
+        ----------
+        other: TTNS
+            The other TTNS to be added
+
+        Returns
+        -------
+        The new TTNS.
+        """
+        new = self.metacopy()
+        for new_node, node1, node2 in zip(new, self, other):
+            new_shape = []
+            indices1 = []
+            indices2 = []
+            for i, (dim1, dim2) in enumerate(zip(node1.shape, node2.shape)):
+                is_physical_idx = len(node1.children) <= i and i != node1.tensor.ndim - 1
+                is_parent_idx = i == node1.tensor.ndim - 1
+                if is_physical_idx or (is_parent_idx and node1 is self.root):
+                    assert dim1 == dim2
+                    new_shape.append(dim1)
+                    indices1.append(slice(0, dim1))
+                    indices2.append(slice(0, dim1))
+                else:
+                    # virtual indices
+                    new_shape.append(dim1 + dim2)
+                    indices1.append(slice(0, dim1))
+                    indices2.append(slice(dim1, dim1 + dim2))
+            dtype = np.promote_types(node1.tensor.dtype, node2.tensor.dtype)
+            new_node.tensor = np.zeros(new_shape, dtype=dtype)
+            indices1 = tuple(indices1)
+            indices2 = tuple(indices2)
+            new_node.tensor[indices1] = node1.tensor
+            new_node.tensor[indices2] = node2.tensor
+            if node1 is self.root:
+                np.testing.assert_allclose(node1.qn, node2.qn)
+                new_node.qn = node1.qn.copy()
+            else:
+                new_node.qn = np.concatenate([node1.qn, node2.qn], axis=0)
+        new.check_shape()
+        # assert new.check_canonical()
+        return new
+
+    def normalize(self, kind):
+        r"""normalize the wavefunction
+
+        Parameters
+        ----------
+        kind: str
+            "ttns_only": the ttns part is normalized and coeff is not modified;
+            "ttns_norm_to_coeff": the ttns part is normalized and the norm is multiplied to coeff;
+            "ttns_and_coeff": both ttns and coeff is normalized
+
+        Returns
+        -------
+        ``self`` is overwritten.
+        """
+
+        return normalize(self, kind)
+
+    def evolve(self, ttno: TTNO, tau: Union[complex, float], normalize: bool = True):
+        imag_time = np.iscomplex(tau)
+        # trick to avoid complex algebra
+        # exp{coeff * H * tau}
+        # coef and tau are different from MPS implementation
+        if imag_time:
+            coeff = 1
+            tau = tau.imag
+            ttns = self
+        else:
+            coeff = -1j
+            ttns = self.to_complex()
+        method = EVOLVE_METHODS[self.evolve_config.method]
+        new_ttns = method(ttns, ttno, coeff, tau)
+        if normalize:
+            if imag_time:
+                new_ttns.normalize("mps_and_coeff")
+            else:
+                new_ttns.normalize("mps_only")
+        return new_ttns
+
+    def metacopy(self):
+        # node tensor and qn not set
+        new = self.__class__(self.basis)
+        new.coeff = self.coeff
+        new.optimize_config = self.optimize_config.copy()
+        new.evolve_config = self.evolve_config.copy()
+        new.compress_config = self.compress_config.copy()
+        return new
+
+    def copy(self):
+        new = self.metacopy()
+        for node1, node2 in zip(new, self):
+            node1.tensor = node2.tensor.copy()
+            node1.qn = node2.qn.copy()
+        return new
+
+    def to_complex(self, inplace: bool = False) -> "TTNS":
+        """
+        Convert the data type from real to complex
+        .
+        Parameters
+        ----------
+        inplace: bool
+            whether convert in place
+        Returns
+        -------
+        The new TTNS
+        """
+        if inplace:
+            new = self
+        else:
+            new = self.metacopy()
+        for node1, node2 in zip(self, new):
+            node2.tensor = np.array(node1.tensor, dtype=complex)
+            node2.qn = node1.qn.copy()
+        return new
+
+    def todense(self, order: List[BasisSet] = None) -> np.ndarray:
+        """
+        Convert the TTNS to dense vector
+
+        Parameters
+        ----------
+        order: list
+            the order of the basis sets for the resulting tensor
+
+        Returns
+        -------
+        The dense vector
+        """
+        _id = str(id(self))
+        args = self.to_contract_args()
+        if order is None:
+            order = self.basis.basis_list
+        indices_up = []
+        for basis in order:
+            indices_up.append(("down", str(basis.dofs)))
+        output_indices = indices_up
+        args.append(output_indices)
+        res = oe.contract(*asxp_oe_args(args))
+        # to be consistent with the behavior of MPS/MPO
+        res = asnumpy(res)
+        return res
+
+    def to_contract_args(self, conj: bool = False):
+        """
+        Get the arguments for contraction based on ``opt_einsum`` of the whole TTNO.
+
+        Parameters
+        ----------
+        conj: bool
+            deprecated.
+        Returns
+        -------
+        The arguments for contraction
+        """
+        args = []
+        for node in self.node_list:
+            assert isinstance(node, TreeNodeTensor)
+            indices = self.get_node_indices(node, conj)
+            tensor = node.tensor
+            if conj:
+                tensor = tensor.conj()
+            indices = [indices[i] for i, s in enumerate(tensor.shape) if s != 1]
+            tensor = tensor.squeeze()
+            assert len(indices) == tensor.ndim
+            args.extend([tensor, indices])
+        return args
+
+    def decompose_to_parent(self, node: TreeNodeTensor) -> np.ndarray:
+        """
+        QR decompose the node toward the parent direction.
+        Set the node to the orthogonal matrix Q
+        and return the triangular matrix R.
+
+        Parameters
+        ----------
+        node: TreeNodeTensor
+            The node to be decomposed
+        Returns
+        -------
+        The triangular matrix R
+        """
+        assert node.parent
+        qnbigl, qnbigr, _ = self.get_qnmat(node, include_parent=False)
+        tensor = node.tensor.reshape(-1, node.shape[-1])
+        u, qnlnew, v, qnrnew = svd_qn(tensor, qnbigl, qnbigr, self.qntot, QR=True, system="L", full_matrices=False)
+        # could shrink during QR
+        node.tensor = u.reshape(list(node.shape[:-1]) + [u.shape[1]])
+        node.qn = np.array(qnlnew)
+        return v
+
+    def merge_to_parent(self, node: TreeNodeTensor, v: np.ndarray):
+        """
+        Merge the coefficient tensor to the parent of the node.
+
+        Parameters
+        ----------
+        node: TreeNodeTensor
+            The node whose parent will be merged with ``v``
+        v: np.ndarray
+            The coefficient tensor to be merged
+        """
+        # contract parent
+        parent_indices = self.get_node_indices(node.parent)
+        args = [node.parent.tensor, parent_indices]
+        child_idx1 = parent_indices[node.idx_as_child]  # old parent->child index
+        child_idx2 = tuple(list(child_idx1) + ["_idx2"])  # new parent->child index
+        args.extend([v, (child_idx1, child_idx2)])
+        output_indices = parent_indices.copy()
+        output_indices[node.idx_as_child] = child_idx2
+        args.append(output_indices)
+        node.parent.tensor = oe.contract(*asxp_oe_args(args))
+
+    def push_cano_to_parent(self, node: TreeNodeTensor):
+        """
+        Push the canonical center to the parent of the node.
+
+        Parameters
+        ----------
+        node: TreeNodeTensor
+            The current canonical center node
+        """
+        assert node.parent
+        # move the cano center to parent
+        v = self.decompose_to_parent(node)
+        self.merge_to_parent(node, v)
+
+    def decompose_to_child(self, node: TreeNodeTensor, ichild: int) -> np.ndarray:
+        """
+        QR decompose the node towards the children direction.
+        Set the node to the orthogonal matrix Q
+        and return the triangular matrix R.
+
+        Parameters
+        ----------
+        node: TreeNodeTensor
+            The node to be decomposed
+        ichild: int
+            The index for the child
+        Returns
+        -------
+        The triangular matrix R
+        """
+        qnbigl, qnbigr, tensor, shape = moveaxis(self, node, ichild)
+
+        # u for node and v for child
+        u, qnl, v, qnr = svd_qn(tensor, qnbigl, qnbigr, self.qntot, QR=True, system="L", full_matrices=False)
+        # XXX: temp debug code for consistency with MPS
+        """
+        # u for node and v for child
+        u, qnl, v, qnr = svd_qn(
+            tensor.T,
+            qnbigr,
+            qnbigl,
+            self.qntot,
+            QR=True,
+            system="R",
+            full_matrices=False
+        )
+        v, qnr, u, qnl = u, qnl, v, qnr
+        """
+        shape[-1] = u.shape[-1]
+        node.tensor = np.moveaxis(u.reshape(shape), -1, ichild)
+        node.children[ichild].qn = qnr
+        return v
+
+    def merge_to_child(self, node: TreeNodeTensor, ichild: int, v: np.ndarray):
+        """
+        Merge the coefficient tensor to one of the children of the node.
+
+        Parameters
+        ----------
+        node: TreeNodeTensor
+            The node whose child will be merged with ``v``
+        ichild: int
+            The index for the child
+        v: np.ndarray
+            The coefficient tensor to be merged
+        """
+        child = node.children[ichild]
+        child.tensor = tensordot(child.tensor, v, axes=[-1, 0])
+
+    def push_cano_to_child(self, node: TreeNodeTensor, ichild: int):
+        """
+        Push the canonical center to the child of the node.
+
+        Parameters
+        ----------
+        node: TreeNodeTensor
+            The current canonical center node
+        ichild: int
+            The index for the child
+        """
+        v = self.decompose_to_child(node, ichild)
+        self.merge_to_child(node, ichild, v)
+
+    def compress_node(
+        self, node: TreeNodeTensor, ichild: int, temp_m_trunc: int = None, cano_child: bool = True
+    ) -> np.ndarray:
+        """
+        Compress the bond between node and one of its child based on SVD
+
+        Parameters
+        ----------
+        node: TreeNodeTensor
+            The node to be compressed
+        ichild: int
+            The index for the child
+        temp_m_trunc: int
+            The truncation of the singular values
+        cano_child: bool
+            Whether the canonical center should be set to the child node
+
+        Returns
+        -------
+        The singular values before truncation
+        """
+        qnbigl, qnbigr, tensor, shape = moveaxis(self, node, ichild)
+        # u for node and v for child
+        u, s, qnl, v, s, qnr = svd_qn(
+            tensor,
+            qnbigl,
+            qnbigr,
+            self.qntot,
+            full_matrices=False,
+        )
+        if temp_m_trunc is None:
+            idx = self.node_idx[node.children[ichild]]
+            m_trunc = self.compress_config.compute_m_trunc(s, idx, left=False)
+        else:
+            m_trunc = min(temp_m_trunc, len(s))
+        orig_s = s.copy()
+        u, s, v, qnl, qnr = truncate_tensors(u, s, v, qnl, qnr, m_trunc)
+
+        if cano_child:
+            v *= s.reshape(1, -1)
+        else:
+            u *= s.reshape(1, -1)
+
+        shape[-1] = min(m_trunc, u.shape[-1])
+        node.tensor = np.moveaxis(u.reshape(shape), -1, ichild)
+        child = node.children[ichild]
+        child.tensor = tensordot(child.tensor, v, axes=[-1, 0])
+        child.qn = qnr
+        return orig_s
+
+    def get_qnmat(self, node: TreeNodeTensor, include_parent: bool = False):
+        qnbigl = np.zeros(self.basis.qn_size, dtype=int)
+        for child in node.children:
+            qnbigl = add_outer(qnbigl, child.qn)
+        for b in self.tn2bn[node].basis_sets:
+            qnbigl = add_outer(qnbigl, b.sigmaqn)
+        if not include_parent:
+            qnbigr = self.qntot - node.qn
+            # single site
+            qnmat = add_outer(qnbigl, qnbigr)
+            return qnbigl, qnbigr, qnmat
+        # two site
+        qnbigr = np.zeros(self.basis.qn_size, dtype=int)
+        assert node.parent is not None
+        for child in node.parent.children:
+            if child is node:
+                continue
+            qnbigr = add_outer(qnbigr, child.qn)
+        for b in self.tn2bn[node.parent].basis_sets:
+            qnbigr = add_outer(qnbigr, b.sigmaqn)
+        qnbigr = add_outer(qnbigr, self.qntot - node.parent.qn)
+        qnmat = add_outer(qnbigl, qnbigr)
+        return qnbigl, qnbigr, qnmat
+
+    def get_qnmask(self, node, include_parent=False):
+        qnmat = self.get_qnmat(node, include_parent)[-1]
+        return get_qn_mask(qnmat, self.qntot)
+
+    def update_2site(self, node, tensor, m: int, percent: float = 0, cano_parent: bool = True):
+        """cano_parent: set canonical center at parent. to_right = True"""
+        parent = node.parent
+        assert parent is not None
+        qnbigl, qnbigr, _ = self.get_qnmat(node, include_parent=True)
+        dim1 = np.prod(qnbigl.shape)
+        tensor = asnumpy(tensor.reshape(dim1, -1))
+        # u for snode and v for parent
+        # duplicate with MatrixProduct._udpate_mps. Should consider merging when doing e.g. state averaged algorithm.
+        u, su, qnlnew, v, sv, qnrnew = svd_qn(tensor, qnbigl, qnbigr, self.qntot)
+        if cano_parent:
+            m_node, msdim, msqn, m_parent = select_basis(u, su, qnlnew, v, m, percent=percent)
+        else:
+            m_parent, msdim, msqn, m_node = select_basis(v, sv, qnrnew, u, m, percent=percent)
+        m_parent = m_parent.T
+        node.tensor = m_node.reshape(list(node.shape[:-1]) + [-1])
+        if cano_parent:
+            node.qn = msqn
+        else:
+            node.qn = self.qntot - msqn
+        assert len(node.qn) == node.shape[-1]
+        shape = list(parent.tensor.shape)
+        ichild = parent.children.index(node)
+        del shape[ichild]
+        shape = [-1] + shape
+        parent.tensor = np.moveaxis(m_parent.reshape(shape), 0, ichild)
+
+    def get_node_indices(
+        self, node: TreeNodeTensor, conj: bool = False, include_parent: bool = False, ttno: TTNO = None
+    ) -> List[Tuple]:
+        if include_parent:
+            snode_indices = self.get_node_indices(node, conj, ttno=ttno)
+            parent_indices = self.get_node_indices(node.parent, conj, ttno=ttno)
+            indices = snode_indices + parent_indices
+            shared_bond = snode_indices[-1]
+            for i in range(2):
+                indices.remove(shared_bond)
+            return indices
+
+        # here up and down are defined based on TTNO convention
+        if not conj:
+            _id = str(id(self))
+        else:
+            _id = str(id(self)) + "_conj"
+        skip_pidx = get_skip_pidx(node, self, ttno)
+
+        all_dofs = self.tn2dofs[node]
+        indices = []
+        for child in node.children:
+            indices.append((_id, str(all_dofs), str(self.tn2dofs[child])))
+        for i, dofs in enumerate(all_dofs):
+            if not conj and i not in skip_pidx:
+                ud = "down"
+            else:
+                ud = "up"
+            indices.append((ud, str(dofs)))
+        if node.parent is None:
+            indices.append((_id, "root", str(all_dofs)))
+        else:
+            indices.append((_id, str(self.tn2dofs[node.parent]), str(all_dofs)))
+        assert len(indices) == node.tensor.ndim
+        return indices
+
+    def merge_with_parent(self, node):
+        # merge a node with its parent
+        args = []
+        snode_indices = self.get_node_indices(node)
+        parent_indices = self.get_node_indices(node.parent)
+        args.extend([node.tensor, snode_indices])
+        args.extend([node.parent.tensor, parent_indices])
+        output_indices = self.get_node_indices(node, include_parent=True)
+        args.append(output_indices)
+        return oe.contract(*asxp_oe_args(args))
+
+    @property
+    def norm(self):
+        """duplicate with Mps"""
+        return np.linalg.norm(self.coeff) * self.ttns_norm
+
+    @property
+    def ttns_norm(self):
+        res = self.expectation(TTNO.identity(self.basis)).real
+
+        if res < 0:
+            assert np.abs(res) < 1e-8
+            res = 0
+        res = np.sqrt(res)
+        return float(res)
+
+    def scale(self, val, inplace=False):
+        # no need to be canonical
+        # self.check_canonical()
+        if inplace:
+            new_mp = self
+        else:
+            new_mp = self.copy()
+        if np.iscomplex(val):
+            new_mp.to_complex(inplace=True)
+        else:
+            val = val.real
+
+        new_mp.root.tensor *= val
+        return new_mp
+
+    def dump(self, fname, other_attrs=None):
+        if other_attrs is None:
+            other_attrs = []
+        other_attrs = other_attrs + ["coeff"]
+        super().dump(fname, other_attrs)
+
+    def __add__(self, other: "TTNS"):
+        return self.add(other)
+
+
+class TTNEnviron(Tree):
+    def __init__(self, ttns: TTNS, ttno: TTNO, build_environ=True):
+        self.basis_ttns = ttns.basis
+        self.basis_ttno = ttno.basis
+        enodes: List[TreeNodeEnviron] = [TreeNodeEnviron() for _ in range(ttns.size)]
+        copy_connection(ttns.node_list, enodes)
+        super().__init__(enodes[0])
+        assert self.root.parent is None
+        self.root.environ_parent = np.array([1], dtype=backend.real_dtype).reshape([1, 1, 1])
+        # tensor node to basis node. todo: remove duplication?
+        self.tn2dofs_ttns = {tn: bn.dofs for tn, bn in zip(self.node_list, self.basis_ttns.node_list)}
+        self.tn2dofs_ttno = {tn: bn.dofs for tn, bn in zip(self.node_list, self.basis_ttno.node_list)}
+        if build_environ:
+            self.build_children_environ(ttns, ttno)
+            self.build_parent_environ(ttns, ttno)
+
+    def build_children_environ(self, ttns, ttno):
+        # first run, children environment to the parent.
+        # set enode.environ_children
+        snodes: List[TreeNodeTensor] = ttns.postorder_list()
+        for snode in snodes:
+            self.build_children_environ_node(snode, ttns, ttno)
+
+    def build_parent_environ(self, ttns, ttno):
+        # second run, parent environment to children
+        # set enode.environ_parent
+        snodes: List[TreeNodeTensor] = ttns.node_list
+        for snode in snodes:
+            for ichild in range(len(snode.children)):
+                self.build_parent_environ_node(snode, ichild, ttns, ttno)
+
+    def update_1bond(self, snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO):
+        # update environ for the bond between snode and snode.parent
+        self.build_children_environ_node(snode, ttns, ttno)
+        self.build_parent_environ_node(snode.parent, snode.idx_as_child, ttns, ttno)
+
+    def update_1site(self, snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO):
+        # update environ based on snode
+        self.build_children_environ_node(snode, ttns, ttno)
+        for ichild in range(len(snode.children)):
+            self.build_parent_environ_node(snode, ichild, ttns, ttno)
+
+    def update_2site(self, snode, ttns, ttno):
+        # update environ based on snode and its parent
+        self.build_children_environ_node(snode, ttns, ttno)
+        self.build_children_environ_node(snode.parent, ttns, ttno)
+        for ichild in range(len(snode.parent.children)):
+            self.build_parent_environ_node(snode.parent, ichild, ttns, ttno)
+        for ichild in range(len(snode.children)):
+            self.build_parent_environ_node(snode, ichild, ttns, ttno)
+
+    def build_children_environ_node(self, snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO):
+        # build the environment from snode to its parent and store the environment in its parent
+        if snode.parent is None:
+            return
+        enode = self.node_list[ttns.node_idx[snode]]
+        onode = ttno.node_list[ttns.node_idx[snode]]
+        args = []
+        for i, child_tensor in enumerate(enode.environ_children):
+            indices = self.get_child_indices(enode, i, ttns, ttno)
+            args.extend([child_tensor, indices])
+
+        args.append(snode.tensor.conj())
+        args.append(ttns.get_node_indices(snode, conj=True))
+
+        args.append(onode.tensor)
+        args.append(ttno.get_node_indices(onode))
+
+        args.append(snode.tensor)
+        args.append(ttns.get_node_indices(snode, ttno=ttno))
+
+        # indices for the resulting tensor
+        indices = self.get_parent_indices(enode, ttns, ttno)
+        args.append(indices)
+        res = oe.contract(*asxp_oe_args(args))
+        if len(enode.parent.environ_children) != len(enode.parent.children):
+            # first run
+            enode.parent.environ_children.append(asnumpy(res))
+        else:
+            # updating
+            ichild = snode.parent.children.index(snode)
+            enode.parent.environ_children[ichild] = asnumpy(res)
+
+    def build_parent_environ_node(self, snode: TreeNodeTensor, ichild: int, ttns: TTNS, ttno: TTNO):
+        # build the environment from snode to the ith child of snode and store the environment in the child
+        enode = self.node_list[ttns.node_idx[snode]]
+        onode = ttno.node_list[ttns.node_idx[snode]]
+        args = []
+        # children tensor
+        for j, child_tensor in enumerate(enode.environ_children):
+            if j == ichild:
+                continue
+            indices = self.get_child_indices(enode, j, ttns, ttno)
+            args.extend([child_tensor, indices])
+
+        # parent tensor
+        indices = self.get_parent_indices(enode, ttns, ttno)
+        args.extend([enode.environ_parent, indices])
+
+        args.append(snode.tensor.conj())
+        args.append(ttns.get_node_indices(snode, conj=True))
+
+        args.append(onode.tensor)
+        args.append(ttno.get_node_indices(onode))
+
+        args.append(snode.tensor)
+        args.append(ttns.get_node_indices(snode, ttno=ttno))
+
+        # indices for the resulting tensor
+        indices = self.get_child_indices(enode, ichild, ttns, ttno)
+
+        args.append(indices)
+        res = oe.contract(*asxp_oe_args(args))
+        enode.children[ichild].environ_parent = asnumpy(res)
+
+    def get_child_indices(self, enode, i, ttns, ttno):
+        dofs_ttns = self.tn2dofs_ttns[enode]
+        dofs_child_ttns = self.tn2dofs_ttns[enode.children[i]]
+        dofs_ttno = self.tn2dofs_ttno[enode]
+        dofs_child_ttno = self.tn2dofs_ttno[enode.children[i]]
+        indices = [
+            (str(id(ttns)) + "_conj", str(dofs_ttns), str(dofs_child_ttns)),
+            (str(id(ttno)), str(dofs_ttno), str(dofs_child_ttno)),
+            (str(id(ttns)), str(dofs_ttns), str(dofs_child_ttns)),
+        ]
+        return indices
+
+    def get_parent_indices(self, enode, ttns, ttno):
+        dofs_ttns = self.tn2dofs_ttns[enode]
+        dofs_ttno = self.tn2dofs_ttno[enode]
+        if enode.parent is not None:
+            dofs_parent_ttns = self.tn2dofs_ttns[enode.parent]
+            dofs_parent_ttno = self.tn2dofs_ttno[enode.parent]
+        else:
+            dofs_parent_ttns = dofs_parent_ttno = "root"
+        indices = [
+            (str(id(ttns)) + "_conj", str(dofs_parent_ttns), str(dofs_ttns)),
+            (str(id(ttno)), str(dofs_parent_ttno), str(dofs_ttno)),
+            (str(id(ttns)), str(dofs_parent_ttns), str(dofs_ttns)),
+        ]
+        return indices
+
+
+def from_mps(mps: Mps) -> Tuple[BasisTree, TTNS, TTNO]:
+    # useful function for comparing results with MPS
+    mps = mps.copy()
+    mps.ensure_left_canonical()
+    mps.move_qnidx(len(mps) + 1)
+    # take reverse because in `node` the order of the indices is
+    # children + physical + parent
+    # |    |    |     |
+    # o -> o -> o -> root (canonical center)
+    basis = BasisTree.linear(mps.model.basis[::-1])
+    ttns = TTNS(basis)
+    for i in range(len(mps)):
+        node = ttns.node_list[::-1][i]
+        node.tensor = mps[i].array
+        node.qn = mps.qn[i + 1]
+        if i == 0:
+            # remove the empty children index
+            node.tensor = node.tensor[0, ...]
+    ttns.check_shape()
+    ttns.check_canonical()
+    ttno = TTNO(basis, mps.model.ham_terms)
+    return basis, ttns, ttno
+
+
+def compress_recursion(snode: TreeNodeTensor, ttns: TTNS, s_dict: Dict, temp_m_trunc: int = None):
+    assert snode.children, "can't compress a single tree node"
+    for ichild, child in enumerate(snode.children):
+        cano_child = bool(child.children)
+        # compress snode - child
+        s = ttns.compress_node(snode, ichild, temp_m_trunc, cano_child)
+        s_dict[child] = s
+
+        if cano_child:
+            compress_recursion(child, ttns, s_dict, temp_m_trunc)
+            # cano to snode
+            ttns.push_cano_to_parent(child)
+
+
+def truncate_tensors(u, s, v, qnl, qnr, m):
+    u = u[:, :m]
+    s = s[:m]
+    v = v[:, :m]
+    qnl = qnl[:m]
+    qnr = qnr[:m]
+    return u, s, v, qnl, qnr
+
+
+def moveaxis(ttns: TTNS, node: TreeNodeTensor, ichild: int):
+    # move one of the children indices to the end
+
+    # left indices: other children + physical bonds + parent
+    qnbigl = np.zeros(ttns.basis.qn_size, dtype=int)
+    # other children
+    for child in node.children:
+        if child == node.children[ichild]:
+            continue
+        qnbigl = add_outer(qnbigl, child.qn)
+    # physical bonds
+    for b in ttns.tn2bn[node].basis_sets:
+        qnbigl = add_outer(qnbigl, b.sigmaqn)
+    # parent
+    qnbigl = add_outer(qnbigl, ttns.qntot - node.qn)
+    # right indices: the ith child
+    qnbigr = node.children[ichild].qn
+    # 2d tensor (node, child)
+    tensor = np.moveaxis(node.tensor, ichild, -1)
+    shape = list(tensor.shape)
+    tensor = tensor.reshape(-1, node.shape[ichild])
+    return qnbigl, qnbigr, tensor, shape
+
+
+def get_skip_pidx(snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO) -> List[int]:
+    # get the indices of the physical bonds that should be skipped
+    # because of missing dofs in ttno
+    if ttno is None:
+        return []
+    idx = ttns.node_idx[snode]
+    basis_ttns: TreeNodeBasis = ttns.basis.node_list[idx]
+    basis_ttno: TreeNodeBasis = ttno.basis.node_list[idx]
+    if basis_ttns.dofs == basis_ttno.dofs:
+        return []
+    assert len(basis_ttno.dofs) < len(basis_ttns.dofs)
+    for dof in basis_ttno.dofs:
+        assert dof in basis_ttns.dofs
+    skip_pidx = []
+    for i, dof in enumerate(basis_ttns.dofs):
+        if dof not in basis_ttno.dofs:
+            skip_pidx.append(i)
+    return skip_pidx
diff --git a/renormalizer/tn/treebase.py b/renormalizer/tn/treebase.py
new file mode 100644
index 00000000..34a652fa
--- /dev/null
+++ b/renormalizer/tn/treebase.py
@@ -0,0 +1,257 @@
+from itertools import chain
+from typing import List, Sequence
+
+import numpy as np
+from print_tree import print_tree
+
+from renormalizer import Op
+from renormalizer.model.basis import BasisSet, BasisDummy
+from renormalizer.tn.node import NodeUnion, TreeNodeBasis, copy_connection
+
+
+class Tree:
+    def __init__(self, root: NodeUnion):
+        assert root.parent is None
+        self.root = root
+        self.node_list = self.preorder_list()
+        self.node_idx = {node: i for i, node in enumerate(self.node_list)}
+
+    def preorder_list(self, func=None) -> List[NodeUnion]:
+        def recursion(node: NodeUnion):
+            if func is None:
+                ret = [node]
+            else:
+                ret = [func(node)]
+            if not node.children:
+                return ret
+            for child in node.children:
+                ret += recursion(child)
+            return ret
+
+        return recursion(self.root)
+
+    def postorder_list(self) -> List[NodeUnion]:
+        def recursion(node: NodeUnion):
+            if not node.children:
+                return [node]
+            ret = []
+            for child in node.children:
+                ret += recursion(child)
+            ret.append(node)
+            return ret
+
+        return recursion(self.root)
+
+    @property
+    def size(self):
+        return len(self.node_list)
+
+    def __len__(self):
+        return self.size
+
+    def __iter__(self):
+        return iter(self.node_list)
+
+
+class BasisTree(Tree):
+    """Tree of basis sets."""
+
+    @classmethod
+    def linear(cls, basis_list: List[BasisSet]):
+        node_list = [TreeNodeBasis([basis]) for basis in basis_list]
+        for i in range(len(node_list) - 1):
+            node_list[i].add_child(node_list[i + 1])
+        return cls(node_list[0])
+
+    @classmethod
+    def binary(cls, basis_list: List[BasisSet]):
+        node_list = [TreeNodeBasis([basis]) for basis in basis_list]
+
+        def binary_recursion(node: TreeNodeBasis, offspring: List[TreeNodeBasis]):
+            if len(offspring) == 0:
+                return
+            node.add_child(offspring[0])
+            if len(offspring) == 1:
+                return
+            node.add_child(offspring[1])
+            new_offspring = offspring[2:]
+            mid_idx = len(new_offspring) // 2
+            binary_recursion(offspring[0], new_offspring[:mid_idx])
+            binary_recursion(offspring[1], new_offspring[mid_idx:])
+
+        binary_recursion(node_list[0], node_list[1:])
+        return cls(node_list[0])
+
+    @classmethod
+    def general_mctdh(
+        cls,
+        basis_list: List[BasisSet],
+        tree_order: int,
+        contract_primitive: bool = False,
+        contract_label: Sequence[bool] = None,
+        dummy_label="MCTDH virtual",
+    ):
+        # `contract_primitive` means a 2-index tensor for each primitive basis
+        #       o
+        #      / \
+        #      o  o
+        # d -> |  | <- d
+        # not contract primitive
+        #        o
+        # d ->  / \ <- d
+        # d means physical bond
+        # `contract_label` decides whether we do contraction for a particular basis
+        assert len(basis_list) > 1
+
+        # prepare elementary nodes
+        elementary_nodes: List[TreeNodeBasis] = []
+        if not contract_primitive:
+            assert contract_label is None, "providing label makes sense only when primitives are contracted"
+            while tree_order < len(basis_list):
+                node = TreeNodeBasis(basis_list[:tree_order])
+                elementary_nodes.append(node)
+                basis_list = basis_list[tree_order:]
+            elementary_nodes.append(TreeNodeBasis(basis_list))
+        else:
+            if contract_label is None:
+                for basis in basis_list:
+                    node1 = TreeNodeBasis([basis])
+                    elementary_nodes.append(node1)
+            else:
+                assert len(contract_label) == len(basis_list)
+                i = 0
+                while i != len(basis_list):
+                    if contract_label[i]:
+                        elementary_nodes.append(TreeNodeBasis([basis_list[i]]))
+                        i += 1
+                    else:
+                        for j in range(1, tree_order + 1):
+                            if i + j == len(contract_label) or contract_label[i + j]:
+                                break
+                        elementary_nodes.append(TreeNodeBasis(basis_list[i : i + j]))
+                        i += j
+
+        # recursive tree construction
+        def recursion(elementary_nodes_: List[TreeNodeBasis]) -> TreeNodeBasis:
+            nonlocal dummy_i
+            node = TreeNodeBasis([BasisDummy((dummy_label, dummy_i))])
+            dummy_i += 1
+            if len(elementary_nodes_) <= tree_order:
+                node.add_child(elementary_nodes_)
+                return node
+            for group in approximate_partition(elementary_nodes_, tree_order):
+                node.add_child(recursion(group))
+            return node
+
+        dummy_i = 0
+        root = recursion(elementary_nodes)
+        return cls(root)
+
+    @classmethod
+    def binary_mctdh(
+        cls, basis_list: List[BasisSet], contract_primitive=False, contract_label=None, dummy_label="MCTDH virtual"
+    ):
+        return cls.general_mctdh(basis_list, 2, contract_primitive, contract_label, dummy_label)
+
+    @classmethod
+    def ternary_mctdh(
+        cls, basis_list: List[BasisSet], contract_primitive=False, contract_label=None, dummy_label="MCTDH virtual"
+    ):
+        return cls.general_mctdh(basis_list, 3, contract_primitive, contract_label, dummy_label)
+
+    @classmethod
+    def t3ns(cls, basis_list: List[BasisSet], t3ns_label="T3NS virtual"):
+        def recursion(parent, basis_list_: List[BasisSet]):
+            nonlocal dummy_i
+            if len(basis_list_) == 0:
+                return
+            if len(basis_list_) == 1:
+                parent.add_child(TreeNodeBasis(basis_list_))
+                return
+            if len(basis_list_) == 2:
+                node1 = TreeNodeBasis(basis_list_[:1])
+                parent.add_child(node1)
+                node2 = TreeNodeBasis(basis_list_[1:])
+                node1.add_child(node2)
+                return
+            node1 = TreeNodeBasis(basis_list_[:1])
+            parent.add_child(node1)
+            node2 = TreeNodeBasis([BasisDummy((t3ns_label, dummy_i))])
+            dummy_i += 1
+            node1.add_child(node2)
+            for partition_ in approximate_partition(basis_list_[1:], 2):
+                recursion(node2, partition_)
+
+        dummy_i = 0
+        root = TreeNodeBasis([BasisDummy((t3ns_label, dummy_i))])
+        dummy_i += 1
+        for partition in approximate_partition(basis_list, 3):
+            recursion(root, partition)
+        return cls(root)
+
+    def __init__(self, root: TreeNodeBasis):
+        super().__init__(root)
+        for node in self.node_list:
+            assert isinstance(node, TreeNodeBasis)
+        qn_size_list = [n.qn_size for n in self.node_list]
+        if len(set(qn_size_list)) != 1:
+            raise ValueError(f"Inconsistent quantum number size: {set(qn_size_list)}")
+        self.qn_size: int = qn_size_list[0]
+        # identity operator
+        self.identity_op: Op = Op("I", self.root.dofs[0][0])
+        # identity ttno
+        self.identity_ttno = None
+
+    def print(self, print_function=None):
+        class print_tn_basis(print_tree):
+            def get_children(self, node):
+                return node.children
+
+            def get_node_str(self, node):
+                return str([b.dofs for b in node.basis_sets])
+
+        tree = print_tn_basis(self.root)
+        if print_function is not None:
+            for row in tree.rows:
+                print_function(row)
+
+    @property
+    def basis_list(self) -> List[BasisSet]:
+        return list(chain(*[n.basis_sets for n in self.node_list]))
+
+    @property
+    def basis_list_postorder(self) -> List[BasisSet]:
+        return list(chain(*[n.basis_sets for n in self.postorder_list()]))
+
+    @property
+    def pbond_dims(self) -> List[List[int]]:
+        return [b.pbond_dims for b in self.node_list]
+
+    def add_auxiliary_space(self, auxiliary_label="Q") -> "BasisTree":
+        # make a new basis tree with auxiliary basis
+        node2_list = []
+        for node in self:
+            basis_set2_list = []
+            for basis in node.basis_sets:
+                # the P space
+                basis_set2_list.append(basis)
+                if not isinstance(basis, BasisDummy):
+                    # the Q space
+                    basis_q: BasisSet = basis.copy((auxiliary_label, basis.dofs))
+                    # set to zero for know. could change to more complicated case in the future
+                    basis_q.sigmaqn = np.zeros_like(basis.sigmaqn)
+                    basis_set2_list.append(basis_q)
+            node2_list.append(TreeNodeBasis(basis_set2_list))
+        copy_connection(self.node_list, node2_list)
+        basis_tree2 = BasisTree(node2_list[0])
+        return basis_tree2
+
+
+def approximate_partition(sequence, ngroups):
+    size = (len(sequence) - 1) // ngroups + 1
+    ret = []
+    for i in range(ngroups):
+        start = i * size
+        end = min((i + 1) * size, len(sequence))
+        ret.append(sequence[start:end])
+    return ret
diff --git a/renormalizer/tn/utils_eph.py b/renormalizer/tn/utils_eph.py
new file mode 100644
index 00000000..4bd1c851
--- /dev/null
+++ b/renormalizer/tn/utils_eph.py
@@ -0,0 +1,41 @@
+from renormalizer import Op
+from renormalizer.tn import BasisTree, TTNS, TTNO, TreeNodeTensor, copy_connection
+from renormalizer.mps.backend import np
+
+
+def max_entangled_ex(basis_tree: BasisTree):
+    # basis_tree should have both P and Q space
+    node_tensor_list = []
+    for basis_node in basis_tree:
+        nbas = basis_node.basis_sets[0].nbas
+        if basis_node.n_sets == 1:
+            # dummy
+            assert nbas == 1
+            tensor = np.ones([1] * len(basis_node.children) + [1, 1])
+        elif basis_node.n_sets == 2 and basis_node.basis_sets[0].is_phonon:
+            shape = [1] * len(basis_node.children) + [nbas, nbas, 1]
+            tensor = np.eye(nbas).reshape(shape)
+            tensor /= np.sqrt(nbas)
+        elif basis_node.n_sets == 2 and basis_node.basis_sets[0].is_electron:
+            shape = [1] * len(basis_node.children) + [nbas, nbas, 1]
+            tensor = np.zeros(shape)
+            tensor[..., 0, 0, 0] = 1
+        else:
+            assert False
+        node = TreeNodeTensor(tensor, np.array([0]).reshape(1, basis_tree.qn_size))
+        node_tensor_list.append(node)
+
+    copy_connection(basis_tree.node_list, node_tensor_list)
+    ttns = TTNS(basis_tree, root=node_tensor_list[0])
+    ex_ops = []
+    for b in basis_tree.basis_list:
+        # skip Q space
+        if isinstance(b.dof, tuple) and len(b.dof) == 2 and b.dof[0] == "Q":
+            continue
+        if b.is_electron:
+            op = Op(r"a^\dagger a^\dagger", [b.dof, ("Q", b.dofs)], qn=[1, 0])
+            ex_ops.append(op)
+    ex_ttno = TTNO(basis_tree, ex_ops)
+    ttns = ex_ttno.apply(ttns)
+    ttns.normalize("ttns_and_coeff")
+    return ttns
diff --git a/renormalizer/utils/configs.py b/renormalizer/utils/configs.py
index a8b08c1d..3f0aa8df 100644
--- a/renormalizer/utils/configs.py
+++ b/renormalizer/utils/configs.py
@@ -328,6 +328,11 @@ def __init__(self, procedure=None):
         # -1.0 to get the largest eigenvalue
         self.inverse = 1.0
 
+    def copy(self):
+        new = self.__class__.__new__(self.__class__)
+        new.__dict__ = self.__dict__.copy()
+        new.procedure = self.procedure.copy()
+        return new
 
 class EvolveMethod(Enum):
     """
diff --git a/renormalizer/utils/quantity.py b/renormalizer/utils/quantity.py
index 8f94082b..fce2e110 100644
--- a/renormalizer/utils/quantity.py
+++ b/renormalizer/utils/quantity.py
@@ -4,6 +4,7 @@
 
 from __future__ import division
 
+import math
 import logging
 
 from renormalizer.utils import constant
@@ -49,6 +50,8 @@ def as_unit(self, unit):
 
     # kelvin to beta  au^-1
     def to_beta(self):
+        if self.value == 0:
+            return math.inf
         return 1.0 / self.as_au()
 
     # a simplified and incomplete model for + - * /
diff --git a/requirements.txt b/requirements.txt
index ce4ce55c..c736f3dd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,3 +9,4 @@ furo==2022.9.*
 nbsphinx==0.8.*
 opt_einsum==3.3.*
 qutip==4.6.*
+print-tree2==0.9.10
diff --git a/setup.py b/setup.py
index 79c600b1..b6e571da 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,8 @@
 req = ["numpy",
        "scipy",
        "h5py",
-       "opt_einsum"
+       "opt_einsum",
+       "sympy"
        ]
 
 setuptools.setup(


Test patch:
diff --git a/pytest.ini b/pytest.ini
index df3eb518..3fb7d1fd 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -1,2 +1,3 @@
 [pytest]
 addopts = --doctest-modules
+norecursedirs = doc example
diff --git a/renormalizer/mps/tests/test_mpo.py b/renormalizer/mps/tests/test_mpo.py
index cd73d3ae..9e6a8774 100644
--- a/renormalizer/mps/tests/test_mpo.py
+++ b/renormalizer/mps/tests/test_mpo.py
@@ -129,6 +129,7 @@ def test_scheme4():
     mps3.append(np.array([0, 1]).reshape((1,2,1)))
     mps3.append(np.array([0.707, 0.707]).reshape((1,2,1)))
     mps3.append(np.array([1, 0]).reshape((1,2,1)))
+    mps3.build_empty_qn()
     e3 = mps3.expectation(mpo3)
     assert pytest.approx(e4) == e3
 
diff --git a/renormalizer/sbm/tests/test_sbm.py b/renormalizer/sbm/tests/test_sbm.py
index 4735f466..c865e3e4 100644
--- a/renormalizer/sbm/tests/test_sbm.py
+++ b/renormalizer/sbm/tests/test_sbm.py
@@ -14,6 +14,7 @@ def test_sdf():
     omega_c = Quantity(5)
     sdf = SpectralDensityFunction(alpha, omega_c)
     omega_list, displacement_list = sdf.trapz(200, 0.0, 50)
+    omega_list, displacement_list = sdf.post_process(omega_list, displacement_list)
 
     ph_list = [Phonon.simplest_phonon(o, d) for o,d in zip(omega_list, displacement_list)]
     mol_reor = sum(ph.reorganization_energy.as_au() for ph in ph_list)
diff --git a/renormalizer/tn/tests/test_evolve.py b/renormalizer/tn/tests/test_evolve.py
new file mode 100644
index 00000000..c2fde3ec
--- /dev/null
+++ b/renormalizer/tn/tests/test_evolve.py
@@ -0,0 +1,172 @@
+import os
+from  typing import List
+
+import pytest
+
+from renormalizer import Op, Quantity
+from renormalizer.mps.backend import np
+from renormalizer.mps.mps import expand_bond_dimension_general
+from renormalizer.mps.tests.test_evolve import QUTIP_STEP, qutip_expectations, init_mps
+from renormalizer.tests.parameter_exact import model
+from renormalizer.tests import parameter
+from renormalizer.tn import BasisTree, TTNO, TTNS
+from renormalizer.tn.tree import from_mps
+from renormalizer.tn.node import TreeNodeBasis
+from renormalizer.tn.utils_eph import max_entangled_ex
+from renormalizer.utils import EvolveConfig, EvolveMethod, CompressConfig, CompressCriteria
+
+
+def add_ttno_offset(ttns: TTNS, ttno: TTNO):
+    e = ttns.expectation(ttno)
+    ham_terms = ttno.terms.copy()
+    ham_terms.append(ttns.basis.identity_op * (-e))
+    return TTNO(ttno.basis, ham_terms)
+
+
+def construct_ttns_and_ttno_chain():
+    basis, ttns, ttno = from_mps(init_mps)
+    op_n_list = [TTNO(basis, [Op(r"a^\dagger a", i)]) for i in range(3)]
+    ttno = add_ttno_offset(ttns, ttno)
+    return ttns, ttno, op_n_list
+
+
+def construct_ttns_and_ttno_tree():
+    node_list = [TreeNodeBasis([basis]) for basis in model.basis]
+    # 0 - 2 - 4
+    # |   |   |
+    # 1   3   5
+    root = node_list[2]
+    root.add_child(node_list[0])
+    root.add_child(node_list[3])
+    root.add_child(node_list[4])
+    node_list[0].add_child(node_list[1])
+    node_list[4].add_child(node_list[5])
+    basis = BasisTree(root)
+    ttno = TTNO(basis, model.ham_terms)
+    op_n_list = [TTNO(basis, [Op(r"a^\dagger a", i)]) for i in range(3)]
+    ttns = TTNS(basis, {0: 1})
+    ttno = add_ttno_offset(ttns, ttno)
+    return ttns, ttno, op_n_list
+
+
+def construct_ttns_and_ttno_tree_mctdh():
+    basis = BasisTree.binary_mctdh(model.basis)
+    op_n_list = [TTNO(basis, [Op(r"a^\dagger a", i)]) for i in range(3)]
+    ttns = TTNS(basis, {0: 1})
+    ttno = TTNO(basis, model.ham_terms)
+    ttno = add_ttno_offset(ttns, ttno)
+    return ttns, ttno, op_n_list
+
+
+init_chain = construct_ttns_and_ttno_chain()
+init_tree = construct_ttns_and_ttno_tree()
+init_tree_mctdh = construct_ttns_and_ttno_tree_mctdh()
+
+
+def check_result(ttns: TTNS, ttno: TTNO, time_step: float, final_time: float, op_n_list: List, atol: float=1e-4):
+    expectations = [[ttns.expectation(o) for o in op_n_list]]
+    for i in range(round(final_time / time_step)):
+        ttns = ttns.evolve(ttno, time_step)
+        es = [ttns.expectation(o) for o in op_n_list]
+        expectations.append(es)
+    expectations = np.array(expectations)
+    qutip_end = round(final_time / QUTIP_STEP) + 1
+    qutip_interval = round(time_step / QUTIP_STEP)
+    # more strict than mcd (the error criteria used for mps tests)
+    np.testing.assert_allclose(expectations, qutip_expectations[:qutip_end:qutip_interval], atol=atol)
+    diff = np.max(np.abs(expectations - qutip_expectations[:qutip_end:qutip_interval]), axis=0)
+    print(diff)
+    return ttns
+
+
+@pytest.mark.parametrize("ttns_and_ttno", [init_chain, init_tree, init_tree_mctdh])
+def test_tdvp_vmf(ttns_and_ttno):
+    ttns, ttno, op_n_list = ttns_and_ttno
+    # expand bond dimension
+    ttns = ttns + ttns.random(ttns.basis, 1, 5).scale(1e-5, inplace=True)
+    ttns.canonicalise()
+    ttns.evolve_config = EvolveConfig(EvolveMethod.tdvp_vmf, ivp_rtol=1e-4, ivp_atol=1e-7, force_ovlp=False)
+    check_result(ttns, ttno, 0.5, 2, op_n_list)
+
+
+@pytest.mark.parametrize("ttns_and_ttno", [init_chain, init_tree, init_tree_mctdh])
+def test_pc(ttns_and_ttno):
+    ttns, ttno, op_n_list = ttns_and_ttno
+    ttns = ttns.copy()
+    ttns.evolve_config = EvolveConfig(EvolveMethod.prop_and_compress_tdrk4)
+    ttns.compress_config = CompressConfig(CompressCriteria.fixed)
+    check_result(ttns, ttno, 0.2, 5, op_n_list, 5e-4)
+
+
+@pytest.mark.parametrize("ttns_and_ttno", [init_chain, init_tree, init_tree_mctdh])
+@pytest.mark.parametrize("method", [EvolveMethod.tdvp_ps, EvolveMethod.tdvp_ps2])
+def test_tdvp_ps(ttns_and_ttno, method):
+    ttns, ttno, op_n_list = ttns_and_ttno
+    if ttns_and_ttno is init_chain:
+        ttns = ttns.copy()
+    else:
+        # expand bond dimension
+        ttns = ttns + ttns.random(ttns.basis, 1, 5).scale(1e-5, inplace=True)
+        ttns.canonicalise()
+    ttns.evolve_config = EvolveConfig(method)
+    ttns.compress_config = CompressConfig(CompressCriteria.fixed)
+    if method is EvolveMethod.tdvp_ps:
+        check_result(ttns, ttno, 0.4, 5, op_n_list)
+    else:
+        assert method is EvolveMethod.tdvp_ps2
+        check_result(ttns, ttno, 2, 10, op_n_list, 5e-4)
+
+
+def test_thermalprop():
+    # imaginary time evolution on the P space. Q space untouched
+    holstein_model = parameter.holstein_model
+
+    basis_tree = BasisTree.binary_mctdh(holstein_model.basis, contract_primitive=True)
+    basis_tree2 = basis_tree.add_auxiliary_space()
+
+    ttns = max_entangled_ex(basis_tree2)
+    ttns.compress_config.bond_dim_max_value = 12
+    ttno = TTNO(basis_tree, holstein_model.ham_terms)
+    ttns = expand_bond_dimension_general(ttns, hint_mpo=ttno)
+    ttns.evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
+
+
+    beta = Quantity(298, "K").to_beta()
+    evolve_time = beta / 2j
+    nsteps = 50
+    dbeta = evolve_time / nsteps
+    for i in range(nsteps):
+        ttns.evolve(ttno, dbeta)
+        e = ttns.expectation(ttno)
+
+    ne_ttno_list = [TTNO(basis_tree, Op(r"a^\dagger a", b.dof)) for b in holstein_model.basis if b.is_electron]
+    occ = [ttns.expectation(ttno) for ttno in ne_ttno_list]
+
+    etot_std = 0.0853388 + parameter.holstein_model.gs_zpe
+    occ_std = [0.20896541050347484, 0.35240029674394463, 0.4386342927525734]
+    # used small M and large dt
+    # as long as the result is not absurdly wrong, it's fine
+    rtol = 5e-3
+    assert np.allclose(occ, occ_std, rtol=rtol)
+    assert np.allclose(e, etot_std, rtol=rtol)
+
+
+@pytest.mark.parametrize("ttns_and_ttno", [init_chain, init_tree, init_tree_mctdh])
+def test_save_load(ttns_and_ttno):
+    ttns, ttno, op_n_list = ttns_and_ttno
+    ttns = ttns + ttns.random(ttns.basis, 1, 5).scale(1e-5, inplace=True)
+    ttns.canonicalise()
+    tau = 0.5
+    ttns1 = ttns.copy()
+    for i in range(2):
+        ttns1 = ttns1.evolve(ttno, tau)
+    exp1 = [ttns1.expectation(o) for o in op_n_list]
+    ttns2 = ttns.evolve(ttno, tau)
+    fname = f"{id(ttns2)}.npz"
+    ttns2.dump(fname)
+    ttns2 = TTNS.load(ttns.basis, fname)
+    ttns2 = ttns2.evolve(ttno, tau)
+    assert ttns2.coeff == ttns1.coeff
+    exp2 = [ttns2.expectation(o) for o in op_n_list]
+    np.testing.assert_allclose(exp2, exp1)
+    os.remove(fname)
diff --git a/renormalizer/tn/tests/test_tn.py b/renormalizer/tn/tests/test_tn.py
new file mode 100644
index 00000000..c7f55095
--- /dev/null
+++ b/renormalizer/tn/tests/test_tn.py
@@ -0,0 +1,205 @@
+import pytest
+
+from renormalizer import BasisHalfSpin, Model, Mpo, Mps
+from renormalizer.mps.backend import np
+from renormalizer.model.model import heisenberg_ops
+from renormalizer.tn.node import TreeNodeBasis
+from renormalizer.tn.tree import TTNO, TTNS, TTNEnviron, from_mps
+from renormalizer.tn.treebase import BasisTree
+from renormalizer.tn.gs import optimize_ttns
+from renormalizer.tests.parameter import holstein_model
+from renormalizer.tests.parameter_exact import model
+
+
+def multi_basis_tree(basis_list):
+    #         3
+    #        / \
+    #       2
+    #    /     \
+    #  0,1    4, 5, 6
+    node1 = TreeNodeBasis([basis_list[0], basis_list[1]])
+    node2 = TreeNodeBasis([basis_list[2]])
+    node3 = TreeNodeBasis([basis_list[3]])
+    node4 = TreeNodeBasis([basis_list[4], basis_list[5], basis_list[6]])
+    node3.add_child(node2)
+    node2.add_child(node1)
+    node2.add_child(node4)
+    basis = BasisTree(node3)
+    return basis
+
+
+nspin = 7
+basis_list = [BasisHalfSpin(i) for i in range(nspin)]
+basis_binary = BasisTree.binary(basis_list)
+basis_multi_basis = multi_basis_tree(basis_list)
+
+
+def holstein_scheme3() -> BasisTree:
+    model = holstein_model
+    node_list = [TreeNodeBasis([basis]) for basis in model.basis]
+    root = node_list[3]
+    root.add_child(node_list[0])
+    root.add_child(node_list[6])
+    for i in range(3):
+        node_list[3 * i].add_child(node_list[3 * i + 1])
+        node_list[3 * i + 1].add_child(node_list[3 * i + 2])
+    return BasisTree(root)
+
+
+@pytest.mark.parametrize("basis", [basis_binary, basis_multi_basis])
+def test_ttno(basis):
+    ham_terms = heisenberg_ops(nspin)
+
+    ttno = TTNO(basis, ham_terms)
+    dense = ttno.todense(basis_list)
+
+    dense2 = Mpo(Model(basis_list, ham_terms)).todense()
+    np.testing.assert_allclose(dense, dense2)
+
+
+@pytest.mark.parametrize("basis", [basis_binary, basis_multi_basis])
+def test_ttns(basis):
+    ham_terms = heisenberg_ops(nspin)
+    condition = {1:1, 3:1}
+    ttns = TTNS(basis, condition)
+    ttno = TTNO(basis, ham_terms)
+    e1 = ttns.expectation(ttno)
+    model = Model([BasisHalfSpin(i) for i in range(nspin)], ham_terms)
+    mps = Mps.hartree_product_state(model, condition)
+    mpo = Mpo(model)
+    e2 = mps.expectation(mpo)
+    np.testing.assert_allclose(e1, e2)
+    env = TTNEnviron(ttns, ttno)
+    for node in env.node_list:
+        for child, environ_child in zip(node.children, node.environ_children):
+            e3 = environ_child.ravel() @ child.environ_parent.ravel()
+            np.testing.assert_allclose(e3, e2)
+
+
+@pytest.mark.parametrize("basis", [basis_binary, basis_multi_basis])
+def test_push_cano(basis):
+    ttns = TTNS.random(basis, 0, 5, 1)
+    s1 = ttns.todense()
+    ttns.push_cano_to_child(ttns.root, 0)
+    s2 = ttns.todense()
+    np.testing.assert_allclose(s2, s1)
+    ttns.push_cano_to_parent(ttns.root.children[0])
+    s3 = ttns.todense()
+    np.testing.assert_allclose(s3, s1)
+
+
+def test_from_mps():
+    mps = Mps.random(model, 1, 10)
+    mpo = Mpo(model)
+    e_ref = mps.expectation(mpo)
+    basis, ttns, ttno = from_mps(mps)
+    e = ttns.expectation(ttno)
+    np.testing.assert_allclose(e, e_ref)
+
+
+@pytest.mark.parametrize("basis_tree", [basis_binary, basis_multi_basis])
+@pytest.mark.parametrize("ite", [False, True])
+def test_gs_heisenberg(basis_tree, ite):
+    ham_terms = heisenberg_ops(4)
+    ttns = TTNS.random(basis_tree, qntot=0, m_max=20)
+    ttno = TTNO(basis_tree, ham_terms)
+    if not ite:
+        e1 = optimize_ttns(ttns, ttno)
+        e1 = min(e1)
+    else:
+        # imaginary time evolution for the ground state
+        for i in range(10):
+            ttns.check_canonical()
+            ttns = ttns.evolve(ttno, -2j)
+        e1 = ttns.expectation(ttno)
+    h = ttno.todense()
+    e2 = np.linalg.eigh(h)[0][0]
+    np.testing.assert_allclose(e1, e2)
+
+
+@pytest.mark.parametrize("scheme", [3, 4])
+def test_gs_holstein(scheme):
+    if scheme == 3:
+        model = holstein_model
+        basis = holstein_scheme3()
+    else:
+        assert scheme == 4
+        model = holstein_model.switch_scheme(4)
+        node_list = [TreeNodeBasis([basis]) for basis in model.basis]
+        root = node_list.pop(2)
+        assert len(node_list) == 6
+        for i in range(3):
+            root.add_child(node_list[2*i])
+            node_list[2*i].add_child(node_list[2*i+1])
+        basis = BasisTree(root)
+    m = 4
+    ttns = TTNS.random(basis, qntot=1, m_max=m)
+    ttno = TTNO(basis, model.ham_terms)
+    procedure = [[m, 0.4], [m, 0.2], [m, 0.1], [m, 0], [m, 0]]
+    e1 = optimize_ttns(ttns, ttno, procedure)
+    e2 = 0.08401412 + model.gs_zpe
+    np.testing.assert_allclose(min(e1), e2)
+
+
+@pytest.mark.parametrize("basis_tree", [basis_binary, basis_multi_basis])
+def test_add(basis_tree):
+    ttns1 = TTNS.random(basis_tree, qntot=0, m_max=4)
+    ttns2 = TTNS.random(basis_tree, qntot=0, m_max=2).scale(1j)
+    ttns3 = ttns1.add(ttns2)
+    s1 = ttns1.todense()
+    s2 = ttns2.todense()
+    assert np.iscomplexobj(s2)
+    s3 = ttns3.todense()
+    np.testing.assert_allclose(s1 + s2, s3)
+
+
+@pytest.mark.parametrize("basis_tree", [basis_binary, basis_multi_basis])
+def test_apply(basis_tree):
+    ttns1 = TTNS.random(basis_tree, qntot=0, m_max=4)
+    ttno = TTNO(basis_tree, heisenberg_ops(nspin))
+    ttns2 = ttno.apply(ttns1)
+    s1 = ttns1.todense()
+    s2 = ttns2.todense()
+    op = ttno.todense()
+    np.testing.assert_allclose(s2.ravel(), op @ s1.ravel())
+
+
+def test_compress():
+    m1 = 5
+    m2 = 4
+    basis = holstein_scheme3()
+    ttno = TTNO(basis, holstein_model.ham_terms)
+    ttns = TTNS.random(basis, 1, m1)
+    procedure1, procedure2 = [[[m, 0.4], [m, 0.2], [m, 0.1], [m, 0], [m, 0]] for m in [m1, m2]]
+    optimize_ttns(ttns, ttno, procedure1)
+    ttns2 = ttns.copy().compress(m2)
+    optimize_ttns(ttns, ttno, procedure2)
+    s1 = ttns.todense().ravel()
+    s2 = ttns2.todense().ravel()
+
+    np.testing.assert_allclose(np.abs(s1 @ s2), 1, atol=1e-5)
+
+
+@pytest.mark.parametrize("basis_tree", [basis_binary, basis_multi_basis])
+def test_partial_ttno(basis_tree):
+    # apply ttno on ttns when ttns has more dofs than ttno
+    basis_tree2 = basis_tree.add_auxiliary_space()
+    ttns = TTNS.random(basis_tree2, qntot=0, m_max=4)
+    ttno = TTNO(basis_tree, heisenberg_ops(nspin))
+    ttno2 = TTNO(basis_tree2, heisenberg_ops(nspin))
+    # test apply
+    new_ttns = ttno.apply(ttns)
+    new_ttns2 = ttno2.apply(ttns)
+    np.testing.assert_allclose(new_ttns.todense(), new_ttns2.todense())
+    # test expectation
+    e = ttns.expectation(ttno)
+    e2 = ttns.expectation(ttno2)
+    np.testing.assert_allclose(e, e2)
+
+
+@pytest.mark.parametrize("basis", [basis_binary, basis_multi_basis])
+def test_print(basis):
+    basis.print()
+    ttns = TTNS.random(basis, 0, 5, 1)
+    ttns.print_shape()
+    ttns.print_vn_entropy()