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::iter::{RestartableIterator, StatefulIterator};
use crate::loc::Loc;
use crate::maputil::{map2, map4};
use crate::sets::Cube;
use crate::types::*;
use serde::{Deserialize, Serialize};

// 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`]`<N, F>` 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<const N: usize, F: Float = f64> = LinSpace<Loc<N, F>, [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<N, F>,
    count: &[usize; N],
) -> LinSpace<Loc<N, F>, [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<N, F>,
    count: &[usize; N],
) -> LinSpace<Loc<N, F>, [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<N, F>, [I; N]>
    for LinSpace<Loc<N, F>, [I; N]>
{
    #[inline]
    fn entry_linear_unchecked(&self, i_: usize) -> Loc<N, F> {
        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<N, F> {
        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<N, F>, [I; N]>
    for LinSpaceIterator<Loc<N, F>, [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()
    }
}