Implemented untested modular distribution samplers

Author: juntyr

necsim/core/bond/src/closed_open_unit_f64.rs

@@ -45,6 +45,14 @@ impl From<ClosedOpenUnitF64> for f64 {
     }
 }
 
+impl TryFrom<ClosedUnitF64> for ClosedOpenUnitF64 {
+    type Error = ClosedOpenUnitF64Error;
+
+    fn try_from(value: ClosedUnitF64) -> Result<Self, Self::Error> {
+        Self::new(value.get())
+    }
+}
+
 impl fmt::Debug for ClosedOpenUnitF64 {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         struct ClosedOpenUnitF64Range(f64);

necsim/core/bond/src/open_closed_unit_f64.rs

@@ -8,7 +8,7 @@ use core::{
 use necsim_core_maths::MathsCore;
 use serde::{Deserialize, Serialize};
 
-use crate::NonPositiveF64;
+use crate::{ClosedUnitF64, NonPositiveF64};
 
 #[derive(Debug)]
 #[allow(clippy::module_name_repetitions)]
@@ -46,6 +46,14 @@ impl From<OpenClosedUnitF64> for f64 {
     }
 }
 
+impl TryFrom<ClosedUnitF64> for OpenClosedUnitF64 {
+    type Error = OpenClosedUnitF64Error;
+
+    fn try_from(value: ClosedUnitF64) -> Result<Self, Self::Error> {
+        Self::new(value.get())
+    }
+}
+
 impl fmt::Debug for OpenClosedUnitF64 {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         struct OpenClosedUnitF64Range(f64);

necsim/impls/no-std/src/cogs/distribution/bernoulli_64b.rs

@@ -0,0 +1,32 @@
+use necsim_core::cogs::{distribution::Bernoulli, DistributionSampler, MathsCore, RngCore};
+use necsim_core_bond::ClosedUnitF64;
+
+#[allow(clippy::module_name_repetitions)]
+pub struct Bernoulli64BitSampler;
+
+impl<M: MathsCore, R: RngCore, S> DistributionSampler<M, R, S, Bernoulli>
+    for Bernoulli64BitSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(&self, rng: &mut R, _samplers: &S, probability: ClosedUnitF64) -> bool {
+        #[allow(clippy::cast_precision_loss)]
+        const SCALE: f64 = 2.0 * (1u64 << 63) as f64;
+
+        // Safety:
+        //  (a) 0 <= probability < 1: probability * SCALE is in [0, 2^64)
+        //                            since 1 - 2^-53 is before 1.0
+        //  (b) probability == 1    : p_u64 is undefined
+        //                            this case is checked for in the return
+        let p_u64 = unsafe { (probability.get() * SCALE).to_int_unchecked::<u64>() };
+
+        #[allow(clippy::float_cmp)]
+        {
+            (rng.sample_u64() < p_u64) || (probability == 1.0_f64)
+        }
+    }
+}

necsim/impls/no-std/src/cogs/distribution/bernoulli_from_unit.rs

@@ -0,0 +1,24 @@
+use necsim_core::cogs::{
+    distribution::{Bernoulli, RawDistribution, UniformClosedOpenUnit},
+    DistributionSampler, MathsCore, RngCore,
+};
+use necsim_core_bond::ClosedUnitF64;
+
+#[allow(clippy::module_name_repetitions)]
+pub struct BernoulliFromUnitSampler;
+
+impl<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, UniformClosedOpenUnit>>
+    DistributionSampler<M, R, S, Bernoulli> for BernoulliFromUnitSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(&self, rng: &mut R, samplers: &S, probability: ClosedUnitF64) -> bool {
+        // if probability == 1, then U[0, 1) always < 1.0
+        // if probability == 0, then U[0, 1) never < 0.0
+        UniformClosedOpenUnit::sample_raw(rng, samplers) < probability
+    }
+}

necsim/impls/no-std/src/cogs/distribution/exp_inversion.rs

