tred

 90def tsvdm(
 91    A,
 92    M=None,
 93    Minv=None,
 94    *,
 95    keep_hats=False,
 96    full_frontal_slices=True,
 97    svals_matrix_form=False,
 98):
 99    """Return the t-SVDM decomposition from Kilmer et al. (2021).
100
101    This is a modified version, inspired by the implementation at
102    https://github.com/UriaMorP/mprod_package
103
104    *NOTE: For now, unlike some other implementations (Numpy, Scipy), we
105    will return the tensor $V$ NOT $V^T$.*
106
107    Parameters
108    ----------
109    A : ndarray of shape (n, p, t)
110        Data tensor
111
112    M : Callable[[ndarray], ndarray] or None, default=None
113        A function which expects an order-3 tensor as input, and returns the
114        image under a m-transform. If unspecified TPCA will use the Discrete
115        Cosine Transform (ii) from `scipy.fft`.
116
117    MInv : Callable[[ndarray], ndarray] or None, default=None
118        A function implementing the inverse transform of `M`.
119
120    keep_hats : bool, default=False
121        Setting to `True` will return the tSVDM factors in the m-transform
122        space, under the specified `M`
123
124    full_frontal_slices : bool, default=True
125        To reconstruct `A`, one only needs the first $k$ columns of
126        $U$ and $V$.
127
128        Setting this to False will return the truncated tensors.
129        See: https://numpy.org/doc/stable/reference/generated/numpy.linalg.svd.html
130
131    svals_matrix_form : bool, default=False
132        Setting to `True` will return a compressed version of $S$, where the
133        singular values of each f-diagonal frontal slice becomes the column of
134        a matrix, with `t` columns total.
135
136    Returns
137    -------
138    U_tens : ndarray of shape (n, n, t)
139        If `full_frontal_slices==False` shape is (n, k, t) instead
140
141    S_tens : ndarray of shape (n, p, t)
142        If `full_frontal_slices==False` shape is (k, k, t) instead
143        If `svals_matrix_form==True`, `S_mat` of shape (k, t) returned instead
144
145    V_tens : ndarray of shape (p, p, t)
146        If `full_frontal_slices==False` shape is (p, k, t) instead
147
148    References
149    ----------
150    Kilmer, M.E., Horesh, L., Avron, H. and Newman, E., 2021. Tensor-tensor
151    algebra for optimal representation and compression of multiway data.
152    Proceedings of the National Academy of Sciences, 118(28), p.e2015851118.
153    """
154
155    assert len(A.shape) == 3, "Ensure order-3 tensor input"
156    assert not (
157        callable(M) ^ callable(Minv)
158    ), "If explicitly defined, both M and its inverse must be defined"
159
160    if not callable(M):  # and Minv is not defined - guaranteed by assertion
161        M, Minv = generate_default_m_transform_pair(A.shape[-1])
162
163    # transform the tensor to new space via the mode-3 product
164    hatA = M(A)
165
166    # an appropriate transposition allows Numpys array broadcasting to do
167    # facewise svd's S_mat contains the singular values per matrix in the
168    # input stack of matrices (the transpose tensor stacks top to bottom,
169    # with t slices of size n by p)
170    U_stack, S_mat, Vt_stack = np.linalg.svd(
171        hatA.transpose(2, 0, 1), full_matrices=full_frontal_slices
172    )
173
174    hatU = U_stack.transpose(1, 2, 0)
175    S_mat = S_mat.transpose()
176    # the following is a call to .transpose(1, 2, 0) followed by a facewise
177    # transpose defined by .transpose(1, 0, 2)
178    hatV = Vt_stack.transpose(2, 1, 0)
179
180    # if we are transforming scipy's singular values matrix back into tensor
181    # form, make sure we use the correct dimensions corresponding to whether
182    # or not the tensor faces were truncated during svd
183    if not svals_matrix_form:
184        if full_frontal_slices:
185            desired_S_tens_shape = A.shape
186        else:
187            n, p, t = A.shape
188            k = min(n, p)
189            desired_S_tens_shape = (k, k, t)
190
191    if keep_hats:
192        return (
193            hatU,
194            # by default return S as n,p,t f-diagonal tensor (or) convert into
195            # compressed matrix of singular values of shape (k,t)
196            S_mat
197            if svals_matrix_form
198            else _singular_vals_mat_to_tensor(S_mat, *desired_S_tens_shape),
199            hatV,
200        )
201    else:
202        return (
203            Minv(hatU),
204            # by default return S as n,p,t f-diagonal tensor (or) convert into
205            # compressed matrix of singular values of shape (k,t)
206            Minv(S_mat)
207            if svals_matrix_form
208            else _singular_vals_mat_to_tensor(Minv(S_mat), *desired_S_tens_shape),
209            Minv(hatV),
210        )

 25class TPCA(ClassNamePrefixFeaturesOutMixin, TransformerMixin, BaseEstimator):
 26    """Tensor analogue of PCA, introduced by Mor et al. (2022).
 27
 28    t-SVDM tensor analogue of PCA using explicit rank truncation with explicit
 29    rank truncation from Mor et al. (2022), and underlying m-product framework
 30    from Kilmer et al. (2021). Takes in an $n \\times p \\times t$ input
 31    tensor, and transforms into a $n \\times$ `n_components` matrix of 2D
 32    transformed projections.
 33
 34    The input tensor is centred into Mean Deviation Form (by location), but
 35    not normalized (by scale).
 36
 37    Parameters
 38    ----------
 39    n_components : int, float, or None, default=None
 40        Control number of components to keep. If n_components is not set at
 41        all, or set to `None`, ``n_components == min(n, p) * t``
 42
 43        If `n_components` is an integer, the TPCA will return the number of
 44        loadings, provided that for the tensor data passed into `fit`,
 45        satisfies ``1 <= n_components <= min(n, p) * t``
 46
 47        If ``0 < n_components < 1``, TPCA will select the number of
 48        compononents such that the amount of variance that needs to be
 49        explained is greater than the percentage specified.
 50
 51    copy : bool, default=True
 52        If False, data passed to fit are overwritten and running
 53        `fit(X).transform(X)` will not yield the expected results. Use
 54        `fit_transform(X)` instead.
 55
 56    M : Callable[[ndarray], ndarray] or None, default=None
 57        A function which expects an order-3 tensor as input, and returns the
 58        image under a m-transform. If unspecified TPCA will use the Discrete
 59        Cosine Transform (ii) from `scipy.fft`.
 60
 61    Minv : Callable[[ndarray], ndarray] or None, default=None
 62        A function implementing the inverse transform of `M`.
 63
 64    centre : bool, default=True
 65        If False, the data tensor will not be centralized into Mean Deviation
 66        Form. By default, the mean horizontal slice of the tensor is
 67        subtracted, so that all of the horizontal slices sum to 0, analagous
 68        to centering the data in PCA.
 69
 70    Attributes
 71    ----------
 72    n_, p_, t_, k_ : int
 73        The dimensions of the training data. ``k_ == min(n_, p_)``
 74
 75    M_, MInv_ : Callable[[ndarray], ndarray]
 76        The m-transform pair (forward and inverse) used for the underlying
 77        tensor-tensor m-product.
 78
 79    n_components_ : int
 80        The estimated number of components. If `n_components` was explicitly
 81        set by an integer value, this will be the same as that. If
 82        `n_components` was a number between 0 and 1, this number is estimated
 83        from input data. Otherwise, if not set (defaults to None), it will
 84        default to $k \\times t$ in the training data.
 85
 86    explained_variance_ratio_ : ndarray of shape (n_components_,)
 87        Percentage of total variance explained by each of the selected
 88        components. The selected components are selected so that this is
 89        returned in descending order.
 90
 91        If `n_components` is not set then all components are stored and
 92        the sum of this ratios array is equal to 1.0.
 93
 94    singular_values_ : ndarray of shape (n_components,)
 95        The singular values corresponding to each of the selected components.
 96
 97    mean_ : ndarray of shape (p_, t_)
 98        Per-feature, per-timepoint empirical mean, estimated from the
 99        training set. This is used to normalize any new data passed to
100        `transform(X)`, unless centre is explicitly turned off via
101        ``centre==False`` during object instantiation.
102
103    rho_ : ndarray of shape (t,)
104        The rho used in multi-rank truncation to achieve the desired explicit
105        rank of ``n_components``. See Mor et al. (2022) for detail.
106
107    loadings_matrix_ : ndarray of shape (n_components_, p_)
108        The i-th row corresponds to the column of $\\hat{V}$ which contains
109        the feature weights applied to the data (in the hat-space) to get the
110        i-th TPCA component.
111
112    References
113    ----------
114    Mor, U., Cohen, Y., Valdés-Mas, R., Kviatcovsky, D., Elinav, E. and Avron,
115    H., 2022. Dimensionality reduction of longitudinal’omics data using modern
116    tensor factorizations. PLoS Computational Biology, 18(7), p.e1010212.
117    """
118
119    def __init__(self, n_components=None, *, copy=True, M=None, Minv=None, centre=True):
120        """@private
121        Hacky way (for now) to suppress pdoc documentation being generated
122        for the instance variables and the constructor"""
123        # as per sklearn conventions, we perform any and all parameter
124        # validation inside fit, and none in __init__
125        self.n_components = n_components
126        """@private"""
127        self.copy = copy
128        """@private"""
129        self.M = M
130        """@private"""
131        self.Minv = Minv
132        """@private"""
133        self.centre = centre
134        """@private"""
135
136    def fit(self, X, y=None):
137        """Fit the model with X.
138
139        Parameters
140        ----------
141        X : ndarray of shape (n, p, t)
142            Training data, where `n` is the number of samples, `p` is the
143            number of features, as `t` is the number of time points.
144
145        y : Ignored
146            Ignored.
147
148        Returns
149        -------
150        self : object
151            Returns the instance itself, after being fitted.
152        """
153        self._fit(X)
154        return self
155
156    def transform(self, X):
157        """Apply dimensionality reduction to X.
158
159        See the TCAM algorithm from Mor et al. (2022)
160
161        Parameters
162        ----------
163        X : ndarray of shape (n, p, t)
164            Training data, where `n` is the number of samples, `p` is the
165            number of features, as `t` is the number of time points.
166
167        y : Ignored
168            Ignored.
169
170        Returns
171        -------
172        X_transformed : ndarray of shape (n, n_components)
173            TCAM projections in 2D transformed space.
174        """
175
176        check_is_fitted(self)
177        assert len(X.shape) == 3, "Ensure order-3 tensor input"
178        assert (
179            X.shape[1] == self.p_ and X.shape[2] == self.t_
180        ), "Ensure the number of features, and time points, matches the model fit data"
181
182        if self.centre:
183            X = X - self.mean_
184
185        # in the interest of efficiency, V was returned in the m-transformed
186        # space from tsvdm saving a pair of roundabout calls to M and Minv,
187        # and pick out the top i_q and j_q indexes, as notated in
188        # Mor et al. (2022)
189        return facewise_product(self.M_(X), self._hatV)[
190            :, self._k_t_flatten_sort[0], self._k_t_flatten_sort[1]
191        ]
192
193    def fit_transform(self, X, y=None):
194        """Fit the model with X and apply the dimensionality reduction on X.
195
196        The output here will not be identical to calling fix(X).transform(X).
197        But, the two give the same results up to machine precision. We
198        provide a brief discussion of this below.
199
200        Parameters
201        ----------
202        X : ndarray of shape (n, p, t)
203            Training data, where `n` is the number of samples, `p` is the
204            number of features, as `t` is the number of time points.
205
206        y : Ignored
207            Ignored.
208
209        Returns
210        -------
211        X_transformed : ndarray of shape (n, n_components)
212            TCAM projections in 2D transformed space.
213
214        Notes
215        -----
216        The tensor m-product from Kilmer et al. (2021) has a notion of tensor
217        inverse, and tensor orthogonality.
218
219        We benchmarked an alternative approach as taken by sklearn in their
220        PCA class. If we right multiply A's tSVDM by $V$ we note that it
221        cancels the $V^T$ giving us:
222            $$
223                Z = A *_M V = U *_M S
224            $$
225        It appears that computing the final term is more computationally
226        efficient, even if we have to convert $S$ into its full (sparse)
227        tensor representation.
228
229        By contrast, transform(X) will simply compute
230            $$
231                Z = A *_M V
232            $$
233
234        In both cases, $Z$ still needs to be converted by $\\times_3 M^{-1}$
235        and 'compressed' before being returned.
236        For these details see Mor et al. (2022).
237        """
238        # note that these tensors do NOT have full face-wise matrices
239        hatU, hatS_mat, _ = self._fit(X)
240        hatS = _singular_vals_mat_to_tensor(hatS_mat, self.k_, self.k_, self.t_)
241        return facewise_product(hatU, hatS)[
242            :, self._k_t_flatten_sort[0], self._k_t_flatten_sort[1]
243        ]
244
245    def _fit(self, X):
246        """Fit the model by computing the full SVD on X. Implementation
247        loosely modelled around `_fit_full()` method from sklearn's PCA.
248        In the future, we could potentially explore different svd solvers
249        which lets one directly pass truncation specifications into the
250        low-level solver...?
251        """
252        assert not (
253            callable(self.M) ^ callable(self.Minv)
254        ), "If explicitly defined, both M and its inverse must be defined"
255
256        assert len(X.shape) == 3, "Ensure order-3 tensor input"
257        n, p, t = X.shape
258        k = min(n, p)
259
260        # center the data into mean deviation form, see Mor et al. (2022)
261        # similar to sklearns PCA, we choose to implement this within the
262        # class and just store the mean slice for subsequent transform calls
263        self.mean_ = np.mean(X, axis=0)
264        if self.copy:
265            X = X.copy()
266        if self.centre:
267            X -= self.mean_
268
269        # if there is no explicitly defined transform in __init__, assign
270        # functions to perform a default transformation
271        if not callable(self.M):
272            self.M_, self.Minv_ = generate_default_m_transform_pair(X.shape[2])
273        else:
274            self.M_, self.Minv_ = self.M, self.Minv
275
276        # perform tensor decomposition via Kilmer's tSVDM
277        hatU, hatS_mat, hatV = tsvdm(
278            X,
279            self.M_,
280            self.Minv_,
281            keep_hats=True,
282            full_frontal_slices=False,
283            svals_matrix_form=True,
284        )
285
286        # we flatten out the compressed singular value matrix in Fortran memory
287        # style (column-wise). tensor-wise, we can interpret this as stacking
288        # the diagonals of each tensor face in S next to each other in the
289        # flattened array, where the singular values are grouped by face
290        singular_values_ = hatS_mat.flatten(order="F")
291        self._k_t_flatten_sort = singular_values_.argsort()[::-1]
292        singular_values_ = singular_values_[self._k_t_flatten_sort]
293
294        # get variance explained by singular values
295        # note that we are not yet aware of any notion of 'variance' for random
296        # tensors so we do not have sklearn PCA's self.explained_variance_
297        # however we may find literature for this in the future to include it
298        squared_singular_values = singular_values_**2
299        total_var = np.sum(squared_singular_values)
300        explained_variance_ratio_ = squared_singular_values / total_var
301
302        # process n_components input
303        if self.n_components is None:
304            n_components = k * t
305        elif isinstance(self.n_components, RealNotInt):
306            if 0 < self.n_components < 1.0:
307                # retrieve the integer number of components required to explain
308                # this proportion of the total squared sum of singular values
309                ratio_cumsum = np.cumsum(explained_variance_ratio_)
310                # np.searchsorted call returns the mininum i
311                #   s.t. n_components < ratio_cumsum[i] (see numpy docs)
312                # which means that the (i+1)th element in ratio_cumsum[i]
313                # strictly exceeds the user's specified variance ratio
314                n_components = (
315                    np.searchsorted(ratio_cumsum, self.n_components, side="right") + 1
316                )
317            else:
318                raise ValueError(
319                    "For non-integer inputs, ensure that 0 < n_components < 1"
320                )
321        elif isinstance(self.n_components, Integral):
322            if 1 <= self.n_components <= k * t:
323                n_components = self.n_components
324            else:
325                raise ValueError(
326                    f"Integer inputs must satisfy 1 <= n_components <= min(n, p)*t={k*t}"
327                )
328        else:
329            raise TypeError(
330                "n_components must be an integer, float, or None"
331                f"Got {type(self.n_components)} instead"
332            )
333
334        # convert the argsort indexes back into the two dimensional indexes.
335        # in the same tensor semantics as hatS_mat, the rows (tuple[0]) in this
336        # multindex correspond to the p-dimension location, and the
337        # columns (tuple[1]) in the multindex correspond to the t-dimension
338        # location. these are now the collection of i_h's and j_h's
339        # in Mor et al. (2022)
340        self._k_t_flatten_sort = np.unravel_index(
341            self._k_t_flatten_sort[:n_components], shape=hatS_mat.shape, order="F"
342        )
343
344        # perform truncation. this speeds up subsequent calls to transform,
345        # but is not required to obtain the correct results. we include it
346        # because it also computes rho, at basically no extra cost
347        rho = _rank_q_truncation_zero_out(
348            hatU, hatS_mat, hatV, sigma_q=singular_values_[n_components - 1]
349        )
350
351        # we store the features tensor, in the tSVDM decomposition, in the
352        # transforemd space, saving roundabout calls to M and Minv when
353        # performing m-product with new data
354        self._hatV = hatV
355
356        # store public attributes; as per sklearn conventions, we use trailing
357        # underscores to indicate that they have been populated following a
358        # call to fit()
359        self.n_, self.p_, self.t_, self.k_ = n, p, t, k
360        self.n_components_ = n_components
361        self.explained_variance_ratio_ = explained_variance_ratio_[:n_components]
362        self.singular_values_ = singular_values_[:n_components]
363        self.rho_ = rho
364        self.loadings_matrix_ = hatV[
365            :, self._k_t_flatten_sort[0], self._k_t_flatten_sort[1]
366        ].T
367
368        return hatU, hatS_mat, hatV
369
370    @property
371    def _n_features_out(self):
372        """Number of transformed output features.
373
374        [CAN IGNORE - NOT MATHEMATICALLY RELEVANT]:
375        See sklearn/decompositions/_base.py
376        """
377        return self.n_components_

