MultiDistribution (#18)

Co-authored-by: Diggory Hardy <git@dhardy.name>

Author: benjamin-lieser

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [0.6.0] — Unreleased
 - Bump to MSRV 1.85.0 and Edition 2024 in line with `rand` (#27)
 
+### Additions
+- `MultiDistribution` trait to sample more efficiently from multi-dimensional distributions (#18)
+
+### Changes
+- Moved `Dirichlet` into the new `multi` module and implement `MultiDistribution` for it (#18)
+
 ## [0.5.2]
 
 ### API Changes

src/lib.rs

@@ -72,7 +72,7 @@
 //!   - [`Beta`] distribution
 //!   - [`Triangular`] distribution
 //! - Multivariate probability distributions
-//!   - [`Dirichlet`] distribution
+//!   - [`multi::Dirichlet`] distribution
 //!   - [`UnitSphere`] distribution
 //!   - [`UnitBall`] distribution
 //!   - [`UnitCircle`] distribution
@@ -100,8 +100,6 @@ pub use self::beta::{Beta, Error as BetaError};
 pub use self::binomial::{Binomial, Error as BinomialError};
 pub use self::cauchy::{Cauchy, Error as CauchyError};
 pub use self::chi_squared::{ChiSquared, Error as ChiSquaredError};
-#[cfg(feature = "alloc")]
-pub use self::dirichlet::{Dirichlet, Error as DirichletError};
 pub use self::exponential::{Error as ExpError, Exp, Exp1};
 pub use self::fisher_f::{Error as FisherFError, FisherF};
 pub use self::frechet::{Error as FrechetError, Frechet};
@@ -130,6 +128,8 @@ pub use student_t::StudentT;
 
 pub use num_traits;
 
+#[cfg(feature = "alloc")]
+pub mod multi;
 #[cfg(feature = "alloc")]
 pub mod weighted;
 
@@ -188,7 +188,6 @@ mod beta;
 mod binomial;
 mod cauchy;
 mod chi_squared;
-mod dirichlet;
 mod exponential;
 mod fisher_f;
 mod frechet;

src/dirichlet.rs renamed to src/multi/dirichlet.rs

@@ -10,7 +10,7 @@
 //! The dirichlet distribution `Dirichlet(α₁, α₂, ..., αₙ)`.
 
 #![cfg(feature = "alloc")]
-use crate::{Beta, Distribution, Exp1, Gamma, Open01, StandardNormal};
+use crate::{Beta, Distribution, Exp1, Gamma, Open01, StandardNormal, multi::MultiDistribution};
 use core::fmt;
 use num_traits::{Float, NumCast};
 use rand::Rng;
@@ -68,26 +68,29 @@ where
     }
 }
 
-impl<F, const N: usize> Distribution<[F; N]> for DirichletFromGamma<F, N>
+impl<F, const N: usize> MultiDistribution<F> for DirichletFromGamma<F, N>
 where
     F: Float,
     StandardNormal: Distribution<F>,
     Exp1: Distribution<F>,
     Open01: Distribution<F>,
 {
-    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> [F; N] {
-        let mut samples = [F::zero(); N];
+    fn sample_len(&self) -> usize {
+        N
+    }
+    fn sample_to_slice<R: Rng + ?Sized>(&self, rng: &mut R, output: &mut [F]) {
+        assert_eq!(output.len(), N);
+
         let mut sum = F::zero();
 
-        for (s, g) in samples.iter_mut().zip(self.samplers.iter()) {
+        for (s, g) in output.iter_mut().zip(self.samplers.iter()) {
             *s = g.sample(rng);
             sum = sum + *s;
         }
         let invacc = F::one() / sum;
-        for s in samples.iter_mut() {
+        for s in output.iter_mut() {
             *s = *s * invacc;
         }
-        samples
     }
 }
 
@@ -149,24 +152,27 @@ where
     }
 }
 
-impl<F, const N: usize> Distribution<[F; N]> for DirichletFromBeta<F, N>
+impl<F, const N: usize> MultiDistribution<F> for DirichletFromBeta<F, N>
 where
     F: Float,
     StandardNormal: Distribution<F>,
     Exp1: Distribution<F>,
     Open01: Distribution<F>,
 {
-    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> [F; N] {
-        let mut samples = [F::zero(); N];
+    fn sample_len(&self) -> usize {
+        N
+    }
+    fn sample_to_slice<R: Rng + ?Sized>(&self, rng: &mut R, output: &mut [F]) {
+        assert_eq!(output.len(), N);
+
         let mut acc = F::one();
 
-        for (s, beta) in samples.iter_mut().zip(self.samplers.iter()) {
+        for (s, beta) in output.iter_mut().zip(self.samplers.iter()) {
             let beta_sample = beta.sample(rng);
             *s = acc * beta_sample;
             acc = acc * (F::one() - beta_sample);
         }
-        samples[N - 1] = acc;
-        samples
+        output[N - 1] = acc;
     }
 }
 
@@ -208,7 +214,8 @@ where
 ///
 /// ```
 /// use rand::prelude::*;
-/// use rand_distr::Dirichlet;
+/// use rand_distr::multi::Dirichlet;
+/// use rand_distr::multi::MultiDistribution;
 ///
 /// let dirichlet = Dirichlet::new([1.0, 2.0, 3.0]).unwrap();
 /// let samples = dirichlet.sample(&mut rand::rng());
@@ -259,7 +266,7 @@ impl fmt::Display for Error {
                 "failed to create required Gamma distribution for Dirichlet distribution"
             }
             Error::FailedToCreateBeta => {
-                "failed to create required Beta distribition for Dirichlet distribution"
+                "failed to create required Beta distribution for Dirichlet distribution"
             }
         })
     }
@@ -315,21 +322,34 @@ where
     }
 }
 
