reformat ttn

renormalizer/tn/gs.py

@@ -28,12 +28,13 @@ def optimize_ttns(ttns: TTNS, ttno: TTNO, procedure=None):
     return e_list
 
 
-def optimize_recursion(snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, m:int, percent:float=0) -> List[float]:
+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)
@@ -57,7 +58,6 @@ def optimize_recursion(snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNE
 
 
 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()
@@ -70,7 +70,7 @@ def hop(x):
         return asnumpy(ret)
 
     assert ttns.optimize_config.nroots == 1
-    algo:str = ttns.optimize_config.algo
+    algo: str = ttns.optimize_config.algo
     e, c = eigh_iterative(hop, hdiag, cguess, algo)
     c = vec2tensor(c, qn_mask)
     return e, c
@@ -84,9 +84,7 @@ def eigh_iterative(hop, hdiag, cguess, algo):
     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
-        )
+        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")

renormalizer/tn/node.py

@@ -123,7 +123,7 @@ def copy_connection(source_node_list: List[NodeUnion], target_node_list: List[No
     NodeUnion
         The root node of the target tree.
     """
-    node2idx: Dict[NodeUnion, int] = {n:i for i, n in enumerate(source_node_list)}
+    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:
@@ -133,5 +133,3 @@ def copy_connection(source_node_list: List[NodeUnion], target_node_list: List[No
             root = target_node
     assert root is not None
     return root
-
-

renormalizer/tn/symbolic_mpo.py

@@ -49,10 +49,10 @@ def symbolic_mo_to_numeric_mo_general(basis_sets: List[BasisSet], mo, dtype):
             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)
+    return np.moveaxis(mo_tensor, mo.ndim - 1, -1)
 
 
-def construct_symbolic_mpo(tn:BasisTree, terms: List[Op], const:float=0):
+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]))
@@ -70,8 +70,8 @@ def construct_symbolic_mpo(tn:BasisTree, terms: List[Op], const:float=0):
         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:]
+            table_row = table[:, : k + 1]
+            table_col = table[:, k + 1 :]
             in_ops_list = [dummy_in_ops]
         else:
             # the children must have been visited
@@ -81,10 +81,11 @@ def construct_symbolic_mpo(tn:BasisTree, terms: List[Op], const:float=0):
             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)
+            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)

renormalizer/tn/tests/test_evolve.py

@@ -162,7 +162,7 @@ def test_save_load(ttns_and_ttno):
         ttns1 = ttns1.evolve(ttno, tau)
     exp1 = [ttns1.expectation(o) for o in op_n_list]
     ttns2 = ttns.evolve(ttno, tau)
-    fname = "test.npz"
+    fname = f"{id(ttns2)}.npz"
     ttns2.dump(fname)
     ttns2 = TTNS.load(ttns.basis, fname)
     ttns2 = ttns2.evolve(ttno, tau)

renormalizer/tn/time_evolution.py

@@ -52,11 +52,11 @@ def regularized_inversion(m, eps):
     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 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
@@ -67,7 +67,7 @@ def ivp_func(t, params):
     return new_ttns
 
 
-def evolve_prop_and_compress_tdrk4(ttns:TTNS, ttno:TTNO, coeff:Union[complex, float], tau:float):
+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]))
@@ -76,7 +76,7 @@ def evolve_prop_and_compress_tdrk4(ttns:TTNS, ttno:TTNO, coeff:Union[complex, fl
     return compressed_sum(termlist)
 
 
-def evolve_tdvp_ps(ttns:TTNS, ttno:TTNO, coeff:Union[complex, float], tau:float):
+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)
@@ -99,11 +99,7 @@ def evolve_tdvp_ps(ttns:TTNS, ttno:TTNO, coeff:Union[complex, float], tau:float)
     return ttns
 
 
-def _tdvp_ps_forward(ttns: TTNS,
-                     ttno: TTNO,
-                     ttne: TTNEnviron,
-                     coeff: Union[complex, float],
-                     tau: float) -> List[int]:
+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)]
@@ -144,11 +140,7 @@ def _tdvp_ps_forward(ttns: TTNS,
     return local_steps
 
 
-def _tdvp_ps_backward(ttns: TTNS,
-                     ttno: TTNO,
-                     ttne: TTNEnviron,
-                     coeff: Union[complex, float],
-                     tau: float) -> List[int]:
+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)]
@@ -182,7 +174,7 @@ def _tdvp_ps_backward(ttns: TTNS,
     return local_steps
 
 
-def evolve_tdvp_ps2(ttns:TTNS, ttno:TTNO, coeff:Union[complex, float], tau:float):
+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)
@@ -195,12 +187,9 @@ def evolve_tdvp_ps2(ttns:TTNS, ttno:TTNO, coeff:Union[complex, float], tau:float
     return ttns
 
 
-def _tdvp_ps2_recursion_forward(snode: TreeNodeTensor,
-                                ttns: TTNS,
-                                ttno: TTNO,
-                                ttne: TTNEnviron,
-                                coeff:Union[complex, float],
-                                tau:float) -> List[int]:
+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"""
@@ -209,7 +198,6 @@ def _tdvp_ps2_recursion_forward(snode: TreeNodeTensor,
     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)
@@ -237,12 +225,9 @@ def _tdvp_ps2_recursion_forward(snode: TreeNodeTensor,
     return local_steps
 
 
-def _tdvp_ps2_recursion_backward(snode: TreeNodeTensor,
-                                 ttns: TTNS,
-                                 ttno: TTNO,
-                                 ttne: TTNEnviron,
-                                 coeff:Union[complex, float],
-                                 tau:float) -> List[int]:
+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"""
@@ -278,36 +263,36 @@ def _tdvp_ps2_recursion_backward(snode: TreeNodeTensor,
     return local_steps
 
 
-def evolve_2site(snode: TreeNodeTensor, ttns: TTNS, ttno: TTNO, ttne: TTNEnviron, coeff:Union[complex, float], tau:float):
+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()
-    )
+    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):
+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()
-    )
+    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):
+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()
-    )
+    ms_t, j = expm_krylov(lambda y: hop(y.reshape(ms.shape)).ravel(), coeff * tau, ms.ravel())
     return ms_t, j