1use crate::sampling::lut::LutF64;
26use polydat::numeric::special::{inv_regularized_beta, inv_regularized_gamma_p, probit};
27
28pub const DEFAULT_RESOLUTION: usize = 1000;
30
31#[polydat::polydat_node(category = Conversions)]
48fn unit_interval(input: u64) -> f64 {
49 input as f64 / u64::MAX as f64
50}
51
52#[polydat::polydat_node(category = Conversions)]
70fn clamp_f64(
71 input: f64,
72 #[poly_default(f64::MIN)] min: Const<f64>,
73 #[poly_default(f64::MAX)] max: Const<f64>,
74) -> f64 {
75 input.clamp(*min, *max)
76}
77
78pub fn dist_normal_lut(mean: f64, stddev: f64, resolution: usize) -> LutF64 {
84 LutF64::from_fn(|p| mean + stddev * probit(p), resolution)
85}
86
87pub fn dist_exponential_lut(rate: f64, resolution: usize) -> LutF64 {
89 LutF64::from_fn(|p| -(1.0 - p).ln() / rate, resolution)
90}
91
92pub fn dist_uniform_lut(min: f64, max: f64, resolution: usize) -> LutF64 {
94 LutF64::from_fn(|p| min + p * (max - min), resolution)
95}
96
97pub fn dist_pareto_lut(scale: f64, shape: f64, resolution: usize) -> LutF64 {
99 LutF64::from_fn(|p| scale / (1.0 - p).powf(1.0 / shape), resolution)
100}
101
102pub fn dist_lognormal_lut(mean: f64, stddev: f64, resolution: usize) -> LutF64 {
104 LutF64::from_fn(|p| (mean + stddev * probit(p)).exp(), resolution)
105}
106
107pub fn dist_weibull_lut(shape: f64, scale: f64, resolution: usize) -> LutF64 {
109 LutF64::from_fn(|p| scale * (-(1.0 - p).ln()).powf(1.0 / shape), resolution)
110}
111
112pub fn dist_cauchy_lut(location: f64, scale: f64, resolution: usize) -> LutF64 {
114 LutF64::from_fn(
115 |p| location + scale * (std::f64::consts::PI * (p - 0.5)).tan(),
116 resolution,
117 )
118}
119
120pub fn dist_laplace_lut(location: f64, scale: f64, resolution: usize) -> LutF64 {
122 LutF64::from_fn(
123 |p| {
124 if p <= 0.5 {
125 location + scale * (2.0 * p).ln()
126 } else {
127 location - scale * (2.0 * (1.0 - p)).ln()
128 }
129 },
130 resolution,
131 )
132}
133
134pub fn dist_beta_lut(alpha: f64, beta: f64, resolution: usize) -> LutF64 {
136 LutF64::from_fn(|p| inv_regularized_beta(p, alpha, beta), resolution)
137}
138
139pub fn dist_gamma_lut(shape: f64, scale: f64, resolution: usize) -> LutF64 {
141 LutF64::from_fn(|p| scale * inv_regularized_gamma_p(p, shape), resolution)
142}
143
144pub fn dist_zipf_lut(n: u64, exponent: f64, resolution: usize) -> LutF64 {
154 let mut cdf = polydat::derive_support::buffer_for(n.saturating_add(1), "dist_zipf");
159 let harmonic: f64 = (1..=n).map(|k| 1.0 / (k as f64).powf(exponent)).sum();
161 cdf.push(0.0);
162 let mut cumulative = 0.0;
163 for k in 1..=n {
164 cumulative += (1.0 / (k as f64).powf(exponent)) / harmonic;
165 cdf.push(cumulative);
166 }
167
168 LutF64::from_fn(
170 |p| {
171 let p = p.clamp(0.0, 1.0);
172 match cdf.binary_search_by(|v| v.partial_cmp(&p).unwrap()) {
173 Ok(idx) => idx as f64,
174 Err(idx) => (idx as f64).max(1.0).min(n as f64),
175 }
176 },
177 resolution,
178 )
179}
180
181pub fn dist_poisson_lut(lambda: f64, resolution: usize) -> LutF64 {
185 let upper = (lambda + 6.0 * lambda.sqrt() + 10.0).ceil() as u64;
189
190 let mut cdf = polydat::derive_support::buffer_for(upper.saturating_add(2), "dist_poisson");
192 cdf.push(0.0);
193 let mut cumulative = 0.0;
194 let mut pmf = (-lambda).exp(); for k in 0..=upper {
196 cumulative += pmf;
197 cdf.push(cumulative.min(1.0));
198 pmf *= lambda / (k + 1) as f64;
199 }
200
201 LutF64::from_fn(
202 |p| {
203 let p = p.clamp(0.0, 1.0);
204 match cdf.binary_search_by(|v| v.partial_cmp(&p).unwrap()) {
205 Ok(idx) => idx.saturating_sub(1) as f64,
206 Err(idx) => idx.saturating_sub(1) as f64,
207 }
208 },
209 resolution,
210 )
211}
212
213pub fn dist_binomial_lut(trials: u64, prob: f64, resolution: usize) -> LutF64 {
215 let mut cdf = polydat::derive_support::buffer_for(trials.saturating_add(2), "dist_binomial");
217 let n = trials as usize; cdf.push(0.0);
219 let mut cumulative = 0.0;
220 let mut pmf = match i32::try_from(n) {
226 Ok(e) => (1.0 - prob).powi(e),
227 Err(_) => (1.0 - prob).powf(n as f64),
228 };
229 for k in 0..=n {
230 cumulative += pmf;
231 cdf.push(cumulative.min(1.0));
232 if k < n {
233 pmf *= prob / (1.0 - prob) * ((n - k) as f64) / ((k + 1) as f64);
234 }
235 }
236
237 LutF64::from_fn(
238 |p| {
239 let p = p.clamp(0.0, 1.0);
240 match cdf.binary_search_by(|v| v.partial_cmp(&p).unwrap()) {
241 Ok(idx) => idx.saturating_sub(1) as f64,
242 Err(idx) => idx.saturating_sub(1) as f64,
243 }
244 },
245 resolution,
246 )
247}
248
249pub fn dist_geometric_lut(prob: f64, resolution: usize) -> LutF64 {
253 let ln_q = (1.0 - prob).ln();
254 LutF64::from_fn(
255 |p| {
256 if p <= 0.0 {
257 return 1.0;
258 }
259 if p >= 1.0 {
260 return f64::INFINITY;
261 }
262 ((1.0 - p).ln() / ln_q).ceil().max(1.0)
263 },
264 resolution,
265 )
266}
267
268pub fn dist_empirical_lut(data: &[f64], resolution: usize) -> LutF64 {
278 assert!(!data.is_empty(), "data must not be empty");
279 let mut sorted = data.to_vec();
280 sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
281
282 LutF64::from_fn(
283 |p| {
284 let pos = p * (sorted.len() - 1) as f64;
285 let idx = pos as usize;
286 let idx = idx.min(sorted.len() - 2);
287 let frac = pos - idx as f64;
288 sorted[idx] * (1.0 - frac) + sorted[idx + 1] * frac
289 },
290 resolution,
291 )
292}
293
294pub fn dist_empirical_weighted_lut(values: &[f64], weights: &[f64], resolution: usize) -> LutF64 {
298 assert_eq!(values.len(), weights.len());
299 assert!(!values.is_empty());
300
301 let mut pairs: Vec<(f64, f64)> = values
303 .iter()
304 .copied()
305 .zip(weights.iter().copied())
306 .collect();
307 pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
308
309 let total: f64 = pairs.iter().map(|(_, w)| w).sum();
310 let mut cdf_points: Vec<(f64, f64)> = Vec::new(); let mut cumulative = 0.0;
312 for (val, weight) in &pairs {
313 cumulative += weight / total;
314 cdf_points.push((cumulative, *val));
315 }
316
317 LutF64::from_fn(
319 |p| match cdf_points.binary_search_by(|&(cp, _)| cp.partial_cmp(&p).unwrap()) {
320 Ok(idx) => cdf_points[idx].1,
321 Err(idx) => {
322 if idx >= cdf_points.len() {
323 cdf_points.last().unwrap().1
324 } else {
325 cdf_points[idx].1
326 }
327 }
328 },
329 resolution,
330 )
331}
332
333fn build_normal_lut(mean: f64, stddev: f64) -> LutF64 {
340 dist_normal_lut(mean, stddev, DEFAULT_RESOLUTION)
341}
342
343fn dist_normal_jit_constants(node: &DistNormal) -> Vec<u64> {
347 vec![node.lut.as_ptr() as u64, node.lut.len() as u64]
348}
349
350#[polydat::polydat_node(category = Distributions, jit_constants = dist_normal_jit_constants)]
356fn dist_normal(
357 input: f64,
358 mean: Const<f64>,
359 stddev: Const<f64>,
360 #[poly_const(build_normal_lut, from = (mean, stddev))] lut: &LutF64,
361) -> f64 {
362 let _ = mean;
363 let _ = stddev;
364 lut.sample(input)
365}
366
367fn icd_normal_jit_constants(node: &IcdNormal) -> Vec<u64> {
368 vec![node.lut.as_ptr() as u64, node.lut.len() as u64]
369}
370
371#[polydat::polydat_node(category = Distributions, jit_constants = icd_normal_jit_constants)]
375fn icd_normal(
376 input: f64,
377 mean: Const<f64>,
378 stddev: Const<f64>,
379 #[poly_const(build_normal_lut, from = (mean, stddev))] lut: &LutF64,
380) -> f64 {
381 let _ = mean;
382 let _ = stddev;
383 lut.sample(input)
384}
385
386fn build_exponential_lut(rate: f64) -> LutF64 {
387 dist_exponential_lut(rate, DEFAULT_RESOLUTION)
388}
389
390fn dist_exponential_jit_constants(node: &DistExponential) -> Vec<u64> {
391 vec![node.lut.as_ptr() as u64, node.lut.len() as u64]
392}
393
394#[polydat::polydat_node(category = Distributions, jit_constants = dist_exponential_jit_constants)]
396fn dist_exponential(
397 input: f64,
398 rate: Const<f64>,
399 #[poly_const(build_exponential_lut, from = rate)] lut: &LutF64,
400) -> f64 {
401 let _ = rate;
402 lut.sample(input)
403}
404
405fn icd_exponential_jit_constants(node: &IcdExponential) -> Vec<u64> {
406 vec![node.lut.as_ptr() as u64, node.lut.len() as u64]
407}
408
409#[polydat::polydat_node(category = Distributions, jit_constants = icd_exponential_jit_constants)]
411fn icd_exponential(
412 input: f64,
413 rate: Const<f64>,
414 #[poly_const(build_exponential_lut, from = rate)] lut: &LutF64,
415) -> f64 {
416 let _ = rate;
417 lut.sample(input)
418}
419
420fn build_uniform_lut(min: f64, max: f64) -> LutF64 {
421 dist_uniform_lut(min, max, DEFAULT_RESOLUTION)
422}
423
424fn dist_uniform_jit_constants(node: &DistUniform) -> Vec<u64> {
425 vec![node.lut.as_ptr() as u64, node.lut.len() as u64]
426}
427
428#[polydat::polydat_node(category = Distributions, jit_constants = dist_uniform_jit_constants)]
430fn dist_uniform(
431 input: f64,
432 min: Const<f64>,
433 max: Const<f64>,
434 #[poly_const(build_uniform_lut, from = (min, max))] lut: &LutF64,
435) -> f64 {
436 let _ = min;
437 let _ = max;
438 lut.sample(input)
439}
440
441fn build_pareto_lut(scale: f64, shape: f64) -> LutF64 {
442 dist_pareto_lut(scale, shape, DEFAULT_RESOLUTION)
443}
444
445fn dist_pareto_jit_constants(node: &DistPareto) -> Vec<u64> {
446 vec![node.lut.as_ptr() as u64, node.lut.len() as u64]
447}
448
449#[polydat::polydat_node(category = Distributions, jit_constants = dist_pareto_jit_constants)]
451fn dist_pareto(
452 input: f64,
453 scale: Const<f64>,
454 shape: Const<f64>,
455 #[poly_const(build_pareto_lut, from = (scale, shape))] lut: &LutF64,
456) -> f64 {
457 let _ = scale;
458 let _ = shape;
459 lut.sample(input)
460}
461
462fn build_zipf_lut(n: u64, exponent: f64) -> LutF64 {
463 dist_zipf_lut(n, exponent, DEFAULT_RESOLUTION)
464}
465
466fn dist_zipf_jit_constants(node: &DistZipf) -> Vec<u64> {
467 vec![node.lut.as_ptr() as u64, node.lut.len() as u64]
468}
469
470#[polydat::polydat_node(category = Distributions, jit_constants = dist_zipf_jit_constants)]
472fn dist_zipf(
473 input: f64,
474 n: Const<u64>,
475 exponent: Const<f64>,
476 #[poly_const(build_zipf_lut, from = (n, exponent))] lut: &LutF64,
477) -> f64 {
478 let _ = n;
479 let _ = exponent;
480 lut.sample(input)
481}
482
483#[cfg(test)]
492mod tests {
493 use super::*;
494 use polydat::ast::{PolydatNode, Value};
495
496 #[test]
497 fn unit_interval_range() {
498 let node = UnitInterval::new();
499 let mut out = [Value::None];
500 node.eval(&[Value::U64(0)], &mut out);
501 assert_eq!(out[0].as_f64(), 0.0);
502 node.eval(&[Value::U64(u64::MAX)], &mut out);
503 assert!((0.999..=1.0).contains(&out[0].as_f64()));
504 }
505
506 #[test]
507 fn normal_symmetry() {
508 let lut = dist_normal_lut(0.0, 1.0, 1000);
509 assert!(lut.sample(0.5).abs() < 0.01);
510 assert!((lut.sample(0.25) + lut.sample(0.75)).abs() < 0.01);
511 }
512
513 #[test]
514 fn normal_mean_stddev() {
515 let lut = dist_normal_lut(100.0, 10.0, 1000);
516 assert!((lut.sample(0.5) - 100.0).abs() < 0.5);
517 }
518
519 #[test]
520 fn exponential_median() {
521 let lut = dist_exponential_lut(1.0, 1000);
522 assert!((lut.sample(0.5) - 0.693).abs() < 0.01);
523 }
524
525 #[test]
526 fn exponential_positive() {
527 let lut = dist_exponential_lut(1.0, 1000);
528 for i in 1..1000 {
529 assert!(lut.sample(i as f64 / 1000.0) >= 0.0);
530 }
531 }
532
533 #[test]
534 fn uniform_linear() {
535 let lut = dist_uniform_lut(10.0, 20.0, 1000);
536 assert!((lut.sample(0.0) - 10.0).abs() < 0.1);
537 assert!((lut.sample(0.5) - 15.0).abs() < 0.1);
538 assert!((lut.sample(0.999) - 20.0).abs() < 0.1);
539 }
540
541 #[test]
542 fn pareto_heavy_tail() {
543 let lut = dist_pareto_lut(1.0, 1.0, 1000);
544 assert!((lut.sample(0.5) - 2.0).abs() < 0.1);
545 assert!(lut.sample(0.99) > 50.0);
546 }
547
548 #[test]
549 fn cauchy_symmetric() {
550 let lut = dist_cauchy_lut(0.0, 1.0, 1000);
551 assert!(lut.sample(0.5).abs() < 0.1);
552 assert!((lut.sample(0.25) + lut.sample(0.75)).abs() < 0.1);
553 }
554
555 #[test]
556 fn laplace_symmetric() {
557 let lut = dist_laplace_lut(5.0, 2.0, 1000);
558 assert!((lut.sample(0.5) - 5.0).abs() < 0.1);
559 }
560
561 #[test]
562 fn beta_bounded_01() {
563 let lut = dist_beta_lut(2.0, 5.0, 1000);
564 for i in 0..=1000 {
565 let v = lut.sample(i as f64 / 1000.0);
566 assert!((0.0..=1.0).contains(&v), "beta out of [0,1]: {v}");
567 }
568 }
569
570 #[test]
571 fn beta_symmetric_at_half() {
572 let lut = dist_beta_lut(2.0, 2.0, 1000);
574 assert!(
575 (lut.sample(0.5) - 0.5).abs() < 0.1,
576 "beta(2,2) median={}, expected ~0.5",
577 lut.sample(0.5)
578 );
579 }
580
581 #[test]
582 fn gamma_positive() {
583 let lut = dist_gamma_lut(2.0, 1.0, 1000);
584 for i in 1..1000 {
585 assert!(lut.sample(i as f64 / 1000.0) > 0.0);
586 }
587 }
588
589 #[test]
590 fn gamma_mean() {
591 let lut = dist_gamma_lut(3.0, 2.0, 1000);
593 assert!((lut.sample(0.5) - 5.0).abs() < 1.5); }
595
596 #[test]
597 fn weibull_positive() {
598 let lut = dist_weibull_lut(2.0, 1.0, 1000);
599 for i in 1..1000 {
600 assert!(lut.sample(i as f64 / 1000.0) >= 0.0);
601 }
602 }
603
604 #[test]
605 fn zipf_range() {
606 let lut = dist_zipf_lut(100, 1.0, 1000);
607 for i in 1..1000 {
608 let v = lut.sample(i as f64 / 1000.0);
609 assert!((1.0..=100.0).contains(&v), "zipf out of [1,100]: {v}");
610 }
611 }
612
613 #[test]
614 fn zipf_skewed() {
615 let lut = dist_zipf_lut(100, 1.0, 1000);
617 let low_quantile = lut.sample(0.5);
618 assert!(
619 low_quantile < 20.0,
620 "median of Zipf(100,1) should be low, got {low_quantile}"
621 );
622 }
623
624 #[test]
625 fn poisson_mean() {
626 let lut = dist_poisson_lut(5.0, 1000);
628 let median = lut.sample(0.5);
629 assert!(
630 (median - 5.0).abs() < 1.0,
631 "poisson median={median}, expected ~5"
632 );
633 }
634
635 #[test]
636 fn poisson_nonnegative() {
637 let lut = dist_poisson_lut(3.0, 1000);
638 for i in 0..=1000 {
639 assert!(lut.sample(i as f64 / 1000.0) >= 0.0);
640 }
641 }
642
643 #[test]
644 fn binomial_range() {
645 let lut = dist_binomial_lut(20, 0.5, 1000);
646 for i in 0..=1000 {
647 let v = lut.sample(i as f64 / 1000.0);
648 assert!((0.0..=20.0).contains(&v), "binomial out of [0,20]: {v}");
649 }
650 }
651
652 #[test]
653 fn binomial_mean() {
654 let lut = dist_binomial_lut(20, 0.5, 1000);
656 let median = lut.sample(0.5);
657 assert!(
658 (median - 10.0).abs() < 1.5,
659 "binomial median={median}, expected ~10"
660 );
661 }
662
663 #[test]
664 fn geometric_starts_at_one() {
665 let lut = dist_geometric_lut(0.5, 1000);
666 assert!(lut.sample(0.001) >= 1.0);
667 }
668
669 #[test]
670 fn geometric_mean() {
671 let lut = dist_geometric_lut(0.5, 1000);
673 let median = lut.sample(0.5);
674 assert!(
675 (median - 1.0).abs() < 1.0,
676 "geometric median={median}, expected ~1-2"
677 );
678 }
679
680 #[test]
681 fn dist_normal_node_eval() {
682 let node = DistNormal::new(0.0, 1.0);
683 let mut out = [Value::None];
684 node.eval(&[Value::F64(0.5)], &mut out);
685 assert!(out[0].as_f64().abs() < 0.01);
686 }
687
688 #[test]
689 fn full_pipeline_hash_normalize_sample() {
690 use xxhash_rust::xxh3::xxh3_64;
691
692 let lut = dist_normal_lut(72.0, 5.0, 1000);
693 let mut values = Vec::new();
694 for i in 0..10_000u64 {
695 let hashed = xxh3_64(&i.to_le_bytes());
696 let u = hashed as f64 / u64::MAX as f64;
697 values.push(lut.sample(u));
698 }
699 let mean = values.iter().sum::<f64>() / values.len() as f64;
700 let variance = values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / values.len() as f64;
701 let stddev = variance.sqrt();
702 assert!((mean - 72.0).abs() < 0.5, "mean={mean}");
703 assert!((stddev - 5.0).abs() < 0.5, "stddev={stddev}");
704 }
705}