-impl<F, const N: usize> Distribution<[F; N]> for Dirichlet<F, N>
+impl<F, const N: usize> MultiDistribution<F> 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) -> [F; N] {
+    fn sample_len(&self) -> usize {
+        N
+    }
+    fn sample_to_slice<R: Rng + ?Sized>(&self, rng: &mut R, output: &mut [F]) {
         match &self.repr {
-            DirichletRepr::FromGamma(dirichlet) => dirichlet.sample(rng),
-            DirichletRepr::FromBeta(dirichlet) => dirichlet.sample(rng),
+            DirichletRepr::FromGamma(dirichlet) => dirichlet.sample_to_slice(rng, output),
+            DirichletRepr::FromBeta(dirichlet) => dirichlet.sample_to_slice(rng, output),
         }
     }
 }
 
+impl<F, const N: usize> Distribution<Vec<F>> for Dirichlet<F, N>
+where
+    F: Float + Default,
+    StandardNormal: Distribution<F>,
+    Exp1: Distribution<F>,
+    Open01: Distribution<F>,
+{
+    distribution_impl!(F);
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -403,7 +423,7 @@ mod test {
         let alpha_sum: f64 = alpha.iter().sum();
         let expected_mean = alpha.map(|x| x / alpha_sum);
         for i in 0..N {
-            assert_almost_eq!(sample_mean[i], expected_mean[i], rtol);
+            average::assert_almost_eq!(sample_mean[i], expected_mean[i], rtol);
         }
     }
 

src/multi/mod.rs

@@ -0,0 +1,37 @@
+// Copyright 2025 Developers of the Rand project.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+//! Contains Multi-dimensional distributions.
+//!
+//! We provide a trait `MultiDistribution` which allows to sample from a multi-dimensional distribution without extra allocations.
+//! All multi-dimensional distributions implement `MultiDistribution` instead of the `Distribution` trait.
+
+use rand::Rng;
+
+/// A standard abstraction for distributions with multi-dimensional results
+pub trait MultiDistribution<T> {
+    /// returns the length of one sample (dimension of the distribution)
+    fn sample_len(&self) -> usize;
+    /// samples from the distribution and writes the result to `output`
+    fn sample_to_slice<R: Rng + ?Sized>(&self, rng: &mut R, output: &mut [T]);
+}
+
+macro_rules! distribution_impl {
+    ($scalar:ident) => {
+        fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<$scalar> {
+            use crate::multi::MultiDistribution;
+            let mut buf = vec![Default::default(); self.sample_len()];
+            self.sample_to_slice(rng, &mut buf);
+            buf
+        }
+    };
+}
+
+pub use dirichlet::Dirichlet;
+
+mod dirichlet;

tests/value_stability.rs

@@ -502,11 +502,13 @@ 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()),
+        multi::Dirichlet::new([1.0, 2.0, 3.0])
+            .unwrap()
+            .sample(&mut rng),
         [0.12941567177708177, 0.4702121891675036, 0.4003721390554146]
     );
     assert_eq!(
-        rng.sample(Dirichlet::new([8.0; 5]).unwrap()),
+        multi::Dirichlet::new([8.0; 5]).unwrap().sample(&mut rng),
         [
             0.17684200044809556,
             0.29915953935953055,
@@ -517,7 +519,9 @@ fn dirichlet_stability() {
     );
     // Test stability for the case where all alphas are less than 0.1.
     assert_eq!(
-        rng.sample(Dirichlet::new([0.05, 0.025, 0.075, 0.05]).unwrap()),
+        multi::Dirichlet::new([0.05, 0.025, 0.075, 0.05])
+            .unwrap()
+            .sample(&mut rng),
         [
             0.00027580456855692104,
             2.296135759821706e-20,