Nick/tensor scale (#221)

* CP_ALS: Should work with TTensor with no changes
* Update interface for mypy to check
* Update simple tests to confirm

* Tensor Scale:
* Add scale and tests
* Didn't add exclude dims, but could to match MATLAB
* Skipped the function handle since that seems redudant with other functionality we have

* Tensor Scale:
* Forgot negative test

* Tensor Scale:
* Typo in my TODO

Author: ntjohnson1

pyttb/cp_als.py

@@ -12,7 +12,7 @@
 
 
 def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
-    input_tensor: Union[ttb.tensor, ttb.sptensor],
+    input_tensor: Union[ttb.tensor, ttb.sptensor, ttb.ttensor],
     rank: int,
     stoptol: float = 1e-4,
     maxiters: int = 1000,

pyttb/sptensor.py

@@ -1055,14 +1055,18 @@ def reshape(
             tuple(np.concatenate((keep_shape, new_shape))),
         )
 
-    def scale(self, factor: np.ndarray, dims: Union[float, np.ndarray]) -> sptensor:
+    def scale(
+        self,
+        factor: Union[np.ndarray, ttb.tensor, ttb.sptensor],
+        dims: Union[float, np.ndarray],
+    ) -> sptensor:
         """
         Scale along specified dimensions for sparse tensors
 
         Parameters
         ----------
-        factor:
-        dims:
+        factor: Scaling factor
+        dims: Dimensions to scale
 
         Returns
         -------

pyttb/tensor.py

@@ -401,7 +401,7 @@ def find(self) -> Tuple[np.ndarray, np.ndarray]:
          [3 4]]
         >>> T_threshold = T > 2
         >>> subs, vals = T_threshold.find()
-        >>> subs
+        >>> subs.astype(int)
         array([[1, 0],
                [1, 1]])
         >>> vals
@@ -1019,6 +1019,73 @@ def reshape(self, shape: Tuple[int, ...]) -> tensor:
 
         return ttb.tensor(np.reshape(self.data, shape, order="F"), shape)
 
+    def scale(
+        self,
+        factor: Union[np.ndarray, ttb.tensor],
+        dims: Union[float, np.ndarray],
+    ) -> tensor:
+        """
+        Scale along specified dimensions for tensors.
+
+        Parameters
+        ----------
+        factor: Scaling factor
+        dims: Dimensions to scale
+
+        Returns
+        -------
+        Scaled Tensor.
+
+        Examples
+        --------
+        >>> T = ttb.tenones((3, 4, 5))
+        >>> S = np.arange(5)
+        >>> Y = T.scale(S, 2)
+        >>> Y.data[0, 0, :]
+        array([0., 1., 2., 3., 4.])
+        >>> S = ttb.tensor(np.arange(5))
+        >>> Y = T.scale(S, 2)
+        >>> Y.data[0, 0, :]
+        array([0., 1., 2., 3., 4.])
+        >>> S = ttb.tensor(np.arange(12), shape=(3, 4))
+        >>> Y = T.scale(S, [0, 1])
+        >>> Y.data[:, :, 0]
+        array([[ 0.,  3.,  6.,  9.],
+               [ 1.,  4.,  7., 10.],
+               [ 2.,  5.,  8., 11.]])
+        """
+        if isinstance(dims, list):
+            dims = np.array(dims)
+        elif isinstance(dims, (float, int, np.generic)):
+            dims = np.array([dims])
+
+        # TODO update tt_dimscheck overload so I don't need explicit
+        #   Nones to appease mypy
+        dims, _ = tt_dimscheck(self.ndims, None, dims, None)
+        remdims = np.setdiff1d(np.arange(0, self.ndims), dims)
+
+        if not np.array_equal(factor.shape, np.array(self.shape)[dims]):
+            raise ValueError(
+                f"Scaling factor has shape {factor.shape}, but dimensions "
+                f"to scale had shape {np.array(self.shape)[dims]}"
+            )
+        if isinstance(factor, np.ndarray):
+            if len(factor.shape) == 1:
+                factor = factor[:, None]
+            factor = ttb.tensor(factor, copy=False)
+        # TODO this should probably be doable directly as a numpy view
+        #   where I think this is currently a copy
+        vector_factor = ttb.tenmat.from_tensor_type(
+            factor, np.arange(factor.ndims)
+        ).double()
+        vector_self = ttb.tenmat.from_tensor_type(self, dims, remdims).double()
+        # Numpy broadcasting should be equivalent to bsxfun
+        result = vector_self * vector_factor
+        # TODO why do we need this transpose for things to work?
+        if len(dims) == 1:
+            result = result.transpose()
+        return ttb.tenmat.from_data(result, dims, remdims, self.shape).to_tensor()
+
     def squeeze(self) -> Union[tensor, np.ndarray, float]:
         """
         Removes singleton dimensions from the tensor.

tests/test_cp_als.py

@@ -27,12 +27,37 @@ def sample_sptensor():
     return data, sptensorInstance
 
 
+@pytest.fixture()
+def sample_ttensor():
+    """Simple TTENSOR to verify by hand"""
+    core = ttb.tensor(np.ones((2, 3)))
+    factors = [
+        np.ones((5, 2)),
+        np.ones((2, 3)),
+    ]
+    ttensorInstance = ttb.ttensor(core, factors)
+    return ttensorInstance
+
+
+@pytest.fixture()
+def random_ttensor():
+    """Arbitrary TTENSOR to verify consistency between alternative operations"""
+    core = ttb.tensor(np.random.random((2, 3, 4)))
+    factors = [
+        np.random.random((5, 2)),
+        np.random.random((2, 3)),
+        np.random.random((4, 4)),
+    ]
+    ttensorInstance = ttb.ttensor(core, factors)
+    return ttensorInstance
+
+
 @pytest.mark.indevelopment
 def test_cp_als_tensor_default_init(capsys, sample_tensor):
     (data, T) = sample_tensor
     (M, Minit, output) = ttb.cp_als(T, 2)
     capsys.readouterr()
-    assert pytest.approx(output["fit"], 1) == 0
+    assert pytest.approx(output["fit"]) == 1
 
 
 @pytest.mark.indevelopment
@@ -49,7 +74,7 @@ def test_cp_als_tensor_ktensor_init(capsys, sample_tensor):
     KInit = ttb.ktensor.from_function(np.random.random_sample, T.shape, 2)
     (M, Minit, output) = ttb.cp_als(T, 2, init=KInit)
     capsys.readouterr()
-    assert pytest.approx(output["fit"], 1) == 0
+    assert pytest.approx(output["fit"]) == 1
 
 
 @pytest.mark.indevelopment
@@ -76,7 +101,7 @@ def test_cp_als_sptensor_default_init(capsys, sample_sptensor):
     (data, T) = sample_sptensor
     (M, Minit, output) = ttb.cp_als(T, 2)
     capsys.readouterr()
-    assert pytest.approx(output["fit"], 1) == 0
+    assert pytest.approx(output["fit"]) == 1
 
 
 @pytest.mark.indevelopment
@@ -93,7 +118,26 @@ def test_cp_als_sptensor_ktensor_init(capsys, sample_sptensor):
     KInit = ttb.ktensor.from_function(np.random.random_sample, T.shape, 2)
     (M, Minit, output) = ttb.cp_als(T, 2, init=KInit)
     capsys.readouterr()
-    assert pytest.approx(output["fit"], 1) == 0
+    assert pytest.approx(output["fit"]) == 1
+
+
+@pytest.mark.indevelopment
+def test_cp_als_ttensor_default_init(capsys, sample_ttensor):
+    T = sample_ttensor
+    (M, Minit, output) = ttb.cp_als(T, 1)
+    capsys.readouterr()
+    assert pytest.approx(output["fit"]) == 1
+
+
+@pytest.mark.indevelopment
+def test_cp_als_ttensor_default_init_consistency(capsys, random_ttensor):
+    T = random_ttensor
+    KInit = ttb.ktensor.from_function(np.random.random_sample, T.shape, 2)
+    _, _, output = ttb.cp_als(T, 2, init=KInit)
+    capsys.readouterr()
+    _, _, dense_output = ttb.cp_als(T.full(), 2, init=KInit)
+    capsys.readouterr()
+    assert pytest.approx(output["fit"]) == dense_output["fit"]
 
 
 @pytest.mark.indevelopment
@@ -102,19 +146,15 @@ def test_cp_als_tensor_dimorder(capsys, sample_tensor):
 
     # default dimorder
     dimorder = [i for i in range(T.ndims)]
-    print(dimorder)
-    print(dimorder.__class__)
     (M, Minit, output) = ttb.cp_als(T, 2, dimorder=dimorder)
     capsys.readouterr()
-    assert pytest.approx(output["fit"], 1) == 0
+    assert pytest.approx(output["fit"]) == 1
 
     # reverse should work
     dimorder = [T.ndims - i - 1 for i in range(T.ndims)]
-    print(dimorder)
-    print(dimorder.__class__)
     (M, Minit, output) = ttb.cp_als(T, 2, dimorder=dimorder)
     capsys.readouterr()
-    assert pytest.approx(output["fit"], 1) == 0
+    assert pytest.approx(output["fit"]) == 1
 
     # dimorder not a list
     with pytest.raises(AssertionError) as excinfo:

tests/test_tensor.py

@@ -1090,6 +1090,30 @@ def test_tensor_mask(sample_tensor_2way):
     assert "Mask cannot be bigger than the data tensor" in str(excinfo)
 
 
+def test_tensor_scale():
+    T = ttb.tenones((3, 4, 5))
+    S = np.arange(5, dtype=float)
+    Y = T.scale(S, 2)
+    assert np.array_equal(Y.data[0, 0, :], S)
+
+    S = ttb.tensor(np.arange(5, dtype=float))
+    Y = T.scale(S, 2)
+    assert np.array_equal(Y.data[0, 0, :], S.data)
+
+    S = ttb.tensor(np.arange(12, dtype=float), shape=(3, 4))
+    Y = T.scale(S, [0, 1])
+    assert np.array_equal(Y.data[:, :, 0], S.data)
+
+    S = ttb.tensor(np.arange(60, dtype=float), shape=(3, 4, 5))
+    Y = T.scale(S, [0, 1, 2])
+    assert np.array_equal(Y.data, S.data)
+
+    # Negative test
+    with pytest.raises(ValueError):
+        S = ttb.tensor(np.arange(60, dtype=float), shape=(3, 4, 5))
+        Y = T.scale(S, 0)
+
+
 def test_tensor_squeeze(sample_tensor_2way):
     (params, tensorInstance) = sample_tensor_2way