ENH: Get and set coordinates as a array of floats (#44)

docs/coordinates.rst

@@ -0,0 +1,9 @@
+:mod:`pygeos.coordinates`
+=========================
+
+.. automodule:: pygeos.coordinates
+   :members:
+   :undoc-members:
+   :special-members:
+   :inherited-members:
+   :show-inheritance:

docs/index.rst

@@ -62,6 +62,7 @@ API Reference
    measurement
    predicates
    set_operations
+   coordinates
 
 
 Indices and tables

pygeos/__init__.py

@@ -8,7 +8,9 @@
 from .measurement import *
 from .set_operations import *
 from .linear import *
+from .coordinates import *
 
 from ._version import get_versions
-__version__ = get_versions()['version']
+
+__version__ = get_versions()["version"]
 del get_versions

pygeos/coordinates.py

@@ -0,0 +1,127 @@
+from . import lib, Geometry
+import numpy as np
+
+__all__ = ["apply", "count_coordinates", "get_coordinates", "set_coordinates"]
+
+
+def apply(geometry, transformation):
+    """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.
+
+    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.
+
+    Examples
+    --------
+    >>> apply(Geometry("POINT (0 0)"), lambda x: x + 1)
+    <pygeos.Geometry POINT (1 1)>
+    >>> apply(Geometry("LINESTRING (2 2, 4 4)"), lambda x: x * [2, 3])
+    <pygeos.Geometry LINESTRING (4 6, 8 12)>
+    >>> apply(None, lambda x: x) is None
+    True
+    >>> apply([Geometry("POINT (0 0)"), None], lambda x: x).tolist()
+    [<pygeos.Geometry POINT (0 0)>, None]
+    """
+    geometry_arr = np.array(geometry, dtype=np.object)  # makes a copy
+    coordinates = lib.get_coordinates(geometry_arr)
+    new_coordinates = transformation(coordinates)
+    # check the array to yield understandable error messages
+    if not isinstance(new_coordinates, np.ndarray):
+        raise ValueError("The provided transformation did not return a numpy array")
+    if new_coordinates.dtype != np.float64:
+        raise ValueError(
+            "The provided transformation returned an array with an unexpected "
+            "dtype ({})".format(new_coordinates.dtype)
+        )
+    if new_coordinates.shape != coordinates.shape:
+        # if the shape is too small we will get a segfault
+        raise ValueError(
+            "The provided transformation returned an array with an unexpected "
+            "shape ({})".format(new_coordinates.shape)
+        )
+    geometry_arr = lib.set_coordinates(geometry_arr, new_coordinates)
+    if geometry_arr.ndim == 0 and not isinstance(geometry, np.ndarray):
+        return geometry_arr.item()
+    return geometry_arr
+
+
+def count_coordinates(geometry):
+    """Counts the number of coordinate pairs in a geometry array.
+
+    Parameters
+    ----------
+    geometry : Geometry or array_like
+
+    Examples
+    --------
+    >>> count_coordinates(Geometry("POINT (0 0)"))
+    1
+    >>> count_coordinates(Geometry("LINESTRING (2 2, 4 4)"))
+    2
+    >>> count_coordinates(None)
+    0
+    >>> count_coordinates([Geometry("POINT (0 0)"), None])
+    1
+    """
+    return lib.count_coordinates(np.asarray(geometry, dtype=np.object))
+
+
+def get_coordinates(geometry):
+    """Gets coordinates from a geometry array as an array of floats.
+
+    The shape of the returned array is (N, 2), with N being the number of
+    coordinate pairs. Three-dimensional data is ignored.
+
+    Parameters
+    ----------
+    geometry : Geometry or array_like
+
+    Examples
+    --------
+    >>> get_coordinates(Geometry("POINT (0 0)")).tolist()
+    [[0.0, 0.0]]
+    >>> get_coordinates(Geometry("LINESTRING (2 2, 4 4)")).tolist()
+    [[2.0, 2.0], [4.0, 4.0]]
+    >>> get_coordinates(None)
+    array([], shape=(0, 2), dtype=float64)
+    """
+    return lib.get_coordinates(np.asarray(geometry, dtype=np.object))
+
+
+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.
+
+    Parameters
+    ----------
+    geometry : Geometry or array_like
+    coordinates: array_like
+
+    Examples
+    --------
+    >>> set_coordinates(Geometry("POINT (0 0)"), [[1, 1]])
+    <pygeos.Geometry POINT (1 1)>
+    >>> set_coordinates([Geometry("POINT (0 0)"), Geometry("LINESTRING (0 0, 0 0)")], [[1, 2], [3, 4], [5, 6]]).tolist()
+    [<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)>]
+    """
+    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):
+        raise ValueError(
+            "The coordinate array has an invalid shape {}".format(coordinates.shape)
+        )
+    lib.set_coordinates(geometry_arr, coordinates)
+    if geometry_arr.ndim == 0 and not isinstance(geometry, np.ndarray):
+        return geometry_arr.item()
+    return geometry_arr

pygeos/test/test_constructive.py

@@ -4,7 +4,7 @@
 
 from pygeos import Geometry, GEOSException
 
-from .common import point, all_types, geometry_collection
+from .common import point, all_types
 
 CONSTRUCTIVE_NO_ARGS = (
     pygeos.boundary,

pygeos/test/test_coordinates.py

@@ -0,0 +1,159 @@
+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 .common import empty
+from .common import point
+from .common import point_z
+from .common import line_string
+from .common import linear_ring
+from .common import polygon
+from .common import polygon_with_hole
+from .common import multi_point
+from .common import multi_line_string
+from .common import multi_polygon
+from .common import geometry_collection
+
+nested_2 = pygeos.geometrycollections([geometry_collection, point])
+nested_3 = pygeos.geometrycollections([nested_2, point])
+
+
+@pytest.mark.parametrize(
+    "geoms,count",
+    [
+        ([], 0),
+        ([empty], 0),
+        ([point, empty], 1),
+        ([empty, point, empty], 1),
+        ([point, None], 1),
+        ([None, point, None], 1),
+        ([point, point], 2),
+        ([point, point_z], 2),
+        ([line_string, linear_ring], 8),
+        ([polygon], 5),
+        ([polygon_with_hole], 10),
+        ([multi_point, multi_line_string], 4),
+        ([multi_polygon], 10),
+        ([geometry_collection], 3),
+        ([nested_2], 4),
+        ([nested_3], 5),
+    ],
+)
+def test_count_coords(geoms, count):
+    actual = count_coordinates(np.array(geoms, np.object))
+    assert actual == count
+
+
+# fmt: off
+@pytest.mark.parametrize(
+    "geoms,x,y",
+    [
+        ([], [], []),
+        ([empty], [], []),
+        ([point, empty], [2], [3]),
+        ([empty, point, empty], [2], [3]),
+        ([point, None], [2], [3]),
+        ([None, point, None], [2], [3]),
+        ([point, point], [2, 2], [3, 3]),
+        ([point, point_z], [2, 1], [3, 1]),
+        ([line_string, linear_ring], [0, 1, 1, 0, 1, 1, 0, 0], [0, 0, 1, 0, 0, 1, 1, 0]),
+        ([polygon], [0, 2, 2, 0, 0], [0, 0, 2, 2, 0]),
+        ([polygon_with_hole], [0, 0, 10, 10, 0, 2, 2, 4, 4, 2], [0, 10, 10, 0, 0, 2, 4, 4, 2, 2]),
+        ([multi_point, multi_line_string], [0, 1, 0, 1], [0, 2, 0, 2]),
+        ([multi_polygon], [0, 1, 1, 0, 0, 2.1, 2.2, 2.2, 2.1, 2.1], [0, 0, 1, 1, 0, 2.1, 2.1, 2.2, 2.2, 2.1]),
+        ([geometry_collection], [51, 52, 49], [-1, -1, 2]),
+        ([nested_2], [51, 52, 49, 2], [-1, -1, 2, 3]),
+        ([nested_3], [51, 52, 49, 2, 2], [-1, -1, 2, 3, 3]),
+    ],
+)  # fmt: on
+def test_get_coords(geoms, x, y):
+    actual = get_coordinates(np.array(geoms, np.object))
+    expected = np.array([x, y], np.float64).T
+    assert_equal(actual, expected)
+
+
+@pytest.mark.parametrize(
+    "geoms,count,has_ring",
+    [
+        ([], 0, False),
+        ([empty], 0, False),
+        ([point, empty], 1, False),
+        ([empty, point, empty], 1, False),
+        ([point, None], 1, False),
+        ([None, point, None], 1, False),
+        ([point, point], 2, False),
+        ([point, point_z], 2, False),
+        ([line_string, linear_ring], 8, True),
+        ([polygon], 5, True),
+        ([polygon_with_hole], 10, True),
+        ([multi_point, multi_line_string], 4, False),
+        ([multi_polygon], 10, True),
+        ([geometry_collection], 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_nan():
+    geoms = np.array([point])
+    coords = np.array([[np.nan, np.inf]])
+    new_geoms = set_coordinates(geoms, coords)
+    assert_equal(coords, get_coordinates(new_geoms))
+
+
+def test_set_coords_breaks_ring():
+    with pytest.raises(pygeos.GEOSException):
+        set_coordinates(linear_ring, np.random.random((5, 2)))
+
+
+def test_set_coords_0dim():
+    # a geometry input returns a geometry
+    actual = set_coordinates(point, [[1, 1]])
+    assert isinstance(actual, pygeos.Geometry)
+    # a 0-dim array input returns a 0-dim array
+    actual = set_coordinates(np.asarray(point), [[1, 1]])
+    assert isinstance(actual, np.ndarray)
+    assert actual.ndim == 0
+
+
+@pytest.mark.parametrize(
+    "geoms",
+    [[], [empty], [None, point, None], [nested_3]],
+)
+def test_apply(geoms):
+    geoms = np.array(geoms, np.object)
+    coordinates_before = get_coordinates(geoms)
+    new_geoms = apply(geoms, lambda x: x + 1)
+    assert new_geoms is not geoms
+    coordinates_after = get_coordinates(new_geoms)
+    assert_equal(coordinates_before + 1, coordinates_after)
+
+
+def test_apply_0dim():
+    # a geometry input returns a geometry
+    actual = apply(point, lambda x: x + 1)
+    assert isinstance(actual, pygeos.Geometry)
+    # a 0-dim array input returns a 0-dim array
+    actual = apply(np.asarray(point), lambda x: x + 1)
+    assert isinstance(actual, np.ndarray)
+    assert actual.ndim == 0
+
+
+def test_apply_check_shape():
+    def remove_coord(arr):
+        return arr[:-1]
+
+    with pytest.raises(ValueError):
+        apply(linear_ring, remove_coord)

setup.py

@@ -109,7 +109,7 @@ def finalize_options(self):
 
 module_lib = Extension(
     "pygeos.lib",
-    sources=["src/lib.c", "src/geos.c", "src/pygeom.c", "src/ufuncs.c"],
+    sources=["src/lib.c", "src/geos.c", "src/pygeom.c", "src/ufuncs.c", "src/coords.c"],
     **get_geos_paths()
 )
 
@@ -135,7 +135,10 @@ def finalize_options(self):
     packages=["pygeos"],
     setup_requires=["numpy"],
     install_requires=["numpy>=1.10"],
-    extras_require={"test": ["pytest"], "docs": ["sphinx", "numpydoc"]},
+    extras_require={
+        "test": ["pytest"],
+        "docs": ["sphinx", "numpydoc"],
+    },
     python_requires=">=3",
     include_package_data=True,
     ext_modules=[module_lib],