ENH: prepared geometry as additional (cached) attribute on GeometryObject

pygeos/creation.py

@@ -13,6 +13,8 @@
     "multipolygons",
     "geometrycollections",
     "box",
+    "prepare",
+    "destroy_prepared",
 ]
 
 
@@ -170,3 +172,46 @@ def geometrycollections(geometries):
         An array of geometries
     """
     return lib.create_collection(geometries, GeometryType.GEOMETRYCOLLECTION)
+
+
+def prepare(geometry, **kwargs):
+    """Compute a prepared geometry.
+
+    A prepared geometry is a normal geometry with added information such as an
+    index on the line segments. This improves the performance of the following operations:
+    contains, contains_properly, covered_by, covers, crosses, disjoint, intersects,
+    overlaps, touches, and within.
+
+    Note that if a prepared geometry is modified, the newly created Geometry object is
+    not prepared. In that case, ``prepare`` should be called again.
+
+    This function does not recompute previously prepared geometries; 
+    it is efficient to call this function on an array that partially contains prepared geometries. 
+
+    Parameters
+    ----------
+    geometry : Geometry or array_like
+
+    See also
+    --------
+    destroy_prepared
+    """
+    return lib.prepare(geometry, **kwargs)
+
+
+def destroy_prepared(geometry, **kwargs):
+    """Destroy a previously prepared geometry, freeing up memory.
+
+    Note that the prepared geometry will always be cleaned up if the geometry itself
+    is dereferenced. This function needs only be called in very specific circumstances,
+    such as freeing up memory without losing the geometries, or benchmarking.
+
+    Parameters
+    ----------
+    geometry : Geometry or array_like
+
+    See also
+    --------
+    prepare
+    """
+    return lib.destroy_prepared(geometry, **kwargs)

pygeos/predicates.py

@@ -345,6 +345,10 @@ def crosses(a, b, **kwargs):
     ----------
     a, b : Geometry or array_like
 
+    See also
+    --------
+    prepare : improve performance by preparing ``a`` (the first argument)
+
     Examples
     --------
     >>> line = Geometry("LINESTRING(0 0, 1 1)")
@@ -383,6 +387,7 @@ def contains(a, b, **kwargs):
     See also
     --------
     within : ``contains(A, B) == within(B, A)``
+    prepare : improve performance by preparing ``a`` (the first argument)
 
     Examples
     --------
@@ -424,6 +429,7 @@ def covered_by(a, b, **kwargs):
     See also
     --------
     covers : ``covered_by(A, B) == covers(B, A)``
+    prepare : improve performance by preparing ``a`` (the first argument)
 
     Examples
     --------
@@ -465,6 +471,7 @@ def covers(a, b, **kwargs):
     See also
     --------
     covered_by : ``covers(A, B) == covered_by(B, A)``
+    prepare : improve performance by preparing ``a`` (the first argument)
 
     Examples
     --------
@@ -509,6 +516,7 @@ def disjoint(a, b, **kwargs):
     See also
     --------
     intersects : ``disjoint(A, B) == ~intersects(A, B)``
+    prepare : improve performance by preparing ``a`` (the first argument)
 
     Examples
     --------
@@ -574,6 +582,7 @@ def intersects(a, b, **kwargs):
     See also
     --------
     disjoint : ``intersects(A, B) == ~disjoint(A, B)``
+    prepare : improve performance by preparing ``a`` (the first argument)
 
     Examples
     --------
@@ -599,6 +608,10 @@ def overlaps(a, b, **kwargs):
     ----------
     a, b : Geometry or array_like
 
+    See also
+    --------
+    prepare : improve performance by preparing ``a`` (the first argument)
+
     Examples
     --------
     >>> line = Geometry("LINESTRING(0 0, 1 1)")
@@ -625,6 +638,10 @@ def touches(a, b, **kwargs):
     ----------
     a, b : Geometry or array_like
 
+    See also
+    --------
+    prepare : improve performance by preparing ``a`` (the first argument)
+
     Examples
     --------
     >>> line = Geometry("LINESTRING(0 2, 2 0)")
@@ -663,6 +680,7 @@ def within(a, b, **kwargs):
     See also
     --------
     contains : ``within(A, B) == contains(B, A)``
+    prepare : improve performance by preparing ``a`` (the first argument)
 
     Examples
     --------

pygeos/test/test_creation.py

@@ -65,12 +65,15 @@ def test_linestrings_from_xyz():
     assert str(actual) == "LINESTRING Z (0 2 0, 1 3 0)"
 
 
-@pytest.mark.parametrize("shape", [
-    (2, 1, 2),  # 2 linestrings of 1 2D point
-    (1, 1, 2),  # 1 linestring of 1 2D point
-    (1, 2),  # 1 linestring of 1 2D point (scalar)
-    (2, ),  # 1 2D point (scalar)
-])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (2, 1, 2),  # 2 linestrings of 1 2D point
+        (1, 1, 2),  # 1 linestring of 1 2D point
+        (1, 2),  # 1 linestring of 1 2D point (scalar)
+        (2,),  # 1 2D point (scalar)
+    ],
+)
 def test_linestrings_invalid_shape(shape):
     with pytest.raises(ValueError):
         pygeos.linestrings(np.ones(shape))
@@ -91,18 +94,21 @@ def test_linearrings_unclosed():
     assert str(actual) == "LINEARRING (1 0, 1 1, 0 1, 0 0, 1 0)"
 
 
-@pytest.mark.parametrize("shape", [
-    (2, 1, 2),  # 2 linearrings of 1 2D point
-    (1, 1, 2),  # 1 linearring of 1 2D point
-    (1, 2),  # 1 linearring of 1 2D point (scalar)
-    (2, 2, 2),  # 2 linearrings of 2 2D points
-    (1, 2, 2),  # 1 linearring of 2 2D points
-    (2, 2),  # 1 linearring of 2 2D points (scalar)
-    (2, 3, 2),  # 2 linearrings of 3 2D points
-    (1, 3, 2),  # 1 linearring of 3 2D points
-    (3, 2),  # 1 linearring of 3 2D points (scalar)
-    (2, ),  # 1 2D point (scalar)
-])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (2, 1, 2),  # 2 linearrings of 1 2D point
+        (1, 1, 2),  # 1 linearring of 1 2D point
+        (1, 2),  # 1 linearring of 1 2D point (scalar)
+        (2, 2, 2),  # 2 linearrings of 2 2D points
+        (1, 2, 2),  # 1 linearring of 2 2D points
+        (2, 2),  # 1 linearring of 2 2D points (scalar)
+        (2, 3, 2),  # 2 linearrings of 3 2D points
+        (1, 3, 2),  # 1 linearring of 3 2D points
+        (3, 2),  # 1 linearring of 3 2D points (scalar)
+        (2,),  # 1 2D point (scalar)
+    ],
+)
 def test_linearrings_invalid_shape(shape):
     coords = np.ones(shape)
     with pytest.raises(ValueError):
@@ -113,6 +119,7 @@ def test_linearrings_invalid_shape(shape):
     with pytest.raises(ValueError):
         pygeos.linearrings(coords)
 
+
 def test_polygon_from_linearring():
     actual = pygeos.polygons(pygeos.linearrings(box_tpl(0, 0, 1, 1)))
     assert str(actual) == "POLYGON ((1 0, 1 1, 0 1, 0 0, 1 0))"
@@ -163,46 +170,52 @@ def test_2_polygons_with_different_holes():
     assert pygeos.area(actual).tolist() == [96.0, 24.0]
 
 
-@pytest.mark.parametrize("shape", [
-    (2, 1, 2),  # 2 linearrings of 1 2D point
-    (1, 1, 2),  # 1 linearring of 1 2D point
-    (1, 2),  # 1 linearring of 1 2D point (scalar)
-    (2, 2, 2),  # 2 linearrings of 2 2D points
-    (1, 2, 2),  # 1 linearring of 2 2D points
-    (2, 2),  # 1 linearring of 2 2D points (scalar)
-    (2, 3, 2),  # 2 linearrings of 3 2D points
-    (1, 3, 2),  # 1 linearring of 3 2D points
-    (3, 2),  # 1 linearring of 3 2D points (scalar)
-    (2, ),  # 1 2D point (scalar)
-])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (2, 1, 2),  # 2 linearrings of 1 2D point
+        (1, 1, 2),  # 1 linearring of 1 2D point
+        (1, 2),  # 1 linearring of 1 2D point (scalar)
+        (2, 2, 2),  # 2 linearrings of 2 2D points
+        (1, 2, 2),  # 1 linearring of 2 2D points
+        (2, 2),  # 1 linearring of 2 2D points (scalar)
+        (2, 3, 2),  # 2 linearrings of 3 2D points
+        (1, 3, 2),  # 1 linearring of 3 2D points
+        (3, 2),  # 1 linearring of 3 2D points (scalar)
+        (2,),  # 1 2D point (scalar)
+    ],
+)
 def test_polygons_not_enough_points_in_shell(shape):
     coords = np.ones(shape)
     with pytest.raises(ValueError):
         pygeos.polygons(coords)
-    
+
     # make sure the first coordinate != second coordinate
     coords[..., 1] += 1
     with pytest.raises(ValueError):
         pygeos.polygons(coords)
 
 
-@pytest.mark.parametrize("shape", [
-    (2, 1, 2),  # 2 linearrings of 1 2D point
-    (1, 1, 2),  # 1 linearring of 1 2D point
-    (1, 2),  # 1 linearring of 1 2D point (scalar)
-    (2, 2, 2),  # 2 linearrings of 2 2D points
-    (1, 2, 2),  # 1 linearring of 2 2D points
-    (2, 2),  # 1 linearring of 2 2D points (scalar)
-    (2, 3, 2),  # 2 linearrings of 3 2D points
-    (1, 3, 2),  # 1 linearring of 3 2D points
-    (3, 2),  # 1 linearring of 3 2D points (scalar)
-    (2, ),  # 1 2D point (scalar)
-])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (2, 1, 2),  # 2 linearrings of 1 2D point
+        (1, 1, 2),  # 1 linearring of 1 2D point
+        (1, 2),  # 1 linearring of 1 2D point (scalar)
+        (2, 2, 2),  # 2 linearrings of 2 2D points
+        (1, 2, 2),  # 1 linearring of 2 2D points
+        (2, 2),  # 1 linearring of 2 2D points (scalar)
+        (2, 3, 2),  # 2 linearrings of 3 2D points
+        (1, 3, 2),  # 1 linearring of 3 2D points
+        (3, 2),  # 1 linearring of 3 2D points (scalar)
+        (2,),  # 1 2D point (scalar)
+    ],
+)
 def test_polygons_not_enough_points_in_holes(shape):
     coords = np.ones(shape)
     with pytest.raises(ValueError):
         pygeos.polygons(np.ones((1, 4, 2)), coords)
-    
+
     # make sure the first coordinate != second coordinate
     coords[..., 1] += 1
     with pytest.raises(ValueError):
@@ -320,6 +333,33 @@ def test_subclasses(with_point_in_registry):
         assert point.x == 1
 
 
+def test_prepare():
+    arr = np.array([pygeos.points(1, 1), None, pygeos.box(0, 0, 1, 1)])
+    assert arr[0]._ptr_prepared == 0
+    assert arr[2]._ptr_prepared == 0
+    pygeos.prepare(arr)
+    assert arr[0]._ptr_prepared != 0
+    assert arr[1] is None
+    assert arr[2]._ptr_prepared != 0
+
+    # preparing again actually does nothing
+    original = arr[0]._ptr_prepared
+    pygeos.prepare(arr)
+    assert arr[0]._ptr_prepared == original
+
+
+def test_destroy_prepared():
+    arr = np.array([pygeos.points(1, 1), None, pygeos.box(0, 0, 1, 1)])
+    pygeos.prepare(arr)
+    assert arr[0]._ptr_prepared != 0
+    assert arr[2]._ptr_prepared != 0
+    pygeos.destroy_prepared(arr)
+    assert arr[0]._ptr_prepared == 0
+    assert arr[1] is None
+    assert arr[2]._ptr_prepared == 0
+    pygeos.destroy_prepared(arr)  # does not error
+
+
 def test_subclass_is_geometry(with_point_in_registry):
     assert pygeos.is_geometry(Point("POINT (1 1)"))
 

pygeos/test/test_predicates.py

@@ -31,6 +31,10 @@
     pygeos.equals_exact,
 )
 
+BINARY_PREPARED_PREDICATES = tuple(
+    set(BINARY_PREDICATES) - {pygeos.equals, pygeos.equals_exact}
+)
+
 
 @pytest.mark.parametrize("geometry", all_types)
 @pytest.mark.parametrize("func", UNARY_PREDICATES)
@@ -122,3 +126,12 @@ def test_relate_none(g1, g2):
 ])
 def test_is_ccw(geom, expected):
     assert pygeos.is_ccw(geom) == expected
+
+
+@pytest.mark.parametrize("a", all_types)
+@pytest.mark.parametrize("func", BINARY_PREPARED_PREDICATES)
+def test_binary_prepared(a, func):
+    actual = func(a, point)
+    pygeos.lib.prepare(a)
+    result = func(a, point)
+    assert actual == result

src/fast_loop_macros.h

@@ -10,6 +10,16 @@
  *     npy_intp *steps
  */
 
+
+/** (ip1) -> () */
+#define NO_OUTPUT_LOOP\
+    char *ip1 = args[0];\
+    npy_intp is1 = steps[0];\
+    npy_intp n = dimensions[0];\
+    npy_intp i;\
+    for(i = 0; i < n; i++, ip1 += is1)
+
+
 /** (ip1) -> (op1) */
 #define UNARY_LOOP                         \
   char *ip1 = args[0], *op1 = args[1];     \