Function train_test_split

pub fn train_test_split(
    dataset: &Dataset,
    test_size: f64,
    random_seed: Option<u64>,
) -> Result<(Dataset, Dataset)>

Split a dataset into training and test sets

This function creates a random split of the dataset while preserving the metadata and feature information in both resulting datasets.

Arguments

• dataset - The dataset to split
• test_size - Fraction of samples to include in test set (0.0 to 1.0)
• random_seed - Optional random seed for reproducible splits

Returns

A tuple of (train_dataset, test_dataset)

Examples

use ndarray::Array2;
use scirs2_datasets::utils::{Dataset, train_test_split};

let data = Array2::from_shape_vec((10, 3), (0..30).map(|x| x as f64).collect()).unwrap();
let dataset = Dataset::new(data, None);

let (train, test) = train_test_split(&dataset, 0.3, Some(42)).unwrap();
assert_eq!(train.n_samples() + test.n_samples(), 10);

Examples found in repository

examples/dataset_loaders.rs (line 31)

fn main() {
    // Check if a CSV file is provided as a command-line argument
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("Usage: {} <path_to_csv_file>", args[0]);
        println!("Example: {} examples/sample_data.csv", args[0]);
        return;
    }

    let file_path = &args[1];

    // Verify the file exists
    if !Path::new(file_path).exists() {
        println!("Error: File '{}' does not exist", file_path);
        return;
    }

    // Load CSV file
    println!("Loading CSV file: {}", file_path);
    match loaders::load_csv(file_path, true, None) {
        Ok(dataset) => {
            print_dataset_info(&dataset, "Loaded CSV");

            // Split the dataset for demonstration
            println!("\nDemonstrating train-test split...");
            match train_test_split(&dataset, 0.2, Some(42)) {
                Ok((train, test)) => {
                    println!("Training set: {} samples", train.n_samples());
                    println!("Test set: {} samples", test.n_samples());

                    // Save as JSON for demonstration
                    let json_path = format!("{}.json", file_path);
                    println!("\nSaving training dataset to JSON: {}", json_path);
                    if let Err(e) = loaders::save_json(&train, &json_path) {
                        println!("Error saving JSON: {}", e);
                    } else {
                        println!("Successfully saved JSON file");

                        // Load back the JSON file
                        println!("\nLoading back from JSON file...");
                        match loaders::load_json(&json_path) {
                            Ok(loaded) => {
                                print_dataset_info(&loaded, "Loaded JSON");
                            }
                            Err(e) => println!("Error loading JSON: {}", e),
                        }
                    }
                }
                Err(e) => println!("Error splitting dataset: {}", e),
            }
        }
        Err(e) => println!("Error loading CSV: {}", e),
    }
}

examples/data_generators.rs (line 23)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Creating synthetic datasets...\n");

    // Generate classification dataset
    let n_samples = 100;
    let n_features = 5;

    let classification_data = make_classification(
        n_samples,
        n_features,
        3,        // 3 classes
        2,        // 2 clusters per class
        3,        // 3 informative features
        Some(42), // random seed
    )?;

    // Train-test split
    let (train, test) = train_test_split(&classification_data, 0.2, Some(42))?;

    println!("Classification dataset:");
    println!("  Total samples: {}", classification_data.n_samples());
    println!("  Features: {}", classification_data.n_features());
    println!("  Training samples: {}", train.n_samples());
    println!("  Test samples: {}", test.n_samples());

    // Generate regression dataset
    let regression_data = make_regression(
        n_samples,
        n_features,
        3,   // 3 informative features
        0.5, // noise level
        Some(42),
    )?;

    println!("\nRegression dataset:");
    println!("  Samples: {}", regression_data.n_samples());
    println!("  Features: {}", regression_data.n_features());

    // Normalize the data (in-place)
    let mut data_copy = regression_data.data.clone();
    normalize(&mut data_copy);
    println!("  Data normalized successfully");

    // Generate clustering data (blobs)
    let clustering_data = make_blobs(
        n_samples,
        2,   // 2 features for easy visualization
        4,   // 4 clusters
        0.8, // cluster standard deviation
        Some(42),
    )?;

    println!("\nClustering dataset (blobs):");
    println!("  Samples: {}", clustering_data.n_samples());
    println!("  Features: {}", clustering_data.n_features());

    // Find the number of clusters by finding the max value of target
    let num_clusters = clustering_data.target.as_ref().map_or(0, |t| {
        let mut max_val = -1.0;
        for &val in t.iter() {
            if val > max_val {
                max_val = val;
            }
        }
        (max_val as usize) + 1
    });

    println!("  Clusters: {}", num_clusters);

    // Generate time series data
    let time_series = make_time_series(
        100,  // 100 time steps
        3,    // 3 features/variables
        true, // with trend
        true, // with seasonality
        0.2,  // noise level
        Some(42),
    )?;

    println!("\nTime series dataset:");
    println!("  Time steps: {}", time_series.n_samples());
    println!("  Features: {}", time_series.n_features());

    Ok(())
}