spatial_join/
rtrees.rs

1use geo::algorithm::bounding_rect::BoundingRect;
2use geo::{Line, LineString, Point, Polygon, Rect, Triangle};
3use rstar::RTree;
4
5use crate::SplitGeoSeq;
6
7type RTreeEnvelope = rstar::AABB<[f64; 2]>;
8
9#[derive(Debug)]
10pub struct FakeRegion {
11    pub id: usize,
12    pub bbox: RTreeEnvelope,
13}
14
15impl rstar::RTreeObject for FakeRegion {
16    type Envelope = RTreeEnvelope;
17
18    fn envelope(&self) -> Self::Envelope {
19        self.bbox
20    }
21}
22
23impl SplitGeoSeq {
24    pub fn to_rtrees(&self, max_distance: f64) -> [RTree<FakeRegion>; 6] {
25        // Why duplicate? Because bounding_rect isn't defined for
26        // Point and for the geos it is defined for, it sometimes
27        // gives you a Rect and sometimes Option<Rect>
28        [
29            RTree::bulk_load(
30                self.geos
31                    .points
32                    .iter()
33                    .enumerate()
34                    .map(|(index, pt)| FakeRegion {
35                        id: index,
36                        bbox: cheap_buffer(pt.to_env(), max_distance),
37                    })
38                    .collect(),
39            ),
40            RTree::bulk_load(
41                self.geos
42                    .lines
43                    .iter()
44                    .enumerate()
45                    .map(|(index, ln)| FakeRegion {
46                        id: index,
47                        bbox: cheap_buffer(ln.to_env(), max_distance),
48                    })
49                    .collect(),
50            ),
51            RTree::bulk_load(
52                self.geos
53                    .polys
54                    .iter()
55                    .enumerate()
56                    .map(|(index, poly)| FakeRegion {
57                        id: index,
58                        bbox: cheap_buffer(poly.to_env(), max_distance),
59                    })
60                    .collect(),
61            ),
62            RTree::bulk_load(
63                self.geos
64                    .line_strings
65                    .iter()
66                    .enumerate()
67                    .map(|(index, ls)| FakeRegion {
68                        id: index,
69                        bbox: cheap_buffer(ls.to_env(), max_distance),
70                    })
71                    .collect(),
72            ),
73            RTree::bulk_load(
74                self.geos
75                    .rects
76                    .iter()
77                    .enumerate()
78                    .map(|(index, rect)| FakeRegion {
79                        id: index,
80                        bbox: cheap_buffer(rect.to_env(), max_distance),
81                    })
82                    .collect(),
83            ),
84            RTree::bulk_load(
85                self.geos
86                    .tris
87                    .iter()
88                    .enumerate()
89                    .map(|(index, tri)| FakeRegion {
90                        id: index,
91                        bbox: cheap_buffer(tri.to_env(), max_distance),
92                    })
93                    .collect(),
94            ),
95        ]
96    }
97}
98
99pub trait Envelope {
100    fn to_env(&self) -> RTreeEnvelope;
101}
102
103impl Envelope for Point<f64> {
104    fn to_env(&self) -> RTreeEnvelope {
105        RTreeEnvelope::from_point([self.x(), self.y()])
106    }
107}
108
109impl Envelope for Line<f64> {
110    fn to_env(&self) -> RTreeEnvelope {
111        let bounds = self.bounding_rect();
112        RTreeEnvelope::from_corners(
113            [bounds.min().x, bounds.min().y],
114            [bounds.max().x, bounds.max().y],
115        )
116    }
117}
118
119impl Envelope for Polygon<f64> {
120    fn to_env(&self) -> RTreeEnvelope {
121        let bounds = self
122            .bounding_rect()
123            .expect("invalid bounding_rect for Polygon");
124        RTreeEnvelope::from_corners(
125            [bounds.min().x, bounds.min().y],
126            [bounds.max().x, bounds.max().y],
127        )
128    }
129}
130
131impl Envelope for LineString<f64> {
132    fn to_env(&self) -> RTreeEnvelope {
133        let bounds = self
134            .bounding_rect()
135            .expect("invalid bounding_rect for LineString");
136        RTreeEnvelope::from_corners(
137            [bounds.min().x, bounds.min().y],
138            [bounds.max().x, bounds.max().y],
139        )
140    }
141}
142
143impl Envelope for Rect<f64> {
144    fn to_env(&self) -> RTreeEnvelope {
145        RTreeEnvelope::from_corners([self.min().x, self.min().y], [self.max().x, self.max().y])
146    }
147}
148
149impl Envelope for Triangle<f64> {
150    fn to_env(&self) -> RTreeEnvelope {
151        let bounds = self.bounding_rect();
152        RTreeEnvelope::from_corners(
153            [bounds.min().x, bounds.min().y],
154            [bounds.max().x, bounds.max().y],
155        )
156    }
157}
158
159pub fn cheap_buffer(bbox: RTreeEnvelope, distance: f64) -> RTreeEnvelope {
160    let lower = bbox.lower();
161    let upper = bbox.upper();
162    RTreeEnvelope::from_corners(
163        [lower[0] - distance, lower[1] - distance],
164        [upper[0] + distance, upper[1] + distance],
165    )
166}