src/lingrid.rs

/// Production of equally spaced nodes within intervals.

use crate::types::*;
use crate::loc::Loc;
use crate::sets::Cube;
use crate::iter::{RestartableIterator, StatefulIterator};
use crate::maputil::{map2, map4};
use serde::Serialize;

// TODO: rewrite this using crate::sets::Cube.

/// The interval `[start, end]` divided into `count` nodes.
/// Implementes `IntoIterator` to iterate over the nodes of the `LinSpace`.
/// It is similar to [`num::range_step_inclusive`], but [restartable][RestartableIterator] and
/// parametrised by the number of nodes instead of a step. This way it can be ensured that the
/// last item produced is equal to `end`. To use other types than `float` and `f32`, create
/// the `LinSpace` struct directly.
#[derive(Clone,Copy,Debug,Serialize,Eq,PartialEq)]
pub struct LinSpace<F, I> {
    pub start : F,
    pub end : F,
    pub count : I,
}

#[allow(type_alias_bounds)] // Need it to access F::CompatibleSize.
pub type LinGrid<F : Float, const N : usize> = LinSpace<Loc<F, N>, [usize; N]>;

/// Creates a [`LinSpace`] on the real line.
pub fn linspace<F : Float>(start : F, end : F, count : usize) -> LinSpace<F, usize> {
    LinSpace{ start : start, end : end, count : count }
}

/// Create a multi-dimensional linear grid.
///
/// The first and last point in each dimension are the boundaries of the corresponding
/// dimensions of `cube`, and there are `count` nodes along each dimension.
pub fn lingrid<F : Float, const N : usize>(
    cube : &Cube<F, N>,
    count : &[usize; N]
) -> LinSpace<Loc<F, N>, [usize; N]> {
    LinSpace{ start : cube.span_start(), end : cube.span_end(), count : *count }
}

/// Create a multi-dimensional linear grid with centered nodes.
///
/// There are `count` along each dimension and each node has equally-sized subcube surrounding it
/// inside `cube`. Thus, if $w_i$ is the width of the cube along dimension $i$, and $n_i$ the number
/// of nodes, the width of the subcube along this dimension is $h_i = w_i/(n_i+1)$, and the first
/// and last nodes are at a distance $h_i/2$ from the closest boundary.
pub fn lingrid_centered<F : Float, const N : usize>(
    cube : &Cube<F, N>,
    count : &[usize; N]
) -> LinSpace<Loc<F, N>, [usize; N]> {
    let h_div_2 = map2(cube.width(), count, |w, &n| w / F::cast_from(2 * (n + 1)));
    let span_start = map2(cube.span_start(), &h_div_2, |a, &t| a + t).into();
    let span_end   = map2(cube.span_end(),   &h_div_2, |b, &t| b - t).into();
    LinSpace{ start : span_start, end : span_end, count : *count }
}


/// Iterator over a `LinSpace`.
#[derive(Clone, Debug)]
pub struct LinSpaceIterator<F, I> {
    lingrid : LinSpace<F,I>,
    current : Option<I>,
}

pub trait Grid<F, I> {
    fn entry_linear_unchecked(&self, i : usize) -> F;
    //fn entry(&self, i : I) -> Option<F>;
    fn entry_unchecked(&self, i : &I) -> F;
}

pub trait GridIteration<F, I> {
    fn next_index(&mut self) -> Option<I>;
}

impl<F : Float + CastFrom<I>, I : Unsigned> Grid<F, I> for LinSpace<F, I> {
    /*fn entry(&self, i : I) -> Option<F> {
         if i < self.count {
            Some(self.entry_unchecked(i))
         } else {
            None
         }
    }*/

    #[inline]
    fn entry_linear_unchecked(&self, i : usize) -> F {
        self.entry_unchecked(&I::cast_from(i))
    }

    #[inline]
    fn entry_unchecked(&self, i : &I) -> F {
        let idx = F::cast_from(*i);
        let scale = F::cast_from(self.count-I::ONE);
        self.start + (self.end-self.start)*idx/scale
    }
}

impl<F : Float + CastFrom<I>, I : Unsigned> GridIteration<F, I>
for LinSpaceIterator<F, I> {
    #[inline]
    fn next_index(&mut self) -> Option<I> {
        match self.current {
            None if I::ZERO < self.lingrid.count     => { self.current = Some(I::ZERO); self.current }
            Some(v) if v+I::ONE < self.lingrid.count => { self.current = Some(v+I::ONE); self.current }
            _                                        => { None }
        }
    }
}

impl<F : Float + CastFrom<I>, I : Unsigned, const N : usize> Grid<Loc<F,N>, [I; N]>
for LinSpace<Loc<F,N>, [I; N]> {
    #[inline]
    fn entry_linear_unchecked(&self, i_ : usize) -> Loc<F, N> {
        let mut i = I::cast_from(i_);
        let mut tmp = [I::ZERO; N];
        for k in 0..N {
            tmp[k] = i % self.count[k];
            i /= self.count[k];
        }
        self.entry_unchecked(&tmp)
    }

    #[inline]
    fn entry_unchecked(&self, i : &[I; N]) -> Loc<F, N> {
        let LinSpace{ start, end, count } = self;
        map4(i, start, end, count, |&ik, &sk, &ek, &ck| {
            let idx = F::cast_from(ik);
            let scale = F::cast_from(ck-I::ONE);
            sk + (ek - sk) * idx / scale
        }).into()
    }
}

impl<F : Float + CastFrom<I>, I : Unsigned, const N : usize> GridIteration<Loc<F,N>, [I; N]>
for LinSpaceIterator<Loc<F,N>, [I; N]> {

    #[inline]
    fn next_index(&mut self) -> Option<[I; N]> {
        match self.current {
            None if self.lingrid.count.iter().all(|v| I::ZERO < *v)  => {
                self.current = Some([I::ZERO; N]);
                self.current
            },
            Some(ref mut v) => {
                for k in 0..N {
                    let a = v[k] + I::ONE;
                    if a < self.lingrid.count[k] {
                        v[k] = a;
                        return self.current
                    } else {
                        v[k] = I::ZERO;
                    }
                }
                None
            },
            _ => None
        }
    }
}

impl<F, I> IntoIterator for LinSpace<F,I>
where LinSpace<F, I> : Grid<F, I>,
      LinSpaceIterator<F, I> : GridIteration<F, I> {
    type Item = F;
    type IntoIter = LinSpaceIterator<F,I>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        LinSpaceIterator { lingrid : self, current : None }
    }
}

impl<F, I> Iterator for LinSpaceIterator<F,I>
where LinSpace<F, I> : Grid<F, I>,
      LinSpaceIterator<F, I> : GridIteration<F, I> {
    type Item = F;
    #[inline]
    fn next(&mut self) -> Option<F> {
        self.next_index().map(|v| self.lingrid.entry_unchecked(&v))
    }
}

impl<F, I> StatefulIterator for LinSpaceIterator<F,I>
where LinSpace<F, I> : Grid<F, I>,
      LinSpaceIterator<F, I> : GridIteration<F, I> {
    #[inline]
    fn current(&self) -> Option<F> {
        self.current.as_ref().map(|c| self.lingrid.entry_unchecked(c))
    }
}


impl<F, I> RestartableIterator for LinSpaceIterator<F,I>
where LinSpace<F, I> : Grid<F, I>,
      LinSpaceIterator<F, I> : GridIteration<F, I> {
    #[inline]
    fn restart(&mut self) -> Option<F> {
        self.current = None;
        self.next()
    }
}