h3ron_polars/algorithm/chunkedarray/
grid_disk.rs

1use crate::algorithm::chunkedarray::util::list_map_cells;
2use crate::{Error, FromIndexIterator, IndexChunked};
3use h3ron::H3Cell;
4use polars_core::prelude::{ListChunked, UInt64Chunked};
5
6/// Produces all cells within k distance of the origin cell.
7pub trait H3GridDisk {
8    /// Produces all cells within k distance of the origin cell.
9    ///
10    /// k=0 is defined as the origin cell, k=1 is defined as k=0 + all
11    /// neighboring cells, and so on.
12    fn h3_grid_disk(&self, k: u32) -> Result<ListChunked, Error>;
13}
14
15impl<'a> H3GridDisk for IndexChunked<'a, H3Cell> {
16    fn h3_grid_disk(&self, k: u32) -> Result<ListChunked, Error> {
17        list_map_cells(self, |cell| {
18            cell.grid_disk(k)
19                .map(|cells| UInt64Chunked::from_index_iter(&cells))
20                .map_err(Error::from)
21        })
22    }
23}
24
25#[cfg(test)]
26mod tests {
27    use crate::algorithm::chunkedarray::{
28        H3ChangeResolution, H3CompactCells, H3GridDisk, H3Resolution,
29    };
30    use crate::from::{FromIndexIterator, NamedFromIndexes};
31    use crate::AsH3CellChunked;
32    use h3ron::{H3Cell, Index};
33    use polars::prelude::{ChunkExplode, TakeRandom, UInt64Chunked};
34
35    #[test]
36    fn cell_resolution() {
37        let expected_res = 6;
38        let cell = H3Cell::from_coordinate((4.5, 1.3).into(), expected_res).unwrap();
39        let ca = UInt64Chunked::from_index_iter([
40            Some(cell),
41            Some(H3Cell::new(55)), // invalid
42            None,
43        ]);
44        assert_eq!(ca.len(), 3);
45        let resolution_ca = ca.h3cell().h3_resolution();
46
47        assert_eq!(resolution_ca.get(0), Some(expected_res));
48        assert_eq!(resolution_ca.get(1), None);
49        assert_eq!(resolution_ca.get(2), None);
50    }
51
52    #[test]
53    fn cell_change_resolution_to_child() {
54        let cell = H3Cell::from_coordinate((4.5, 1.3).into(), 6).unwrap();
55        let ca = UInt64Chunked::new_from_indexes("", vec![cell]);
56        assert_eq!(ca.len(), 1);
57
58        let changed = ca.h3cell().h3_change_resolution(7).unwrap();
59        assert_eq!(changed.len(), 1);
60        let exploded = changed.explode().unwrap().unique().unwrap();
61        assert_eq!(exploded.len(), 7);
62    }
63
64    #[test]
65    fn cell_compact() {
66        let cell = H3Cell::from_coordinate((4.5, 1.3).into(), 6).unwrap();
67
68        let ca = UInt64Chunked::from_index_iter(&cell.get_children(7).unwrap());
69        assert_eq!(ca.len(), 7);
70
71        let changed = ca.h3cell().h3_compact_cells().unwrap();
72        assert_eq!(changed.len(), 1);
73        assert_eq!(changed.h3cell().get(0), Some(cell));
74    }
75
76    #[test]
77    fn cell_grid_disk() {
78        let cell = H3Cell::from_coordinate((4.5, 1.3).into(), 6).unwrap();
79
80        let ca = UInt64Chunked::new_from_indexes("", vec![cell]);
81        assert_eq!(ca.len(), 1);
82
83        let disk: Vec<_> = ca
84            .h3cell()
85            .h3_grid_disk(1)
86            .unwrap()
87            .explode()
88            .unwrap()
89            .sort(false)
90            .u64()
91            .unwrap()
92            .h3cell()
93            .iter_indexes_nonvalidated()
94            .flatten()
95            .collect();
96
97        let mut expected: Vec<_> = cell.grid_disk(1).unwrap().iter().collect();
98        expected.sort_unstable();
99
100        assert_eq!(disk, expected);
101    }
102}
h3ron_polars/algorithm/chunkedarray/grid_disk.rs

h3ron_polars/algorithm/chunkedarray/
grid_disk.rs
