hiepy 0.1.0__py3-none-any.whl

hiepy/__init__.py

@@ -0,0 +1,117 @@
+"""
+hiepy - Hierarchical Tucker (HT) tensor library in PyTorch.
+
+Overview
+--------
+hiepy provides efficient representations and operations for tensors in the
+Hierarchical Tucker format, including algebraic operations, truncation,
+and conversions from dense PyTorch tensors.
+
+Main Component
+---------------
+HTTensor
+    Hierarchical Tucker tensor representation.
+
+Key Features
+-----------
+- Efficient tensor arithmetic (addition, contraction, etc.)
+- Rank truncation and orthogonalization utilities
+- Conversions from dense PyTorch tensors
+- Random tensor generation in HT format
+"""
+
+from importlib.metadata import version as _version
+
+from .algebra.add import add
+from .algebra.contract import contract
+from .algebra.ele_mode_mul import ele_mode_mul
+from .algebra.ele_mul import ele_mul
+from .algebra.mode_mul import mode_mul
+from .algebra.scalar_mul import scalar_mul
+from .algebra.sub import sub
+from .factory.arange import arange
+from .factory.constant import constant
+from .factory.copy import clone, copy, detach, requires_grad_, to
+from .factory.hilbert import hilbert
+from .factory.ones import ones
+from .factory.randn import randn
+from .factory.to_htt import to_htt
+from .factory.zeros import zeros
+from .httensor.full import full
+from .httensor.get_item import get
+from .httensor.getter import device, dtype, nbytes, order, rank, requires_grad, shape
+from .httensor.gramians import gramians
+from .httensor.httensor import HTTensor
+from .httensor.io import load, save
+from .linalg.left_svd_gramian import left_svd_gramian
+from .linalg.local_left_svd import local_left_svd
+from .linalg.matricize import dematricize, matricize
+from .linalg.mode_mul_torch import mode_mul_torch, multi_mode_mul_torch
+from .linalg.norm import norm
+from .linalg.vectorize import devectorize, vectorize
+from .transformation.orthog import orthog
+from .transformation.squeeze import squeeze
+from .transformation.truncate import truncate, truncate_gram, truncate_svd
+from .utils.data_types import ScalarType, Shape, TensorDict, TruncOpts
+from .utils.dimtree import DimTree, Node
+from .utils.trunc_rank import trunc_rank
+
+# Version of the installed package, from the single source in pyproject.toml
+__version__ = _version("hiepy")
+
+__all__ = [
+    "__version__",
+    "nbytes",
+    "dtype",
+    "device",
+    "requires_grad",
+    "order",
+    "rank",
+    "shape",
+    "gramians",
+    "HTTensor",
+    "full",
+    "get",
+    "load",
+    "save",
+    "truncate",
+    "truncate_gram",
+    "truncate_svd",
+    "orthog",
+    "squeeze",
+    "left_svd_gramian",
+    "local_left_svd",
+    "norm",
+    "TruncOpts",
+    "ScalarType",
+    "Shape",
+    "TensorDict",
+    "trunc_rank",
+    "DimTree",
+    "Node",
+    "matricize",
+    "dematricize",
+    "vectorize",
+    "devectorize",
+    "add",
+    "sub",
+    "mode_mul_torch",
+    "multi_mode_mul_torch",
+    "contract",
+    "ele_mul",
+    "ele_mode_mul",
+    "mode_mul",
+    "scalar_mul",
+    "zeros",
+    "hilbert",
+    "ones",
+    "randn",
+    "arange",
+    "constant",
+    "to_htt",
+    "clone",
+    "copy",
+    "detach",
+    "requires_grad_",
+    "to",
+]

hiepy/algebra/add.py

