diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 73c5e4c..af32812 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -5,6 +5,8 @@ on: branches: - main pull_request: + branches: + - main jobs: lint-test: diff --git a/.gitignore b/.gitignore index 4fe000d..76b13c0 100644 --- a/.gitignore +++ b/.gitignore @@ -25,3 +25,5 @@ target/ # Added by cargo /target + +.idea/ diff --git a/examples/distance_metrics.rs b/examples/distance_metrics.rs new file mode 100644 index 0000000..bb1deff --- /dev/null +++ b/examples/distance_metrics.rs @@ -0,0 +1,137 @@ +//! Example demonstrating different distance metrics for spatial queries. +//! +//! This example shows how to use Euclidean, Haversine, and Spheroid distance metrics +//! for finding nearest neighbors in spatial datasets. + +use geo_index::rtree::distance::{EuclideanDistance, HaversineDistance, SpheroidDistance}; +use geo_index::rtree::sort::HilbertSort; +use geo_index::rtree::{RTreeBuilder, RTreeIndex}; + +fn main() { + println!("=== Distance Metrics Example ===\n"); + + // Example 1: Euclidean Distance (for planar coordinates) + println!("1. Euclidean Distance (planar coordinates):"); + euclidean_distance_example(); + + // Example 2: Haversine Distance (for geographic coordinates) + println!("\n2. Haversine Distance (geographic coordinates):"); + haversine_distance_example(); + + // Example 3: Spheroid Distance (for high-precision geographic coordinates) + println!("\n3. Spheroid Distance (high-precision geographic coordinates):"); + spheroid_distance_example(); + + // Example 4: Comparison of different metrics + println!("\n4. Comparison of different distance metrics:"); + comparison_example(); +} + +fn euclidean_distance_example() { + let mut builder = RTreeBuilder::::new(4); + + // Add some points in a planar coordinate system + builder.add(0., 0., 1., 1.); // Point A + builder.add(3., 4., 4., 5.); // Point B + builder.add(6., 8., 7., 9.); // Point C + builder.add(1., 1., 2., 2.); // Point D + + let tree = builder.finish::(); + + let euclidean = EuclideanDistance; + let query_point = (0., 0.); + + println!(" Query point: {:?}", query_point); + let results = + tree.neighbors_with_distance(query_point.0, query_point.1, Some(3), None, &euclidean); + + println!(" Nearest neighbors (by insertion order): {:?}", results); + println!(" Distance metric: Euclidean (straight-line distance)"); +} + +fn haversine_distance_example() { + let mut builder = RTreeBuilder::::new(5); + + // Add some major cities (longitude, latitude) + builder.add(-74.0, 40.7, -74.0, 40.7); // New York + builder.add(-0.1, 51.5, -0.1, 51.5); // London + builder.add(139.7, 35.7, 139.7, 35.7); // Tokyo + builder.add(-118.2, 34.1, -118.2, 34.1); // Los Angeles + builder.add(2.3, 48.9, 2.3, 48.9); // Paris + + let tree = builder.finish::(); + + let haversine = HaversineDistance::default(); + let query_point = (-74.0, 40.7); // New York + + println!(" Query point: New York {:?}", query_point); + let results = + tree.neighbors_with_distance(query_point.0, query_point.1, Some(3), None, &haversine); + + println!(" Nearest neighbors (by insertion order): {:?}", results); + println!(" Distance metric: Haversine (great-circle distance on sphere)"); + println!(" Earth radius: {} meters", haversine.earth_radius); +} + +fn spheroid_distance_example() { + let mut builder = RTreeBuilder::::new(5); + + // Add some major cities (longitude, latitude) + builder.add(-74.0, 40.7, -74.0, 40.7); // New York + builder.add(-0.1, 51.5, -0.1, 51.5); // London + builder.add(139.7, 35.7, 139.7, 35.7); // Tokyo + builder.add(-118.2, 34.1, -118.2, 34.1); // Los Angeles + builder.add(2.3, 48.9, 2.3, 48.9); // Paris + + let tree = builder.finish::(); + + let spheroid = SpheroidDistance::default(); // WGS84 ellipsoid + let query_point = (-74.0, 40.7); // New York + + println!(" Query point: New York {:?}", query_point); + let results = + tree.neighbors_with_distance(query_point.0, query_point.1, Some(3), None, &spheroid); + + println!(" Nearest neighbors (by insertion order): {:?}", results); + println!(" Distance metric: Spheroid (distance on ellipsoid)"); + println!(" Semi-major axis: {} meters", spheroid.semi_major_axis); + println!(" Semi-minor axis: {} meters", spheroid.semi_minor_axis); +} + +fn comparison_example() { + let mut builder = RTreeBuilder::::new(3); + + // Add points with different characteristics + builder.add(0., 0., 1., 1.); // Origin + builder.add(1., 1., 2., 2.); // Close point + builder.add(10., 10., 11., 11.); // Distant point + + let tree = builder.finish::(); + + let query_point = (0., 0.); + + // Test with different distance metrics + let euclidean = EuclideanDistance; + let haversine = HaversineDistance::default(); + let spheroid = SpheroidDistance::default(); + + println!(" Query point: {:?}", query_point); + + let euclidean_results = + tree.neighbors_with_distance(query_point.0, query_point.1, Some(2), None, &euclidean); + println!(" Euclidean results: {:?}", euclidean_results); + + let haversine_results = + tree.neighbors_with_distance(query_point.0, query_point.1, Some(2), None, &haversine); + println!(" Haversine results: {:?}", haversine_results); + + let spheroid_results = + tree.neighbors_with_distance(query_point.0, query_point.1, Some(2), None, &spheroid); + println!(" Spheroid results: {:?}", spheroid_results); + + // Test backward compatibility + let original_results = tree.neighbors(query_point.0, query_point.1, Some(2), None); + println!(" Original method results: {:?}", original_results); + + println!(" Note: For small distances, all metrics should give similar ordering."); +} diff --git a/src/kdtree/index.rs b/src/kdtree/index.rs index 84a192f..d76a1fa 100644 --- a/src/kdtree/index.rs +++ b/src/kdtree/index.rs @@ -62,7 +62,7 @@ impl KDTreeMetadata { let version = version_and_type >> 4; if version != KDBUSH_VERSION { return Err(GeoIndexError::General( - format!("Got v{} data when expected v{}.", version, KDBUSH_VERSION).to_string(), + format!("Got v{version} data when expected v{KDBUSH_VERSION}.").to_string(), )); } diff --git a/src/rtree/distance.rs b/src/rtree/distance.rs new file mode 100644 index 0000000..bcb53dd --- /dev/null +++ b/src/rtree/distance.rs @@ -0,0 +1,332 @@ +//! Distance metrics for spatial queries. +//! +//! This module provides different distance calculation methods for spatial queries, +//! including Euclidean, Haversine, and Spheroid distance calculations. + +use crate::r#type::IndexableNum; +use std::f64::consts::PI; + +/// A trait for calculating distances between two points. +pub trait DistanceMetric { + /// Calculate the distance between two points (x1, y1) and (x2, y2). + fn distance(&self, x1: N, y1: N, x2: N, y2: N) -> N; + + /// Calculate the distance from a point to a bounding box. + /// This is used for spatial index optimization. + fn distance_to_bbox(&self, x: N, y: N, min_x: N, min_y: N, max_x: N, max_y: N) -> N; + + /// Return the maximum distance value for this metric. + fn max_distance(&self) -> N { + N::max_value() + } +} + +/// Euclidean distance metric. +/// +/// This is the standard straight-line distance calculation suitable for +/// planar coordinate systems. When working with longitude/latitude coordinates, +/// the unit of distance will be degrees. +#[derive(Debug, Clone, Copy, Default)] +pub struct EuclideanDistance; + +impl DistanceMetric for EuclideanDistance { + #[inline] + fn distance(&self, x1: N, y1: N, x2: N, y2: N) -> N { + let dx = x1 - x2; + let dy = y1 - y2; + (dx * dx + dy * dy).sqrt().unwrap_or(N::max_value()) + } + + #[inline] + fn distance_to_bbox(&self, x: N, y: N, min_x: N, min_y: N, max_x: N, max_y: N) -> N { + let dx = axis_dist(x, min_x, max_x); + let dy = axis_dist(y, min_y, max_y); + (dx * dx + dy * dy).sqrt().unwrap_or(N::max_value()) + } +} + +/// Haversine distance metric. +/// +/// This calculates the great-circle distance between two points on a sphere. +/// It's more accurate for geographic distances than Euclidean distance. +/// The input coordinates should be in longitude/latitude (degrees), and +/// the output distance is in meters. +#[derive(Debug, Clone, Copy)] +pub struct HaversineDistance { + /// Earth's radius in meters + pub earth_radius: f64, +} + +impl Default for HaversineDistance { + fn default() -> Self { + Self { + earth_radius: 6378137.0, // WGS84 equatorial radius in meters + } + } +} + +impl HaversineDistance { + /// Create a new Haversine distance metric with custom Earth radius. + pub fn with_radius(earth_radius: f64) -> Self { + Self { earth_radius } + } +} + +impl DistanceMetric for HaversineDistance { + fn distance(&self, lon1: N, lat1: N, lon2: N, lat2: N) -> N { + let lat1_rad = lat1.to_f64().unwrap_or(0.0) * PI / 180.0; + let lat2_rad = lat2.to_f64().unwrap_or(0.0) * PI / 180.0; + let delta_lat = (lat2.to_f64().unwrap_or(0.0) - lat1.to_f64().unwrap_or(0.0)) * PI / 180.0; + let delta_lon = (lon2.to_f64().unwrap_or(0.0) - lon1.to_f64().unwrap_or(0.0)) * PI / 180.0; + + let a = (delta_lat / 2.0).sin().powi(2) + + lat1_rad.cos() * lat2_rad.cos() * (delta_lon / 2.0).sin().powi(2); + let c = 2.0 * a.sqrt().atan2((1.0 - a).sqrt()); + + N::from_f64(self.earth_radius * c).unwrap_or(N::max_value()) + } + + fn distance_to_bbox( + &self, + lon: N, + lat: N, + min_lon: N, + min_lat: N, + max_lon: N, + max_lat: N, + ) -> N { + // For bbox distance with Haversine, we approximate using the closest point on the bbox + let closest_lon = if lon < min_lon { + min_lon + } else if lon > max_lon { + max_lon + } else { + lon + }; + + let closest_lat = if lat < min_lat { + min_lat + } else if lat > max_lat { + max_lat + } else { + lat + }; + + self.distance(lon, lat, closest_lon, closest_lat) + } +} + +/// Spheroid distance metric. +/// +/// This calculates the shortest distance between two points on the surface +/// of a spheroid (ellipsoid), providing a more accurate Earth model than +/// a simple sphere. The input coordinates should be in longitude/latitude +/// (degrees), and the output distance is in meters. +#[derive(Debug, Clone, Copy)] +pub struct SpheroidDistance { + /// Semi-major axis (equatorial radius) in meters + pub semi_major_axis: f64, + /// Semi-minor axis (polar radius) in meters + pub semi_minor_axis: f64, +} + +impl Default for SpheroidDistance { + fn default() -> Self { + Self { + semi_major_axis: 6378137.0, // WGS84 equatorial radius + semi_minor_axis: 6356752.314245, // WGS84 polar radius + } + } +} + +impl SpheroidDistance { + /// Create a new Spheroid distance metric with custom ellipsoid parameters. + pub fn with_ellipsoid(semi_major_axis: f64, semi_minor_axis: f64) -> Self { + Self { + semi_major_axis, + semi_minor_axis, + } + } + + /// Create a new Spheroid distance metric for GRS80 ellipsoid. + pub fn grs80() -> Self { + Self { + semi_major_axis: 6378137.0, + semi_minor_axis: 6356752.314140, + } + } +} + +impl DistanceMetric for SpheroidDistance { + fn distance(&self, lon1: N, lat1: N, lon2: N, lat2: N) -> N { + // Vincenty's formulae for distance on ellipsoid + let lat1 = match lat1.to_f64() { + Some(value) => value * PI / 180.0, + None => return N::zero(), // Return a default value if conversion fails + }; + let lat2 = match lat2.to_f64() { + Some(value) => value * PI / 180.0, + None => return N::zero(), + }; + let delta_lon = match (lon2.to_f64(), lon1.to_f64()) { + (Some(lon2_value), Some(lon1_value)) => (lon2_value - lon1_value) * PI / 180.0, + _ => return N::zero(), + }; + + let a = self.semi_major_axis; + let b = self.semi_minor_axis; + let f = (a - b) / a; // flattening + + let u1 = ((1.0 - f) * lat1.tan()).atan(); + let u2 = ((1.0 - f) * lat2.tan()).atan(); + + let sin_u1 = u1.sin(); + let cos_u1 = u1.cos(); + let sin_u2 = u2.sin(); + let cos_u2 = u2.cos(); + + let mut lambda = delta_lon; + let mut lambda_prev; + let mut iter_limit = 100; + + let (sin_sigma, cos_sigma, sigma, _sin_alpha, cos_sq_alpha, cos_2sigma_m) = loop { + let sin_lambda = lambda.sin(); + let cos_lambda = lambda.cos(); + + let sin_sigma = ((cos_u2 * sin_lambda).powi(2) + + (cos_u1 * sin_u2 - sin_u1 * cos_u2 * cos_lambda).powi(2)) + .sqrt(); + + if sin_sigma == 0.0 { + // Co-incident points + return N::zero(); + } + + let cos_sigma = sin_u1 * sin_u2 + cos_u1 * cos_u2 * cos_lambda; + let sigma = sin_sigma.atan2(cos_sigma); + + let sin_alpha = cos_u1 * cos_u2 * sin_lambda / sin_sigma; + let cos_sq_alpha = 1.0 - sin_alpha * sin_alpha; + + let cos_2sigma_m = if cos_sq_alpha == 0.0 { + 0.0 // Equatorial line + } else { + cos_sigma - 2.0 * sin_u1 * sin_u2 / cos_sq_alpha + }; + + let c = f / 16.0 * cos_sq_alpha * (4.0 + f * (4.0 - 3.0 * cos_sq_alpha)); + + lambda_prev = lambda; + lambda = delta_lon + + (1.0 - c) + * f + * sin_alpha + * (sigma + + c * sin_sigma + * (cos_2sigma_m + + c * cos_sigma * (-1.0 + 2.0 * cos_2sigma_m * cos_2sigma_m))); + + iter_limit -= 1; + if iter_limit == 0 || (lambda - lambda_prev).abs() < 1e-12 { + break ( + sin_sigma, + cos_sigma, + sigma, + sin_alpha, + cos_sq_alpha, + cos_2sigma_m, + ); + } + }; + + let u_sq = cos_sq_alpha * (a * a - b * b) / (b * b); + let big_a = + 1.0 + u_sq / 16384.0 * (4096.0 + u_sq * (-768.0 + u_sq * (320.0 - 175.0 * u_sq))); + let big_b = u_sq / 1024.0 * (256.0 + u_sq * (-128.0 + u_sq * (74.0 - 47.0 * u_sq))); + + let delta_sigma = big_b + * sin_sigma + * (cos_2sigma_m + + big_b / 4.0 + * (cos_sigma * (-1.0 + 2.0 * cos_2sigma_m * cos_2sigma_m) + - big_b / 6.0 + * cos_2sigma_m + * (-3.0 + 4.0 * sin_sigma * sin_sigma) + * (-3.0 + 4.0 * cos_2sigma_m * cos_2sigma_m))); + + let distance = b * big_a * (sigma - delta_sigma); + + N::from_f64(distance).unwrap_or(N::max_value()) + } + + fn distance_to_bbox( + &self, + lon: N, + lat: N, + min_lon: N, + min_lat: N, + max_lon: N, + max_lat: N, + ) -> N { + // For bbox distance with Spheroid, we approximate using the closest point on the bbox + let closest_lon = if lon < min_lon { + min_lon + } else if lon > max_lon { + max_lon + } else { + lon + }; + + let closest_lat = if lat < min_lat { + min_lat + } else if lat > max_lat { + max_lat + } else { + lat + }; + + self.distance(lon, lat, closest_lon, closest_lat) + } +} + +/// 1D distance from a value to a range. +#[inline] +fn axis_dist(k: N, min: N, max: N) -> N { + if k < min { + min - k + } else if k <= max { + N::zero() + } else { + k - max + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_euclidean_distance() { + let metric = EuclideanDistance; + let distance = metric.distance(0.0f64, 0.0f64, 3.0f64, 4.0f64); + assert!((distance - 5.0f64).abs() < 1e-10); + } + + #[test] + fn test_haversine_distance() { + let metric = HaversineDistance::default(); + // Distance between New York and London (approximately) + let distance = metric.distance(-74.0f64, 40.7f64, -0.1f64, 51.5f64); + // Should be approximately 5585 km + assert!((distance - 5585000.0f64).abs() < 50000.0f64); + } + + #[test] + fn test_spheroid_distance() { + let metric = SpheroidDistance::default(); + // Distance between New York and London (approximately) + let distance = metric.distance(-74.0f64, 40.7f64, -0.1f64, 51.5f64); + // Should be approximately 5585 km (slightly different from Haversine) + assert!((distance - 5585000.0f64).abs() < 50000.0f64); + } +} diff --git a/src/rtree/index.rs b/src/rtree/index.rs index 7db9c94..b91e2de 100644 --- a/src/rtree/index.rs +++ b/src/rtree/index.rs @@ -66,7 +66,7 @@ impl RTreeMetadata { let version = version_and_type >> 4; if version != VERSION { return Err(GeoIndexError::General( - format!("Got v{} data when expected v{}.", version, VERSION).to_string(), + format!("Got v{version} data when expected v{VERSION}.").to_string(), )); } diff --git a/src/rtree/mod.rs b/src/rtree/mod.rs index 08243a7..9897a1f 100644 --- a/src/rtree/mod.rs +++ b/src/rtree/mod.rs @@ -56,6 +56,7 @@ mod builder; mod constants; +pub mod distance; mod index; pub mod sort; mod r#trait; @@ -63,6 +64,7 @@ mod traversal; pub mod util; pub use builder::{RTreeBuilder, DEFAULT_RTREE_NODE_SIZE}; +pub use distance::{DistanceMetric, EuclideanDistance, HaversineDistance, SpheroidDistance}; pub use index::{RTree, RTreeMetadata, RTreeRef}; pub use r#trait::RTreeIndex; pub use traversal::Node; diff --git a/src/rtree/trait.rs b/src/rtree/trait.rs index 19be9be..95df690 100644 --- a/src/rtree/trait.rs +++ b/src/rtree/trait.rs @@ -6,6 +6,7 @@ use geo_traits::{CoordTrait, RectTrait}; use crate::error::Result; use crate::indices::Indices; use crate::r#type::IndexableNum; +use crate::rtree::distance::{DistanceMetric, EuclideanDistance}; use crate::rtree::index::{RTree, RTreeRef}; use crate::rtree::traversal::{IntersectionIterator, Node}; use crate::rtree::util::upper_bound; @@ -134,6 +135,9 @@ pub trait RTreeIndex: Sized { /// Search items in order of distance from the given point. /// + /// This method uses Euclidean distance by default. For other distance metrics, + /// use [`neighbors_with_distance`]. + /// /// ``` /// use geo_index::rtree::{RTreeBuilder, RTreeIndex, RTreeRef}; /// use geo_index::rtree::sort::HilbertSort; @@ -154,15 +158,48 @@ pub trait RTreeIndex: Sized { y: N, max_results: Option, max_distance: Option, + ) -> Vec { + // Use Euclidean distance by default for backward compatibility + let euclidean_distance = EuclideanDistance; + self.neighbors_with_distance(x, y, max_results, max_distance, &euclidean_distance) + } + + /// Search items in order of distance from the given point using a custom distance metric. + /// + /// This method allows you to specify a custom distance calculation method, such as + /// Euclidean, Haversine, or Spheroid distance. + /// + /// ``` + /// use geo_index::rtree::{RTreeBuilder, RTreeIndex}; + /// use geo_index::rtree::distance::{EuclideanDistance, HaversineDistance}; + /// use geo_index::rtree::sort::HilbertSort; + /// + /// // Create an RTree with geographic coordinates (longitude, latitude) + /// let mut builder = RTreeBuilder::::new(3); + /// builder.add(-74.0, 40.7, -74.0, 40.7); // New York + /// builder.add(-0.1, 51.5, -0.1, 51.5); // London + /// builder.add(139.7, 35.7, 139.7, 35.7); // Tokyo + /// let tree = builder.finish::(); + /// + /// // Find nearest neighbors using Haversine distance (great-circle distance) + /// let haversine = HaversineDistance::default(); + /// let results = tree.neighbors_with_distance(-74.0, 40.7, Some(2), None, &haversine); + /// ``` + fn neighbors_with_distance( + &self, + x: N, + y: N, + max_results: Option, + max_distance: Option, + distance_metric: &dyn DistanceMetric, ) -> Vec { let boxes = self.boxes(); let indices = self.indices(); - let max_distance = max_distance.unwrap_or(N::max_value()); + let max_distance = max_distance.unwrap_or(distance_metric.max_distance()); let mut outer_node_index = Some(boxes.len() - 4); let mut queue = BinaryHeap::new(); let mut results: Vec = vec![]; - let max_dist_squared = max_distance * max_distance; 'outer: while let Some(node_index) = outer_node_index { // find the end index of the node @@ -173,10 +210,17 @@ pub trait RTreeIndex: Sized { for pos in (node_index..end).step_by(4) { let index = indices.get(pos >> 2); - let dx = axis_dist(x, boxes[pos], boxes[pos + 2]); - let dy = axis_dist(y, boxes[pos + 1], boxes[pos + 3]); - let dist = dx * dx + dy * dy; - if dist > max_dist_squared { + // Use the custom distance metric for bbox distance calculation + let dist = distance_metric.distance_to_bbox( + x, + y, + boxes[pos], + boxes[pos + 1], + boxes[pos + 2], + boxes[pos + 3], + ); + + if dist > max_distance { continue; } @@ -188,9 +232,14 @@ pub trait RTreeIndex: Sized { })); } else { // leaf item (use odd id) + // For leaf items, calculate distance to the center of the bounding box + let center_x = (boxes[pos] + boxes[pos + 2]) / (N::one() + N::one()); + let center_y = (boxes[pos + 1] + boxes[pos + 3]) / (N::one() + N::one()); + let leaf_dist = distance_metric.distance(x, y, center_x, center_y); + queue.push(Reverse(NeighborNode { id: (index << 1) + 1, - dist, + dist: leaf_dist, })); } } @@ -198,7 +247,7 @@ pub trait RTreeIndex: Sized { // pop items from the queue while !queue.is_empty() && queue.peek().is_some_and(|val| (val.0.id & 1) != 0) { let dist = queue.peek().unwrap().0.dist; - if dist > max_dist_squared { + if dist > max_distance { break 'outer; } let item = queue.pop().unwrap(); @@ -228,6 +277,23 @@ pub trait RTreeIndex: Sized { self.neighbors(coord.x(), coord.y(), max_results, max_distance) } + /// Search items in order of distance from the given coordinate using a custom distance metric. + fn neighbors_coord_with_distance( + &self, + coord: &impl CoordTrait, + max_results: Option, + max_distance: Option, + distance_metric: &dyn DistanceMetric, + ) -> Vec { + self.neighbors_with_distance( + coord.x(), + coord.y(), + max_results, + max_distance, + distance_metric, + ) + } + /// Returns an iterator over the indexes of objects in this and another tree that intersect. /// /// Each returned object is of the form `(u32, u32)`, where the first is the positional @@ -339,4 +405,92 @@ mod test { assert_eq!(results, expected); } } + + mod distance_metrics { + use crate::rtree::distance::{EuclideanDistance, HaversineDistance, SpheroidDistance}; + use crate::rtree::sort::HilbertSort; + use crate::rtree::{RTreeBuilder, RTreeIndex}; + + #[test] + fn test_euclidean_distance_neighbors() { + let mut builder = RTreeBuilder::::new(3); + builder.add(0., 0., 1., 1.); + builder.add(2., 2., 3., 3.); + builder.add(4., 4., 5., 5.); + let tree = builder.finish::(); + + let euclidean = EuclideanDistance; + let results = tree.neighbors_with_distance(0., 0., None, None, &euclidean); + + // Should return items in order of distance from (0,0) + assert_eq!(results, vec![0, 1, 2]); + } + + #[test] + fn test_haversine_distance_neighbors() { + let mut builder = RTreeBuilder::::new(3); + // Add some geographic points (longitude, latitude) + builder.add(-74.0, 40.7, -74.0, 40.7); // New York + builder.add(-0.1, 51.5, -0.1, 51.5); // London + builder.add(139.7, 35.7, 139.7, 35.7); // Tokyo + let tree = builder.finish::(); + + let haversine = HaversineDistance::default(); + let results = tree.neighbors_with_distance(-74.0, 40.7, None, None, &haversine); + + // From New York, should find New York first, then London, then Tokyo + assert_eq!(results, vec![0, 1, 2]); + } + + #[test] + fn test_spheroid_distance_neighbors() { + let mut builder = RTreeBuilder::::new(3); + // Add some geographic points (longitude, latitude) + builder.add(-74.0, 40.7, -74.0, 40.7); // New York + builder.add(-0.1, 51.5, -0.1, 51.5); // London + builder.add(139.7, 35.7, 139.7, 35.7); // Tokyo + let tree = builder.finish::(); + + let spheroid = SpheroidDistance::default(); + let results = tree.neighbors_with_distance(-74.0, 40.7, None, None, &spheroid); + + // From New York, should find New York first, then London, then Tokyo + assert_eq!(results, vec![0, 1, 2]); + } + + #[test] + fn test_backward_compatibility() { + let mut builder = RTreeBuilder::::new(3); + builder.add(0., 0., 1., 1.); + builder.add(2., 2., 3., 3.); + builder.add(4., 4., 5., 5.); + let tree = builder.finish::(); + + // Test that original neighbors method still works + let results_original = tree.neighbors(0., 0., None, None); + + // Test that new method with Euclidean distance gives same results + let euclidean = EuclideanDistance; + let results_new = tree.neighbors_with_distance(0., 0., None, None, &euclidean); + + assert_eq!(results_original, results_new); + } + + #[test] + fn test_max_distance_filtering() { + let mut builder = RTreeBuilder::::new(3); + builder.add(0., 0., 1., 1.); + builder.add(2., 2., 3., 3.); + builder.add(10., 10., 11., 11.); + let tree = builder.finish::(); + + let euclidean = EuclideanDistance; + // Only find neighbors within distance 5 + let results = tree.neighbors_with_distance(0., 0., None, Some(5.0), &euclidean); + + // Should only find first two items, not the distant third one + assert_eq!(results.len(), 2); + assert_eq!(results, vec![0, 1]); + } + } } diff --git a/src/type.rs b/src/type.rs index 3330f5e..5a773d4 100644 --- a/src/type.rs +++ b/src/type.rs @@ -19,6 +19,29 @@ pub trait IndexableNum: const TYPE_INDEX: u8; /// The number of bytes per element const BYTES_PER_ELEMENT: usize; + + /// Convert to f64 for distance calculations + fn to_f64(self) -> Option { + NumCast::from(self) + } + + /// Convert from f64 for distance calculations + fn from_f64(value: f64) -> Option { + NumCast::from(value) + } + + /// Get the square root of this value + fn sqrt(self) -> Option { + self.to_f64() + .and_then(|value| { + if value >= 0.0 { + Some(value.sqrt()) + } else { + None + } + }) + .and_then(NumCast::from) + } } impl IndexableNum for i8 { @@ -122,7 +145,7 @@ impl CoordType { u32::TYPE_INDEX => CoordType::UInt32, f32::TYPE_INDEX => CoordType::Float32, f64::TYPE_INDEX => CoordType::Float64, - t => return Err(GeoIndexError::General(format!("Unexpected type {}.", t))), + t => return Err(GeoIndexError::General(format!("Unexpected type {t}."))), }; Ok(result) }