ENH: add projection support in geoarrow output

CMakeLists.txt

@@ -80,7 +80,7 @@ set(CPP_SOURCES
   src/spherely.cpp)
 
 if(${s2geography_VERSION} VERSION_GREATER_EQUAL "0.2.0")
-    set(CPP_SOURCES ${CPP_SOURCES} src/geoarrow.cpp)
+    set(CPP_SOURCES ${CPP_SOURCES} src/geoarrow.cpp src/projections.cpp)
 endif()
 
 add_library(spherely MODULE ${CPP_SOURCES})

src/geoarrow.cpp

@@ -4,6 +4,7 @@
 #include "constants.hpp"
 #include "creation.hpp"
 #include "geography.hpp"
+#include "projections.hpp"
 #include "pybind11.hpp"
 
 namespace py = pybind11;
@@ -197,6 +198,7 @@ class ArrowArrayHolder {
 
 ArrowArrayHolder to_geoarrow(py::array_t<PyObjectGeography> input,
                              py::object output_schema,
+                             Projection projection,
                              bool planar,
                              float tessellate_tolerance,
                              int precision) {
@@ -207,11 +209,10 @@ ArrowArrayHolder to_geoarrow(py::array_t<PyObjectGeography> input,
 
     s2geog::geoarrow::ExportOptions options;
     options.set_precision(precision);
+    options.set_projection(projection.s2_projection());
     if (planar) {
         auto tol = S1Angle::Radians(tessellate_tolerance / EARTH_RADIUS_METERS);
         options.set_tessellate_tolerance(tol);
-        // TODO make this configurable
-        // options.set_projection(s2geog::geoarrow::mercator());
     }
 
     if (!output_schema.is(py::none())) {
@@ -315,6 +316,7 @@ void init_geoarrow(py::module& m) {
           py::pos_only(),
           py::kw_only(),
           py::arg("output_schema") = py::none(),
+          py::arg("projection") = Projection::lnglat(),
           py::arg("planar") = false,
           py::arg("tessellate_tolerance") = 100.0,
           py::arg("precision") = 6,
@@ -338,6 +340,10 @@ void init_geoarrow(py::module& m) {
             can be a ``pyarrow.DataType`` or ``pyarrow.Field``, and you can
             use the ``geoarrow.pyarrow`` package to construct such geoarrow
             extension types.
+        projection : spherely.Projection, default Projection.lnglat()
+            The projection to use when converting the geographies to the output
+            encoding. By default, uses longitude/latitude coordinates ("plate
+            carree" projection).
         planar : bool, default False
             If set to True, the edges of linestrings and polygons in the output
             are assumed to be linear on the plane. In that case, additional

src/projections.cpp

@@ -0,0 +1,35 @@
+#include "projections.hpp"
+
+#include <s2/s2projections.h>
+#include <s2geography.h>
+
+#include "geography.hpp"
+#include "pybind11.hpp"
+
+namespace py = pybind11;
+namespace s2geog = s2geography;
+using namespace spherely;
+
+// std::tuple<double, double> project_mercator(py::object obj) {
+//     auto geog = (static_cast<PyObjectGeography&>(obj)).as_geog_ptr();
+//     if (geog->geog_type() != GeographyType::Point) {
+//         throw py::value_error("test");
+//     }
+//     auto point = static_cast<Point*>(geog);
+//     auto s2point = point->s2point();
+
+//     auto projection = S2::MercatorProjection(20037508.3427892);
+//     auto r2point = projection.Project(s2point);
+//     double x = r2point.x();
+//     double y = r2point.y();
+
+//     return std::make_tuple(x, y);
+// }
+
+void init_projections(py::module& m) {
+    // m.def("project_mercator", &project_mercator);
+    py::class_<Projection>(m, "Projection")
+        .def("lnglat", &Projection::lnglat)
+        .def("pseudo_mercator", &Projection::pseudo_mercator)
+        .def("orthographic", &Projection::orthographic);
+}

src/projections.hpp

@@ -0,0 +1,38 @@
+#ifndef SPHERELY_PROJECTIONS_H_
+#define SPHERELY_PROJECTIONS_H_
+
+#include <s2/s2latlng.h>
+#include <s2/s2projections.h>
+#include <s2geography.h>
+
+#include "pybind11.hpp"
+
+namespace py = pybind11;
+namespace s2geog = s2geography;
+using namespace spherely;
+
+class Projection {
+public:
+    Projection(std::shared_ptr<S2::Projection> projection)
+        : m_s2_projection(std::move(projection)) {}
+
+    std::shared_ptr<S2::Projection> s2_projection() {
+        return m_s2_projection;
+    }
+
+    static Projection lnglat() {
+        return Projection(std::move(s2geog::lnglat()));
+    }
+    static Projection pseudo_mercator() {
+        return Projection(std::move(s2geog::pseudo_mercator()));
+    }
+    static Projection orthographic(double longitude, double latitude) {
+        return Projection(
+            std::move(s2geog::orthographic(S2LatLng::FromDegrees(latitude, longitude))));
+    }
+
+private:
+    std::shared_ptr<S2::Projection> m_s2_projection;
+};
+
+#endif  // SPHERELY_PROJECTIONS_H_

src/spherely.cpp

@@ -13,6 +13,7 @@ void init_io(py::module&);
 #if defined(S2GEOGRAPHY_VERSION_MAJOR) && \
     (S2GEOGRAPHY_VERSION_MAJOR >= 1 || S2GEOGRAPHY_VERSION_MINOR >= 2)
 void init_geoarrow(py::module&);
+void init_projections(py::module&);
 #endif
 
 PYBIND11_MODULE(spherely, m) {
@@ -31,6 +32,7 @@ PYBIND11_MODULE(spherely, m) {
     init_io(m);
 #if defined(S2GEOGRAPHY_VERSION_MAJOR) && \
     (S2GEOGRAPHY_VERSION_MAJOR >= 1 || S2GEOGRAPHY_VERSION_MINOR >= 2)
+    init_projections(m);
     init_geoarrow(m);
 #endif
 

tests/test_geoarrow.py

@@ -182,3 +182,55 @@ def test_to_geoarrow_invalid_output_schema():
 
     with pytest.raises(ValueError, match="Did you pass a valid schema"):
         spherely.to_geoarrow(arr, output_schema=pa.schema([("test", pa.int64())]))
+
+
+def test_to_geoarrow_projected():
+    arr = spherely.points([1, 2, 3], [1, 2, 3])
+    point_schema = ga.point().with_coord_type(ga.CoordType.INTERLEAVED)
+    result = pa.array(
+        spherely.to_geoarrow(
+            arr, output_schema=point_schema, projection=spherely.Projection.lnglat()
+        )
+    )
+
+    coords = np.asarray(result.storage.values)
+    expected = np.array([1, 1, 2, 2, 3, 3], dtype="float64")
+    np.testing.assert_allclose(coords, expected)
+
+    # import pyproj
+    # trans = pyproj.Transformer.from_crs("EPSG:4326", "EPSG:3857", always_xy=True)
+    # trans.transform([1, 2, 3], [1, 2, 3])
+    result = pa.array(
+        spherely.to_geoarrow(
+            arr,
+            output_schema=point_schema,
+            projection=spherely.Projection.pseudo_mercator(),
+        )
+    )
+    coords = np.asarray(result.storage.values)
+    expected = np.array(
+        [
+            111319.49079327357,
+            111325.1428663851,
+            222638.98158654713,
+            222684.20850554405,
+            333958.4723798207,
+            334111.1714019596,
+        ],
+        dtype="float64",
+    )
+    np.testing.assert_allclose(coords, expected)
+
+    result = pa.array(
+        spherely.to_geoarrow(
+            arr,
+            output_schema=point_schema,
+            projection=spherely.Projection.orthographic(0.0, 0.0),
+        )
+    )
+    coords = np.asarray(result.storage.values)
+    expected = np.array(
+        [0.01744975, 0.01745241, 0.03487824, 0.0348995, 0.05226423, 0.05233596],
+        dtype="float64",
+    )
+    np.testing.assert_allclose(coords, expected, rtol=1e-06)