polydat_nodes/sampling/
histribution.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Histribution: inline discrete histogram distribution.
5//!
6//! Parse a frequency spec string into an alias table at init time.
7//! The name is a portmanteau of "histogram" + "distribution."
8//!
9//! Two formats:
10//! - Implicit labels: `"50 25 13 12"` → outcomes 0,1,2,3 with those weights
11//! - Explicit labels: `"234:50 33:25 17:13 3:12"` → outcomes 234,33,17,3
12
13use crate::sampling::alias::AliasTableU64;
14use polydat::derive_support::PolydatSetup;
15
16/// Parsed histribution: labels + weighted alias table. The
17/// `#[poly_const]` setup returns this single struct so the macro
18/// can hand the body a borrow of one cached field rather than a
19/// tuple. Labeled by index `i` → outcome `labels[i]`; sampling
20/// goes through `table`.
21pub struct ParsedHistribution {
22    /// The outcome per index.
23    pub labels: Vec<u64>,
24    /// The alias table sampled.
25    pub table: AliasTableU64,
26}
27
28impl PolydatSetup for ParsedHistribution {}
29
30/// Parse a histribution spec and build an alias table.
31///
32/// Returns `(labels, table)` where `labels[i]` is the outcome for
33/// alias table index `i`. Kept as a free function for tests that
34/// want the parsed form without constructing the node.
35pub fn parse_histribution(spec: &str) -> (Vec<u64>, AliasTableU64) {
36    let labeled = spec.contains(':');
37    let mut labels = Vec::new();
38    let mut weights = Vec::new();
39
40    for (i, elem) in spec.split([' ', ',', ';']).enumerate() {
41        let elem = elem.trim();
42        if elem.is_empty() {
43            continue;
44        }
45        if labeled {
46            let parts: Vec<&str> = elem.splitn(2, ':').collect();
47            assert_eq!(parts.len(), 2, "all elements must be labeled: {elem}");
48            labels.push(parts[0].parse::<u64>().expect("invalid label"));
49            weights.push(parts[1].parse::<f64>().expect("invalid weight"));
50        } else {
51            labels.push(i as u64);
52            weights.push(elem.parse::<f64>().expect("invalid weight"));
53        }
54    }
55
56    assert!(!weights.is_empty(), "histribution spec must not be empty");
57    let table = AliasTableU64::from_weights(&weights);
58    (labels, table)
59}
60
61/// `#[poly_const]` setup: parse the spec into a `ParsedHistribution`
62/// struct once at construction time.
63fn parse_histribution_setup(spec: &str) -> ParsedHistribution {
64    let (labels, table) = parse_histribution(spec);
65    ParsedHistribution { labels, table }
66}
67
68/// Sample from a histogram-spec distribution. The input should
69/// be hashed (uniform); the output is one of the labeled
70/// outcomes, selected by weighted alias sampling.
71#[polydat::polydat_node(category = Probability)]
72fn histribution(
73    input: u64,
74    spec: polydat::derive_support::Const<&str>,
75    #[poly_const(parse_histribution_setup, from = spec)] parsed: &ParsedHistribution,
76) -> u64 {
77    let idx = parsed.table.sample(input) as usize;
78    parsed.labels[idx]
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use polydat::ast::{PolydatNode, Value};
85    use xxhash_rust::xxh3::xxh3_64;
86
87    #[test]
88    fn parse_implicit_labels() {
89        let (labels, table) = parse_histribution("50 25 13 12");
90        assert_eq!(labels, vec![0, 1, 2, 3]);
91        assert_eq!(table.len(), 4);
92    }
93
94    #[test]
95    fn parse_explicit_labels() {
96        let (labels, table) = parse_histribution("234:50 33:25 17:13 3:12");
97        assert_eq!(labels, vec![234, 33, 17, 3]);
98        assert_eq!(table.len(), 4);
99    }
100
101    #[test]
102    fn parse_comma_separated() {
103        let (labels, _) = parse_histribution("10,20,30");
104        assert_eq!(labels, vec![0, 1, 2]);
105    }
106
107    #[test]
108    fn parse_semicolon_separated() {
109        let (labels, _) = parse_histribution("10;20;30");
110        assert_eq!(labels, vec![0, 1, 2]);
111    }
112
113    #[test]
114    fn histribution_samples_valid_labels() {
115        let node = Histribution::new("234:50 33:25 17:13 3:12".to_string());
116        let valid = [234u64, 33, 17, 3];
117        let mut out = [Value::None];
118        for i in 0..1000u64 {
119            let hashed = xxh3_64(&i.to_le_bytes());
120            node.eval(&[Value::U64(hashed)], &mut out);
121            assert!(
122                valid.contains(&out[0].as_u64()),
123                "unexpected outcome: {}",
124                out[0].as_u64()
125            );
126        }
127    }
128
129    #[test]
130    fn histribution_weighted() {
131        // Outcome 0 has weight 100, others have weight 1 each
132        let node = Histribution::new("100 1 1".to_string());
133        let mut counts = [0u64; 3];
134        for i in 0..10_000u64 {
135            let hashed = xxh3_64(&i.to_le_bytes());
136            let mut out = [Value::None];
137            node.eval(&[Value::U64(hashed)], &mut out);
138            counts[out[0].as_u64() as usize] += 1;
139        }
140        let ratio = counts[0] as f64 / 10_000.0;
141        assert!(ratio > 0.90, "outcome 0 should dominate, got {ratio}");
142    }
143
144    // `histribution` has no named JitOp; it lowers through its
145    // slot kit. The typed-eval path above covers correctness.
146
147    #[test]
148    fn histribution_deterministic() {
149        let node = Histribution::new("50 25 13 12".to_string());
150        let mut out1 = [Value::None];
151        let mut out2 = [Value::None];
152        let hashed = xxh3_64(&42u64.to_le_bytes());
153        node.eval(&[Value::U64(hashed)], &mut out1);
154        node.eval(&[Value::U64(hashed)], &mut out2);
155        assert_eq!(out1[0].as_u64(), out2[0].as_u64());
156    }
157}
polydat_nodes/sampling/histribution.rs

polydat_nodes/sampling/
histribution.rs
