laddu_data/data/
event.rs

1use std::mem::size_of;
2use std::sync::Arc;
3
4use laddu_physics::vectors::RealVec4;
5
6use crate::{LadduDataError, LadduDataResult, schema::Schema};
7
8/// Immutable columnar batch of events sharing one schema.
9#[derive(Clone, Debug)]
10pub struct EventBatch {
11    schema: Arc<Schema>,
12    len: usize,
13    p4s: Arc<[Arc<[RealVec4]>]>,
14    scalars: Arc<[Arc<[f64]>]>,
15    weights: Option<Arc<[f64]>>,
16}
17
18impl EventBatch {
19    /// Validates column counts and lengths and constructs a batch.
20    ///
21    /// # Errors
22    ///
23    /// Returns [`LadduDataError`] when column counts do not match `schema` or
24    /// column and weight lengths are inconsistent.
25    pub fn new(
26        schema: Arc<Schema>,
27        p4s: Vec<Arc<[RealVec4]>>,
28        scalars: Vec<Arc<[f64]>>,
29        weights: Option<Arc<[f64]>>,
30    ) -> LadduDataResult<Self> {
31        if p4s.len() != schema.n_p4s() {
32            return Err(LadduDataError::Schema(
33                "wrong number of vec4 columns".into(),
34            ));
35        }
36
37        if scalars.len() != schema.n_scalars() {
38            return Err(LadduDataError::Schema(
39                "wrong number of scalar columns".into(),
40            ));
41        }
42
43        let len = infer_len(&p4s, &scalars, weights.as_deref())?;
44
45        Ok(Self {
46            schema,
47            len,
48            p4s: p4s.into(),
49            scalars: scalars.into(),
50            weights,
51        })
52    }
53
54    /// Collects owned row events into a columnar batch.
55    ///
56    /// # Errors
57    ///
58    /// Returns [`LadduDataError`] when an event has the wrong number of values
59    /// or weighted and unweighted events are mixed.
60    pub fn from_events<I>(schema: Arc<Schema>, events: I) -> LadduDataResult<Self>
61    where
62        I: IntoIterator<Item = OwnedEvent>,
63    {
64        let mut builder = EventBatchBuilder::new(schema);
65        builder.extend(events)?;
66        builder.finish()
67    }
68
69    /// Returns the shared logical schema.
70    pub fn schema(&self) -> &Arc<Schema> {
71        &self.schema
72    }
73
74    /// Returns the number of rows.
75    pub fn len(&self) -> usize {
76        self.len
77    }
78
79    /// Returns the logical payload bytes per event represented by this batch.
80    pub fn bytes_per_event(&self) -> usize {
81        self.p4s.len() * size_of::<RealVec4>()
82            + self.scalars.len() * size_of::<f64>()
83            + usize::from(self.weights.is_some()) * size_of::<f64>()
84    }
85
86    /// Returns the retained column payload size in bytes.
87    ///
88    /// Shared schema metadata, allocation headers, and other owners of shared
89    /// columns are not included.
90    pub fn resident_bytes(&self) -> usize {
91        self.p4s
92            .iter()
93            .map(|column| column.len() * size_of::<RealVec4>())
94            .sum::<usize>()
95            + self
96                .scalars
97                .iter()
98                .map(|column| column.len() * size_of::<f64>())
99                .sum::<usize>()
100            + self
101                .weights
102                .as_ref()
103                .map_or(0, |weights| weights.len() * size_of::<f64>())
104    }
105
106    /// Returns whether the batch contains no rows.
107    pub fn is_empty(&self) -> bool {
108        self.len == 0
109    }
110
111    /// Returns a four-momentum column by index.
112    pub fn vec4_column(&self, index: usize) -> &[RealVec4] {
113        &self.p4s[index]
114    }
115
116    /// Returns a scalar column by index.
117    pub fn scalar_column(&self, index: usize) -> &[f64] {
118        &self.scalars[index]
119    }
120
121    /// Returns the optional explicit weight column.
122    pub fn weights_column(&self) -> Option<&[f64]> {
123        self.weights.as_deref()
124    }
125
126    /// Returns a four-momentum column by logical name.
127    pub fn vec4_column_named(&self, name: &str) -> Option<&[RealVec4]> {
128        let i = self.schema.p4_index(name)?;
129        Some(self.vec4_column(i))
130    }
131
132    /// Returns a scalar column by logical name.
133    pub fn scalar_column_named(&self, name: &str) -> Option<&[f64]> {
134        let i = self.schema.scalar_index(name)?;
135        Some(self.scalar_column(i))
136    }
137
138    /// Returns one four-momentum cell.
139    pub fn p4_at(&self, col: usize, row: usize) -> RealVec4 {
140        self.p4s[col][row]
141    }
142
143    /// Returns one scalar cell.
144    pub fn scalar_at(&self, col: usize, row: usize) -> f64 {
145        self.scalars[col][row]
146    }
147
148    /// Returns the explicit row weight, or one when weights are absent.
149    pub fn weights_at(&self, row: usize) -> f64 {
150        self.weights.as_ref().map_or(1.0, |w| w[row])
151    }
152
153    /// Returns a borrowed view of one row.
154    pub fn event(&self, row: usize) -> BatchEvent<'_> {
155        BatchEvent { batch: self, row }
156    }
157
158    /// Iterates over borrowed event views.
159    pub fn iter(&self) -> impl Iterator<Item = BatchEvent<'_>> {
160        (0..self.len()).map(|i| self.event(i))
161    }
162
163    /// Copies selected rows into a new batch in the requested order.
164    pub fn select(&self, rows: &[usize]) -> Self {
165        let p4s = self
166            .p4s
167            .iter()
168            .map(|col| rows.iter().map(|&i| col[i]).collect())
169            .collect();
170
171        let scalars = self
172            .scalars
173            .iter()
174            .map(|col| rows.iter().map(|&i| col[i]).collect())
175            .collect();
176
177        let weights = self
178            .weights
179            .as_ref()
180            .map(|w| rows.iter().map(|&i| w[i]).collect());
181
182        Self {
183            schema: Arc::clone(&self.schema),
184            len: rows.len(),
185            p4s,
186            scalars,
187            weights,
188        }
189    }
190
191    /// Copies rows satisfying `keep` into a new batch.
192    pub fn filter<F>(&self, keep: F) -> Self
193    where
194        F: Fn(BatchEvent<'_>) -> bool,
195    {
196        let rows: Vec<usize> = (0..self.len).filter(|&i| keep(self.event(i))).collect();
197
198        self.select(&rows)
199    }
200
201    /// Returns a batch sharing value columns with newly computed weights.
202    pub fn reweight<F>(&self, f: F) -> Self
203    where
204        F: Fn(usize, f64) -> f64,
205    {
206        let weights: Arc<[f64]> = (0..self.len).map(|i| f(i, self.weights_at(i))).collect();
207
208        Self {
209            schema: Arc::clone(&self.schema),
210            len: self.len,
211            p4s: Arc::clone(&self.p4s),
212            scalars: Arc::clone(&self.scalars),
213            weights: Some(weights),
214        }
215    }
216
217    /// Copies the half-open row range `start..end` into a new batch.
218    ///
219    /// # Panics
220    ///
221    /// Panics when `start > end` or `end` exceeds the batch length.
222    pub fn slice(&self, start: usize, end: usize) -> Self {
223        assert!(start <= end);
224        assert!(end <= self.len);
225
226        if start == 0 && end == self.len {
227            return self.clone();
228        }
229
230        let p4s = self
231            .p4s
232            .iter()
233            .map(|col| Arc::<[RealVec4]>::from(&col[start..end]))
234            .collect();
235
236        let scalars = self
237            .scalars
238            .iter()
239            .map(|col| Arc::<[f64]>::from(&col[start..end]))
240            .collect();
241
242        let weights = self
243            .weights
244            .as_ref()
245            .map(|w| Arc::<[f64]>::from(&w[start..end]));
246
247        Self {
248            schema: Arc::clone(&self.schema),
249            len: end - start,
250            p4s,
251            scalars,
252            weights,
253        }
254    }
255
256    /// Concatenates schema-compatible batches.
257    ///
258    /// # Errors
259    ///
260    /// Returns [`LadduDataError`] when `batches` is empty or contains
261    /// incompatible schemas.
262    pub fn concat(batches: &[Self]) -> LadduDataResult<Self> {
263        if batches.is_empty() {
264            return Err(LadduDataError::InvalidArgument(
265                "cannot concatenate zero batches",
266            ));
267        }
268
269        let schema = Arc::clone(&batches[0].schema);
270        let len: usize = batches.iter().map(|b| b.len).sum();
271
272        for batch in batches {
273            if schema != batch.schema {
274                return Err(LadduDataError::Schema(
275                    "cannot concatenate batches with different schemas".into(),
276                ));
277            }
278        }
279
280        let mut p4s = Vec::with_capacity(schema.n_p4s());
281
282        for col in 0..schema.n_p4s() {
283            let mut out = Vec::with_capacity(len);
284            for batch in batches {
285                out.extend_from_slice(batch.vec4_column(col));
286            }
287            p4s.push(Arc::from(out));
288        }
289
290        let mut scalars = Vec::with_capacity(schema.n_scalars());
291
292        for col in 0..schema.n_scalars() {
293            let mut out = Vec::with_capacity(len);
294            for batch in batches {
295                out.extend_from_slice(batch.scalar_column(col));
296            }
297            scalars.push(Arc::from(out));
298        }
299
300        let any_weights = batches.iter().any(|b| b.weights.is_some());
301
302        let weights = if any_weights {
303            let mut out = Vec::with_capacity(len);
304            for batch in batches {
305                for i in 0..batch.len {
306                    out.push(batch.weights_at(i));
307                }
308            }
309            Some(Arc::from(out))
310        } else {
311            None
312        };
313
314        Ok(Self {
315            schema,
316            len,
317            p4s: p4s.into(),
318            scalars: scalars.into(),
319            weights,
320        })
321    }
322}
323
324fn infer_len(
325    vec4s: &[Arc<[RealVec4]>],
326    scalars: &[Arc<[f64]>],
327    weight: Option<&[f64]>,
328) -> LadduDataResult<usize> {
329    let len = vec4s
330        .first()
331        .map(|c| c.len())
332        .or_else(|| scalars.first().map(|c| c.len()))
333        .or_else(|| weight.map(|w| w.len()))
334        .unwrap_or(0);
335
336    for col in vec4s {
337        if col.len() != len {
338            return Err(LadduDataError::Schema(
339                "inconsistent vec4 column length".into(),
340            ));
341        }
342    }
343
344    for col in scalars {
345        if col.len() != len {
346            return Err(LadduDataError::Schema(
347                "inconsistent scalar column length".into(),
348            ));
349        }
350    }
351
352    if let Some(w) = weight
353        && w.len() != len
354    {
355        return Err(LadduDataError::Schema("inconsistent weight length".into()));
356    }
357
358    Ok(len)
359}
360
361/// Borrowed view of one row in an [`EventBatch`].
362#[derive(Copy, Clone, Debug)]
363pub struct BatchEvent<'a> {
364    batch: &'a EventBatch,
365    row: usize,
366}
367
368impl<'a> BatchEvent<'a> {
369    /// Returns the row index.
370    pub fn row(&self) -> usize {
371        self.row
372    }
373
374    /// Returns the backing batch.
375    pub fn batch(&self) -> &'a EventBatch {
376        self.batch
377    }
378
379    /// Returns a four-momentum value by column index.
380    pub fn p4(&self, col: usize) -> RealVec4 {
381        self.batch.p4_at(col, self.row)
382    }
383
384    /// Returns a scalar value by column index.
385    pub fn scalar(&self, col: usize) -> f64 {
386        self.batch.scalar_at(col, self.row)
387    }
388
389    /// Returns the row weight, defaulting to one.
390    pub fn weight(&self) -> f64 {
391        self.batch.weights_at(self.row)
392    }
393
394    /// Returns a four-momentum value by logical name.
395    pub fn p4_named(&self, name: &str) -> Option<RealVec4> {
396        let col = self.batch.schema.p4_index(name)?;
397        Some(self.p4(col))
398    }
399
400    /// Returns a scalar value by logical name.
401    pub fn scalar_named(&self, name: &str) -> Option<f64> {
402        let col = self.batch.schema.scalar_index(name)?;
403        Some(self.scalar(col))
404    }
405}
406
407/// Borrowed event view with a possibly transformed weight.
408#[derive(Copy, Clone, Debug)]
409pub struct Event<'a> {
410    pub(super) batch: &'a EventBatch,
411    pub(super) row: usize,
412    pub(super) weight: f64,
413}
414
415impl<'a> Event<'a> {
416    /// Returns the row index in the backing batch.
417    pub fn row(&self) -> usize {
418        self.row
419    }
420
421    /// Returns a four-momentum value by column index.
422    pub fn p4(&self, col: usize) -> RealVec4 {
423        self.batch.p4_at(col, self.row)
424    }
425
426    /// Returns a scalar value by column index.
427    pub fn scalar(&self, col: usize) -> f64 {
428        self.batch.scalar_at(col, self.row)
429    }
430
431    /// Returns this view's effective weight.
432    pub fn weight(&self) -> f64 {
433        self.weight
434    }
435
436    /// Returns a four-momentum value by logical name.
437    pub fn p4_named(&self, name: &str) -> Option<RealVec4> {
438        let col = self.batch.schema.p4_index(name)?;
439        Some(self.p4(col))
440    }
441
442    /// Returns a scalar value by logical name.
443    pub fn scalar_named(&self, name: &str) -> Option<f64> {
444        let col = self.batch.schema.scalar_index(name)?;
445        Some(self.scalar(col))
446    }
447}
448
449/// Owned row-oriented event used while constructing batches.
450#[derive(Clone, Debug)]
451pub struct OwnedEvent {
452    /// Four-momentum values in schema order.
453    pub p4s: Vec<RealVec4>,
454    /// Scalar values in schema order.
455    pub scalars: Vec<f64>,
456    /// Optional explicit event weight.
457    pub weight: Option<f64>,
458}
459
460impl OwnedEvent {
461    /// Creates an unweighted owned event.
462    pub fn new(p4s: Vec<RealVec4>, scalars: Vec<f64>) -> Self {
463        Self {
464            p4s,
465            scalars,
466            weight: None,
467        }
468    }
469
470    /// Creates an owned event with an explicit weight.
471    pub fn weighted(p4s: Vec<RealVec4>, scalars: Vec<f64>, weight: f64) -> Self {
472        Self {
473            p4s,
474            scalars,
475            weight: Some(weight),
476        }
477    }
478}
479
480/// Incremental builder for a columnar [`EventBatch`].
481pub struct EventBatchBuilder {
482    schema: Arc<Schema>,
483    p4s: Vec<Vec<RealVec4>>,
484    scalars: Vec<Vec<f64>>,
485    weights: Option<Vec<f64>>,
486    len: usize,
487}
488
489impl EventBatchBuilder {
490    /// Creates an empty builder.
491    pub fn new(schema: Arc<Schema>) -> Self {
492        let p4s = (0..schema.n_p4s()).map(|_| Vec::new()).collect();
493        let scalars = (0..schema.n_scalars()).map(|_| Vec::new()).collect();
494
495        Self {
496            schema,
497            p4s,
498            scalars,
499            weights: None,
500            len: 0,
501        }
502    }
503
504    /// Creates an empty builder with per-column capacity.
505    pub fn with_capacity(schema: Arc<Schema>, capacity: usize) -> Self {
506        let p4s = (0..schema.n_p4s())
507            .map(|_| Vec::with_capacity(capacity))
508            .collect();
509
510        let scalars = (0..schema.n_scalars())
511            .map(|_| Vec::with_capacity(capacity))
512            .collect();
513
514        Self {
515            schema,
516            p4s,
517            scalars,
518            weights: None,
519            len: 0,
520        }
521    }
522
523    /// Appends an unweighted event from ordered values.
524    ///
525    /// # Errors
526    ///
527    /// Returns [`LadduDataError`] when value counts do not match the schema or
528    /// the builder already contains weighted events.
529    pub fn push<P, S>(&mut self, p4s: P, scalars: S) -> LadduDataResult<&mut Self>
530    where
531        P: IntoIterator<Item = RealVec4>,
532        S: IntoIterator<Item = f64>,
533    {
534        self.push_event(OwnedEvent::new(
535            p4s.into_iter().collect(),
536            scalars.into_iter().collect(),
537        ))
538    }
539
540    /// Appends a weighted event from ordered values.
541    ///
542    /// # Errors
543    ///
544    /// Returns [`LadduDataError`] when value counts do not match the schema or
545    /// the builder already contains unweighted events.
546    pub fn push_weighted<P, S>(
547        &mut self,
548        p4s: P,
549        scalars: S,
550        weight: f64,
551    ) -> LadduDataResult<&mut Self>
552    where
553        P: IntoIterator<Item = RealVec4>,
554        S: IntoIterator<Item = f64>,
555    {
556        self.push_event(OwnedEvent::weighted(
557            p4s.into_iter().collect(),
558            scalars.into_iter().collect(),
559            weight,
560        ))
561    }
562
563    /// Validates and appends one owned event.
564    ///
565    /// # Errors
566    ///
567    /// Returns [`LadduDataError`] when the event shape does not match the
568    /// schema or its weight presence differs from prior events.
569    pub fn push_event(&mut self, event: OwnedEvent) -> LadduDataResult<&mut Self> {
570        if event.p4s.len() != self.schema.n_p4s() {
571            return Err(LadduDataError::Schema(
572                "wrong number of event vec4 values".into(),
573            ));
574        }
575
576        if event.scalars.len() != self.schema.n_scalars() {
577            return Err(LadduDataError::Schema(
578                "wrong number of event scalar values".into(),
579            ));
580        }
581
582        match (&mut self.weights, event.weight) {
583            (Some(weights), Some(weight)) => weights.push(weight),
584
585            (Some(_), None) => {
586                return Err(LadduDataError::InvalidArgument(
587                    "cannot mix weighted and unweighted events in one batch",
588                ));
589            }
590
591            (None, Some(weight)) if self.len == 0 => {
592                self.weights = Some(vec![weight]);
593            }
594
595            (None, Some(_)) => {
596                return Err(LadduDataError::InvalidArgument(
597                    "cannot mix unweighted and weighted events in one batch",
598                ));
599            }
600
601            (None, None) => {}
602        }
603
604        for (col, value) in event.p4s.into_iter().enumerate() {
605            self.p4s[col].push(value);
606        }
607
608        for (col, value) in event.scalars.into_iter().enumerate() {
609            self.scalars[col].push(value);
610        }
611
612        self.len += 1;
613
614        Ok(self)
615    }
616
617    /// Appends all owned events from an iterator.
618    ///
619    /// # Errors
620    ///
621    /// Returns the first [`LadduDataError`] produced by an event whose shape or
622    /// weight presence is incompatible with the builder.
623    pub fn extend<I>(&mut self, events: I) -> LadduDataResult<&mut Self>
624    where
625        I: IntoIterator<Item = OwnedEvent>,
626    {
627        for event in events {
628            self.push_event(event)?;
629        }
630
631        Ok(self)
632    }
633
634    /// Finalizes the builder into an immutable batch.
635    ///
636    /// # Errors
637    ///
638    /// Returns [`LadduDataError`] if the accumulated columns or weights have
639    /// inconsistent lengths.
640    pub fn finish(self) -> LadduDataResult<EventBatch> {
641        let p4s = self.p4s.into_iter().map(Arc::from).collect();
642        let scalars = self.scalars.into_iter().map(Arc::from).collect();
643        let weights = self.weights.map(Arc::from);
644
645        EventBatch::new(self.schema, p4s, scalars, weights)
646    }
647}
648
649#[cfg(test)]
650mod tests {
651    use super::*;
652
653    fn v(x: f64) -> RealVec4 {
654        RealVec4 {
655            e: x + 0.3,
656            px: x,
657            py: x + 0.1,
658            pz: x + 0.2,
659        }
660    }
661
662    fn schema_with_weight() -> Arc<Schema> {
663        Arc::new(Schema::new(["p"], ["x"], true).unwrap())
664    }
665
666    fn weighted_batch(start: usize, len: usize) -> EventBatch {
667        let schema = schema_with_weight();
668
669        let events = (start..start + len)
670            .map(|i| OwnedEvent::weighted(vec![v(i as f64)], vec![i as f64], 10.0 + i as f64));
671
672        EventBatch::from_events(schema, events).unwrap()
673    }
674
675    fn scalar_values(batch: &EventBatch) -> Vec<f64> {
676        batch.scalar_column(0).to_vec()
677    }
678
679    #[test]
680    fn event_batch_rejects_shape_mismatches_and_builder_rejects_mixed_weights() {
681        let schema = schema_with_weight();
682
683        let bad_vec4_count = EventBatch::new(
684            Arc::clone(&schema),
685            vec![],
686            vec![Arc::from([1.0, 2.0])],
687            Some(Arc::from([1.0, 2.0])),
688        );
689
690        assert!(matches!(bad_vec4_count, Err(LadduDataError::Schema(_))));
691
692        let bad_lengths = EventBatch::new(
693            Arc::clone(&schema),
694            vec![Arc::from([v(1.0), v(2.0)])],
695            vec![Arc::from([1.0])],
696            Some(Arc::from([1.0, 2.0])),
697        );
698
699        assert!(matches!(bad_lengths, Err(LadduDataError::Schema(_))));
700
701        let mut builder = EventBatchBuilder::new(schema);
702        builder.push([v(1.0)], [1.0]).unwrap();
703
704        let mixed = builder.push_weighted([v(2.0)], [2.0], 2.0);
705        assert!(matches!(mixed, Err(LadduDataError::InvalidArgument(_))));
706    }
707
708    #[test]
709    fn select_slice_filter_reweight_and_concat_preserve_columns_and_weight_semantics() {
710        let weighted = weighted_batch(0, 4);
711        let selected = weighted.select(&[3, 1]);
712
713        assert_eq!(scalar_values(&selected), vec![3.0, 1.0]);
714        assert_eq!(selected.weights_column().unwrap(), &[13.0, 11.0]);
715        assert_eq!(selected.p4_at(0, 0).px, 3.0);
716        assert_eq!(selected.p4_at(0, 1).e, 1.3);
717
718        let sliced = weighted.slice(1, 3);
719        assert_eq!(scalar_values(&sliced), vec![1.0, 2.0]);
720        assert_eq!(sliced.weights_column().unwrap(), &[11.0, 12.0]);
721
722        let filtered = weighted.filter(|ev| ev.scalar(0) >= 2.0);
723        assert_eq!(scalar_values(&filtered), vec![2.0, 3.0]);
724
725        let reweighted = filtered.reweight(|i, w| w + 100.0 + i as f64);
726        assert_eq!(reweighted.weights_column().unwrap(), &[112.0, 114.0]);
727
728        let schema = schema_with_weight();
729
730        let unweighted_with_weight_schema = EventBatch::from_events(
731            Arc::clone(&schema),
732            [
733                OwnedEvent::new(vec![v(100.0)], vec![100.0]),
734                OwnedEvent::new(vec![v(101.0)], vec![101.0]),
735            ],
736        )
737        .unwrap();
738
739        let weighted_tail = EventBatch::from_events(
740            schema,
741            [
742                OwnedEvent::weighted(vec![v(200.0)], vec![200.0], 5.0),
743                OwnedEvent::weighted(vec![v(201.0)], vec![201.0], 6.0),
744            ],
745        )
746        .unwrap();
747
748        let concatenated =
749            EventBatch::concat(&[unweighted_with_weight_schema, weighted_tail]).unwrap();
750
751        assert_eq!(
752            scalar_values(&concatenated),
753            vec![100.0, 101.0, 200.0, 201.0]
754        );
755        assert_eq!(
756            concatenated.weights_column().unwrap(),
757            &[1.0, 1.0, 5.0, 6.0]
758        );
759    }
760}
laddu_data/data/event.rs

laddu_data/data/
event.rs
