Implement more robust index sampling

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/core/src/cogs/active_lineage_sampler.rs

@@ -3,8 +3,8 @@ use core::ops::ControlFlow;
 use necsim_core_bond::{NonNegativeF64, PositiveF64};
 
 use super::{
-    CoalescenceSampler, DispersalSampler, EmigrationExit, EventSampler, Habitat, ImmigrationEntry,
-    LineageStore, MathsCore, RngCore, SpeciationProbability, TurnoverRate,
+    Backup, CoalescenceSampler, DispersalSampler, EmigrationExit, EventSampler, Habitat,
+    ImmigrationEntry, LineageStore, MathsCore, Rng, SpeciationProbability, TurnoverRate,
 };
 
 use crate::{lineage::Lineage, simulation::partial::active_lineage_sampler::PartialSimulation};
@@ -15,7 +15,7 @@ use crate::{lineage::Lineage, simulation::partial::active_lineage_sampler::Parti
 pub trait ActiveLineageSampler<
     M: MathsCore,
     H: Habitat<M>,
-    G: RngCore<M>,
+    G: Rng<M>,
     S: LineageStore<M, H>,
     X: EmigrationExit<M, H, G, S>,
     D: DispersalSampler<M, H, G>,
@@ -24,7 +24,7 @@ pub trait ActiveLineageSampler<
     N: SpeciationProbability<M, H>,
     E: EventSampler<M, H, G, S, X, D, C, T, N>,
     I: ImmigrationEntry<M>,
->: crate::cogs::Backup + core::fmt::Debug
+>: Backup + core::fmt::Debug
 {
     type LineageIterator<'a>: Iterator<Item = &'a Lineage>
     where

necsim/core/src/cogs/coalescence_sampler.rs

@@ -1,21 +1,20 @@
-use core::cmp::{Ord, Ordering};
-
-use necsim_core_bond::ClosedOpenUnitF64;
+use core::{
+    cmp::{Ord, Ordering},
+    num::NonZeroU32,
+};
 
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    cogs::{Backup, MathsCore, RngCore},
+    cogs::{Backup, Habitat, LineageStore, MathsCore, Rng, RngCore},
     landscape::{IndexedLocation, Location},
     lineage::LineageInteraction,
 };
 
-use super::{Habitat, LineageStore};
-
 #[allow(clippy::inline_always, clippy::inline_fn_without_body)]
 #[contract_trait]
 pub trait CoalescenceSampler<M: MathsCore, H: Habitat<M>, S: LineageStore<M, H>>:
-    crate::cogs::Backup + core::fmt::Debug
+    Backup + core::fmt::Debug
 {
     #[must_use]
     #[debug_requires(habitat.get_habitat_at_location(&location) > 0, "location is habitable")]
@@ -31,7 +30,7 @@ pub trait CoalescenceSampler<M: MathsCore, H: Habitat<M>, S: LineageStore<M, H>>
 #[allow(clippy::unsafe_derive_deserialize)]
 #[derive(Debug, PartialEq, Serialize, Deserialize, TypeLayout)]
 #[repr(transparent)]
-pub struct CoalescenceRngSample(ClosedOpenUnitF64);
+pub struct CoalescenceRngSample(u64);
 
 #[contract_trait]
 impl Backup for CoalescenceRngSample {
@@ -57,24 +56,20 @@ impl Eq for CoalescenceRngSample {}
 impl CoalescenceRngSample {
     #[must_use]
     #[inline]
-    pub fn new<M: MathsCore, G: RngCore<M>>(rng: &mut G) -> Self {
-        use crate::cogs::RngSampler;
-
-        Self(rng.sample_uniform_closed_open())
+    pub fn new<M: MathsCore, G: Rng<M>>(rng: &mut G) -> Self {
+        Self(rng.generator().sample_u64())
     }
 
     #[must_use]
     #[inline]
-    #[debug_ensures(ret < length, "samples U(0, length - 1)")]
-    pub fn sample_coalescence_index<M: MathsCore>(self, length: u32) -> u32 {
-        // attributes on expressions are experimental
-        // see https://github.com/rust-lang/rust/issues/15701
-        #[allow(
-            clippy::cast_precision_loss,
-            clippy::cast_possible_truncation,
-            clippy::cast_sign_loss
-        )]
-        let index = M::floor(self.0.get() * f64::from(length)) as u32;
-        index
+    #[debug_ensures(ret < length.get(), "samples U(0, length - 1)")]
+    pub fn sample_coalescence_index(self, 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(self.0) * u128::from(length.get())) >> 64) & u128::from(!0_u32)) as u32
+        }
     }
 }

