def __init__(self, locations: np.ndarray):
self.lats = locations[:, 0]
self.lons = locations[:, 1]
dist0 = vec_haversine(self.lats, self.lons, 0.0, 0.0)
dist1 = vec_haversine(self.lats, self.lons, 90.0, 0.0)
self.idx0 = np.argsort(dist0)
self.idx1 = np.argsort(dist1)
self.sorted0 = dist0[self.idx0]
self.sorted1 = dist1[self.idx1]
def query_radius(self, location: np.ndarray, r: float) -> np.ndarray:
"""
Selects the indices of the points that lie within a given distance from
a given location.
:param location: Location to query in [lat, lon] format
:param r: Radius in meters
:return: Array of indices
"""
lat = location[0, 0]
lon = location[0, 1]
d0 = num_haversine(lat, lon, 0.0, 0.0)
d1 = num_haversine(lat, lon, 90.0, 0.0)
i0 = np.searchsorted(self.sorted0, d0 - r)
i1 = np.searchsorted(self.sorted0, d0 + r)
match0 = self.idx0[i0:i1 + 1]
i0 = np.searchsorted(self.sorted1, d1 - r)
i1 = np.searchsorted(self.sorted1, d1 + r)
match1 = self.idx1[i0:i1 + 1]
intersect = np.intersect1d(match0, match1)
dist = vec_haversine(self.lats[intersect], self.lons[intersect], lat,
lon)
return intersect[dist <= r]
def query_radius(self, location: np.ndarray, r: float) -> np.ndarray:
"""
Selects the indices of the points that lie within a given distance from
a given location.
:param location: Location to query in [lat, lon] format
:param r: Radius in meters
:return: Array of indices
"""
lat = location[0, 0]
lon = location[0, 1]
dist = vec_haversine(self.lats, self.lons, lat, lon)
return np.argwhere(dist <= r)
def calculate_dx(self, df: pd.DataFrame) -> pd.DataFrame:
"""
Calculates the consecutive distance in meters.
:param df: Input DataFrame
:return: DataFrame with added 'dx' column.
"""
lat0 = df[self.lat_col][:-1].to_numpy()
lon0 = df[self.lon_col][:-1].to_numpy()
lat1 = df[self.lat_col][1:].to_numpy()
lon1 = df[self.lon_col][1:].to_numpy()
dist = vec_haversine(lat0, lon0, lat1, lon1)
df[self.dx_col] = np.insert(dist, 0, 0.0)
return df
def get_knn(self, lat: float, lon: float, k: int = 5) -> List[OSMNode]:
"""
Calculates the k-nearest neighbors of the given point.
Note that this function uses a brute-force approach not recommended
for large datasets
:param lat: Query point latitude
:param lon: Query point longitude
:param k: Number of nearest neighbors
:return: List of the (up to) k nearest neighbors
"""
lats = np.array([n.lat for n in self.ref_nodes])
lons = np.array([n.lon for n in self.ref_nodes])
dist = vec_haversine(lats, lons, lat, lon)
return np.argsort(dist)[:k]