Fix SeparableAliasSelfDispersal CUDA repr

juntyr · Jun 3, 2024 · 7ae8180 · 7ae8180
1 parent a58fbfa
commit 7ae8180
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 18 deletions.
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
@@ -23,7 +23,7 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> DispersalSampler<M, H, G>
 
         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;
-        let location_index =
+        let self_dispersal_index =
             location_row * usize::from(habitat.get_extent().width()) + location_column;
 
         // Only safe as trait precondition that `location` is inside `habitat`
@@ -41,13 +41,13 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> DispersalSampler<M, H, G>
         let dispersal_target_index: usize = if self_dispersal.self_dispersal == ClosedUnitF64::one()
         {
             // guaranteed self-dispersal
-            location_index
+            self_dispersal_index
         } else if (
             // guaranteed non-self-dispersal
             self_dispersal.self_dispersal == ClosedUnitF64::zero()
         ) || (
             // self-dispersal with an underfull atom, so included
-            self_dispersal.non_self_dispersal_event.is_some()
+            self_dispersal.non_self_dispersal_event != self_dispersal_index
         ) || (
             // excluded self-dispersal, but we sampled non-self-dispersal
             rng.sample_uniform_closed_open() >= self_dispersal.self_dispersal
@@ -62,7 +62,7 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> DispersalSampler<M, H, G>
             AliasMethodSamplerAtom::sample_event(alias_dispersals_at_location, rng)
         } else {
             // excluded self-dispersal, and we sampled it
-            location_index
+            self_dispersal_index
         };
 
         #[allow(clippy::cast_possible_truncation)]
@@ -90,7 +90,7 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> SeparableDispersalSampler<M, H,
     ) -> 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;
-        let location_index =
+        let self_dispersal_index =
             location_row * usize::from(habitat.get_extent().width()) + location_column;
 
         // Only safe as trait precondition that `location` is inside `habitat`
@@ -120,7 +120,7 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> SeparableDispersalSampler<M, H,
             // if non_self_dispersal_event is None, self-dispersal is already
             //  excluded from the alias sampler and so we can sample the full
             //  CDF
-            if self_dispersal.non_self_dispersal_event.is_none() {
+            if self_dispersal.non_self_dispersal_event == self_dispersal_index {
                 ClosedUnitF64::one()
             } else {
                 self_dispersal.self_dispersal.one_minus()
@@ -129,8 +129,8 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> SeparableDispersalSampler<M, H,
 
         // if rounding errors caused self-dispersal, replace with the non-self-dispersal
         //  event
-        if dispersal_target_index == location_index {
-            dispersal_target_index = self_dispersal.non_self_dispersal_event.unwrap();
+        if dispersal_target_index == self_dispersal_index {
+            dispersal_target_index = self_dispersal.non_self_dispersal_event;
         }
 
         #[allow(clippy::cast_possible_truncation)]

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
@@ -38,9 +38,9 @@ pub struct SeparableAliasSelfDispersal {
     self_dispersal: ClosedUnitF64,
     // non-self-dispersal event to sample in case rounding errors cause
     //  self-dispersal to be sampled in no-self-dispersal mode
-    // if `Some(x)`, then x is the event
-    // if `None`, then self-dispersal is not part of the alias sampler
-    non_self_dispersal_event: Option<usize>,
+    // if it is set to self-dispersal, then it marks self-dispersal not
+    //  being part of the alias sampler
+    non_self_dispersal_event: usize,
 }
 
 #[allow(clippy::module_name_repetitions)]
@@ -74,14 +74,16 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> InMemoryDispersalSampler<M, H,
 
         let mut alias_dispersal_buffer = Vec::new();
 
-        let mut self_dispersal = VecArray2D::filled_with(
-            SeparableAliasSelfDispersal {
+        let mut self_dispersal = VecArray2D::from_iter_row_major(
+            (0..).map(|self_dispersal_event| SeparableAliasSelfDispersal {
                 self_dispersal: ClosedUnitF64::zero(),
-                non_self_dispersal_event: None,
-            },
+                // no self-dispersal
+                non_self_dispersal_event: self_dispersal_event,
+            }),
             usize::from(habitat_extent.height()),
             usize::from(habitat_extent.width()),
-        );
+        )
+        .unwrap();
 
         let alias_dispersal_ranges = Array2D::from_iter_row_major(
             dispersal.rows_iter().enumerate().map(|(row_index, row)| {
@@ -140,7 +142,7 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> InMemoryDispersalSampler<M, H,
                         / total_weight;
 
                     let mut atoms;
-                    let mut non_self_dispersal_event = None;
+                    let mut non_self_dispersal_event = row_index;
 
                     if self_dispersal_u < 1.0 {
                         atoms = AliasMethodSamplerAtom::create(&event_weights);
@@ -160,7 +162,7 @@ impl<M: MathsCore, H: Habitat<M>, G: RngCore<M>> InMemoryDispersalSampler<M, H,
                             if self_dispersal_atom.e() == row_index {
                                 self_dispersal_atom.flip();
                             }
-                            non_self_dispersal_event = Some(self_dispersal_atom.e());
+                            non_self_dispersal_event = self_dispersal_atom.e();
                             let last_atom_index = atoms.len() - 1;
                             atoms.swap(self_dispersal_index, last_atom_index);
                         };