epoint_core/
point_data_columns.rs

1use crate::Error::{NoData, ShapeMisMatch};
2use crate::{Error, PointDataColumnType};
3use chrono::{DateTime, Timelike, Utc};
4use nalgebra::Point3;
5use palette::Srgb;
6use polars::datatypes::DataType;
7use polars::frame::DataFrame;
8use polars::prelude::{Column, NamedFrom};
9
10pub struct PointDataColumns {
11    pub point: Vec<Point3<f64>>,
12    pub id: Option<Vec<u64>>,
13    pub frame_id: Option<Vec<String>>,
14    pub timestamp: Option<Vec<DateTime<Utc>>>,
15    pub intensity: Option<Vec<f32>>,
16    pub sensor_translation: Option<Vec<Point3<f64>>>,
17    pub color: Option<Vec<Srgb<u16>>>,
18}
19
20impl PointDataColumns {
21    pub fn new(
22        point: Vec<Point3<f64>>,
23        id: Option<Vec<u64>>,
24        frame_id: Option<Vec<String>>,
25        timestamp: Option<Vec<DateTime<Utc>>>,
26        intensity: Option<Vec<f32>>,
27        sensor_translation: Option<Vec<Point3<f64>>>,
28        color: Option<Vec<Srgb<u16>>>,
29    ) -> Result<Self, Error> {
30        if point.is_empty() {
31            return Err(NoData("point"));
32        }
33        let total_length = point.len();
34
35        if let Some(id) = &id
36            && id.len() != total_length
37        {
38            return Err(ShapeMisMatch(
39                "id vector has a different length than the point vector",
40            ));
41        }
42        if let Some(frame_id_entries) = &frame_id
43            && frame_id_entries.len() != total_length
44        {
45            return Err(ShapeMisMatch(
46                "frame_id vector has a different length than the point vector",
47            ));
48        }
49        if let Some(timestamp_entries) = &timestamp
50            && timestamp_entries.len() != total_length
51        {
52            return Err(ShapeMisMatch(
53                "frame_id vector has a different length than the point vector",
54            ));
55        }
56        if let Some(intensity_entries) = &intensity
57            && intensity_entries.len() != total_length
58        {
59            return Err(ShapeMisMatch(
60                "intensity vector has a different length than the point vector",
61            ));
62        }
63        if let Some(sensor_translation_entries) = &sensor_translation
64            && sensor_translation_entries.len() != total_length
65        {
66            return Err(ShapeMisMatch(
67                "sensor_translation vector has a different length than the point vector",
68            ));
69        }
70        if let Some(color_entries) = &color
71            && color_entries.len() != total_length
72        {
73            return Err(ShapeMisMatch(
74                "color vector has a different length than the point vector",
75            ));
76        }
77
78        Ok(Self {
79            point,
80            id,
81            frame_id,
82            timestamp,
83            intensity,
84            sensor_translation,
85            color,
86        })
87    }
88
89    pub fn is_empty(&self) -> bool {
90        self.point.is_empty()
91    }
92
93    pub fn len(&self) -> usize {
94        self.point.len()
95    }
96
97    pub fn get_as_data_frame(&self) -> DataFrame {
98        let mut columns = self.get_xyz_columns();
99
100        if let Some(id) = &self.id {
101            let id_column = Column::new(PointDataColumnType::Id.into(), id);
102            columns.push(id_column);
103        }
104
105        if let Some(frame_id) = &self.frame_id {
106            let frame_id_column = Column::new(PointDataColumnType::FrameId.into(), frame_id);
107            let frame_id_column = frame_id_column
108                .cast(&DataType::Categorical(None, Default::default()))
109                .unwrap();
110            columns.push(frame_id_column);
111        }
112
113        if let Some(timestamp_entries) = &self.timestamp {
114            let timestamp_seconds_column = Column::new(
115                PointDataColumnType::TimestampSeconds.into(),
116                timestamp_entries
117                    .iter()
118                    .map(|t| t.timestamp())
119                    .collect::<Vec<i64>>(),
120            );
121            columns.push(timestamp_seconds_column);
122
123            let timestamp_nanoseconds_column = Column::new(
124                PointDataColumnType::TimestampNanoSeconds.into(),
125                timestamp_entries
126                    .iter()
127                    .map(|t| t.nanosecond())
128                    .collect::<Vec<u32>>(),
129            );
130            columns.push(timestamp_nanoseconds_column);
131        }
132
133        if let Some(intensity) = &self.intensity {
134            let intensity_column = Column::new(PointDataColumnType::Intensity.into(), intensity);
135            columns.push(intensity_column);
136        }
137
138        if self.sensor_translation.is_some() {
139            columns.append(&mut self.get_sensor_translation_xyz_columns().unwrap());
140        }
141
142        if self.color.is_some() {
143            columns.append(&mut self.get_color_columns().unwrap());
144        }
145
146        DataFrame::new(columns).unwrap()
147    }
148
149    pub fn get_xyz_columns(&self) -> Vec<Column> {
150        vec![
151            Column::new(
152                PointDataColumnType::X.into(),
153                self.point.iter().map(|p| p.x).collect::<Vec<f64>>(),
154            ),
155            Column::new(
156                PointDataColumnType::Y.into(),
157                self.point.iter().map(|p| p.y).collect::<Vec<f64>>(),
158            ),
159            Column::new(
160                PointDataColumnType::Z.into(),
161                self.point.iter().map(|p| p.z).collect::<Vec<f64>>(),
162            ),
163        ]
164    }
165
166    pub fn get_sensor_translation_xyz_columns(&self) -> Option<Vec<Column>> {
167        let sensor_translation = self.sensor_translation.as_ref()?;
168
169        Some(vec![
170            Column::new(
171                PointDataColumnType::SensorTranslationX.into(),
172                sensor_translation.iter().map(|p| p.x).collect::<Vec<f64>>(),
173            ),
174            Column::new(
175                PointDataColumnType::SensorTranslationY.into(),
176                sensor_translation.iter().map(|p| p.y).collect::<Vec<f64>>(),
177            ),
178            Column::new(
179                PointDataColumnType::SensorTranslationZ.into(),
180                sensor_translation.iter().map(|p| p.z).collect::<Vec<f64>>(),
181            ),
182        ])
183    }
184
185    pub fn get_color_columns(&self) -> Option<Vec<Column>> {
186        let color = self.color.as_ref()?;
187
188        Some(vec![
189            Column::new(
190                PointDataColumnType::ColorRed.into(),
191                color.iter().map(|c| c.red).collect::<Vec<u16>>(),
192            ),
193            Column::new(
194                PointDataColumnType::ColorGreen.into(),
195                color.iter().map(|c| c.green).collect::<Vec<u16>>(),
196            ),
197            Column::new(
198                PointDataColumnType::ColorBlue.into(),
199                color.iter().map(|c| c.blue).collect::<Vec<u16>>(),
200            ),
201        ])
202    }
203}
epoint_core/point_data_columns.rs

epoint_core/
point_data_columns.rs
