391 lines
11 KiB
Rust
391 lines
11 KiB
Rust
//! K-D Tree implementation for spatial queries.
|
|
//!
|
|
//! This module provides an efficient data structure for nearest-neighbor
|
|
//! searches in 3D space, used for mapping MRE properties to mesh nodes.
|
|
|
|
use nalgebra::{Point3, Vector3};
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// A node in the K-D tree.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
struct KdNode {
|
|
/// The point stored at this node.
|
|
point: Point3<f64>,
|
|
/// Index of this point in the original data.
|
|
index: usize,
|
|
/// Left child (points with coordinate < split value).
|
|
left: Option<Box<KdNode>>,
|
|
/// Right child (points with coordinate >= split value).
|
|
right: Option<Box<KdNode>>,
|
|
/// Split dimension (0=x, 1=y, 2=z).
|
|
split_dim: usize,
|
|
}
|
|
|
|
/// A 3D K-D Tree for efficient nearest-neighbor queries.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct KdTree3D {
|
|
root: Option<Box<KdNode>>,
|
|
size: usize,
|
|
}
|
|
|
|
impl KdTree3D {
|
|
/// Build a K-D tree from a set of points.
|
|
pub fn build(points: &[Point3<f64>]) -> Self {
|
|
if points.is_empty() {
|
|
return Self {
|
|
root: None,
|
|
size: 0,
|
|
};
|
|
}
|
|
|
|
let indices: Vec<usize> = (0..points.len()).collect();
|
|
let root = Self::build_recursive(points, &indices, 0);
|
|
|
|
Self {
|
|
root: Some(root),
|
|
size: points.len(),
|
|
}
|
|
}
|
|
|
|
fn build_recursive(points: &[Point3<f64>], indices: &[usize], depth: usize) -> Box<KdNode> {
|
|
let dim = depth % 3;
|
|
|
|
// Sort indices by the current dimension
|
|
let mut sorted_indices = indices.to_vec();
|
|
sorted_indices.sort_by(|&a, &b| {
|
|
let val_a = points[a].coords[dim];
|
|
let val_b = points[b].coords[dim];
|
|
val_a.partial_cmp(&val_b).unwrap()
|
|
});
|
|
|
|
let mid = sorted_indices.len() / 2;
|
|
let median_idx = sorted_indices[mid];
|
|
|
|
let left = if mid > 0 {
|
|
Some(Self::build_recursive(
|
|
points,
|
|
&sorted_indices[..mid],
|
|
depth + 1,
|
|
))
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let right = if mid + 1 < sorted_indices.len() {
|
|
Some(Self::build_recursive(
|
|
points,
|
|
&sorted_indices[mid + 1..],
|
|
depth + 1,
|
|
))
|
|
} else {
|
|
None
|
|
};
|
|
|
|
Box::new(KdNode {
|
|
point: points[median_idx],
|
|
index: median_idx,
|
|
left,
|
|
right,
|
|
split_dim: dim,
|
|
})
|
|
}
|
|
|
|
/// Get the number of points in the tree.
|
|
pub fn len(&self) -> usize {
|
|
self.size
|
|
}
|
|
|
|
/// Check if the tree is empty.
|
|
pub fn is_empty(&self) -> bool {
|
|
self.size == 0
|
|
}
|
|
|
|
/// Find the nearest neighbor to a query point.
|
|
///
|
|
/// Returns (index, distance) of the nearest point.
|
|
pub fn nearest(&self, query: &Point3<f64>) -> Option<(usize, f64)> {
|
|
let root = self.root.as_ref()?;
|
|
let mut best = (root.index, (root.point - query).norm());
|
|
Self::nearest_recursive(root, query, &mut best);
|
|
Some(best)
|
|
}
|
|
|
|
fn nearest_recursive(node: &KdNode, query: &Point3<f64>, best: &mut (usize, f64)) {
|
|
let dist = (node.point - query).norm();
|
|
if dist < best.1 {
|
|
*best = (node.index, dist);
|
|
}
|
|
|
|
let dim = node.split_dim;
|
|
let diff = query.coords[dim] - node.point.coords[dim];
|
|
|
|
// Determine which child to search first
|
|
let (first, second) = if diff < 0.0 {
|
|
(&node.left, &node.right)
|
|
} else {
|
|
(&node.right, &node.left)
|
|
};
|
|
|
|
// Search the closer subtree first
|
|
if let Some(child) = first {
|
|
Self::nearest_recursive(child, query, best);
|
|
}
|
|
|
|
// Only search the other subtree if it could contain a closer point
|
|
if diff.abs() < best.1 {
|
|
if let Some(child) = second {
|
|
Self::nearest_recursive(child, query, best);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Find all points within a given radius of a query point.
|
|
///
|
|
/// Returns a vector of (index, distance) pairs.
|
|
pub fn within_radius(&self, query: &Point3<f64>, radius: f64) -> Vec<(usize, f64)> {
|
|
let mut results = Vec::new();
|
|
if let Some(root) = &self.root {
|
|
Self::within_radius_recursive(root, query, radius, &mut results);
|
|
}
|
|
results
|
|
}
|
|
|
|
fn within_radius_recursive(
|
|
node: &KdNode,
|
|
query: &Point3<f64>,
|
|
radius: f64,
|
|
results: &mut Vec<(usize, f64)>,
|
|
) {
|
|
let dist = (node.point - query).norm();
|
|
if dist <= radius {
|
|
results.push((node.index, dist));
|
|
}
|
|
|
|
let dim = node.split_dim;
|
|
let diff = query.coords[dim] - node.point.coords[dim];
|
|
|
|
// Check if we need to search left subtree
|
|
if diff - radius <= 0.0 {
|
|
if let Some(left) = &node.left {
|
|
Self::within_radius_recursive(left, query, radius, results);
|
|
}
|
|
}
|
|
|
|
// Check if we need to search right subtree
|
|
if diff + radius >= 0.0 {
|
|
if let Some(right) = &node.right {
|
|
Self::within_radius_recursive(right, query, radius, results);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Find the k nearest neighbors to a query point.
|
|
///
|
|
/// Returns a vector of (index, distance) pairs, sorted by distance.
|
|
pub fn k_nearest(&self, query: &Point3<f64>, k: usize) -> Vec<(usize, f64)> {
|
|
if k == 0 || self.is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
let mut results: Vec<(usize, f64)> = Vec::with_capacity(k);
|
|
if let Some(root) = &self.root {
|
|
Self::k_nearest_recursive(root, query, k, &mut results);
|
|
}
|
|
|
|
results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
|
results
|
|
}
|
|
|
|
fn k_nearest_recursive(
|
|
node: &KdNode,
|
|
query: &Point3<f64>,
|
|
k: usize,
|
|
results: &mut Vec<(usize, f64)>,
|
|
) {
|
|
let dist = (node.point - query).norm();
|
|
|
|
// Check if we should add this point
|
|
if results.len() < k {
|
|
results.push((node.index, dist));
|
|
results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
|
} else if dist < results.last().unwrap().1 {
|
|
results.pop();
|
|
results.push((node.index, dist));
|
|
results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
|
|
}
|
|
|
|
let dim = node.split_dim;
|
|
let diff = query.coords[dim] - node.point.coords[dim];
|
|
let max_dist = if results.len() < k {
|
|
f64::INFINITY
|
|
} else {
|
|
results.last().unwrap().1
|
|
};
|
|
|
|
// Determine which child to search first
|
|
let (first, second) = if diff < 0.0 {
|
|
(&node.left, &node.right)
|
|
} else {
|
|
(&node.right, &node.left)
|
|
};
|
|
|
|
// Search the closer subtree first
|
|
if let Some(child) = first {
|
|
Self::k_nearest_recursive(child, query, k, results);
|
|
}
|
|
|
|
// Update max distance after searching first subtree
|
|
let max_dist = if results.len() < k {
|
|
f64::INFINITY
|
|
} else {
|
|
results.last().unwrap().1
|
|
};
|
|
|
|
// Only search the other subtree if it could contain a closer point
|
|
if diff.abs() < max_dist {
|
|
if let Some(child) = second {
|
|
Self::k_nearest_recursive(child, query, k, results);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Compute the bounding box of a set of points.
|
|
pub fn bounding_box(points: &[Point3<f64>]) -> Option<(Point3<f64>, Point3<f64>)> {
|
|
if points.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let mut min = points[0];
|
|
let mut max = points[0];
|
|
|
|
for p in points.iter().skip(1) {
|
|
for i in 0..3 {
|
|
if p.coords[i] < min.coords[i] {
|
|
min.coords[i] = p.coords[i];
|
|
}
|
|
if p.coords[i] > max.coords[i] {
|
|
max.coords[i] = p.coords[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
Some((min, max))
|
|
}
|
|
|
|
/// Compute the center of mass of a set of points.
|
|
pub fn center_of_mass(points: &[Point3<f64>]) -> Option<Point3<f64>> {
|
|
if points.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let sum: Vector3<f64> = points.iter().map(|p| p.coords).sum();
|
|
Some(Point3::from(sum / points.len() as f64))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn create_test_points() -> Vec<Point3<f64>> {
|
|
vec![
|
|
Point3::new(0.0, 0.0, 0.0),
|
|
Point3::new(1.0, 0.0, 0.0),
|
|
Point3::new(0.0, 1.0, 0.0),
|
|
Point3::new(0.0, 0.0, 1.0),
|
|
Point3::new(1.0, 1.0, 0.0),
|
|
Point3::new(1.0, 0.0, 1.0),
|
|
Point3::new(0.0, 1.0, 1.0),
|
|
Point3::new(1.0, 1.0, 1.0),
|
|
]
|
|
}
|
|
|
|
#[test]
|
|
fn test_build_tree() {
|
|
let points = create_test_points();
|
|
let tree = KdTree3D::build(&points);
|
|
|
|
assert_eq!(tree.len(), 8);
|
|
assert!(!tree.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_empty_tree() {
|
|
let tree = KdTree3D::build(&[]);
|
|
assert!(tree.is_empty());
|
|
assert_eq!(tree.len(), 0);
|
|
assert!(tree.nearest(&Point3::new(0.0, 0.0, 0.0)).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_nearest_exact_match() {
|
|
let points = create_test_points();
|
|
let tree = KdTree3D::build(&points);
|
|
|
|
// Query for a point that exists in the tree
|
|
let result = tree.nearest(&Point3::new(1.0, 1.0, 1.0)).unwrap();
|
|
assert_eq!(result.0, 7); // Index of (1,1,1)
|
|
assert!(result.1 < 1e-10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_nearest_interpolated() {
|
|
let points = create_test_points();
|
|
let tree = KdTree3D::build(&points);
|
|
|
|
// Query for a point not in the tree
|
|
let result = tree.nearest(&Point3::new(0.1, 0.1, 0.1)).unwrap();
|
|
assert_eq!(result.0, 0); // Closest to origin
|
|
}
|
|
|
|
#[test]
|
|
fn test_within_radius() {
|
|
let points = create_test_points();
|
|
let tree = KdTree3D::build(&points);
|
|
|
|
// All points at distance <= sqrt(3) from origin
|
|
let results = tree.within_radius(&Point3::new(0.0, 0.0, 0.0), 2.0);
|
|
assert_eq!(results.len(), 8);
|
|
|
|
// Only origin at distance 0
|
|
let results = tree.within_radius(&Point3::new(0.0, 0.0, 0.0), 0.1);
|
|
assert_eq!(results.len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_k_nearest() {
|
|
let points = create_test_points();
|
|
let tree = KdTree3D::build(&points);
|
|
|
|
// Find 3 nearest to origin
|
|
let results = tree.k_nearest(&Point3::new(0.0, 0.0, 0.0), 3);
|
|
assert_eq!(results.len(), 3);
|
|
|
|
// First should be origin
|
|
assert_eq!(results[0].0, 0);
|
|
assert!(results[0].1 < 1e-10);
|
|
|
|
// Results should be sorted by distance
|
|
for i in 1..results.len() {
|
|
assert!(results[i].1 >= results[i - 1].1);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_bounding_box() {
|
|
let points = create_test_points();
|
|
let (min, max) = bounding_box(&points).unwrap();
|
|
|
|
assert!((min - Point3::new(0.0, 0.0, 0.0)).norm() < 1e-10);
|
|
assert!((max - Point3::new(1.0, 1.0, 1.0)).norm() < 1e-10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_center_of_mass() {
|
|
let points = create_test_points();
|
|
let com = center_of_mass(&points).unwrap();
|
|
|
|
assert!((com - Point3::new(0.5, 0.5, 0.5)).norm() < 1e-10);
|
|
}
|
|
}
|