Teneye: Add preliminary implementation and tests (#222)

* Teneye: Add preliminary implementation and tests

* Black: Re-run black
* Added to local test suite in other PR

---------

Co-authored-by: Danny Dunlavy <[email protected]>

Author: ntjohnson1

pyttb/__init__.py

@@ -24,7 +24,7 @@
 from pyttb.symktensor import symktensor
 from pyttb.symtensor import symtensor
 from pyttb.tenmat import tenmat
-from pyttb.tensor import tendiag, tenones, tenrand, tensor, tenzeros
+from pyttb.tensor import tendiag, teneye, tenones, tenrand, tensor, tenzeros
 from pyttb.ttensor import ttensor
 from pyttb.tucker_als import tucker_als
 
@@ -57,6 +57,7 @@ def ignore_warnings(ignore=True):
     sumtensor.__name__,
     symktensor.__name__,
     symtensor.__name__,
+    teneye.__name__,
     tenmat.__name__,
     tendiag.__name__,
     tenones.__name__,

pyttb/tensor.py

@@ -7,7 +7,7 @@
 
 import logging
 from collections.abc import Iterable
-from itertools import permutations
+from itertools import combinations_with_replacement, permutations
 from math import factorial
 from typing import Any, Callable, List, Optional, Tuple, Union
 
@@ -2553,6 +2553,53 @@ def tendiag(elements: np.ndarray, shape: Optional[Tuple[int, ...]] = None) -> te
     return X
 
 
+def teneye(order: int, size: int) -> tensor:
+    """Create identity tensor of specified shape.
+
+    T is an "identity tensor if T.ttsv(x, skip_dim=0) = x for all x such that
+    norm(x) == 1.
+
+    An identity tensor only exists if order is even.
+    This method is resource intensive
+    for even moderate orders or sizes (>=6).
+
+    Parameters
+    ----------
+    order: Number of dimensions of tensor.
+    size: Number of elements in any dimension of the tensor.
+
+    Examples
+    --------
+    >>> ttb.teneye(2, 3)
+    tensor of shape (3, 3)
+    data[:, :] =
+    [[1. 0. 0.]
+     [0. 1. 0.]
+     [0. 0. 1.]]
+    >>> x = np.ones((5,))
+    >>> x /= np.linalg.norm(x)
+    >>> T = ttb.teneye(4, 5)
+    >>> np.allclose(T.ttsv(x, 0), x)
+    True
+
+    Returns
+    -------
+    Identity tensor.
+    """
+    if order % 2 != 0:
+        raise ValueError(f"Order must be even but received {order}")
+    idx_iterator = combinations_with_replacement(range(size), order)
+    A = tenzeros((size,) * order)
+    s = np.zeros((factorial(order), order // 2))
+    for _i, indices in enumerate(idx_iterator):
+        p = np.array(list(permutations(indices)))
+        for j in range(order // 2):
+            s[:, j] = p[:, 2 * j - 1] == p[:, 2 * j]
+        v = np.sum(np.sum(s, axis=1) == order // 2)
+        A[tuple(zip(*p))] = v / factorial(order)
+    return A
+
+
 def mttv_left(W_in: np.ndarray, U1: np.ndarray) -> np.ndarray:
     """
     Contract leading mode in partial MTTKRP W_in using factor matrix U1.

tests/test_tensor.py

@@ -1708,6 +1708,17 @@ def test_tendiag():
         assert X[diag_index] == i
 
 
+def test_teneye():
+    with pytest.raises(ValueError):
+        ttb.teneye(1, 0)
+    size = 5
+    order = 4
+    T = ttb.teneye(order, size)
+    x = np.random.random((size,))
+    x = x / np.linalg.norm(x)
+    np.testing.assert_almost_equal(T.ttsv(x, 0), x)
+
+
 def test_mttv_left():
     m1 = 2
     mi = [range(1, 4)]