<commit_before>import  numpy as np
import  scipy as sp
import pandas as pd
import scipy.sparse
import numbers

from .helper import SparseTensor

def make_train_test(Y, ntest, shape = None, seed = None):
    """Splits a sparse matrix Y into a train and a test matrix.

    Parameters
    ----------
        Y : scipy sparse matrix (coo_matrix, csr_matrix or csc_matrix)
             or
            numpy dense ndarray 
             or
            pandas DataFrame or smurff.SparseTensor

            Matrix/Array/Tensor to split

        ntest : float <1.0 or integer.
           - if float, then indicates the ratio of test cells
           - if integer, then indicates the number of test cells

    Returns
    -------
        Ytrain : csr_matrix
            train part

        Ytest : csr_matrix
            test part
    """
    if isinstance(Y, pd.DataFrame) or isinstance(Y, SparseTensor):
        return make_train_test_df(Y, ntest, shape)

    if isinstance(Y, np.ndarray):
        nmodes = len(Y.shape)
        if (nmodes > 2):
            Ysparse = SparseTensor(Y)
        else:
            Ysparse = sp.sparse.coo_matrix(Y)

        return make_train_test(Ysparse, ntest, shape)
    
    if not sp.sparse.issparse(Y):
        raise TypeError("Unsupported Y type: " + str(type(Y)))

    if not isinstance(ntest, numbers.Real) or ntest < 0:
        raise TypeError("ntest has to be a non-negative number (number or ratio of test samples).")

    Y = Y.tocoo(copy = False)
    if ntest < 1:
        ntest = Y.nnz * ntest
    ntest = int(round(ntest))
    np.random.seed(seed)
    rperm = np.random.permutation(Y.nnz)
    train = rperm[ntest:]
    test  = rperm[0:ntest]
    if shape is None:
        shape = Y.shape

    Ytrain = sp.sparse.coo_matrix( (Y.data[train], (Y.row[train], Y.col[train])), shape=shape )
    Ytest  = sp.sparse.coo_matrix( (Y.data[test],  (Y.row[test],  Y.col[test])),  shape=shape )
    return Ytrain, Ytest

def make_train_test_df(Y, ntest, shape = None):
    """Splits rows of dataframe Y into a train and a test dataframe.
       Y      pandas dataframe
       ntest  either a float below 1.0 or integer.
              if float, then indicates the ratio of test cells
              if integer, then indicates the number of test cells
       returns Ytrain, Ytest (type SparseTensor)
    """
    if  isinstance(Y, pd.DataFrame):
        return make_train_test_df(SparseTensor(Y), ntest, Y.shape)

    if not isinstance(Y, SparseTensor):
        raise TypeError("Y should be DataFrame or SparseTensor.")

    if not isinstance(ntest, numbers.Real) or ntest < 0:
        raise TypeError("ntest has to be a non-negative number (number or ratio of test samples).")

    ## randomly spliting train-test
    if ntest < 1:
        ntest = Y.nnz * ntest

    ntest  = int(round(ntest))
    rperm  = np.random.permutation(Y.nnz)
    train  = rperm[ntest:]
    test   = rperm[0:ntest]
<commit_msg>FIX: make sure we have always at least one element in test-set<commit_after>import  numpy as np
import  scipy as sp
import pandas as pd
import scipy.sparse
import numbers

from .helper import SparseTensor

def make_train_test(Y, ntest, shape = None, seed = None):
    """Splits a sparse matrix Y into a train and a test matrix.

    Parameters
    ----------
        Y : scipy sparse matrix (coo_matrix, csr_matrix or csc_matrix)
             or
            numpy dense ndarray 
             or
            pandas DataFrame or smurff.SparseTensor

            Matrix/Array/Tensor to split

        ntest : float <1.0 or integer.
           - if float, then indicates the ratio of test cells
           - if integer, then indicates the number of test cells

    Returns
    -------
        Ytrain : csr_matrix
            train part

        Ytest : csr_matrix
            test part
    """
    if isinstance(Y, pd.DataFrame) or isinstance(Y, SparseTensor):
        return make_train_test_df(Y, ntest, shape)

    if isinstance(Y, np.ndarray):
        nmodes = len(Y.shape)
        if (nmodes > 2):
            Ysparse = SparseTensor(Y)
        else:
            Ysparse = sp.sparse.coo_matrix(Y)

        return make_train_test(Ysparse, ntest, shape)
    
    if not sp.sparse.issparse(Y):
        raise TypeError("Unsupported Y type: " + str(type(Y)))

    if not isinstance(ntest, numbers.Real) or ntest < 0:
        raise TypeError("ntest has to be a non-negative number (number or ratio of test samples).")

    Y = Y.tocoo(copy = False)
    if ntest < 1:
        ntest = Y.nnz * ntest
    ntest = int(round(ntest))
    ntest = max(1,ntest)

    np.random.seed(seed)
    rperm = np.random.permutation(Y.nnz)
    train = rperm[ntest:]
    test  = rperm[0:ntest]
    if shape is None:
        shape = Y.shape

    Ytrain = sp.sparse.coo_matrix( (Y.data[train], (Y.row[train], Y.col[train])), shape=shape )
    Ytest  = sp.sparse.coo_matrix( (Y.data[test],  (Y.row[test],  Y.col[test])),  shape=shape )
    return Ytrain, Ytest

def make_train_test_df(Y, ntest, shape = None):
    """Splits rows of dataframe Y into a train and a test dataframe.
       Y      pandas dataframe
       ntest  either a float below 1.0 or integer.
              if float, then indicates the ratio of test cells
              if integer, then indicates the number of test cells
       returns Ytrain, Ytest (type SparseTensor)
    """
    if  isinstance(Y, pd.DataFrame):
        return make_train_test_df(SparseTensor(Y), ntest, Y.shape)

    if not isinstance(Y, SparseTensor):
        raise TypeError("Y should be DataFrame or SparseTensor.")

    if not isinstance(ntest, numbers.Real) or ntest < 0:
        raise TypeError("ntest has to be a non-negative number (number or ratio of test samples).")

    ## randomly spliting train-test
    if ntest < 1:
        ntest = Y.nnz * ntest

    ntest = int(round(ntest))
    ntest = max(1,ntest)

    rperm  = np.random.permutation(Y.nnz)
    train  = rperm[ntest:]
    test   = rperm[0:ntest]

{'fraction_non_alphanumeric': 0.09720820719811638, 'fraction_numerical': 0.004036326942482341, 'mean_word_length': 2.2783902976846746, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 5, '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': '6477892', 'n_tokens_mistral': 1956, 'n_tokens_neox': 1903, 'n_words': 603}