Skip to content

Commit

Permalink
add pretty print for Bond
Browse files Browse the repository at this point in the history
  • Loading branch information
kaihsin committed Apr 15, 2024
1 parent 40fb2ee commit ca928cb
Showing 1 changed file with 32 additions and 0 deletions.
32 changes: 32 additions & 0 deletions src/cytnx_torch/bond.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from dataclasses import dataclass, field
from beartype.typing import List
from enum import Enum
from abc import abstractmethod
from .symmetry import Symmetry


Expand Down Expand Up @@ -30,6 +31,11 @@ class AbstractBond:
dim: int
bond_type: BondType = field(default=BondType.NONE)

@property
@abstractmethod
def nsym(self) -> int:
raise NotImplementedError("not implement for abstract type trait.")


@dataclass
class Bond(AbstractBond):
Expand All @@ -38,6 +44,15 @@ def __post_init__(self):
if self.dim < 0:
raise ValueError(f"dim should be non-negative. got {self.dim}")

def __str__(self):
out = f"Dim = {self.dim} |\n"
out += f"{self.bond_type} :\n"
return out

@property
def nsym(self) -> int:
return 0


@dataclass(init=False)
class SymBond(AbstractBond):
Expand All @@ -52,6 +67,10 @@ class SymBond(AbstractBond):
)
_syms: List[Symmetry] = field(default_factory=list)

@property
def nsym(self) -> int:
return len(self._syms)

def _check_qnums(self) -> None:
for i, s in enumerate(self._syms):
if not s.check_qnums(self._qnums[:, i]):
Expand Down Expand Up @@ -83,3 +102,16 @@ def __init__(self, bond_type: BondType, qnums: List[Qs], syms: List[Symmetry]):

# checking qnums are consistent!
self._check_qnums()

def __str__(self):

out = f"Dim = {self.dim} |\n"
out += f"{self.bond_type} :\n"

for n in range(self.nsym):
out += f" {str(self._syms[n])}:: "
for idim in range(len(self._qnums)):
out += " %+d" % (self._qnums[idim, n])
out += "\n "

return out

0 comments on commit ca928cb

Please sign in to comment.