antecedent_data/
vector_vars.rs

1//! pinned baseline-style vector variable groups.
2//!
3//! Components are separate Float64 columns. Logical discovery nodes are the
4//! first component of each group; [`column_blocks_for_frame`] builds per-lag
5//! frame-index blocks for pairwise multivariate CI.
6//!
7//! SPDX-License-Identifier: MIT OR Apache-2.0
8
9use std::collections::HashSet;
10use std::num::NonZeroU32;
11use std::sync::Arc;
12
13use antecedent_core::{CausalSchemaBuilder, Lag, MeasurementSpec, ValueType, VariableId};
14
15use crate::column::{Float64Column, OwnedColumn};
16use crate::dataset::TimeSeriesData;
17use crate::error::DataError;
18use crate::lagged_frame::LaggedFrame;
19use crate::storage::OwnedColumnarStorage;
20use crate::table::TableView;
21
22/// Ordered groups of component variable ids (pinned baseline `vector_vars` style).
23///
24/// The first id in each group is the **logical** discovery node; remaining ids
25/// are CI-only components excluded from the search variable list.
26#[derive(Clone, Debug, Default, Eq, PartialEq)]
27pub struct VectorVariableGroups {
28    groups: Arc<[Arc<[VariableId]>]>,
29}
30
31impl VectorVariableGroups {
32    /// Empty groups (scalar CI).
33    #[must_use]
34    pub fn empty() -> Self {
35        Self { groups: Arc::from([]) }
36    }
37
38    /// Construct from component groups.
39    ///
40    /// # Errors
41    ///
42    /// Empty group or duplicate component across groups.
43    pub fn try_new(groups: impl Into<Arc<[Arc<[VariableId]>]>>) -> Result<Self, DataError> {
44        let groups = groups.into();
45        let mut seen = HashSet::new();
46        for g in groups.iter() {
47            if g.is_empty() {
48                return Err(DataError::InvalidArgument {
49                    message: "vector variable group must be non-empty".into(),
50                });
51            }
52            for &id in g.iter() {
53                if !seen.insert(id) {
54                    return Err(DataError::InvalidArgument {
55                        message: format!(
56                            "vector variable component {id} appears in multiple groups"
57                        ),
58                    });
59                }
60            }
61        }
62        Ok(Self { groups })
63    }
64
65    /// Whether any group is present.
66    #[must_use]
67    pub fn is_empty(&self) -> bool {
68        self.groups.is_empty()
69    }
70
71    /// Borrow groups.
72    #[must_use]
73    pub fn groups(&self) -> &[Arc<[VariableId]>] {
74        &self.groups
75    }
76
77    /// Logical discovery nodes (first component of each group).
78    #[must_use]
79    pub fn logical_ids(&self) -> Vec<VariableId> {
80        self.groups.iter().filter_map(|g| g.first().copied()).collect()
81    }
82
83    /// Component ids that are not logical heads (excluded from search vars).
84    #[must_use]
85    pub fn secondary_component_ids(&self) -> Vec<VariableId> {
86        let mut out = Vec::new();
87        for g in self.groups.iter() {
88            for &id in g.iter().skip(1) {
89                out.push(id);
90            }
91        }
92        out
93    }
94
95    /// Filter `variables` to logical search nodes (drop secondary components).
96    #[must_use]
97    pub fn filter_search_variables(&self, variables: &[VariableId]) -> Vec<VariableId> {
98        let secondary: HashSet<VariableId> = self.secondary_component_ids().into_iter().collect();
99        variables.iter().copied().filter(|v| !secondary.contains(v)).collect()
100    }
101}
102
103/// Build per-lag column blocks for a lagged frame from vector-variable groups.
104///
105/// For each multi-component group and each lag `0..=max_lag`, emits one block of
106/// frame column indexes `[idx(c0,τ), idx(c1,τ), …]`.
107///
108/// # Errors
109///
110/// Missing lagged columns for a component.
111pub fn column_blocks_for_frame(
112    groups: &VectorVariableGroups,
113    frame: &LaggedFrame,
114) -> Result<Arc<[Arc<[usize]>]>, DataError> {
115    if groups.is_empty() {
116        return Ok(Arc::from([]));
117    }
118    let max_lag = frame.max_lag();
119    let mut blocks: Vec<Arc<[usize]>> = Vec::new();
120    for g in groups.groups() {
121        if g.len() < 2 {
122            continue;
123        }
124        for lag in 0..=max_lag {
125            let lag = Lag::from_raw(lag);
126            let mut block = Vec::with_capacity(g.len());
127            for &id in g.iter() {
128                let Some(idx) = frame.column_index(id, lag) else {
129                    return Err(DataError::InvalidArgument {
130                        message: format!(
131                            "vector component {id} lag {lag:?} missing from lagged frame"
132                        ),
133                    });
134                };
135                block.push(idx);
136            }
137            blocks.push(Arc::from(block));
138        }
139    }
140    Ok(Arc::from(blocks))
141}
142
143/// Expand every [`FixedVectorColumn`] in `data` into `dim` Float64 columns and
144/// register them as vector-variable groups.
145///
146/// The original vector variable id becomes the first (logical) component; additional
147/// component ids are appended to the schema as `{name}__c{k}`.
148///
149/// # Errors
150///
151/// Schema construction failure, width mismatch, or unsupported non-float companion columns.
152pub fn expand_fixed_vector_columns(
153    data: &TimeSeriesData,
154) -> Result<(TimeSeriesData, VectorVariableGroups), DataError> {
155    let storage = data.storage();
156    let n = storage.row_count();
157    let schema = storage.schema();
158    let mut builder = CausalSchemaBuilder::new();
159    let mut new_cols: Vec<OwnedColumn> = Vec::new();
160    let mut groups: Vec<Arc<[VariableId]>> = Vec::new();
161    let mut next_id = 0u32;
162
163    for old in storage.columns() {
164        match old {
165            OwnedColumn::FixedVector(fv) => {
166                let width = NonZeroU32::new(u32::try_from(fv.dim).map_err(|_| {
167                    DataError::InvalidArgument { message: "fixed vector dim overflow".into() }
168                })?)
169                .ok_or_else(|| DataError::InvalidArgument {
170                    message: "fixed vector dim must be > 0".into(),
171                })?;
172                let old_schema = schema.get(fv.id).map_err(|e| DataError::Schema(e.to_string()))?;
173                if let ValueType::Vector { width: declared, .. } = &old_schema.value_type {
174                    if declared != &width {
175                        return Err(DataError::InvalidArgument {
176                            message: format!(
177                                "FixedVector dim {} != schema Vector width {}",
178                                fv.dim,
179                                declared.get()
180                            ),
181                        });
182                    }
183                }
184                let mut comp_ids = Vec::with_capacity(fv.dim);
185                for k in 0..fv.dim {
186                    let name = if k == 0 {
187                        Arc::clone(&old_schema.name)
188                    } else {
189                        Arc::<str>::from(format!("{}__c{k}", old_schema.name))
190                    };
191                    builder
192                        .add_variable(
193                            name,
194                            ValueType::Continuous,
195                            old_schema.role_hints,
196                            old_schema.unit.clone(),
197                            None,
198                            MeasurementSpec::default(),
199                        )
200                        .map_err(|e| DataError::Schema(e.to_string()))?;
201                    let id = VariableId::from_raw(next_id);
202                    next_id += 1;
203                    comp_ids.push(id);
204                    let mut values = vec![0.0; n];
205                    for (row, slot) in values.iter_mut().enumerate() {
206                        *slot = fv.values[row * fv.dim + k];
207                    }
208                    new_cols.push(OwnedColumn::Float64(Float64Column::new(
209                        id,
210                        Arc::from(values),
211                        fv.validity.clone(),
212                    )?));
213                }
214                groups.push(Arc::from(comp_ids));
215            }
216            OwnedColumn::Float64(c) => {
217                let old_schema = schema.get(c.id).map_err(|e| DataError::Schema(e.to_string()))?;
218                builder
219                    .add_variable(
220                        Arc::clone(&old_schema.name),
221                        old_schema.value_type.clone(),
222                        old_schema.role_hints,
223                        old_schema.unit.clone(),
224                        old_schema.category_domain,
225                        old_schema.measurement.clone(),
226                    )
227                    .map_err(|e| DataError::Schema(e.to_string()))?;
228                let new_id = VariableId::from_raw(next_id);
229                next_id += 1;
230                new_cols.push(OwnedColumn::Float64(Float64Column::new(
231                    new_id,
232                    c.values.clone(),
233                    c.validity.clone(),
234                )?));
235            }
236            other => {
237                return Err(DataError::InvalidArgument {
238                    message: format!(
239                        "expand_fixed_vector_columns: unsupported column kind for id {}",
240                        other.id()
241                    ),
242                });
243            }
244        }
245    }
246
247    let new_schema = builder.build().map_err(|e| DataError::Schema(e.to_string()))?;
248    let new_storage = OwnedColumnarStorage::try_new(
249        new_schema,
250        new_cols,
251        storage.analysis_mask().cloned(),
252        storage.weights().map(Arc::from),
253    )?;
254    let series = TimeSeriesData::try_new(new_storage, data.time_index().clone())?;
255    let groups = VectorVariableGroups::try_new(Arc::from(groups))?;
256    Ok((series, groups))
257}
258
259#[cfg(test)]
260#[allow(clippy::cast_precision_loss)]
261mod tests {
262    use super::*;
263    use crate::column::{FixedVectorColumn, ValidityBitmap};
264    use crate::testing::float_series;
265
266    #[test]
267    fn column_blocks_per_lag() {
268        let data = float_series(20, 3);
269        let vars = [VariableId::from_raw(0), VariableId::from_raw(1), VariableId::from_raw(2)];
270        let frame = LaggedFrame::from_series(
271            &data,
272            &vars,
273            1,
274            &antecedent_core::KernelPolicy::default_policy(),
275        )
276        .unwrap();
277        let groups = VectorVariableGroups::try_new(Arc::from([Arc::from([
278            VariableId::from_raw(0),
279            VariableId::from_raw(1),
280        ])]))
281        .unwrap();
282        let blocks = column_blocks_for_frame(&groups, &frame).unwrap();
283        // 2 lags (0,1) × 1 group
284        assert_eq!(blocks.len(), 2);
285        assert_eq!(blocks[0].len(), 2);
286        assert_eq!(
287            blocks[0].as_ref(),
288            &[
289                frame.column_index(VariableId::from_raw(0), Lag::CONTEMPORANEOUS).unwrap(),
290                frame.column_index(VariableId::from_raw(1), Lag::CONTEMPORANEOUS).unwrap(),
291            ]
292        );
293    }
294
295    #[test]
296    fn filter_drops_secondary_components() {
297        let groups = VectorVariableGroups::try_new(Arc::from([Arc::from([
298            VariableId::from_raw(0),
299            VariableId::from_raw(1),
300        ])]))
301        .unwrap();
302        let search = groups.filter_search_variables(&[
303            VariableId::from_raw(0),
304            VariableId::from_raw(1),
305            VariableId::from_raw(2),
306        ]);
307        assert_eq!(search, vec![VariableId::from_raw(0), VariableId::from_raw(2)]);
308    }
309
310    #[test]
311    fn expand_fixed_vector_round_trip() {
312        use crate::temporal::{SamplingRegularity, TimeIndex};
313        use antecedent_core::{
314            CausalSchemaBuilder, MeasurementSpec, RoleHint, SmallRoleSet, ValueType,
315        };
316
317        let n = 10usize;
318        let dim = 2usize;
319        let mut b = CausalSchemaBuilder::new();
320        b.add_variable(
321            "v",
322            ValueType::Vector {
323                width: NonZeroU32::new(2).unwrap(),
324                element: antecedent_core::ScalarType::Float64,
325            },
326            SmallRoleSet::from_hint(RoleHint::Context),
327            None,
328            None,
329            MeasurementSpec::default(),
330        )
331        .unwrap();
332        b.add_variable(
333            "y",
334            ValueType::Continuous,
335            SmallRoleSet::from_hint(RoleHint::Context),
336            None,
337            None,
338            MeasurementSpec::default(),
339        )
340        .unwrap();
341        let schema = b.build().unwrap();
342        let mut fv_vals = vec![0.0; n * dim];
343        let mut y = vec![0.0; n];
344        for t in 0..n {
345            fv_vals[t * dim] = t as f64;
346            fv_vals[t * dim + 1] = (t as f64) + 100.0;
347            y[t] = t as f64;
348        }
349        let cols = vec![
350            OwnedColumn::FixedVector(
351                FixedVectorColumn::new(
352                    VariableId::from_raw(0),
353                    dim,
354                    Arc::from(fv_vals),
355                    ValidityBitmap::all_valid(n),
356                )
357                .unwrap(),
358            ),
359            OwnedColumn::Float64(
360                Float64Column::new(
361                    VariableId::from_raw(1),
362                    Arc::from(y),
363                    ValidityBitmap::all_valid(n),
364                )
365                .unwrap(),
366            ),
367        ];
368        let storage = OwnedColumnarStorage::try_new(schema, cols, None, None).unwrap();
369        let data = TimeSeriesData::try_new(
370            storage,
371            TimeIndex { regularity: SamplingRegularity::Regular { interval_ns: 1 }, length: n },
372        )
373        .unwrap();
374        let (expanded, groups) = expand_fixed_vector_columns(&data).unwrap();
375        assert_eq!(expanded.storage().columns().len(), 3);
376        assert_eq!(groups.groups().len(), 1);
377        assert_eq!(groups.groups()[0].len(), 2);
378        assert_eq!(groups.logical_ids(), vec![VariableId::from_raw(0)]);
379    }
380}
antecedent_data/vector_vars.rs

antecedent_data/
vector_vars.rs
