1use crate::error::Error;
2use crate::{PointCloudInfo, PointDataColumnType, PointDataColumns};
3use chrono::{DateTime, Utc};
4use std::collections::{HashMap, HashSet};
5
6use ecoord::{FrameId, ReferenceFrames, TransformId};
7use nalgebra;
8use nalgebra::Point3;
9
10use polars::prelude::DataFrame;
11
12use crate::Error::{
13 MultipleFrameIdDefinitions, NoFrameIdDefinition, NoFrameIdDefinitions, NoIdColumn,
14};
15use crate::point_data::PointData;
16use polars::prelude::*;
17use rayon::prelude::*;
18
19#[derive(Debug, Clone, PartialEq)]
20pub struct PointCloud {
21 pub point_data: PointData,
22 pub info: PointCloudInfo,
23 pub reference_frames: ReferenceFrames,
24}
25
26impl PointCloud {
27 pub fn new(
28 point_data: PointDataColumns,
29 info: PointCloudInfo,
30 reference_frames: ReferenceFrames,
31 ) -> Result<Self, Error> {
32 let point_data = point_data.get_as_data_frame();
33 if point_data
34 .column(PointDataColumnType::FrameId.as_str())
35 .is_ok()
36 && info.frame_id.is_some()
37 {
38 return Err(MultipleFrameIdDefinitions);
39 }
40
41 Ok(Self {
42 point_data: PointData::new(point_data)?,
43 info,
44 reference_frames,
45 })
46 }
47
48 pub fn from_data_frame(
49 point_data: DataFrame,
50 info: PointCloudInfo,
51 reference_frames: ReferenceFrames,
52 ) -> Result<Self, Error> {
53 if point_data
54 .column(PointDataColumnType::FrameId.as_str())
55 .is_ok()
56 && info.frame_id.is_some()
57 {
58 return Err(MultipleFrameIdDefinitions);
59 }
60
61 Ok(Self {
62 point_data: PointData::new(point_data)?,
63 info,
64 reference_frames,
65 })
66 }
67}
68
69impl PointCloud {
70 pub fn point_data(&self) -> &PointData {
71 &self.point_data
72 }
73
74 pub fn info(&self) -> &PointCloudInfo {
75 &self.info
76 }
77
78 pub fn reference_frames(&self) -> &ReferenceFrames {
79 &self.reference_frames
80 }
81
82 pub fn size(&self) -> usize {
83 self.point_data.height()
84 }
85
86 pub fn info_frame_id(&self) -> Option<&FrameId> {
87 self.info.frame_id.as_ref()
88 }
89
90 pub fn get_distinct_frame_ids(&self) -> Option<HashSet<FrameId>> {
91 if let Some(frame_id) = &self.info.frame_id {
92 let frame_ids = HashSet::from([frame_id.clone()]);
93 return Some(frame_ids);
94 }
95
96 if self.point_data.contains_frame_id_column() {
97 let frame_ids = self
98 .point_data
99 .get_distinct_frame_ids()
100 .expect("must be derivable");
101 return Some(frame_ids);
102 }
103
104 None
105 }
106}
107
108impl PointCloud {
109 pub fn contains_ids(&self) -> bool {
110 self.point_data.contains_id_column()
111 }
112
113 pub fn contains_frame_ids(&self) -> bool {
114 self.info.frame_id.is_some() || !self.point_data.contains_frame_id_column()
115 }
116 pub fn contains_timestamps(&self) -> bool {
117 self.point_data.contains_timestamps()
118 }
119
120 pub fn contains_sensor_translation(&self) -> bool {
121 self.point_data.contains_sensor_translation()
122 }
123
124 pub fn contains_colors(&self) -> bool {
125 self.point_data.contains_colors()
126 }
127
128 pub fn set_reference_frames(&mut self, reference_frames: ReferenceFrames) {
129 self.reference_frames = reference_frames;
130 }
131
132 pub fn set_info_frame_id(&mut self, frame_id: Option<FrameId>) {
133 self.info.frame_id = frame_id;
134 }
135
136 pub fn update_points(
137 &mut self,
138 points: Vec<Point3<f64>>,
139 frame_id: Option<FrameId>,
140 ) -> Result<(), Error> {
141 self.point_data.update_points_in_place(points)?;
142 self.info.frame_id = frame_id;
143
144 Ok(())
145 }
146
147 pub fn derive_spherical_points(&mut self) -> Result<(), Error> {
148 self.point_data.derive_spherical_points()?;
149
150 Ok(())
151 }
152
153 pub fn filter_by_id_range(
154 &self,
155 id_min: Option<u64>,
156 id_max: Option<u64>,
157 ) -> Result<PointCloud, Error> {
158 if !self.contains_ids() {
159 return Err(NoIdColumn);
160 }
161
162 let mut filter_predicate = col(PointDataColumnType::Id.as_str());
163 if let Some(id_min) = id_min {
164 filter_predicate = filter_predicate.gt_eq(lit(id_min));
165 }
166 if let Some(id_max) = id_max {
167 filter_predicate =
168 filter_predicate.and(col(PointDataColumnType::Id.as_str()).lt_eq(id_max));
169 }
170
171 let point_data = self
172 .point_data
173 .data_frame
174 .clone()
175 .lazy()
176 .filter(filter_predicate)
177 .collect()?;
178
179 let filtered_point_cloud = PointCloud::from_data_frame(
180 point_data,
181 self.info.clone(),
182 self.reference_frames.clone(),
183 )?;
184 Ok(filtered_point_cloud)
185 }
186
187 pub fn filter_by_frame_id(&self, frame_id: &FrameId) -> Result<PointCloud, Error> {
188 if !self
189 .get_distinct_frame_ids()
190 .ok_or(NoFrameIdDefinitions)?
191 .contains(frame_id)
192 {
193 return Err(NoFrameIdDefinition(frame_id.clone()));
194 }
195
196 let filter_predicate = col(PointDataColumnType::FrameId.as_str())
197 .cast(DataType::String)
198 .eq(lit(frame_id.clone().to_string().as_str()));
199
200 let point_data = self
201 .point_data
202 .data_frame
203 .clone()
204 .lazy()
205 .filter(filter_predicate)
206 .collect()?;
207
208 let filtered_point_cloud = PointCloud::from_data_frame(
209 point_data,
210 self.info.clone(),
211 self.reference_frames.clone(),
212 )?;
213 Ok(filtered_point_cloud)
214 }
215
216 pub fn filter_by_row_indices(&self, row_indices: HashSet<usize>) -> Result<PointCloud, Error> {
217 let filtered_point_data = self.point_data.filter_by_row_indices(row_indices)?;
218
219 let filtered_point_cloud = PointCloud::from_data_frame(
220 filtered_point_data.data_frame,
221 self.info.clone(),
222 self.reference_frames.clone(),
223 )?;
224 Ok(filtered_point_cloud)
225 }
226
227 pub fn filter_by_boolean_mask(&self, mask: &Vec<bool>) -> Result<PointCloud, Error> {
228 let mask_series: Series = mask.iter().collect();
229 let filtered_point_data = self
230 .point_data
231 .filter_by_boolean_mask(mask_series.bool()?)?;
232
233 let filtered_point_cloud = PointCloud::from_data_frame(
234 filtered_point_data.data_frame,
235 self.info.clone(),
236 self.reference_frames.clone(),
237 )?;
238 Ok(filtered_point_cloud)
239 }
240
241 pub fn filter_by_bounds(
242 &self,
243 bound_min: Point3<f64>,
244 bound_max: Point3<f64>,
245 ) -> Result<Option<PointCloud>, Error> {
246 let filtered_point_data = self.point_data.filter_by_bounds(bound_min, bound_max)?;
247
248 let result = if let Some(filtered_point_data) = filtered_point_data {
249 let filtered_point_cloud = PointCloud::from_data_frame(
250 filtered_point_data.data_frame,
251 self.info.clone(),
252 self.reference_frames.clone(),
253 )?;
254 Some(filtered_point_cloud)
255 } else {
256 None
257 };
258
259 Ok(result)
260 }
261
262 pub fn filter_by_beam_length(
263 &self,
264 beam_length_min: f64,
265 beam_length_max: f64,
266 ) -> Result<Option<PointCloud>, Error> {
267 let filtered_point_data = self
268 .point_data
269 .filter_by_beam_length(beam_length_min, beam_length_max)?;
270
271 let result = if let Some(filtered_point_data) = filtered_point_data {
272 let filtered_point_cloud = PointCloud::from_data_frame(
273 filtered_point_data.data_frame,
274 self.info.clone(),
275 self.reference_frames.clone(),
276 )?;
277 Some(filtered_point_cloud)
278 } else {
279 None
280 };
281
282 Ok(result)
283 }
284
285 pub fn filter_by_x_min(&self, x_min: f64) -> Result<Option<PointCloud>, Error> {
286 let filtered_point_data = self.point_data.filter_by_x_min(x_min)?;
287
288 let result = if let Some(filtered_point_data) = filtered_point_data {
289 let filtered_point_cloud = PointCloud::from_data_frame(
290 filtered_point_data.data_frame,
291 self.info.clone(),
292 self.reference_frames.clone(),
293 )?;
294 Some(filtered_point_cloud)
295 } else {
296 None
297 };
298
299 Ok(result)
300 }
301
302 pub fn filter_by_x_max(&self, x_max: f64) -> Result<Option<PointCloud>, Error> {
303 let filtered_point_data = self.point_data.filter_by_x_max(x_max)?;
304
305 let result = if let Some(filtered_point_data) = filtered_point_data {
306 let filtered_point_cloud = PointCloud::from_data_frame(
307 filtered_point_data.data_frame,
308 self.info.clone(),
309 self.reference_frames.clone(),
310 )?;
311 Some(filtered_point_cloud)
312 } else {
313 None
314 };
315
316 Ok(result)
317 }
318
319 pub fn filter_by_y_min(&self, y_min: f64) -> Result<Option<PointCloud>, Error> {
320 let filtered_point_data = self.point_data.filter_by_y_min(y_min)?;
321
322 let result = if let Some(filtered_point_data) = filtered_point_data {
323 let filtered_point_cloud = PointCloud::from_data_frame(
324 filtered_point_data.data_frame,
325 self.info.clone(),
326 self.reference_frames.clone(),
327 )?;
328 Some(filtered_point_cloud)
329 } else {
330 None
331 };
332
333 Ok(result)
334 }
335
336 pub fn filter_by_y_max(&self, y_max: f64) -> Result<Option<PointCloud>, Error> {
337 let filtered_point_data = self.point_data.filter_by_y_max(y_max)?;
338
339 let result = if let Some(filtered_point_data) = filtered_point_data {
340 let filtered_point_cloud = PointCloud::from_data_frame(
341 filtered_point_data.data_frame,
342 self.info.clone(),
343 self.reference_frames.clone(),
344 )?;
345 Some(filtered_point_cloud)
346 } else {
347 None
348 };
349
350 Ok(result)
351 }
352
353 pub fn filter_by_z_min(&self, z_min: f64) -> Result<Option<PointCloud>, Error> {
354 let filtered_point_data = self.point_data.filter_by_z_min(z_min)?;
355
356 let result = if let Some(filtered_point_data) = filtered_point_data {
357 let filtered_point_cloud = PointCloud::from_data_frame(
358 filtered_point_data.data_frame,
359 self.info.clone(),
360 self.reference_frames.clone(),
361 )?;
362 Some(filtered_point_cloud)
363 } else {
364 None
365 };
366
367 Ok(result)
368 }
369
370 pub fn filter_by_z_max(&self, z_max: f64) -> Result<Option<PointCloud>, Error> {
371 let filtered_point_data = self.point_data.filter_by_z_max(z_max)?;
372
373 let result = if let Some(filtered_point_data) = filtered_point_data {
374 let filtered_point_cloud = PointCloud::from_data_frame(
375 filtered_point_data.data_frame,
376 self.info.clone(),
377 self.reference_frames.clone(),
378 )?;
379 Some(filtered_point_cloud)
380 } else {
381 None
382 };
383
384 Ok(result)
385 }
386
387 pub fn filter_by_spherical_range_min(
388 &self,
389 spherical_range_min: f64,
390 ) -> Result<Option<PointCloud>, Error> {
391 let filtered_point_data = self
392 .point_data
393 .filter_by_spherical_range_min(spherical_range_min)?;
394
395 let result = if let Some(filtered_point_data) = filtered_point_data {
396 let filtered_point_cloud = PointCloud::from_data_frame(
397 filtered_point_data.data_frame,
398 self.info.clone(),
399 self.reference_frames.clone(),
400 )?;
401 Some(filtered_point_cloud)
402 } else {
403 None
404 };
405
406 Ok(result)
407 }
408
409 pub fn filter_by_spherical_range_max(
410 &self,
411 spherical_range_max: f64,
412 ) -> Result<Option<PointCloud>, Error> {
413 let filtered_point_data = self
414 .point_data
415 .filter_by_spherical_range_max(spherical_range_max)?;
416
417 let result = if let Some(filtered_point_data) = filtered_point_data {
418 let filtered_point_cloud = PointCloud::from_data_frame(
419 filtered_point_data.data_frame,
420 self.info.clone(),
421 self.reference_frames.clone(),
422 )?;
423 Some(filtered_point_cloud)
424 } else {
425 None
426 };
427
428 Ok(result)
429 }
430}
431
432impl PointCloud {
433 pub fn append_reference_frames(
434 &mut self,
435 reference_frames: ReferenceFrames,
436 ) -> Result<(), ecoord::Error> {
437 let merged_reference_frames =
438 ecoord::merge(&[self.reference_frames.clone(), reference_frames])?;
439 self.set_reference_frames(merged_reference_frames);
440 Ok(())
441 }
442
443 pub fn resolve_to_frame(&mut self, target_frame_id: FrameId) -> Result<(), Error> {
450 if self.info.frame_id.is_none() && !self.point_data.contains_frame_id_column() {
451 return Err(NoFrameIdDefinitions);
452 }
453
454 let mut partition_columns: Vec<&str> = Vec::new();
455 if self.point_data.contains_frame_id_column() {
456 partition_columns.push(PointDataColumnType::FrameId.as_str());
457 }
458 if self.point_data.contains_timestamp_sec_column() {
459 partition_columns.push(PointDataColumnType::TimestampSecond.as_str());
460 }
461 if self.point_data.contains_timestamp_nanosec_column() {
462 partition_columns.push(PointDataColumnType::TimestampNanoSecond.as_str());
463 }
464
465 let partitioned_data_frames: Vec<DataFrame> = if partition_columns.is_empty() {
466 vec![self.point_data.data_frame.clone()]
467 } else {
468 self.point_data
469 .data_frame
470 .clone()
471 .partition_by(partition_columns, true)?
472 };
473
474 let partitioned_data_frames: HashMap<(FrameId, Option<DateTime<Utc>>), DataFrame> =
475 partitioned_data_frames
476 .into_iter()
477 .map(|df| {
478 let point_data = PointData::new_unchecked(df);
479 let frame_ids_series = point_data.get_all_frame_ids();
481 let frame_id = frame_ids_series.map_or_else(
482 |_| self.info.frame_id.clone().unwrap(),
483 |f| f.first().unwrap().clone(),
484 );
485
486 let timestamp = point_data
487 .get_all_timestamps()
488 .ok()
489 .map(|t| *t.first().unwrap());
490 ((frame_id, timestamp), point_data.data_frame)
491 })
492 .collect();
493
494 let partitioned_data_frames: Vec<DataFrame> = partitioned_data_frames
495 .iter()
496 .map(|((current_frame_id, current_timestamp), df)| {
497 let mut point_data = PointData::new_unchecked(df.clone());
498 let r = point_data.resolve_data_frame(
499 &self.reference_frames,
500 current_timestamp,
501 current_frame_id,
502 &target_frame_id,
503 );
504 if r.is_err() {
505 print!(
506 "should work: current_frame_id = {current_frame_id:?}, target_frame_id: {target_frame_id:?}"
507 );
508 println!(
509 "all available frame ids: {:?}",
510 &self.reference_frames.get_frame_ids()
511 );
512 panic!("error");
513 }
514
515 point_data.data_frame
516 })
517 .collect();
518
519 let partitioned_lazy: Vec<LazyFrame> = partitioned_data_frames
520 .iter()
521 .map(|d| d.clone().lazy())
522 .collect();
523
524 let mut merged_again = concat(partitioned_lazy, Default::default())
525 .expect("concatenation should work")
526 .collect()
527 .expect("concatenation should work");
528
529 if merged_again
531 .column(PointDataColumnType::TimestampSecond.as_str())
532 .is_ok()
533 && merged_again
534 .column(PointDataColumnType::TimestampNanoSecond.as_str())
535 .is_ok()
536 && merged_again
537 .column(PointDataColumnType::FrameId.as_str())
538 .is_err()
539 {
540 merged_again = merged_again
541 .sort(
542 [
543 PointDataColumnType::TimestampSecond.as_str(),
544 PointDataColumnType::TimestampNanoSecond.as_str(),
545 ],
546 SortMultipleOptions::default().with_maintain_order(true),
547 )
548 .expect("");
549 }
550 if merged_again
552 .column(PointDataColumnType::Id.as_str())
553 .is_ok()
554 {
555 merged_again = merged_again
556 .sort(
557 [PointDataColumnType::Id.as_str()],
558 SortMultipleOptions::default().with_maintain_order(true),
559 )
560 .expect("");
561 }
562
563 self.point_data.data_frame = merged_again;
564 self.info.frame_id = Some(target_frame_id);
565
566 Ok(())
567 }
568
569 pub fn add_sensor_poses_from_frame(&mut self, frame_id: FrameId) -> Result<(), Error> {
575 let timestamps: Vec<DateTime<Utc>> = self.point_data.get_all_timestamps()?;
576 let transform_id = TransformId::new(self.info.frame_id.clone().unwrap(), frame_id);
577
578 let unique_timestamps: HashSet<_> = timestamps.iter().collect();
579
580 let isometry_map: HashMap<_, _> = unique_timestamps
581 .into_par_iter()
582 .map(|current_timestamp| {
583 let isometry_graph = self
584 .reference_frames
585 .derive_transform_graph(&None, &Some(*current_timestamp))?;
586 let isometry = isometry_graph.get_isometry(&transform_id)?;
587 Ok((current_timestamp, isometry))
588 })
589 .collect::<Result<HashMap<_, _>, Error>>()?;
590
591 let isometries: Vec<_> = timestamps.iter().map(|ts| isometry_map[ts]).collect();
592
593 self.point_data.add_sensor_poses(isometries)?;
594
595 Ok(())
596 }
597}