Skip to main content

spatialrust_episode/
augment.rs

1//! Deterministic episode augmentation operators.
2
3use spatialrust_sync::MemoryEpisode;
4
5use crate::{Episode, EpisodeResult};
6
7/// Augmentation operators for embodied datasets.
8#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
9pub enum AugmentationOp {
10    /// Drop every Nth record to thin the episode.
11    TemporalSubsample {
12        /// Keep stride (must be >= 1).
13        stride: usize,
14    },
15    /// Reverse record order (for stress tests).
16    Reverse,
17}
18
19/// Applies augmentation ops while preserving schema contracts.
20#[derive(Clone, Copy, Debug, Default)]
21pub struct EpisodeAugmentor;
22
23impl EpisodeAugmentor {
24    /// Applies one operator to an episode and returns a new episode id suffix.
25    pub fn apply(&self, episode: &Episode, op: AugmentationOp) -> EpisodeResult<Episode> {
26        let records = episode.memory.records();
27        let next = match op {
28            AugmentationOp::TemporalSubsample { stride } => {
29                if stride == 0 {
30                    return Err(crate::EpisodeError::InvalidConfiguration(
31                        "stride must be positive".into(),
32                    ));
33                }
34                records.iter().step_by(stride).cloned().collect::<Vec<_>>()
35            }
36            AugmentationOp::Reverse => records.iter().rev().cloned().collect::<Vec<_>>(),
37        };
38        let memory = MemoryEpisode::from_records(next);
39        let mut out = Episode::try_new(format!("{}::aug", episode.id.0), memory)?;
40        out.annotations = episode.annotations.clone();
41        out.provenance = episode.provenance.clone();
42        Ok(out)
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::{AugmentationOp, EpisodeAugmentor};
49    use crate::Episode;
50    use spatialrust_core::{
51        PointBuffer, PointBufferSet, PointCloud, SpatialMetadata, StandardSchemas, Timestamp,
52    };
53    use spatialrust_records::{SchemaVersion, SpatialRecord};
54    use spatialrust_sync::{ClockDomain, MemoryEpisode, StampedRecord, StampedTime};
55
56    fn sample(nanos: u64) -> StampedRecord {
57        let mut buffers = PointBufferSet::new();
58        buffers.insert("x", PointBuffer::from_f32(vec![0.0]));
59        buffers.insert("y", PointBuffer::from_f32(vec![0.0]));
60        buffers.insert("z", PointBuffer::from_f32(vec![0.0]));
61        let cloud = PointCloud::try_from_parts(
62            StandardSchemas::point_xyz(),
63            buffers,
64            SpatialMetadata::default(),
65        )
66        .unwrap();
67        let record =
68            SpatialRecord::try_from_cloud("point", SchemaVersion::new(1, 0), cloud).unwrap();
69        StampedRecord::new(
70            "lidar",
71            StampedTime::exact("host", ClockDomain::HostSteady, Timestamp::from_nanos(nanos)),
72            record,
73        )
74    }
75
76    #[test]
77    fn subsamples_episode() {
78        let memory = MemoryEpisode::from_records(vec![sample(1), sample(2), sample(3), sample(4)]);
79        let episode = Episode::try_new("ep0", memory).unwrap();
80        let out = EpisodeAugmentor
81            .apply(&episode, AugmentationOp::TemporalSubsample { stride: 2 })
82            .unwrap();
83        assert_eq!(out.memory.records().len(), 2);
84    }
85}