compb-dla-model/src/system/spaces/nalg.rs
2023-03-19 10:26:33 +00:00

149 lines
3.8 KiB
Rust

use std::ops::Add;
use itertools::Itertools;
use nalgebra::{EuclideanNorm, LpNorm, Matrix, Norm, OMatrix, SVector};
use num_traits::Pow;
use serde::{Serialize, Deserialize};
use crate::system::{GriddedPosition, Position, Storage};
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Gridded<const D: usize>(SVector<i32, D>);
#[derive(Clone, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Continuous<const D: usize>(SVector<f32, D>);
impl<const D: usize> Add for Continuous<D> {
type Output = Continuous<D>;
fn add(self, rhs: Self) -> Self::Output {
Continuous(self.0 + rhs.0)
}
}
impl<const D: usize> Position for Continuous<D> {
const DIM: usize = 0;
fn zero() -> Self {
Continuous(SVector::<f32, D>::zeros())
}
fn abs(&self) -> f32 {
self.0.norm()
}
fn from_cartesian(cartesian: &[f32]) -> Self {
Continuous(SVector::<f32, D>::from_fn(|i, _| cartesian[i]))
}
fn to_cartesian(&self) -> Vec<f32> {
self.0.as_slice().to_vec()
}
}
impl<const D: usize> Add for Gridded<D> {
type Output = Gridded<D>;
fn add(self, rhs: Self) -> Self::Output {
Gridded(self.0 + rhs.0)
}
}
impl<const D: usize> Position for Gridded<D> {
const DIM: usize = 0;
fn zero() -> Self {
Gridded(SVector::<i32, D>::zeros())
}
fn abs(&self) -> f32 {
(self.0.fold(0, |r, c| r + c.pow(2)) as f32).sqrt()
}
fn from_cartesian(cartesian: &[f32]) -> Self {
Gridded(SVector::<i32, D>::from_fn(|i, _| cartesian[i] as i32))
}
fn to_cartesian(&self) -> Vec<f32> {
self.0.as_slice()
.iter()
.map(|a| *a as f32)
.collect_vec()
}
}
pub struct KDSpace<const N: usize>(pub(crate) kiddo::KdTree<f32, (), N>);
impl<const D: usize> Storage<Gridded<D>> for KDSpace<D> {
fn is_occupied(&self, position: &Gridded<D>) -> bool {
let a = self.0.best_n_within(
&position.0.data.0[0].map(|i| i as f32),
0f32,
1,
&|a, b| {
LpNorm(1).metric_distance(
&SVector::<f32, D>::from_row_slice(a),
&SVector::<f32, D>::from_row_slice(b),
)
},
).unwrap();
!a.is_empty()
}
fn deposit(&mut self, position: &Gridded<D>) {
self.0.add(&position.0.data.0[0].map(|i| i as f32), ())
.expect("Failed to write to space")
}
}
impl<const D: usize> Storage<Continuous<D>> for KDSpace<D> {
fn is_occupied(&self, position: &Continuous<D>) -> bool {
let a = self.0.best_n_within(
&position.0.data.0[0],
0f32,
1,
&|a, b| {
EuclideanNorm.metric_distance(
&SVector::<f32, D>::from_row_slice(a),
&SVector::<f32, D>::from_row_slice(b),
)
},
).unwrap();
!a.is_empty()
}
fn deposit(&mut self, position: &Continuous<D>) {
self.0.add(&position.0.data.0[0], ())
.expect("Failed to write to space")
}
}
pub struct VectorStorage {
backing: Vec<bool>,
grid_size: usize,
}
impl<const D: usize> Storage<Gridded<D>> for VectorStorage {
fn is_occupied(&self, position: &Gridded<D>) -> bool {
let mut index: usize = 0;
for i in 0..D {
index += (position.0[i] + (self.grid_size as i32 / 2)) as usize * self.grid_size * i;
}
return self.backing[index];
}
fn deposit(&mut self, position: &Gridded<D>) {
let mut index: usize = 0;
for i in 0..D {
index += (position.0[i] + (self.grid_size as i32 / 2)) as usize * self.grid_size * i;
}
self.backing[index] = true;
}
}