Skip to main content

antecedent_model/
model_collection.rs

1//! Weighted collection of fitted per-graph causal models.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5use std::sync::Arc;
6
7use crate::compile::CompiledCausalModel;
8use crate::error::ModelError;
9
10/// Collection of fitted models weighted by graph posterior mass.
11#[derive(Clone, Debug)]
12pub struct ModelCollection {
13    /// Per-graph compiled models.
14    pub models: Arc<[CompiledCausalModel]>,
15    /// Graph keys aligned with `models`.
16    pub graph_keys: Arc<[u64]>,
17    /// Normalized weights (sum to 1 over identified graphs).
18    pub weights: Arc<[f64]>,
19}
20
21impl ModelCollection {
22    /// Build from parallel arrays.
23    ///
24    /// # Errors
25    ///
26    /// Length mismatch or non-positive weight sum.
27    pub fn new(
28        models: impl Into<Arc<[CompiledCausalModel]>>,
29        graph_keys: impl Into<Arc<[u64]>>,
30        weights: impl Into<Arc<[f64]>>,
31    ) -> Result<Self, ModelError> {
32        let models = models.into();
33        let graph_keys = graph_keys.into();
34        let weights = weights.into();
35        if models.len() != graph_keys.len() || models.len() != weights.len() {
36            return Err(ModelError::Shape { message: "ModelCollection length mismatch".into() });
37        }
38        let sum: f64 = weights.iter().sum();
39        if sum.partial_cmp(&0.0) != Some(std::cmp::Ordering::Greater) {
40            return Err(ModelError::Shape {
41                message: "ModelCollection weights non-positive".into(),
42            });
43        }
44        let weights: Arc<[f64]> = Arc::from(weights.iter().map(|w| w / sum).collect::<Vec<_>>());
45        Ok(Self { models, graph_keys, weights })
46    }
47
48    /// Number of graphs.
49    #[must_use]
50    pub fn len(&self) -> usize {
51        self.models.len()
52    }
53
54    /// Empty check.
55    #[must_use]
56    pub fn is_empty(&self) -> bool {
57        self.models.is_empty()
58    }
59}