Merge pull request #30 from rust-random/push-ywmuottvwtvt

Remove const-generic size parameter from Dirichlet distribution

Author: dhardy

CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changes
 - Moved `Dirichlet` into the new `multi` module and implement `MultiDistribution` for it (#18)
+- `Dirichlet` no longer uses `const` generics, which means that its size is not required at compile time. Essentially a revert of rand#1292. (#15)
+- Add `Dirichlet::new_with_size` constructor (#15)
 
 ## [0.5.2]
 

src/multi/dirichlet.rs

@@ -21,27 +21,24 @@ use alloc::{boxed::Box, vec, vec::Vec};
 
 #[derive(Clone, Debug, PartialEq)]
 #[cfg_attr(feature = "serde", serde_as)]
-struct DirichletFromGamma<F, const N: usize>
+struct DirichletFromGamma<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
     Exp1: Distribution<F>,
     Open01: Distribution<F>,
 {
-    samplers: [Gamma<F>; N],
+    samplers: Vec<Gamma<F>>,
 }
 
 /// Error type returned from [`DirchletFromGamma::new`].
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 enum DirichletFromGammaError {
     /// Gamma::new(a, 1) failed.
     GammmaNewFailed,
-
-    /// gamma_dists.try_into() failed (in theory, this should not happen).
-    GammaArrayCreationFailed,
 }
 
-impl<F, const N: usize> DirichletFromGamma<F, N>
+impl<F> DirichletFromGamma<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
@@ -53,33 +50,32 @@ where
     /// This function is part of a private implementation detail.
     /// It assumes that the input is correct, so no validation of alpha is done.
     #[inline]
-    fn new(alpha: [F; N]) -> Result<DirichletFromGamma<F, N>, DirichletFromGammaError> {
+    fn new(alpha: &[F]) -> Result<DirichletFromGamma<F>, DirichletFromGammaError> {
         let mut gamma_dists = Vec::new();
         for a in alpha {
             let dist =
-                Gamma::new(a, F::one()).map_err(|_| DirichletFromGammaError::GammmaNewFailed)?;
+                Gamma::new(*a, F::one()).map_err(|_| DirichletFromGammaError::GammmaNewFailed)?;
             gamma_dists.push(dist);
         }
         Ok(DirichletFromGamma {
-            samplers: gamma_dists
-                .try_into()
-                .map_err(|_| DirichletFromGammaError::GammaArrayCreationFailed)?,
+            samplers: gamma_dists,
         })
     }
 }
 
-impl<F, const N: usize> MultiDistribution<F> for DirichletFromGamma<F, N>
+impl<F> MultiDistribution<F> for DirichletFromGamma<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
     Exp1: Distribution<F>,
     Open01: Distribution<F>,
 {
+    #[inline]
     fn sample_len(&self) -> usize {
-        N
+        self.samplers.len()
     }
     fn sample_to_slice<R: Rng + ?Sized>(&self, rng: &mut R, output: &mut [F]) {
-        assert_eq!(output.len(), N);
+        assert_eq!(output.len(), self.sample_len());
 
         let mut sum = F::zero();
 
@@ -96,7 +92,7 @@ where
 
 #[derive(Clone, Debug, PartialEq)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-struct DirichletFromBeta<F, const N: usize>
+struct DirichletFromBeta<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
@@ -113,7 +109,7 @@ enum DirichletFromBetaError {
     BetaNewFailed,
 }
 
-impl<F, const N: usize> DirichletFromBeta<F, N>
+impl<F> DirichletFromBeta<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
@@ -125,15 +121,16 @@ where
     /// This function is part of a private implementation detail.
     /// It assumes that the input is correct, so no validation of alpha is done.
     #[inline]
-    fn new(alpha: [F; N]) -> Result<DirichletFromBeta<F, N>, DirichletFromBetaError> {
+    fn new(alpha: &[F]) -> Result<DirichletFromBeta<F>, DirichletFromBetaError> {
         // `alpha_rev_csum` is the reverse of the cumulative sum of the
         // reverse of `alpha[1..]`.  E.g. if `alpha = [a0, a1, a2, a3]`, then
         // `alpha_rev_csum` is `[a1 + a2 + a3, a2 + a3, a3]`.
         // Note that instances of DirichletFromBeta will always have N >= 2,
         // so the subtractions of 1, 2 and 3 from N in the following are safe.
-        let mut alpha_rev_csum = vec![alpha[N - 1]; N - 1];
-        for k in 0..(N - 2) {
-            alpha_rev_csum[N - 3 - k] = alpha_rev_csum[N - 2 - k] + alpha[N - 2 - k];
+        let n = alpha.len();
+        let mut alpha_rev_csum = vec![alpha[n - 1]; n - 1];
+        for k in 0..(n - 2) {
+            alpha_rev_csum[n - 3 - k] = alpha_rev_csum[n - 2 - k] + alpha[n - 2 - k];
         }
 
         // Zip `alpha[..(N-1)]` and `alpha_rev_csum`; for the example
@@ -142,7 +139,7 @@ where
         // Then pass each tuple to `Beta::new()` to create the `Beta`
         // instances.
         let mut beta_dists = Vec::new();
-        for (&a, &b) in alpha[..(N - 1)].iter().zip(alpha_rev_csum.iter()) {
+        for (&a, &b) in alpha[..(n - 1)].iter().zip(alpha_rev_csum.iter()) {
             let dist = Beta::new(a, b).map_err(|_| DirichletFromBetaError::BetaNewFailed)?;
             beta_dists.push(dist);
         }
@@ -152,18 +149,19 @@ where
     }
 }
 
-impl<F, const N: usize> MultiDistribution<F> for DirichletFromBeta<F, N>
+impl<F> MultiDistribution<F> for DirichletFromBeta<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
     Exp1: Distribution<F>,
     Open01: Distribution<F>,
 {
+    #[inline]
     fn sample_len(&self) -> usize {
-        N
+        self.samplers.len() + 1
     }
     fn sample_to_slice<R: Rng + ?Sized>(&self, rng: &mut R, output: &mut [F]) {
-        assert_eq!(output.len(), N);
+        assert_eq!(output.len(), self.sample_len());
 
         let mut acc = F::one();
 
@@ -172,24 +170,24 @@ where
             *s = acc * beta_sample;
             acc = acc * (F::one() - beta_sample);
         }
-        output[N - 1] = acc;
+        output[output.len() - 1] = acc;
     }
 }
 
 #[derive(Clone, Debug, PartialEq)]
 #[cfg_attr(feature = "serde", serde_as)]
-enum DirichletRepr<F, const N: usize>
+enum DirichletRepr<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
     Exp1: Distribution<F>,
     Open01: Distribution<F>,
 {
     /// Dirichlet distribution that generates samples using the Gamma distribution.
-    FromGamma(DirichletFromGamma<F, N>),
+    FromGamma(DirichletFromGamma<F>),
 
     /// Dirichlet distribution that generates samples using the Beta distribution.
-    FromBeta(DirichletFromBeta<F, N>),
+    FromBeta(DirichletFromBeta<F>),
 }
 
 /// The [Dirichlet distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution) `Dirichlet(α₁, α₂, ..., αₖ)`.
@@ -217,20 +215,20 @@ where
 /// use rand_distr::multi::Dirichlet;
 /// use rand_distr::multi::MultiDistribution;
 ///
-/// 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::rng());
-/// println!("{:?} is from a Dirichlet([1.0, 2.0, 3.0]) distribution", samples);
+/// println!("{:?} is from a Dirichlet(&[1.0, 2.0, 3.0]) distribution", samples);
 /// ```
 #[cfg_attr(feature = "serde", serde_as)]
 #[derive(Clone, Debug, PartialEq)]
-pub struct Dirichlet<F, const N: usize>
+pub struct Dirichlet<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
     Exp1: Distribution<F>,
     Open01: Distribution<F>,
 {
-    repr: DirichletRepr<F, N>,
+    repr: DirichletRepr<F>,
 }
 
 /// Error type returned from [`Dirichlet::new`].
@@ -275,7 +273,7 @@ impl fmt::Display for Error {
 #[cfg(feature = "std")]
 impl std::error::Error for Error {}
 
-impl<F, const N: usize> Dirichlet<F, N>
+impl<F> Dirichlet<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
@@ -287,8 +285,8 @@ where
     /// Requires `alpha.len() >= 2`, and each value in `alpha` must be positive,
     /// finite and not subnormal.
     #[inline]
-    pub fn new(alpha: [F; N]) -> Result<Dirichlet<F, N>, Error> {
-        if N < 2 {
+    pub fn new(alpha: &[F]) -> Result<Dirichlet<F>, Error> {
+        if alpha.len() < 2 {
             return Err(Error::AlphaTooShort);
         }
         for &ai in alpha.iter() {
@@ -322,15 +320,19 @@ where
     }
 }
 
-impl<F, const N: usize> MultiDistribution<F> for Dirichlet<F, N>
+impl<F> MultiDistribution<F> for Dirichlet<F>
 where
     F: Float,
     StandardNormal: Distribution<F>,
     Exp1: Distribution<F>,
     Open01: Distribution<F>,
 {
+    #[inline]
     fn sample_len(&self) -> usize {
-        N
+        match &self.repr {
+            DirichletRepr::FromGamma(dirichlet) => dirichlet.sample_len(),
+            DirichletRepr::FromBeta(dirichlet) => dirichlet.sample_len(),
+        }
     }
     fn sample_to_slice<R: Rng + ?Sized>(&self, rng: &mut R, output: &mut [F]) {
         match &self.repr {
@@ -340,7 +342,7 @@ where
     }
 }
 
-impl<F, const N: usize> Distribution<Vec<F>> for Dirichlet<F, N>
+impl<F> Distribution<Vec<F>> for Dirichlet<F>
 where
     F: Float + Default,
     StandardNormal: Distribution<F>,
@@ -356,7 +358,7 @@ mod test {
 
     #[test]
     fn test_dirichlet() {
-        let d = Dirichlet::new([1.0, 2.0, 3.0]).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);
         assert!(samples.into_iter().all(|x: f64| x > 0.0));
@@ -365,42 +367,42 @@ mod test {
     #[test]
     #[should_panic]
     fn test_dirichlet_invalid_length() {
-        Dirichlet::new([0.5]).unwrap();
+        Dirichlet::new(&[0.5]).unwrap();
     }
 
     #[test]
     #[should_panic]
     fn test_dirichlet_alpha_zero() {
-        Dirichlet::new([0.1, 0.0, 0.3]).unwrap();
+        Dirichlet::new(&[0.1, 0.0, 0.3]).unwrap();
     }
 
     #[test]
     #[should_panic]
     fn test_dirichlet_alpha_negative() {
-        Dirichlet::new([0.1, -1.5, 0.3]).unwrap();
+        Dirichlet::new(&[0.1, -1.5, 0.3]).unwrap();
     }
 
     #[test]
     #[should_panic]
     fn test_dirichlet_alpha_nan() {
-        Dirichlet::new([0.5, f64::NAN, 0.25]).unwrap();
+        Dirichlet::new(&[0.5, f64::NAN, 0.25]).unwrap();
     }
 
     #[test]
     #[should_panic]
     fn test_dirichlet_alpha_subnormal() {
-        Dirichlet::new([0.5, 1.5e-321, 0.25]).unwrap();
+        Dirichlet::new(&[0.5, 1.5e-321, 0.25]).unwrap();
     }
 
     #[test]
     #[should_panic]
     fn test_dirichlet_alpha_inf() {
-        Dirichlet::new([0.5, f64::INFINITY, 0.25]).unwrap();
+        Dirichlet::new(&[0.5, f64::INFINITY, 0.25]).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]));
     }
 
     /// Check that the means of the components of n samples from
@@ -410,7 +412,7 @@ mod test {
     /// This is a crude statistical test, but it will catch egregious
     /// mistakes.  It will also also fail if any samples contain nan.
     fn check_dirichlet_means<const N: usize>(alpha: [f64; N], n: i32, rtol: f64, seed: u64) {
-        let d = Dirichlet::new(alpha).unwrap();
+        let d = Dirichlet::new(&alpha).unwrap();
         let mut rng = crate::test::rng(seed);
         let mut sums = [0.0; N];
         for _ in 0..n {

tests/value_stability.rs

@@ -502,13 +502,13 @@ fn weibull_stability() {
 fn dirichlet_stability() {
     let mut rng = get_rng(223);
     assert_eq!(
-        multi::Dirichlet::new([1.0, 2.0, 3.0])
+        multi::Dirichlet::new(&[1.0, 2.0, 3.0])
             .unwrap()
             .sample(&mut rng),
         [0.12941567177708177, 0.4702121891675036, 0.4003721390554146]
     );
     assert_eq!(
-        multi::Dirichlet::new([8.0; 5]).unwrap().sample(&mut rng),
+        multi::Dirichlet::new(&[8.0; 5]).unwrap().sample(&mut rng),
         [
             0.17684200044809556,
             0.29915953935953055,
@@ -519,7 +519,7 @@ fn dirichlet_stability() {
     );
     // Test stability for the case where all alphas are less than 0.1.
     assert_eq!(
-        multi::Dirichlet::new([0.05, 0.025, 0.075, 0.05])
+        multi::Dirichlet::new(&[0.05, 0.025, 0.075, 0.05])
             .unwrap()
             .sample(&mut rng),
         [