fix dump load coeff bug

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

renormalizer/tn/tests/test_evolve.py

@@ -23,7 +23,6 @@ def add_ttno_offset(ttns: TTNS, ttno: TTNO):
     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)]
@@ -167,6 +166,7 @@ def test_save_load(ttns_and_ttno):
     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(exp1, exp2)
+    np.testing.assert_allclose(exp2, exp1)
     os.remove(fname)

renormalizer/tn/tree.py

@@ -27,7 +27,7 @@ class TTNBase(Tree):
     # A tree whose tree node is TreeNodeTensor
 
     @classmethod
-    def load(cls, basis: BasisTree, fname: str):
+    def load(cls, basis: BasisTree, fname: str, other_attrs=None):
         npload = np.load(fname, allow_pickle=True)
         assert npload["version"] == "0.1"
 
@@ -38,7 +38,10 @@ def load(cls, basis: BasisTree, fname: str):
             qn = npload[f"qn_{i}"]
             nodes.append(TreeNodeTensor(tensor, qn))
         copy_connection(basis.node_list, nodes)
-        return cls(basis, root=nodes[0])
+        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
@@ -47,13 +50,15 @@ def __init__(self, basis: BasisTree, root: TreeNodeTensor):
     def dump(self, fname: str, other_attrs=None):
         if other_attrs is None:
             other_attrs = []
-        elif isinstance(other_attrs, str):
-            other_attrs = [other_attrs]
-        assert isinstance(other_attrs, list)
+
         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
@@ -321,6 +326,13 @@ def __matmul__(self, other):
 
 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):
         """
@@ -1069,6 +1081,12 @@ def scale(self, val, inplace=False):
         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)