Update Edge type output format.

USTC-TNS · Nov 22, 2023 · 2839e88 · 2839e88
1 parent 1932b16
commit 2839e88
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/tat/edge.py b/tat/edge.py
@@ -197,10 +197,15 @@ def __str__(self: Edge) -> str:
         # pylint: disable=no-else-return
         if any(self.fermion):
             # Fermionic edge
-            return f"(dimension={self.dimension}, arrow={self.arrow}, fermion={self.fermion}, symmetry={self.symmetry})"
+            fermion = ','.join(str(sub_fermion) for sub_fermion in self.fermion)
+            symmetry = ','.join(
+                f"[{','.join(str(sub_sym.item()) for sub_sym in sub_symmetry)}]" for sub_symmetry in self.symmetry)
+            return f"(dimension={self.dimension}, arrow={self.arrow}, fermion=({fermion}), symmetry=({symmetry}))"
         elif self.fermion:
             # Bosonic edge
-            return f"(dimension={self.dimension}, symmetry={self.symmetry})"
+            symmetry = ','.join(
+                f"[{','.join(str(sub_sym.item()) for sub_sym in sub_symmetry)}]" for sub_symmetry in self.symmetry)
+            return f"(dimension={self.dimension}, symmetry=({symmetry}))"
         else:
             # Trivial edge
             return f"(dimension={self.dimension})"

diff --git a/tests/test_edge.py b/tests/test_edge.py
@@ -29,10 +29,10 @@ def test_edge_conjugate_and_equal() -> None:
 
 def test_repr() -> None:
     a = Edge(fermion=(False, True), symmetry=(torch.tensor([False, True]), torch.tensor([0, 1])), arrow=True)
-    repr_a = "Edge(dimension=2, arrow=True, fermion=(False, True), symmetry=(tensor([False,  True]), tensor([0, 1])))"
+    repr_a = "Edge(dimension=2, arrow=True, fermion=(False,True), symmetry=([False,True],[0,1]))"
     assert repr_a == repr(a)
     b = Edge(symmetry=(torch.tensor([False, True]), torch.tensor([0, 1])))
-    repr_b = "Edge(dimension=2, symmetry=(tensor([False,  True]), tensor([0, 1])))"
+    repr_b = "Edge(dimension=2, symmetry=([False,True],[0,1]))"
     assert repr_b == repr(b)
     c = Edge(dimension=4)
     repr_c = "Edge(dimension=4)"