@@ -0,0 +1,271 @@
+"""HTTensor addition module.
+
+This module provides functionality to compute the formal sum of two HT tensors
+by embedding their transfer tensors and leaf matrices into block-diagonal structures.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from hiepy.httensor.httensor import HTTensor
+from hiepy.utils.data_types import TensorDict
+from hiepy.utils.dtype_device import check_device, result_dtype
+
+
+def add(x: HTTensor, y: HTTensor) -> HTTensor:
+    """
+    Adds two HT tensors via block-diagonal embedding.
+
+    Constructs an HT  tensor representing the exact sum of two HT tensors by embedding their
+    leaf matrices and transfer tensors into a block-diagonal HT representation.
+
+    The resulting tensor $Z = x + y$ is not computed in full tensor space, but
+    represented exactly in HT format as:
+
+    - Leaf matrices:
+        $U_i^Z = [U_i^x, U_i^y]$
+
+    - Transfer tensors:
+        $B_t^Z = diag(B_t^x, B_t^y)$
+
+    where the block-diagonal structure is applied at each node of the dimension tree.
+
+    Parameters
+    ----------
+    x : HTTensor
+        First HT tensor. Must share the same dimension tree and
+        mode sizes as ``y``.
+
+    y : HTTensor
+        Second HT tensor. Must share the same dimension tree and
+        mode sizes as ``x``.
+
+    Returns
+    -------
+    HTTensor
+        A new HT tensor representing the sum x + y in block-diagonal form.
+
+    Raises
+    ------
+    ValueError
+        If ``x`` and ``y`` do not share the same dimension tree or shape.
+
+    RuntimeError
+        If ``x`` and ``y`` live on different devices.
+
+    Notes
+    -----
+    - The result has the promoted dtype of ``x`` and ``y`` (e.g. ``float32`` and ``float64`` give ``float64``).
+    - The dimension tree is preserved.
+    - The tensor shape is preserved.
+    - Rank growth is additive:
+          $r_t^{out} = r_t^x + r_t^y$
+    - Orthogonality is not preserved.
+    - The result typically requires recompression.
+    - The arithmetic cost of the operation is negligible, as it primarily involves the
+      concatenation of leaf matrices and the creation of block-diagonal transfer tensors.
+
+    Examples
+    --------
+    >>> import hiepy
+    >>> import torch
+    >>> torch.set_default_dtype(torch.float64)
+    >>> x = hiepy.randn(shape=(8, 8, 8))
+    >>> y = hiepy.randn(shape=(8, 8, 8))
+    >>> sum_hiepy = hiepy.add(x,y)
+    >>> x_t = x.full()
+    >>> y_t = y.full()
+    >>> sum_torch = x_t + y_t
+    >>> torch.allclose(sum_torch, sum_hiepy.full())
+    True
+    """
+    # Validity checks
+    if not x.dtree == y.dtree:
+        raise ValueError("Dimension trees of 'x' and 'y' are not compatible")
+    if x.shape != y.shape:
+        raise ValueError(f"Shape of 'y' {y.shape} does not match shape of 'x' {x.shape}")
+    device = check_device(x=x, y=y)
+    dtype = result_dtype(x, y)
+    x, y = x.to(dtype=dtype), y.to(dtype=dtype)
+
+    U: TensorDict = {}
+    B: TensorDict = {}
+    dtree = x.dtree.copy()
+
+    # Concatenate leaf matrices
+    for leaf in dtree.get_leaves():
+        U[leaf] = torch.cat((x.U[leaf], y.U[leaf]), dim=1)
+
+    # Embed the transfer tensors in a 3D blockdiagonal way
+    for t in dtree.get_inner_nodes():
+        if x.dtree.is_root(t):
+            xl, xr, _ = x.B[t].shape  # rank of node x children
+            yl, yr, _ = y.B[t].shape  # rank of node y children
+            Bt = torch.zeros(xl + yl, xr + yr, 1, dtype=dtype, device=device)
+            Bt[:xl, :xr] = x.B[t]
+            Bt[xl:, xr:] = y.B[t]
+        else:
+            xl, xr, xk = x.B[t].shape  # rank of node children and node x (left, right, node)
+            yl, yr, yk = y.B[t].shape  # rank of node children and node y (left, right, node)
+            Bt = torch.zeros(xl + yl, xr + yr, xk + yk, dtype=dtype, device=device)
+            Bt[:xl, :xr, :xk] = x.B[t]
+            Bt[xl:, xr:, xk:] = y.B[t]
+        B[t] = Bt
+
+    # Create resulting HTTensor
+    res = HTTensor(U=U, B=B, dtree=dtree, is_orthog=False)
+    return res
+
+
+# trunc_sum is currently not maintained
+
+# def trunc_sum(summands: list[HTTensor], opts: HTTensorOpts) -> HTTensor:
+#     """Add HTTensors contained in summands and truncate sum on the fly."""
+#     # input validity checks
+#     if not isinstance(summands, list):
+#         raise TypeError(f"Argument 'summands' must be of type list, but got {type(summands)}.")
+
+#     if not all(isinstance(item, HTTensor) for item in summands):
+#         raise TypeError("Argument 'summands' must only contain HTTensor objects.")
+
+#     if len(summands) < 2:
+#         raise ValueError(f"Argument 'summands' must contain at least two elements, but got {len(summands)}.")
+
+#     if not all(summands[0].dtree.is_equal(item.dtree) for item in summands):
+#         raise ValueError(
+#             "All HTTensors in 'summands' must share the same dimension tree, "
+#             "but at least one tensor has a non-matching tree."
+#         )
+#     check_opts(opts)
+
+#     # Copy first summand if required
+#     U, B, dtree, opts = {}, {}, deepcopy(summands[0].dtree), deepcopy(opts)
+
+#     # Anpassen der Fehlertoleranzen in opts
+#     # Soll global der Fehler e eingehalten werden, muss der Kuerzungsfehler pro Knoten
+#     # kleiner gleich e / sqrt((Tensorordnung * 2 - 2)) bleiben
+#     adjusted_opts = deepcopy(opts)
+#     if "atol" in adjusted_opts:
+#         adjusted_opts["atol"] = adjusted_opts["atol"] / (x.order * 2 - 3) ** 0.5
+#     if "rtol" in adjusted_opts:
+#         adjusted_opts["rtol"] = adjusted_opts["rtol"] / (x.order * 2 - 3) ** 0.5
+
+#     # Berechne die reduzierten Gram'schen Matrizen der impliziten Summe
+#     G = gramians_sum(summands)
+
+#     # Berechne gekuerzte Blattmatrizen der impliziten Summe
+#     # Update dabei on the fly den Transfertensor des Elternknotens
+#     B_upd = {node.dims: {} for node in dtree.get_inner_nodes()}
+#     for leaf in dtree.get_leaves():
+#         dim = leaf.dims
+#         # Konkatenieren der Blattmatrizen
+#         U_cat = torch.hstack(tuple(item.U[dim] for item in summands))
+#         # QR Zerlegung
+#         Q, R = torch.linalg.qr(U_cat, mode="reduced")
+#         # Update reduzierte Gram'sche Matrix
+#         G_upd = R @ G[dim] @ R.T
+#         # Berechne davon linke Singulaervektoren
+#         S, sv = left_svd_gramian(G_upd)
+#         # Berechne basierend auf opts wie viele Spalten behalten werden
+#         if adjusted_opts:
+#             rank = trunc_rank(sv, **adjusted_opts)
+#         else:
+#             rank = len(sv)
+#         S = S[:, :rank]
+#         # Berechne schliesslich finale Blattmatrix
+#         U[dim] = Q @ S
+
+#         # Update Transfertensor des Elternknotens
+#         # Adaptiere R entsprechend der Kuerzung
+#         R = S.T @ R
+#         # Teile dafuer zunaechst R in R=[R1 | R2 | ... | Rn]
+#         # wobei die Spaltenanzahl von Ri durch den entsprechenden Rang des i-ten Summanden
+#         # vorgegeben ist
+#         R = torch.split(R, [item.rank[dim] for item in summands], dim=1)
+#         # Multipliziere nun R[i] in den Transfertensor des Elternknotens des i-ten Summanden
+#         assert leaf.father is not None
+#         fdims = leaf.father.dims
+#         for i in range(len(summands)):
+#             if leaf.is_left():
+#                 if i in B_upd[fdims]:
+#                     B_upd[fdims][i] = torch.tensordot(R[i], B_upd[fdims][i], dims=([1], [0]))  # type: ignore
+#                 else:
+#                     B_upd[fdims][i] = torch.tensordot(R[i], summands[i].B[fdims], dims=([1], [0]))  # type: ignore
+#             else:
+#                 if i in B_upd[fdims]:
+#                     B_upd[fdims][i] = torch.tensordot(R[i], B_upd[fdims][i], dims=([1], [1]))  # type: ignore
+#                     B_upd[fdims][i] = torch.movedim(B_upd[fdims][i], source=0, destination=1)
+#                 else:
+#                     B_upd[fdims][i] = torch.tensordot(R[i], summands[i].B[fdims], dims=([1], [1]))  # type: ignore
+#                     B_upd[fdims][i] = torch.movedim(B_upd[fdims][i], source=0, destination=1)
+
+#     # Berechne nun gekuerzte Transfertensoren der impliziten Summe
+#     # Durschreite den Dimensionsbaum dazu bottom-up
+#     for level in dtree.levels(order="bottom-up"):
+#         for node in level:
+#             if node.is_leaf() or node.is_root():
+#                 continue
+#             dims = node.dims
+#             # Konkateniere die aktualisierten Transfertensoren der Summanden
+#             B_cat = torch.cat([B_upd[dims][i] for i in range(len(summands))], dim=2)
+#             # Berechne QR Zerlegung der Matrizierung davon
+#             Q, R = torch.linalg.qr(matricise(B_cat, t=(0, 1)), mode="reduced")
+#             # Dematriziere Q zu 3D Transfertensor
+#             Q = dematricise(Q, shape=(B_cat.shape[0], B_cat.shape[1], Q.shape[1]), t=(0, 1))
+
+#             # Aktualisiere reduzierte Gram'sche Matrix
+#             G_upd = R @ G[dims] @ R.T
+#             # Berechne davon linke Singulaervektoren
+#             S, sv = left_svd_gramian(G_upd)
+#             # Berechne basierend auf opts wie viele Spalten behalten werden
+#             if adjusted_opts:
+#                 rank = trunc_rank(sv, **adjusted_opts)
+#             else:
+#                 rank = len(sv)
+#             S = S[:, :rank]
+#             # Berechne schliesslich finalen Transfertensor
+#             B[dims] = torch.tensordot(Q, S, dims=([2], [0]))  # type: ignore
+
+#             # Update Transfertensor des Elternknotens
+#             # Adaptiere R entsprechend der Kuerzung
+#             R = S.T @ R
+#             # Teile dafuer zunaechst R in R=[R1 | R2 | ... | Rn]
+#             # wobei die Spaltenanzahl von Ri durch den entsprechenden Rang des i-ten Summanden
+#             # vorgegeben ist
+#             R = torch.split(R, [item.rank[dims] for item in summands], dim=1)
+#             # Multipliziere nun R[i] in den Transfertensor des Elternknotens des i-ten Summanden
+#             assert node.father is not None
+#             fdims = node.father.dims
+#             for i in range(len(summands)):
+#                 if node.is_left():
+#                     if i in B_upd[fdims]:
+#                         B_upd[fdims][i] = torch.tensordot(R[i], B_upd[fdims][i], dims=([1], [0]))  # type: ignore
+#                     else:
+#                         B_upd[fdims][i] = torch.tensordot(
+#                             R[i],
+#                             summands[i].B[fdims],
+#                             dims=([1], [0]),  # type: ignore
+#                         )
+#                 else:
+#                     if i in B_upd[fdims]:
+#                         B_upd[fdims][i] = torch.tensordot(R[i], B_upd[fdims][i], dims=([1], [1]))  # type: ignore
+#                         B_upd[fdims][i] = torch.movedim(B_upd[fdims][i], source=0, destination=1)
+#                     else:
+#                         B_upd[fdims][i] = torch.tensordot(
+#                             R[i],
+#                             summands[i].B[fdims],
+#                             dims=([1], [1]),  # type: ignore
+#                         )
+#                         B_upd[fdims][i] = torch.movedim(B_upd[fdims][i], source=0, destination=1)
+#             # Gebe B_upd[node] frei
+#             del B_upd[dims]
+
+#     # Wurzelbehandlung..
+#     B_root = torch.sum(torch.concat([item for item in B_upd[dtree.root.dims].values()], dim=2), dim=2)
+#     B[dtree.root.dims] = B_root.reshape(B_root.shape[0], B_root.shape[1], 1)
+
+#     # Erstelle HTucker Tensor der Summe
+#     z = HTTensor(U=U, B=B, dtree=dtree, is_orthog=False)
+
+#     return z

hiepy/algebra/contract.py

@@ -0,0 +1,101 @@
+"""HTTensor contraction module.
+
+This module provides functionality to compute the full tensor contraction
+(inner product) of two HT tensors sharing the same dimension tree.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from hiepy.httensor.httensor import HTTensor
+from hiepy.utils.dtype_device import check_device, result_dtype
+
+
+def contract(x: HTTensor, y: HTTensor) -> torch.Tensor:
+    r"""
+    Computes the full inner product of two HT tensors.
+
+    This operation evaluates the scalar contraction
+        $\langle X, Y \rangle = \sum_{i_1,\dots,i_d} X_{i_1\dots i_d} Y_{i_1\dots i_d}$
+    without unfolding the full tensors, by performing a recursive contraction along the dimension tree.
+
+    Parameters
+    ----------
+    x : HTTensor
+        First HT tensor. Must share the same dimension tree and mode sizes as ``y``.
+
+    y : HTTensor
+        Second HT tensor. Must share the same dimension tree and mode sizes as ``x``.
+
+    Returns
+    -------
+    torch.Tensor
+        Scalar tensor representing the full inner product of ``x``and ``y``.
+
+    Raises
+    ------
+    ValueError
+        If the dimension trees of ``x`` and ``y`` differ.
+
+    ValueError
+        If the tensor shapes of ``x`` and ``y`` do not match.
+
+    RuntimeError
+        If ``x`` and ``y`` live on different devices.
+
+    Notes
+    -----
+    - The result has the promoted dtype of ``x`` and ``y`` (e.g. ``float32`` and ``float64`` give ``float64``).
+    - The operation's stability is improved when ``x`` and ``y`` are orthogonal.
+    - No recompression is required, as the result is scalar-valued.
+    - The arithmetic cost of the operation scales as $\mathcal{O}(dnr_xr_y + dr_x^2r_y^2)$,
+      where $d$ is the tensor order, $n$ the maximum mode size, and $r_x$, $r_y$ the maximum
+      HT ranks of ``x`` and ``y``, respectively.
+
+    Examples
+    --------
+    >>> import torch
+    >>> import hiepy
+    >>> torch.set_default_dtype(torch.float64)
+    >>> shape = 8,8,8,8
+    >>> x = torch.randn(shape)
+    >>> y = torch.randn(shape)
+    >>> prod_torch = torch.sum(x*y)
+    >>> opts = hiepy.TruncOpts({"atol":1e-10})
+    >>> x_ht = hiepy.to_htt(x,opts)
+    >>> y_ht = hiepy.to_htt(y,opts)
+    >>> prod_ht = hiepy.contract(x_ht,y_ht)
+    >>> torch.allclose(prod_torch, prod_ht)
+    True
+    """
+    # Validity checks
+    if not x.dtree == y.dtree:
+        raise ValueError("Dimension trees of 'x' and 'y' are not compatible")
+    if x.shape != y.shape:
+        raise ValueError(f"Shape of 'y' {y.shape} does not match shape of 'x' {x.shape}")
+    check_device(x=x, y=y)
+    dtype = result_dtype(x, y)
+    x, y = x.to(dtype=dtype), y.to(dtype=dtype)
+
+    dtree = x.dtree
+    Ux, Uy = x.U, y.U
+    Bx, By = x.B, y.B
+    M = {}
+
+    # Contract leaf matrices
+    for leaf in dtree.get_leaves():
+        M[leaf] = Ux[leaf].T @ Uy[leaf]
+
+    # Contract M and transfer tensors bottom-up along the tree
+    for level in dtree.get_levels(order="bottom-up"):
+        for t in level:
+            if dtree.is_leaf(t):
+                continue
+            left, right = dtree.get_children(t)
+            M_ = torch.tensordot(M[left], By[t], dims=([1], [0]))  # type:ignore
+            M_ = torch.tensordot(Bx[t], M_, dims=([0], [0]))  # type:ignore
+            M_ = torch.tensordot(M_, M[right], dims=([0, 2], [0, 1]))  # type:ignore
+            M[t] = M_
+
+    return M[dtree.root][0, 0]