@@ -0,0 +1,30 @@
+use necsim_core::cogs::{
+    distribution::{Exponential, Lambda, RawDistribution, UniformOpenClosedUnit},
+    DistributionSampler, MathsCore, RngCore,
+};
+use necsim_core_bond::NonNegativeF64;
+
+#[allow(clippy::module_name_repetitions)]
+pub struct ExponentialInverseTransformSampler;
+
+impl<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, UniformOpenClosedUnit>>
+    DistributionSampler<M, R, S, Exponential> for ExponentialInverseTransformSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        samplers: &S,
+        Lambda(lambda): Lambda,
+    ) -> NonNegativeF64 {
+        let u01 = UniformOpenClosedUnit::sample_raw(rng, samplers);
+
+        // Inverse transform sample: X = -ln(U(0,1]) / lambda
+        -u01.ln::<M>() / lambda
+    }
+}

necsim/impls/no-std/src/cogs/distribution/index_from_unit.rs

@@ -0,0 +1,124 @@
+use core::num::{NonZeroU128, NonZeroU32, NonZeroU64, NonZeroUsize};
+
+use necsim_core::cogs::{
+    distribution::{
+        IndexU128, IndexU32, IndexU64, IndexUsize, Length, RawDistribution, UniformClosedOpenUnit,
+    },
+    DistributionSampler, MathsCore, RngCore,
+};
+
+#[allow(clippy::module_name_repetitions)]
+pub struct IndexFromUnitSampler;
+
+impl<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, UniformClosedOpenUnit>>
+    DistributionSampler<M, R, S, IndexUsize> for IndexFromUnitSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        samplers: &S,
+        Length(length): Length<NonZeroUsize>,
+    ) -> usize {
+        let u01 = UniformClosedOpenUnit::sample_raw(rng, samplers);
+
+        // Safety: U[0, 1) * length in [0, 2^[32/64]) is a valid [u32/u64]
+        //         since (1 - 2^-53) * 2^[32/64] <= (2^[32/64] - 1)
+        #[allow(clippy::cast_precision_loss)]
+        let index =
+            unsafe { M::floor(u01.get() * (length.get() as f64)).to_int_unchecked::<usize>() };
+
+        if cfg!(target_pointer_width = "32") {
+            // Note: [0, 2^32) is losslessly represented in f64
+            index
+        } else {
+            // Note: Ensure index < length despite
+            //       usize->f64->usize precision loss
+            index.min(length.get() - 1)
+        }
+    }
+}
+
+impl<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, UniformClosedOpenUnit>>
+    DistributionSampler<M, R, S, IndexU32> for IndexFromUnitSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        samplers: &S,
+        Length(length): Length<NonZeroU32>,
+    ) -> u32 {
+        let u01 = UniformClosedOpenUnit::sample_raw(rng, samplers);
+
+        // Safety: U[0, 1) * length in [0, 2^32) is losslessly represented
+        //         in both f64 and u32
+        unsafe { M::floor(u01.get() * f64::from(length.get())).to_int_unchecked::<u32>() }
+    }
+}
+
+impl<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, UniformClosedOpenUnit>>
+    DistributionSampler<M, R, S, IndexU64> for IndexFromUnitSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        samplers: &S,
+        Length(length): Length<NonZeroU64>,
+    ) -> u64 {
+        let u01 = UniformClosedOpenUnit::sample_raw(rng, samplers);
+
+        // Safety: U[0, 1) * length in [0, 2^64) is a valid u64
+        //         since (1 - 2^-53) * 2^64 <= (2^64 - 1)
+        #[allow(clippy::cast_precision_loss)]
+        let index =
+            unsafe { M::floor(u01.get() * (length.get() as f64)).to_int_unchecked::<u64>() };
+
+        // Note: Ensure index < length despite u64->f64->u64 precision loss
+        index.min(length.get() - 1)
+    }
+}
+
+impl<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, UniformClosedOpenUnit>>
+    DistributionSampler<M, R, S, IndexU128> for IndexFromUnitSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        samplers: &S,
+        Length(length): Length<NonZeroU128>,
+    ) -> u128 {
+        let u01 = UniformClosedOpenUnit::sample_raw(rng, samplers);
+
+        // Safety: U[0, 1) * length in [0, 2^128) is a valid u128
+        //         since (1 - 2^-53) * 2^128 <= (2^128 - 1)
+        #[allow(clippy::cast_precision_loss)]
+        let index =
+            unsafe { M::floor(u01.get() * (length.get() as f64)).to_int_unchecked::<u128>() };
+
+        // Note: Ensure index < length despite u128->f64->u128 precision loss
+        index.min(length.get() - 1)
+    }
+}

