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

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

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