hiepy/algebra/ele_mode_mul.py

@@ -0,0 +1,113 @@
+"""
+Implements the elementwise mode-product of a vector with an HT  tensor.
+
+This module provides functionality to compute the elementwise mode-product of a vector
+and an HT  tensor.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from hiepy.httensor.httensor import HTTensor as HTTensor
+from hiepy.utils.dtype_device import check_device, result_dtype
+
+
+def ele_mode_mul(x: HTTensor, v: torch.Tensor, dim: int) -> HTTensor:
+    r"""
+    Scales an HT  tensor ``x`` along its mode ``dim`` by a vector ``v``.
+
+    The elementwise mode multiplication of a tensor $X\in\mathbb{R}^{n_0\times\dots\times n_{d-1}}$ with a vector
+    $v \in \mathbb{R}^{n_{\mu}}$ along mode $\mu \in \{0,\dots,d-1\}$ is defined as:
+
+    $$(X \star_{\mu} v)_{i_0,\dots,i_{d-1}} = X_{i_0,\dots,i_{d-1}} \cdot v_{i_\mu}$$
+
+
+    Parameters
+    ----------
+    x : HTTensor
+        Input HT tensor.
+
+    v : torch.Tensor
+        One-dimensional vector defining the mode-wise scaling factors.
+        Its length must coincide with the size of mode ``dim``.
+
+    dim : int
+        Mode along which ``x`` is scaled by ``v``.
+
+    Returns
+    -------
+    HTTensor
+        A new HT tensor representing the scaled tensor.
+
+    Raises
+    ------
+    ValueError
+        If ```v``` is not one-dimensional.
+
+    ValueError
+        If ```dim``` is outside the valid mode range.
+
+    ValueError
+        If the length of ```v``` does not match the size of the selected mode.
+
+    RuntimeError
+        If ``x`` and ``v`` live on different devices.
+
+    Notes
+    -----
+    - The result has the promoted dtype of ``x`` and ``v`` (e.g. ``float32`` and ``float64`` give ``float64``).
+    - The dimension tree is preserved.
+    - Hierarchical ranks are preserved.
+    - Orthogonality is not preserved.
+    - Computational cost is linear in mode size and corresponding rank, $\mathcal{O}(nr)$.
+    - All transfer tensors remain unchanged.
+    - All but one leaf matrices remain unchanged.
+
+    Examples
+    --------
+    >>> import torch
+    >>> import hiepy
+    >>> torch.set_default_dtype(torch.float64)
+    >>> x_ht = hiepy.arange((2, 3, 4))
+    >>> x_ht.full()
+    tensor([[[ 0.,  1.,  2.,  3.],
+             [ 4.,  5.,  6.,  7.],
+             [ 8.,  9., 10., 11.]],
+    <BLANKLINE>
+            [[12., 13., 14., 15.],
+             [16., 17., 18., 19.],
+             [20., 21., 22., 23.]]])
+    >>> v = torch.arange(3)
+    >>> hiepy.ele_mode_mul(x_ht,v,dim=1).full()
+    tensor([[[ 0.,  0.,  0.,  0.],
+             [ 4.,  5.,  6.,  7.],
+             [16., 18., 20., 22.]],
+    <BLANKLINE>
+            [[ 0.,  0.,  0.,  0.],
+             [16., 17., 18., 19.],
+             [40., 42., 44., 46.]]])
+    """
+    # Validity checks
+    if v.ndim != 1:
+        raise ValueError(f"Argument 'v' must be a 1D torch.Tensor, got shape {v.shape}")
+    if dim < 0 or dim >= x.order:
+        raise ValueError(f"Argument 'dim' is out of range for this HTTensor (dim={dim}, order={x.order})")
+    if x.shape[dim] != v.shape[0]:
+        raise ValueError(
+            f"Mismatch between dimension {dim} of x (size {x.shape[dim]}) and length of v (size {v.shape[0]})."
+        )
+
+    check_device(x=x, v=v)
+    dtype = result_dtype(x, v)
+    x = x.clone().to(dtype=dtype)
+    v = v.to(dtype=dtype)
+
+    # Update leaf matrix corresponding to dim
+    node = (dim,)
+    x.U[node] = x.U[node] * (v[:, None])
+
+    # Update is_orthog
+    x.is_orthog = False
+
+    return x