1 // Copyright 2014-2016 bluss and ndarray developers.
2 //
3 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6 // option. This file may not be copied, modified, or distributed
7 // except according to those terms.
8 #![cfg(feature = "std")]
9 use num_traits::Float;
10
11 /// An iterator of a sequence of geometrically spaced floats.
12 ///
13 /// Iterator element type is `F`.
14 pub struct Geomspace<F> {
15 sign: F,
16 start: F,
17 step: F,
18 index: usize,
19 len: usize,
20 }
21
22 impl<F> Iterator for Geomspace<F>
23 where
24 F: Float,
25 {
26 type Item = F;
27
28 #[inline]
next(&mut self) -> Option<F>29 fn next(&mut self) -> Option<F> {
30 if self.index >= self.len {
31 None
32 } else {
33 // Calculate the value just like numpy.linspace does
34 let i = self.index;
35 self.index += 1;
36 let exponent = self.start + self.step * F::from(i).unwrap();
37 Some(self.sign * exponent.exp())
38 }
39 }
40
41 #[inline]
size_hint(&self) -> (usize, Option<usize>)42 fn size_hint(&self) -> (usize, Option<usize>) {
43 let n = self.len - self.index;
44 (n, Some(n))
45 }
46 }
47
48 impl<F> DoubleEndedIterator for Geomspace<F>
49 where
50 F: Float,
51 {
52 #[inline]
next_back(&mut self) -> Option<F>53 fn next_back(&mut self) -> Option<F> {
54 if self.index >= self.len {
55 None
56 } else {
57 // Calculate the value just like numpy.linspace does
58 self.len -= 1;
59 let i = self.len;
60 let exponent = self.start + self.step * F::from(i).unwrap();
61 Some(self.sign * exponent.exp())
62 }
63 }
64 }
65
66 impl<F> ExactSizeIterator for Geomspace<F> where Geomspace<F>: Iterator {}
67
68 /// An iterator of a sequence of geometrically spaced values.
69 ///
70 /// The `Geomspace` has `n` geometrically spaced elements from `start` to `end`
71 /// (inclusive).
72 ///
73 /// The iterator element type is `F`, where `F` must implement `Float`, e.g.
74 /// `f32` or `f64`.
75 ///
76 /// Returns `None` if `start` and `end` have different signs or if either one
77 /// is zero. Conceptually, this means that in order to obtain a `Some` result,
78 /// `end / start` must be positive.
79 ///
80 /// **Panics** if converting `n - 1` to type `F` fails.
81 #[inline]
geomspace<F>(a: F, b: F, n: usize) -> Option<Geomspace<F>> where F: Float,82 pub fn geomspace<F>(a: F, b: F, n: usize) -> Option<Geomspace<F>>
83 where
84 F: Float,
85 {
86 if a == F::zero() || b == F::zero() || a.is_sign_negative() != b.is_sign_negative() {
87 return None;
88 }
89 let log_a = a.abs().ln();
90 let log_b = b.abs().ln();
91 let step = if n > 1 {
92 let num_steps = F::from(n - 1).expect("Converting number of steps to `A` must not fail.");
93 (log_b - log_a) / num_steps
94 } else {
95 F::zero()
96 };
97 Some(Geomspace {
98 sign: a.signum(),
99 start: log_a,
100 step,
101 index: 0,
102 len: n,
103 })
104 }
105
106 #[cfg(test)]
107 mod tests {
108 use super::geomspace;
109
110 #[test]
111 #[cfg(feature = "approx")]
valid()112 fn valid() {
113 use crate::{arr1, Array1};
114 use approx::assert_abs_diff_eq;
115
116 let array: Array1<_> = geomspace(1e0, 1e3, 4).unwrap().collect();
117 assert_abs_diff_eq!(array, arr1(&[1e0, 1e1, 1e2, 1e3]), epsilon = 1e-12);
118
119 let array: Array1<_> = geomspace(1e3, 1e0, 4).unwrap().collect();
120 assert_abs_diff_eq!(array, arr1(&[1e3, 1e2, 1e1, 1e0]), epsilon = 1e-12);
121
122 let array: Array1<_> = geomspace(-1e3, -1e0, 4).unwrap().collect();
123 assert_abs_diff_eq!(array, arr1(&[-1e3, -1e2, -1e1, -1e0]), epsilon = 1e-12);
124
125 let array: Array1<_> = geomspace(-1e0, -1e3, 4).unwrap().collect();
126 assert_abs_diff_eq!(array, arr1(&[-1e0, -1e1, -1e2, -1e3]), epsilon = 1e-12);
127 }
128
129 #[test]
iter_forward()130 fn iter_forward() {
131 let mut iter = geomspace(1.0f64, 1e3, 4).unwrap();
132
133 assert!(iter.size_hint() == (4, Some(4)));
134
135 assert!((iter.next().unwrap() - 1e0).abs() < 1e-5);
136 assert!((iter.next().unwrap() - 1e1).abs() < 1e-5);
137 assert!((iter.next().unwrap() - 1e2).abs() < 1e-5);
138 assert!((iter.next().unwrap() - 1e3).abs() < 1e-5);
139 assert!(iter.next().is_none());
140
141 assert!(iter.size_hint() == (0, Some(0)));
142 }
143
144 #[test]
iter_backward()145 fn iter_backward() {
146 let mut iter = geomspace(1.0f64, 1e3, 4).unwrap();
147
148 assert!(iter.size_hint() == (4, Some(4)));
149
150 assert!((iter.next_back().unwrap() - 1e3).abs() < 1e-5);
151 assert!((iter.next_back().unwrap() - 1e2).abs() < 1e-5);
152 assert!((iter.next_back().unwrap() - 1e1).abs() < 1e-5);
153 assert!((iter.next_back().unwrap() - 1e0).abs() < 1e-5);
154 assert!(iter.next_back().is_none());
155
156 assert!(iter.size_hint() == (0, Some(0)));
157 }
158
159 #[test]
zero_lower()160 fn zero_lower() {
161 assert!(geomspace(0.0, 1.0, 4).is_none());
162 }
163
164 #[test]
zero_upper()165 fn zero_upper() {
166 assert!(geomspace(1.0, 0.0, 4).is_none());
167 }
168
169 #[test]
zero_included()170 fn zero_included() {
171 assert!(geomspace(-1.0, 1.0, 4).is_none());
172 }
173 }
174