23def display_tensor_facewise(tens):
24    """Display an order-3 tensor represented by a Numpy ndarray nicely.
25
26    By default Numpy prints order-3 arrays as vertical stacks of order-2
27    arrays in line with their broadcasting rules. This function prints a
28    transpose view so the print output is a more intuitive sequence of frontal
29    slices. It also prints the rounded values.
30
31    We often use this in notebooks.
32
33    UPDATE (v0.1.1): Also works for matrices and vectors now.
34
35    Parameters
36    ----------
37    tens : ndarray of shape (n, p, t)
38        Order-3 tensor.
39
40    Examples
41    --------
42    >>> import numpy as np
43    >>> from tred import display_tensor_facewise as disp
44    >>> test = np.eye(3)[:, :, None] # a 3x3x1 tensor
45    >>> print(test) # Numpy ndarrays default __str__ method is not intuitive
46    [[[1.]
47    [0.]
48    [0.]]
49    [[0.]
50    [1.]
51    [0.]]
52    [[0.]
53    [0.]
54    [1.]]]
55    >>> disp(test)
56    Tensor with dimensions (3, 3, 1)
57    [[[1. 0. 0.]
58    [0. 1. 0.]
59    [0. 0. 1.]]]
60    """
61    assert len(tens.shape) in (
62        1,  # vector
63        2,  # matrix
64        3,  # tensor
65    ), "Expecting 1D (vector), 2D (matrix) or 3D (tensor) input"
66
67    if len(tens.shape) == 3:
68        print(f"Tensor with shape {tens.shape}")
69        print(tens.transpose(2, 0, 1).round(6))
70    else:
71        if len(tens.shape) == 2:
72            print(f"Matrix with shape {tens.shape}")
73        else:
74            print(f"Vector with length {tens.shape}")
75        print(tens.round(6))

