ENH: Add get_z ufunc

pygeos/geometry.py

@@ -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,32 @@ def get_y(point):
     return lib.get_y(point)
 
 
+@multithreading_enabled
+def get_z(point):
+    """Returns the z-coordinate of a point.
+
+    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,9 @@ 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, pygeos.get_z])
 @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 +157,14 @@ def test_get_y():
     assert pygeos.get_y([point, point_z]).tolist() == [3.0, 1.0]
 
 
+def test_get_z():
+    assert pygeos.get_z([point_z]).tolist() == [1.0]
+
+
+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,19 @@ static int GetY(void *context, void *a, double *b) {
     }
 }
 static void *get_y_data[1] = {GetY};
+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 if (GEOSGeom_getCoordinateDimension_r(context, a) != 3) {
+    //     *(double *)b = NPY_NAN;
+    //     return 0;
+    } else {
+        return GEOSGeomGetZ_r(context, a, b);
+    }
+}
+static void *get_z_data[1] = {GetZ};
 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);
@@ -1896,6 +1909,7 @@ int init_ufuncs(PyObject *m, PyObject *d)
 
     DEFINE_Y_d (get_x);
     DEFINE_Y_d (get_y);
+    DEFINE_Y_d (get_z);
     DEFINE_Y_d (area);
     DEFINE_Y_d (length);