Nearest implementation with index

nearest.py

from shapely.geometry import Point, LineString, box


def _expand(geom, amt=0.5):
    bounds = geom.bounds
    d = math.sqrt(2 * (amt ** 2))
    while True:
        yield box(*bounds)
        bounds = (bounds[0] - d, bounds[1] - d, bounds[2] + d, bounds[3] + d)


def _min_distance(geometry, options):
    distances = [geometry.distance(g) for g in options]
    minimum_distance = min(distances)
    for opt, dist in zip(options, distances):
        if dist == minimum_distance:
            yield opt


def nearest(tree, query):
    for idx, bnds in enumerate(_expand(query)):
        print(f"search {idx + 1} with bounds {bnds}")
        results = tree.query(bnds)
        if results:
            return list(_min_distance(query, results))


def main():
    keys = [Point(2.5, 4), Point(0.0, 1.5)]
    tree = STRtree(keys)

    search = LineString([Point(1, 1), Point(3, 3)])
    nearest_geoms = nearest(tree, search)
    print([g.wkt for g in nearest_geoms])

    search2 = LineString([Point(3, 1), Point(0, 4)])
    nearest_geoms = nearest(tree, search2)
    print([g.wkt for g in nearest_geoms])


if __name__ == '__main__':
    main()