29def generate_default_m_transform_pair(t):
30    """Return the default `M, Minv` used by `tred` algorithms.
31
32    We do not guarantee that the default m-transform used in `tred` will remain
33    consistent between versions, as this may change depending on recent
34    research and literature. We highly encourage `tred` users to explicitly
35    define M and Minv in the calling state to 'future-proof' your code.
36
37    Parameters
38    ----------
39    t : int
40        The length of the tubal fibers of the target tensors, i.e. the size of
41        its third dimension.
42
43    Returns
44    -------
45    M : Callable[[ndarray], ndarray]
46        A function which expects an order-3 tensor as input, and returns the
47        image under the default m-transform being used in `tred`
48
49    Minv : Callable[[ndarray], ndarray]
50        A function implementing the inverse transform of `M`.
51
52    References
53    ----------
54    Kilmer, M.E., Horesh, L., Avron, H. and Newman, E., 2021. Tensor-tensor
55    algebra for optimal representation and compression of multiway data.
56    Proceedings of the National Academy of Sciences, 118(28), p.e2015851118.
57
58    Mor, U., Cohen, Y., Valdés-Mas, R., Kviatcovsky, D., Elinav, E. and Avron,
59    H., 2022. Dimensionality reduction of longitudinal’omics data using modern
60    tensor factorizations. PLoS Computational Biology, 18(7), p.e1010212.
61    """
62    return generate_dctii_m_transform_pair(t)

 65def generate_transform_pair_from_matrix(M_mat, Minv_mat=None, *, inplace=False):
 66    """Return `M, Minv` as defined by an orthogonal matrix.
 67
 68    Allows the user to specify any orthogonal matrix, and this function will
 69    infer `t`, and numerically compute the inverse, returning functions `M`
 70    and `Minv` which can be used for `tred` algorithms.
 71
 72    Optionally, the user can also choose to specify the inverse explicitly.
 73
 74    Parameters
 75    ----------
 76    M_mat : ndarray
 77        Orthogonal square matrix.
 78
 79    Minv_mat : ndarray or None, default=None
 80        Square matrix, the inverse of M_mat. If not specified, this function
 81        will numerically evaluate the inverse of `M_mat`.
 82
 83    inplace : bool, default=False
 84        *Placeholder for future development*
 85
 86    Returns
 87    -------
 88    M : Callable[[ndarray], ndarray]
 89        A function which expects an order-3 tensor as input, and applies
 90        `M_mat` to each of the tubal fibres. This preserves the dimensions of
 91        the tensor.
 92
 93    Minv : Callable[[ndarray], ndarray]
 94        A function implementing the inverse transform of `M`.
 95
 96    References
 97    ----------
 98    Kilmer, M.E., Horesh, L., Avron, H. and Newman, E., 2021. Tensor-tensor
 99    algebra for optimal representation and compression of multiway data.
100    Proceedings of the National Academy of Sciences, 118(28), p.e2015851118.
101    """
102    assert len(M_mat.shape) == 2, "Expecting matrix (order-2 array) input"
103    assert M_mat.shape[0] == M_mat.shape[1], "Expecting square matrix input"
104    t_ = M_mat.shape[0]
105
106    # numerically evaluate inverse matrix if not explicitly specified
107    if Minv_mat is None:
108        # be pedantic and check for invertibility
109        # https://stackoverflow.com/questions/13249108/efficient-pythonic-check-for-singular-matrix
110        # NOTE: revisit later, is the following more robust?:
111        # https://stackoverflow.com/questions/17931613/how-to-decide-a-whether-a-matrix-is-singular-in-python-numpy
112        if linalg.cond(M_mat) < 1 / sys.float_info.epsilon:
113            Minv_mat = linalg.inv(M_mat)
114        else:
115            raise ValueError(
116                "Input matrix must be invertible, but appears singular, or close to singular"
117            )
118    else:
119        assert (
120            Minv_mat.shape == M_mat.shape
121        ), "Ensure the shapes of matrix M and its inverse are the same"
122
123    # a quick way of applying M along the tubal fibres of X, using numpy
124    # broadcasting: transpose X into a (n x t x p) vertical stack of (t x p)
125    # matrices. matrix multiplication is broadcast vertically, whereby M left
126    # multiplies each of the t x p matrices, effectively applying the
127    # transform to the columns, which were the original tubal fibres.
128    # then just perform the reverse transposition to get the original
129    # orientation of the tensor
130    def M(X):
131        _assert_t_and_order(X, t_)
132        if len(X.shape) == 3:
133            return (M_mat @ X.transpose(0, 2, 1)).transpose(0, 2, 1)
134        elif len(X.shape) == 2:
135            return (M_mat @ X.T).T
136        else:  # len(X.shape == 1)
137            return M_mat @ X
138
139    def Minv(X):
140        _assert_t_and_order(X, t_)
141        if len(X.shape) == 3:
142            return (Minv_mat @ X.transpose(0, 2, 1)).transpose(0, 2, 1)
143        elif len(X.shape) == 2:
144            return (Minv_mat @ X.T).T
145        else:  # len(X.shape == 1)
146            return Minv_mat @ X
147
148    return M, Minv

