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

epoint_core/
point_data_columns.rs
