XProj integration (#120)

benbovy · web-flow · commit 6c897b09ef7c · 2025-06-12T15:38:11.000+02:00
* add xproj as a required dependency

* remove optional xproj checks

Note: with HAS_XPROJ checks removed it is now possible that accessing
the CRS via `.proj.crs` returns an error (when multiple CRSs are found
in the Xarray object) instead of setting `None` (when HAS_XPROJ was
False). This is possible in theory but quite unlikely in practice,
though.

* make GeometryIndex an xproj's CRS-aware index

GeometryIndex now inherits from `xproj.ProjIndexMixin`.

Reuse functionality that has been implemented in xproj (moved from
xvec). Some warning and/or error messages may slightly change, although
I tried to keep them clear and meaningful.

* more xproj integration

Set or convert the CRS of a GeometryIndex via XProj's API.
diff --git a/pyproject.toml b/pyproject.toml
@@ -31,6 +31,7 @@ dependencies = [
     "pyproj >= 3.0.0",
     "shapely >= 2.0b1",
     "cf_xarray >= 0.9.2",
+    "xproj >= 0.2.0",
 ]
 
 [project.optional-dependencies]
@@ -69,4 +70,4 @@ ignore = ['E501', 'Q000', 'Q001', 'Q002', 'Q003', 'W191', 'C408']
 select = ["E", "F", "W", "I", "UP", "B", "A", "C4", "Q"]
 
 [tool.ruff.lint.per-file-ignores]
-"doc/**" = ["F401"]
+"doc/**" = ["F401"]
diff --git a/xvec/accessor.py b/xvec/accessor.py
@@ -9,10 +9,12 @@
 import pandas as pd
 import shapely
 import xarray as xr
-from pyproj import CRS, Transformer
+import xproj  # noqa: F401
+from pyproj import CRS
 
 from .index import GeometryIndex
 from .plotting import _plot
+from .utils import transform_geom
 from .zonal import (
     _variable_zonal,
     _zonal_stats_exactextract,
@@ -23,13 +25,6 @@
 if TYPE_CHECKING:
     from geopandas import GeoDataFrame
 
-try:
-    import xproj  # noqa: F401
-
-    HAS_XPROJ = True
-except ImportError:
-    HAS_XPROJ = False
-
 
 @xr.register_dataarray_accessor("xvec")
 @xr.register_dataset_accessor("xvec")
@@ -361,7 +356,7 @@ def to_crs(
                     "handling projection information."
                 )
 
-            data = _obj[key]
+            data = _obj[key].data
             data_crs = self._obj.xindexes[key].crs  # type: ignore
 
             # transformation code taken from geopandas (BSD 3-clause license)
@@ -376,25 +371,7 @@ def to_crs(
             if data_crs.is_exact_same(crs):
                 pass
 
-            transformer = Transformer.from_crs(data_crs, crs, always_xy=True)
-
-            has_z = shapely.has_z(data)
-
-            result = np.empty_like(data)
-
-            coordinates = shapely.get_coordinates(data[~has_z], include_z=False)
-            new_coords = transformer.transform(coordinates[:, 0], coordinates[:, 1])
-            result[~has_z] = shapely.set_coordinates(
-                data[~has_z].copy(), np.array(new_coords).T
-            )
-
-            coords_z = shapely.get_coordinates(data[has_z], include_z=True)
-            new_coords_z = transformer.transform(
-                coords_z[:, 0], coords_z[:, 1], coords_z[:, 2]
-            )
-            result[has_z] = shapely.set_coordinates(
-                data[has_z].copy(), np.array(new_coords_z).T
-            )
+            result = transform_geom(data, data_crs, crs)
 
             transformed[key] = (result, crs)
 
@@ -973,9 +950,7 @@ def to_geodataframe(
 
         if geometry is not None:
             if geometry not in self._geom_coords_all:  # variable geometry
-                return df.set_geometry(
-                    geometry, crs=self._obj.proj.crs if HAS_XPROJ else None
-                )
+                return df.set_geometry(geometry, crs=self._obj.proj.crs)
 
             # coordinate geometry
             return df.set_geometry(
@@ -987,7 +962,7 @@ def to_geodataframe(
             name if name else (self._obj.name if hasattr(self._obj, "name") else None)
         )
         if name is not None and shapely.is_valid_input(df[name]).all():
-            return df.set_geometry(name, crs=self._obj.proj.crs if HAS_XPROJ else None)
+            return df.set_geometry(name, crs=self._obj.proj.crs)
 
         warnings.warn(
             "No active geometry column to be set. The resulting object "
@@ -1525,15 +1500,15 @@ def encode_wkb(self) -> xr.DataArray | xr.Dataset:
         if isinstance(obj, xr.DataArray):
             if np.all(shapely.is_valid_input(obj.data)):
                 obj = shapely.to_wkb(obj)
-                if HAS_XPROJ and obj.proj.crs:
+                if obj.proj.crs:
                     obj.attrs["crs"] = obj.proj.crs.to_json()
                 obj.attrs["wkb_encoded_geometry"] = True
 
         else:
             for data in obj.data_vars:
                 if np.all(shapely.is_valid_input(obj[data].data)):
                     obj[data] = shapely.to_wkb(obj[data])
-                    if HAS_XPROJ and obj[data].proj.crs:
+                    if obj[data].proj.crs:
                         obj[data].attrs["crs"] = obj[data].proj.crs.to_json()
                     obj[data].attrs["wkb_encoded_geometry"] = True
 
@@ -1565,7 +1540,7 @@ def decode_wkb(self) -> xr.DataArray | xr.Dataset:
         if isinstance(obj, xr.DataArray):
             if obj.attrs.get("wkb_encoded_geometry", False):
                 obj.data = shapely.from_wkb(obj)
-                if HAS_XPROJ and "crs" in obj.attrs:
+                if "crs" in obj.attrs:
                     obj = obj.proj.assign_crs(
                         spatial_ref=json.loads(obj.attrs.pop("crs")),
                         allow_override=True,
@@ -1576,7 +1551,7 @@ def decode_wkb(self) -> xr.DataArray | xr.Dataset:
             for data in obj.data_vars:
                 if obj[data].attrs.get("wkb_encoded_geometry", False):
                     obj[data].data = shapely.from_wkb(obj[data])
-                    if HAS_XPROJ and "crs" in obj[data].attrs:
+                    if "crs" in obj[data].attrs:
                         obj = obj.proj.assign_crs(
                             spatial_ref=json.loads(obj[data].attrs.pop("crs")),
                             allow_override=True,
@@ -1675,9 +1650,7 @@ def _union(x: xr.DataArray) -> xr.DataArray:
         return (
             self._obj.assign_coords(summary_geometry=(dim, summary))
             .set_xindex("summary_geometry")
-            .xvec.set_geom_indexes(
-                "summary_geometry", crs=self._obj.proj.crs if HAS_XPROJ else None
-            )
+            .xvec.set_geom_indexes("summary_geometry", crs=self._obj.proj.crs)
         )
 
     def plot(
diff --git a/xvec/index.py b/xvec/index.py
@@ -7,46 +7,16 @@
 import numpy as np
 import pandas as pd
 import shapely
+import xproj
 from pyproj import CRS
 from xarray import DataArray, Variable, get_options
 from xarray.core.indexing import IndexSelResult
 from xarray.indexes import Index, PandasIndex
 
+from xvec.utils import transform_geom
 
-def _format_crs(crs: CRS | None, max_width: int = 50) -> str:
-    if crs is not None:
-        srs = crs.to_string()
-    else:
-        srs = "None"
 
-    return srs if len(srs) <= max_width else " ".join([srs[:max_width], "..."])
-
-
-def _get_common_crs(crs_set: set[CRS | None]) -> CRS | None:
-    # code taken from geopandas (BSD-3 Licence)
-
-    crs_not_none = [crs for crs in crs_set if crs is not None]
-    names = [crs.name for crs in crs_not_none]
-
-    if len(crs_not_none) == 0:
-        return None
-    if len(crs_not_none) == 1:
-        if len(crs_set) != 1:
-            warnings.warn(  # noqa: B028
-                "CRS not set for some of the concatenation inputs. "
-                f"Setting output's CRS as {names[0]} "
-                "(the single non-null CRS provided).",
-                stacklevel=3,
-            )
-        return crs_not_none[0]
-
-    raise ValueError(
-        f"cannot determine common CRS for concatenation inputs, got {names}. "
-        # "Use `to_crs()` to transform geometries to the same CRS before merging."
-    )
-
-
-class GeometryIndex(Index):
+class GeometryIndex(Index, xproj.ProjIndexMixin):
     """An CRS-aware, Xarray-compatible index for vector geometries.
 
     This index can be set from any 1-dimensional coordinate of
@@ -89,6 +59,30 @@ def crs(self) -> CRS | None:
         """
         return self._crs
 
+    def _proj_set_crs(
+        self: GeometryIndex, spatial_ref: Hashable, crs: CRS
+    ) -> GeometryIndex:
+        # Returns a geometry index shallow copy with a replaced CRS, without transformation
+        # (XProj integration via xproj.ProjIndexMixin)
+        # Note: XProj already handles the case of overriding any existing CRS
+        return GeometryIndex(self._index, crs=crs)
+
+    def _proj_to_crs(
+        self: GeometryIndex, spatial_ref: Hashable, crs: CRS
+    ) -> GeometryIndex:
+        # Returns a new geometry index with a replaced CRS and transformed geometries
+        # (XProj integration via xproj.ProjIndexMixin)
+        # Note: XProj already handles the case of overriding any existing CRS
+
+        # XProj redirects to `._proj_set_crs()` if this index's CRS is undefined
+        assert self.crs is not None
+
+        result = transform_geom(np.asarray(self._index.index), self.crs, crs)
+        index = PandasIndex(
+            result, self._index.dim, coord_dtype=self._index.coord_dtype
+        )
+        return GeometryIndex(index, crs=crs)
+
     @property
     def sindex(self) -> shapely.STRtree:
         """Returns the spatial index, i.e., a :class:`shapely.STRtree` object.
@@ -99,25 +93,14 @@ def sindex(self) -> shapely.STRtree:
             self._sindex = shapely.STRtree(self._index.index)
         return self._sindex
 
-    def _check_crs(self, other_crs: CRS | None, allow_none: bool = False) -> bool:
-        """Check if the index's projection is the same than the given one.
-        If allow_none is True, empty CRS is treated as the same.
-        """
-        if allow_none:
-            if self.crs is None or other_crs is None:
-                return True
-        if not self.crs == other_crs:
-            return False
-        return True
-
     def _crs_mismatch_raise(
         self, other_crs: CRS | None, warn: bool = False, stacklevel: int = 3
     ) -> None:
         """Raise a CRS mismatch error or warning with the information
         on the assigned CRS.
         """
-        srs = _format_crs(self.crs, max_width=50)
-        other_srs = _format_crs(other_crs, max_width=50)
+        srs = xproj.format_crs(self.crs, max_width=50)
+        other_srs = xproj.format_crs(other_crs, max_width=50)
 
         # TODO: expand message with reproject suggestion
         msg = (
@@ -152,7 +135,7 @@ def concat(
         positions: Iterable[Iterable[int]] | None = None,
     ) -> GeometryIndex:
         crs_set = {idx.crs for idx in indexes}
-        crs = _get_common_crs(crs_set)
+        crs = xproj.get_common_crs(crs_set)
 
         indexes_ = [idx._index for idx in indexes]
         index = PandasIndex.concat(indexes_, dim, positions)
@@ -232,14 +215,14 @@ def equals(
     ) -> bool:
         if not isinstance(other, GeometryIndex):
             return False
-        if not self._check_crs(other.crs, allow_none=True):
+        if not self._proj_crs_equals(other, allow_none=True):
             return False
         return self._index.equals(other._index, exclude=exclude)
 
     def join(
         self: GeometryIndex, other: GeometryIndex, how: str = "inner"
     ) -> GeometryIndex:
-        if not self._check_crs(other.crs, allow_none=True):
+        if not self._proj_crs_equals(other, allow_none=True):
             self._crs_mismatch_raise(other.crs)
 
         index = self._index.join(other._index, how=how)
@@ -251,7 +234,7 @@ def reindex_like(
         method: str | None = None,
         tolerance: int | float | Iterable[int | float] | None = None,
     ) -> dict[Hashable, Any]:
-        if not self._check_crs(other.crs, allow_none=True):
+        if not self._proj_crs_equals(other, allow_none=True):
             self._crs_mismatch_raise(other.crs)
 
         return self._index.reindex_like(
@@ -273,11 +256,11 @@ def _repr_inline_(self, max_width: int) -> str:
         if max_width is None:
             max_width = get_options()["display_width"]
 
-        srs = _format_crs(self.crs, max_width=max_width)
+        srs = xproj.format_crs(self.crs, max_width=max_width)
         return f"{self.__class__.__name__} (crs={srs})"
 
     def __repr__(self) -> str:
-        srs = _format_crs(self.crs)
+        srs = xproj.format_crs(self.crs)
         shape = self._index.index.shape[0]
         if shape == 0:
             return f"GeometryIndex([], crs={srs})"
diff --git a/xvec/plotting.py b/xvec/plotting.py
@@ -4,15 +4,9 @@
 import numpy as np
 import shapely
 import xarray as xr
+import xproj  # noqa: F401
 from xarray.plot.utils import _determine_cmap_params, label_from_attrs
 
-try:
-    import xproj  # noqa: F401
-
-    HAS_XPROJ = True
-except ImportError:
-    HAS_XPROJ = False
-
 
 def _setup_axes(n_rows, n_cols, arr, geometry, crs, subplot_kws, figsize):
     import matplotlib.pyplot as plt
@@ -67,15 +61,9 @@ def _get_crs(arr, geometry=None):
         if geometry:
             return arr[geometry].crs
         elif np.all(shapely.is_valid_input(arr.data)):
-            return arr.proj.crs if HAS_XPROJ else None
+            return arr.proj.crs
         return arr.xindexes[list(arr.xvec._geom_coords_all)[0]].crs
-    return (
-        arr[geometry].crs
-        if hasattr(arr[geometry], "crs")
-        else arr.proj.crs
-        if HAS_XPROJ
-        else None
-    )
+    return arr[geometry].crs if hasattr(arr[geometry], "crs") else arr.proj.crs
 
 
 def _setup_legend(fig, cmap_params, label=None):
diff --git a/xvec/tests/test_index.py b/xvec/tests/test_index.py
@@ -51,9 +51,7 @@ def test_concat(geom_dataset, geom_array, geom_dataset_no_index, geom_dataset_no
     xr.testing.assert_identical(actual, expected)
 
     # mixed CRS / no CRS
-    with pytest.warns(
-        UserWarning, match="CRS not set for some of the concatenation inputs"
-    ):
+    with pytest.warns(UserWarning, match="CRS is undefined for some of the inputs"):
         xr.concat([geom_dataset, geom_dataset_no_crs], "geom")
 
 
diff --git a/xvec/tests/test_xproj_integration.py b/xvec/tests/test_xproj_integration.py
@@ -0,0 +1,26 @@
+import shapely
+import xproj  # noqa: F401
+
+
+def test_xproj_crs(geom_dataset):
+    assert geom_dataset.xindexes["geom"].crs.equals(geom_dataset.proj("geom").crs)
+
+
+def test_xproj_map_crs(geom_dataset_no_crs):
+    ds_with_crsindex = geom_dataset_no_crs.proj.assign_crs(spatial_ref=26915)
+
+    # set crs
+    ds = ds_with_crsindex.proj.map_crs(spatial_ref=["geom"])
+    assert ds.proj("geom").crs.equals(ds.proj("spatial_ref").crs)
+
+    # to crs
+    geom_array_4326 = shapely.from_wkt(
+        ["POINT (-97.488735 0.000018)", "POINT (-97.488717 0.000036)"]
+    )
+
+    ds_with_crsindex2 = ds.proj.assign_crs(spatial_ref=4326, allow_override=True)
+    ds2 = ds_with_crsindex2.proj.map_crs(
+        spatial_ref=["geom"], transform=True, allow_override=True
+    )
+    assert ds2.proj("geom").crs.equals(ds2.proj("spatial_ref").crs)
+    assert shapely.equals_exact(geom_array_4326, ds2.geom, tolerance=1e-5).all()
diff --git a/xvec/utils.py b/xvec/utils.py
