[Done] Force 2D/3D (#396)

CHANGELOG.rst

@@ -11,6 +11,8 @@ Version 0.11 (unreleased)
   geometries from ``indices`` (#380).
 * Added ``pygeos.empty`` to create a geometry array pre-filled with None or
   with empty geometries (#381).
+* Added ``pygeos.force_2d`` and ``pygeos.force_3d`` to change the dimensionality of
+  the coordinates in a geometry (#396).
 
 **API Changes**
 

pygeos/geometry.py

@@ -28,6 +28,8 @@
     "get_rings",
     "get_precision",
     "set_precision",
+    "force_2d",
+    "force_3d",
 ]
 
 
@@ -118,7 +120,8 @@ def get_dimensions(geometry, **kwargs):
 def get_coordinate_dimension(geometry, **kwargs):
     """Returns the dimensionality of the coordinates in a geometry (2 or 3).
 
-    Returns -1 for not-a-geometry values.
+    Returns -1 for missing geometries (``None`` values). Note that if the first Z
+    coordinate equals ``nan``, this function will return ``2``.
 
     Parameters
     ----------
@@ -135,6 +138,8 @@ def get_coordinate_dimension(geometry, **kwargs):
     3
     >>> get_coordinate_dimension(None)
     -1
+    >>> get_coordinate_dimension(Geometry("POINT Z (0 0 nan)"))
+    2
     """
     return lib.get_coordinate_dimension(geometry, **kwargs)
 
@@ -735,3 +740,61 @@ def set_precision(geometry, grid_size, preserve_topology=False, **kwargs):
     """
 
     return lib.set_precision(geometry, grid_size, preserve_topology, **kwargs)
+
+
+@multithreading_enabled
+def force_2d(geometry, **kwargs):
+    """Forces the dimensionality of a geometry to 2D.
+
+    Parameters
+    ----------
+    geometry : Geometry or array_like
+    **kwargs
+        For other keyword-only arguments, see the
+        `NumPy ufunc docs <https://numpy.org/doc/stable/reference/ufuncs.html#ufuncs-kwargs>`_.
+
+    Examples
+    --------
+    >>> force_2d(Geometry("POINT Z (0 0 0)"))
+    <pygeos.Geometry POINT (0 0)>
+    >>> force_2d(Geometry("POINT (0 0)"))
+    <pygeos.Geometry POINT (0 0)>
+    >>> force_2d(Geometry("LINESTRING (0 0 0, 0 1 1, 1 1 2)"))
+    <pygeos.Geometry LINESTRING (0 0, 0 1, 1 1)>
+    >>> force_2d(Geometry("POLYGON Z EMPTY"))
+    <pygeos.Geometry POLYGON EMPTY>
+    >>> force_2d(None) is None
+    True
+    """
+    return lib.force_2d(geometry, **kwargs)
+
+
+@multithreading_enabled
+def force_3d(geometry, z=0.0, **kwargs):
+    """Forces the dimensionality of a geometry to 3D.
+
+    2D geometries will get the provided Z coordinate; Z coordinates of 3D geometries
+    are unchanged (unless they are nan).
+
+    Parameters
+    ----------
+    geometry : Geometry or array_like
+    z : float or array_like, default 0.0
+    **kwargs
+        For other keyword-only arguments, see the
+        `NumPy ufunc docs <https://numpy.org/doc/stable/reference/ufuncs.html#ufuncs-kwargs>`_.
+
+    Examples
+    --------
+    >>> force_3d(Geometry("POINT (0 0)"), z=3)
+    <pygeos.Geometry POINT Z (0 0 3)>
+    >>> force_3d(Geometry("POINT Z (0 0 0)"), z=3)
+    <pygeos.Geometry POINT Z (0 0 0)>
+    >>> force_3d(Geometry("LINESTRING (0 0, 0 1, 1 1)"))
+    <pygeos.Geometry LINESTRING Z (0 0 0, 0 1 0, 1 1 0)>
+    >>> force_3d(None) is None
+    True
+    """
+    if np.isnan(z).any():
+        raise ValueError("It is not allowed to set the Z coordinate to NaN.")
+    return lib.force_3d(geometry, z, **kwargs)

pygeos/predicates.py

@@ -38,19 +38,28 @@
 def has_z(geometry, **kwargs):
     """Returns True if a geometry has a Z coordinate.
 
+    Note that this function returns False if the (first) Z coordinate equals NaN or
+    if the geometry is empty.
+
     Parameters
     ----------
     geometry : Geometry or array_like
     **kwargs
         For other keyword-only arguments, see the
         `NumPy ufunc docs <https://numpy.org/doc/stable/reference/ufuncs.html#ufuncs-kwargs>`_.
 
+    See also
+    --------
+    get_coordinate_dimension
+
     Examples
     --------
     >>> has_z(Geometry("POINT (0 0)"))
     False
     >>> has_z(Geometry("POINT Z (0 0 0)"))
     True
+    >>> has_z(Geometry("POINT Z(0 0 nan)"))
+    False
     """
     return lib.has_z(geometry, **kwargs)
 

pygeos/tests/common.py

@@ -26,19 +26,31 @@
 geometry_collection = pygeos.geometrycollections(
     [pygeos.points(51, -1), pygeos.linestrings([(52, -1), (49, 2)])]
 )
-point_z = pygeos.points(1.0, 1.0, 1.0)
-line_string_z = pygeos.linestrings([(0, 0, 0), (1, 0, 1), (1, 1, 2)])
-polygon_z = pygeos.polygons([(0, 0, 0), (2, 0, 1), (2, 2, 2), (0, 2, 3), (0, 0, 0)])
+point_z = pygeos.points(2, 3, 4)
+line_string_z = pygeos.linestrings([(0, 0, 4), (1, 0, 4), (1, 1, 4)])
+polygon_z = pygeos.polygons([(0, 0, 4), (2, 0, 4), (2, 2, 4), (0, 2, 4), (0, 0, 4)])
 geometry_collection_z = pygeos.geometrycollections([point_z, line_string_z])
 polygon_with_hole = pygeos.Geometry(
     "POLYGON((0 0, 0 10, 10 10, 10 0, 0 0), (2 2, 2 4, 4 4, 4 2, 2 2))"
 )
 empty_point = pygeos.Geometry("POINT EMPTY")
 empty_point_z = pygeos.Geometry("POINT Z EMPTY")
 empty_line_string = pygeos.Geometry("LINESTRING EMPTY")
+empty_line_string_z = pygeos.Geometry("LINESTRING Z EMPTY")
 empty_polygon = pygeos.Geometry("POLYGON EMPTY")
 empty = pygeos.Geometry("GEOMETRYCOLLECTION EMPTY")
 line_string_nan = pygeos.linestrings([(np.nan, np.nan), (np.nan, np.nan)])
+multi_point_z = pygeos.multipoints([(0, 0, 4), (1, 2, 4)])
+multi_line_string_z = pygeos.multilinestrings([[(0, 0, 4), (1, 2, 4)]])
+multi_polygon_z = pygeos.multipolygons(
+    [
+        [(0, 0, 4), (1, 0, 4), (1, 1, 4), (0, 1, 4), (0, 0, 4)],
+        [(2.1, 2.1, 4), (2.2, 2.1, 4), (2.2, 2.2, 4), (2.1, 2.2, 4), (2.1, 2.1, 4)],
+    ]
+)
+polygon_with_hole_z = pygeos.Geometry(
+    "POLYGON Z((0 0 4, 0 10 4, 10 10 4, 10 0 4, 0 0 4), (2 2 4, 2 4 4, 4 4 4, 4 2 4, 2 2 4))"
+)
 
 all_types = (
     point,
@@ -78,6 +90,6 @@ def assert_geometries_equal(actual, expected):
     expected = np.broadcast_to(expected, actual.shape)
     mask = pygeos.is_geometry(expected)
     if np.any(mask):
-        assert pygeos.equals(actual[mask], expected[mask]).all()
+        assert pygeos.equals_exact(actual[mask], expected[mask]).all()
     if np.any(~mask):
         assert_array_equal(actual[~mask], expected[~mask])

pygeos/tests/test_coordinates.py

@@ -7,7 +7,9 @@
 
 from .common import (
     empty,
+    empty_line_string_z,
     empty_point,
+    empty_point_z,
     geometry_collection,
     geometry_collection_z,
     line_string,
@@ -119,10 +121,10 @@ def test_get_coords_index_multidim(order):
 @pytest.mark.parametrize(
     "geoms,x,y,z",
     [
-        ([point, point_z], [2, 1], [3, 1], [np.nan, 1]),
-        ([line_string_z], [0, 1, 1], [0, 0, 1], [0, 1, 2]),
-        ([polygon_z], [0, 2, 2, 0, 0], [0, 0, 2, 2, 0], [0, 1, 2, 3, 0]),
-        ([geometry_collection_z], [1, 0, 1, 1], [1, 0, 0, 1], [1, 0, 1, 2]),
+        ([point, point_z], [2, 2], [3, 3], [np.nan, 4]),
+        ([line_string_z], [0, 1, 1], [0, 0, 1], [4, 4, 4]),
+        ([polygon_z], [0, 2, 2, 0, 0], [0, 0, 2, 2, 0], [4, 4, 4, 4, 4]),
+        ([geometry_collection_z], [2, 0, 1, 1], [3, 0, 0, 1], [4, 4, 4, 4]),
     ],
 )  # fmt: on
 def test_get_coords_3d(geoms, x, y, z, include_z):
@@ -242,3 +244,23 @@ def test_apply_correct_coordinate_dimension():
     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
+
+
+@pytest.mark.parametrize("geom", [
+    pytest.param(empty_point_z, marks=pytest.mark.skipif(pygeos.geos_version < (3, 9, 0), reason="Empty points don't have a dimensionality before GEOS 3.9")),
+    empty_line_string_z,
+])
+def test_apply_empty_preserve_z(geom):
+    assert pygeos.get_coordinate_dimension(geom) == 3
+    new_geom = apply(geom, lambda x: x + 1, include_z=True)
+    assert pygeos.get_coordinate_dimension(new_geom) == 3
+
+
+@pytest.mark.parametrize("geom", [
+    pytest.param(empty_point_z, marks=pytest.mark.skipif(pygeos.geos_version < (3, 9, 0), reason="Empty points don't have a dimensionality before GEOS 3.9")),
+    empty_line_string_z,
+])
+def test_apply_remove_z(geom):
+    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

pygeos/tests/test_geometry.py

@@ -3,23 +3,32 @@
 
 import pygeos
 
-from .common import all_types
+from .common import all_types, assert_geometries_equal
 from .common import empty as empty_geometry_collection
 from .common import (
     empty_line_string,
+    empty_line_string_z,
     empty_point,
+    empty_point_z,
     empty_polygon,
     geometry_collection,
+    geometry_collection_z,
     line_string,
     line_string_nan,
+    line_string_z,
     linear_ring,
     multi_line_string,
+    multi_line_string_z,
     multi_point,
+    multi_point_z,
     multi_polygon,
+    multi_polygon_z,
     point,
     point_z,
     polygon,
     polygon_with_hole,
+    polygon_with_hole_z,
+    polygon_z,
 )
 
 
@@ -175,16 +184,16 @@ def test_get_xyz_no_point(func, geom):
 
 
 def test_get_x():
-    assert pygeos.get_x([point, point_z]).tolist() == [2.0, 1.0]
+    assert pygeos.get_x([point, point_z]).tolist() == [2.0, 2.0]
 
 
 def test_get_y():
-    assert pygeos.get_y([point, point_z]).tolist() == [3.0, 1.0]
+    assert pygeos.get_y([point, point_z]).tolist() == [3.0, 3.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]
+    assert pygeos.get_z([point_z]).tolist() == [4.0]
 
 
 @pytest.mark.skipif(pygeos.geos_version < (3, 7, 0), reason="GEOS < 3.7")
@@ -560,3 +569,73 @@ def test_empty():
     """Compatibility with empty_like, see GH373"""
     g = np.empty_like(np.array([None, None]))
     assert pygeos.is_missing(g).all()
+
+
+# corresponding to geometry_collection_z:
+geometry_collection_2 = pygeos.geometrycollections([point, line_string])
+empty_point_mark = pytest.mark.skipif(
+    pygeos.geos_version < (3, 9, 0),
+    reason="Empty points don't have a dimensionality before GEOS 3.9",
+)
+
+
+@pytest.mark.parametrize(
+    "geom,expected",
+    [
+        (point, point),
+        (point_z, point),
+        pytest.param(empty_point, empty_point, marks=empty_point_mark),
+        pytest.param(empty_point_z, empty_point, marks=empty_point_mark),
+        (line_string, line_string),
+        (line_string_z, line_string),
+        (empty_line_string, empty_line_string),
+        (empty_line_string_z, empty_line_string),
+        (polygon, polygon),
+        (polygon_z, polygon),
+        (polygon_with_hole, polygon_with_hole),
+        (polygon_with_hole_z, polygon_with_hole),
+        (multi_point, multi_point),
+        (multi_point_z, multi_point),
+        (multi_line_string, multi_line_string),
+        (multi_line_string_z, multi_line_string),
+        (multi_polygon, multi_polygon),
+        (multi_polygon_z, multi_polygon),
+        (geometry_collection_2, geometry_collection_2),
+        (geometry_collection_z, geometry_collection_2),
+    ],
+)
+def test_force_2d(geom, expected):
+    actual = pygeos.force_2d(geom)
+    assert pygeos.get_coordinate_dimension(actual) == 2
+    assert_geometries_equal(actual, expected)
+
+
+@pytest.mark.parametrize(
+    "geom,expected",
+    [
+        (point, point_z),
+        (point_z, point_z),
+        pytest.param(empty_point, empty_point_z, marks=empty_point_mark),
+        pytest.param(empty_point_z, empty_point_z, marks=empty_point_mark),
+        (line_string, line_string_z),
+        (line_string_z, line_string_z),
+        (empty_line_string, empty_line_string_z),
+        (empty_line_string_z, empty_line_string_z),
+        (polygon, polygon_z),
+        (polygon_z, polygon_z),
+        (polygon_with_hole, polygon_with_hole_z),
+        (polygon_with_hole_z, polygon_with_hole_z),
+        (multi_point, multi_point_z),
+        (multi_point_z, multi_point_z),
+        (multi_line_string, multi_line_string_z),
+        (multi_line_string_z, multi_line_string_z),
+        (multi_polygon, multi_polygon_z),
+        (multi_polygon_z, multi_polygon_z),
+        (geometry_collection_2, geometry_collection_z),
+        (geometry_collection_z, geometry_collection_z),
+    ],
+)
+def test_force_3d(geom, expected):
+    actual = pygeos.force_3d(geom, z=4)
+    assert pygeos.get_coordinate_dimension(actual) == 3
+    assert_geometries_equal(actual, expected)