1#![cfg(feature = "std")]
9use num_traits::Float;
10
11pub struct Linspace<F> {
15 start: F,
16 step: F,
17 index: usize,
18 len: usize,
19}
20
21impl<F> Iterator for Linspace<F>
22where
23 F: Float,
24{
25 type Item = F;
26
27 #[inline]
28 fn next(&mut self) -> Option<F> {
29 if self.index >= self.len {
30 None
31 } else {
32 let i = self.index;
34 self.index += 1;
35 Some(self.start + self.step * F::from(i).unwrap())
36 }
37 }
38
39 #[inline]
40 fn size_hint(&self) -> (usize, Option<usize>) {
41 let n = self.len - self.index;
42 (n, Some(n))
43 }
44}
45
46impl<F> DoubleEndedIterator for Linspace<F>
47where
48 F: Float,
49{
50 #[inline]
51 fn next_back(&mut self) -> Option<F> {
52 if self.index >= self.len {
53 None
54 } else {
55 self.len -= 1;
57 let i = self.len;
58 Some(self.start + self.step * F::from(i).unwrap())
59 }
60 }
61}
62
63impl<F> ExactSizeIterator for Linspace<F> where Linspace<F>: Iterator {}
64
65#[inline]
74pub fn linspace<F>(a: F, b: F, n: usize) -> Linspace<F>
75where
76 F: Float,
77{
78 let step = if n > 1 {
79 let num_steps = F::from(n - 1).expect("Converting number of steps to `A` must not fail.");
80 (b - a) / num_steps
81 } else {
82 F::zero()
83 };
84 Linspace {
85 start: a,
86 step,
87 index: 0,
88 len: n,
89 }
90}
91
92#[inline]
102pub fn range<F>(a: F, b: F, step: F) -> Linspace<F>
103where
104 F: Float,
105{
106 let len = b - a;
107 let steps = F::ceil(len / step);
108 Linspace {
109 start: a,
110 step,
111 len: steps.to_usize().expect(
112 "Converting the length to `usize` must not fail. The most likely \
113 cause of this failure is if the sign of `end - start` is \
114 different from the sign of `step`.",
115 ),
116 index: 0,
117 }
118}