conspire/geometry/ntree/from/grid/
mod.rs

1#[cfg(test)]
2mod test;
3
4use crate::{
5    geometry::{
6        grid::Grid,
7        ntree::{
8            Orthotree,
9            balance::Balancing,
10            node::{Kind, Node, split::Split},
11            pair::Pairing,
12            rescale::Rescaling,
13        },
14    },
15    math::Scalar,
16};
17use std::{array::from_fn, ops::Add};
18
19type Pyramid<const D: usize, V> = Vec<([usize; D], Vec<Option<V>>)>;
20
21enum Cell<V> {
22    Empty,
23    Uniform(V),
24    Mixed,
25}
26
27impl<const D: usize, const L: usize, const M: usize, const N: usize, T, U, V> From<Grid<D, V>>
28    for Orthotree<D, L, M, N, T, U, V>
29where
30    T: Add<Output = T> + Copy + From<u16> + Into<usize> + Split,
31    U: Copy + From<usize> + Into<usize>,
32    V: Copy + PartialEq,
33{
34    fn from(grid: Grid<D, V>) -> Self {
35        let nel = *grid.nel();
36        let max = nel.iter().copied().max().unwrap_or(0).max(1);
37        let mut root_length = 1u16;
38        while (root_length as usize) < max {
39            root_length <<= 1;
40        }
41        let half = root_length as Scalar / 2.0;
42        let mut tree = Self {
43            balanced: Balancing::None,
44            nodes: vec![Node {
45                corner: from_fn(|_| T::from(0)),
46                length: T::from(root_length),
47                facets: [None; M],
48                kind: Kind::Leaf,
49                value: None,
50            }],
51            paired: Pairing::None,
52            rescale: Rescaling {
53                center: [half; D],
54                cell: 1.0,
55                half,
56            },
57        };
58        let pyramid = pyramid(
59            &nel,
60            root_length.trailing_zeros(),
61            grid.data_col_major().into_owned(),
62        );
63        let mut index = 0;
64        while index < tree.len() {
65            let node = &tree.nodes[index];
66            let corner = from_fn(|ax| node.corner[ax].into());
67            let length = node.length.into();
68            match classify(corner, length, &nel, &pyramid) {
69                Cell::Uniform(value) => tree.nodes[index].value = Some(value),
70                Cell::Mixed => {
71                    tree.subdivide(U::from(index)).ok();
72                }
73                Cell::Empty => {}
74            }
75            index += 1;
76        }
77        tree
78    }
79}
80
81fn pyramid<const D: usize, V: Copy + PartialEq>(
82    nel: &[usize; D],
83    levels: u32,
84    data: Vec<V>,
85) -> Pyramid<D, V> {
86    let mut out: Pyramid<D, V> = vec![(*nel, data.into_iter().map(Some).collect())];
87    for _ in 0..levels {
88        let (dim, prev) = out.last().unwrap();
89        let dim = *dim;
90        let next_dim: [usize; D] = from_fn(|ax| dim[ax].div_ceil(2));
91        let mut next = vec![None; next_dim.iter().product()];
92        for (cell, slot) in next.iter_mut().enumerate() {
93            let base = unflatten(cell, &next_dim);
94            let mut value = None;
95            let mut uniform = true;
96            'gather: for child in 0..(1usize << D) {
97                let coord = from_fn(|ax| 2 * base[ax] + ((child >> ax) & 1));
98                if (0..D).any(|ax| coord[ax] >= dim[ax]) {
99                    continue;
100                }
101                match prev[flatten(&coord, &dim)] {
102                    Some(entry) if value.is_none_or(|seen| seen == entry) => value = Some(entry),
103                    _ => {
104                        uniform = false;
105                        break 'gather;
106                    }
107                }
108            }
109            *slot = uniform.then_some(value).flatten();
110        }
111        out.push((next_dim, next));
112    }
113    out
114}
115
116fn flatten<const D: usize>(coord: &[usize; D], dim: &[usize; D]) -> usize {
117    let mut offset = 0;
118    let mut stride = 1;
119    for (c, n) in coord.iter().zip(dim) {
120        offset += c * stride;
121        stride *= n;
122    }
123    offset
124}
125
126fn unflatten<const D: usize>(mut index: usize, dim: &[usize; D]) -> [usize; D] {
127    from_fn(|ax| {
128        let coord = index % dim[ax];
129        index /= dim[ax];
130        coord
131    })
132}
133
134fn classify<const D: usize, V: Copy>(
135    corner: [usize; D],
136    length: usize,
137    nel: &[usize; D],
138    pyramid: &Pyramid<D, V>,
139) -> Cell<V> {
140    if (0..D).any(|ax| corner[ax] >= nel[ax]) {
141        return Cell::Empty;
142    }
143    if (0..D).any(|ax| corner[ax] + length > nel[ax]) {
144        return Cell::Mixed;
145    }
146    let (dim, data) = &pyramid[length.trailing_zeros() as usize];
147    let cell = from_fn(|ax| corner[ax] / length);
148    match data[flatten(&cell, dim)] {
149        Some(value) => Cell::Uniform(value),
150        None => Cell::Mixed,
151    }
152}
conspire/geometry/ntree/from/grid/mod.rs

conspire/geometry/ntree/from/grid/
mod.rs
