Skip to main content

rustpix_core/
soa.rs

1//! Structure of Arrays (`SoA`) types for efficient processing.
2//!
3//! This module defines the `HitBatch` structure which stores hit data
4//! in parallel vectors (`SoA` layout) rather than an array of structs (`AoS`).
5//! This layout works better with modern CPU caches and SIMD instructions.
6
7#[cfg(feature = "serde")]
8use serde::{Deserialize, Serialize};
9
10/// A batch of hits stored in Structure of Arrays (`SoA`) format.
11#[derive(Debug, Clone, Default, PartialEq)]
12#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
13pub struct HitBatch {
14    /// Columnar storage for X coordinates.
15    pub x: Vec<u16>,
16    /// Columnar storage for Y coordinates.
17    pub y: Vec<u16>,
18    /// Columnar storage for Time-of-Flight (corrected).
19    pub tof: Vec<u32>,
20    /// Columnar storage for Time-over-Threshold.
21    pub tot: Vec<u16>,
22    /// Columnar storage for global timestamps.
23    pub timestamp: Vec<u32>,
24    /// Columnar storage for Chip IDs (optional if batch is per-chip).
25    pub chip_id: Vec<u8>,
26    /// Cluster assignments (output of clustering).
27    pub cluster_id: Vec<i32>,
28}
29
30/// Tuple-based hit representation used to push into a batch without `AoS` storage.
31pub type HitRecord = (u16, u16, u32, u16, u32, u8);
32
33impl HitBatch {
34    /// Creates a new empty batch with specified capacity.
35    #[must_use]
36    pub fn with_capacity(capacity: usize) -> Self {
37        Self {
38            x: Vec::with_capacity(capacity),
39            y: Vec::with_capacity(capacity),
40            tof: Vec::with_capacity(capacity),
41            tot: Vec::with_capacity(capacity),
42            timestamp: Vec::with_capacity(capacity),
43            chip_id: Vec::with_capacity(capacity),
44            cluster_id: Vec::with_capacity(capacity),
45        }
46    }
47
48    /// Returns the number of hits in the batch.
49    #[must_use]
50    pub fn len(&self) -> usize {
51        self.x.len()
52    }
53
54    /// Returns true if the batch is empty.
55    #[must_use]
56    pub fn is_empty(&self) -> bool {
57        self.x.is_empty()
58    }
59
60    /// Clears all vectors in the batch.
61    pub fn clear(&mut self) {
62        self.x.clear();
63        self.y.clear();
64        self.tof.clear();
65        self.tot.clear();
66        self.timestamp.clear();
67        self.chip_id.clear();
68        self.cluster_id.clear();
69    }
70
71    /// Appends all hits from another batch to this one.
72    pub fn append(&mut self, other: &HitBatch) {
73        self.x.extend_from_slice(&other.x);
74        self.y.extend_from_slice(&other.y);
75        self.tof.extend_from_slice(&other.tof);
76        self.tot.extend_from_slice(&other.tot);
77        self.timestamp.extend_from_slice(&other.timestamp);
78        self.chip_id.extend_from_slice(&other.chip_id);
79        self.cluster_id.extend_from_slice(&other.cluster_id);
80    }
81
82    /// Pushes a single hit into the batch.
83    pub fn push(&mut self, hit: HitRecord) {
84        let (x, y, tof, tot, timestamp, chip_id) = hit;
85        self.x.push(x);
86        self.y.push(y);
87        self.tof.push(tof);
88        self.tot.push(tot);
89        self.timestamp.push(timestamp);
90        self.chip_id.push(chip_id);
91        self.cluster_id.push(-1); // Default unclustered
92    }
93
94    /// Sorts all hits by TOF, keeping columns aligned.
95    pub fn sort_by_tof(&mut self) {
96        let len = self.len();
97        if len <= 1 {
98            return;
99        }
100
101        let mut indices: Vec<usize> = (0..len).collect();
102        indices.sort_unstable_by_key(|&i| self.tof[i]);
103
104        let mut x = Vec::with_capacity(len);
105        let mut y = Vec::with_capacity(len);
106        let mut tof = Vec::with_capacity(len);
107        let mut tot = Vec::with_capacity(len);
108        let mut timestamp = Vec::with_capacity(len);
109        let mut chip_id = Vec::with_capacity(len);
110        let mut cluster_id = Vec::with_capacity(len);
111
112        for i in indices {
113            x.push(self.x[i]);
114            y.push(self.y[i]);
115            tof.push(self.tof[i]);
116            tot.push(self.tot[i]);
117            timestamp.push(self.timestamp[i]);
118            chip_id.push(self.chip_id[i]);
119            cluster_id.push(self.cluster_id[i]);
120        }
121
122        self.x = x;
123        self.y = y;
124        self.tof = tof;
125        self.tot = tot;
126        self.timestamp = timestamp;
127        self.chip_id = chip_id;
128        self.cluster_id = cluster_id;
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn test_hit_batch_operations() {
138        let mut batch = HitBatch::with_capacity(10);
139        assert!(batch.is_empty());
140
141        batch.push((10, 20, 1000, 5, 123_456, 0));
142        assert_eq!(batch.len(), 1);
143        assert_eq!(batch.x[0], 10);
144        assert_eq!(batch.cluster_id[0], -1);
145
146        batch.push((11, 21, 1001, 6, 123_457, 0));
147        assert_eq!(batch.len(), 2);
148
149        batch.clear();
150        assert!(batch.is_empty());
151        assert_eq!(batch.len(), 0);
152    }
153}