diff --git a/rioxarray/_io.py b/rioxarray/_io.py index 151dd415..b875ee3f 100644 --- a/rioxarray/_io.py +++ b/rioxarray/_io.py @@ -22,7 +22,7 @@ from packaging import version from rasterio.errors import NotGeoreferencedWarning from rasterio.vrt import WarpedVRT -from xarray import Dataset, IndexVariable +from xarray import Coordinates, Dataset, IndexVariable from xarray.backends.common import BackendArray from xarray.backends.file_manager import CachingFileManager, FileManager from xarray.backends.locks import SerializableLock @@ -1212,6 +1212,7 @@ def open_rasterio( # Get geospatial coordinates if parse_coordinates: + coords = Coordinates(coords) coords.update( _generate_spatial_coords( affine=riods.transform, width=riods.width, height=riods.height diff --git a/rioxarray/_options.py b/rioxarray/_options.py index 1dc55ffa..9931630e 100644 --- a/rioxarray/_options.py +++ b/rioxarray/_options.py @@ -10,15 +10,18 @@ EXPORT_GRID_MAPPING = "export_grid_mapping" SKIP_MISSING_SPATIAL_DIMS = "skip_missing_spatial_dims" +USE_RASTER_INDEX = "use_raster_index" OPTIONS = { EXPORT_GRID_MAPPING: True, SKIP_MISSING_SPATIAL_DIMS: False, + USE_RASTER_INDEX: False, } OPTION_NAMES = set(OPTIONS) VALIDATORS = { EXPORT_GRID_MAPPING: lambda choice: isinstance(choice, bool), + USE_RASTER_INDEX: lambda choice: isinstance(choice, bool), } @@ -60,6 +63,13 @@ class set_options: # pylint: disable=invalid-name If True, it will not perform spatial operations on variables within a :class:`xarray.Dataset` if the spatial dimensions are not found. + use_raster_index: bool, default=False + If True, this option will generate (lazy) spatial coordinates wrapping the + affine :meth:`~rioxarray.rioxarray.XRasterBase.transform` and will associate + them with a custom ``RasterIndex`` Xarray index. + Otherwise this will generate spatial coordinates using explict values + computed by forward transformation and will associate each of the 1-dimensional + spatial coordinates with a default (pandas) index. Usage as a context manager:: diff --git a/rioxarray/rioxarray.py b/rioxarray/rioxarray.py index e939aab7..f4ae4b8f 100644 --- a/rioxarray/rioxarray.py +++ b/rioxarray/rioxarray.py @@ -121,23 +121,39 @@ def affine_to_coords( def _generate_spatial_coords( *, affine: Affine, width: int, height: int -) -> dict[Hashable, Any]: +) -> xarray.Coordinates: """get spatial coords in new transform""" - new_spatial_coords = affine_to_coords(affine, width, height) - if new_spatial_coords["x"].ndim == 1: - return { - "x": xarray.IndexVariable("x", new_spatial_coords["x"]), - "y": xarray.IndexVariable("y", new_spatial_coords["y"]), - } - return { - "xc": (("y", "x"), new_spatial_coords["x"]), - "yc": (("y", "x"), new_spatial_coords["y"]), - } + + if get_option("use_raster_index"): + try: + from rasterix import RasterIndex + except (ImportError, ModuleNotFoundError) as err: + raise ImportError( + "The rasterix package is needed for generating " + "the spatial coordinates with a RasterIndex" + ) from err + + raster_index = RasterIndex.from_transform(affine, width, height) + return xarray.Coordinates.from_xindex(raster_index) + + else: + new_spatial_coords = affine_to_coords(affine, width, height) + if new_spatial_coords["x"].ndim == 1: + coords = { + "x": xarray.IndexVariable("x", new_spatial_coords["x"]), + "y": xarray.IndexVariable("y", new_spatial_coords["y"]), + } + else: + coords = { + "xc": (("y", "x"), new_spatial_coords["x"]), + "yc": (("y", "x"), new_spatial_coords["y"]), + } + return xarray.Coordinates(coords) def _get_nonspatial_coords( src_data_array: Union[xarray.DataArray, xarray.Dataset] -) -> dict[Hashable, Union[xarray.Variable, xarray.IndexVariable]]: +) -> xarray.Coordinates: coords: dict[Hashable, Union[xarray.Variable, xarray.IndexVariable]] = {} for coord in set(src_data_array.coords) - { src_data_array.rio.x_dim, @@ -162,7 +178,7 @@ def _get_nonspatial_coords( src_data_array[coord].values, src_data_array[coord].attrs, ) - return coords + return xarray.Coordinates(coords) def _make_coords( @@ -172,7 +188,7 @@ def _make_coords( dst_width: int, dst_height: int, force_generate: bool = False, -) -> dict[Hashable, Any]: +) -> xarray.Coordinates: """Generate the coordinates of the new projected `xarray.DataArray`""" coords = _get_nonspatial_coords(src_data_array) if (