necsim/core/src/cogs/dispersal_sampler.rs

@@ -1,7 +1,7 @@
 use necsim_core_bond::ClosedUnitF64;
 
 use crate::{
-    cogs::{MathsCore, RngCore},
+    cogs::{Backup, MathsCore, Rng},
     landscape::Location,
 };
 
@@ -11,8 +11,8 @@ use super::Habitat;
 #[allow(clippy::no_effect_underscore_binding)]
 #[allow(clippy::module_name_repetitions)]
 #[contract_trait]
-pub trait DispersalSampler<M: MathsCore, H: Habitat<M>, G: RngCore<M>>:
-    crate::cogs::Backup + core::fmt::Debug
+pub trait DispersalSampler<M: MathsCore, H: Habitat<M>, G: Rng<M>>:
+    Backup + core::fmt::Debug
 {
     #[must_use]
     #[debug_requires(habitat.is_location_habitable(location), "location is habitable")]
@@ -29,7 +29,7 @@ pub trait DispersalSampler<M: MathsCore, H: Habitat<M>, G: RngCore<M>>:
 #[allow(clippy::no_effect_underscore_binding)]
 #[allow(clippy::module_name_repetitions)]
 #[contract_trait]
-pub trait SeparableDispersalSampler<M: MathsCore, H: Habitat<M>, G: RngCore<M>>:
+pub trait SeparableDispersalSampler<M: MathsCore, H: Habitat<M>, G: Rng<M>>:
     DispersalSampler<M, H, G>
 {
     #[must_use]

necsim/core/src/cogs/distribution.rs

@@ -0,0 +1,162 @@
+use core::num::{NonZeroU128, NonZeroU32, NonZeroU64, NonZeroUsize};
+
+use necsim_core_bond::{
+    ClosedOpenUnitF64, ClosedUnitF64, NonNegativeF64, OpenClosedUnitF64, PositiveF64,
+};
+
+use crate::cogs::{MathsCore, RngCore, Samples};
+
+#[allow(clippy::module_name_repetitions)]
+pub trait DistributionCore {
+    type Parameters;
+    type Sample;
+}
+
+pub trait Distribution: DistributionCore {
+    fn sample_with<M: MathsCore, R: Samples<M, Self>>(
+        rng: &mut R,
+        params: Self::Parameters,
+    ) -> Self::Sample;
+
+    fn sample<M: MathsCore, R: Samples<M, Self>>(rng: &mut R) -> Self::Sample
+    where
+        Self: DistributionCore<Parameters = ()>,
+    {
+        Self::sample_with(rng, ())
+    }
+}
+
+impl<D: DistributionCore> Distribution for D {
+    fn sample_with<M: MathsCore, R: Samples<M, Self>>(
+        rng: &mut R,
+        params: Self::Parameters,
+    ) -> Self::Sample {
+        rng.sample_with(params)
+    }
+}
+
+#[allow(clippy::module_name_repetitions)]
+pub trait RawDistribution: DistributionCore {
+    fn sample_raw_with<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, Self>>(
+        rng: &mut R,
+        samplers: &S,
+        params: Self::Parameters,
+    ) -> Self::Sample;
+
+    fn sample_raw<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, Self>>(
+        rng: &mut R,
+        samplers: &S,
+    ) -> Self::Sample
+    where
+        Self: DistributionCore<Parameters = ()>,
+    {
+        Self::sample_raw_with(rng, samplers, ())
+    }
+}
+
+impl<D: DistributionCore> RawDistribution for D {
+    fn sample_raw_with<M: MathsCore, R: RngCore, S: DistributionSampler<M, R, S, Self>>(
+        rng: &mut R,
+        samplers: &S,
+        params: Self::Parameters,
+    ) -> Self::Sample {
+        samplers.sample_distribution(rng, samplers, params)
+    }
+}
+
+#[allow(clippy::module_name_repetitions)]
+pub trait DistributionSampler<M: MathsCore, R: RngCore, S, D: DistributionCore + ?Sized> {
+    type ConcreteSampler: DistributionSampler<M, R, S, D>;
+
+    #[must_use]
+    fn concrete(&self) -> &Self::ConcreteSampler;
+
+    #[must_use]
+    fn sample_distribution(&self, rng: &mut R, samplers: &S, params: D::Parameters) -> D::Sample;
+}
+
+pub enum UniformClosedOpenUnit {}
+
+impl DistributionCore for UniformClosedOpenUnit {
+    type Parameters = ();
+    type Sample = ClosedOpenUnitF64;
+}
+
+pub enum UniformOpenClosedUnit {}
+
+impl DistributionCore for UniformOpenClosedUnit {
+    type Parameters = ();
+    type Sample = OpenClosedUnitF64;
+}
+
+pub enum IndexUsize {}
+
+pub struct Length<T>(pub T);
+
+impl DistributionCore for IndexUsize {
+    type Parameters = Length<NonZeroUsize>;
+    type Sample = usize;
+}
+
+pub enum IndexU32 {}
+
+impl DistributionCore for IndexU32 {
+    type Parameters = Length<NonZeroU32>;
+    type Sample = u32;
+}
+
+pub enum IndexU64 {}
+
+impl DistributionCore for IndexU64 {
+    type Parameters = Length<NonZeroU64>;
+    type Sample = u64;
+}
+
+pub enum IndexU128 {}
+
+impl DistributionCore for IndexU128 {
+    type Parameters = Length<NonZeroU128>;
+    type Sample = u128;
+}
+
+pub struct Lambda(pub PositiveF64);
+
+pub enum Exponential {}
+
+impl DistributionCore for Exponential {
+    type Parameters = Lambda;
+    type Sample = NonNegativeF64;
+}
+
+pub enum Poisson {}
+
+impl DistributionCore for Poisson {
+    type Parameters = Lambda;
+    type Sample = u64;
+}
+
+pub enum Bernoulli {}
+
+impl DistributionCore for Bernoulli {
+    type Parameters = ClosedUnitF64;
+    type Sample = bool;
+}
+
+pub enum StandardNormal2D {}
+
+impl DistributionCore for StandardNormal2D {
+    type Parameters = ();
+    type Sample = (f64, f64);
+}
+
+pub struct Normal {
+    pub mu: f64,
+    pub sigma: NonNegativeF64,
+}
+
+pub enum Normal2D {}
+
+impl DistributionCore for Normal2D {
+    type Parameters = Normal;
+    type Sample = (f64, f64);
+}

necsim/core/src/cogs/emigration_exit.rs

@@ -1,7 +1,7 @@
 use necsim_core_bond::{NonNegativeF64, PositiveF64};
 
 use crate::{
-    cogs::{Habitat, LineageStore, MathsCore, RngCore},
+    cogs::{Backup, Habitat, LineageStore, MathsCore, Rng},
     landscape::{IndexedLocation, Location},
     lineage::GlobalLineageReference,
     simulation::partial::emigration_exit::PartialSimulation,
@@ -14,8 +14,8 @@ use crate::{
 )]
 #[allow(clippy::no_effect_underscore_binding)]
 #[contract_trait]
-pub trait EmigrationExit<M: MathsCore, H: Habitat<M>, G: RngCore<M>, S: LineageStore<M, H>>:
-    crate::cogs::Backup + core::fmt::Debug
+pub trait EmigrationExit<M: MathsCore, H: Habitat<M>, G: Rng<M>, S: LineageStore<M, H>>:
+    Backup + core::fmt::Debug
 {
     #[must_use]
     #[debug_ensures(match &ret {