ENH: Add get_z ufunc (#175)

CHANGELOG.rst

@@ -18,6 +18,7 @@ Version 0.8 (unreleased)
 * Addition of a ``build_area()`` function for GEOS >= 3.8 (#141)
 * Addition of a ``normalize()`` function (#136)
 * Addition of a ``make_valid()`` function for GEOS >= 3.8 (#107)
+* Addition of a ``get_z()`` function for GEOS >= 3.7 (#175)
 * The ``get_coordinate_dimensions()`` function was renamed to
   ``get_coordinate_dimension()`` for consistency with GEOS (#176)
 * Addition of ``covers``, ``covered_by``, ``contains_properly`` predicates

pygeos/geometry.py

@@ -2,7 +2,7 @@
 import numpy as np
 from . import lib
 from . import Geometry  # NOQA
-from .decorators import multithreading_enabled
+from .decorators import multithreading_enabled, requires_geos
 
 __all__ = [
     "GeometryType",
@@ -14,6 +14,7 @@
     "set_srid",
     "get_x",
     "get_y",
+    "get_z",
     "get_exterior_ring",
     "get_num_points",
     "get_num_interior_rings",
@@ -212,7 +213,7 @@ def get_x(point):
 
     See also
     --------
-    get_y
+    get_y, get_z
 
     Examples
     --------
@@ -235,7 +236,7 @@ def get_y(point):
 
     See also
     --------
-    get_x
+    get_x, get_z
 
     Examples
     --------
@@ -247,6 +248,35 @@ def get_y(point):
     return lib.get_y(point)
 
 
+@requires_geos("3.7.0")
+@multithreading_enabled
+def get_z(point):
+    """Returns the z-coordinate of a point.
+
+    Requires at least GEOS 3.7.0.
+
+    Parameters
+    ----------
+    point : Geometry or array_like
+        Non-point geometries or geometries without 3rd dimension will result
+        in NaN being returned.
+
+    See also
+    --------
+    get_x, get_y
+
+    Examples
+    --------
+    >>> get_z(Geometry("POINT Z (1 2 3)"))
+    3.0
+    >>> get_z(Geometry("POINT (1 2)"))
+    nan
+    >>> get_z(Geometry("MULTIPOINT Z (1 1 1, 2 2 2)"))
+    nan
+    """
+    return lib.get_z(point)
+
+
 # linestrings
 
 @multithreading_enabled

pygeos/test/test_geometry.py

@@ -143,9 +143,21 @@ def test_get_set_srid():
     assert pygeos.get_srid(actual) == 4326
 
 
-@pytest.mark.parametrize("func", [pygeos.get_x, pygeos.get_y])
+@pytest.mark.parametrize(
+    "func",
+    [
+        pygeos.get_x,
+        pygeos.get_y,
+        pytest.param(
+            pygeos.get_z,
+            marks=pytest.mark.skipif(
+                pygeos.geos_version < (3, 7, 0), reason="GEOS < 3.7"
+            ),
+        ),
+    ],
+)
 @pytest.mark.parametrize("geom", all_types[1:])
-def test_get_xy_no_point(func, geom):
+def test_get_xyz_no_point(func, geom):
     assert np.isnan(func(geom))
 
 
@@ -157,6 +169,16 @@ def test_get_y():
     assert pygeos.get_y([point, point_z]).tolist() == [3.0, 1.0]
 
 
+@pytest.mark.skipif(pygeos.geos_version < (3, 7, 0), reason="GEOS < 3.7")
+def test_get_z():
+    assert pygeos.get_z([point_z]).tolist() == [1.0]
+
+
+@pytest.mark.skipif(pygeos.geos_version < (3, 7, 0), reason="GEOS < 3.7")
+def test_get_z_2d():
+    assert np.isnan(pygeos.get_z(point))
+
+
 @pytest.mark.parametrize("geom", all_types)
 def test_new_from_wkt(geom):
     actual = pygeos.Geometry(str(geom))

src/ufuncs.c

@@ -699,6 +699,18 @@ static int GetY(void *context, void *a, double *b) {
     }
 }
 static void *get_y_data[1] = {GetY};
+#if GEOS_SINCE_3_7_0
+    static int GetZ(void *context, void *a, double *b) {
+        char typ = GEOSGeomTypeId_r(context, a);
+        if (typ != 0) {
+            *(double *)b = NPY_NAN;
+            return 1;
+        } else {
+            return GEOSGeomGetZ_r(context, a, b);
+        }
+    }
+    static void *get_z_data[1] = {GetZ};
+#endif
 static void *area_data[1] = {GEOSArea_r};
 static void *length_data[1] = {GEOSLength_r};
 typedef int FuncGEOS_Y_d(void *context, void *a, double *b);
@@ -1937,6 +1949,7 @@ int init_ufuncs(PyObject *m, PyObject *d)
     DEFINE_CUSTOM (from_shapely, 1);
 
     #if GEOS_SINCE_3_7_0
+      DEFINE_Y_d (get_z);
       DEFINE_YY_d (frechet_distance);
       DEFINE_YYd_d (frechet_distance_densify);
     #endif