Add box ufunc based on create_box C function (#308)

CHANGELOG.rst

@@ -8,6 +8,10 @@ Version 0.10 (unreleased)
 
 * Addition of ``nearest`` and ``nearest_all`` functions to ``STRtree`` for
   GEOS >= 3.6 to find the nearest neighbors (#272).
+* Updated ``box`` ufunc to use internal C function for creating polygon
+  (about 2x faster) and added ``ccw`` parameter to create polygon in
+  counterclockwise (default) or clockwise direction (#308).
+
 
 **API Changes**
 

benchmarks/benchmarks.py

@@ -48,7 +48,8 @@ class ConstructiveSuite:
     """Benchmarks constructive functions on a set of 10,000 points"""
 
     def setup(self):
-        self.points = pygeos.points(np.random.random((10000, 2)))
+        self.coords = np.random.random((10000, 2))
+        self.points = pygeos.points(self.coords)
 
     def time_voronoi_polygons(self):
         pygeos.voronoi_polygons(self.points)
@@ -62,6 +63,9 @@ def time_convex_hull(self):
     def time_delaunay_triangles(self):
         pygeos.delaunay_triangles(self.points)
 
+    def time_box(self):
+        pygeos.box(*np.hstack([self.coords, self.coords + 100]).T)
+
 
 class ClipSuite:
     """Benchmarks for different methods of clipping geometries by boxes"""
@@ -247,12 +251,14 @@ def time_tree_nearest_points_equidistant_manual_all(self):
         # result
         l, r = self.grid_point_tree.nearest(self.grid_points)
         # calculate distance to nearest neighbor
-        dist = pygeos.distance(self.grid_points.take(l), self.grid_point_tree.geometries.take(r))
+        dist = pygeos.distance(
+            self.grid_points.take(l), self.grid_point_tree.geometries.take(r)
+        )
         # include a slight epsilon to ensure nearest are within this radius
         b = pygeos.buffer(self.grid_points, dist + 1e-8)
 
         # query the tree for others in the same buffer distance
-        left, right = self.grid_point_tree.query_bulk(b, predicate='intersects')
+        left, right = self.grid_point_tree.query_bulk(b, predicate="intersects")
         dist = pygeos.distance(
             self.grid_points.take(left), self.grid_point_tree.geometries.take(right)
         )

pygeos/creation.py

@@ -101,21 +101,30 @@ def polygons(shells, holes=None):
     return lib.polygons_with_holes(shells, holes)
 
 
-def box(x1, y1, x2, y2):
+def box(xmin, ymin, xmax, ymax, ccw=True, **kwargs):
     """Create box polygons.
 
     Parameters
     ----------
-    x1 : array_like
-    y1 : array_like
-    x2 : array_like
-    y2 : array_like
+    xmin : array_like
+    ymin : array_like
+    xmax : array_like
+    ymax : array_like
+    ccw : bool (default: True)
+        If True, box will be created in counterclockwise direction starting
+        from bottom right coordinate (xmax, ymin).
+        If False, box will be created in clockwise direction starting from
+        bottom left coordinate (xmin, ymin).
+
+    Examples
+    --------
+    >>> box(0, 0, 1, 1)
+    <pygeos.Geometry POLYGON ((1 0, 1 1, 0 1, 0 0, 1 0))>
+    >>> box(0, 0, 1, 1, ccw=False)
+    <pygeos.Geometry POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0))>
+
     """
-    x1, y1, x2, y2 = np.broadcast_arrays(x1, y1, x2, y2)
-    rings = np.array(((x2, y1), (x2, y2), (x1, y2), (x1, y1)))
-    # bring first two axes to the last two positions
-    rings = rings.transpose(list(range(2, rings.ndim)) + [0, 1])
-    return polygons(rings)
+    return lib.box(xmin, ymin, xmax, ymax, ccw, **kwargs)
 
 
 def multipoints(geometries):
@@ -188,8 +197,8 @@ def prepare(geometry, **kwargs):
     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. 
+    This function does not recompute previously prepared geometries;
+    it is efficient to call this function on an array that partially contains prepared geometries.
 
     Parameters
     ----------

pygeos/test/test_creation.py

@@ -296,15 +296,56 @@ def test_create_collection_wrong_geom_type(func, sub_geom):
         func([sub_geom])
 
 
-def test_box():
-    actual = pygeos.box(0, 0, 1, 1)
-    assert str(actual) == "POLYGON ((1 0, 1 1, 0 1, 0 0, 1 0))"
+@pytest.mark.parametrize(
+    "coords,ccw,expected",
+    [
+        ((0, 0, 1, 1), True, pygeos.Geometry("POLYGON ((1 0, 1 1, 0 1, 0 0, 1 0))")),
+        ((0, 0, 1, 1), False, pygeos.Geometry("POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0))")),
+    ],
+)
+def test_box(coords, ccw, expected):
+    actual = pygeos.box(*coords, ccw=ccw)
+    assert pygeos.equals(actual, expected)
 
 
-def test_box_multiple():
-    actual = pygeos.box(0, 0, [1, 2], [1, 2])
-    assert str(actual[0]) == "POLYGON ((1 0, 1 1, 0 1, 0 0, 1 0))"
-    assert str(actual[1]) == "POLYGON ((2 0, 2 2, 0 2, 0 0, 2 0))"
+@pytest.mark.parametrize(
+    "coords,ccw,expected",
+    [
+        (
+            (0, 0, [1, 2], [1, 2]),
+            True,
+            [
+                pygeos.Geometry("POLYGON ((1 0, 1 1, 0 1, 0 0, 1 0))"),
+                pygeos.Geometry("POLYGON ((2 0, 2 2, 0 2, 0 0, 2 0))"),
+            ],
+        ),
+        (
+            (0, 0, [1, 2], [1, 2]),
+            [True, False],
+            [
+                pygeos.Geometry("POLYGON ((1 0, 1 1, 0 1, 0 0, 1 0))"),
+                pygeos.Geometry("POLYGON ((0 0, 0 2, 2 2, 2 0, 0 0))"),
+            ],
+        ),
+    ],
+)
+def test_box_array(coords, ccw, expected):
+    actual = pygeos.box(*coords, ccw=ccw)
+    assert pygeos.equals(actual, expected).all()
+
+
+@pytest.mark.parametrize(
+    "coords",
+    [
+        [np.nan, np.nan, np.nan, np.nan],
+        [np.nan, 0, 1, 1],
+        [0, np.nan, 1, 1],
+        [0, 0, np.nan, 1],
+        [0, 0, 1, np.nan],
+    ],
+)
+def test_box_nan(coords):
+    assert pygeos.box(*coords) is None
 
 
 class BaseGeometry(pygeos.Geometry):

src/geos.c

@@ -4,8 +4,8 @@
 #include "geos.h"
 
 #include <Python.h>
-#include <structmember.h>
 #include <numpy/npy_math.h>
+#include <structmember.h>
 
 /* This initializes a globally accessible GEOSException object */
 PyObject* geos_exception[1] = {NULL};
@@ -354,7 +354,6 @@ void geos_notice_handler(const char* message, void* userdata) {
   snprintf(userdata, 1024, "%s", message);
 }
 
-
 /* Extract bounds from geometry.
  *
  * Bounds coordinates will be set to NPY_NAN if geom is NULL, empty, or does not have an
@@ -378,7 +377,6 @@ void geos_notice_handler(const char* message, void* userdata) {
  */
 int get_bounds(GEOSContextHandle_t ctx, GEOSGeometry* geom, double* xmin, double* ymin,
                double* xmax, double* ymax) {
-
   int retval = 1;
 
   if (geom == NULL || GEOSisEmpty_r(ctx, geom)) {
@@ -475,15 +473,17 @@ int get_bounds(GEOSContextHandle_t ctx, GEOSGeometry* geom, double* xmin, double
  * ymin: minimum Y value
  * xmax: maximum X value
  * ymax: maximum Y value
+ * ccw: if 1, box will be created in counterclockwise direction from bottom right;
+ *  otherwise will be created in clockwise direction from bottom left.
  *
  * Returns
  * -------
  * GEOSGeometry* on success (owned by caller) or
  * NULL on failure or NPY_NAN coordinates
  */
 GEOSGeometry* create_box(GEOSContextHandle_t ctx, double xmin, double ymin, double xmax,
-                         double ymax) {
-  if (xmin == NPY_NAN || ymin == NPY_NAN || xmax == NPY_NAN || ymax == NPY_NAN) {
+                         double ymax, char ccw) {
+  if (npy_isnan(xmin) | npy_isnan(ymin) | npy_isnan(xmax) | npy_isnan(ymax)) {
     return NULL;
   }
 
@@ -497,41 +497,56 @@ GEOSGeometry* create_box(GEOSContextHandle_t ctx, double xmin, double ymin, doub
     return NULL;
   }
 
-  if (!(GEOSCoordSeq_setX_r(ctx, coords, 0, xmin) &&
-        GEOSCoordSeq_setY_r(ctx, coords, 0, ymin) &&
-        GEOSCoordSeq_setX_r(ctx, coords, 1, xmax) &&
-        GEOSCoordSeq_setY_r(ctx, coords, 1, ymin) &&
-        GEOSCoordSeq_setX_r(ctx, coords, 2, xmax) &&
-        GEOSCoordSeq_setY_r(ctx, coords, 2, ymax) &&
-        GEOSCoordSeq_setX_r(ctx, coords, 3, xmin) &&
-        GEOSCoordSeq_setY_r(ctx, coords, 3, ymax) &&
-        GEOSCoordSeq_setX_r(ctx, coords, 4, xmin) &&
-        GEOSCoordSeq_setY_r(ctx, coords, 4, ymin))) {
-    if (coords != NULL) {
-      GEOSCoordSeq_destroy_r(ctx, coords);
+  if (ccw) {
+    // Start from bottom right (xmax, ymin) to match shapely
+    if (!(GEOSCoordSeq_setX_r(ctx, coords, 0, xmax) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 0, ymin) &&
+          GEOSCoordSeq_setX_r(ctx, coords, 1, xmax) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 1, ymax) &&
+          GEOSCoordSeq_setX_r(ctx, coords, 2, xmin) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 2, ymax) &&
+          GEOSCoordSeq_setX_r(ctx, coords, 3, xmin) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 3, ymin) &&
+          GEOSCoordSeq_setX_r(ctx, coords, 4, xmax) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 4, ymin))) {
+      if (coords != NULL) {
+        GEOSCoordSeq_destroy_r(ctx, coords);
+      }
+
+      return NULL;
     }
+  } else {
+    // Start from bottom left (min, ymin) to match shapely
+    if (!(GEOSCoordSeq_setX_r(ctx, coords, 0, xmin) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 0, ymin) &&
+          GEOSCoordSeq_setX_r(ctx, coords, 1, xmin) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 1, ymax) &&
+          GEOSCoordSeq_setX_r(ctx, coords, 2, xmax) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 2, ymax) &&
+          GEOSCoordSeq_setX_r(ctx, coords, 3, xmax) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 3, ymin) &&
+          GEOSCoordSeq_setX_r(ctx, coords, 4, xmin) &&
+          GEOSCoordSeq_setY_r(ctx, coords, 4, ymin))) {
+      if (coords != NULL) {
+        GEOSCoordSeq_destroy_r(ctx, coords);
+      }
 
-    return NULL;
+      return NULL;
+    }
   }
 
   // Construct linear ring then use to construct polygon
-  // Note: coords are owned by ring
+  // Note: coords are owned by ring; if ring fails to construct, it will
+  // automatically clean up the coords
   ring = GEOSGeom_createLinearRing_r(ctx, coords);
   if (ring == NULL) {
-    if (coords != NULL) {
-      GEOSCoordSeq_destroy_r(ctx, coords);
-    }
-
     return NULL;
   }
 
-  // Note: ring is owned by polygon
+  // Note: ring is owned by polygon; if polygon fails to construct, it will
+  // automatically clean up the ring
   geom = GEOSGeom_createPolygon_r(ctx, ring, NULL, 0);
   if (geom == NULL) {
-    if (ring != NULL) {
-      GEOSGeom_destroy_r(ctx, ring);
-    }
-
     return NULL;
   }
 

src/geos.h

@@ -141,10 +141,9 @@ extern char geos_interpolate_checker(GEOSContextHandle_t ctx, GEOSGeometry* geom
 
 extern int init_geos(PyObject* m);
 
-
 int get_bounds(GEOSContextHandle_t ctx, GEOSGeometry* geom, double* xmin, double* ymin,
                double* xmax, double* ymax);
 GEOSGeometry* create_box(GEOSContextHandle_t ctx, double xmin, double ymin, double xmax,
-                         double ymax);
+                         double ymax, char ccw);
 
-#endif // _GEOS_H
+#endif  // _GEOS_H

src/strtree.c

@@ -923,7 +923,7 @@ static PyObject* STRtree_nearest_all(STRtreeObject* self, PyObject* args) {
       }
 
       envelope = create_box(ctx, xmin - max_distance, ymin - max_distance,
-                            xmax + max_distance, ymax + max_distance);
+                            xmax + max_distance, ymax + max_distance, 1);
       if (envelope == NULL) {
         errstate = PGERR_GEOS_EXCEPTION;
         kv_destroy(src_indexes);

src/ufuncs.c

@@ -1249,6 +1249,48 @@ static PyUFuncGenericFunction YYd_Y_funcs[1] = {&YYd_Y_func};
 #endif
 
 /* Define functions with unique call signatures */
+static char box_dtypes[6] = {NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE,
+                             NPY_DOUBLE, NPY_BOOL, NPY_OBJECT};
+static void box_func(char** args, npy_intp* dimensions, npy_intp* steps, void* data) {
+  char *ip1 = args[0], *ip2 = args[1], *ip3 = args[2], *ip4 = args[3], *ip5 = args[4];
+  npy_intp is1 = steps[0], is2 = steps[1], is3 = steps[2], is4 = steps[3], is5 = steps[4];
+  npy_intp n = dimensions[0];
+  npy_intp i;
+  GEOSGeometry** geom_arr;
+
+  CHECK_NO_INPLACE_OUTPUT(6);
+
+  // allocate a temporary array to store output GEOSGeometry objects
+  geom_arr = malloc(sizeof(void*) * n);
+  CHECK_ALLOC(geom_arr);
+
+  GEOS_INIT_THREADS;
+
+  for (i = 0; i < n; i++, ip1 += is1, ip2 += is2, ip3 += is3, ip4 += is4, ip5 += is5) {
+    geom_arr[i] = create_box(ctx, *(double*)ip1, *(double*)ip2, *(double*)ip3,
+                             *(double*)ip4, *(char*)ip5);
+    if (geom_arr[i] == NULL) {
+      // result will be NULL for any nan coordinates, which is OK;
+      // otherwise raise an error
+      if (!(npy_isnan(*(double*)ip1) | npy_isnan(*(double*)ip2) |
+            npy_isnan(*(double*)ip3) | npy_isnan(*(double*)ip4))) {
+        errstate = PGERR_GEOS_EXCEPTION;
+        destroy_geom_arr(ctx, geom_arr, i - 1);
+        break;
+      }
+    }
+  }
+
+  GEOS_FINISH_THREADS;
+
+  // fill the numpy array with PyObjects while holding the GIL
+  if (errstate == PGERR_SUCCESS) {
+    geom_arr_to_npy(geom_arr, args[5], steps[5], dimensions[0]);
+  }
+  free(geom_arr);
+}
+
+static PyUFuncGenericFunction box_funcs[1] = {&box_func};
 
 static void* null_data[1] = {NULL};
 static char buffer_inner(void* ctx, GEOSBufferParams* params, void* ip1, void* ip2,
@@ -2744,6 +2786,7 @@ int init_ufuncs(PyObject* m, PyObject* d) {
 
   DEFINE_YYd_d(hausdorff_distance_densify);
 
+  DEFINE_CUSTOM(box, 5);
   DEFINE_CUSTOM(buffer, 7);
   DEFINE_CUSTOM(offset_curve, 5);
   DEFINE_CUSTOM(snap, 3);