Skip to main content

antecedent_kernels/
lib.rs

1//! Low-level borrowed views and numerical kernels for causal-library.
2//!
3//! Scalar kernels are the correctness reference. Portable-optimized and
4//! architecture-specific paths must pass the same differential tests
5//!
6//! SPDX-License-Identifier: MIT OR Apache-2.0
7
8#![deny(missing_docs)]
9// Unsafe is confined to validated unchecked view indexing in hot loops.
10#![allow(unsafe_code)]
11#![cfg_attr(test, allow(clippy::cast_precision_loss))]
12
13pub mod dispatch;
14pub mod parcorr;
15pub mod portable;
16pub mod posterior_reduce;
17pub mod rng;
18pub mod scalar;
19pub mod special;
20pub mod view;
21
22pub use dispatch::{
23    KernelImpl, accumulate_contingency, accumulate_contingency_rows, arch_simd_available, copy_vec,
24    gather, masked_covariance, masked_mean, masked_sum, masked_variance, pairwise_l1_fill,
25    partial_correlation, select_impl, standardize_inplace, weighted_dot, weighted_mean,
26    weighted_sum,
27};
28pub use parcorr::{
29    ParCorrMode, ParCorrQuery, ParCorrWorkspace, partial_correlation_batch, pearson,
30};
31pub use posterior_reduce::{PosteriorReduceOp, reduce_posterior_draws};
32pub use rng::{
33    categorical_from_u, fill_standard_normal, sample_categorical, shuffle, standard_normal,
34    standard_normal_pair, unbiased_index,
35};
36pub use scalar::sanitize_weight;
37pub use special::{erf, erfc, norm_cdf, norm_inv, norm_pdf};
38pub use view::{BitMaskView, F64MatrixView, F64VectorView, ViewError};
39
40#[cfg(test)]
41mod tests {
42    use antecedent_core::{KernelPolicy, ToleranceClass};
43
44    use super::*;
45
46    fn sample_data() -> Vec<f64> {
47        (0..128).map(|i| f64::from(i) * 0.5 - 10.0).collect()
48    }
49
50    #[test]
51    fn scalar_and_portable_agree_on_reductions() {
52        let data = sample_data();
53        let x = F64VectorView::contiguous(&data);
54        let mask_bytes = vec![0b1010_1011u8; data.len().div_ceil(8)];
55        let mask = BitMaskView::new(&mask_bytes, data.len()).unwrap();
56
57        let s_sum = scalar::masked_sum(x, Some(mask));
58        let p_sum = portable::masked_sum(x, Some(mask));
59        assert!(ToleranceClass::StableFloat.close(s_sum, p_sum));
60
61        let s_mean = scalar::masked_mean(x, Some(mask)).unwrap();
62        let p_mean = portable::masked_mean(x, Some(mask)).unwrap();
63        assert!(ToleranceClass::StableFloat.close(s_mean, p_mean));
64
65        let s_var = scalar::masked_variance(x, Some(mask)).unwrap();
66        let p_var = portable::masked_variance(x, Some(mask)).unwrap();
67        assert!(ToleranceClass::StableFloat.close(s_var, p_var));
68    }
69
70    #[test]
71    fn scalar_and_portable_agree_on_covariance() {
72        let x_data = sample_data();
73        let y_data: Vec<f64> = x_data.iter().map(|v| v * 0.3 + 1.0).collect();
74        let x = F64VectorView::contiguous(&x_data);
75        let y = F64VectorView::contiguous(&y_data);
76        let mask_bytes = vec![0b1111_0000u8; x_data.len().div_ceil(8)];
77        let mask = BitMaskView::new(&mask_bytes, x_data.len()).unwrap();
78
79        let s = scalar::masked_covariance(x, y, Some(mask)).unwrap();
80        let p = portable::masked_covariance(x, y, Some(mask)).unwrap();
81        assert!(ToleranceClass::StableFloat.close(s, p));
82
83        let s_full = scalar::masked_covariance(x, y, None).unwrap();
84        let p_full = portable::masked_covariance(x, y, None).unwrap();
85        assert!(ToleranceClass::StableFloat.close(s_full, p_full));
86    }
87
88    #[test]
89    fn scalar_and_portable_agree_on_standardize() {
90        let mut a = sample_data();
91        let mut b = a.clone();
92        let (ms, ss) = scalar::standardize_inplace(&mut a, 1e-12);
93        let (mp, sp) = portable::standardize_inplace(&mut b, 1e-12);
94        assert!(ToleranceClass::StableFloat.close(ms, mp));
95        assert!(ToleranceClass::StableFloat.close(ss, sp));
96        for (sa, sb) in a.iter().zip(b.iter()) {
97            assert!(ToleranceClass::StableFloat.close(*sa, *sb));
98        }
99    }
100
101    #[test]
102    fn scalar_and_portable_agree_on_pairwise_l1() {
103        let x = sample_data();
104        let n = x.len();
105        let mut out_s = vec![0.0; n * n];
106        let mut out_p = vec![0.0; n * n];
107        scalar::pairwise_l1_fill(&x, &mut out_s);
108        portable::pairwise_l1_fill(&x, &mut out_p);
109        assert_eq!(out_s, out_p);
110    }
111
112    #[test]
113    fn scalar_and_portable_agree_on_contingency() {
114        let x_codes: Vec<u32> = (0u32..64).map(|i| i % 4).collect();
115        let y_codes: Vec<u32> = (0u32..64).map(|i| i % 3).collect();
116        let mut out_s = vec![0.0; 4 * 3];
117        let mut out_p = vec![0.0; 4 * 3];
118        scalar::accumulate_contingency(&x_codes, &y_codes, &mut out_s, 3);
119        portable::accumulate_contingency(&x_codes, &y_codes, &mut out_p, 3);
120        assert_eq!(out_s, out_p);
121    }
122
123    #[test]
124    fn scalar_and_portable_agree_on_weighted() {
125        let x = sample_data();
126        let y: Vec<f64> = x.iter().map(|v| v + 1.0).collect();
127        let w: Vec<f64> =
128            (0..x.len()).map(|i| if i % 5 == 0 { 0.0 } else { 1.0 + (i as f64) * 0.01 }).collect();
129        let s_sum = scalar::weighted_sum(&x, &w);
130        let p_sum = portable::weighted_sum(&x, &w);
131        assert!(ToleranceClass::StableFloat.close(s_sum, p_sum));
132        let s_mean = scalar::weighted_mean(&x, &w).unwrap();
133        let p_mean = portable::weighted_mean(&x, &w).unwrap();
134        assert!(ToleranceClass::StableFloat.close(s_mean, p_mean));
135        let s_dot = scalar::weighted_dot(&x, &y, &w);
136        let p_dot = portable::weighted_dot(&x, &y, &w);
137        assert!(ToleranceClass::StableFloat.close(s_dot, p_dot));
138    }
139
140    #[test]
141    fn sanitize_weight_drops_nonfinite_and_negative() {
142        assert!(ToleranceClass::Exact.close(sanitize_weight(1.5), 1.5));
143        assert!(ToleranceClass::Exact.close(sanitize_weight(0.0), 0.0));
144        assert!(ToleranceClass::Exact.close(sanitize_weight(-2.0), 0.0));
145        assert!(ToleranceClass::Exact.close(sanitize_weight(f64::NAN), 0.0));
146        assert!(ToleranceClass::Exact.close(sanitize_weight(f64::INFINITY), 0.0));
147        assert!(ToleranceClass::Exact.close(sanitize_weight(f64::NEG_INFINITY), 0.0));
148        let x = [1.0, 2.0, 3.0];
149        let w = [1.0, f64::NAN, -1.0];
150        assert!(ToleranceClass::StableFloat.close(scalar::weighted_mean(&x, &w).unwrap(), 1.0));
151        assert!(ToleranceClass::StableFloat.close(portable::weighted_mean(&x, &w).unwrap(), 1.0));
152    }
153
154    #[test]
155    fn gather_differential_contiguous_and_strided() {
156        let data = sample_data();
157        let contiguous = F64VectorView::contiguous(&data);
158        let strided = F64VectorView::strided(&data, 32, 2).unwrap();
159        let indices: Vec<usize> = (0..32).collect();
160        let mut out_s = vec![0.0; 32];
161        let mut out_p = vec![0.0; 32];
162        scalar::gather(contiguous, &indices, &mut out_s);
163        portable::gather(contiguous, &indices, &mut out_p);
164        assert_eq!(out_s, out_p);
165
166        let mut out_strided = vec![0.0; 16];
167        let idx: Vec<usize> = (0..16).collect();
168        scalar::gather(strided, &idx, &mut out_strided);
169        let mut out_strided_p = vec![0.0; 16];
170        portable::gather(strided, &idx, &mut out_strided_p);
171        assert_eq!(out_strided, out_strided_p);
172    }
173
174    #[test]
175    fn dispatch_respects_force_scalar() {
176        let policy = KernelPolicy::scalar_only();
177        assert_eq!(select_impl(&policy), KernelImpl::Scalar);
178        let default = KernelPolicy::default_policy();
179        // Without simd-runtime, default selects portable even when allow_arch_simd is true.
180        assert_eq!(select_impl(&default), KernelImpl::PortableOptimized);
181        assert!(!arch_simd_available());
182    }
183
184    #[test]
185    fn dispatch_honors_allow_arch_simd_false() {
186        let mut policy = KernelPolicy::default_policy();
187        policy.allow_arch_simd = false;
188        assert_eq!(select_impl(&policy), KernelImpl::PortableOptimized);
189        policy.allow_portable_optimized = false;
190        assert_eq!(select_impl(&policy), KernelImpl::Scalar);
191    }
192
193    #[test]
194    fn dispatch_entries_match_scalar_under_force() {
195        let data = sample_data();
196        let x = F64VectorView::contiguous(&data);
197        let policy = KernelPolicy::scalar_only();
198        let mut out = vec![0.0; 8];
199        let idx = [0usize, 1, 2, 3, 4, 5, 6, 7];
200        gather(&policy, x, &idx, &mut out);
201        let mut expected = vec![0.0; 8];
202        scalar::gather(x, &idx, &mut expected);
203        assert_eq!(out, expected);
204
205        let y: Vec<f64> = data.iter().map(|v| v * 2.0).collect();
206        let yv = F64VectorView::contiguous(&y);
207        let cov = masked_covariance(&policy, x, yv, None).unwrap();
208        let cov_s = scalar::masked_covariance(x, yv, None).unwrap();
209        assert!(ToleranceClass::StableFloat.close(cov, cov_s));
210    }
211
212    /// : scalar ≡ portable under random lengths, strides/tails, masks, NaNs.
213    #[test]
214    fn property_scalar_portable_random_strides_masks_nans() {
215        let mut rng = antecedent_core::CausalRng::from_seed(28_02);
216        for trial in 0..80 {
217            let len = 1 + usize::try_from(rng.next_u64() % 48).unwrap_or(0); // 1..=48
218            let stride = 1 + usize::try_from(rng.next_u64() % 4).unwrap_or(0); // 1..=4
219            let tail = usize::try_from(rng.next_u64() % 5).unwrap_or(0); // unused padding
220            let need = (len - 1) * stride + 1 + tail;
221            let mut x_buf = vec![0.0f64; need];
222            let mut y_buf = vec![0.0f64; need];
223            for i in 0..need {
224                // Mix finite values with occasional NaNs (inject into logical slots).
225                let u = rng.next_u64();
226                let base = ((u % 1000) as f64) * 0.01 - 5.0;
227                x_buf[i] = if u % 17 == 0 { f64::NAN } else { base };
228                y_buf[i] = if (u / 17) % 19 == 0 { f64::NAN } else { base * 0.3 + 1.0 };
229            }
230            let x = F64VectorView::strided(&x_buf, len, stride).unwrap();
231            let y = F64VectorView::strided(&y_buf, len, stride).unwrap();
232
233            let mut bits = vec![0u8; len.div_ceil(8)];
234            let use_mask = rng.next_u64() % 3 != 0;
235            if use_mask {
236                for i in 0..len {
237                    if rng.next_u64() & 1 == 1 {
238                        bits[i / 8] |= 1 << (i % 8);
239                    }
240                }
241            }
242            // Ensure at least one valid bit when masked so mean/var/cov have a chance.
243            if use_mask && bits.iter().all(|&b| b == 0) {
244                bits[0] |= 1;
245            }
246            let mask = if use_mask { Some(BitMaskView::new(&bits, len).unwrap()) } else { None };
247
248            let s_sum = scalar::masked_sum(x, mask);
249            let p_sum = portable::masked_sum(x, mask);
250            assert!(
251                floats_agree(s_sum, p_sum),
252                "trial={trial} sum scalar={s_sum} portable={p_sum} len={len} stride={stride}"
253            );
254
255            let s_mean = scalar::masked_mean(x, mask);
256            let p_mean = portable::masked_mean(x, mask);
257            assert!(
258                options_agree(s_mean, p_mean),
259                "trial={trial} mean scalar={s_mean:?} portable={p_mean:?}"
260            );
261
262            let s_var = scalar::masked_variance(x, mask);
263            let p_var = portable::masked_variance(x, mask);
264            assert!(
265                options_agree(s_var, p_var),
266                "trial={trial} var scalar={s_var:?} portable={p_var:?}"
267            );
268
269            let s_cov = scalar::masked_covariance(x, y, mask);
270            let p_cov = portable::masked_covariance(x, y, mask);
271            assert!(
272                options_agree(s_cov, p_cov),
273                "trial={trial} cov scalar={s_cov:?} portable={p_cov:?}"
274            );
275        }
276    }
277
278    fn floats_agree(a: f64, b: f64) -> bool {
279        (a.is_nan() && b.is_nan()) || ToleranceClass::StableFloat.close(a, b)
280    }
281
282    fn options_agree(a: Option<f64>, b: Option<f64>) -> bool {
283        match (a, b) {
284            (None, None) => true,
285            (Some(x), Some(y)) => floats_agree(x, y),
286            _ => false,
287        }
288    }
289
290    #[test]
291    fn gather_hot_path_reuses_output_buffer() {
292        let n = 8_000usize;
293        let data: Vec<f64> = (0..n).map(|i| i as f64).collect();
294        let src = F64VectorView::contiguous(&data);
295        let indices: Vec<usize> = (0..n).step_by(8).collect();
296        let mut out = vec![0.0; indices.len()];
297        let policy = KernelPolicy::default_policy();
298        let ptr = out.as_ptr();
299        let cap = out.capacity();
300        for _ in 0..200 {
301            gather(&policy, src, &indices, &mut out);
302            assert_eq!(out.as_ptr(), ptr);
303            assert_eq!(out.capacity(), cap);
304        }
305    }
306
307    #[test]
308    fn pairwise_and_contingency_reuse_output_buffers() {
309        let policy = KernelPolicy::default_policy();
310        let x: Vec<f64> = (0..64).map(f64::from).collect();
311        let mut out = vec![0.0; 64 * 64];
312        let ptr = out.as_ptr();
313        let cap = out.capacity();
314        for _ in 0..50 {
315            pairwise_l1_fill(&policy, &x, &mut out);
316            assert_eq!(out.as_ptr(), ptr);
317            assert_eq!(out.capacity(), cap);
318        }
319
320        let xc: Vec<u32> = (0u32..256).map(|i| i % 5).collect();
321        let yc: Vec<u32> = (0u32..256).map(|i| i % 4).collect();
322        let mut table = vec![0.0; 5 * 4];
323        let tptr = table.as_ptr();
324        let tcap = table.capacity();
325        for _ in 0..50 {
326            table.fill(0.0);
327            accumulate_contingency(&policy, &xc, &yc, &mut table, 4);
328            assert_eq!(table.as_ptr(), tptr);
329            assert_eq!(table.capacity(), tcap);
330        }
331    }
332}