1#![deny(missing_docs)]
9#![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 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 #[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); let stride = 1 + usize::try_from(rng.next_u64() % 4).unwrap_or(0); let tail = usize::try_from(rng.next_u64() % 5).unwrap_or(0); 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 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 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}