Fri, 13 Oct 2023 13:32:46 -0500
feature(binary_heap_retain) is stable now
/*! 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() } }