Add range implementations of linspace and logspace

Author: Kwonunn

benches/construct.rs

@@ -21,7 +21,7 @@ fn zeros_f64(bench: &mut Bencher)
 #[bench]
 fn map_regular(bench: &mut test::Bencher)
 {
-    let a = Array::linspace(0., 127., 128)
+    let a = Array::linspace(0.0..=127.0, 128)
         .into_shape_with_order((8, 16))
         .unwrap();
     bench.iter(|| a.map(|&x| 2. * x));
@@ -31,7 +31,7 @@ fn map_regular(bench: &mut test::Bencher)
 #[bench]
 fn map_stride(bench: &mut test::Bencher)
 {
-    let a = Array::linspace(0., 127., 256)
+    let a = Array::linspace(0.0..=127.0, 256)
         .into_shape_with_order((8, 32))
         .unwrap();
     let av = a.slice(s![.., ..;2]);

src/impl_constructors.rs

@@ -44,7 +44,8 @@ use rawpointer::PointerExt;
 ///
 /// ## Constructor methods for one-dimensional arrays.
 impl<S, A> ArrayBase<S, Ix1>
-where S: DataOwned<Elem = A>
+where
+    S: DataOwned<Elem = A>,
 {
     /// Create a one-dimensional array from a vector (no copying needed).
     ///
@@ -55,13 +56,9 @@ where S: DataOwned<Elem = A>
     ///
     /// let array = Array::from_vec(vec![1., 2., 3., 4.]);
     /// ```
-    pub fn from_vec(v: Vec<A>) -> Self
-    {
+    pub fn from_vec(v: Vec<A>) -> Self {
         if mem::size_of::<A>() == 0 {
-            assert!(
-                v.len() <= isize::MAX as usize,
-                "Length must fit in `isize`.",
-            );
+            assert!(v.len() <= isize::MAX as usize, "Length must fit in `isize`.",);
         }
         unsafe { Self::from_shape_vec_unchecked(v.len() as Ix, v) }
     }
@@ -76,8 +73,7 @@ where S: DataOwned<Elem = A>
     /// let array = Array::from_iter(0..10);
     /// ```
     #[allow(clippy::should_implement_trait)]
-    pub fn from_iter<I: IntoIterator<Item = A>>(iterable: I) -> Self
-    {
+    pub fn from_iter<I: IntoIterator<Item = A>>(iterable: I) -> Self {
         Self::from_vec(iterable.into_iter().collect())
     }
 
@@ -99,10 +95,12 @@ where S: DataOwned<Elem = A>
     /// assert!(array == arr1(&[0.0, 0.25, 0.5, 0.75, 1.0]))
     /// ```
     #[cfg(feature = "std")]
-    pub fn linspace(start: A, end: A, n: usize) -> Self
-    where A: Float
+    pub fn linspace<R>(range: R, n: usize) -> Self
+    where
+        R: std::ops::RangeBounds<A>,
+        A: Float,
     {
-        Self::from(to_vec(linspace::linspace(start, end, n)))
+        Self::from(to_vec(linspace::linspace(range, n)))
     }
 
     /// Create a one-dimensional array with elements from `start` to `end`
@@ -118,7 +116,8 @@ where S: DataOwned<Elem = A>
     /// ```
     #[cfg(feature = "std")]
     pub fn range(start: A, end: A, step: A) -> Self
-    where A: Float
+    where
+        A: Float,
     {
         Self::from(to_vec(linspace::range(start, end, step)))
     }
@@ -145,10 +144,12 @@ where S: DataOwned<Elem = A>
     /// # }
     /// ```
     #[cfg(feature = "std")]
-    pub fn logspace(base: A, start: A, end: A, n: usize) -> Self
-    where A: Float
+    pub fn logspace<R>(base: A, range: R, n: usize) -> Self
+    where
+        R: std::ops::RangeBounds<A>,
+        A: Float,
     {
-        Self::from(to_vec(logspace::logspace(base, start, end, n)))
+        Self::from(to_vec(logspace::logspace(base, range, n)))
     }
 
     /// Create a one-dimensional array with `n` geometrically spaced elements
@@ -180,15 +181,17 @@ where S: DataOwned<Elem = A>
     /// ```
     #[cfg(feature = "std")]
     pub fn geomspace(start: A, end: A, n: usize) -> Option<Self>
-    where A: Float
+    where
+        A: Float,
     {
         Some(Self::from(to_vec(geomspace::geomspace(start, end, n)?)))
     }
 }
 
 /// ## Constructor methods for two-dimensional arrays.
 impl<S, A> ArrayBase<S, Ix2>
-where S: DataOwned<Elem = A>
+where
+    S: DataOwned<Elem = A>,
 {
     /// Create an identity matrix of size `n` (square 2D array).
     ///
@@ -470,14 +473,14 @@ where
     /// );
     /// ```
     pub fn from_shape_vec<Sh>(shape: Sh, v: Vec<A>) -> Result<Self, ShapeError>
-    where Sh: Into<StrideShape<D>>
+    where
+        Sh: Into<StrideShape<D>>,
     {
         // eliminate the type parameter Sh as soon as possible
         Self::from_shape_vec_impl(shape.into(), v)
     }
 
-    fn from_shape_vec_impl(shape: StrideShape<D>, v: Vec<A>) -> Result<Self, ShapeError>
-    {
+    fn from_shape_vec_impl(shape: StrideShape<D>, v: Vec<A>) -> Result<Self, ShapeError> {
         let dim = shape.dim;
         let is_custom = shape.strides.is_custom();
         dimension::can_index_slice_with_strides(&v, &dim, &shape.strides, dimension::CanIndexCheckMode::OwnedMutable)?;
@@ -513,16 +516,16 @@ where
     /// 5. The strides must not allow any element to be referenced by two different
     ///    indices.
     pub unsafe fn from_shape_vec_unchecked<Sh>(shape: Sh, v: Vec<A>) -> Self
-    where Sh: Into<StrideShape<D>>
+    where
+        Sh: Into<StrideShape<D>>,
     {
         let shape = shape.into();
         let dim = shape.dim;
         let strides = shape.strides.strides_for_dim(&dim);
         Self::from_vec_dim_stride_unchecked(dim, strides, v)
     }
 
-    unsafe fn from_vec_dim_stride_unchecked(dim: D, strides: D, mut v: Vec<A>) -> Self
-    {
+    unsafe fn from_vec_dim_stride_unchecked(dim: D, strides: D, mut v: Vec<A>) -> Self {
         // debug check for issues that indicates wrong use of this constructor
         debug_assert!(dimension::can_index_slice(&v, &dim, &strides, CanIndexCheckMode::OwnedMutable).is_ok());
 
@@ -595,7 +598,8 @@ where
     /// # let _ = shift_by_two;
     /// ```
     pub fn uninit<Sh>(shape: Sh) -> ArrayBase<S::MaybeUninit, D>
-    where Sh: ShapeBuilder<Dim = D>
+    where
+        Sh: ShapeBuilder<Dim = D>,
     {
         unsafe {
             let shape = shape.into_shape_with_order();

src/lib.rs

@@ -201,7 +201,7 @@ mod layout;
 mod linalg_traits;
 mod linspace;
 #[cfg(feature = "std")]
-pub use crate::linspace::{linspace, linspace_exclusive, range, Linspace};
+pub use crate::linspace::{linspace, range, Linspace};
 mod logspace;
 #[cfg(feature = "std")]
 pub use crate::logspace::{logspace, Logspace};

src/linspace.rs

@@ -6,27 +6,29 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 #![cfg(feature = "std")]
+
+use std::ops::{Bound, RangeBounds};
+
 use num_traits::Float;
 
 /// An iterator of a sequence of evenly spaced floats.
 ///
 /// Iterator element type is `F`.
-pub struct Linspace<F>
-{
+pub struct Linspace<F> {
     start: F,
     step: F,
     index: usize,
     len: usize,
 }
 
 impl<F> Iterator for Linspace<F>
-where F: Float
+where
+    F: Float,
 {
     type Item = F;
 
     #[inline]
-    fn next(&mut self) -> Option<F>
-    {
+    fn next(&mut self) -> Option<F> {
         if self.index >= self.len {
             None
         } else {
@@ -38,19 +40,18 @@ where F: Float
     }
 
     #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>)
-    {
+    fn size_hint(&self) -> (usize, Option<usize>) {
         let n = self.len - self.index;
         (n, Some(n))
     }
 }
 
 impl<F> DoubleEndedIterator for Linspace<F>
-where F: Float
+where
+    F: Float,
 {
     #[inline]
-    fn next_back(&mut self) -> Option<F>
-    {
+    fn next_back(&mut self) -> Option<F> {
         if self.index >= self.len {
             None
         } else {
@@ -71,43 +72,31 @@ impl<F> ExactSizeIterator for Linspace<F> where Linspace<F>: Iterator {}
 /// The iterator element type is `F`, where `F` must implement [`Float`], e.g.
 /// [`f32`] or [`f64`].
 ///
-/// **Panics** if converting `n - 1` to type `F` fails.
+/// ## Panics
+/// - If called with a range type other than `a..b` or `a..=b`.
+/// - If converting `n` to type `F` fails.
 #[inline]
-pub fn linspace<F>(a: F, b: F, n: usize) -> Linspace<F>
-where F: Float
+pub fn linspace<R, F>(range: R, n: usize) -> Linspace<F>
+where
+    R: RangeBounds<F>,
+    F: Float,
 {
-    let step = if n > 1 {
-        let num_steps = F::from(n - 1).expect("Converting number of steps to `A` must not fail.");
-        (b - a) / num_steps
-    } else {
-        F::zero()
+    let (a, b, num_steps) = match (range.start_bound(), range.end_bound()) {
+        (Bound::Included(a), Bound::Included(b)) => {
+            (*a, *b, F::from(n - 1).expect("Converting number of steps to `A` must not fail."))
+        }
+        (Bound::Included(a), Bound::Excluded(b)) => {
+            (*a, *b, F::from(n).expect("Converting number of steps to `A` must not fail."))
+        }
+        _ => panic!("Only a..b and a..=b ranges are supported."),
     };
-    Linspace {
-        start: a,
-        step,
-        index: 0,
-        len: n,
-    }
-}
 
-/// Return an iterator of evenly spaced floats.
-///
-/// The `Linspace` has `n` elements from `a` to `b` (exclusive).
-///
-/// The iterator element type is `F`, where `F` must implement [`Float`], e.g.
-/// [`f32`] or [`f64`].
-///
-/// **Panics** if converting `n` to type `F` fails.
-#[inline]
-pub fn linspace_exclusive<F>(a: F, b: F, n: usize) -> Linspace<F>
-where F: Float
-{
-    let step = if n > 1 {
-        let num_steps = F::from(n).expect("Converting number of steps to `A` must not fail.");
+    let step = if num_steps > F::zero() {
         (b - a) / num_steps
     } else {
         F::zero()
     };
+
     Linspace {
         start: a,
         step,
@@ -127,7 +116,8 @@ where F: Float
 /// **Panics** if converting `((b - a) / step).ceil()` to type `F` fails.
 #[inline]
 pub fn range<F>(a: F, b: F, step: F) -> Linspace<F>
-where F: Float
+where
+    F: Float,
 {
     let len = b - a;
     let steps = F::ceil(len / step);

src/logspace.rs

@@ -6,6 +6,8 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 #![cfg(feature = "std")]
+
+use std::ops::{Bound, RangeBounds};
 use num_traits::Float;
 
 /// An iterator of a sequence of logarithmically spaced number.
@@ -79,15 +81,27 @@ impl<F> ExactSizeIterator for Logspace<F> where Logspace<F>: Iterator {}
 ///
 /// **Panics** if converting `n - 1` to type `F` fails.
 #[inline]
-pub fn logspace<F>(base: F, a: F, b: F, n: usize) -> Logspace<F>
-where F: Float
+pub fn logspace<R, F>(base: F, range: R, n: usize) -> Logspace<F>
+where 
+    R: RangeBounds<F>,
+    F: Float,
 {
-    let step = if n > 1 {
-        let num_steps = F::from(n - 1).expect("Converting number of steps to `A` must not fail.");
+    let (a, b, num_steps) = match (range.start_bound(), range.end_bound()) {
+        (Bound::Included(a), Bound::Included(b)) => {
+            (*a, *b, F::from(n - 1).expect("Converting number of steps to `A` must not fail."))
+        }
+        (Bound::Included(a), Bound::Excluded(b)) => {
+            (*a, *b, F::from(n).expect("Converting number of steps to `A` must not fail."))
+        }
+        _ => panic!("Only a..b and a..=b ranges are supported."),
+    };
+
+    let step = if num_steps > F::zero() {
         (b - a) / num_steps
     } else {
         F::zero()
     };
+
     Logspace {
         sign: base.signum(),
         base: base.abs(),
@@ -110,23 +124,23 @@ mod tests
         use crate::{arr1, Array1};
         use approx::assert_abs_diff_eq;
 
-        let array: Array1<_> = logspace(10.0, 0.0, 3.0, 4).collect();
+        let array: Array1<_> = logspace(10.0, 0.0..=3.0, 4).collect();
         assert_abs_diff_eq!(array, arr1(&[1e0, 1e1, 1e2, 1e3]), epsilon = 1e-12);
 
-        let array: Array1<_> = logspace(10.0, 3.0, 0.0, 4).collect();
+        let array: Array1<_> = logspace(10.0, 3.0..=0.0, 4).collect();
         assert_abs_diff_eq!(array, arr1(&[1e3, 1e2, 1e1, 1e0]), epsilon = 1e-12);
 
-        let array: Array1<_> = logspace(-10.0, 3.0, 0.0, 4).collect();
+        let array: Array1<_> = logspace(-10.0, 3.0..=0.0, 4).collect();
         assert_abs_diff_eq!(array, arr1(&[-1e3, -1e2, -1e1, -1e0]), epsilon = 1e-12);
 
-        let array: Array1<_> = logspace(-10.0, 0.0, 3.0, 4).collect();
+        let array: Array1<_> = logspace(-10.0, 0.0..=3.0, 4).collect();
         assert_abs_diff_eq!(array, arr1(&[-1e0, -1e1, -1e2, -1e3]), epsilon = 1e-12);
     }
 
     #[test]
     fn iter_forward()
     {
-        let mut iter = logspace(10.0f64, 0.0, 3.0, 4);
+        let mut iter = logspace(10.0f64, 0.0..=3.0, 4);
 
         assert!(iter.size_hint() == (4, Some(4)));
 
@@ -142,7 +156,7 @@ mod tests
     #[test]
     fn iter_backward()
     {
-        let mut iter = logspace(10.0f64, 0.0, 3.0, 4);
+        let mut iter = logspace(10.0f64, 0.0..=3.0, 4);
 
         assert!(iter.size_hint() == (4, Some(4)));