parzen/
distribution.rs

1// Copyright 2026 Thomas Santerre and Moderately AI Inc.
2//
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Validated parameter distributions.
6
7use crate::{ParamValue, ParzenError};
8
9/// A categorical, floating-point, or integer search distribution.
10#[derive(Debug, Clone, PartialEq)]
11#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12pub enum Distribution {
13    /// Categorical indices in `0..num_choices`.
14    Categorical(CategoricalDistribution),
15    /// A bounded floating-point distribution.
16    Float(FloatDistribution),
17    /// A bounded integer distribution.
18    Int(IntDistribution),
19}
20
21/// A categorical distribution represented by choice indices.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23#[cfg_attr(feature = "serde", derive(serde::Serialize))]
24pub struct CategoricalDistribution {
25    num_choices: u32,
26}
27
28impl CategoricalDistribution {
29    /// Create a distribution with indices in `0..num_choices`.
30    pub fn new(num_choices: u32) -> Result<Self, ParzenError> {
31        if num_choices == 0 {
32            return Err(ParzenError::InvalidDistribution(
33                "categorical choice count must be positive".into(),
34            ));
35        }
36        Ok(Self { num_choices })
37    }
38
39    /// Number of categorical choices.
40    #[must_use]
41    pub const fn num_choices(self) -> u32 {
42        self.num_choices
43    }
44}
45
46/// Scale used by a floating-point distribution.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
49pub enum FloatScale {
50    /// Uniform linear coordinate space.
51    Linear,
52    /// Natural-log coordinate space.
53    Log,
54}
55
56/// A bounded floating-point distribution.
57#[derive(Debug, Clone, Copy, PartialEq)]
58#[cfg_attr(feature = "serde", derive(serde::Serialize))]
59pub struct FloatDistribution {
60    low: f64,
61    high: f64,
62    scale: FloatScale,
63    step: Option<f64>,
64}
65
66impl FloatDistribution {
67    /// Create a linear continuous distribution.
68    pub fn linear(low: f64, high: f64) -> Result<Self, ParzenError> {
69        Self::new(low, high, FloatScale::Linear)
70    }
71
72    /// Create a log-scaled continuous distribution.
73    pub fn log(low: f64, high: f64) -> Result<Self, ParzenError> {
74        Self::new(low, high, FloatScale::Log)
75    }
76
77    fn new(low: f64, high: f64, scale: FloatScale) -> Result<Self, ParzenError> {
78        if !low.is_finite() || !high.is_finite() || low >= high {
79            return Err(ParzenError::InvalidDistribution(
80                "float bounds must be finite and low < high".into(),
81            ));
82        }
83        if scale == FloatScale::Log && low <= 0.0 {
84            return Err(ParzenError::InvalidDistribution(
85                "log-float bounds must be positive".into(),
86            ));
87        }
88        let transformed_low = match scale {
89            FloatScale::Linear => low,
90            FloatScale::Log => low.ln(),
91        };
92        let transformed_high = match scale {
93            FloatScale::Linear => high,
94            FloatScale::Log => high.ln(),
95        };
96        if !transformed_low.is_finite()
97            || !transformed_high.is_finite()
98            || transformed_low >= transformed_high
99            || !(transformed_high - transformed_low).is_finite()
100        {
101            return Err(ParzenError::InvalidDistribution(
102                "float bounds must define a finite, non-empty transformed interval".into(),
103            ));
104        }
105        Ok(Self {
106            low,
107            high,
108            scale,
109            step: None,
110        })
111    }
112
113    /// Quantize samples to a positive step anchored at `low`.
114    pub fn with_step(mut self, step: f64) -> Result<Self, ParzenError> {
115        if self.scale == FloatScale::Log {
116            return Err(ParzenError::InvalidDistribution(
117                "log-float distributions cannot have a step".into(),
118            ));
119        }
120        if !step.is_finite() || step <= 0.0 {
121            return Err(ParzenError::InvalidDistribution(
122                "float step must be finite and positive".into(),
123            ));
124        }
125        let units = ((self.high - self.low) / step).floor();
126        const MAX_EXACT_INTEGER: f64 = 9_007_199_254_740_991.0;
127        if !units.is_finite() || units > MAX_EXACT_INTEGER {
128            return Err(ParzenError::InvalidDistribution(
129                "float step creates too many exactly representable grid points".into(),
130            ));
131        }
132        let highest = step.mul_add(units, self.low);
133        if !highest.is_finite() || highest < self.low || highest > self.high {
134            return Err(ParzenError::InvalidDistribution(
135                "float step produces invalid grid arithmetic".into(),
136            ));
137        }
138        if units >= 1.0 && step.mul_add(1.0, self.low) <= self.low {
139            return Err(ParzenError::InvalidDistribution(
140                "float step does not produce distinct grid values at this magnitude".into(),
141            ));
142        }
143        let adapted_low = self.low - step / 2.0;
144        let adapted_high = highest + step / 2.0;
145        if !adapted_low.is_finite()
146            || !adapted_high.is_finite()
147            || adapted_low >= adapted_high
148            || !(adapted_high - adapted_low).is_finite()
149        {
150            return Err(ParzenError::InvalidDistribution(
151                "float step must define a finite discrete kernel domain".into(),
152            ));
153        }
154        self.step = Some(step);
155        Ok(self)
156    }
157
158    /// Inclusive lower bound.
159    #[must_use]
160    pub const fn low(self) -> f64 {
161        self.low
162    }
163    /// Inclusive upper bound after clamping.
164    #[must_use]
165    pub const fn high(self) -> f64 {
166        self.high
167    }
168    /// Coordinate scale.
169    #[must_use]
170    pub const fn scale(self) -> FloatScale {
171        self.scale
172    }
173    /// Optional quantization step.
174    #[must_use]
175    pub const fn step(self) -> Option<f64> {
176        self.step
177    }
178
179    pub(crate) fn transform(self, value: f64) -> f64 {
180        match self.scale {
181            FloatScale::Linear => value,
182            FloatScale::Log => value.ln(),
183        }
184    }
185
186    pub(crate) fn untransform(self, value: f64) -> f64 {
187        let raw = match self.scale {
188            FloatScale::Linear => value,
189            FloatScale::Log => value.exp(),
190        };
191        self.quantize(raw)
192    }
193
194    pub(crate) fn quantize(self, value: f64) -> f64 {
195        let value = if let Some(step) = self.step {
196            let highest_unit = ((self.high - self.low) / step).floor();
197            let unit = ((value - self.low) / step).round().clamp(0.0, highest_unit);
198            step.mul_add(unit, self.low)
199        } else {
200            value
201        };
202        value.clamp(self.low, self.high)
203    }
204
205    pub(crate) fn max_step_index(self) -> Option<u64> {
206        self.step
207            .map(|step| ((self.high - self.low) / step).floor() as u64)
208    }
209
210    pub(crate) fn grid_value(self, index: u64) -> f64 {
211        self.step
212            .map_or(self.low, |step| step.mul_add(index as f64, self.low))
213    }
214}
215
216/// Scale used by an integer distribution.
217#[derive(Debug, Clone, Copy, PartialEq, Eq)]
218#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
219pub enum IntScale {
220    Linear,
221    Log,
222}
223
224/// A bounded integer distribution.
225#[derive(Debug, Clone, Copy, PartialEq, Eq)]
226#[cfg_attr(feature = "serde", derive(serde::Serialize))]
227pub struct IntDistribution {
228    low: i64,
229    high: i64,
230    scale: IntScale,
231    step: u64,
232}
233
234impl IntDistribution {
235    /// Create a linear integer distribution with unit step.
236    pub fn linear(low: i64, high: i64) -> Result<Self, ParzenError> {
237        Self::new(low, high, IntScale::Linear)
238    }
239    /// Create a positive logarithmic integer distribution with unit step.
240    pub fn log(low: i64, high: i64) -> Result<Self, ParzenError> {
241        Self::new(low, high, IntScale::Log)
242    }
243    fn new(low: i64, high: i64, scale: IntScale) -> Result<Self, ParzenError> {
244        if low > high {
245            return Err(ParzenError::InvalidDistribution(
246                "integer low must be <= high".into(),
247            ));
248        }
249        if scale == IntScale::Log && low <= 0 {
250            return Err(ParzenError::InvalidDistribution(
251                "log-integer bounds must be positive".into(),
252            ));
253        }
254        Ok(Self {
255            low,
256            high,
257            scale,
258            step: 1,
259        })
260    }
261    /// Set a positive linear step anchored at `low`.
262    pub fn with_step(mut self, step: u64) -> Result<Self, ParzenError> {
263        if step == 0 {
264            return Err(ParzenError::InvalidDistribution(
265                "integer step must be positive".into(),
266            ));
267        }
268        if self.scale == IntScale::Log && step != 1 {
269            return Err(ParzenError::InvalidDistribution(
270                "log-integer distributions require unit step".into(),
271            ));
272        }
273        self.step = step;
274        Ok(self)
275    }
276    #[must_use]
277    pub const fn low(self) -> i64 {
278        self.low
279    }
280    #[must_use]
281    pub const fn high(self) -> i64 {
282        self.high
283    }
284    #[must_use]
285    pub const fn scale(self) -> IntScale {
286        self.scale
287    }
288    #[must_use]
289    pub const fn step(self) -> u64 {
290        self.step
291    }
292
293    pub(crate) fn transform(self, value: i64) -> f64 {
294        let value = value as f64;
295        match self.scale {
296            IntScale::Linear => value,
297            IntScale::Log => value.ln(),
298        }
299    }
300
301    pub(crate) fn untransform(self, value: f64) -> i64 {
302        let raw = match self.scale {
303            IntScale::Linear => value,
304            IntScale::Log => value.exp(),
305        };
306        let low = self.low as i128;
307        let rounded = raw.round().clamp(self.low as f64, self.high as f64) as i128;
308        let step = i128::from(self.step);
309        let aligned = low + ((rounded - low + step / 2) / step) * step;
310        let high = i128::from(self.high);
311        let highest_legal = low + ((high - low) / step) * step;
312        aligned.clamp(low, highest_legal) as i64
313    }
314
315    pub(crate) fn max_step_index(self) -> u64 {
316        let width = i128::from(self.high) - i128::from(self.low);
317        (width / i128::from(self.step)) as u64
318    }
319
320    pub(crate) fn grid_value(self, index: u64) -> i64 {
321        (i128::from(self.low) + i128::from(index) * i128::from(self.step)) as i64
322    }
323}
324
325impl Distribution {
326    pub(crate) fn canonicalize(&self, value: ParamValue) -> Option<ParamValue> {
327        match (self, value) {
328            (Self::Categorical(dist), ParamValue::Categorical(choice)) => {
329                (choice < dist.num_choices).then_some(ParamValue::Categorical(choice))
330            }
331            (Self::Float(dist), ParamValue::Float(value)) => {
332                if !value.is_finite() || value < dist.low || value > dist.high {
333                    return None;
334                }
335                let canonical = dist.quantize(value);
336                if dist.step.is_none() || ulp_distance(value, canonical) <= 4 {
337                    Some(ParamValue::Float(canonical))
338                } else {
339                    None
340                }
341            }
342            (Self::Int(dist), ParamValue::Int(value)) => {
343                (value >= dist.low && value <= dist.high && {
344                    let delta = i128::from(value) - i128::from(dist.low);
345                    delta % i128::from(dist.step) == 0
346                })
347                .then_some(ParamValue::Int(value))
348            }
349            _ => None,
350        }
351    }
352}
353
354fn ulp_distance(left: f64, right: f64) -> u64 {
355    fn ordered(value: f64) -> u64 {
356        let bits = value.to_bits();
357        if bits >> 63 == 0 {
358            bits | (1 << 63)
359        } else {
360            !bits
361        }
362    }
363    ordered(left).abs_diff(ordered(right))
364}
365
366#[cfg(feature = "serde")]
367mod deserialize {
368    use serde::{Deserialize, Deserializer, de::Error as _};
369
370    use super::*;
371
372    #[derive(Deserialize)]
373    #[serde(deny_unknown_fields)]
374    struct CategoricalWire {
375        num_choices: u32,
376    }
377
378    impl<'de> Deserialize<'de> for CategoricalDistribution {
379        fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
380            let wire = CategoricalWire::deserialize(deserializer)?;
381            Self::new(wire.num_choices).map_err(D::Error::custom)
382        }
383    }
384
385    #[derive(Deserialize)]
386    #[serde(deny_unknown_fields)]
387    struct FloatWire {
388        low: f64,
389        high: f64,
390        scale: FloatScale,
391        step: Option<f64>,
392    }
393
394    impl<'de> Deserialize<'de> for FloatDistribution {
395        fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
396            let wire = FloatWire::deserialize(deserializer)?;
397            let distribution = match wire.scale {
398                FloatScale::Linear => Self::linear(wire.low, wire.high),
399                FloatScale::Log => Self::log(wire.low, wire.high),
400            }
401            .map_err(D::Error::custom)?;
402            wire.step.map_or(Ok(distribution), |step| {
403                distribution.with_step(step).map_err(D::Error::custom)
404            })
405        }
406    }
407
408    #[derive(Deserialize)]
409    #[serde(deny_unknown_fields)]
410    struct IntWire {
411        low: i64,
412        high: i64,
413        scale: IntScale,
414        step: u64,
415    }
416
417    impl<'de> Deserialize<'de> for IntDistribution {
418        fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
419            let wire = IntWire::deserialize(deserializer)?;
420            let distribution = match wire.scale {
421                IntScale::Linear => Self::linear(wire.low, wire.high),
422                IntScale::Log => Self::log(wire.low, wire.high),
423            }
424            .map_err(D::Error::custom)?;
425            distribution.with_step(wire.step).map_err(D::Error::custom)
426        }
427    }
428}
parzen/distribution.rs

parzen/
distribution.rs
