<commit_before>import numpy as np
from numpy import newaxis
from ffthompy.sparse.objects.tensors import SparseTensorFuns
from ffthompy.tensors import Tensor
from ffthompy.sparse.objects.tensors import fft_form_default

class CanoTensor(SparseTensorFuns):
    kind='cano'

    def __init__(self, name='unnamed', val=None, core=None, basis=None, orthogonal=False,
                 Fourier=False, r=None, N=None, randomise=False, fft_form=fft_form_default):

        self.name=name
        self.Fourier=Fourier # TODO: dtype instead of Fourier
        self.orthogonal=orthogonal

        if val is not None:
            if len(val.shape)==2:
                u, s, vt=np.linalg.svd(val, full_matrices=0)

                self.order=2
                self.basis=[u.T, vt]
                self.core=s
                self.r=self.core.shape[0]

                self.N=np.empty(self.order, dtype=np.int)
                for ii in range(self.order):
                    self.N[ii]=self.basis[ii].shape[1]
            else:
                raise ValueError("Canonical format not yet applicable to tensors higher than 2 dimensional.")

        elif core is not None and basis is not None:
            self.order=basis.__len__()
            self.basis=basis

            if len(core.shape)==2:
                self.core=np.diag(core)
            else:
                self.core=core

            self.r=basis[0].shape[0] # since all basis share the same r

            self.N=np.empty(self.order, dtype=np.int)
            for ii in range(self.order):
                self.N[ii]=basis[ii].shape[1]
        else:
            self.r=3
            self.N=[5, 5]
            self.order=self.N.__len__()
            if randomise:
                self.randomise()
            else:
                self.core=np.zeros(r)
                self.basis=[np.zeros([self.r[ii], self.N[ii]]) for ii in range(self.order)]

        self._set_fft(fft_form)

    def set_fft_form(self, fft_form=fft_form_default, copy=False):
        if copy:
            R=self.copy()
        else:
            R=self

        if self.fft_form==fft_form:
            return R

        if R.Fourier:
            for i in range(R.order):
                R.basis[i]= R.ifft(R.basis[i], R.N[i])

            R._set_fft(fft_form)
            for i in range(R.order):
                R.basis[i]= R.fft(R.basis[i], R.N[i])
            R.orthogonal=False # FFT modes makes the basis no longer orthonormal
        else:
            R._set_fft(fft_form)

        return R

    def randomise(self):
        self.core=np.random.random((self.r,))
        self.basis=[np.random.random([self.r, self.N[ii]]) for ii in range(self.order)]

    def orthogonalise(self):
        """re-orthogonalise the basis"""
        if self.orthogonal:
            return self
        else:
            # re-orthogonalise the basis by QR and SVD
            qa, ra=np.linalg.qr(self.basis[0].T, 'reduced')
            qb, rb=np.linalg.qr(self.basis[1].T, 'reduced')

            core=ra.real*self.core[np.newaxis, :]
            core=np.dot(core, rb.T.real)

            u, s, vt=np.linalg.svd(core,full_matrices=False)

            newA=np.dot(qa, u)
            newB=np.dot(vt, qb.T)

            newBasis=[newA.T, newB]
            return CanoTensor(name=self.name, core=s, basis=newBasis, orthogonal=True,
                              Fourier=self.Fourier, fft_form=self.fft_form)

    def __add__(self, Y):
        X=self
        assert(X.Fourier==Y.Fourier)
        core=np.hstack([X.core, Y.core])
        basis=[np.vstack([X.basis[ii], Y.basis[ii]]) for ii in range(self.order)]

        return self.copy(name=X.name+'+'+Y.name, core=core, basis=basis)

    def __mul__(self, Y):
        "element-wise multiplication of two canonical tensors"

        if isinstance(Y, (float, int)) or np.isscalar(Y) or (isinstance(Y, np.ndarray) and Y.shape==()):
            R=self.copy()
            R.core=self.core*Y
            return R
        elif isinstance(Y, np.ndarray) and Y.shape==(1,):
            R=self.copy()
            R.core=self.core*Y[0]
            return R
        else:
            coeff=np.kron(self.core, Y.core)
            newBasis=[None]*self.order
            if self.Fourier and self.fft_form=='sr': #product of scipy rfft tensors need a special multiplication
                for d in range(0, self.order):
                    B=np.empty((self.r*Y.r, self.N[d]))
                    B[:,0]=np.kron(self.basis[d][:,0],Y.basis[d][:,0])
                    if self.N[d]%2 != 0:
                        ar=self.basis[d][:,1::2]
                        ai=self.basis[d][:,2::2]
                        br=Y.basis[d][:,1::2]
                        bi=Y.basis[d][:,2::2]
                        B[:,1::2]=(ar[:, newaxis, :]*br[newaxis, :, :]-ai[:, newaxis, :]*bi[newaxis, :, :]).reshape(self.r*Y.r,-1)
                        B[:,2::2]=(ar[:, newaxis, :]*bi[newaxis, :, :]+ai[:, newaxis, :]*br[newaxis, :, :]).reshape(self.r*Y.r,-1)
                    else:
                        B[:,-1]=np.kron(self.basis[d][:,-1],Y.basis[d][:,-1])
                        ar=self.basis[d][:,1:-1:2]
                        ai=self.basis[d][:,2:-1:2]
                        br=Y.basis[d][:,1:-1:2]
                        bi=Y.basis[d][:,2:-1:2]
                        B[:,1:-1:2]=(ar[:, newaxis, :]*br[newaxis, :, :]-ai[:, newaxis, :]*bi[newaxis, :, :]).reshape(self.r*Y.r,-1)
                        B[:,2:-1:2]=(ar[:, newaxis, :]*bi[newaxis, :, :]+ai[:, newaxis, :]*br[newaxis, :, :]).reshape(self.r*Y.r,-1)
                    newBasis[d]=B
            else:
                for d in range(0, self.order):
                    newBasis[d]=np.multiply(self.basis[d][:, newaxis, :], Y.basis[d][newaxis, :, :])
                    newBasis[d]=np.reshape(newBasis[d], (-1, self.N[d]))

            return self.copy(name=self.name+'*'+Y.name, core=coeff, basis=newBasis, orthogonal=False)

    def full(self, **kwargs):
        "return a full tensor object"

        if self.order==2:
            if self.Fourier:
                res=self.fourier(copy=True)
            else:
                res=self

            # generate Tensor in real domain
            val=np.einsum('i,ik,il->kl', res.core, res.basis[0], res.basis[1])
            kwargsT=dict(name=res.name, val=val, order=0, N=val.shape, Fourier=False, fft_form='c')
            kwargsT.update(kwargs)
            T=Tensor(**kwargsT)

            if self.Fourier:
                T.fourier()
            return T
        else:
            raise NotImplementedError()

    def truncate(self, rank=None, tol=None):
        "return truncated tensor"
        # tol is the maximum "portion" of the core trace to be lost, e.g. tol=0.01 means at most 1 percent could be lost in the truncation.
        # if tol is not none, it will override rank as the truncation criteria.
        if tol is None and rank is None:
            return self

        if rank is not None:
            if rank>=self.r:
                #print ("Warning: Rank of the truncation not smaller than the original rank, truncation aborted!")
                return self

        self=self.orthogonalise()

        basis=list(self.basis)
        core=self.core

        if tol is not None:
            # determine the truncation rank so that (1.0-tol)*100% of the trace of the core is perserved.
            rank=np.searchsorted(np.cumsum(np.abs(core))/np.sum(np.abs(core)), 1.0-tol)+1<commit_msg>Add a fast option to the truncation function of canonical format  that does a preliminary non-optimal truncation before the optimal truncation.<commit_after>import numpy as np
