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

//! 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 (see [split_parts](./fn.split_parts.html)).
//!
//! It also provides some handy wrappers to generate [train-test](./fn.train_test_split.html) and
//! [train-test-validation](./fn.ttv_split.html) splits.
use rand::seq::index::sample;

/// 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());
    }
}