evoxel_transform/
filter.rs

1use crate::Error;
2use evoxel_core::{VoxelDataColumnType, VoxelGrid};
3use nalgebra::Point3;
4use polars::datatypes::PlSmallStr;
5use polars::frame::DataFrame;
6
7use crate::Error::LowerCornerMustBeBelowUpperCorner;
8use polars::prelude::{ChunkCompareIneq, IntoLazy, Selector, all, col, len};
9
10pub fn aggregate_by_index(voxel_grid: &VoxelGrid) -> Result<VoxelGrid, Error> {
11    let voxel_data = voxel_grid.voxel_data();
12    let partition_columns = vec![
13        VoxelDataColumnType::X.as_str(),
14        VoxelDataColumnType::Y.as_str(),
15        VoxelDataColumnType::Z.as_str(),
16    ];
17
18    let partitioned: DataFrame = voxel_data
19        .clone()
20        .lazy()
21        .group_by(partition_columns)
22        .agg([all().as_expr(), len()])
23        //.limit(5)
24        .collect()?;
25    //dbg!("{:?}", &partitioned);
26
27    /*println!("partitions: {}", partitioned.len());
28    let parit_filtered: Vec<DataFrame> = partitioned
29        .into_iter()
30        .filter(|d| d.height() >= minimum)
31        .collect();
32
33    let partitioned_lazy: Vec<LazyFrame> =
34        parit_filtered.iter().map(|d| d.clone().lazy()).collect();
35
36    let merged_again = concat(partitioned_lazy, true, true)
37        .unwrap()
38        .collect()
39        .unwrap();*/
40
41    let filtered_voxel_grid = VoxelGrid::new(
42        partitioned,
43        voxel_grid.info().clone(),
44        voxel_grid.reference_frames().clone(),
45    )?;
46    Ok(filtered_voxel_grid)
47}
48
49pub fn explode(voxel_grid: &VoxelGrid) -> Result<VoxelGrid, Error> {
50    let voxel_data = voxel_grid.voxel_data();
51
52    let column_names: Vec<&str> = voxel_data
53        .get_columns()
54        .iter()
55        .filter(|s| s.dtype().inner_dtype().is_some()) // if contains inner, it's a list
56        .map(|s| s.name().as_str())
57        .collect();
58
59    let df: DataFrame = voxel_data
60        .clone()
61        .lazy()
62        .explode(Selector::ByName {
63            names: column_names.into_iter().map(PlSmallStr::from_str).collect(),
64            strict: false,
65        })
66        //.limit(5)
67        .collect()?;
68
69    let filtered_voxel_grid = VoxelGrid::new(
70        df,
71        voxel_grid.info().clone(),
72        voxel_grid.reference_frames().clone(),
73    )?;
74    Ok(filtered_voxel_grid)
75}
76
77pub fn filter_by_count(voxel_grid: &VoxelGrid, minimum: usize) -> Result<VoxelGrid, Error> {
78    let voxel_data = voxel_grid.voxel_data().clone();
79
80    let mask = voxel_data
81        .column(VoxelDataColumnType::Count.as_str())?
82        .as_materialized_series()
83        .gt_eq(minimum as i32)?;
84
85    let filtered_voxel_data = voxel_data.filter(&mask)?;
86    let filtered_voxel_grid = VoxelGrid::new(
87        filtered_voxel_data,
88        voxel_grid.info().clone(),
89        voxel_grid.reference_frames().clone(),
90    )?;
91
92    Ok(filtered_voxel_grid)
93}
94
95pub fn filter_by_index_bounds(
96    voxel_grid: &VoxelGrid,
97    lower_corner: Point3<i64>,
98    upper_corner: Point3<i64>,
99) -> Result<VoxelGrid, Error> {
100    if lower_corner >= upper_corner {
101        return Err(LowerCornerMustBeBelowUpperCorner(""));
102    }
103
104    let voxel_data = voxel_grid.voxel_data().clone();
105
106    let filtered_voxel_data = voxel_data
107        .lazy()
108        .filter(
109            col(VoxelDataColumnType::X.as_str())
110                .gt_eq(lower_corner.x)
111                .and(col(VoxelDataColumnType::X.as_str()).lt_eq(upper_corner.x)),
112        )
113        .filter(
114            col(VoxelDataColumnType::Y.as_str())
115                .gt_eq(lower_corner.y)
116                .and(col(VoxelDataColumnType::Y.as_str()).lt_eq(upper_corner.y)),
117        )
118        .filter(
119            col(VoxelDataColumnType::Z.as_str())
120                .gt_eq(lower_corner.z)
121                .and(col(VoxelDataColumnType::Z.as_str()).lt_eq(upper_corner.z)),
122        )
123        .collect()?;
124    let filtered_voxel_grid = VoxelGrid::new(
125        filtered_voxel_data,
126        voxel_grid.info().clone(),
127        voxel_grid.reference_frames().clone(),
128    )?;
129
130    Ok(filtered_voxel_grid)
131}
evoxel_transform/filter.rs

evoxel_transform/
filter.rs
