Repository: pygeos/pygeos

Issue:
Support get/set_coordinates/apply with more than 2 dimensions
Currently, the `apply` function (based on the `get_coordinates` / `set_coordinates` functions), explicitly states it is only for 2 dimensions and others will be dropped:

https://github.com/pygeos/pygeos/blob/b176e297c673671cd78f82738f1a899310501dd9/pygeos/coordinates.py#L7-L12

However, for using this function for coordinate transformations in GeoPandas, we need to be able to work with more than 2 dimensions (at least right now with Shapely, that is supported). 

@caspervdw do you remember if there was a fundamental reason this could not be done, or rather to keep the initial PR simpler?
 
I suppose the main problem is that there is no guarantee that all geometries in the array will have the same number of dimensions. 
However, I suppose there could be an option to specify how many dimensions should be used, and then in `get_coordinates`, those geometries that have less dimensions as the specified dimension can be filled with 0's or nan's for those missing dimensions?

Discussion:

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 6913085..b18283f 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,6 +19,8 @@ Version 0.9 (unreleased)
   types, now actually requiring line and point geometries (#216)
 * Addition of ``get_parts`` function to get individual parts of an array of multipart
   geometries (#197).
+* The ``apply`` function for coordinate transformations and the ``set_coordinates``
+  function now support geometries with z-coordinates (#131).
 * Addition of Cython and internal PyGEOS C API to enable easier development of internal
   functions (previously all significant internal functions were developed in C) (#51).
 
diff --git a/README.rst b/README.rst
index d8b6687..67651d2 100644
--- a/README.rst
+++ b/README.rst
@@ -20,6 +20,12 @@ PyGEOS
 	:alt: Appveyor CI status
 	:target: https://ci.appveyor.com/project/caspervdw/pygeos-3e5cu
 
+.. Github Actions status — https://github.com/pygeos/pygeos/actions
+
+.. image:: https://github.com/pygeos/pygeos/workflows/Conda/badge.svg
+	:alt: Github Actions status
+	:target: https://github.com/pygeos/pygeos/actions?query=workflow%3AConda
+
 .. PyPI
 
 .. image:: https://badge.fury.io/py/pygeos.svg
diff --git a/pygeos/coordinates.py b/pygeos/coordinates.py
index 559ad16..0cb40f6 100644
--- a/pygeos/coordinates.py
+++ b/pygeos/coordinates.py
@@ -4,19 +4,26 @@ import numpy as np
 __all__ = ["apply", "count_coordinates", "get_coordinates", "set_coordinates"]
 
 
-def apply(geometry, transformation):
+def apply(geometry, transformation, include_z=False):
     """Returns a copy of a geometry array with a function applied to its
     coordinates.
 
-    All returned geometries will be two-dimensional; the third dimension will
-    be discarded, if present.
+    With the default of ``include_z=False``, all returned geometries will be
+    two-dimensional; the third dimension will be discarded, if present.
+    When specifying ``include_z=True``, the returned geometries preserve
+    the dimensionality of the respective input geometries.
 
     Parameters
     ----------
     geometry : Geometry or array_like
     transformation : function
-        A function that transforms a (N, 2) ndarray of float64 to another
-        (N, 2) ndarray of float64.
+        A function that transforms a (N, 2) or (N, 3) ndarray of float64 to
+        another (N, 2) or (N, 3) ndarray of float64.
+    include_z : bool, default False
+        Whether to include the third dimension in the coordinates array
+        that is passed to the `transformation` function. If True, and a
+        geometry has no third dimension, the z-coordinates passed to the
+        function will be NaN.
 
     Examples
     --------
@@ -28,9 +35,16 @@ def apply(geometry, transformation):
     True
     >>> apply([Geometry("POINT (0 0)"), None], lambda x: x).tolist()
     [<pygeos.Geometry POINT (0 0)>, None]
+
+    By default, the third dimension is ignored:
+
+    >>> apply(Geometry("POINT Z (0 0 0)"), lambda x: x + 1)
+    <pygeos.Geometry POINT (1 1)>
+    >>> apply(Geometry("POINT Z (0 0 0)"), lambda x: x + 1, include_z=True)
+    <pygeos.Geometry POINT Z (1 1 1)>
     """
     geometry_arr = np.array(geometry, dtype=np.object)  # makes a copy
-    coordinates = lib.get_coordinates(geometry_arr, False)
+    coordinates = lib.get_coordinates(geometry_arr, include_z)
     new_coordinates = transformation(coordinates)
     # check the array to yield understandable error messages
     if not isinstance(new_coordinates, np.ndarray):
@@ -110,8 +124,10 @@ def get_coordinates(geometry, include_z=False):
 def set_coordinates(geometry, coordinates):
     """Returns a copy of a geometry array with different coordinates.
 
-    All returned geometries will be two-dimensional; the third dimension will
-    be discarded, if present.
+    If the coordinates array has shape (N, 2), all returned geometries
+    will be two-dimensional, and the third dimension will be discarded,
+    if present. If the coordinates array has shape (N, 3), the returned
+    geometries preserve the dimensionality of the input geometries.
 
     Parameters
     ----------
@@ -126,10 +142,27 @@ def set_coordinates(geometry, coordinates):
     [<pygeos.Geometry POINT (1 2)>, <pygeos.Geometry LINESTRING (3 4, 5 6)>]
     >>> set_coordinates([None, Geometry("POINT (0 0)")], [[1, 2]]).tolist()
     [None, <pygeos.Geometry POINT (1 2)>]
+
+    Third dimension of input geometry is discarded if coordinates array does
+    not include one:
+
+    >>> set_coordinates(Geometry("POINT Z (0 0 0)"), [[1, 1]])
+    <pygeos.Geometry POINT (1 1)>
+    >>> set_coordinates(Geometry("POINT Z (0 0 0)"), [[1, 1, 1]])
+    <pygeos.Geometry POINT Z (1 1 1)>
     """
     geometry_arr = np.asarray(geometry, dtype=np.object)
     coordinates = np.atleast_2d(np.asarray(coordinates)).astype(np.float64)
-    if coordinates.shape != (lib.count_coordinates(geometry_arr), 2):
+    n_coords = lib.count_coordinates(geometry_arr)
+    if coordinates.ndim != 2:
+        raise ValueError(
+            "The coordinate array should have dimension of 2 "
+            "(has {})".format(coordinates.ndim)
+        )
+    if (
+        (coordinates.shape[0] != lib.count_coordinates(geometry_arr))
+        or (coordinates.shape[1] not in {2, 3})
+    ):
         raise ValueError(
             "The coordinate array has an invalid shape {}".format(coordinates.shape)
         )
diff --git a/src/coords.c b/src/coords.c
index ae7959f..acd7c19 100644
--- a/src/coords.c
+++ b/src/coords.c
@@ -18,7 +18,7 @@
 static char get_coordinates(GEOSContextHandle_t, GEOSGeometry*, PyArrayObject*, npy_intp*,
                             int);
 static void* set_coordinates(GEOSContextHandle_t, GEOSGeometry*, PyArrayObject*,
-                             npy_intp*);
+                             npy_intp*, int);
 
 /* Get coordinates from a point, linestring or linearring and puts them at
 position `cursor` in the array `out`. Increases the cursor correspondingly.
@@ -133,9 +133,10 @@ static char get_coordinates(GEOSContextHandle_t context, GEOSGeometry* geom,
 new coordinates set from position `cursor` in the array `out`. The value of the
 cursor is increased correspondingly. Returns NULL on error,*/
 static void* set_coordinates_simple(GEOSContextHandle_t context, GEOSGeometry* geom,
-                                    int type, PyArrayObject* coords, npy_intp* cursor) {
+                                    int type, PyArrayObject* coords, npy_intp* cursor,
+                                    int include_z) {
   unsigned int n, i, dims;
-  double *x, *y;
+  double *x, *y, *z;
   GEOSGeometry* ret;
 
   /* Special case for POINT EMPTY (Point coordinate list cannot be 0-length) */
@@ -155,6 +156,12 @@ static void* set_coordinates_simple(GEOSContextHandle_t context, GEOSGeometry* g
     return NULL;
   }
 
+  /* If we have CoordSeq with z dim, but new coordinates only are 2D,
+   * create new CoordSeq that is also only 2D */
+  if ((dims == 3) && !include_z) {
+    dims = 2;
+  }
+
   /* Create a new one to fill with the new coordinates */
   GEOSCoordSequence* seq_new = GEOSCoordSeq_create_r(context, n, dims);
   if (seq_new == NULL) {
@@ -170,6 +177,12 @@ static void* set_coordinates_simple(GEOSContextHandle_t context, GEOSGeometry* g
     if (GEOSCoordSeq_setY_r(context, seq_new, i, *y) == 0) {
       goto fail;
     }
+    if (dims == 3) {
+      z = PyArray_GETPTR2(coords, *cursor, 2);
+      if (GEOSCoordSeq_setZ_r(context, seq_new, i, *z) == 0) {
+        goto fail;
+      }
+    }
   }
 
   /* Construct a new geometry */
@@ -194,7 +207,8 @@ fail:
 `set_coordinates_simple` on the linearrings that make the polygon.
 Returns NULL on error,*/
 static void* set_coordinates_polygon(GEOSContextHandle_t context, GEOSGeometry* geom,
-                                     PyArrayObject* coords, npy_intp* cursor) {
+                                     PyArrayObject* coords, npy_intp* cursor,
+                                     int include_z) {
   int i;
   GEOSGeometry *shell, *hole, *result = NULL;
   int n = GEOSGetNumInteriorRings_r(context, geom);
@@ -209,7 +223,7 @@ static void* set_coordinates_polygon(GEOSContextHandle_t context, GEOSGeometry*
   if (shell == NULL) {
     goto finish;
   }
-  shell = set_coordinates_simple(context, shell, 2, coords, cursor);
+  shell = set_coordinates_simple(context, shell, 2, coords, cursor, include_z);
   if (shell == NULL) {
     goto finish;
   }
@@ -219,7 +233,7 @@ static void* set_coordinates_polygon(GEOSContextHandle_t context, GEOSGeometry*
     if (hole == NULL) {
       goto finish;
     }
-    hole = set_coordinates_simple(context, hole, 2, coords, cursor);
+    hole = set_coordinates_simple(context, hole, 2, coords, cursor, include_z);
     if (hole == NULL) {
       goto finish;
     }
@@ -238,8 +252,8 @@ finish:
 /* Returns a copy of the input collection with new coordinates set by calling
 `set_coordinates` on the constituent subgeometries. Returns NULL on error,*/
 static void* set_coordinates_collection(GEOSContextHandle_t context, GEOSGeometry* geom,
-                                        int type, PyArrayObject* coords,
-                                        npy_intp* cursor) {
+                                        int type, PyArrayObject* coords, npy_intp* cursor,
+                                        int include_z) {
   int i;
   GEOSGeometry *sub_geom, *result = NULL;
   int n = GEOSGetNumGeometries_r(context, geom);
@@ -253,7 +267,7 @@ static void* set_coordinates_collection(GEOSContextHandle_t context, GEOSGeometr
     if (sub_geom == NULL) {
       goto finish;
     }
-    sub_geom = set_coordinates(context, sub_geom, coords, cursor);
+    sub_geom = set_coordinates(context, sub_geom, coords, cursor, include_z);
     if (sub_geom == NULL) {
       goto finish;
     }
@@ -272,14 +286,14 @@ finish:
 `cursor` in the array `out`. The value of the cursor is increased
 correspondingly. Returns NULL on error,*/
 static void* set_coordinates(GEOSContextHandle_t context, GEOSGeometry* geom,
-                             PyArrayObject* coords, npy_intp* cursor) {
+                             PyArrayObject* coords, npy_intp* cursor, int include_z) {
   int type = GEOSGeomTypeId_r(context, geom);
   if ((type == 0) | (type == 1) | (type == 2)) {
-    return set_coordinates_simple(context, geom, type, coords, cursor);
+    return set_coordinates_simple(context, geom, type, coords, cursor, include_z);
   } else if (type == 3) {
-    return set_coordinates_polygon(context, geom, coords, cursor);
+    return set_coordinates_polygon(context, geom, coords, cursor, include_z);
   } else if ((type >= 4) & (type <= 7)) {
-    return set_coordinates_collection(context, geom, type, coords, cursor);
+    return set_coordinates_collection(context, geom, type, coords, cursor, include_z);
   } else {
     return NULL;
   }
@@ -431,6 +445,8 @@ PyObject* SetCoords(PyArrayObject* geoms, PyArrayObject* coords) {
   NpyIter_IterNextFunc* iternext;
   char** dataptr;
   npy_intp cursor;
+  npy_intp* coords_shape;
+  int include_z;
   GeometryObject* obj;
   PyObject* new_obj;
   GEOSGeometry *geom, *new_geom;
@@ -442,6 +458,9 @@ PyObject* SetCoords(PyArrayObject* geoms, PyArrayObject* coords) {
     return (PyObject*)geoms;
   }
 
+  coords_shape = PyArray_SHAPE(coords);
+  include_z = (coords_shape[1] == 3);
+
   /* We use the Numpy iterator C-API here.
   The iterator exposes an "iternext" function which updates a "dataptr"
   see also: https://docs.scipy.org/doc/numpy/reference/c-api.iterator.html */
@@ -474,7 +493,7 @@ PyObject* SetCoords(PyArrayObject* geoms, PyArrayObject* coords) {
       continue;
     }
     /* create a new geometry with coordinates from "coords" array */
-    new_geom = set_coordinates(ctx, geom, coords, &cursor);
+    new_geom = set_coordinates(ctx, geom, coords, &cursor, include_z);
     if (new_geom == NULL) {
       errstate = PGERR_GEOS_EXCEPTION;
       goto finish;
@@ -556,6 +575,10 @@ PyObject* PySetCoords(PyObject* self, PyObject* args) {
     PyErr_SetString(PyExc_TypeError, "Coordinate array should be of float64 dtype");
     return NULL;
   }
+  if ((PyArray_NDIM((PyArrayObject*)coords)) != 2) {
+    PyErr_SetString(PyExc_ValueError, "Coordinate array should be 2-dimensional");
+    return NULL;
+  }
   geoms = SetCoords((PyArrayObject*)geoms, (PyArrayObject*)coords);
   if (geoms == Py_None) {
     return NULL;


Test patch:
diff --git a/.github/workflows/test-conda.yml b/.github/workflows/test-conda.yml
index ce5172f..3f68939 100644
--- a/.github/workflows/test-conda.yml
+++ b/.github/workflows/test-conda.yml
@@ -1,10 +1,10 @@
-name: Test using Conda environments
+name: Conda
 
 on: [push, pull_request]
 
 jobs:
   TestConda:
-    name: ${{ matrix.os }} - Python ${{ matrix.python }} (GEOS ${{ matrix.geos || 'latest' }}, numpy ${{ matrix.numpy || 'latest' }}, conda ${{ matrix.channel || 'defaults' }})
+    name: ${{ matrix.os }} - ${{ matrix.channel || 'defaults' }} - Python ${{ matrix.python }} (GEOS ${{ matrix.geos || 'latest' }}, numpy ${{ matrix.numpy || 'latest' }})
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -16,42 +16,27 @@ jobs:
           - os: ubuntu-latest
             python: 3.8
             channel: conda-forge
-          - os: macos-latest
-            python: 3.6
-            geos: 3.6
-            numpy: 1.13
-          - os: macos-latest
-            python: 3.7
-            geos: 3.7
-            numpy: 1.15
-          - os: macos-latest
-            python: 3.8
-            geos: 3.8
-            numpy: 1.17
-          - os: macos-latest
+          - os: windows-latest
             python: 3.8
-            channel: defaults
+            channel: conda-forge
           - os: macos-latest
             python: 3.8
             channel: conda-forge
-          - os: windows-latest
+          - os: macos-latest
             python: 3.6
             geos: 3.6
             numpy: 1.13
-          - os: windows-latest
+          - os: macos-latest
             python: 3.7
             geos: 3.7
             numpy: 1.15
-          - os: windows-latest
+          - os: macos-latest
             python: 3.8
             geos: 3.8
             numpy: 1.17
-          - os: windows-latest
+          - os: macos-latest
             python: 3.8
             channel: defaults
-          - os: windows-latest
-            python: 3.8
-            channel: conda-forge
 
     steps:
       - uses: actions/checkout@v2
diff --git a/pygeos/test/test_coordinates.py b/pygeos/test/test_coordinates.py
index 376557b..a4c46b0 100644
--- a/pygeos/test/test_coordinates.py
+++ b/pygeos/test/test_coordinates.py
@@ -2,7 +2,7 @@ import pytest
 import pygeos
 from pygeos import apply, count_coordinates, get_coordinates, set_coordinates
 import numpy as np
-from numpy.testing import assert_equal
+from numpy.testing import assert_equal, assert_allclose
 
 from .common import empty
 from .common import point
@@ -101,6 +101,7 @@ def test_get_coords_3d(geoms, x, y, z, include_z):
     assert_equal(actual, expected)
 
 
+@pytest.mark.parametrize("include_z", [True, False])
 @pytest.mark.parametrize(
     "geoms,count,has_ring",
     [
@@ -114,24 +115,24 @@ def test_get_coords_3d(geoms, x, y, z, include_z):
         ([point, point], 2, False),
         ([point, point_z], 2, False),
         ([line_string, linear_ring], 8, True),
+        ([line_string_z], 3, True),
         ([polygon], 5, True),
+        ([polygon_z], 5, True),
         ([polygon_with_hole], 10, True),
         ([multi_point, multi_line_string], 4, False),
         ([multi_polygon], 10, True),
         ([geometry_collection], 3, False),
+        ([geometry_collection_z], 3, False),
         ([nested_2], 4, False),
         ([nested_3], 5, False),
     ],
 )
-def test_set_coords(geoms, count, has_ring):
-    geoms = np.array(geoms, np.object)
-    if has_ring:
-        # do not randomize; linearrings / polygons should stay closed
-        coords = get_coordinates(geoms) + np.random.random((1, 2))
-    else:
-        coords = np.random.random((count, 2))
-    new_geoms = set_coordinates(geoms, coords)
-    assert_equal(coords, get_coordinates(new_geoms))
+def test_set_coords(geoms, count, has_ring, include_z):
+    arr_geoms = np.array(geoms, np.object)
+    n = 3 if include_z else 2
+    coords = get_coordinates(arr_geoms, include_z=include_z) + np.random.random((1, n))
+    new_geoms = set_coordinates(arr_geoms, coords)
+    assert_equal(coords, get_coordinates(new_geoms, include_z=include_z))
 
 
 def test_set_coords_nan():
@@ -156,17 +157,32 @@ def test_set_coords_0dim():
     assert actual.ndim == 0
 
 
+@pytest.mark.parametrize("include_z", [True, False])
+def test_set_coords_mixed_dimension(include_z):
+    geoms = np.array([point, point_z], dtype=object)
+    coords = get_coordinates(geoms, include_z=include_z)
+    new_geoms = set_coordinates(geoms, coords * 2)
+    if include_z:
+        # preserve original dimensionality
+        assert not pygeos.has_z(new_geoms[0])
+        assert pygeos.has_z(new_geoms[1])
+    else:
+        # all 2D
+        assert not pygeos.has_z(new_geoms).any()
+
+
+@pytest.mark.parametrize("include_z", [True, False])
 @pytest.mark.parametrize(
     "geoms",
-    [[], [empty], [None, point, None], [nested_3]],
+    [[], [empty], [None, point, None], [nested_3], [point, point_z], [line_string_z]],
 )
-def test_apply(geoms):
+def test_apply(geoms, include_z):
     geoms = np.array(geoms, np.object)
-    coordinates_before = get_coordinates(geoms)
-    new_geoms = apply(geoms, lambda x: x + 1)
+    coordinates_before = get_coordinates(geoms, include_z=include_z)
+    new_geoms = apply(geoms, lambda x: x + 1, include_z=include_z)
     assert new_geoms is not geoms
-    coordinates_after = get_coordinates(new_geoms)
-    assert_equal(coordinates_before + 1, coordinates_after)
+    coordinates_after = get_coordinates(new_geoms, include_z=include_z)
+    assert_allclose(coordinates_before + 1, coordinates_after, equal_nan=True)
 
 
 def test_apply_0dim():
@@ -185,3 +201,11 @@ def test_apply_check_shape():
 
     with pytest.raises(ValueError):
         apply(linear_ring, remove_coord)
+
+
+def test_apply_correct_coordinate_dimension():
+    # ensure that new geometry is 2D with include_z=False
+    geom = line_string_z
+    assert pygeos.get_coordinate_dimension(geom) == 3
+    new_geom = apply(geom, lambda x: x + 1, include_z=False)
+    assert pygeos.get_coordinate_dimension(new_geom) == 2