Skip to main content

antecedent_data/
panel.rs

1//! Panel data: unit partitions with per-unit time indexes.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5use std::sync::Arc;
6
7use antecedent_core::{CausalSchema, VariableId};
8
9use crate::column::ColumnView;
10use crate::dataset::TimeSeriesData;
11use crate::error::DataError;
12use crate::table::TableView;
13
14/// One panel unit: a time series for a single cross-sectional unit.
15#[derive(Clone, Debug)]
16pub struct PanelUnit {
17    /// Unit label / dense index.
18    pub unit_id: u32,
19    /// Per-unit time series (borrowed schema must match the panel).
20    pub series: TimeSeriesData,
21}
22
23/// Panel dataset: multiple units sharing one schema, each with its own time index.
24///
25/// Units are stored by `Arc` so planners can reference them without cloning series.
26#[derive(Clone, Debug)]
27pub struct PanelData {
28    schema: Arc<CausalSchema>,
29    units: Arc<[PanelUnit]>,
30}
31
32impl PanelData {
33    /// Construct from units; all must share an identical schema.
34    ///
35    /// # Errors
36    ///
37    /// Empty list or schema mismatch.
38    pub fn try_new(units: impl Into<Arc<[PanelUnit]>>) -> Result<Self, DataError> {
39        let units = units.into();
40        if units.is_empty() {
41            return Err(DataError::InvalidArgument {
42                message: "panel data needs ≥1 unit".into()
43            });
44        }
45        let schema = Arc::new(units[0].series.schema().clone());
46        for u in units.iter().skip(1) {
47            if u.series.schema() != schema.as_ref() {
48                return Err(DataError::InvalidArgument {
49                    message: "panel unit schemas must match".into(),
50                });
51            }
52        }
53        Ok(Self { schema, units })
54    }
55
56    /// Shared schema.
57    #[must_use]
58    pub fn schema(&self) -> &CausalSchema {
59        &self.schema
60    }
61
62    /// Number of units.
63    #[must_use]
64    pub fn unit_count(&self) -> usize {
65        self.units.len()
66    }
67
68    /// Borrow unit `i`.
69    ///
70    /// # Errors
71    ///
72    /// Out of range.
73    pub fn unit(&self, i: usize) -> Result<&PanelUnit, DataError> {
74        self.units
75            .get(i)
76            .ok_or(DataError::InvalidArgument { message: "panel unit index out of range".into() })
77    }
78
79    /// All units.
80    #[must_use]
81    pub fn units(&self) -> &[PanelUnit] {
82        &self.units
83    }
84
85    /// Total rows across all units (sum of lengths; not a contiguous table).
86    #[must_use]
87    pub fn total_rows(&self) -> usize {
88        self.units.iter().map(|u| u.series.row_count()).sum()
89    }
90
91    /// View panel units as multi-environment series (shared schema, one env per unit).
92    ///
93    /// Used for pooled discovery (e.g. J-PCMCI+) without cloning payloads.
94    ///
95    /// # Errors
96    ///
97    /// Propagates [`MultiEnvironmentData::try_new`] failures (should not occur for a
98    /// well-formed panel).
99    pub fn as_multi_env(&self) -> Result<crate::multi_env::MultiEnvironmentData, DataError> {
100        let series: Vec<TimeSeriesData> = self.units.iter().map(|u| u.series.clone()).collect();
101        crate::multi_env::MultiEnvironmentData::try_new(Arc::from(series))
102    }
103}
104
105/// Borrowed view of a single panel unit as a table.
106pub struct PanelUnitView<'a> {
107    unit: &'a PanelUnit,
108}
109
110impl<'a> PanelUnitView<'a> {
111    /// Wrap a unit.
112    #[must_use]
113    pub fn new(unit: &'a PanelUnit) -> Self {
114        Self { unit }
115    }
116}
117
118impl TableView for PanelUnitView<'_> {
119    fn schema(&self) -> &CausalSchema {
120        self.unit.series.schema()
121    }
122
123    fn row_count(&self) -> usize {
124        self.unit.series.row_count()
125    }
126
127    fn column(&self, id: VariableId) -> Result<ColumnView<'_>, DataError> {
128        self.unit.series.column(id)
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::testing::float_series;
136
137    #[test]
138    fn rejects_empty() {
139        assert!(PanelData::try_new(Arc::from([])).is_err());
140    }
141
142    #[test]
143    fn builds_two_unit_panel() {
144        let panel = PanelData::try_new(Arc::from([
145            PanelUnit { unit_id: 0, series: float_series(10, 2) },
146            PanelUnit { unit_id: 1, series: float_series(12, 2) },
147        ]))
148        .unwrap();
149        assert_eq!(panel.unit_count(), 2);
150        assert_eq!(panel.total_rows(), 22);
151    }
152}