burn_dataset/transform/
window.rs

1use std::{cmp::max, marker::PhantomData, num::NonZeroUsize};
2
3use crate::{Dataset, DatasetError};
4
5/// Functionality to create a window.
6pub trait Window<I> {
7    /// Creates a window of a collection.
8    ///
9    /// # Returns
10    ///
11    /// A `Vec<I>` representing the window, or `None` if the window doesn't fit within the
12    /// collection.
13    fn window(&self, current: usize, size: NonZeroUsize) -> Result<Option<Vec<I>>, DatasetError>;
14}
15
16impl<I, T: Dataset<I> + ?Sized> Window<I> for T {
17    fn window(&self, current: usize, size: NonZeroUsize) -> Result<Option<Vec<I>>, DatasetError> {
18        let end = current + size.get();
19        if end > self.len() {
20            return Ok(None);
21        }
22        let items = self.get_many((current..end).collect())?;
23        Ok(Some(items))
24    }
25}
26
27/// Functionality to create a `WindowsIterator`.
28pub trait Windows<I> {
29    /// Creates and returns an iterator over all the windows of length `size`.
30    fn windows(&self, size: usize) -> WindowsIterator<'_, I>;
31}
32
33impl<I, T: Dataset<I>> Windows<I> for T {
34    /// Is empty if the `Dataset` is shorter than `size`.
35    ///
36    /// # Panics
37    ///
38    /// Panics if `size` is 0.    
39    ///
40    /// # Examples
41    ///
42    /// ```
43    /// use crate::burn_dataset::{
44    ///    transform::{Windows, WindowsDataset},
45    ///    Dataset, InMemDataset,
46    /// };
47    ///
48    /// let items = [1, 2, 3, 4].to_vec();
49    /// let dataset = InMemDataset::new(items.clone());
50    ///
51    /// for window in dataset.windows(2) {
52    ///  // window is a Result<Vec<I>, DatasetError>
53    /// }
54    /// ```
55    fn windows(&self, size: usize) -> WindowsIterator<'_, I> {
56        let size = NonZeroUsize::new(size).expect("window size must be non-zero");
57        WindowsIterator::new(self, size)
58    }
59}
60
61/// Overlapping windows iterator.
62pub struct WindowsIterator<'a, I> {
63    /// The size of the windows.
64    pub size: NonZeroUsize,
65    current: usize,
66    len: usize,
67    dataset: &'a dyn Dataset<I>,
68}
69
70impl<'a, I> WindowsIterator<'a, I> {
71    /// Creates a new `WindowsIterator` instance. The windows overlap.
72    /// Is empty if the input `Dataset` is shorter than `size`.
73    ///
74    /// # Parameters
75    ///
76    /// - `dataset`: The dataset over which windows will be created.
77    /// - `size`: The size of the windows.
78    pub fn new(dataset: &'a dyn Dataset<I>, size: NonZeroUsize) -> Self {
79        let len = max(dataset.len() as isize - size.get() as isize + 1, 0) as usize;
80
81        WindowsIterator {
82            current: 0,
83            len,
84            dataset,
85            size,
86        }
87    }
88}
89
90impl<I> Iterator for WindowsIterator<'_, I> {
91    type Item = Result<Vec<I>, DatasetError>;
92
93    fn next(&mut self) -> Option<Self::Item> {
94        while self.current < self.len {
95            let window = self.dataset.window(self.current, self.size);
96            self.current += 1;
97            match window {
98                Ok(Some(window)) => return Some(Ok(window)),
99                Ok(None) => continue,
100                Err(err) => return Some(Err(err)),
101            }
102        }
103        None
104    }
105}
106
107impl<I> Clone for WindowsIterator<'_, I> {
108    fn clone(&self) -> Self {
109        WindowsIterator {
110            size: self.size,
111            dataset: self.dataset,
112            current: self.current,
113            len: self.len,
114        }
115    }
116}
117
118/// Dataset designed to work with overlapping windows of data.
119pub struct WindowsDataset<D, I> {
120    /// The size of the windows.
121    pub size: NonZeroUsize,
122    dataset: D,
123    input: PhantomData<I>,
124}
125
126impl<D, I> WindowsDataset<D, I>
127where
128    D: Dataset<I>,
129{
130    /// Creates a new `WindowsDataset` instance. The windows overlap.
131    /// Is empty if the input `Dataset` is shorter than `size`.
132    ///
133    /// # Parameters
134    ///
135    /// - `dataset`: The dataset over which windows will be created.
136    /// - `size`: The size of the windows.
137    pub fn new(dataset: D, size: usize) -> Self
138    where
139        D:,
140    {
141        let size = NonZeroUsize::new(size).expect("window size must be non-zero");
142        WindowsDataset::<D, I> {
143            size,
144            dataset,
145            input: PhantomData,
146        }
147    }
148}
149
150impl<D, I> Dataset<Vec<I>> for WindowsDataset<D, I>
151where
152    D: Dataset<I>,
153    I: Send + Sync,
154{
155    /// Retrieves a window of items from the dataset.
156    ///
157    /// # Parameters
158    ///
159    /// - `index`: The index of the window.
160    ///
161    /// # Returns
162    ///
163    /// A vector representing the window.
164    ///
165    /// # Panics
166    ///
167    /// Panics if `index >= len()`.
168    fn get(&self, index: usize) -> Result<Vec<I>, DatasetError> {
169        match self.dataset.window(index, self.size)? {
170            Some(window) => Ok(window),
171            None => panic!(
172                "Index out of bounds for WindowsDataset: {index} >= {}",
173                self.len()
174            ),
175        }
176    }
177
178    /// Retrieves the number of windows in the dataset.
179    ///
180    /// # Returns
181    ///
182    /// A size representing the number of windows.
183    fn len(&self) -> usize {
184        let len = self.dataset.len() as isize - self.size.get() as isize + 1;
185        max(len, 0) as usize
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use rstest::rstest;
192
193    use crate::{
194        Dataset, InMemDataset,
195        transform::{Windows, WindowsDataset},
196    };
197
198    #[rstest]
199    pub fn windows_should_be_equal_to_vec_windows() {
200        let items = [1, 2, 3, 4, 5].to_vec();
201        let dataset = InMemDataset::new(items.clone());
202        let expected = items
203            .windows(3)
204            .map(|x| x.to_vec())
205            .collect::<Vec<Vec<i32>>>();
206
207        let result = dataset
208            .windows(3)
209            .map(Result::unwrap)
210            .collect::<Vec<Vec<i32>>>();
211
212        assert_eq!(result, expected);
213    }
214
215    #[rstest]
216    pub fn windows_dataset_should_be_equal_to_vec_windows() {
217        let items = [1, 2, 3, 4, 5].to_vec();
218        let dataset = InMemDataset::new(items.clone());
219        let expected = items
220            .windows(3)
221            .map(|x| x.to_vec())
222            .collect::<Vec<Vec<i32>>>();
223
224        let result = WindowsDataset::new(dataset, 3)
225            .iter()
226            .map(Result::unwrap)
227            .collect::<Vec<Vec<i32>>>();
228
229        assert_eq!(result, expected);
230    }
231
232    #[rstest]
233    pub fn cloned_iterator_should_be_equal() {
234        let items = [1, 2, 3, 4, 5].to_vec();
235        let dataset = InMemDataset::new(items.clone());
236        let original = dataset.windows(4);
237
238        let cloned = original.clone();
239
240        assert!(std::ptr::eq(cloned.dataset, original.dataset));
241        assert_eq!(cloned.size, original.size);
242        assert_eq!(cloned.current, original.current);
243    }
244
245    #[rstest]
246    pub fn cloned_iterator_should_be_unaffected() {
247        let items = [1, 2, 3, 4, 5].to_vec();
248        let dataset = InMemDataset::new(items.clone());
249        let mut original = dataset.windows(4);
250
251        let cloned = original.clone();
252        original.current = 2;
253
254        assert_ne!(cloned.current, original.current);
255    }
256
257    #[rstest]
258    #[should_panic(expected = "window size must be non-zero")]
259    pub fn windows_should_panic() {
260        let items = [1, 2].to_vec();
261        let dataset = InMemDataset::new(items.clone());
262
263        dataset.windows(0);
264    }
265
266    #[rstest]
267    #[should_panic(expected = "window size must be non-zero")]
268    pub fn new_window_dataset_should_panic() {
269        let items = [1, 2].to_vec();
270        let dataset = InMemDataset::new(items.clone());
271
272        WindowsDataset::new(dataset, 0);
273    }
274
275    #[rstest]
276    pub fn window_dataset_len_should_be_equal() {
277        let dataset = InMemDataset::new([1, 2, 3, 4].to_vec());
278
279        let result = WindowsDataset::new(dataset, 2).len();
280
281        assert_eq!(result, 3);
282    }
283
284    #[rstest]
285    pub fn window_iterator_should_be_empty() {
286        let dataset = InMemDataset::new([1, 2].to_vec());
287        let mut peekable = dataset.windows(4).peekable();
288
289        let result = peekable.peek();
290
291        assert!(result.is_none());
292    }
293
294    #[rstest]
295    pub fn window_dataset_len_should_be_zero() {
296        let dataset = InMemDataset::new([1, 2].to_vec());
297
298        let result = WindowsDataset::new(dataset, 4).len();
299
300        assert_eq!(result, 0);
301    }
302
303    #[rstest]
304    pub fn window_dataset_get_should_be_equal() {
305        let dataset = InMemDataset::new([1, 2, 3, 4].to_vec());
306        let expected = [1, 2, 3].to_vec();
307
308        let result = WindowsDataset::new(dataset, 3).get(0).unwrap();
309
310        assert_eq!(result, expected);
311    }
312
313    #[rstest]
314    #[should_panic(expected = "Index out of bounds for WindowsDataset")]
315    pub fn window_dataset_get_should_be_none() {
316        let dataset = InMemDataset::new([1, 2].to_vec());
317
318        WindowsDataset::new(dataset, 4).get(0).unwrap();
319    }
320}
burn_dataset/transform/window.rs

burn_dataset/transform/
window.rs