necsim/impls/no-std/src/cogs/distribution/index_from_wide.rs

@@ -0,0 +1,135 @@
+use core::num::{NonZeroU128, NonZeroU32, NonZeroU64, NonZeroUsize};
+
+use necsim_core::cogs::{
+    distribution::{IndexU128, IndexU32, IndexU64, IndexUsize, Length, RawDistribution},
+    DistributionSampler, MathsCore, RngCore,
+};
+
+#[allow(clippy::module_name_repetitions)]
+pub struct IndexFromWideMulSampler;
+
+impl<M: MathsCore, R: RngCore, S> DistributionSampler<M, R, S, IndexUsize>
+    for IndexFromWideMulSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    #[inline]
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        _samplers: &S,
+        Length(length): Length<NonZeroUsize>,
+    ) -> usize {
+        #[allow(clippy::cast_possible_truncation)]
+        if cfg!(target_pointer_width = "32") {
+            IndexU32::sample_raw_with::<M, R, Self>(
+                rng,
+                self,
+                Length(unsafe { NonZeroU32::new_unchecked(length.get() as u32) }),
+            ) as usize
+        } else {
+            IndexU64::sample_raw_with::<M, R, Self>(
+                rng,
+                self,
+                Length(unsafe { NonZeroU64::new_unchecked(length.get() as u64) }),
+            ) as usize
+        }
+    }
+}
+
+impl<M: MathsCore, R: RngCore, S> DistributionSampler<M, R, S, IndexU32>
+    for IndexFromWideMulSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        _samplers: &S,
+        Length(length): Length<NonZeroU32>,
+    ) -> u32 {
+        // Sample U(0, length - 1) using a widening multiplication
+        // Note: Some slight bias is traded for only needing one u64 sample
+        // Note: Should optimise to a single 64 bit (high-only) multiplication
+        #[allow(clippy::cast_possible_truncation)]
+        {
+            (((u128::from(rng.sample_u64()) * u128::from(length.get())) >> 64) & u128::from(!0_u32))
+                as u32
+        }
+    }
+}
+
+impl<M: MathsCore, R: RngCore, S> DistributionSampler<M, R, S, IndexU64>
+    for IndexFromWideMulSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        _samplers: &S,
+        Length(length): Length<NonZeroU64>,
+    ) -> u64 {
+        // Sample U(0, length - 1) using a widening multiplication
+        // Note: Some slight bias is traded for only needing one u64 sample
+        // Note: Should optimise to a single 64 bit (high-only) multiplication
+        #[allow(clippy::cast_possible_truncation)]
+        {
+            (((u128::from(rng.sample_u64()) * u128::from(length.get())) >> 64) & u128::from(!0_u64))
+                as u64
+        }
+    }
+}
+
+impl<M: MathsCore, R: RngCore, S> DistributionSampler<M, R, S, IndexU128>
+    for IndexFromWideMulSampler
+{
+    type ConcreteSampler = Self;
+
+    fn concrete(&self) -> &Self::ConcreteSampler {
+        self
+    }
+
+    fn sample_distribution(
+        &self,
+        rng: &mut R,
+        _samplers: &S,
+        Length(length): Length<NonZeroU128>,
+    ) -> u128 {
+        // Sample U(0, length - 1) using a widening multiplication
+        // Note: Some slight bias is traded for only needing one u128 sample
+
+        const LOWER_MASK: u128 = !0 >> 64;
+
+        let raw_hi = u128::from(rng.sample_u64());
+        let raw_lo = u128::from(rng.sample_u64());
+
+        // 256-bit multiplication (hi, lo) = (raw_hi, raw_lo) * length
+        let mut low = raw_lo * (length.get() & LOWER_MASK);
+        let mut t = low >> 64;
+        low &= LOWER_MASK;
+        t += raw_hi * (length.get() & LOWER_MASK);
+        low += (t & LOWER_MASK) << 64;
+        let mut high = t >> 64;
+        t = low >> 64;
+        // low-only: low &= LOWER_MASK;
+        t += (length.get() >> 64) * raw_lo;
+        // low-only: low += (t & LOWER_MASK) << 64;
+        high += t >> 64;
+        high += raw_hi * (length.get() >> 64);
+
+        high
+    }
+}