antecedent_data/
multi_env.rs1use std::sync::Arc;
10
11use antecedent_core::CausalSchema;
12
13use crate::dataset::TimeSeriesData;
14use crate::error::DataError;
15use crate::table::TableView;
16
17#[derive(Clone, Debug)]
19pub struct MultiEnvironmentData {
20 schema: Arc<CausalSchema>,
21 environments: Arc<[TimeSeriesData]>,
22}
23
24impl MultiEnvironmentData {
25 pub fn try_new(environments: impl Into<Arc<[TimeSeriesData]>>) -> Result<Self, DataError> {
31 let environments = environments.into();
32 if environments.is_empty() {
33 return Err(DataError::InvalidArgument {
34 message: "multi-environment data needs ≥1 environment".into(),
35 });
36 }
37 let schema = Arc::new(environments[0].schema().clone());
38 for env in environments.iter().skip(1) {
39 if env.schema() != schema.as_ref() {
40 return Err(DataError::InvalidArgument {
41 message: "environment schemas must match".into(),
42 });
43 }
44 }
45 Ok(Self { schema, environments })
46 }
47
48 #[must_use]
50 pub fn schema(&self) -> &CausalSchema {
51 &self.schema
52 }
53
54 #[must_use]
56 pub fn env_count(&self) -> usize {
57 self.environments.len()
58 }
59
60 pub fn environment(&self, i: usize) -> Result<&TimeSeriesData, DataError> {
66 self.environments
67 .get(i)
68 .ok_or(DataError::InvalidArgument { message: "environment index out of range".into() })
69 }
70
71 #[must_use]
73 pub fn environments(&self) -> &[TimeSeriesData] {
74 &self.environments
75 }
76}
77
78#[cfg(test)]
79#[allow(clippy::cast_precision_loss)]
80mod tests {
81 use std::sync::Arc;
82
83 use super::*;
84 use crate::testing::float_series;
85
86 #[test]
87 fn multi_env_requires_matching_schema() {
88 let a = float_series(10, 1);
89 let b = float_series(20, 1);
90 let m = MultiEnvironmentData::try_new(Arc::from([a, b])).unwrap();
91 assert_eq!(m.env_count(), 2);
92 assert_eq!(m.environment(1).unwrap().row_count(), 20);
93 }
94}