added tensor module

jetpack/__init__.py

@@ -9,11 +9,12 @@
 __version__ = '0.0.3'
 
 
-from . timepiece import *
-from . signals import *
-from . chart import *
-from . ionic import *
-from . capsules import *
+from .timepiece import *
+from .signals import *
+from .chart import *
+from .ionic import *
+from .capsules import *
+from .tensor import *
 
 __all__ = timepiece.__all__ + signals.__all__ + chart.__all__ + \
-    ionic.__all__ + capsules.__all__
+    ionic.__all__ + capsules.__all__ + tensor.__all__

jetpack/tensor.py

@@ -0,0 +1,70 @@
+"""
+Tensor
+------
+
+tensor unfolding
+
+"""
+
+import numpy as np
+
+__all__ = ['Tensor', 'UnfoldedTensor']
+
+
+class Tensor(np.ndarray):
+
+    def __new__(cls, arr):
+        obj = np.asarray(arr).view(cls)
+        return obj
+
+    def unfold(self, ax):
+        assert ax in range(self.ndim), "ax less than ndim"
+        orig_shape = self.shape
+        rolled_axes = pullax(list(range(self.ndim)), ax)
+        unfolded = self.transpose(rolled_axes).reshape(orig_shape[ax], -1)
+        return UnfoldedTensor(unfolded, ax, orig_shape)
+
+    @property
+    def norm(self):
+        return np.linalg.norm(self.ravel(), ord=2)
+
+
+class UnfoldedTensor(np.ndarray):
+
+    def __new__(cls, arr, axis, shape):
+        obj = np.asarray(arr).view(cls)
+        obj.orig_shape = shape
+        obj.axis = axis
+        return obj
+
+    @property
+    def norm(self):
+        return np.linalg.norm(self, ord='fro')
+
+    @property
+    def nucnorm(self):
+        return np.sum(np.linalg.svd(self, compute_uv=False))
+
+    def svt(self, threshold):
+        u, s, v = np.linalg.svd(self, full_matrices=False)
+        sthr = np.diag(np.maximum(s - threshold, 0))
+        return UnfoldedTensor(u.dot(sthr).dot(v), self.axis, self.orig_shape)
+
+    def __array_finalize__(self, obj):
+
+        if obj is None:
+            print('unfolded tensor None')
+            return
+
+        self.orig_shape = getattr(obj, 'orig_shape', None)
+        self.axis = getattr(obj, 'axis', None)
+
+    def fold(self):
+        rolled_axes = pullax(list(range(len(self.orig_shape))), self.axis)
+        folded = self.reshape(tuple(self.orig_shape[i] for i in rolled_axes))
+        return Tensor(folded.transpose(np.argsort(rolled_axes)))
+
+
+def pullax(values, idx):
+    values.insert(0, values.pop(idx))
+    return tuple(values)

tests/test_tensor.py

@@ -0,0 +1,37 @@
+"""
+Tests for the tensor module
+
+"""
+
+from jetpack.tensor import Tensor
+import numpy as np
+
+
+def test_unfolding():
+
+    t = Tensor(np.random.randn(3, 4, 5))
+
+    t0 = t.unfold(0)
+    t1 = t.unfold(1)
+    t2 = t.unfold(2)
+
+    assert t0.shape == (3, 20)
+    assert np.allclose(t, t0.fold())
+
+    assert t1.shape == (4, 15)
+    assert np.allclose(t, t1.fold())
+
+    assert t2.shape == (5, 12)
+    assert np.allclose(t, t2.fold())
+
+
+def test_norms():
+
+    # test Frobenius norm
+    t = Tensor(np.arange(12).reshape(2, 3, 2))
+    assert np.allclose(np.linalg.norm(np.arange(12)), t.norm)
+    assert np.allclose(np.linalg.norm(np.arange(12)), t.unfold(2).norm)
+
+    # test nuclear norm
+    t = Tensor(np.eye(5, 10).reshape(5, 2, 5))
+    assert np.allclose(t.unfold(0).nucnorm, 5.)