215 lines
7.0 KiB
Rust
215 lines
7.0 KiB
Rust
//! Multi-node data aggregation.
|
|
//!
|
|
//! Collects [`NeuralDataPacket`]s from multiple ESP32 nodes and assembles them
|
|
//! into a unified [`MultiChannelTimeSeries`] once all nodes have reported for
|
|
//! a given time window.
|
|
|
|
use ruv_neural_core::signal::MultiChannelTimeSeries;
|
|
use ruv_neural_core::{Result, RuvNeuralError};
|
|
|
|
use crate::protocol::NeuralDataPacket;
|
|
|
|
/// Aggregates data packets from multiple ESP32 sensor nodes.
|
|
///
|
|
/// Packets are buffered per-node. When every node has contributed at least one
|
|
/// packet, [`try_assemble`](NodeAggregator::try_assemble) combines them into a
|
|
/// single time series — matching packets by timestamp within the configured
|
|
/// sync tolerance.
|
|
pub struct NodeAggregator {
|
|
node_count: usize,
|
|
buffers: Vec<Vec<NeuralDataPacket>>,
|
|
sync_tolerance_us: u64,
|
|
}
|
|
|
|
impl NodeAggregator {
|
|
/// Create a new aggregator expecting `node_count` distinct nodes.
|
|
pub fn new(node_count: usize) -> Self {
|
|
Self {
|
|
node_count,
|
|
buffers: vec![Vec::new(); node_count],
|
|
sync_tolerance_us: 1_000, // 1 ms default
|
|
}
|
|
}
|
|
|
|
/// Buffer a packet from a specific node.
|
|
pub fn receive_packet(
|
|
&mut self,
|
|
node_id: usize,
|
|
packet: NeuralDataPacket,
|
|
) -> Result<()> {
|
|
if node_id >= self.node_count {
|
|
return Err(RuvNeuralError::Sensor(format!(
|
|
"Node ID {node_id} out of range (max {})",
|
|
self.node_count - 1
|
|
)));
|
|
}
|
|
self.buffers[node_id].push(packet);
|
|
Ok(())
|
|
}
|
|
|
|
/// Try to assemble a [`MultiChannelTimeSeries`] from the buffered packets.
|
|
///
|
|
/// Returns `Some` when every node has at least one packet whose timestamps
|
|
/// are within `sync_tolerance_us` of each other. The matching packets are
|
|
/// consumed from the buffers.
|
|
pub fn try_assemble(&mut self) -> Option<MultiChannelTimeSeries> {
|
|
// Check that every node has at least one packet
|
|
if self.buffers.iter().any(|b| b.is_empty()) {
|
|
return None;
|
|
}
|
|
|
|
// Use the first node's earliest packet as the reference timestamp
|
|
let ref_ts = self.buffers[0][0].header.timestamp_us;
|
|
|
|
// Find a matching packet in each buffer
|
|
let mut indices: Vec<usize> = Vec::with_capacity(self.node_count);
|
|
for buf in &self.buffers {
|
|
let found = buf.iter().position(|p| {
|
|
let diff = if p.header.timestamp_us >= ref_ts {
|
|
p.header.timestamp_us - ref_ts
|
|
} else {
|
|
ref_ts - p.header.timestamp_us
|
|
};
|
|
diff <= self.sync_tolerance_us
|
|
});
|
|
match found {
|
|
Some(idx) => indices.push(idx),
|
|
None => return None,
|
|
}
|
|
}
|
|
|
|
// Remove matched packets and merge channel data
|
|
let mut all_data: Vec<Vec<f64>> = Vec::new();
|
|
let mut sample_rate = 1000.0_f64;
|
|
|
|
for (buf_idx, &pkt_idx) in indices.iter().enumerate() {
|
|
let pkt = self.buffers[buf_idx].remove(pkt_idx);
|
|
sample_rate = pkt.header.sample_rate_hz as f64;
|
|
for ch in &pkt.channels {
|
|
let channel_data: Vec<f64> = ch
|
|
.samples
|
|
.iter()
|
|
.map(|&s| s as f64 * ch.scale_factor as f64)
|
|
.collect();
|
|
all_data.push(channel_data);
|
|
}
|
|
}
|
|
|
|
if all_data.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let timestamp = ref_ts as f64 / 1_000_000.0;
|
|
MultiChannelTimeSeries::new(all_data, sample_rate, timestamp).ok()
|
|
}
|
|
|
|
/// Set the timestamp tolerance in microseconds for matching packets
|
|
/// across nodes.
|
|
pub fn set_sync_tolerance(&mut self, tolerance_us: u64) {
|
|
self.sync_tolerance_us = tolerance_us;
|
|
}
|
|
|
|
/// Returns the number of buffered packets for a given node.
|
|
pub fn buffered_count(&self, node_id: usize) -> usize {
|
|
self.buffers.get(node_id).map_or(0, |b| b.len())
|
|
}
|
|
|
|
/// Returns the total number of expected nodes.
|
|
pub fn node_count(&self) -> usize {
|
|
self.node_count
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::protocol::{ChannelData, NeuralDataPacket, PacketHeader, PACKET_MAGIC, PROTOCOL_VERSION};
|
|
|
|
fn make_packet(num_channels: u8, timestamp_us: u64, samples: Vec<i16>) -> NeuralDataPacket {
|
|
let channels = (0..num_channels)
|
|
.map(|id| ChannelData {
|
|
channel_id: id,
|
|
samples: samples.clone(),
|
|
scale_factor: 1.0,
|
|
})
|
|
.collect();
|
|
|
|
NeuralDataPacket {
|
|
header: PacketHeader {
|
|
magic: PACKET_MAGIC,
|
|
version: PROTOCOL_VERSION,
|
|
packet_id: 0,
|
|
timestamp_us,
|
|
num_channels,
|
|
samples_per_channel: samples.len() as u16,
|
|
sample_rate_hz: 1000,
|
|
},
|
|
channels,
|
|
quality: vec![255; num_channels as usize],
|
|
checksum: 0,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_assemble_two_nodes() {
|
|
let mut agg = NodeAggregator::new(2);
|
|
|
|
let p0 = make_packet(1, 1000, vec![10, 20, 30]);
|
|
let p1 = make_packet(1, 1000, vec![40, 50, 60]);
|
|
|
|
agg.receive_packet(0, p0).unwrap();
|
|
// Only one node has reported — assembly requires all nodes
|
|
assert!(agg.try_assemble().is_none());
|
|
|
|
agg.receive_packet(1, p1).unwrap();
|
|
let ts = agg.try_assemble().unwrap();
|
|
assert_eq!(ts.num_channels, 2);
|
|
assert_eq!(ts.num_samples, 3);
|
|
assert!((ts.data[0][0] - 10.0).abs() < 1e-6);
|
|
assert!((ts.data[1][2] - 60.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_assemble_with_tolerance() {
|
|
let mut agg = NodeAggregator::new(2);
|
|
agg.set_sync_tolerance(500);
|
|
|
|
let p0 = make_packet(1, 1000, vec![1, 2]);
|
|
let p1 = make_packet(1, 1400, vec![3, 4]); // Within 500 us tolerance
|
|
|
|
agg.receive_packet(0, p0).unwrap();
|
|
agg.receive_packet(1, p1).unwrap();
|
|
assert!(agg.try_assemble().is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_assemble_exceeds_tolerance() {
|
|
let mut agg = NodeAggregator::new(2);
|
|
agg.set_sync_tolerance(100);
|
|
|
|
let p0 = make_packet(1, 1000, vec![1, 2]);
|
|
let p1 = make_packet(1, 2000, vec![3, 4]); // 1000 us apart > 100 us tolerance
|
|
|
|
agg.receive_packet(0, p0).unwrap();
|
|
agg.receive_packet(1, p1).unwrap();
|
|
assert!(agg.try_assemble().is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_receive_invalid_node() {
|
|
let mut agg = NodeAggregator::new(2);
|
|
let p = make_packet(1, 0, vec![1]);
|
|
assert!(agg.receive_packet(5, p).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_buffers_consumed_after_assembly() {
|
|
let mut agg = NodeAggregator::new(1);
|
|
let p = make_packet(1, 0, vec![1, 2, 3]);
|
|
agg.receive_packet(0, p).unwrap();
|
|
assert_eq!(agg.buffered_count(0), 1);
|
|
agg.try_assemble().unwrap();
|
|
assert_eq!(agg.buffered_count(0), 0);
|
|
}
|
|
}
|