CODEWITHSUNDEEP
rustrust-ndarray
Rufus
Rufus

Reputation: 5566

linspace for arbitrary length 1D array

In numpy, one can do

>>> np.linspace([1,2,3], [4,5,6], 4)
array([[1., 2., 3.],
       [2., 3., 4.],
       [3., 4., 5.],
       [4., 5., 6.]])

However, it seems this isn't available from rust ndarray's linspace which only accepts Float.

How do I generalize rust-ndarray's linspace to accept 1D arrays (Array1<T>) of arbitrary size?

Upvotes: 0

Views: 109

Answers (1)

Rufus
Rufus

Reputation: 5566

May not be optimized, but one way is the following

fn linspace<A:Float>(start: Array1<A>, end: Array1<A>, n: usize) -> Array<A, Dim<[usize; 2]>>{
  assert!(start.len() == end.len());
  Array::from_shape_vec((start.len(), n), (0..start.len()).map(|idx| Array::linspace(start[idx], end[idx], n)).flatten().collect()).unwrap().reversed_axes()
}

Basically,

  1. assert!(start.len() == end.len()); asserts same length
  2. (start.len(), n), (0..start.len()).map(|idx| Array::linspace(start[idx], end[idx], n)) repeatedly calls Array::linspace on each index which results in a nested Vec<Array<...>>
  3. .flatten().collect() flattens the Vec<Array<...>> into a 1D Vec<...>
  4. Array::from_shape_vec().unwrap() converts the Vec<...> into a (start.len(), n) sized Array (this is required since flatten arranges the elements as such)
  5. .reversed_axes() transposes the array into shape (n, start.len()) as is the case in numpy's linspace

This passes the following test

mod tests {
  use ndarray::array;

  use super::*;

  #[test]
  fn test_linspace() {
    assert_eq!(linspace(array![1.,2.,3.], array![4.,5.,6.], 4), array![[1.,2.,3.],[2., 3., 4.], [3.,4.,5.], [4.,5.,6.]])
  }

}

Upvotes: 0

Related Questions