Skip to main content

eventcv_core/representation/
mcts.rs

1use super::{
2    age_ms, event_index, frame_len, reference_time, validate_positive, EventFrame, EventFrameData,
3    Representation, RepresentationError, RepresentationKind,
4};
5use crate::EventStream;
6
7const WINDOW_COUNT: usize = 5;
8
9#[derive(Clone, Copy, Debug)]
10pub struct Mcts {
11    max_window_ms: f64,
12}
13
14impl Mcts {
15    pub fn new(max_window_ms: f64) -> Self {
16        Self { max_window_ms }
17    }
18
19    fn windows(&self) -> [f64; WINDOW_COUNT] {
20        let ratio = self.max_window_ms.powf(1.0 / (WINDOW_COUNT - 1) as f64);
21        std::array::from_fn(|index| ratio.powi(index as i32))
22    }
23}
24
25impl Default for Mcts {
26    fn default() -> Self {
27        Self::new(30.0)
28    }
29}
30
31impl Representation for Mcts {
32    type Output = EventFrame;
33
34    fn generate(&self, stream: &EventStream) -> Result<EventFrame, RepresentationError> {
35        validate_positive(self.max_window_ms, "max_window_ms")?;
36        if self.max_window_ms < 1.0 {
37            return Err(RepresentationError::InvalidParameter("max_window_ms"));
38        }
39        let (width, height, length) = frame_len(stream, WINDOW_COUNT * 2)?;
40        let plane_len = width * height;
41        let windows = self.windows();
42        let mut values = vec![0_f32; length];
43
44        if let Some(reference) = reference_time(stream) {
45            for event in stream.iter() {
46                let index = event_index(event, width, height)?;
47                let age = age_ms(stream, reference, event.timestamp);
48                let channel_offset = if event.polarity { WINDOW_COUNT } else { 0 };
49                for (window_index, window) in windows.iter().copied().enumerate() {
50                    if age <= window {
51                        let value = (1.0 - age / window) as f32;
52                        let output_index = (channel_offset + window_index) * plane_len + index;
53                        values[output_index] = values[output_index].max(value);
54                    }
55                }
56            }
57        }
58
59        let channel_names = ["negative", "positive"]
60            .into_iter()
61            .flat_map(|polarity| {
62                windows
63                    .iter()
64                    .map(move |window| format!("{polarity}_{window:.3}ms"))
65            })
66            .collect();
67
68        Ok(EventFrame {
69            data: EventFrameData::F32(values),
70            channels: WINDOW_COUNT * 2,
71            width,
72            height,
73            kind: RepresentationKind::Mcts,
74            channel_names,
75        })
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use ndarray::array;
82
83    use super::{Mcts, Representation};
84    use crate::{representation::EventFrameData, EventStream};
85
86    #[test]
87    fn builds_logarithmic_windows_for_each_polarity() {
88        let stream =
89            EventStream::from_array2(array![[0, 0, 14_000, 0], [0, 0, 16_000, 1]], 1, 1, 0.001);
90
91        let frame = Mcts::new(16.0).generate(&stream).unwrap();
92
93        assert_eq!(
94            frame.data(),
95            &EventFrameData::F32(vec![0.0, 0.0, 0.5, 0.75, 0.875, 1.0, 1.0, 1.0, 1.0, 1.0])
96        );
97    }
98}