From 9f74ec2ef03d4996f1074456dd93bd529df06427 Mon Sep 17 00:00:00 2001
From: Joshua Coles <josh@coles.to>
Date: Thu, 9 Mar 2023 08:13:38 +0000
Subject: [PATCH] KD impl 1

---
 src/cli/cli.rs               |  1 +
 src/main.rs                  | 15 ++++++++
 src/system/spaces/kd_grid.rs | 74 ++++++++++++++++++++++++++++++++++++
 src/system/spaces/mod.rs     |  1 +
 src/system/sticker.rs        |  1 -
 5 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100644 src/system/spaces/kd_grid.rs

diff --git a/src/cli/cli.rs b/src/cli/cli.rs
index 4ad7d3d..a2246b4 100644
--- a/src/cli/cli.rs
+++ b/src/cli/cli.rs
@@ -37,6 +37,7 @@ pub struct SurfaceProbabilityMeasureCli {
 pub enum PCM {
     Initial(InitialCli),
     StickProbability(StickProbabilityCli),
+    Grid3KD(StickProbabilityCli),
     Grid3(StickProbabilityCli),
     Hex(StickProbabilityCli),
     Balls(BallsCli),
diff --git a/src/main.rs b/src/main.rs
index 0b5b9c5..ba94bc3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -12,6 +12,7 @@ use crate::surface_probability_measure::{LoggerSticker, ReadOnlyVectorStorage};
 use crate::system::model::DLASystem;
 use crate::system::spaces::continuous::{ContinuousSticker, ContinuousStorage, ContinuousWalker};
 use crate::system::spaces::hexagonal::HexPosition;
+use crate::system::spaces::kd_grid::{KDGrid, KDGrid2Sticker};
 use crate::system::spaces::square_grid::{Grid2D, Grid3D};
 use crate::system::spaces::VectorStorage;
 use crate::system::spawner::UniformSpawner;
@@ -56,6 +57,20 @@ fn main() {
             write(&sys, cli.format, &cli.output);
         }
 
+        PCM::Grid3KD(StickProbabilityCli { grid_size, stick_probability }) => {
+            let mut sys = DLASystem::<_, Grid2D, _, _, _, _>::new(
+                SmallRng::seed_from_u64(cli.seed),
+                KDGrid { inner: kiddo::KdTree::new() },
+                LocalRandomWalker,
+                UniformSpawner,
+                KDGrid2Sticker { stick_probability },
+                cli.max_particles,
+            );
+
+            drive_system(&mut sys, cli.max_frames, cli.notify_every);
+            write(&sys, cli.format, &cli.output);
+        }
+
         PCM::Grid3(StickProbabilityCli { grid_size, stick_probability }) => {
             let mut sys = DLASystem::<_, Grid3D, _, _, _, _>::new(
                 SmallRng::seed_from_u64(cli.seed),
diff --git a/src/system/spaces/kd_grid.rs b/src/system/spaces/kd_grid.rs
new file mode 100644
index 0000000..f483308
--- /dev/null
+++ b/src/system/spaces/kd_grid.rs
@@ -0,0 +1,74 @@
+use rand::Rng;
+use crate::system::{GriddedPosition, Position, Storage};
+use crate::system::spaces::square_grid::{Grid2D, Grid3D};
+use crate::system::sticker::Sticker;
+
+fn taxicab_grid2(a: &[f32; 2], b: &[f32; 2]) -> f32 {
+    (a[0] - b[0]).abs() + (a[1] - b[1]).abs()
+}
+
+fn taxicab_grid3(a: &[f32; 3], b: &[f32; 3]) -> f32 {
+    (a[0] - b[0]).abs() + (a[1] - b[1]).abs() + (a[2] - b[2]).abs()
+}
+
+pub struct KDGrid<const N: usize> {
+    pub(crate) inner: kiddo::KdTree<f32, (), N>,
+}
+
+impl Storage<Grid2D> for KDGrid<2> {
+    fn is_occupied(&self, position: &Grid2D) -> bool {
+        let (dist_sq, _) = self.inner.nearest_one(&[position.x as f32, position.y as f32], &taxicab_grid2).unwrap();
+        dist_sq == 0f32
+    }
+
+    fn deposit(&mut self, position: &Grid2D) {
+        self.inner.add(&[position.x as f32, position.y as f32], ())
+            .expect("Failed to write to space")
+    }
+}
+
+impl Storage<Grid3D> for KDGrid<3> {
+    fn is_occupied(&self, position: &Grid3D) -> bool {
+        let (dist_sq, _) = self.inner.nearest_one(&[position.x as f32, position.y as f32, position.z as f32], &taxicab_grid3).unwrap();
+        dist_sq == 0f32
+    }
+
+    fn deposit(&mut self, position: &Grid3D) {
+        self.inner.add(&[position.x as f32, position.y as f32, position.z as f32], ())
+            .expect("Failed to write to space")
+    }
+}
+
+
+pub struct KDGrid2Sticker {
+    pub(crate) stick_probability: f32,
+}
+
+impl Sticker<Grid2D, KDGrid<2>> for KDGrid2Sticker {
+    fn should_stick<R: Rng>(&self, rng: &mut R, space: &KDGrid<2>, position: &Grid2D) -> bool {
+        let a = space.inner.nearest(&[position.x as f32, position.y as f32], Grid2D::NEIGHBOURS as usize, &taxicab_grid2)
+            .unwrap();
+
+        if a.len() == 0 {
+            return false;
+        }
+
+        let b = a.iter().filter(|(a, _)| *a <= 1f32)
+            .count();
+
+        let q = (1f32 - self.stick_probability);
+        let a = q.powi(b as i32);
+        rng.gen_range(0f32..1f32) > a
+    }
+}
+
+impl Sticker<Grid3D, KDGrid<3>> for KDGrid2Sticker {
+    fn should_stick<R: Rng>(&self, rng: &mut R, space: &KDGrid<3>, position: &Grid3D) -> bool {
+        let a = space.inner.nearest(&[position.x as f32, position.y as f32, position.z as f32], Grid3D::NEIGHBOURS as usize, &taxicab_grid3)
+            .unwrap();
+
+        let q = (1f32 - self.stick_probability);
+        let a = q.powi(a.len() as i32);
+        rng.gen_range(0f32..1f32) > a
+    }
+}
diff --git a/src/system/spaces/mod.rs b/src/system/spaces/mod.rs
index b62ff55..870fc9d 100644
--- a/src/system/spaces/mod.rs
+++ b/src/system/spaces/mod.rs
@@ -2,6 +2,7 @@ pub mod vector_storage;
 pub use vector_storage::VectorStorage;
 
 pub mod square_grid;
+pub mod kd_grid;
 pub mod hexagonal;
 
 pub mod continuous;
diff --git a/src/system/sticker.rs b/src/system/sticker.rs
index 0136293..e096100 100644
--- a/src/system/sticker.rs
+++ b/src/system/sticker.rs
@@ -26,4 +26,3 @@ impl<P: GriddedPosition, S: Storage<P>> Sticker<P, S> for ProbabilisticSticking
             .any(|neighbour| space.is_occupied(&neighbour) && rng.gen_range(0.0f32..=1.0) < self.stick_probability)
     }
 }
-