Skip to main content

antecedent_data/
multi_env.rs

1//! Multi-environment / multi-dataset container.
2//!
3//! Typed list of series sharing one schema. Sample planning without per-env
4//! full copies lives in [`crate::multi_env_plan`]. J-PCMCI+ discovery
5//! constraints are wired in discovery.
6//!
7//! SPDX-License-Identifier: MIT OR Apache-2.0
8
9use std::sync::Arc;
10
11use antecedent_core::CausalSchema;
12
13use crate::dataset::TimeSeriesData;
14use crate::error::DataError;
15use crate::table::TableView;
16
17/// Collection of time-series environments with a shared schema.
18#[derive(Clone, Debug)]
19pub struct MultiEnvironmentData {
20    schema: Arc<CausalSchema>,
21    environments: Arc<[TimeSeriesData]>,
22}
23
24impl MultiEnvironmentData {
25    /// Construct from one or more environments; all must share an identical schema.
26    ///
27    /// # Errors
28    ///
29    /// Empty list or schema mismatch across environments.
30    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    /// Shared schema.
49    #[must_use]
50    pub fn schema(&self) -> &CausalSchema {
51        &self.schema
52    }
53
54    /// Number of environments.
55    #[must_use]
56    pub fn env_count(&self) -> usize {
57        self.environments.len()
58    }
59
60    /// Borrow environment `i`.
61    ///
62    /// # Errors
63    ///
64    /// Out of range.
65    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    /// All environments.
72    #[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}