itertools_num/
linspace.rs

1use num_traits::Float;
2
3/// An iterator of a sequence of evenly spaced floats.
4///
5/// Iterator element type is `F`.
6#[derive(Clone, Debug)]
7pub struct Linspace<F> {
8    start: F,
9    step: F,
10    index: usize,
11    len: usize,
12}
13
14impl<F> Iterator for Linspace<F>
15    where F: Float
16{
17    type Item = F;
18
19    #[inline]
20    fn next(&mut self) -> Option<F> {
21        if self.index >= self.len {
22            None
23        } else {
24            // Calculate the value just like numpy.linspace does
25            let i = self.index;
26            self.index += 1;
27            Some(self.start + self.step * F::from(i).unwrap())
28        }
29    }
30
31    #[inline]
32    fn size_hint(&self) -> (usize, Option<usize>) {
33        let n = self.len - self.index;
34        (n, Some(n))
35    }
36}
37
38impl<F> DoubleEndedIterator for Linspace<F>
39    where F: Float,
40{
41    #[inline]
42    fn next_back(&mut self) -> Option<F> {
43        if self.index >= self.len {
44            None
45        } else {
46            // Calculate the value just like numpy.linspace does
47            self.len -= 1;
48            let i = self.len;
49            Some(self.start + self.step * F::from(i).unwrap())
50        }
51    }
52}
53
54impl<F> ExactSizeIterator for Linspace<F>
55    where Linspace<F>: Iterator
56{}
57
58/// Return an iterator of evenly spaced floats.
59///
60/// The `Linspace` has `n` elements, where the first
61/// element is `a` and the last element is `b`.
62///
63/// Iterator element type is `F`, where `F` must be
64/// either `f32` or `f64`.
65///
66/// ```
67/// extern crate itertools;
68/// extern crate itertools_num;
69///
70/// use itertools_num::linspace;
71///
72/// # fn main() {
73/// itertools::assert_equal(linspace::<f32>(0., 1., 5),
74///                         vec![0., 0.25, 0.5, 0.75, 1.0]);
75/// # }
76/// ```
77#[inline]
78pub fn linspace<F>(a: F, b: F, n: usize) -> Linspace<F>
79    where F: Float
80{
81    let step = if n > 1 {
82        let nf: F = F::from(n).unwrap();
83        (b - a) / (nf - F::one())
84    } else {
85        F::zero()
86    };
87    Linspace {
88        start: a,
89        step: step,
90        index: 0,
91        len: n,
92    }
93}