151def generate_dctii_m_transform_pair(t, *, inplace=False, norm="ortho"):
152    """Return `M, Minv` as defined by the Discrete Cosine Transform of
153    length `t`.
154
155    Bascially a wrapper around scipy.fft to generate functions to perform
156    fourier transform based $\\times_3 M$ operations, as used by
157    Mor et al. (2022)
158
159    Parameters
160    ----------
161    t : int
162        The length of the transform
163
164    inplace : bool, default=False
165        Control whether or not the generated functions modify the input tensor
166        in-place, or return a copy with the m-transform applied
167
168    norm : {“backward”, “ortho”, “forward”}, default="ortho"
169        See https://docs.scipy.org/doc/scipy/reference/generated/scipy.fft.dct.html#scipy.fft.dct
170
171    Returns
172    -------
173    M : Callable[[ndarray], ndarray]
174        A function which expects an order-3 tensor as input, and applies the
175        DCT-II transorm to each of the tubal fibres. This preserves the
176        dimensions of the tensor.
177
178    Minv : Callable[[ndarray], ndarray]
179        A function implementing the inverse transform of `M`.
180
181    References
182    ----------
183    Mor, U., Cohen, Y., Valdés-Mas, R., Kviatcovsky, D., Elinav, E. and Avron,
184    H., 2022. Dimensionality reduction of longitudinal’omics data using modern
185    tensor factorizations. PLoS Computational Biology, 18(7), p.e1010212.
186    """
187
188    def M(X):
189        _assert_t_and_order(X, t)
190        return dct(
191            X,
192            type=2,
193            n=t,
194            axis=-1,
195            norm=norm,
196            workers=2 * os.cpu_count(),
197            overwrite_x=inplace,
198        )
199
200    def Minv(X):
201        _assert_t_and_order(X, t)
202        return idct(
203            X,
204            type=2,
205            n=t,
206            axis=-1,
207            norm=norm,
208            workers=2 * os.cpu_count(),
209            overwrite_x=inplace,
210        )
211
212    return M, Minv

