src/lingrid.rs

/*!
Linear grids.

These are multi-dimensional intervals $\prod_{i=1}^N [a_i, b_i]$ divided along each dimension
into n_i equally-spaced nodes, with $a_i$ the first node and $b_i$ the last node along each
dimension.

The [`LinSpace`]s provided by this module are similar to [`num::range_step_inclusive`], but as an
iterator they are [restartable][RestartableIterator] and parametrised by the number of nodes
instead of a step. This way it can be ensured that $a_i$ and $b_i$ are the last and the first node.

The starting points for the use of this module are the [`linspace`], [`lingrid`], and
[`lingrid_centered`] functions. They return a [`LinSpace`]s that implements [`IntoIterator`] for
iteration over the grid. Additional utility functions are in the [`Grid`] trait.
*/

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

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

/// An abstraction of possibly multi-dimensional linear grids.
///
/// `U` is typically a `F` for a `Float` `F` for one-dimensional grids created by `linspace`,
/// or [`Loc`]`<F, N>` for multi-dimensional grids created by `lingrid`.
/// In the first case `count` of nodes is `usize`, and in the second case `[usize; N]`.
#[derive(Clone, Copy, Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct LinSpace<U, I> {
    pub start : U,
    pub end : U,
    pub count : I,
}

/// A `N`-dimensional interval divided into an indicated number of equally-spaced nodes along
/// each dimension.
#[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 }
}

/// Creates 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>,
}

/// Abstraction of a linear grid over space `U` with multi-dimensional index set `I`.
pub trait Grid<U, I> {
    /// Converts a linear index `i` into a grid point.
    fn entry_linear_unchecked(&self, i : usize) -> U;
    // Converts a multi-dimensional index `i` into a grid point.
    fn entry_unchecked(&self, i : &I) -> U;

    // fn entry(&self, i : I) -> Option<F>
}

/// Helper trait for iteration of [`Grid`]s.
pub trait GridIteration<F, I> {
    /// Returns the next multi-dimensional index (not yet converted into grid point).
    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()
    }
}