glmm/
spec.rs

1//! `glmm`'s owned model-spec input vocabulary. Minimal copy of the subset of
2//! MCPower's `engine_contract` cluster types that the fit kernels actually read
3//! MCPower converts its `ClusterSpec` into this via a
4//! conversion fn on its side — `glmm` never sees `engine_contract`.
5
6/// Predictor column index (mirrors `engine_contract::ColumnId`'s underlying type).
7pub type ColumnId = u32;
8
9/// Cluster sizing regime. Mirror of `engine_contract::ClusterSizing`.
10#[derive(Debug, Clone, PartialEq)]
11pub enum Sizing {
12    FixedClusters { n_clusters: u32 },
13    FixedSize { cluster_size: u32 },
14}
15
16impl Sizing {
17    /// Smallest legal increment in total N (the grid atom).
18    pub fn atom(&self) -> usize {
19        match self {
20            Sizing::FixedClusters { n_clusters } => (*n_clusters).max(1) as usize,
21            Sizing::FixedSize { cluster_size } => (*cluster_size).max(1) as usize,
22        }
23    }
24    pub fn n_clusters_at(&self, n: usize) -> usize {
25        match self {
26            Sizing::FixedClusters { n_clusters } => (*n_clusters).max(1) as usize,
27            Sizing::FixedSize { cluster_size } => n / (*cluster_size).max(1) as usize,
28        }
29    }
30    pub fn cluster_of_row(&self, i: usize) -> usize {
31        match self {
32            Sizing::FixedClusters { n_clusters } => i % (*n_clusters).max(1) as usize,
33            Sizing::FixedSize { cluster_size } => i / (*cluster_size).max(1) as usize,
34        }
35    }
36}
37
38/// One random slope. Mirror of `engine_contract::SlopeTerm`.
39#[derive(Debug, Clone, PartialEq)]
40pub struct SlopeTerm {
41    pub column: ColumnId,
42    pub variance: f64,
43    pub corr_with_intercept: f64,
44    pub corr_with: Vec<f64>,
45}
46
47/// An extra grouping factor. Mirror of `engine_contract::GroupingSpec`.
48#[derive(Debug, Clone, PartialEq)]
49pub struct Grouping {
50    pub relation: GroupingRelation,
51    pub tau_squared: f64,
52    pub slopes: Vec<SlopeTerm>,
53}
54
55impl Grouping {
56    /// The `q_g×q_g` RE *correlation* matrix `R_g` over `[intercept, slope_0, …]`
57    /// (`q_g = 1 + slopes.len()`). Identical recipe to the primary's
58    /// [`ModelSpec::re_correlation_matrix`]; `cluster_theta_truth` reads it to form
59    /// each extra factor's relative-covariance factor `Λ_g = chol(D_g)`.
60    pub fn re_correlation_matrix(&self) -> (usize, Vec<f64>) {
61        re_correlation_from_slopes(&self.slopes)
62    }
63}
64
65/// Build the `q×q` RE correlation matrix over `[intercept, slope_0, …]` from a
66/// slope list (`q = 1 + slopes.len()`), row-major: diagonal 1;
67/// `R[0][k+1]=R[k+1][0]=slopes[k].corr_with_intercept`;
68/// `R[i+1][k+1]=R[k+1][i+1]=slopes[k].corr_with[i]` for `i < k`. Shared by the
69/// primary factor ([`ModelSpec::re_correlation_matrix`]) and each extra
70/// [`Grouping::re_correlation_matrix`] — one source, no drift.
71fn re_correlation_from_slopes(slopes: &[SlopeTerm]) -> (usize, Vec<f64>) {
72    let q = 1 + slopes.len();
73    let mut r = vec![0.0; q * q];
74    for d in 0..q {
75        r[d * q + d] = 1.0;
76    }
77    for (k, s) in slopes.iter().enumerate() {
78        r[k + 1] = s.corr_with_intercept; // R[0][k+1]
79        r[(k + 1) * q] = s.corr_with_intercept; // R[k+1][0]
80        for (i, &cik) in s.corr_with.iter().enumerate() {
81            r[(i + 1) * q + (k + 1)] = cik;
82            r[(k + 1) * q + (i + 1)] = cik;
83        }
84    }
85    (q, r)
86}
87
88#[derive(Debug, Clone, PartialEq)]
89pub enum GroupingRelation {
90    Crossed { n_clusters: u32 },
91    NestedWithin { n_per_parent: u32 },
92}
93
94/// Estimator class. Mirror of `engine_contract::EstimatorSpec`.
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96pub enum Estimator {
97    Ols,
98    Glm,
99    Mle,
100}
101
102/// GLMM fixed-effect Wald-SE denominator. Mirror of `engine_contract::WaldSe`.
103#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
104pub enum WaldSe {
105    #[default]
106    Hessian,
107    Rx,
108}
109
110/// The kernels' model-spec input. Exactly the fields the fit side reads.
111#[derive(Debug, Clone, PartialEq)]
112pub struct ModelSpec {
113    pub sizing: Sizing,
114    pub tau_squared: f64,
115    pub slopes: Vec<SlopeTerm>,
116    pub extra_groupings: Vec<Grouping>,
117    pub estimator: Estimator,
118    pub wald_se: WaldSe,
119}
120
121impl ModelSpec {
122    /// The `q_p×q_p` RE *correlation* matrix `R` over `[intercept, slope_0, …]`,
123    /// row-major (`q_p = 1 + slopes.len()`). Diagonal 1; `R[0][k+1]=R[k+1][0]=
124    /// slopes[k].corr_with_intercept`; `R[i+1][k+1]=R[k+1][i+1]=slopes[k].corr_with[i]`
125    /// for `i < k`. Multiply by `diag(τ)` on both sides for `D`. Verbatim mirror of
126    /// `engine_contract::ClusterSpec::re_correlation_matrix` — `cluster_theta_truth`
127    /// reads it after the caller converts its `ClusterSpec` to `&ModelSpec`.
128    pub fn re_correlation_matrix(&self) -> (usize, Vec<f64>) {
129        re_correlation_from_slopes(&self.slopes)
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    #[test]
138    fn model_spec_constructs_and_reports_q() {
139        let spec = ModelSpec {
140            sizing: Sizing::FixedClusters { n_clusters: 30 },
141            tau_squared: 0.5,
142            slopes: vec![SlopeTerm {
143                column: 1,
144                variance: 0.2,
145                corr_with_intercept: 0.1,
146                corr_with: vec![],
147            }],
148            extra_groupings: vec![Grouping {
149                relation: GroupingRelation::Crossed { n_clusters: 12 },
150                tau_squared: 0.3,
151                slopes: vec![],
152            }],
153            estimator: Estimator::Mle,
154            wald_se: WaldSe::Hessian,
155        };
156        // primary RE width q_p = 1 (intercept) + #slopes
157        assert_eq!(1 + spec.slopes.len(), 2);
158        assert_eq!(spec.sizing.atom(), 30);
159    }
160}
glmm/spec.rs

glmm/
spec.rs
