Danfo.js has something similar to "sklearn.model_selection.train_test_split" in scikit-learn/Python

[LSW1980]

Hi all,

Does Danfo.js has something similar to "sklearn.model_selection.train_test_split" in scikit-learn/Python?
[https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html]

What function most close to it in Danfo.js ? Example?

Thank you

[risenW]

Danfojs does not have an exact equivalent of sklearn's train_test_split, you can achieve the same functionality by combining the sample() method with some data manipulation. for example:

const dfd = require('danfojs-node')

async function train_test_split(df, testSize = 0.2, seed = 42) {
    // Calculate the number of samples for test
    const numTest = Math.round(df.shape[0] * testSize)
    const numTrain = df.shape[0] - numTest
    
    // Randomly sample for test set
    const testDf = await df.sample(numTest, { seed: seed })
    
    // Get the indices not in test set for training
    const testIndices = new Set(testDf.index)
    const trainIndices = df.index.filter(idx => !testIndices.has(idx))
    
    // Create train set using iloc
    const trainDf = df.iloc({ rows: trainIndices })
    
    return [trainDf, testDf]
}

// Example usage:
async function example() {
    // Create or load your DataFrame
    const data = {
        'feature1': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        'feature2': [10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
        'target': [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
    }
    const df = new dfd.DataFrame(data)
    
    // Split features and target
    const X = df.drop({ columns: ['target'] })
    const y = df['target']
    
    // Perform train-test split
    const [X_train, X_test] = await train_test_split(X)
    const [y_train, y_test] = await train_test_split(y)
    
    // Now you have your training and testing sets
    console.log("Training set shape:", X_train.shape)
    console.log("Testing set shape:", X_test.shape)
}