ENH: Load and write RPCs #837 (#890)

Author: remi-braun

docs/history.rst

@@ -3,6 +3,7 @@ History
 
 Latest
 ------
+- ENH: RPCs: Load and write RPCs (#837)
 
 0.20.0
 ------

rioxarray/_io.py

@@ -1209,8 +1209,12 @@ def open_rasterio(
         coord_name = "band"
         coords[coord_name] = numpy.asarray(riods.indexes)
 
+    # Handle GCPs and RPCs
     has_gcps = riods.gcps[0]
-    if has_gcps:
+    has_rpcs = riods.rpcs
+
+    # Only parse coordinates in case the array is georeferenced
+    if has_gcps or has_rpcs:
         parse_coordinates = False
 
     # Get geospatial coordinates
@@ -1281,6 +1285,8 @@ def open_rasterio(
         result.rio.write_crs(rio_crs, inplace=True)
     if has_gcps:
         result.rio.write_gcps(*riods.gcps, inplace=True)
+    if has_rpcs:
+        result.rio.write_rpcs(riods.rpcs, inplace=True)
 
     if chunks is not None:
         result = _prepare_dask(

rioxarray/raster_array.py

@@ -458,37 +458,26 @@ def reproject(
                 "CRS not found. Please set the CRS with 'rio.write_crs()'."
                 f"{_get_data_var_message(self._obj)}"
             )
-        gcps = self.get_gcps()
-        if gcps:
-            kwargs.setdefault("gcps", gcps)
 
-        use_affine = (
-            "gcps" not in kwargs
-            and "rpcs" not in kwargs
-            and "src_geoloc_array" not in kwargs
-        )
-        src_affine = None if not use_affine else self.transform(recalc=True)
-        if transform is None:
-            dst_affine, dst_width, dst_height = _make_dst_affine(
-                src_data_array=self._obj,
-                src_crs=self.crs,
-                dst_crs=dst_crs,
-                dst_resolution=resolution,
-                dst_shape=shape,
-                **kwargs,
-            )
-        else:
-            dst_affine = transform
-            if shape is not None:
-                dst_height, dst_width = shape
-            else:
-                dst_height, dst_width = self.shape
+        kwargs = self._reproj_update_kwargs(**kwargs)
+
         if isinstance(resampling, str):
             resampling = _convert_str_to_resampling(resampling)
 
-        dst_data = self._create_dst_data(dst_height=dst_height, dst_width=dst_width)
+        # Get source data from inputs
+        src_affine, use_affine = self._reproj_get_src(**kwargs)
 
-        dst_nodata = self._get_dst_nodata(nodata)
+        # Get destination data from inputs
+        dst_data, dst_height, dst_width, dst_affine, dst_nodata = self._reproj_get_dst(
+            dst_crs=dst_crs,
+            resolution=resolution,
+            shape=shape,
+            transform=transform,
+            nodata=nodata,
+            **kwargs,
+        )
+
+        # Do the reprojection using rasterio
         rasterio.warp.reproject(
             source=self._obj.values,
             destination=dst_data,
@@ -501,8 +490,33 @@ def reproject(
             resampling=resampling,
             **kwargs,
         )
+
+        # Convert the ndarray to a xarray
+        return self._reproj_convert_to_xarray(
+            dst_data=dst_data,
+            dst_nodata=dst_nodata,
+            dst_affine=dst_affine,
+            dst_width=dst_width,
+            dst_height=dst_height,
+            dst_crs=dst_crs,
+            use_affine=use_affine,
+        )
+
+    def _reproj_convert_to_xarray(
+        self,
+        *,
+        dst_data: numpy.ndarray,
+        dst_nodata: float,
+        dst_affine: Affine,
+        dst_width: int,
+        dst_height: int,
+        dst_crs: Any,
+        use_affine: bool,
+    ):
+        """Helper function creating a proper xarray (with correct attributes, etc) from the reprojection output"""
         # add necessary attributes
         new_attrs = _generate_attrs(src_data_array=self._obj, dst_nodata=dst_nodata)
+
         # make sure dimensions with coordinates renamed to x,y
         dst_dims: list[Hashable] = []
         for dim in self._obj.dims:
@@ -529,8 +543,63 @@ def reproject(
         xda.rio.write_transform(dst_affine, inplace=True)
         xda.rio.write_crs(dst_crs, inplace=True)
         xda.rio.write_coordinate_system(inplace=True)
+
         return xda
 
+    def _reproj_update_kwargs(self, **kwargs):
+        """Helper function updating kwargs from internal members"""
+        gcps = self.get_gcps()
+        if gcps:
+            kwargs.setdefault("gcps", gcps)
+
+        rpcs = self.get_rpcs()
+        if rpcs:
+            kwargs.setdefault("rpcs", rpcs)
+
+        return kwargs
+
+    def _reproj_get_src(self, **kwargs):
+        """Helper function creating source data from inputs"""
+        use_affine = (
+            "gcps" not in kwargs
+            and "rpcs" not in kwargs
+            and "src_geoloc_array" not in kwargs
+        )
+        src_affine = None if not use_affine else self.transform(recalc=True)
+        return src_affine, use_affine
+
+    def _reproj_get_dst(
+        self,
+        *,
+        dst_crs: Any,
+        resolution: Optional[Union[float, tuple[float, float]]] = None,
+        shape: Optional[tuple[int, int]] = None,
+        transform: Optional[Affine] = None,
+        nodata: Optional[float] = None,
+        **kwargs,
+    ):
+        """Helper function creating destination data from inputs"""
+        if transform is None:
+            dst_affine, dst_width, dst_height = _make_dst_affine(
+                src_data_array=self._obj,
+                src_crs=self.crs,
+                dst_crs=dst_crs,
+                dst_resolution=resolution,
+                dst_shape=shape,
+                **kwargs,
+            )
+        else:
+            dst_affine = transform
+            if shape is not None:
+                dst_height, dst_width = shape
+            else:
+                dst_height, dst_width = self.shape
+
+        dst_data = self._create_dst_data(dst_height=dst_height, dst_width=dst_width)
+        dst_nodata = self._get_dst_nodata(nodata)
+
+        return dst_data, dst_height, dst_width, dst_affine, dst_nodata
+
     def _get_dst_nodata(self, nodata: Optional[float]) -> Optional[float]:
         default_nodata = (
             _NODATA_DTYPE_MAP.get(dtype_rev[self._obj.dtype.name])
@@ -1194,6 +1263,7 @@ def to_raster(
             crs=self.crs,
             transform=self.transform(recalc=recalc_transform),
             gcps=self.get_gcps(),
+            rpcs=self.get_rpcs(),
             nodata=rio_nodata,
             windowed=windowed,
             lock=lock,

rioxarray/raster_writer.py

@@ -10,6 +10,8 @@
 """
 import numpy
 import rasterio
+import xarray
+from rasterio.rpc import RPC
 from rasterio.windows import Window
 from xarray.conventions import encode_cf_variable
 
@@ -184,6 +186,41 @@ def _get_dtypes(*, rasterio_dtype, encoded_rasterio_dtype, dataarray_dtype):
     return rasterio_dtype, numpy_dtype
 
 
+def _to_raster_dtypes(xarray_dataarray: xarray.DataArray, **kwargs):
+    """Helper function managing dtypes in to_raster, managing '_Unsigned' scenario"""
+    kwargs["dtype"], numpy_dtype = _get_dtypes(
+        rasterio_dtype=kwargs["dtype"],
+        encoded_rasterio_dtype=xarray_dataarray.encoding.get("rasterio_dtype"),
+        dataarray_dtype=xarray_dataarray.encoding.get(
+            "dtype", str(xarray_dataarray.dtype)
+        ),
+    )
+    # there is no equivalent for netCDF _Unsigned
+    # across output GDAL formats. It is safest to convert beforehand.
+    # https://github.com/OSGeo/gdal/issues/6352#issuecomment-1245981837
+    if "_Unsigned" in xarray_dataarray.encoding:
+        unsigned_dtype = _get_unsigned_dtype(
+            unsigned=xarray_dataarray.encoding["_Unsigned"] == "true",
+            dtype=numpy_dtype,
+        )
+        if unsigned_dtype is not None:
+            numpy_dtype = unsigned_dtype
+            kwargs["dtype"] = unsigned_dtype
+            xarray_dataarray.encoding["rasterio_dtype"] = str(unsigned_dtype)
+            xarray_dataarray.encoding["dtype"] = str(unsigned_dtype)
+
+    return xarray_dataarray, numpy_dtype, kwargs
+
+
+def _to_raster_validate_rpcs(**kwargs):
+    """Helper function validating RPCs in to_raster: RPCs should be either be a RPC object or a GDAL-compatible dict"""
+    rpcs = kwargs.get("rpcs")
+    if rpcs is not None:
+        assert isinstance(
+            rpcs, (RPC, dict)
+        ), "RPCs must be of type 'rasterio.rpc.RPC' or dict."
+
+
 class RasterioWriter:
     """
 
@@ -254,26 +291,9 @@ def to_raster(self, *, xarray_dataarray, tags, windowed, lock, compute, **kwargs
             Keyword arguments to pass into writing the raster.
         """
         xarray_dataarray = xarray_dataarray.copy()
-        kwargs["dtype"], numpy_dtype = _get_dtypes(
-            rasterio_dtype=kwargs["dtype"],
-            encoded_rasterio_dtype=xarray_dataarray.encoding.get("rasterio_dtype"),
-            dataarray_dtype=xarray_dataarray.encoding.get(
-                "dtype", str(xarray_dataarray.dtype)
-            ),
+        xarray_dataarray, numpy_dtype, kwargs = _to_raster_dtypes(
+            xarray_dataarray, **kwargs
         )
-        # there is no equivalent for netCDF _Unsigned
-        # across output GDAL formats. It is safest to convert beforehand.
-        # https://github.com/OSGeo/gdal/issues/6352#issuecomment-1245981837
-        if "_Unsigned" in xarray_dataarray.encoding:
-            unsigned_dtype = _get_unsigned_dtype(
-                unsigned=xarray_dataarray.encoding["_Unsigned"] == "true",
-                dtype=numpy_dtype,
-            )
-            if unsigned_dtype is not None:
-                numpy_dtype = unsigned_dtype
-                kwargs["dtype"] = unsigned_dtype
-                xarray_dataarray.encoding["rasterio_dtype"] = str(unsigned_dtype)
-                xarray_dataarray.encoding["dtype"] = str(unsigned_dtype)
 
         if kwargs["nodata"] is not None:
             # Ensure dtype of output data matches the expected dtype.
@@ -283,12 +303,17 @@ def to_raster(self, *, xarray_dataarray, tags, windowed, lock, compute, **kwargs
                 original_nodata=kwargs["nodata"], new_dtype=numpy_dtype
             )
 
+        # Check RPC validity: RPCs should be either be a RPC object or a GDAL-compatible dict
+        _to_raster_validate_rpcs(**kwargs)
+
+        # RPCs and GCPs are propagated through **kwargs
         with rasterio.open(self.raster_path, "w", **kwargs) as rds:
             _write_metatata_to_raster(
                 raster_handle=rds, xarray_dataset=xarray_dataarray, tags=tags
             )
+
             if not (lock and is_dask_collection(xarray_dataarray.data)):
-                # write data to raster immmediately if not dask array
+                # write data to raster immediately if not dask array
                 if windowed:
                     window_iter = rds.block_windows(1)
                 else:

rioxarray/rioxarray.py

@@ -20,6 +20,7 @@
 from pyproj.database import query_utm_crs_info
 from rasterio.control import GroundControlPoint
 from rasterio.crs import CRS
+from rasterio.rpc import RPC
 
 from rioxarray._options import EXPORT_GRID_MAPPING, get_option
 from rioxarray.crs import crs_from_user_input
@@ -301,6 +302,7 @@ def __init__(self, xarray_obj: Union[xarray.DataArray, xarray.Dataset]):
         self._width: Optional[int] = None
         self._crs: Union[rasterio.crs.CRS, None, Literal[False]] = None
         self._gcps: Optional[list[GroundControlPoint]] = None
+        self._rpcs: Optional[RPC] = None
 
     @property
     def crs(self) -> Optional[rasterio.crs.CRS]:
@@ -360,6 +362,7 @@ def _get_obj(self, inplace: bool) -> Union[xarray.Dataset, xarray.DataArray]:
         obj_copy.rio._height = self._height
         obj_copy.rio._crs = self._crs
         obj_copy.rio._gcps = self._gcps
+        obj_copy.rio._rpcs = self._rpcs
         return obj_copy
 
     def set_crs(
@@ -1332,6 +1335,76 @@ def _parse_gcp(gcp) -> GroundControlPoint:
         self._gcps = [_parse_gcp(gcp) for gcp in geojson_gcps["features"]]
         return self._gcps
 
+    def write_rpcs(
+        self,
+        rpcs: RPC,
+        *,
+        grid_mapping_name: Optional[str] = None,
+        inplace: bool = False,
+    ) -> xarray.Dataset | xarray.DataArray:
+        """
+        Write the Rational Polynomial Coefficients to the dataset.
+
+        https://rasterio.readthedocs.io/en/latest/topics/georeferencing.html#rational-polynomial-coefficients
+
+        Parameters
+        ----------
+        rpcs: :obj:`rasterio.rpc.RPC`
+            The Rational Polynomial Coefficients to integrate to the dataset.
+        grid_mapping_name: str, optional
+            Name of the grid_mapping coordinate to store the RPCs information in.
+            Default is the grid_mapping name of the dataset.
+        inplace: bool, optional
+            If True, it will write to the existing dataset. Default is False.
+
+        Returns
+        -------
+        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
+            Modified dataset with Rational Polynomial Coefficients written.
+        """
+        grid_mapping_name = (
+            self.grid_mapping if grid_mapping_name is None else grid_mapping_name
+        )
+        data_obj = self._get_obj(inplace=True)
+
+        # RPC CRS is always 4326
+        data_obj = data_obj.rio.write_crs(
+            "epsg:4326", grid_mapping_name=grid_mapping_name, inplace=inplace
+        )
+        try:
+            grid_map_attrs = data_obj.coords[grid_mapping_name].attrs.copy()
+        except KeyError:
+            data_obj.coords[grid_mapping_name] = xarray.Variable((), 0)
+            grid_map_attrs = data_obj.coords[grid_mapping_name].attrs.copy()
+
+        # Store the RPCCs
+        grid_map_attrs["rpcs"] = json.dumps(rpcs.to_dict())
+        data_obj.coords[grid_mapping_name].rio.set_attrs(grid_map_attrs, inplace=True)
+        self._rpcs = rpcs
+
+        return data_obj
+
+    def get_rpcs(self) -> Optional[RPC]:
+        """
+        Get the Rational Polynomial Coefficients from the dataset.
+
+        https://rasterio.readthedocs.io/en/latest/topics/georeferencing.html#rational-polynomial-coefficients
+
+        Returns
+        -------
+        :obj:`rasterio.rpc.RPC` or None
+            The Rational Polynomial Coefficients from the dataset or None if not applicable
+        """
+        if self._rpcs is not None:
+            return self._rpcs
+        try:
+            json_rpcs = json.loads(self._obj.coords[self.grid_mapping].attrs["rpcs"])
+        except (KeyError, AttributeError):
+            return None
+
+        self._rpcs = RPC(**json_rpcs)
+        return self._rpcs
+
 
 def _convert_gcps_to_geojson(
     gcps: Iterable[GroundControlPoint],