215def generate_dstii_m_transform_pair(t, *, inplace=False, norm="ortho"):
216    """Return `M, Minv` as defined by the Discrete Sine Transform of
217    length `t`.
218
219    Bascially a wrapper around scipy.fft to generate functions to perform
220    fourier transform based $\\times_3 M$ operations.
221
222    Parameters
223    ----------
224    t : int
225        The length of the transform
226
227    inplace : bool, default=False
228        Control whether or not the generated functions modify the input tensor
229        in-place, or return a copy with the m-transform applied
230
231    norm : {“backward”, “ortho”, “forward”}, default="ortho"
232        See https://docs.scipy.org/doc/scipy/reference/generated/scipy.fft.dst.html#scipy.fft.dst
233
234    Returns
235    -------
236    M : Callable[[ndarray], ndarray]
237        A function which expects an order-3 tensor as input, and applies the
238        DST-II transorm to each of the tubal fibres. This preserves the
239        dimensions of the tensor.
240
241    Minv : Callable[[ndarray], ndarray]
242        A function implementing the inverse transform of `M`.
243    """
244
245    def M(X):
246        _assert_t_and_order(X, t)
247        return dst(
248            X,
249            type=2,
250            n=t,
251            axis=-1,
252            norm=norm,
253            workers=2 * os.cpu_count(),
254            overwrite_x=inplace,
255        )
256
257    def Minv(X):
258        _assert_t_and_order(X, t)
259        return idst(
260            X,
261            type=2,
262            n=t,
263            axis=-1,
264            norm=norm,
265            workers=2 * os.cpu_count(),
266            overwrite_x=inplace,
267        )
268
269    return M, Minv

