yggdrasil/rand
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Use const generics in Dirichlet (#1292)

Browse Source 
* Use const generics in Dirichlet

* Serialize const arrays with serde_with
This commit is contained in:
Virgile Andreani
2023-03-18 05:21:59 -04:00
committed by GitHub
parent 0f3ecedef0
commit 0f5af6672d
4 changed files with 25 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rand_distr/CHANGELOG.md
+2
View File
@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- Upgrade Rand
- Fix Knuth's method so `Poisson` doesn't return -1.0 for small lambda
- Fix `Poisson` distribution instantiation so it return an error if lambda is infinite
- `Dirichlet` now uses `const` generics, which means that its size is required at compile time (#1292)
- The `Dirichlet::new_with_size` constructor was removed (#1292)
## [0.4.3] - 2021-12-30
- Fix `no_std` build (#1208)
rand_distr/Cargo.toml
+1
View File
@@ -27,6 +27,7 @@ serde1 = ["serde", "rand/serde1"]
rand = { path = "..", version = "0.9.0", default-features = false }
num-traits = { version = "0.2", default-features = false, features = ["libm"] }
serde = { version = "1.0.103", features = ["derive"], optional = true }
serde_with = { version = "1.14.0", optional = true }
[dev-dependencies]
rand_pcg = { version = "0.4.0", path = "../rand_pcg" }
rand_distr/src/dirichlet.rs
+19 -49
View File
@@ -13,7 +13,8 @@ use num_traits::Float;
use crate::{Distribution, Exp1, Gamma, Open01, StandardNormal};
use rand::Rng;
use core::fmt;
use alloc::{boxed::Box, vec, vec::Vec};
#[cfg(feature = "serde_with")]
use serde_with::serde_as;
/// The Dirichlet distribution `Dirichlet(alpha)`.
///
@@ -27,14 +28,14 @@ use alloc::{boxed::Box, vec, vec::Vec};
/// use rand::prelude::*;
/// use rand_distr::Dirichlet;
///
/// let dirichlet = Dirichlet::new(&[1.0, 2.0, 3.0]).unwrap();
/// let dirichlet = Dirichlet::new([1.0, 2.0, 3.0]).unwrap();
/// let samples = dirichlet.sample(&mut rand::thread_rng());
/// println!("{:?} is from a Dirichlet([1.0, 2.0, 3.0]) distribution", samples);
/// ```
#[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
#[cfg_attr(feature = "serde_with", serde_as)]
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct Dirichlet<F>
pub struct Dirichlet<F, const N: usize>
where
F: Float,
StandardNormal: Distribution<F>,
@@ -42,7 +43,8 @@ where
Open01: Distribution<F>,
{
/// Concentration parameters (alpha)
alpha: Box<[F]>,
#[cfg_attr(feature = "serde_with", serde_as(as = "[_; N]"))]
alpha: [F; N],
}
/// Error type returned from `Dirchlet::new`.
@@ -72,7 +74,7 @@ impl fmt::Display for Error {
#[cfg_attr(doc_cfg, doc(cfg(feature = "std")))]
impl std::error::Error for Error {}
impl<F> Dirichlet<F>
impl<F, const N: usize> Dirichlet<F, N>
where
F: Float,
StandardNormal: Distribution<F>,
@@ -83,8 +85,8 @@ where
///
/// Requires `alpha.len() >= 2`.
#[inline]
pub fn new(alpha: &[F]) -> Result<Dirichlet<F>, Error> {
if alpha.len() < 2 {
pub fn new(alpha: [F; N]) -> Result<Dirichlet<F, N>, Error> {
if N < 2 {
return Err(Error::AlphaTooShort);
}
for &ai in alpha.iter() {
@@ -93,36 +95,19 @@ where
}
}
Ok(Dirichlet { alpha: alpha.to_vec().into_boxed_slice() })
}
/// Construct a new `Dirichlet` with the given shape parameter `alpha` and `size`.
///
/// Requires `size >= 2`.
#[inline]
pub fn new_with_size(alpha: F, size: usize) -> Result<Dirichlet<F>, Error> {
if !(alpha > F::zero()) {
return Err(Error::AlphaTooSmall);
}
if size < 2 {
return Err(Error::SizeTooSmall);
}
Ok(Dirichlet {
alpha: vec![alpha; size].into_boxed_slice(),
})
Ok(Dirichlet { alpha })
}
}
impl<F> Distribution<Vec<F>> for Dirichlet<F>
impl<F, const N: usize> Distribution<[F; N]> for Dirichlet<F, N>
where
F: Float,
StandardNormal: Distribution<F>,
Exp1: Distribution<F>,
Open01: Distribution<F>,
{
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<F> {
let n = self.alpha.len();
let mut samples = vec![F::zero(); n];
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> [F; N] {
let mut samples = [F::zero(); N];
let mut sum = F::zero();
for (s, &a) in samples.iter_mut().zip(self.alpha.iter()) {
@@ -140,27 +125,12 @@ where
#[cfg(test)]
mod test {
use alloc::vec::Vec;
use super::*;
#[test]
fn test_dirichlet() {
let d = Dirichlet::new(&[1.0, 2.0, 3.0]).unwrap();
let mut rng = crate::test::rng(221);
let samples = d.sample(&mut rng);
let _: Vec<f64> = samples
.into_iter()
.map(|x| {
assert!(x > 0.0);
x
})
.collect();
}
#[test]
fn test_dirichlet_with_param() {
let alpha = 0.5f64;
let size = 2;
let d = Dirichlet::new_with_size(alpha, size).unwrap();
let d = Dirichlet::new([1.0, 2.0, 3.0]).unwrap();
let mut rng = crate::test::rng(221);
let samples = d.sample(&mut rng);
let _: Vec<f64> = samples
@@ -175,17 +145,17 @@ mod test {
#[test]
#[should_panic]
fn test_dirichlet_invalid_length() {
Dirichlet::new_with_size(0.5f64, 1).unwrap();
Dirichlet::new([0.5]).unwrap();
}
#[test]
#[should_panic]
fn test_dirichlet_invalid_alpha() {
Dirichlet::new_with_size(0.0f64, 2).unwrap();
Dirichlet::new([0.1, 0.0, 0.3]).unwrap();
}
#[test]
fn dirichlet_distributions_can_be_compared() {
assert_eq!(Dirichlet::new(&[1.0, 2.0]), Dirichlet::new(&[1.0, 2.0]));
assert_eq!(Dirichlet::new([1.0, 2.0]), Dirichlet::new([1.0, 2.0]));
}
}
rand_distr/tests/value_stability.rs
+3 -3
View File
@@ -348,10 +348,10 @@ fn weibull_stability() {
fn dirichlet_stability() {
let mut rng = get_rng(223);
assert_eq!(
rng.sample(Dirichlet::new(&[1.0, 2.0, 3.0]).unwrap()),
vec![0.12941567177708177, 0.4702121891675036, 0.4003721390554146]
rng.sample(Dirichlet::new([1.0, 2.0, 3.0]).unwrap()),
[0.12941567177708177, 0.4702121891675036, 0.4003721390554146]
);
assert_eq!(rng.sample(Dirichlet::new_with_size(8.0, 5).unwrap()), vec![
assert_eq!(rng.sample(Dirichlet::new([8.0; 5]).unwrap()), [
0.17684200044809556,
0.29915953935953055,
0.1832858056608014,