1use 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#[derive(Clone, Debug)]
16pub struct PanelUnit {
17 pub unit_id: u32,
19 pub series: TimeSeriesData,
21}
22
23#[derive(Clone, Debug)]
27pub struct PanelData {
28 schema: Arc<CausalSchema>,
29 units: Arc<[PanelUnit]>,
30}
31
32impl PanelData {
33 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 #[must_use]
58 pub fn schema(&self) -> &CausalSchema {
59 &self.schema
60 }
61
62 #[must_use]
64 pub fn unit_count(&self) -> usize {
65 self.units.len()
66 }
67
68 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 #[must_use]
81 pub fn units(&self) -> &[PanelUnit] {
82 &self.units
83 }
84
85 #[must_use]
87 pub fn total_rows(&self) -> usize {
88 self.units.iter().map(|u| u.series.row_count()).sum()
89 }
90
91 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
105pub struct PanelUnitView<'a> {
107 unit: &'a PanelUnit,
108}
109
110impl<'a> PanelUnitView<'a> {
111 #[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}