19def facewise_product(*tensors):
20    """Compute cumulative facewise product.
21
22    Definition:
23    $$
24        C_{:,:,i} = A_{:,:,i} B_{:,:,i}
25    $$
26
27    Parameters
28    ----------
29    *tensors : ndarray
30        Variable number of tensors, such that all adjacent input tensors have
31        shape (a, b, d) and shape (b, c, d) respectively
32
33    Returns
34    -------
35    C : ndarray of shape (a, c, d)
36        Facewise tensor product
37    """
38    # apply the lambda function cumulatively over the tensor inputs
39    return reduce(_BINARY_FACEWISE, tensors)

42def m_product(*tensors, **transforms):
43    """Kilmer et al. (2021) tensor m-product for order-3 tensors.
44
45    See paper for definition.
46
47    Parameters
48    ----------
49    *tensors : ndarray
50        Variable number of tensors, such that all adjacent input tensors have
51        shape (a, b, d) and shape (b, c, d) respectively
52
53    M : Callable[[ndarray], ndarray] or None, default=None
54        A function which, given some order-3 tensor, returns it under an
55        orthogonal tubal transformation
56
57    MInv : Callable[[ndarray], ndarray] or None, default=None
58        A function implementing the inverse tubal transformation of M
59
60    Returns
61    -------
62    m_product : ndarray of shape (a, c, d)
63        Tensor-tensor m-product as found in Kilmer et al. (2021)
64
65    References
66    ----------
67    Kilmer, M.E., Horesh, L., Avron, H. and Newman, E., 2021. Tensor-tensor
68    algebra for optimal representation and compression of multiway data.
69    Proceedings of the National Academy of Sciences, 118(28), p.e2015851118.
70    """
71    # process the m-transform keyword inputs, applying defaults if needed
72    default_transforms = {"M": None, "Minv": None}
73    transforms = {**default_transforms, **transforms}
74    assert not (
75        callable(transforms["M"]) ^ callable(transforms["Minv"])
76    ), "If explicitly defined, both M and its inverse must be defined"
77
78    if not callable(transforms["M"]):
79        M, Minv = generate_default_m_transform_pair(tensors[0].shape[-1])
80    else:
81        M, Minv = transforms["M"], transforms["Minv"]
82
83    # use a generator representing all the tensors under the m-transform
84    # ('hat space'), allowing it to be lazily reduced over using the binary
85    # facewise product anonymous function. prevents storing all of the
86    # transformed tensors in memory at the same time
87    return Minv(reduce(_BINARY_FACEWISE, (M(tens) for tens in tensors)))