From 4b72fb417a4e31ad4ffde7df01927faf54ac362b Mon Sep 17 00:00:00 2001
From: Joris Van den Bossche <jorisvandenbossche@gmail.com>
Date: Tue, 6 Dec 2022 21:29:10 +0100
Subject: [PATCH] BUG: Handle z coordinates in to_crs for Shapely 2.0 / PyGEOS
 (#2641)

* Handle z coordinates in to_crs for Shapely 2.0 / PyGEOS

* handle mixed dimensions

* add explicit test for empty array

* also support pygeos

* remove include_z

* add whatsnew
---
 CHANGELOG.md                |  5 +++++
 geopandas/_vectorized.py    | 27 ++++++++++++++++++---------
 geopandas/tests/test_crs.py | 27 +++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bd6ede1..f56072d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,11 @@
 ## Development version
 
 
+Bug fixes:
+
+- Correctly handle geometries with Z dimension in ``to_crs()`` when using PyGEOS or
+  Shapely >= 2.0 (previously the z coordinates were lost) (#1345).
+
 ## Version 0.12.1 (October 29, 2022)
 
 Small bug-fix release removing the shapely<2 pin in the installation requirements.
diff --git a/geopandas/_vectorized.py b/geopandas/_vectorized.py
index 45d383d..b1700b4 100644
--- a/geopandas/_vectorized.py
+++ b/geopandas/_vectorized.py
@@ -1096,15 +1096,24 @@ def bounds(data):
 
 
 def transform(data, func):
-    if compat.USE_SHAPELY_20:
-        coords = shapely.get_coordinates(data)
-        new_coords = func(coords[:, 0], coords[:, 1])
-        result = shapely.set_coordinates(data.copy(), np.array(new_coords).T)
-        return result
-    if compat.USE_PYGEOS:
-        coords = pygeos.get_coordinates(data)
-        new_coords = func(coords[:, 0], coords[:, 1])
-        result = pygeos.set_coordinates(data.copy(), np.array(new_coords).T)
+    if compat.USE_SHAPELY_20 or compat.USE_PYGEOS:
+        if compat.USE_SHAPELY_20:
+            has_z = shapely.has_z(data)
+            from shapely import get_coordinates, set_coordinates
+        else:
+            has_z = pygeos.has_z(data)
+            from pygeos import get_coordinates, set_coordinates
+
+        result = np.empty_like(data)
+
+        coords = get_coordinates(data[~has_z], include_z=False)
+        new_coords_z = func(coords[:, 0], coords[:, 1])
+        result[~has_z] = set_coordinates(data[~has_z].copy(), np.array(new_coords_z).T)
+
+        coords_z = get_coordinates(data[has_z], include_z=True)
+        new_coords_z = func(coords_z[:, 0], coords_z[:, 1], coords_z[:, 2])
+        result[has_z] = set_coordinates(data[has_z].copy(), np.array(new_coords_z).T)
+
         return result
     else:
         from shapely.ops import transform
diff --git a/geopandas/tests/test_crs.py b/geopandas/tests/test_crs.py
index 1d9200a..b8983a6 100644
--- a/geopandas/tests/test_crs.py
+++ b/geopandas/tests/test_crs.py
@@ -58,6 +58,13 @@ def test_to_crs_transform__missing_data():
     assert_geodataframe_equal(df, utm, check_less_precise=True)
 
 
+def test_to_crs_transform__empty_data():
+    df = df_epsg26918().iloc[:0]
+    lonlat = df.to_crs(epsg=4326)
+    utm = lonlat.to_crs(epsg=26918)
+    assert_geodataframe_equal(df, utm, check_less_precise=True)
+
+
 def test_to_crs_inplace():
     df = df_epsg26918()
     lonlat = df.to_crs(epsg=4326)
@@ -77,6 +84,26 @@ def test_to_crs_geo_column_name():
     assert_geodataframe_equal(df, utm, check_less_precise=True)
 
 
+def test_to_crs_dimension_z():
+    # preserve z dimension
+    arr = points_from_xy([1, 2], [2, 3], [3, 4], crs=4326)
+    assert arr.has_z.all()
+    result = arr.to_crs(epsg=3857)
+    assert result.has_z.all()
+
+
+def test_to_crs_dimension_mixed():
+    s = GeoSeries([Point(1, 2), LineString([(1, 2, 3), (4, 5, 6)])], crs=2056)
+    result = s.to_crs(epsg=4326)
+    assert not result[0].is_empty
+    assert result.has_z.tolist() == [False, True]
+    roundtrip = result.to_crs(epsg=2056)
+    # TODO replace with assert_geoseries_equal once we expose tolerance keyword
+    # assert_geoseries_equal(roundtrip, s, check_less_precise=True)
+    for a, b in zip(roundtrip, s):
+        np.testing.assert_allclose(a.coords[:], b.coords[:], atol=0.01)
+
+
 # -----------------------------------------------------------------------------
 # Test different supported formats for CRS specification