Wed, 26 Oct 2022 22:16:57 +0300
Added tag unthreaded for changeset d80b87b8acd0
/*! Bisection tree basics, [`BT`] type and the [`BTImpl`] trait. */ use std::slice::{Iter,IterMut}; use std::iter::once; use std::rc::Rc; use serde::{Serialize, Deserialize}; pub(super) use nalgebra::Const; use itertools::izip; use crate::iter::{MapF,Mappable}; use crate::types::{Float, Num}; use crate::coefficients::pow; use crate::maputil::{ array_init, map2, map2_indexed, collect_into_array_unchecked }; use crate::sets::Cube; use crate::loc::Loc; use super::support::*; use super::aggregator::*; /// An enum that indicates whether a [`Node`] of a [`BT`] is uninitialised, leaf, or branch. /// /// For the type and const parametere, see the [module level documentation][super]. #[derive(Clone,Debug)] pub(super) enum NodeOption<F : Num, D, A : Aggregator, const N : usize, const P : usize> { /// Indicates an uninitilised node; may become a branch or a leaf. // TODO: Could optimise Uninitialised away by simply treat Leaf with an empty Vec as // something that can be still replaced with Branches. Uninitialised, /// Indicates a leaf node containing a copy-on-write reference-counted vector /// of data of type `D`. Leaf(Rc<Vec<D>>), /// Indicates a branch node, cotaning a copy-on-write reference to the [`Branches`]. Branches(Rc<Branches<F, D, A, N, P>>), } /// Node of a [`BT`] bisection tree. /// /// For the type and const parameteres, see the [module level documentation][super]. #[derive(Clone,Debug)] pub struct Node<F : Num, D, A : Aggregator, const N : usize, const P : usize> { /// The data or branches under the node. pub(super) data : NodeOption<F, D, A, N, P>, /// Aggregator for `data`. pub(super) aggregator : A, } /// Branching information of a [`Node`] of a [`BT`] bisection tree into `P` subnodes. /// /// For the type and const parameters, see the [module level documentation][super]. #[derive(Clone,Debug)] pub(super) struct Branches<F : Num, D, A : Aggregator, const N : usize, const P : usize> { /// Point for subdivision of the (unstored) [`Cube`] corresponding to the node. pub(super) branch_at : Loc<F, N>, /// Subnodes pub(super) nodes : [Node<F, D, A, N, P>; P], } /// Dirty workaround to broken Rust drop, see [https://github.com/rust-lang/rust/issues/58068](). impl<F : Num, D, A : Aggregator, const N : usize, const P : usize> Drop for Node<F, D, A, N, P> { fn drop(&mut self) { use NodeOption as NO; let process = |brc : Rc<Branches<F, D, A, N, P>>, to_drop : &mut Vec<Rc<Branches<F, D, A, N, P>>>| { // We only drop Branches if we have the only strong reference. Rc::try_unwrap(brc).ok().map(|branches| branches.nodes.map(|mut node| { if let NO::Branches(brc2) = std::mem::replace(&mut node.data, NO::Uninitialised) { to_drop.push(brc2) } })); }; // We mark Self as NodeOption::Uninitialised, extracting the real contents. // If we have subprocess, we need to process them. if let NO::Branches(brc1) = std::mem::replace(&mut self.data, NO::Uninitialised) { // We store a queue of Rc<Branches> to drop into a vector let mut to_drop = Vec::new(); process(brc1, &mut to_drop); // While there are any Branches in the drop queue vector, we continue the process, // pushing all internal branching nodes into the queue. while let Some(brc) = to_drop.pop() { process(brc, &mut to_drop) } } } } /// Trait for the depth of a [`BT`]. /// /// This will generally be either a runtime [`DynamicDepth`] or compile-time [`Const`] depth. pub trait Depth : 'static + Copy + std::fmt::Debug { /// Lower depth type. type Lower : Depth; /// Returns a lower depth, if there still is one. fn lower(&self) -> Option<Self::Lower>; /// Returns a lower depth or self if this is the lowest depth. fn lower_or(&self) -> Self::Lower; } /// Dynamic (runtime) [`Depth`] for a [`BT`]. #[derive(Copy,Clone,Debug,Serialize,Deserialize)] pub struct DynamicDepth( /// The depth pub u8 ); impl Depth for DynamicDepth { type Lower = Self; #[inline] fn lower(&self) -> Option<Self> { if self.0>0 { Some(DynamicDepth(self.0-1)) } else { None } } #[inline] fn lower_or(&self) -> Self { DynamicDepth(if self.0>0 { self.0 - 1 } else { 0 }) } } impl Depth for Const<0> { type Lower = Self; fn lower(&self) -> Option<Self::Lower> { None } fn lower_or(&self) -> Self::Lower { Const } } macro_rules! impl_constdepth { ($($n:literal)*) => { $( impl Depth for Const<$n> { type Lower = Const<{$n-1}>; fn lower(&self) -> Option<Self::Lower> { Some(Const) } fn lower_or(&self) -> Self::Lower { Const } } )* }; } impl_constdepth!(1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32); /// Trait for counting the branching factor of a [`BT`] of dimension `N`. /// /// The const parameter `P` from the [module level documentation][super] is required to satisfy /// `Const<P> : Branchcount<N>`. /// This trait is implemented for `P=pow(2, N)` for small `N`. pub trait BranchCount<const N : usize> {} macro_rules! impl_branchcount { ($($n:literal)*) => { $( impl BranchCount<$n> for Const<{pow(2, $n)}>{} )* } } impl_branchcount!(1 2 3 4 5 6 7 8); impl<F : Float, D, A, const N : usize, const P : usize> Branches<F,D,A,N,P> where Const<P> : BranchCount<N>, A : Aggregator { /// Returns the index in {0, …, `P`-1} for the branch to which the point `x` corresponds. /// /// This only takes the branch subdivision point $d$ into account, so is always succesfull. /// Thus, for this point, each branch corresponds to a quadrant of $ℝ^N$ relative to $d$. fn get_node_index(&self, x : &Loc<F, N>) -> usize { izip!(0..P, x.iter(), self.branch_at.iter()).map(|(i, x_i, branch_i)| if x_i > branch_i { 1<<i } else { 0 } ).sum() } /// Returns the node within `Self` containing the point `x`. /// /// This only takes the branch subdivision point $d$ into account, so is always succesfull. /// Thus, for this point, each branch corresponds to a quadrant of $ℝ^N$ relative to $d$. #[inline] fn get_node(&self, x : &Loc<F,N>) -> &Node<F,D,A,N,P> { &self.nodes[self.get_node_index(x)] } } /// An iterator over the data `D` in a [`BT`]. pub struct BTIter<'a, D> { iter : Iter<'a, D>, } impl<'a, D> Iterator for BTIter<'a,D> { type Item = &'a D; #[inline] fn next(&mut self) -> Option<&'a D> { self.iter.next() } } /// An iterator over the $P=2^N$ subcubes of a [`Cube`] subdivided at a point `d`. pub(super) struct SubcubeIter<'b, F : Float, const N : usize, const P : usize> { domain : &'b Cube<F, N>, branch_at : Loc<F, N>, index : usize, } /// Returns the `i`:th subcube of `domain` subdivided at `branch_at`. #[inline] fn get_subcube<F : Float, const N : usize>( branch_at : &Loc<F, N>, domain : &Cube<F, N>, i : usize ) -> Cube<F, N> { map2_indexed(branch_at, domain, move |j, &branch, &[start, end]| { if i & (1 << j) != 0 { [branch, end] } else { [start, branch] } }).into() } impl<'a, 'b, F : Float, const N : usize, const P : usize> Iterator for SubcubeIter<'b, F, N, P> { type Item = Cube<F, N>; #[inline] fn next(&mut self) -> Option<Self::Item> { if self.index < P { let i = self.index; self.index += 1; Some(get_subcube(&self.branch_at, self.domain, i)) } else { None } } } impl<F : Float, D : Copy, A, const N : usize, const P : usize> Branches<F,D,A,N,P> where Const<P> : BranchCount<N>, A : Aggregator { /// Creates a new node branching structure, subdividing `domain` based on the /// [hint][Support::support_hint] of `support`. pub(super) fn new_with<S : LocalAnalysis <F, A, N>>( domain : &Cube<F,N>, support : &S ) -> Self { let hint = support.bisection_hint(domain); let branch_at = map2(&hint, domain, |h, r| { h.unwrap_or_else(|| (r[0]+r[1])/F::TWO).max(r[0]).min(r[1]) }).into(); Branches{ branch_at : branch_at, nodes : array_init(|| Node::new()), } } /// Returns an iterator over the aggregators of the nodes directly under this branch head. #[inline] pub(super) fn aggregators(&self) -> MapF<Iter<'_, Node<F,D,A,N,P>>, &'_ A> { self.nodes.iter().mapF(Node::get_aggregator) } /// Returns an iterator over the subcubes of `domain` subdivided at the branching point /// of `self`. #[inline] pub(super) fn iter_subcubes<'b>(&self, domain : &'b Cube<F, N>) -> SubcubeIter<'b, F, N, P> { SubcubeIter { domain : domain, branch_at : self.branch_at, index : 0, } } /* /// Returns an iterator over all nodes and corresponding subcubes of `self`. #[inline] pub(super) fn nodes_and_cubes<'a, 'b>(&'a self, domain : &'b Cube<F, N>) -> std::iter::Zip<Iter<'a, Node<F,D,A,N,P>>, SubcubeIter<'b, F, N, P>> { self.nodes.iter().zip(self.iter_subcubes(domain)) } */ /// Mutably iterate over all nodes and corresponding subcubes of `self`. #[inline] pub(super) fn nodes_and_cubes_mut<'a, 'b>(&'a mut self, domain : &'b Cube<F, N>) -> std::iter::Zip<IterMut<'a, Node<F,D,A,N,P>>, SubcubeIter<'b, F, N, P>> { let subcube_iter = self.iter_subcubes(domain); self.nodes.iter_mut().zip(subcube_iter) } /// Insert data into the branch. /// /// The parameters are as follows: /// * `domain` is the cube corresponding to this branch. /// * `d` is the data to be inserted /// * `new_leaf_depth` is the depth relative to `self` at which the data is to be inserted. /// * `support` is the [`Support`] that is used determine with which subcubes of `domain` /// (at subdivision depth `new_leaf_depth`) the data `d` is to be associated with. /// pub(super) fn insert<M : Depth, S : LocalAnalysis<F, A, N>>( &mut self, domain : &Cube<F,N>, d : D, new_leaf_depth : M, support : &S ) { let support_hint = support.support_hint(); for (node, subcube) in self.nodes_and_cubes_mut(&domain) { if support_hint.intersects(&subcube) { node.insert( &subcube, d, new_leaf_depth, support ); } } } /// Construct a new instance of the branch for a different aggregator. /// /// The `generator` is used to convert the data of type `D` of the branch into corresponding /// [`Support`]s. The `domain` is the cube corresponding to `self`. /// The type parameter `ANew´ is the new aggregator, and needs to be implemented for the /// generator's `SupportType`. pub(super) fn convert_aggregator<ANew, G>( self, generator : &G, domain : &Cube<F, N> ) -> Branches<F,D,ANew,N,P> where ANew : Aggregator, G : SupportGenerator<F, N, Id=D>, G::SupportType : LocalAnalysis<F, ANew, N> { let branch_at = self.branch_at; let subcube_iter = self.iter_subcubes(domain); let new_nodes = self.nodes.into_iter().zip(subcube_iter).map(|(node, subcube)| { // TODO: avoid clone node.convert_aggregator(generator, &subcube) }); Branches { branch_at : branch_at, nodes : collect_into_array_unchecked(new_nodes), } } /// Recalculate aggregator after changes to generator. /// /// The `generator` is used to convert the data of type `D` of the branch into corresponding /// [`Support`]s. The `domain` is the cube corresponding to `self`. pub(super) fn refresh_aggregator<G>( &mut self, generator : &G, domain : &Cube<F, N> ) where G : SupportGenerator<F, N, Id=D>, G::SupportType : LocalAnalysis<F, A, N> { for (node, subcube) in self.nodes_and_cubes_mut(domain) { node.refresh_aggregator(generator, &subcube) } } } impl<F : Float, D : Copy, A, const N : usize, const P : usize> Node<F,D,A,N,P> where Const<P> : BranchCount<N>, A : Aggregator { /// Create a new node #[inline] pub(super) fn new() -> Self { Node { data : NodeOption::Uninitialised, aggregator : A::new(), } } /// Get leaf data #[inline] pub(super) fn get_leaf_data(&self, x : &Loc<F, N>) -> Option<&Vec<D>> { match self.data { NodeOption::Uninitialised => None, NodeOption::Leaf(ref data) => Some(data), NodeOption::Branches(ref b) => b.get_node(x).get_leaf_data(x), } } /// Returns a reference to the aggregator of this node #[inline] pub(super) fn get_aggregator(&self) -> &A { &self.aggregator } /// Insert data under the node. /// /// The parameters are as follows: /// * `domain` is the cube corresponding to this branch. /// * `d` is the data to be inserted /// * `new_leaf_depth` is the depth relative to `self` at which new leaves are created. /// * `support` is the [`Support`] that is used determine with which subcubes of `domain` /// (at subdivision depth `new_leaf_depth`) the data `d` is to be associated with. /// /// If `self` is already [`NodeOption::Leaf`], the data is inserted directly in this node. /// If `self` is a [`NodeOption::Branches`], the data is passed to branches whose subcubes /// `support` intersects. If an [`NodeOption::Uninitialised`] node is encountered, a new leaf is /// created at a minimum depth of `new_leaf_depth`. pub(super) fn insert<M : Depth, S : LocalAnalysis <F, A, N>>( &mut self, domain : &Cube<F,N>, d : D, new_leaf_depth : M, support : &S ) { match &mut self.data { NodeOption::Uninitialised => { // Replace uninitialised node with a leaf or a branch self.data = match new_leaf_depth.lower() { None => { let a = support.local_analysis(&domain); self.aggregator.aggregate(once(a)); // TODO: this is currently a dirty hard-coded heuristic; // should add capacity as a parameter let mut vec = Vec::with_capacity(2*P+1); vec.push(d); NodeOption::Leaf(Rc::new(vec)) }, Some(lower) => { let b = Rc::new({ let mut b0 = Branches::new_with(domain, support); b0.insert(domain, d, lower, support); b0 }); self.aggregator.summarise(b.aggregators()); NodeOption::Branches(b) } } }, NodeOption::Leaf(leaf) => { Rc::make_mut(leaf).push(d); let a = support.local_analysis(&domain); self.aggregator.aggregate(once(a)); }, NodeOption::Branches(b) => { Rc::make_mut(b).insert(domain, d, new_leaf_depth.lower_or(), support); self.aggregator.summarise(b.aggregators()); }, } } /// Construct a new instance of the node for a different aggregator /// /// The `generator` is used to convert the data of type `D` of the node into corresponding /// [`Support`]s. The `domain` is the cube corresponding to `self`. /// The type parameter `ANew´ is the new aggregator, and needs to be implemented for the /// generator's `SupportType`. pub(super) fn convert_aggregator<ANew, G>( mut self, generator : &G, domain : &Cube<F, N> ) -> Node<F,D,ANew,N,P> where ANew : Aggregator, G : SupportGenerator<F, N, Id=D>, G::SupportType : LocalAnalysis<F, ANew, N> { // The mem::replace is needed due to the [`Drop`] implementation to extract self.data. match std::mem::replace(&mut self.data, NodeOption::Uninitialised) { NodeOption::Uninitialised => Node { data : NodeOption::Uninitialised, aggregator : ANew::new(), }, NodeOption::Leaf(v) => { let mut anew = ANew::new(); anew.aggregate(v.iter().map(|d| { let support = generator.support_for(*d); support.local_analysis(&domain) })); Node { data : NodeOption::Leaf(v), aggregator : anew, } }, NodeOption::Branches(b) => { // TODO: now with Rc, convert_aggregator should be reference-based. let bnew = Rc::new(Rc::unwrap_or_clone(b).convert_aggregator(generator, domain)); let mut anew = ANew::new(); anew.summarise(bnew.aggregators()); Node { data : NodeOption::Branches(bnew), aggregator : anew, } } } } /// Refresh aggregator after changes to generator. /// /// The `generator` is used to convert the data of type `D` of the node into corresponding /// [`Support`]s. The `domain` is the cube corresponding to `self`. pub(super) fn refresh_aggregator<G>( &mut self, generator : &G, domain : &Cube<F, N> ) where G : SupportGenerator<F, N, Id=D>, G::SupportType : LocalAnalysis<F, A, N> { match &mut self.data { NodeOption::Uninitialised => { }, NodeOption::Leaf(v) => { self.aggregator = A::new(); self.aggregator.aggregate(v.iter().map(|d| { generator.support_for(*d) .local_analysis(&domain) })); }, NodeOption::Branches(ref mut b) => { // TODO: now with Rc, convert_aggregator should be reference-based. Rc::make_mut(b).refresh_aggregator(generator, domain); self.aggregator.summarise(b.aggregators()); } } } } /// Helper trait for working with [`Node`]s without the knowledge of `P`. /// /// This can be removed and the methods implemented directly on [`BT`] once Rust's const generics /// are flexible enough to allow fixing `P=pow(2, N)`. pub trait BTNode<F, D, A, const N : usize> where F : Float, D : 'static + Copy, A : Aggregator { type Node : Clone + std::fmt::Debug; } /// Helper structure for looking up a [`Node`] without the knowledge of `P`. /// /// This can be removed once Rust's const generics are flexible enough to allow fixing /// `P=pow(2, N)`. #[derive(Debug)] pub struct BTNodeLookup; /// Basic interface to a [`BT`] bisection tree. /// /// Further routines are provided by the [`BTSearch`][super::refine::BTSearch] trait. pub trait BTImpl<F : Float, const N : usize> : std::fmt::Debug + Clone + GlobalAnalysis<F, Self::Agg> { /// The data type stored in the tree type Data : 'static + Copy; /// The depth type of the tree type Depth : Depth; /// The type for the [aggregate information][Aggregator] about the `Data` stored in each node /// of the tree. type Agg : Aggregator; /// The type of the tree with the aggregator converted to `ANew`. type Converted<ANew> : BTImpl<F, N, Data=Self::Data, Agg=ANew> where ANew : Aggregator; /// Insert the data `d` into the tree for `support`. /// /// Every leaf node of the tree that intersects the `support` will contain a copy of /// `d`. fn insert<S : LocalAnalysis<F, Self::Agg, N>>( &mut self, d : Self::Data, support : &S ); /// Construct a new instance of the tree for a different aggregator /// /// The `generator` is used to convert the data of type [`Self::Data`] contained in the tree /// into corresponding [`Support`]s. fn convert_aggregator<ANew, G>(self, generator : &G) -> Self::Converted<ANew> where ANew : Aggregator, G : SupportGenerator<F, N, Id=Self::Data>, G::SupportType : LocalAnalysis<F, ANew, N>; /// Refreshes the aggregator of the three after possible changes to the support generator. /// /// The `generator` is used to convert the data of type [`Self::Data`] contained in the tree /// into corresponding [`Support`]s. fn refresh_aggregator<G>(&mut self, generator : &G) where G : SupportGenerator<F, N, Id=Self::Data>, G::SupportType : LocalAnalysis<F, Self::Agg, N>; /// Iterarate all [`Self::Data`] items at the point `x` of the domain. fn iter_at<'a>(&'a self, x : &'a Loc<F,N>) -> BTIter<'a, Self::Data>; /// Create a new tree on `domain` of indicated `depth`. fn new(domain : Cube<F, N>, depth : Self::Depth) -> Self; } /// The main bisection tree structure. /// /// It should be accessed via the [`BTImpl`] trait to hide the `const P : usize` parameter until /// const generics are flexible enough to fix `P=pow(2, N)` and thus also get rid of /// the `BTNodeLookup : BTNode<F, D, A, N>` trait bound. #[derive(Clone,Debug)] pub struct BT< M : Depth, F : Float, D : 'static + Copy, A : Aggregator, const N : usize, > where BTNodeLookup : BTNode<F, D, A, N> { /// The depth of the tree (initial, before refinement) pub(super) depth : M, /// The domain of the toplevel node pub(super) domain : Cube<F, N>, /// The toplevel node of the tree pub(super) topnode : <BTNodeLookup as BTNode<F, D, A, N>>::Node, } macro_rules! impl_bt { ($($n:literal)*) => { $( impl<F, D, A> BTNode<F, D, A, $n> for BTNodeLookup where F : Float, D : 'static + Copy + std::fmt::Debug, A : Aggregator { type Node = Node<F,D,A,$n,{pow(2, $n)}>; } impl<M,F,D,A> BTImpl<F,$n> for BT<M,F,D,A,$n> where M : Depth, F : Float, D : 'static + Copy + std::fmt::Debug, A : Aggregator { type Data = D; type Depth = M; type Agg = A; type Converted<ANew> = BT<M,F,D,ANew,$n> where ANew : Aggregator; fn insert<S: LocalAnalysis<F, A, $n>>( &mut self, d : D, support : &S ) { self.topnode.insert( &self.domain, d, self.depth, support ); } fn convert_aggregator<ANew, G>(self, generator : &G) -> Self::Converted<ANew> where ANew : Aggregator, G : SupportGenerator<F, $n, Id=D>, G::SupportType : LocalAnalysis<F, ANew, $n> { let topnode = self.topnode.convert_aggregator(generator, &self.domain); BT { depth : self.depth, domain : self.domain, topnode } } fn refresh_aggregator<G>(&mut self, generator : &G) where G : SupportGenerator<F, $n, Id=Self::Data>, G::SupportType : LocalAnalysis<F, Self::Agg, $n> { self.topnode.refresh_aggregator(generator, &self.domain); } fn iter_at<'a>(&'a self, x : &'a Loc<F,$n>) -> BTIter<'a,D> { match self.topnode.get_leaf_data(x) { Some(data) => BTIter { iter : data.iter() }, None => BTIter { iter : [].iter() } } } fn new(domain : Cube<F, $n>, depth : M) -> Self { BT { depth : depth, domain : domain, topnode : Node::new(), } } } impl<M,F,D,A> GlobalAnalysis<F,A> for BT<M,F,D,A,$n> where M : Depth, F : Float, D : 'static + Copy + std::fmt::Debug, A : Aggregator { fn global_analysis(&self) -> A { self.topnode.get_aggregator().clone() } } )* } } impl_bt!(1 2 3 4);