(ml5717) Added HabitatToU64Injection

Author: juntyr

necsim-core/src/cogs/habitat.rs

@@ -27,3 +27,17 @@ pub trait Habitat: core::fmt::Debug {
     #[debug_requires(self.get_extent().contains(location), "location is inside habitat extent")]
     fn get_habitat_at_location(&self, location: &Location) -> u32;
 }
+
+#[allow(clippy::inline_always, clippy::inline_fn_without_body)]
+#[allow(clippy::module_name_repetitions)]
+#[contract_trait]
+pub trait HabitatToU64Injection: Habitat {
+    #[must_use]
+    #[debug_requires(self.get_extent().contains(location), "location is inside habitat extent")]
+    #[debug_requires(index_at_location < self.get_habitat_at_location(location) as usize)]
+    fn map_indexed_location_to_u64_injective(
+        &self,
+        location: &Location,
+        index_at_location: usize,
+    ) -> u64;
+}

necsim-core/src/cogs/mod.rs

@@ -1,5 +1,5 @@
 mod habitat;
-pub use habitat::Habitat;
+pub use habitat::{Habitat, HabitatToU64Injection};
 
 mod rng;
 pub use rng::{PrimeableRng, RngCore, RngSampler};

necsim-impls-no-std/src/cogs/active_lineage_sampler/independent/mod.rs

@@ -1,7 +1,7 @@
 use core::marker::PhantomData;
 
 use necsim_core::cogs::{
-    DispersalSampler, Habitat, IncoherentLineageStore, LineageReference, PrimeableRng,
+    DispersalSampler, HabitatToU64Injection, IncoherentLineageStore, LineageReference, PrimeableRng,
 };
 
 mod sampler;
@@ -15,7 +15,7 @@ mod sampler;
 #[cfg_attr(feature = "cuda", r2cBound(S: rust_cuda::common::RustToCuda))]
 #[derive(Debug)]
 pub struct IndependentActiveLineageSampler<
-    H: Habitat,
+    H: HabitatToU64Injection,
     G: PrimeableRng<Prime = [u8; 16]>,
     D: DispersalSampler<H, G>,
     R: LineageReference<H>,
@@ -27,7 +27,7 @@ pub struct IndependentActiveLineageSampler<
 }
 
 impl<
-        H: Habitat,
+        H: HabitatToU64Injection,
         G: PrimeableRng<Prime = [u8; 16]>,
         D: DispersalSampler<H, G>,
         R: LineageReference<H>,
@@ -46,7 +46,7 @@ impl<
 }
 
 impl<
-        H: Habitat,
+        H: HabitatToU64Injection,
         G: PrimeableRng<Prime = [u8; 16]>,
         D: DispersalSampler<H, G>,
         R: LineageReference<H>,

necsim-impls-no-std/src/cogs/active_lineage_sampler/independent/sampler.rs

@@ -1,8 +1,8 @@
 use float_next_after::NextAfter;
 
 use necsim_core::cogs::{
-    ActiveLineageSampler, DispersalSampler, Habitat, IncoherentLineageStore, LineageReference,
-    PrimeableRng,
+    ActiveLineageSampler, DispersalSampler, HabitatToU64Injection, IncoherentLineageStore,
+    LineageReference, PrimeableRng,
 };
 use necsim_core::intrinsics::{exp, floor};
 use necsim_core::landscape::Location;
@@ -15,7 +15,7 @@ use super::IndependentActiveLineageSampler;
 
 #[contract_trait]
 impl<
-        H: Habitat,
+        H: HabitatToU64Injection,
         G: PrimeableRng<Prime = [u8; 16]>,
         D: DispersalSampler<H, G>,
         R: LineageReference<H>,
@@ -88,23 +88,21 @@ impl<
         let mut time_step = floor(time / delta_t) as u64 + 1;
 
         loop {
-            let location_bytes_x = lineage_location.x().to_le_bytes();
-            let location_bytes_y = lineage_location.y().to_le_bytes();
-            let location_bytes_i = lineage_index_at_location.to_le_bytes();
+            let location_bytes = simulation
+                .habitat
+                .map_indexed_location_to_u64_injective(&lineage_location, lineage_index_at_location)
+                .to_le_bytes();
             let time_step_bytes = time_step.to_le_bytes();
 
-            // TODO: This should be more automatic, similar to Hash -> it should also NOT use simple xor
-            // TODO: It would be much better if we could get a location index here, i.e. bijection from
-            //       u64 to Location & index_at_location (this would require a cumulative habitat)
             rng.prime_with([
-                location_bytes_x[0] ^ location_bytes_i[7],
-                location_bytes_x[1] ^ location_bytes_i[6],
-                location_bytes_x[2] ^ location_bytes_i[5],
-                location_bytes_x[3] ^ location_bytes_i[4],
-                location_bytes_y[0] ^ location_bytes_i[3],
-                location_bytes_y[1] ^ location_bytes_i[2],
-                location_bytes_y[2] ^ location_bytes_i[1],
-                location_bytes_y[3] ^ location_bytes_i[0],
+                location_bytes[0],
+                location_bytes[1],
+                location_bytes[2],
+                location_bytes[3],
+                location_bytes[4],
+                location_bytes[5],
+                location_bytes[6],
+                location_bytes[7],
                 time_step_bytes[0],
                 time_step_bytes[1],
                 time_step_bytes[2],

necsim-impls-no-std/src/cogs/habitat/in_memory.rs

@@ -1,6 +1,6 @@
 use array2d::Array2D;
 
-use necsim_core::cogs::Habitat;
+use necsim_core::cogs::{Habitat, HabitatToU64Injection};
 use necsim_core::landscape::{LandscapeExtent, Location};
 
 #[allow(clippy::module_name_repetitions)]
@@ -9,6 +9,8 @@ use necsim_core::landscape::{LandscapeExtent, Location};
 pub struct InMemoryHabitat {
     #[cfg_attr(feature = "cuda", r2cEmbed)]
     habitat: Array2D<u32>,
+    #[cfg_attr(feature = "cuda", r2cEmbed)]
+    u64_injection: Array2D<u64>,
 }
 
 #[contract_trait]
@@ -38,6 +40,19 @@ impl Habitat for InMemoryHabitat {
     }
 }
 
+#[contract_trait]
+impl HabitatToU64Injection for InMemoryHabitat {
+    #[must_use]
+    fn map_indexed_location_to_u64_injective(
+        &self,
+        location: &Location,
+        index_at_location: usize,
+    ) -> u64 {
+        self.u64_injection[(location.y() as usize, location.x() as usize)]
+            + index_at_location as u64
+    }
+}
+
 impl InMemoryHabitat {
     #[must_use]
     #[debug_ensures(
@@ -46,7 +61,22 @@ impl InMemoryHabitat {
         "habitat extent has the dimension of the habitat array"
     )]
     pub fn new(habitat: Array2D<u32>) -> Self {
-        #[allow(clippy::cast_possible_truncation)]
-        Self { habitat }
+        let mut index_acc = 0_u64;
+
+        let u64_injection = Array2D::from_iter_row_major(
+            habitat.elements_row_major_iter().map(|h| {
+                let injection = index_acc;
+                index_acc += u64::from(*h);
+                injection
+            }),
+            habitat.num_rows(),
+            habitat.num_columns(),
+        )
+        .unwrap();
+
+        Self {
+            habitat,
+            u64_injection,
+        }
     }
 }