|
| 1 | +use geo::{Area, BooleanOps, BoundingRect, Geometry, Intersects, Polygon, Rect}; |
| 2 | +use rstar::{primitives::Rectangle, RTree, RTreeObject, AABB}; |
| 3 | +use std::iter::Sum; |
| 4 | + |
| 5 | +pub struct Node<C: rstar::RTreeNum + geo::CoordFloat, D> { |
| 6 | + pub geometry: Geometry<C>, |
| 7 | + pub rectangle: Rectangle<(C, C)>, |
| 8 | + pub data: D, |
| 9 | +} |
| 10 | + |
| 11 | +impl<C: rstar::RTreeNum + geo::CoordFloat, D> Node<C, D> { |
| 12 | + pub fn new(geometry: Geometry<C>, data: D) -> Result<Node<C, D>, String> { |
| 13 | + let rectangle = rect_from_geometry(&geometry)?; |
| 14 | + let result = Node { |
| 15 | + geometry, |
| 16 | + rectangle, |
| 17 | + data, |
| 18 | + }; |
| 19 | + Ok(result) |
| 20 | + } |
| 21 | +} |
| 22 | + |
| 23 | +impl<C: rstar::RTreeNum + geo::CoordFloat, D> RTreeObject for Node<C, D> { |
| 24 | + type Envelope = AABB<(C, C)>; |
| 25 | + |
| 26 | + fn envelope(&self) -> Self::Envelope { |
| 27 | + self.rectangle.envelope() |
| 28 | + } |
| 29 | +} |
| 30 | + |
| 31 | +pub type IntersectionWithArea<'a, C, D> = Vec<(&'a Node<C, D>, C)>; |
| 32 | + |
| 33 | +pub struct PolygonalRTree<C: rstar::RTreeNum + geo::CoordFloat, D>(RTree<Node<C, D>>); |
| 34 | + |
| 35 | +impl<C, D> PolygonalRTree<C, D> |
| 36 | +where |
| 37 | + C: rstar::RTreeNum + geo::GeoNum + geo::CoordFloat + geo::bool_ops::BoolOpsNum + Sum, |
| 38 | +{ |
| 39 | + pub fn new(data: Vec<(Geometry<C>, D)>) -> Result<PolygonalRTree<C, D>, String> { |
| 40 | + let nodes = data |
| 41 | + .into_iter() |
| 42 | + .map(|(g, d)| Node::new(g, d)) |
| 43 | + .collect::<Result<Vec<_>, _>>()?; |
| 44 | + let tree = RTree::bulk_load(nodes); |
| 45 | + Ok(PolygonalRTree(tree)) |
| 46 | + } |
| 47 | + |
| 48 | + /// tests for intersection with polygonal data in this tree. this involves two steps: |
| 49 | + /// 1. finding rtree rectangle envelopes that intersect the incoming geometry |
| 50 | + /// 2. testing intersection for each discovered geometry bounded by it's rectangle |
| 51 | + pub fn intersection<'a>( |
| 52 | + &'a self, |
| 53 | + g: &'a Geometry<C>, |
| 54 | + ) -> Result<Box<dyn Iterator<Item = &'a Node<C, D>> + 'a>, String> { |
| 55 | + let query = rect_from_geometry(g)?; |
| 56 | + let iter = self |
| 57 | + .0 |
| 58 | + .locate_in_envelope_intersecting(&query.envelope()) |
| 59 | + .filter(|node| node.geometry.intersects(g)); |
| 60 | + Ok(Box::new(iter)) |
| 61 | + } |
| 62 | + |
| 63 | + pub fn intersection_with_overlap_area<'a>( |
| 64 | + &'a self, |
| 65 | + query: &'a Geometry<C>, |
| 66 | + ) -> Result<IntersectionWithArea<'a, C, D>, String> { |
| 67 | + // get all polygons in the query geometry |
| 68 | + let query_polygons: Vec<Polygon<C>> = match query { |
| 69 | + Geometry::Polygon(p) => Ok(vec![p.clone()]), |
| 70 | + Geometry::MultiPolygon(mp) => Ok(mp.0.clone()), |
| 71 | + // Geometry::GeometryCollection(geometry_collection) => todo!(), |
| 72 | + _ => Err(String::from( |
| 73 | + "areal proportion query must be performed on polygonal data", |
| 74 | + )), |
| 75 | + }?; |
| 76 | + |
| 77 | + // compute the overlap area for each query polygon for each geometry |
| 78 | + // found to intersect the query in the rtree |
| 79 | + let result = self |
| 80 | + .intersection(query)? |
| 81 | + .map(|node| { |
| 82 | + let overlap_areas = query_polygons |
| 83 | + .iter() |
| 84 | + .map(|p| { |
| 85 | + let area = overlap_area(p, &node.geometry)?; |
| 86 | + Ok(area) |
| 87 | + }) |
| 88 | + .collect::<Result<Vec<C>, String>>()?; |
| 89 | + let overlap_area: C = overlap_areas.into_iter().sum(); |
| 90 | + Ok((node, overlap_area)) |
| 91 | + }) |
| 92 | + .collect::<Result<Vec<(&Node<C, D>, C)>, String>>()?; |
| 93 | + |
| 94 | + Ok(result) |
| 95 | + } |
| 96 | +} |
| 97 | + |
| 98 | +/// helper function to create a rectangular rtree envelope for a given geometry |
| 99 | +fn rect_from_geometry<C: rstar::RTreeNum + geo::CoordFloat>( |
| 100 | + g: &Geometry<C>, |
| 101 | +) -> Result<Rectangle<(C, C)>, String> { |
| 102 | + let bbox_vec: Rect<C> = g |
| 103 | + .bounding_rect() |
| 104 | + .ok_or_else(|| String::from("internal error: cannot get bounds of geometry"))?; |
| 105 | + |
| 106 | + let envelope = Rectangle::from_corners(bbox_vec.min().x_y(), bbox_vec.max().x_y()); |
| 107 | + Ok(envelope) |
| 108 | +} |
| 109 | + |
| 110 | +/// helper for computing the overlap area, which is the area that two geometries have in common |
| 111 | +/// (the intersection). |
| 112 | +fn overlap_area<C>(query: &Polygon<C>, overlapping: &Geometry<C>) -> Result<C, String> |
| 113 | +where |
| 114 | + C: rstar::RTreeNum + geo::CoordFloat + geo::bool_ops::BoolOpsNum + Sum, |
| 115 | +{ |
| 116 | + match overlapping { |
| 117 | + Geometry::Polygon(overlap_p) => { |
| 118 | + let overlap = query.intersection(overlap_p); |
| 119 | + let overlap_area: C = overlap.iter().map(|p| p.unsigned_area()).sum(); |
| 120 | + Ok(overlap_area) |
| 121 | + } |
| 122 | + Geometry::MultiPolygon(mp) => { |
| 123 | + let overlaps = mp |
| 124 | + .iter() |
| 125 | + .map(|overlapping_p| { |
| 126 | + overlap_area(query, &geo::Geometry::Polygon(overlapping_p.clone())) |
| 127 | + }) |
| 128 | + .collect::<Result<Vec<_>, String>>()?; |
| 129 | + let overlap_area: C = overlaps.into_iter().sum(); |
| 130 | + Ok(overlap_area) |
| 131 | + } |
| 132 | + _ => Err(String::from("polygonal rtree node must be polygonal!")), |
| 133 | + } |
| 134 | +} |
0 commit comments