diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml new file mode 100644 index 000000000..7026a548f --- /dev/null +++ b/.github/workflows/CI.yml @@ -0,0 +1,37 @@ +name: CI + +on: [push, pull_request] + +jobs: + CI: + + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] + + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Display Python version + run: python -c "import sys; print(sys.version)" + - name: Install CI tools + run: pip install pylint mypy pytest pytest-cov + - name: Show directory + run: pwd && ls + working-directory: ${{ runner.workspace }} + - name: Install requirements + run: pip install . + working-directory: ${{ runner.workspace }} + - name: Run pytest + run: pytest --cov=tat + working-directory: ${{ runner.workspace }} + - name: Run mypy + run: mypy tat + working-directory: ${{ runner.workspace }} + - name: Run pylint + run: pylint tat + working-directory: ${{ runner.workspace }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 000000000..87bbb542f --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +.coverage +__pycache__ \ No newline at end of file diff --git a/README.org b/README.org new file mode 100644 index 000000000..f288f3920 --- /dev/null +++ b/README.org @@ -0,0 +1,3 @@ +* TAT + +A Fermionic tensor library based on pytorch. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..31708aa34 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,25 @@ +[tool.pylint] +max-line-length = 120 +generated-members = 'torch.*' + +[tool.yapf] +based_on_style = 'google' +column_limit = 120 + +[tool.mypy] +check_untyped_defs = true +disallow_untyped_defs = true + +[project] +name = 'tat' +version = '0.4.0' +description = "A Fermionic tensor library based on pytorch." +requires-python = ">=3.7" +authors = [ + {email = "zh970205@mail.ustc.edu.cn"}, + {name = "Hao Zhang"} +] +dependencies = [ + 'multimethod', + 'torch', +] diff --git a/tat/__init__.py b/tat/__init__.py new file mode 100644 index 000000000..9f8cc4bcb --- /dev/null +++ b/tat/__init__.py @@ -0,0 +1,6 @@ +""" +The tat is a Fermionic tensor library based on pytorch. +""" + +from .edge import Edge +from .tensor import Tensor diff --git a/tat/compat.py b/tat/compat.py new file mode 100644 index 000000000..c3cae443c --- /dev/null +++ b/tat/compat.py @@ -0,0 +1,191 @@ +""" +This file implement a compat layer for legacy TAT interface. +""" + +from __future__ import annotations +import typing +from multimethod import multimethod +import torch +from .edge import Edge as E +from .tensor import Tensor as T + +# pylint: disable=too-few-public-methods +# pylint: disable=too-many-instance-attributes + + +class CompatSymmetry: + """ + The common Symmetry namespace + """ + + def __init__(self: CompatSymmetry, fermion: tuple[bool, ...], dtypes: tuple[torch.dtype, ...]) -> None: + self.fermion: tuple[bool, ...] = fermion + self.dtypes: tuple[torch.dtype, ...] = dtypes + + # pylint: disable=invalid-name + self.S: CompatScalar + self.D: CompatScalar + self.C: CompatScalar + self.Z: CompatScalar + self.float32: CompatScalar + self.float64: CompatScalar + self.float: CompatScalar + self.complex64: CompatScalar + self.complex128: CompatScalar + self.complex: CompatScalar + + self.S = self.float32 = CompatScalar(self, torch.float32) + self.D = self.float64 = self.float = CompatScalar(self, torch.float64) + self.C = self.complex64 = CompatScalar(self, torch.complex64) + self.Z = self.complex128 = self.complex = CompatScalar(self, torch.complex128) + + def _parse_segments(self: CompatSymmetry, segments: list) -> tuple[tuple[torch.Tensor, ...], int]: + # Segments may be [Sym] or [(Sym, Size)] + try: + # try [(Sym, Size)] first + return self._parse_segments_kernel(segments) + except TypeError: + # Cannot unpack is a type error, value[index] is a type error, too + # convert [Sym] to [(Sym, Size)] + return self._parse_segments_kernel([(sym, 1) for sym in segments]) + + def _parse_segments_kernel(self: CompatSymmetry, segments: list) -> tuple[tuple[torch.Tensor, ...], int]: + # [(Sym, Size)] for every element + dimension = sum(dim for _, dim in segments) + symmetry = tuple( + torch.tensor( + sum( + ([self._parse_segments_get_subsymmetry(sym, index)] * dim + for sym, dim in segments), + [], + ), # Concat all segment for this subsymmetry + dtype=sub_symmetry, + ) # Generate subsymmetry one by one + for index, sub_symmetry in enumerate(self.dtypes)) + return symmetry, dimension + + def _parse_segments_get_subsymmetry(self: CompatSymmetry, sym: object, index: int) -> object: + # Most of time, symmetry is a tuple of subsymmetry + # But if there is only ome subsymmetry, it could not be a tuple but subsymmetry itself. + # pylint: disable=no-else-return + if isinstance(sym, tuple): + return sym[index] + else: + if len(self.fermion) == 1: + return sym + else: + raise TypeError(f"{sym=} is not subscriptable") + + @multimethod + def Edge(self: CompatSymmetry, dimension: int) -> E: + """ + Create edge with compat interface. + """ + # pylint: disable=invalid-name + symmetry = tuple(torch.zeros(dimension, dtype=sub_symmetry) for sub_symmetry in self.dtypes) + return E(fermion=self.fermion, dtypes=self.dtypes, symmetry=symmetry, dimension=dimension, arrow=False) + + @Edge.register + def _(self: CompatSymmetry, segments: list, arrow: bool = False) -> E: + symmetry, dimension = self._parse_segments(segments) + return E(fermion=self.fermion, dtypes=self.dtypes, symmetry=symmetry, dimension=dimension, arrow=arrow) + + @Edge.register + def _(self: CompatSymmetry, segments_and_bool: tuple[list, bool]) -> E: + segments, arrow = segments_and_bool + symmetry, dimension = self._parse_segments(segments) + return E(fermion=self.fermion, dtypes=self.dtypes, symmetry=symmetry, dimension=dimension, arrow=arrow) + + +class CompatScalar: + """ + The common Scalar namespace. + """ + + def __init__(self: CompatScalar, symmetry: CompatSymmetry, dtype: torch.dtype) -> None: + self.symmetry: CompatSymmetry = symmetry + self.dtype: torch.dtype = dtype + + @multimethod + def Tensor(self: CompatScalar, names: list[str], edges: list) -> T: + """ + Create tensor with compat names and edges. + """ + # pylint: disable=invalid-name + return T( + tuple(names), + tuple(self.symmetry.Edge(edge) for edge in edges), + fermion=self.symmetry.fermion, + dtypes=self.symmetry.dtypes, + dtype=self.dtype, + ) + + @Tensor.register + def _(self: CompatScalar) -> T: + result = T( + (), + (), + fermion=self.symmetry.fermion, + dtypes=self.symmetry.dtypes, + dtype=self.dtype, + ) + result.data.reshape([-1])[0] = 1 + return result + + @Tensor.register + def _( + self: CompatScalar, + number: typing.Any, + names: list[str] | None = None, + edge_symmetry: list | None = None, + edge_arrow: list[bool] | None = None, + ) -> T: + # Create high rank tensor with only one element + if names is None: + names = [] + if edge_symmetry is None: + edge_symmetry = [None for _ in names] + if edge_arrow is None: + edge_arrow = [False for _ in names] + result = T( + tuple(names), + tuple( + E( + fermion=self.symmetry.fermion, + dtypes=self.symmetry.dtypes, + symmetry=tuple( + torch.tensor([self._create_size1_get_subsymmetry(symmetry, index)], dtype=dtype) + for index, dtype in enumerate(self.symmetry.dtypes)), + dimension=1, + arrow=arrow, + ) + for symmetry, arrow in zip(edge_symmetry, edge_arrow)), + fermion=self.symmetry.fermion, + dtypes=self.symmetry.dtypes, + dtype=self.dtype, + ) + result.data.reshape([-1])[0] = number + return result + + def _create_size1_get_subsymmetry(self: CompatScalar, sym: object, index: int) -> object: + # pylint: disable=no-else-return + if sym is None: + return 0 + elif isinstance(sym, tuple): + return sym[index] + else: + if len(self.symmetry.fermion) == 1: + return sym + else: + raise TypeError(f"{sym=} is not subscriptable") + + +No: CompatSymmetry = CompatSymmetry(fermion=(), dtypes=()) +Z2: CompatSymmetry = CompatSymmetry(fermion=(False,), dtypes=(torch.bool,)) +U1: CompatSymmetry = CompatSymmetry(fermion=(False,), dtypes=(torch.int,)) +Fermi: CompatSymmetry = CompatSymmetry(fermion=(True,), dtypes=(torch.int,)) +FermiZ2: CompatSymmetry = CompatSymmetry(fermion=(True, False), dtypes=(torch.int, torch.bool)) +FermiU1: CompatSymmetry = CompatSymmetry(fermion=(True, False), dtypes=(torch.int, torch.int)) +Parity: CompatSymmetry = CompatSymmetry(fermion=(True,), dtypes=(torch.bool,)) +FermiFermi: CompatSymmetry = CompatSymmetry(fermion=(True, True), dtypes=(torch.int, torch.int)) +Normal: CompatSymmetry = No diff --git a/tat/edge.py b/tat/edge.py new file mode 100644 index 000000000..972bdc5e1 --- /dev/null +++ b/tat/edge.py @@ -0,0 +1,278 @@ +""" +This file contains the definition of tensor edge. +""" + +from __future__ import annotations +import typing +import functools +import operator +import torch + + +class Edge: + """ + The edge type of tensor. + """ + + __slots__ = "_fermion", "_dtypes", "_symmetry", "_dimension", "_arrow", "_parity" + + @property + def fermion(self: Edge) -> tuple[bool, ...]: + """ + A tuple records whether every sub symmetry is fermionic. Its length is the number of sub symmetry. + """ + return self._fermion + + @property + def dtypes(self: Edge) -> tuple[torch.dtype, ...]: + """ + A tuple records the basic dtype of every sub symmetry. Its length is the number of sub symmetry. + """ + return self._dtypes + + @property + def symmetry(self: Edge) -> tuple[torch.Tensor, ...]: + """ + A tuple containing all symmetry of this edge. Its length is the number of sub symmetry. Every element of it is a + sub symmetry. + """ + return self._symmetry + + @property + def dimension(self: Edge) -> int: + """ + The dimension of this edge. + """ + return self._dimension + + @property + def arrow(self: Edge) -> bool: + """ + The arrow of this edge. + """ + return self._arrow + + @property + def parity(self: Edge) -> torch.Tensor: + """ + The parity of this edge. + """ + return self._parity + + def __init__( + self: Edge, + *, + fermion: tuple[bool, ...] = (), + dtypes: tuple[torch.dtype, ...] = (), + symmetry: tuple[torch.Tensor, ...] = (), + dimension: int | None = None, + arrow: bool | None = None, + **kwargs: torch.Tensor, + ) -> None: + """ + Create an edge with essential information. + + Examples: + - Edge(dimension=5) + - Edge(symmetry=(torch.tensor([False, False, True, True]),)) + - Edge(fermion=(False, True), symmetry=(torch.tensor([False, True]), torch.tensor([False, True])), arrow=True) + + Parameters + ---------- + fermion : tuple[bool, ...], default=() + Whether each sub symmetry is fermionic symmetry, its length should be the same to symmetry. But it could be + left empty, if so, a total bosonic edge will be created. + dtypes : tuple[torch.dtype, ...], default=() + The basic dtype to identify each sub symmetry, its length should be the same to symmetry, and it is nothing + but the dtypes of each tensor in the symmetry. It could be left empty, if so, it will be derived from + symmetry. + symmetry : tuple[torch.Tensor, ...], default=() + The symmetry information of every sub symmetry, each of sub symmetry should be a one dimensional tensor with + the same length dimension, and their dtype should be integral type, aka, int or bool. + dimension : int, optional + The dimension of the edge, if not specified, dimension will be detected from symmetry. + arrow : bool, optional + The arrow direction of the edge, it is essential for fermionic edge, aka, an edge with fermionic sub + symmetry. + """ + # Fermion could be empty if it is total bosonic edge + if not fermion: + fermion = tuple(False for _ in symmetry) + + # Dtypes could be empty and derived from symmetry + if not dtypes: + dtypes = tuple(sub_symmetry.dtype for sub_symmetry in symmetry) + # Check dtype is compatible with symmetry + assert all(sub_symmetry.dtype is sub_dtype for sub_symmetry, sub_dtype in zip(symmetry, dtypes)) + # Check dtype is valid, aka, bool or int + assert all(not (sub_symmetry.is_floating_point() or sub_symmetry.is_complex()) for sub_symmetry in symmetry) + + # The fermion, dtypes and symmetry information should have the same length + assert len(fermion) == len(dtypes) == len(symmetry) + + # If dimension not set, get dimension from symmetry + if dimension is None: + dimension = len(symmetry[0]) + # Check if the dimensions of different sub_symmetry mismatch + assert all(sub_symmetry.size() == (dimension,) for sub_symmetry in symmetry) + + if arrow is None: + # Arrow not set, it should be bosonic edge. + arrow = False + assert not any(fermion) + + self._fermion: tuple[bool, ...] = fermion + self._dtypes: tuple[torch.dtype, ...] = dtypes + self._symmetry: tuple[torch.Tensor, ...] = symmetry + self._dimension: int = dimension + self._arrow: bool = arrow + + self._parity: torch.Tensor + if "parity" in kwargs: + self._parity = kwargs.pop("parity") + else: + self._parity = self._generate_parity() + assert not kwargs + + def _generate_parity(self: Edge) -> torch.Tensor: + return functools.reduce( + # Reduce sub parity by xor + torch.logical_xor, + ( + # The parity of sub symmetry + sub_symmetry if sub_symmetry.dtype is torch.bool else sub_symmetry % 2 != 0 + # Loop all sub symmetry + for sub_symmetry, sub_fermion in zip(self.symmetry, self.fermion) + # But only reduce if it is fermion sub symmetry + if sub_fermion), + # Reduce with start as tensor filled with False + torch.zeros(self.dimension, dtype=torch.bool), + ) + + def conjugated(self: Edge) -> Edge: + """ + Get the conjugated edge. + + Returns + ------- + Edge + The conjugated edge. + """ + # The only two difference of conjguated edge is symmetry and arrow + return Edge( + fermion=self.fermion, + dtypes=self.dtypes, + symmetry=tuple( + sub_symmetry if sub_symmetry.dtype is torch.bool else -sub_symmetry # bool -> same, int -> neg + for sub_symmetry in self.symmetry), + dimension=self.dimension, + arrow=not self.arrow, + parity=self.parity, + ) + + def __eq__(self: Edge, other: object) -> bool: + if not isinstance(other, Edge): + return NotImplemented + return (self.dimension == other.dimension and # Compare int dimension and bool arrow first + self.arrow == other.arrow and # Since they are fast to compare + self.fermion == other.fermion and # Then the tuple of bool are compared + self.dtypes == other.dtypes and # Then the tuple of dtypes are compared + all( # All of symmetries are compared at last, since it is biggest + torch.equal(self_sub_symmetry, other_sub_symmetry) + for self_sub_symmetry, other_sub_symmetry in zip(self.symmetry, other.symmetry))) + + 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})" + elif self.fermion: + # Bosonic edge + return f"(dimension={self.dimension}, symmetry={self.symmetry})" + else: + # Trivial edge + return f"(dimension={self.dimension})" + + def __repr__(self: Edge) -> str: + return f"Edge{self.__str__()}" + + @staticmethod + def merge_edges( + edges: tuple[Edge, ...], + fermion: tuple[bool, ...] = (), + dtypes: tuple[torch.dtype, ...] = (), + arrow: bool | None = None, + ) -> Edge: + """ + Merge several edges into one edge. + + Parameters + ---------- + edges : tuple[Edge, ...] + The edges to be merged. + fermion : tuple[bool, ...], default=() + Whether each sub symmetry is fermionic, it could be left empty to derive from edges + dtypes : tuple[torch.dtype, ...], default=() + The base type of sub symmetry, it could be left empty to derive from edges + arrow : bool, optional + The arrow of all the edges, it is useful if edges is empty. + + Returns + ------- + Edge + The result edge merged by edges. + """ + # If fermion left empty, get it from edges + if not fermion: + fermion = edges[0].fermion + # All edge should share the same fermion + assert all(fermion == edge.fermion for edge in edges) + # If dtypes left empty, get it from edges + if not dtypes: + dtypes = edges[0].dtypes + # All edge should share the same dtypes + assert all(dtypes == edge.dtypes for edge in edges) + # If arrow set, check it directly, if not set, set to False or get from edges + if arrow is None: + if any(fermion): + # It is fermionic edge. + arrow = edges[0].arrow + else: + # It is bosonic edge, set to False directly since it is useless. + arrow = False + # All edge should share the same arrow + assert all(arrow == edge.arrow for edge in edges) + + rank = len(edges) + # Merge edge + dimension = functools.reduce(operator.mul, (edge.dimension for edge in edges), 1) + # pylint: disable=duplicate-code + symmetry = tuple( + # Every merged sub symmetry is calculated by reduce and flatten + functools.reduce( + # The reduce operator depend on the dtype of this sub symmetry + typing.cast( + typing.Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + # mypy cannot derive the type of the following expression, cast it manually. + torch.logical_xor if sub_symmetry_dtype is torch.bool else torch.add), + ( + # The sub symmetry of every edge will be reshape to be reduced. + edge.symmetry[sub_symmetry_index].reshape( + # This reshape is like unsqueeze, but for high dimension. + [-1 if index == current_index else 1 for index in range(rank)]) + # The sub symmetry of every edge is reduced one by one + for current_index, edge in enumerate(edges)), + # Reduce from a rank-0 tensor + torch.zeros([], dtype=sub_symmetry_dtype), + ).reshape([-1]) + # Merge every sub symmetry one by one + for sub_symmetry_index, sub_symmetry_dtype in enumerate(dtypes)) + + # parity not set here since it need recalculation + return Edge( + fermion=fermion, + dtypes=dtypes, + symmetry=symmetry, + dimension=dimension, + arrow=arrow, + ) diff --git a/tat/tensor.py b/tat/tensor.py new file mode 100644 index 000000000..bd98c06c2 --- /dev/null +++ b/tat/tensor.py @@ -0,0 +1,799 @@ +""" +This file defined the core tensor type for tat package. +""" + +from __future__ import annotations +import typing +import functools +from multimethod import multimethod +import torch +from .edge import Edge + +# pylint: disable=too-many-public-methods + + +class Tensor: + """ + The main tensor type, which wraps pytorch tensor and provides edge names and Fermionic functions. + """ + + __slots__ = "_fermion", "_dtypes", "_names", "_edges", "_data", "_mask" + + def __str__(self: Tensor) -> str: + return f"(names={self.names}, edges={self.edges}, data={self.data})" + + def __repr__(self: Tensor) -> str: + return f"Tensor(names={self.names}, edges={self.edges})" + + @property + def fermion(self: Tensor) -> tuple[bool, ...]: + """ + A tuple records whether every sub symmetry is fermionic. Its length is the number of sub symmetry. + """ + return self._fermion + + @property + def dtypes(self: Tensor) -> tuple[torch.dtype, ...]: + """ + A tuple records the basic dtype of every sub symmetry. Its length is the number of sub symmetry. + """ + return self._dtypes + + @property + def names(self: Tensor) -> tuple[str, ...]: + """ + The edge names of this tensor. + """ + return self._names + + @property + def edges(self: Tensor) -> tuple[Edge, ...]: + """ + The edges information of this tensor. + """ + return self._edges + + @property + def data(self: Tensor) -> torch.Tensor: + """ + The content data of this tensor. + """ + return self._data + + @property + def mask(self: Tensor) -> torch.Tensor: + """ + The content data mask of this tensor. + """ + return self._mask + + @property + def rank(self: Tensor) -> int: + """ + The rank of this tensor. + """ + return len(self._names) + + @property + def dtype(self: Tensor) -> torch.dtype: + """ + The data type of the content in this tensor. + """ + return self.data.dtype + + @property + def btype(self: Tensor) -> str: + """ + The data type of the content in this tensor, represented in BLAS/LAPACK convention. + """ + if self.dtype is torch.float32: + return 'S' + if self.dtype is torch.float64: + return 'D' + if self.dtype is torch.complex64: + return 'C' + if self.dtype is torch.complex128: + return 'Z' + return '?' + + @property + def is_complex(self: Tensor) -> bool: + """ + Whether it is a complex tensor + """ + return self.dtype.is_complex + + @property + def is_real(self: Tensor) -> bool: + """ + Whether it is a real tensor + """ + return self.dtype.is_floating_point + + def edge_by_name(self: Tensor, name: str) -> Edge: + """ + Get edge by the edge name of this tensor. + + Parameters + ---------- + name : str + The given edge name. + + Returns + ------- + Edge + The edge with the given edge name. + """ + return self.edges[self.names.index(name)] + + def _arithmetic_operator(self: Tensor, other: object, operate: typing.Callable) -> Tensor: + if isinstance(other, Tensor): + # If it is tensor, check same shape and transpose before calculating. + assert self.same_shape_with(other) + new_data = operate(self.data, other.transpose(self.names).data) + else: + # Otherwise treat other as a scalar, mask should be applied later. + new_data = operate(self.data, other) + new_data *= self.mask + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=new_data, + mask=self.mask, + ) + + def __add__(self: Tensor, other: object) -> Tensor: + return self._arithmetic_operator(other, torch.add) + + def __sub__(self: Tensor, other: object) -> Tensor: + return self._arithmetic_operator(other, torch.sub) + + def __mul__(self: Tensor, other: object) -> Tensor: + return self._arithmetic_operator(other, torch.mul) + + def __truediv__(self: Tensor, other: object) -> Tensor: + return self._arithmetic_operator(other, torch.div) + + def _right_arithmetic_operator(self: Tensor, other: object, operate: typing.Callable) -> Tensor: + if isinstance(other, Tensor): + # If it is tensor, check same shape and transpose before calculating. + assert self.same_shape_with(other) + new_data = operate(other.transpose(self.names).data, self.data) + else: + # Otherwise treat other as a scalar, mask should be applied later. + new_data = operate(other, self.data) + new_data *= self.mask + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=new_data, + mask=self.mask, + ) + + def __radd__(self: Tensor, other: object) -> Tensor: + return self._right_arithmetic_operator(other, torch.add) + + def __rsub__(self: Tensor, other: object) -> Tensor: + return self._right_arithmetic_operator(other, torch.sub) + + def __rmul__(self: Tensor, other: object) -> Tensor: + return self._right_arithmetic_operator(other, torch.mul) + + def __rtruediv__(self: Tensor, other: object) -> Tensor: + return self._right_arithmetic_operator(other, torch.div) + + def _inplace_arithmetic_operator(self: Tensor, other: object, operate: typing.Callable) -> Tensor: + if isinstance(other, Tensor): + # If it is tensor, check same shape and transpose before calculating. + assert self.same_shape_with(other) + operate(self.data, other.transpose(self.names).data, out=self.data) + else: + # Otherwise treat other as a scalar, mask should be applied later. + operate(self.data, other, out=self.data) + torch.mul(self.data, self.mask, out=self.data) + return self + + def __iadd__(self: Tensor, other: object) -> Tensor: + return self._inplace_arithmetic_operator(other, torch.add) + + def __isub__(self: Tensor, other: object) -> Tensor: + return self._inplace_arithmetic_operator(other, torch.sub) + + def __imul__(self: Tensor, other: object) -> Tensor: + return self._inplace_arithmetic_operator(other, torch.mul) + + def __itruediv__(self: Tensor, other: object) -> Tensor: + return self._inplace_arithmetic_operator(other, torch.div) + + def __pos__(self: Tensor) -> Tensor: + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=+self.data, + mask=self.mask, + ) + + def __neg__(self: Tensor) -> Tensor: + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=-self.data, + mask=self.mask, + ) + + def __float__(self: Tensor) -> float: + return float(self.data) + + def __complex__(self: Tensor) -> complex: + return complex(self.data) + + def norm(self: Tensor, order: typing.Any) -> float: + """ + Get the norm of the tensor, this function will flatten tensor first before calculate norm. + + Parameters + ---------- + order + The order of norm. + + Returns + ------- + float + The norm of the tensor. + """ + return torch.linalg.vector_norm(self.data, ord=order) + + def norm_max(self: Tensor) -> float: + "max norm" + return self.norm(+torch.inf) + + def norm_min(self: Tensor) -> float: + "min norm" + return self.norm(-torch.inf) + + def norm_num(self: Tensor) -> float: + "0-norm" + return self.norm(0) + + def norm_sum(self: Tensor) -> float: + "1-norm" + return self.norm(1) + + def norm_2(self: Tensor) -> float: + "2-norm" + return self.norm(2) + + def copy(self: Tensor) -> Tensor: + """ + Get a copy of this tensor + + Returns + ------- + Tensor + The copy of this tensor + """ + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=torch.clone(self.data), + mask=self.mask, + ) + + def __copy__(self: Tensor) -> Tensor: + return self.copy() + + def __deepcopy__(self: Tensor, _: typing.Any = None) -> Tensor: + return self.copy() + + def same_shape(self: Tensor) -> Tensor: + """ + Get a tensor with same shape to this tensor + + Returns + ------- + Tensor + A new tensor with the same shape to this tensor + """ + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=torch.zeros_like(self.data), + mask=self.mask, + ) + + def zero_(self: Tensor) -> Tensor: + """ + Set all element to zero in this tensor + + Returns + ------- + Tensor + Return this tensor itself. + """ + self.data.zero_() + return self + + def sqrt(self: Tensor) -> Tensor: + """ + Get the sqrt of the tensor. + + Returns + ------- + Tensor + The sqrt of this tensor. + """ + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=torch.sqrt(torch.abs(self.data)), + mask=self.mask, + ) + + def reciprocal(self: Tensor) -> Tensor: + """ + Get the reciprocal of the tensor. + + Returns + ------- + Tensor + The reciprocal of this tensor. + """ + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=torch.where(self.data == 0, self.data, 1 / self.data), + mask=self.mask, + ) + + def range_(self: Tensor, first: typing.Any = 0, step: typing.Any = 1) -> Tensor: + """ + A useful function to generate simple data in tensor for test. + + Parameters + ---------- + first, step + Parameters to generate data. + + Returns + ------- + Tensor + Returns the tensor itself. + """ + data = torch.cumsum(self.mask.reshape([-1]), dim=0, dtype=self.data.dtype).reshape(self.data.size()) + data = (data - 1) * step + first + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=data, + mask=self.mask, + ) + + def to(self: Tensor, new_type: typing.Any) -> Tensor: + """ + Convert this tensor to other scalar type. + + Parameters + ---------- + new_type + The scalar data type of the new tensor. + """ + # pylint: disable=invalid-name + if isinstance(new_type, str): + if new_type in ["float32", "S"]: + new_type = torch.float32 + elif new_type in ["float64", "float", "D"]: + new_type = torch.float64 + elif new_type in ["complex64", "C"]: + new_type = torch.complex64 + elif new_type in ["complex128", "complex", "Z"]: + new_type = torch.complex128 + return Tensor( + names=self.names, + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=self.data.to(new_type), + mask=self.mask, + ) + + def __init__( + self: Tensor, + names: tuple[str, ...], + edges: tuple[Edge, ...], + *, + dtype: torch.dtype = torch.float, + fermion: tuple[bool, ...] = (), + dtypes: tuple[torch.dtype, ...] = (), + **kwargs: torch.Tensor, + ) -> None: + """ + Create a tensor with specific shape. + + Parameters + ---------- + names : tuple[str, ...] + The edge names of the tensor, which length is just the tensor rank. + edges : tuple[Edge, ...] + The detail information of each edge, which length is just the tensor rank. + dtype : torch.dtype, default=torch.float + The dtype of the tensor. + fermion : tuple[bool, ...], default=() + Whether each sub symmetry is fermionic, it could be left empty to derive from edges + dtypes : tuple[torch.dtype, ...], default=() + The base type of sub symmetry, it could be left empty to derive from edges + """ + # Check the rank is correct in names and edges + assert len(names) == len(edges) + # Check whether there are duplicated names + assert len(set(names)) == len(names) + # If fermion left empty, get it from edges + if not fermion: + fermion = edges[0].fermion + # If dtypes left empty, get it from edges + if not dtypes: + dtypes = edges[0].dtypes + # Check if fermion is correct + assert all(edge.fermion == fermion for edge in edges) + # Check if dtypes is correct + assert all(edge.dtypes == dtypes for edge in edges) + + self._fermion: tuple[bool, ...] = fermion + self._dtypes: tuple[torch.dtype, ...] = dtypes + self._names: tuple[str, ...] = names + self._edges: tuple[Edge, ...] = edges + + self._data: torch.Tensor + if "data" in kwargs: + self._data = kwargs.pop("data") + else: + self._data = torch.zeros( + [edge.dimension for edge in self.edges], + dtype=dtype, + ) + self._mask: torch.Tensor + if "mask" in kwargs: + self._mask = kwargs.pop("mask") + else: + self._mask = self._generate_mask() + assert not kwargs + + def _generate_mask(self: Tensor) -> torch.Tensor: + # pylint: disable=duplicate-code + return functools.reduce( + # Mask is valid if all sub symmetry give valid mask. + torch.logical_and, + ( + # The mask is valid if total symmetry is False or total symmetry is 0 + torch.zeros([], dtype=sub_symmetry_dtype) == + # total sub symmetry is calculated by reduce + functools.reduce( + # The reduce operator depend on the dtype of this sub symmetry + typing.cast( + typing.Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + # mypy cannot derive the type of the following expression, cast it manually. + torch.logical_xor if sub_symmetry_dtype is torch.bool else torch.add), + ( + # The sub symmetry of every edge will be reshape to be reduced. + edge.symmetry[sub_symmetry_index].reshape( + # This reshape is like unsqueeze, but for high dimension. + [-1 if index == current_index else 1 + for index in range(self.rank)]) + # The sub symmetry of every edge is reduced one by one + for current_index, edge in enumerate(self.edges)), + # Reduce from a rank-0 tensor + torch.zeros([], dtype=sub_symmetry_dtype), + ) + # Calculate mask on every sub symmetry one by one + for sub_symmetry_index, sub_symmetry_dtype in enumerate(self.dtypes)), + # Reduce from all true mask. + torch.ones(size=self.data.size(), dtype=torch.bool), + ) + + @multimethod + def _prepare_position(self: Tensor, position: tuple[int, ...]) -> tuple[int, ...]: + return position + + @_prepare_position.register + def _(self: Tensor, position: tuple[slice, ...]) -> tuple[int, ...]: + index_by_name: dict[str, int] = {s.start: s.stop for s in position} + return tuple(index_by_name[name] for name in self.names) + + @_prepare_position.register + def _(self: Tensor, position: dict[str, int]) -> tuple[int, ...]: + return tuple(position[name] for name in self.names) + + def __getitem__(self: Tensor, position: tuple[int, ...] | tuple[slice, ...] | dict[str, int]) -> typing.Any: + """ + Get the element of the tensor + + Parameters + ---------- + position : tuple[int, ...] | tuple[slice, ...] | dict[str, int] + The position of the element, which could be either tuple of index directly or a map from edge name to the + index in the corresponding edge. + """ + indices: tuple[int, ...] = self._prepare_position(position) + assert len(indices) == self.rank + assert all(0 <= index < edge.dimension for edge, index in zip(self.edges, indices)) + return self.data[indices] + + def __setitem__(self: Tensor, position: tuple[int, ...] | tuple[slice, ...] | dict[str, int], + value: typing.Any) -> None: + """ + Set the element of the tensor + + Parameters + ---------- + position : tuple[int, ...] | tuple[slice, ...] | dict[str, int] + The position of the element, which could be either tuple of index directly or a map from edge name to the + index in the corresponding edge. + """ + indices = self._prepare_position(position) + assert len(indices) == self.rank + assert all(0 <= index < edge.dimension for edge, index in zip(self.edges, indices)) + if self.mask[indices]: + self.data[indices] = value + + def clear_symmetry(self: Tensor) -> Tensor: + """ + Clear all symmetry of this tensor. + + Returns + ------- + Tensor + The result tensor with symmetry cleared. + """ + # Mask must be generated again + # pylint: disable=no-else-return + if any(self.fermion): + return Tensor( + names=self.names, + edges=tuple( + Edge( + fermion=(True,), + dtypes=(torch.bool,), + symmetry=(edge.parity,), + dimension=edge.dimension, + arrow=edge.arrow, + parity=edge.parity, + ) for edge in self.edges), + fermion=(True,), + dtypes=(torch.bool,), + data=self.data, + ) + else: + return Tensor( + names=self.names, + edges=tuple( + Edge( + fermion=(), + dtypes=(), + symmetry=(), + dimension=edge.dimension, + arrow=edge.arrow, + parity=edge.parity, + ) for edge in self.edges), + fermion=(), + dtypes=(), + data=self.data, + ) + + def randn_(self: Tensor, mean: float = 0., std: float = 1.) -> Tensor: + """ + Fill the tensor with random number in normal distribution. + + Parameters + ---------- + mean, std : float + The parameter of normal distribution. + + Returns + ------- + Tensor + Return this tensor itself. + """ + self.data.normal_(mean, std) + torch.mul(self.data, self.mask, out=self.data) + return self + + def rand_(self: Tensor, low: float = 0., high: float = 1.) -> Tensor: + """ + Fill the tensor with random number in uniform distribution. + + Parameters + ---------- + low, high : float + The parameter of uniform distribution. + + Returns + ------- + Tensor + Return this tensor itself. + """ + self.data.uniform_(low, high) + torch.mul(self.data, self.mask, out=self.data) + return self + + def same_type_with(self: Tensor, other: Tensor) -> bool: + """ + Check whether two tensor has the same type, that is to say they share the same symmetry type. + """ + return self.fermion == other.fermion and self.dtypes == other.dtypes + + def same_shape_with(self: Tensor, other: Tensor, *, allow_transpose: bool = True) -> bool: + """ + Check whether two tensor has the same shape, that is to say the only differences between them are transpose and + data difference. + """ + if not self.same_type_with(other): + return False + # pylint: disable=no-else-return + if allow_transpose: + return dict(zip(self.names, self.edges)) == dict(zip(other.names, other.edges)) + else: + return self.names == other.names and self.edges == other.edges + + def edge_rename(self: Tensor, name_map: dict[str, str]) -> Tensor: + """ + Rename edge name for this tensor. + + Parameters + ---------- + name_map : dict[str, str] + The name map to be used in renaming edge name. + + Returns + ------- + Tensor + The tensor with names renamed. + """ + return Tensor( + names=tuple(name_map.get(name, name) for name in self.names), + edges=self.edges, + fermion=self.fermion, + dtypes=self.dtypes, + data=self.data, + mask=self.mask, + ) + + def transpose(self: Tensor, names: tuple[str, ...]) -> Tensor: + """ + Transpose the tensor outplace. + + Parameters + ---------- + names : tuple[str, ...] + The new edge order identified by edge names. + + Returns + ------- + Tensor + The transpsoe tensor. + """ + if names == self.names: + return self + assert len(names) == len(self.names) + assert set(names) == set(self.names) + before_by_after = tuple(self.names.index(name) for name in names) + after_by_before = tuple(names.index(name) for name in self.names) + data = self.data.permute(before_by_after) + mask = self.mask.permute(before_by_after) + if any(self.fermion): + # It is fermionic tensor + parities_before_transpose = tuple( + edge.parity.reshape([-1 if index == current_index else 1 + for index in range(self.rank)]) + for current_index, edge in enumerate(self.edges)) + parity = torch.zeros([], dtype=torch.bool) + # Loop every 0 <= i < j < rank + for j in range(self.rank): + for i in range(0, j): + if after_by_before[i] > after_by_before[j]: + parity_addon = torch.logical_and(parities_before_transpose[i], parities_before_transpose[j]) + parity = parity.logical_xor(parity_addon) + # parity True -> -x + # parity False -> +x + data = torch.where(parity, -data, +data) + return Tensor( + names=names, + edges=tuple(self.edges[index] for index in before_by_after), + fermion=self.fermion, + dtypes=self.dtypes, + data=data, + mask=mask, + ) + + def reverse_edge( + self: Tensor, + reversed_names: set[str], + apply_parity: bool = False, + parity_exclude_names: set[str] | None = None, + ) -> Tensor: + """ + Reverse some edge in the tensor. + + Parameters + ---------- + reversed_names : set[str] + The edge names of those edges which will be reversed + apply_parity : bool, default=False + Whether to apply parity caused by reversing edge, since reversing edge will generate half a sign. + parity_exclude_names : set[str] | None, default=None + The name of edges in the different behavior other than default set by apply_parity. + """ + if not any(self.fermion): + return self + if parity_exclude_names is None: + parity_exclude_names = set() + parity = torch.zeros([], dtype=torch.bool) + for current_index, [name, edge] in enumerate(zip(self.names, self.edges)): + if name in reversed_names: + if apply_parity ^ (name in parity_exclude_names): + parity_addon = edge.parity.reshape( + [-1 if index == current_index else 1 for index in range(self.rank)]) + parity = parity.logical_xor(parity_addon) + data = torch.where(parity, -self.data, +self.data) + return Tensor( + names=self.names, + edges=tuple( + Edge( + fermion=edge.fermion, + dtypes=edge.dtypes, + symmetry=edge.symmetry, + dimension=edge.dimension, + arrow=not edge.arrow if self.names[current_index] in reversed_names else edge.arrow, + parity=edge.parity, + ) for current_index, edge in enumerate(self.edges)), + fermion=self.fermion, + dtypes=self.dtypes, + data=data, + mask=self.mask, + ) + + # def split_edge( + # self: Tensor, + # split_map: dict[str, tuple[tuple[str, Edge], ...]], + # apply_parity: bool = False, + # parity_exclude_names: set[str] | None = None, + # ) -> Tensor: + # """ + # Split some edges in this tensor. + + # Parameters + # ---------- + # split_map : dict[str, tuple[tuple[str, Edge], ...]] + # The edge splitting plan. + # apply_parity : bool, default=False + # Whether to apply parity caused by splitting edge, since splitting edge will generate half a sign. + # parity_exclude_names : set[str] | None, default=None + # The name of edges in the different behavior other than default set by apply_parity. + # """ + # if parity_exclude_names is None: + # parity_exclude_names = set() + # raise NotImplementedError() + # #return Tensor( + # # names=names, + # # edges=edges, + # # fermion=self.fermion, + # # dtypes=self.dtypes, + # # data=data, + # # mask=mask, + # #) diff --git a/tests/test_compat.py b/tests/test_compat.py new file mode 100644 index 000000000..be78a74be --- /dev/null +++ b/tests/test_compat.py @@ -0,0 +1,113 @@ +import torch +import tat +import tat.compat as compat + + +def test_edge_from_dimension(): + assert compat.No.Edge(4) == tat.Edge(dimension=4) + assert compat.Fermi.Edge(4) == tat.Edge(fermion=(True,), + symmetry=(torch.tensor([0, 0, 0, 0], dtype=torch.int),), + arrow=False) + assert compat.Z2.Edge(4) == tat.Edge(symmetry=(torch.tensor([False, False, False, False]),)) + + +def test_edge_from_segments(): + assert compat.Z2.Edge([ + (False, 2), + (True, 3), + ]) == tat.Edge(symmetry=(torch.tensor([False, False, True, True, True]),),) + assert compat.Fermi.Edge([ + (-1, 1), + (0, 2), + (+1, 3), + ], True) == tat.Edge( + symmetry=(torch.tensor([-1, 0, 0, +1, +1, +1], dtype=torch.int),), + arrow=True, + fermion=(True,), + ) + assert compat.FermiFermi.Edge([ + ((-1, -2), 1), + ((0, +1), 2), + ((+1, 0), 3), + ], True) == tat.Edge( + symmetry=( + torch.tensor([-1, 0, 0, +1, +1, +1], dtype=torch.int), + torch.tensor([-2, +1, +1, 0, 0, 0], dtype=torch.int), + ), + arrow=True, + fermion=(True, True), + ) + + +def test_edge_from_segments_without_dimension(): + assert compat.Z2.Edge([False, True]) == tat.Edge(symmetry=(torch.tensor([False, True]),)) + assert compat.Fermi.Edge([-1, 0, +1], True) == tat.Edge( + symmetry=(torch.tensor([-1, 0, +1], dtype=torch.int),), + arrow=True, + fermion=(True,), + ) + assert compat.FermiFermi.Edge([ + (-1, -2), + (0, +1), + (+1, 0), + ], True) == tat.Edge( + symmetry=(torch.tensor([-1, 0, +1], dtype=torch.int), torch.tensor([-2, +1, 0], dtype=torch.int)), + arrow=True, + fermion=(True, True), + ) + + +def test_edge_from_tuple(): + assert compat.FermiFermi.Edge(([ + ((-1, -2), 1), + ((0, +1), 2), + ((+1, 0), 3), + ], True)) == tat.Edge( + symmetry=( + torch.tensor([-1, 0, 0, +1, +1, +1], dtype=torch.int), + torch.tensor([-2, +1, +1, 0, 0, 0], dtype=torch.int), + ), + arrow=True, + fermion=(True, True), + ) + assert compat.FermiFermi.Edge(([ + (-1, -2), + (0, +1), + (+1, 0), + ], True)) == tat.Edge( + symmetry=(torch.tensor([-1, 0, +1], dtype=torch.int), torch.tensor([-2, +1, 0], dtype=torch.int)), + arrow=True, + fermion=(True, True), + ) + + +def test_tensor(): + a = compat.FermiZ2.D.Tensor(["i", "j"], [ + [(-1, False), (-1, True), (0, True), (0, False)], + [(+1, True), (+1, False), (0, False), (0, True)], + ]) + b = tat.Tensor( + ( + "i", + "j", + ), + ( + tat.Edge( + fermion=(True, False), + symmetry=( + torch.tensor([-1, -1, 0, 0], dtype=torch.int), + torch.tensor([False, True, True, False]), + ), + arrow=False, + ), + tat.Edge( + fermion=(True, False), + symmetry=( + torch.tensor([+1, +1, 0, 0], dtype=torch.int), + torch.tensor([True, False, False, True]), + ), + arrow=False, + ), + ), + ) + assert a.same_shape_with(b, allow_transpose=False) diff --git a/tests/test_create_tensor.py b/tests/test_create_tensor.py new file mode 100644 index 000000000..079609ffa --- /dev/null +++ b/tests/test_create_tensor.py @@ -0,0 +1,89 @@ +import torch +import tat + + +def test_create_tensor(): + a = tat.Tensor( + ( + "i", + "j", + ), + ( + tat.Edge(symmetry=(torch.tensor([False, False, True]),), fermion=(True,), arrow=True), + tat.Edge(symmetry=(torch.tensor([False, False, False, True, True]),), fermion=(True,), arrow=False), + ), + ) + assert a.rank == 2 + assert a.names == ("i", "j") + assert a.edges[0] == tat.Edge(symmetry=(torch.tensor([False, False, True]),), fermion=(True,), arrow=True) + assert a.edges[1] == tat.Edge(symmetry=(torch.tensor([False, False, False, True, True]),), + fermion=(True,), + arrow=False) + assert a.edges[0] == a.edge_by_name("i") + assert a.edges[1] == a.edge_by_name("j") + + +def test_tensor_get_set_item(): + a = tat.Tensor( + ( + "i", + "j", + ), + ( + tat.Edge(symmetry=(torch.tensor([False, False, True]),), fermion=(True,), arrow=True), + tat.Edge(symmetry=(torch.tensor([False, False, False, True, True]),), fermion=(True,), arrow=False), + ), + ) + a[{"i": 0, "j": 0}] = 1 + assert a[0, 0] == 1 + a["i":2, "j":3] = 2 + assert a[{"i": 2, "j": 3}] == 2 + a[2, 0] = 3 + assert a["i":2, "j":0] == 0 + + b = tat.Tensor( + ( + "i", + "j", + ), + ( + tat.Edge(symmetry=(torch.tensor([0, 0, -1]),), fermion=(False,)), + tat.Edge(symmetry=(torch.tensor([0, 0, 0, +1, +1]),), fermion=(False,)), + ), + ) + b[{"i": 0, "j": 0}] = 1 + assert b[0, 0] == 1 + b["i":2, "j":3] = 2 + assert b[{"i": 2, "j": 3}] == 2 + b[2, 0] = 3 + assert b["i":2, "j":0] == 0 + + +def test_create_randn_tensor(): + a = tat.Tensor( + ("i", "j"), + ( + tat.Edge(symmetry=(torch.tensor([False, True]),)), + tat.Edge(symmetry=(torch.tensor([False, True]),)), + ), + dtype=torch.float16, + ).randn_() + assert a.dtype == torch.float16 + assert a[0, 0] != 0 + assert a[1, 1] != 0 + assert a[0, 1] == 0 + assert a[1, 0] == 0 + + b = tat.Tensor( + ("i", "j"), + ( + tat.Edge(symmetry=(torch.tensor([False, False]), torch.tensor([0, -1]))), + tat.Edge(symmetry=(torch.tensor([False, False]), torch.tensor([0, +1]))), + ), + dtype=torch.float16, + ).randn_() + assert b.dtype == torch.float16 + assert b[0, 0] != 0 + assert b[1, 1] != 0 + assert b[0, 1] == 0 + assert b[1, 0] == 0 diff --git a/tests/test_edge.py b/tests/test_edge.py new file mode 100644 index 000000000..8bbe5c1df --- /dev/null +++ b/tests/test_edge.py @@ -0,0 +1,33 @@ +import torch +from tat import Edge + + +def test_create_edge_and_basic(): + a = Edge(dimension=5) + assert a.arrow == False + assert a.dimension == 5 + b = Edge(symmetry=(torch.tensor([False, False, True, True]),)) + assert b.arrow == False + assert b.dimension == 4 + c = Edge(fermion=(False, True), symmetry=(torch.tensor([False, True]), torch.tensor([False, True])), arrow=True) + assert c.arrow == True + assert c.dimension == 2 + + +def test_edge_conjugated_and_equal(): + a = Edge(fermion=(False, True), symmetry=(torch.tensor([False, True]), torch.tensor([0, 1])), arrow=True) + b = Edge(fermion=(False, True), symmetry=(torch.tensor([False, True]), torch.tensor([0, -1])), arrow=False) + assert a.conjugated() == b + assert a != 2 + + +def test_repr(): + a = Edge(fermion=(False, True), symmetry=(torch.tensor([False, True]), torch.tensor([0, 1])), arrow=True) + repr_a = repr(a) + assert repr_a == "Edge(dimension=2, arrow=True, fermion=(False, True), symmetry=(tensor([False, True]), tensor([0, 1])))" + b = Edge(symmetry=(torch.tensor([False, True]), torch.tensor([0, 1]))) + repr_b = repr(b) + assert repr_b == "Edge(dimension=2, symmetry=(tensor([False, True]), tensor([0, 1])))" + c = Edge(dimension=4) + repr_c = repr(c) + assert repr_c == "Edge(dimension=4)"