James
141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
from pathlib import Path
|
|
from typing import Iterable, Union
|
|
|
|
import numpy as np
|
|
|
|
from timezonefinder.configs import IntegerLike
|
|
|
|
from timezonefinder import utils
|
|
from timezonefinder.coord_accessors import AbstractCoordAccessor, create_coord_accessor
|
|
from timezonefinder.flatbuf.io.polygons import (
|
|
get_coordinate_path,
|
|
)
|
|
from timezonefinder.np_binary_helpers import (
|
|
get_xmax_path,
|
|
get_xmin_path,
|
|
get_ymax_path,
|
|
get_ymin_path,
|
|
read_per_polygon_vector,
|
|
)
|
|
|
|
|
|
class PolygonArray:
|
|
xmin: np.ndarray
|
|
xmax: np.ndarray
|
|
ymin: np.ndarray
|
|
ymax: np.ndarray
|
|
coordinates: AbstractCoordAccessor
|
|
|
|
def __init__(
|
|
self,
|
|
data_location: Union[str, Path],
|
|
in_memory: bool = False,
|
|
):
|
|
"""
|
|
Initialize the PolygonArray.
|
|
:param data_location: The path to the binary data files to use.
|
|
:param in_memory: Whether to completely read and keep the coordinate data in memory as numpy.
|
|
"""
|
|
self.in_memory = in_memory
|
|
self.data_location: Path = Path(data_location)
|
|
|
|
xmin_path = get_xmin_path(self.data_location)
|
|
xmax_path = get_xmax_path(self.data_location)
|
|
ymin_path = get_ymin_path(self.data_location)
|
|
ymax_path = get_ymax_path(self.data_location)
|
|
|
|
# read all per polygon vectors directly into memory (no matter the memory mode)
|
|
self.xmin = read_per_polygon_vector(xmin_path)
|
|
self.xmax = read_per_polygon_vector(xmax_path)
|
|
self.ymin = read_per_polygon_vector(ymin_path)
|
|
self.ymax = read_per_polygon_vector(ymax_path)
|
|
|
|
coordinate_file_path = get_coordinate_path(self.data_location)
|
|
# Initialize the appropriate coordinate accessor based on memory mode
|
|
self.coordinates = create_coord_accessor(coordinate_file_path, self.in_memory)
|
|
|
|
def __del__(self):
|
|
"""Clean up resources when the object is destroyed."""
|
|
del self.coordinates
|
|
del self.xmin
|
|
del self.xmax
|
|
del self.ymin
|
|
del self.ymax
|
|
|
|
def __len__(self) -> int:
|
|
"""
|
|
Get the number of polygons in the collection.
|
|
:return: Number of polygons
|
|
"""
|
|
return len(self.xmin)
|
|
|
|
def outside_bbox(self, poly_id: IntegerLike, x: int, y: int) -> bool:
|
|
"""
|
|
Check if a point is outside the bounding box of a polygon.
|
|
|
|
:param poly_id: Polygon ID
|
|
:param x: X-coordinate of the point
|
|
:param y: Y-coordinate of the point
|
|
:return: True if the point is outside the boundaries, False otherwise
|
|
"""
|
|
if x > self.xmax[poly_id]:
|
|
return True
|
|
if x < self.xmin[poly_id]:
|
|
return True
|
|
if y > self.ymax[poly_id]:
|
|
return True
|
|
if y < self.ymin[poly_id]:
|
|
return True
|
|
return False
|
|
|
|
def coords_of(self, idx: IntegerLike) -> np.ndarray:
|
|
"""
|
|
Get the polygon coordinates for the given index.
|
|
|
|
Args:
|
|
idx: The polygon index
|
|
|
|
Returns:
|
|
A numpy array containing the polygon coordinates
|
|
"""
|
|
return self.coordinates[idx]
|
|
|
|
def pip(self, poly_id: IntegerLike, x: int, y: int) -> bool:
|
|
"""
|
|
Point in polygon (PIP) test.
|
|
|
|
:param poly_id: Polygon ID
|
|
:param x: X-coordinate of the point
|
|
:param y: Y-coordinate of the point
|
|
:return: True if the point is inside the polygon, False otherwise
|
|
"""
|
|
polygon = self.coords_of(poly_id)
|
|
return utils.inside_polygon(x, y, polygon)
|
|
|
|
def pip_with_bbox_check(self, poly_id: IntegerLike, x: int, y: int) -> bool:
|
|
"""
|
|
Point in polygon (PIP) test with bounding box check.
|
|
|
|
:param poly_id: Polygon ID
|
|
:param x: X-coordinate of the point
|
|
:param y: Y-coordinate of the point
|
|
:return: True if the point is inside the polygon, False otherwise
|
|
"""
|
|
if self.outside_bbox(poly_id, x, y):
|
|
return False
|
|
return self.pip(poly_id, x, y)
|
|
|
|
def in_any_polygon(self, poly_ids: Iterable[int], x: int, y: int) -> bool:
|
|
"""
|
|
Check if a point is inside any of the specified polygons.
|
|
|
|
:param poly_ids: An iterable of polygon IDs
|
|
:param x: X-coordinate of the point
|
|
:param y: Y-coordinate of the point
|
|
:return: True if the point is inside any polygon, False otherwise
|
|
"""
|
|
for poly_id in poly_ids:
|
|
if self.pip_with_bbox_check(poly_id, x, y):
|
|
return True
|
|
return False
|