1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

//! This crate achieves the functionality of [sklearn's train_test_split](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html)
//! to generate splits of the data (in this case, a slice), generalized for an
//! arbitrary number of splits. It
//! both provides functions (see [split_parts](https://docs.rs/rand_split/0.2.0/rand_split/fn.split_parts.html)) that work on slices and iterator traits (see [PartsSplit](https://docs.rs/rand_split/0.2.0/rand_split/trait.PartsSplit.html))to work
//! with streams of data.
//!
//! Check out the [examples in the repository](https://github.com/carrascomj/rand_split/tree/trunk/examples)
//! for more information.
use rand::seq::index::sample;

mod stream;
pub use stream::PartsSplit;

mod stream_ttv;
pub use stream_ttv::TTVSplit;

/// Split the elements of a container in randomized sets which contain a
/// a part (in `splits`) of the input.
///
/// # Errors
///
/// Return an error if the parts in `splits` do not sum up to 1.
///
/// # Example
///
/// ```
/// use rand_split::split_parts;
///
/// println!("{:#?}", split_parts(&[1,2,3,4,5,6,8,9,10], &[0.4, 0.2, 0.4]));
/// ```
pub fn split_parts<T>(cont: &[T], splits: &[f32]) -> Result<Vec<Vec<T>>, &'static str>
where
    T: Clone,
{
    if (splits.iter().sum::<f32>() - 1.).abs() > 0.001 {
        return Err("splits must sum 1!");
    }
    let n = cont.len();
    let shuffled = sample(&mut rand::thread_rng(), n, n).into_vec();
    let splits: Vec<usize> = {
        let mut tmp: Vec<usize> = splits
            .iter()
            .map(|x| (x * n as f32) as usize)
            .scan(0, |state, x| {
                *state += x;
                Some(*state)
            })
            .collect();
        // account for rounding errors
        tmp[splits.len() - 1] += n - tmp[tmp.len() - 1];
        tmp
    };
    Ok([0]
        .iter()
        .chain(splits[0..(splits.len() - 1)].iter())
        .zip(splits.iter())
        .map(|(start, end)| {
            shuffled[*start..*end]
                .iter()
                .map(|i| cont[*i].clone())
                .collect::<Vec<T>>()
        })
        .collect())
}

/// Generate train-test splits. Wrapper around [`split_parts`](./split_parts)
/// # Examples
///
/// ```
/// use rand_split::train_test_split;
///
/// let cont = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13];
/// let total_len = cont.len();
/// let result = train_test_split(&cont, 0.8, 0.2).unwrap();
/// assert_eq!(
///     result.iter().map(|inner| inner.len()).sum::<usize>(),
///     total_len
/// );
/// ```
pub fn train_test_split<T>(cont: &[T], train: f32, test: f32) -> Result<Vec<Vec<T>>, &'static str>
where
    T: Clone,
{
    split_parts(cont, &[train, test])
}

/// Generate train-test-validation splits. Wrapper around [`split_parts`](./split_parts)
/// # Examples
///
/// ```
/// use rand_split::ttv_split;
///
/// let cont = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13];
/// let total_len = cont.len();
/// let result = ttv_split(&cont, 0.6, 0.2, 0.2).unwrap();
/// assert_eq!(
///     result.iter().map(|inner| inner.len()).sum::<usize>(),
///     total_len
/// );
/// ```
pub fn ttv_split<T>(
    cont: &[T],
    train: f32,
    test: f32,
    validation: f32,
) -> Result<Vec<Vec<T>>, &'static str>
where
    T: Clone,
{
    split_parts(cont, &[train, test, validation])
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn correct_split_sizes_even_number() {
        let splits = [0.2, 0.7, 0.1];
        let cont = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
        let total_len = cont.len();
        let result = split_parts(&cont, &splits).unwrap();
        assert_eq!(
            result.iter().map(|inner| inner.len()).sum::<usize>(),
            total_len
        );
        assert_eq!(result.len(), splits.len());
    }

    #[test]
    fn correct_split_sizes_odd_number() {
        let splits = [0.2, 0.2, 0.3, 0.3];
        let cont = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13];
        let total_len = cont.len();
        let result = split_parts(&cont, &splits).unwrap();
        assert_eq!(
            result.iter().map(|inner| inner.len()).sum::<usize>(),
            total_len
        );
        assert_eq!(result.len(), splits.len());
    }
}