diff --git a/necsim/impls/no-std/src/alias/packed.rs b/necsim/impls/no-std/src/alias/packed.rs index fcce6f9af..596352209 100644 --- a/necsim/impls/no-std/src/alias/packed.rs +++ b/necsim/impls/no-std/src/alias/packed.rs @@ -22,6 +22,14 @@ struct AliasMethodSamplerAtomRaw { } impl AliasMethodSamplerAtom { + pub fn e(&self) -> &E { + &self.e + } + + pub fn k(&self) -> &E { + &self.k + } + #[allow(clippy::no_effect_underscore_binding)] #[debug_requires(!event_weights.is_empty(), "event_weights is non-empty")] #[debug_requires( diff --git a/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/mod.rs b/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/mod.rs index e046c4daa..c8e9f50c2 100644 --- a/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/mod.rs +++ b/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/mod.rs @@ -7,6 +7,7 @@ pub mod contract; pub mod alias; pub mod cumulative; pub mod packed_alias; +pub mod packed_separable_alias; pub mod separable_alias; use contract::explicit_in_memory_dispersal_check_contract; diff --git a/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/packed_separable_alias/dispersal.rs b/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/packed_separable_alias/dispersal.rs new file mode 100644 index 000000000..98523cbd7 --- /dev/null +++ b/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/packed_separable_alias/dispersal.rs @@ -0,0 +1,53 @@ +use core::ops::Range; + +use necsim_core::{ + cogs::{DispersalSampler, Habitat, MathsCore, RngCore}, + landscape::Location, +}; + +use crate::alias::packed::AliasMethodSamplerAtom; + +use super::InMemoryPackedSeparableAliasDispersalSampler; + +#[contract_trait] +impl, G: RngCore> DispersalSampler + for InMemoryPackedSeparableAliasDispersalSampler +{ + #[must_use] + fn sample_dispersal_from_location( + &self, + location: &Location, + habitat: &H, + rng: &mut G, + ) -> Location { + let location_row = location.y().wrapping_sub(habitat.get_extent().origin().y()) as usize; + let location_column = location.x().wrapping_sub(habitat.get_extent().origin().x()) as usize; + + // Only safe as trait precondition that `location` is inside `habitat` + let alias_range = unsafe { + Range::::from( + self.alias_dispersal_ranges + .get(location_row, location_column) + .unwrap_unchecked() + .clone(), + ) + }; + + // Safe by the construction of `InMemoryPackedAliasDispersalSampler` + let alias_dispersals_at_location = + unsafe { &self.alias_dispersal_buffer.get_unchecked(alias_range) }; + + let dispersal_target_index: usize = + AliasMethodSamplerAtom::sample_event(alias_dispersals_at_location, rng); + + #[allow(clippy::cast_possible_truncation)] + Location::new( + habitat.get_extent().origin().x().wrapping_add( + (dispersal_target_index % usize::from(habitat.get_extent().width())) as u32, + ), + habitat.get_extent().origin().y().wrapping_add( + (dispersal_target_index / usize::from(habitat.get_extent().width())) as u32, + ), + ) + } +} diff --git a/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/packed_separable_alias/mod.rs b/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/packed_separable_alias/mod.rs new file mode 100644 index 000000000..4f246b9f9 --- /dev/null +++ b/necsim/impls/no-std/src/cogs/dispersal_sampler/in_memory/packed_separable_alias/mod.rs @@ -0,0 +1,160 @@ +use alloc::{boxed::Box, vec::Vec}; +use core::{marker::PhantomData, ops::Range}; +use necsim_core_bond::NonNegativeF64; + +use r#final::Final; + +use necsim_core::{ + cogs::{Backup, Habitat, MathsCore, RngCore}, + landscape::Location, +}; + +use crate::{alias::packed::AliasMethodSamplerAtom, array2d::Array2D}; + +mod dispersal; + +use super::InMemoryDispersalSampler; + +#[derive(Clone, Debug, TypeLayout)] +#[allow(clippy::module_name_repetitions)] +#[doc(hidden)] +#[repr(C)] +pub struct SeparableAliasSamplerRange { + start: usize, + ledge: usize, + end: usize, +} + +impl From for Range { + fn from(range: SeparableAliasSamplerRange) -> Self { + range.start..range.end + } +} + +#[allow(clippy::module_name_repetitions)] +#[cfg_attr(feature = "cuda", derive(rust_cuda::lend::LendRustToCuda))] +#[cfg_attr(feature = "cuda", cuda(free = "M", free = "H", free = "G"))] +pub struct InMemoryPackedSeparableAliasDispersalSampler, G: RngCore> +{ + #[cfg_attr(feature = "cuda", cuda(embed))] + alias_dispersal_ranges: Final>, + #[cfg_attr(feature = "cuda", cuda(embed))] + alias_dispersal_buffer: Final]>>, + marker: PhantomData<(M, H, G)>, +} + +#[contract_trait] +impl, G: RngCore> InMemoryDispersalSampler + for InMemoryPackedSeparableAliasDispersalSampler +{ + /// Creates a new `InMemoryPackedSeparableAliasDispersalSampler` from the + /// `dispersal` map and extent of the habitat map. + fn unchecked_new(dispersal: &Array2D, habitat: &H) -> Self { + let habitat_extent = habitat.get_extent(); + + let mut event_weights: Vec<(usize, NonNegativeF64)> = + Vec::with_capacity(dispersal.row_len()); + + let mut alias_dispersal_buffer = Vec::new(); + + let alias_dispersal_ranges = Array2D::from_iter_row_major( + dispersal.rows_iter().enumerate().map(|(row_index, row)| { + event_weights.clear(); + + for (col_index, dispersal_probability) in row.enumerate() { + #[allow(clippy::cast_possible_truncation)] + let location = + Location::new( + habitat_extent.origin().x().wrapping_add( + (col_index % usize::from(habitat_extent.width())) as u32, + ), + habitat_extent.origin().y().wrapping_add( + (col_index / usize::from(habitat_extent.width())) as u32, + ), + ); + + // Multiply all dispersal probabilities by the habitat of their target + let weight = *dispersal_probability + * NonNegativeF64::from(habitat.get_habitat_at_location(&location)); + + if weight > 0.0_f64 { + event_weights.push((col_index, weight)); + } + } + + let range_from = alias_dispersal_buffer.len(); + + if event_weights.is_empty() { + SeparableAliasSamplerRange { + start: range_from, + ledge: range_from, + end: range_from, + } + } else { + // sort the alias sampling atoms to push self-dispersal to the right + let mut atoms = AliasMethodSamplerAtom::create(&event_weights); + atoms.sort_by_key(|a| { + usize::from(*a.e() == row_index) + usize::from(*a.k() == row_index) + }); + + // find the index amongst the atoms of the first atom that includes + // self-dispersal, either with u < 1.0 (uniquely) or u = 1.0 (iff + // no self-dispersal with u < 1.0 exists) + let ledge = match atoms.binary_search_by_key(&1, |a| { + usize::from(*a.e() == row_index) + usize::from(*a.k() == row_index) + }) { + Ok(i) | Err(i) => i, + }; + + alias_dispersal_buffer.append(&mut atoms); + + SeparableAliasSamplerRange { + start: range_from, + ledge: range_from + ledge, + end: alias_dispersal_buffer.len(), + } + } + }), + usize::from(habitat_extent.height()), + usize::from(habitat_extent.width()), + ) + .unwrap(); // infallible by PRE; + + Self { + alias_dispersal_ranges: Final::new(alias_dispersal_ranges), + alias_dispersal_buffer: Final::new(alias_dispersal_buffer.into_boxed_slice()), + marker: PhantomData::<(M, H, G)>, + } + } +} + +impl, G: RngCore> core::fmt::Debug + for InMemoryPackedSeparableAliasDispersalSampler +{ + fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result { + f.debug_struct(stringify!(InMemoryPackedSeparableAliasDispersalSampler)) + .field("alias_dispersal_ranges", &self.alias_dispersal_ranges) + .field( + "alias_dispersal_buffer", + &format_args!( + "Box [ {:p}; {} ]", + &self.alias_dispersal_buffer, + self.alias_dispersal_buffer.len() + ), + ) + .finish() + } +} + +#[contract_trait] +impl, G: RngCore> Backup + for InMemoryPackedSeparableAliasDispersalSampler +{ + unsafe fn backup_unchecked(&self) -> Self { + Self { + alias_dispersal_ranges: Final::new(self.alias_dispersal_ranges.clone()), + alias_dispersal_buffer: Final::new(self.alias_dispersal_buffer.clone()), + marker: PhantomData::<(M, H, G)>, + } + } +}