from numpy import newaxis
from ffthompy.sparse.objects.tensors import SparseTensorFuns
from ffthompy.tensors import Tensor
from ffthompy.sparse.objects.tensors import fft_form_default
from ffthompy.sparse.decompositions import fast_qr

class CanoTensor(SparseTensorFuns):
    kind='cano'

    def __init__(self, name='unnamed', val=None, core=None, basis=None, orthogonal=False,
                 Fourier=False, r=None, N=None, randomise=False, fft_form=fft_form_default):

        self.name=name
        self.Fourier=Fourier # TODO: dtype instead of Fourier
        self.orthogonal=orthogonal

        if val is not None:
            if len(val.shape)==2:
                u, s, vt=np.linalg.svd(val, full_matrices=0)

                self.order=2
                self.basis=[u.T, vt]
                self.core=s
                self.r=self.core.shape[0]

                self.N=np.empty(self.order, dtype=np.int)
                for ii in range(self.order):
                    self.N[ii]=self.basis[ii].shape[1]
            else:
                raise ValueError("Canonical format not yet applicable to tensors higher than 2 dimensional.")

        elif core is not None and basis is not None:
            self.order=basis.__len__()
            self.basis=basis

            if len(core.shape)==2:
                self.core=np.diag(core)
            else:
                self.core=core

            self.r=basis[0].shape[0] # since all basis share the same r

            self.N=np.empty(self.order, dtype=np.int)
            for ii in range(self.order):
                self.N[ii]=basis[ii].shape[1]
        else:
            self.r=3
            self.N=[5, 5]
            self.order=self.N.__len__()
            if randomise:
                self.randomise()
            else:
                self.core=np.zeros(r)
                self.basis=[np.zeros([self.r[ii], self.N[ii]]) for ii in range(self.order)]

        self._set_fft(fft_form)

    def set_fft_form(self, fft_form=fft_form_default, copy=False):
        if copy:
            R=self.copy()
        else:
            R=self

        if self.fft_form==fft_form:
            return R

        if R.Fourier:
            for i in range(R.order):
                R.basis[i]= R.ifft(R.basis[i], R.N[i])

            R._set_fft(fft_form)
            for i in range(R.order):
                R.basis[i]= R.fft(R.basis[i], R.N[i])
            R.orthogonal=False # FFT modes makes the basis no longer orthonormal
        else:
            R._set_fft(fft_form)

        return R

    def randomise(self):
        self.core=np.random.random((self.r,))
        self.basis=[np.random.random([self.r, self.N[ii]]) for ii in range(self.order)]

    def orthogonalise(self, rank=None, tol=None):
        """re-orthogonalise the basis"""
        if self.orthogonal:
            return self

        if (rank is not None) or (tol is not None) :

            norm0= np.linalg.norm(self.basis[0],axis=1)
            norm1= np.linalg.norm(self.basis[1],axis=1)


            self.basis[0]/= norm0[:, np.newaxis]
            self.basis[1]/= norm1[:, np.newaxis]

            self.core *= (norm0*norm1)

            if tol is not None:
                core_abs= np.abs(self.core)
                select_ind = np.argwhere(core_abs >= tol*np.sum(core_abs))
            else:
                ind=np.argpartition(-np.abs(self.core), rank)[:rank]
                rank_th_largest = np.abs(self.core[ind[-1]])
                select_ind = np.argwhere(np.abs(self.core) >= (1e-3)*rank_th_largest)

            select_ind = np.squeeze(select_ind)

            self.core = np.take(self.core, select_ind )
            self.basis[0] = np.take(self.basis[0], select_ind, axis=0 )
            self.basis[1] = np.take(self.basis[1], select_ind, axis=0 )

            qa, ra=fast_qr(self.basis[0].T)
            qb, rb=fast_qr(self.basis[1].T)
        else:
            qa, ra=np.linalg.qr(self.basis[0].T, 'reduced')
            qb, rb=np.linalg.qr(self.basis[1].T, 'reduced')

        core=ra.real*self.core[np.newaxis, :]
        core=np.dot(core, rb.T.real)

        u, s, vt=np.linalg.svd(core,full_matrices=False)


        newA=np.dot(qa, u)
        newB=np.dot(vt, qb.T)

        newBasis=[newA.T, newB]

        return CanoTensor(name=self.name, core=s, basis=newBasis, orthogonal=True,
                          Fourier=self.Fourier, fft_form=self.fft_form)

    def __add__(self, Y):
        X=self
        assert(X.Fourier==Y.Fourier)
        core=np.hstack([X.core, Y.core])
        basis=[np.vstack([X.basis[ii], Y.basis[ii]]) for ii in range(self.order)]

        return self.copy(name=X.name+'+'+Y.name, core=core, basis=basis)

    def __mul__(self, Y):
        "element-wise multiplication of two canonical tensors"

        if isinstance(Y, (float, int)) or np.isscalar(Y) or (isinstance(Y, np.ndarray) and Y.shape==()):
            R=self.copy()
            R.core=self.core*Y
            return R
        elif isinstance(Y, np.ndarray) and Y.shape==(1,):
            R=self.copy()
            R.core=self.core*Y[0]
            return R
        else:
            coeff=np.kron(self.core, Y.core)
            newBasis=[None]*self.order
            if self.Fourier and self.fft_form=='sr': #product of scipy rfft tensors need a special multiplication
                for d in range(0, self.order):
                    B=np.empty((self.r*Y.r, self.N[d]))
                    B[:,0]=np.kron(self.basis[d][:,0],Y.basis[d][:,0])
                    if self.N[d]%2 != 0:
                        ar=self.basis[d][:,1::2]
                        ai=self.basis[d][:,2::2]
                        br=Y.basis[d][:,1::2]
                        bi=Y.basis[d][:,2::2]
                        B[:,1::2]=(ar[:, newaxis, :]*br[newaxis, :, :]-ai[:, newaxis, :]*bi[newaxis, :, :]).reshape(self.r*Y.r,-1)
                        B[:,2::2]=(ar[:, newaxis, :]*bi[newaxis, :, :]+ai[:, newaxis, :]*br[newaxis, :, :]).reshape(self.r*Y.r,-1)
                    else:
                        B[:,-1]=np.kron(self.basis[d][:,-1],Y.basis[d][:,-1])
                        ar=self.basis[d][:,1:-1:2]
                        ai=self.basis[d][:,2:-1:2]
                        br=Y.basis[d][:,1:-1:2]
                        bi=Y.basis[d][:,2:-1:2]
                        B[:,1:-1:2]=(ar[:, newaxis, :]*br[newaxis, :, :]-ai[:, newaxis, :]*bi[newaxis, :, :]).reshape(self.r*Y.r,-1)
                        B[:,2:-1:2]=(ar[:, newaxis, :]*bi[newaxis, :, :]+ai[:, newaxis, :]*br[newaxis, :, :]).reshape(self.r*Y.r,-1)
                    newBasis[d]=B
            else:
                for d in range(0, self.order):
                    newBasis[d]=np.multiply(self.basis[d][:, newaxis, :], Y.basis[d][newaxis, :, :])
                    newBasis[d]=np.reshape(newBasis[d], (-1, self.N[d]))

            return self.copy(name=self.name+'*'+Y.name, core=coeff, basis=newBasis, orthogonal=False)

    def full(self, **kwargs):
        "return a full tensor object"

        if self.order==2:
            if self.Fourier:
                res=self.fourier(copy=True)
            else:
                res=self

            # generate Tensor in real domain
            val=np.einsum('i,ik,il->kl', res.core, res.basis[0], res.basis[1])
            kwargsT=dict(name=res.name, val=val, order=0, N=val.shape, Fourier=False, fft_form='c')
            kwargsT.update(kwargs)
            T=Tensor(**kwargsT)

            if self.Fourier:
                T.fourier()
            return T
        else:
            raise NotImplementedError()

    def truncate(self, rank=None, tol=None, fast=False):
        "return truncated tensor"
        # tol is the maximum "portion" of the core trace to be lost, e.g. tol=0.01 means at most 1 percent could be lost in the truncation.
        # if tol is not none, it will override rank as the truncation criteria.
        if tol is None and rank is None:
            return self

        if rank is not None:
            if rank>=self.r:
                #print ("Warning: Rank of the truncation not smaller than the original rank, truncation aborted!")
                return self
        if fast:
            self=self.orthogonalise(rank=rank, tol=tol)
        else:
            self=self.orthogonalise()

        basis=list(self.basis)
        core=self.core

        if tol is not None:
            # determine the truncation rank so that (1.0-tol)*100% of the trace of the core is perserved.
            rank=np.searchsorted(np.cumsum(np.abs(core))/np.sum(np.abs(core)), 1.0-tol)+1
{'fraction_non_alphanumeric': 0.14570561644683883, 'fraction_numerical': 0.011084277664251657, 'mean_word_length': 1.775390960646159, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3936901', 'n_tokens_mistral': 5688, 'n_tokens_neox': 5333, 'n_words': 1079}