Add Xarray accessor mirroring Raster class

dev-environment.yml

@@ -3,14 +3,18 @@ channels:
   - conda-forge
 dependencies:
   - python>=3.10,<3.15
+
+  # Required dependencies
   - geopandas>=0.12.0
-  - pyproj=3.*
   - rasterio>=1.3
-  - pandas>=1,<3
-  - numpy>=1,<3
   - scipy=1.*
   - xarray>2023
   - rioxarray=0.*
+
+  # Depend on the above
+  - pyproj=3.*
+  - pandas>=1,<3
+  - numpy>=1,<3
   - affine
   - shapely
   - pyogrio

environment.yml

@@ -4,7 +4,7 @@ channels:
 dependencies:
   - python>=3.10,<3.15
 
-  # Core dependencies
+  # Required dependencies
   - geopandas>=0.12.0
   - rasterio>=1.3
   - scipy=1.*

geoutils/__init__.py

@@ -23,7 +23,8 @@
 from geoutils import examples, projtools  # noqa
 from geoutils._config import config  # noqa
 
-from geoutils.raster import Raster  # noqa isort:skip
+from geoutils.raster import Raster, xr_accessor  # noqa isort:skip
+from geoutils.raster.xr_accessor import open_raster  # noqa isort:skip
 from geoutils.vector import Vector  # noqa isort:skip
 from geoutils.pointcloud import PointCloud  # noqa isort:skip
 

geoutils/interface/distance.py

@@ -93,7 +93,15 @@ def _proximity_from_vector_or_raster(
     # If not all pixels are targets, then we compute the distance
     non_targets = np.count_nonzero(mask_boundary)
     if non_targets > 0:
-        proximity = distance_transform_edt(~mask_boundary, sampling=sampling)
+        # For multi-band raster, loop over bands
+        if mask_boundary.ndim == 3:
+            list_band_proxi = []
+            for i in range(mask_boundary.shape[0]):
+                prox = distance_transform_edt(~mask_boundary[i, :, :], sampling=sampling)
+                list_band_proxi.append(prox)
+            proximity = np.stack(list_band_proxi, axis=0)
+        else:
+            proximity = distance_transform_edt(~mask_boundary, sampling=sampling)
     # Otherwise, pass an array full of nodata
     else:
         proximity = np.ones(np.shape(mask_boundary)) * np.nan

geoutils/interface/raster_point.py

@@ -185,14 +185,14 @@ def _raster_to_pointcloud(
 
     # If subsample is the entire array, load it to optimize speed
     if subsample == 1 and not source_raster.is_loaded:
-        source_raster.load(bands=all_bands)
+        source_raster.load()
 
     # Band indexes in the array are band number minus one
     all_indexes = [b - 1 for b in all_bands]
 
     # We do 2D subsampling on the data band only, regardless of valid masks on other bands
     if skip_nodata:
-        if source_raster.is_loaded:
+        if source_raster.is_loaded or source_raster._is_xr:
             if source_raster.count == 1:
                 self_mask = get_mask_from_array(
                     source_raster.data
@@ -220,7 +220,7 @@ def _raster_to_pointcloud(
     indices = subsample_array(array=ma_valid, subsample=subsample, random_state=random_state, return_indices=True)
 
     # If the Raster is loaded, pick from the data while ignoring the mask
-    if source_raster.is_loaded:
+    if source_raster.is_loaded or source_raster._is_xr:
         if source_raster.count == 1:
             pixel_data = source_raster.data[indices[0], indices[1]]
         else:
@@ -234,7 +234,7 @@ def _raster_to_pointcloud(
         # Further below we redefine output coordinates based on point interpretation
         x_coords, y_coords = (np.array(a) for a in source_raster.ij2xy(indices[0], indices[1], force_offset="ul"))
 
-        with rio.open(source_raster.filename) as raster:
+        with rio.open(source_raster.name) as raster:
             # Rasterio uses indexes (starts at 1)
             pixel_data = np.array(list(raster.sample(zip(x_coords, y_coords), indexes=all_bands))).T
 

geoutils/interface/raster_vector.py

@@ -26,8 +26,8 @@
 import affine
 import geopandas as gpd
 import numpy as np
-import rasterio
 import rasterio as rio
+import xarray as xr
 from rasterio import features, warp
 from rasterio.crs import CRS
 from rasterio.features import shapes
@@ -53,31 +53,33 @@ def _polygonize(
             warnings.warn("Raster mask (boolean type) passed, using target value of 1 (True).")
         target_values = True
 
+    nanarray = source_raster.get_nanarray()
+
     # Mask a unique value set by a number
     if isinstance(target_values, (int, float, np.integer, np.floating)):
-        if np.sum(source_raster.data == target_values) == 0:
+        if np.sum(nanarray == target_values) == 0:
             raise ValueError(f"no pixel with in_value {target_values}")
 
-        bool_msk = np.array(source_raster.data == target_values).astype(np.uint8)
+        bool_msk = np.array(nanarray == target_values).astype(np.uint8)
 
     # Mask values within boundaries set by a tuple
     elif isinstance(target_values, tuple):
-        if np.sum((source_raster.data > target_values[0]) & (source_raster.data < target_values[1])) == 0:
+        if np.sum((nanarray > target_values[0]) & (nanarray < target_values[1])) == 0:
             raise ValueError(f"no pixel with in_value between {target_values[0]} and {target_values[1]}")
 
-        bool_msk = ((source_raster.data > target_values[0]) & (source_raster.data < target_values[1])).astype(np.uint8)
+        bool_msk = ((nanarray > target_values[0]) & (nanarray < target_values[1])).astype(np.uint8)
 
     # Mask specific values set by a sequence
     elif isinstance(target_values, list) or isinstance(target_values, np.ndarray):
-        if np.sum(np.isin(source_raster.data, np.array(target_values))) == 0:
+        if np.sum(np.isin(nanarray, np.array(target_values))) == 0:
             raise ValueError("no pixel with in_value " + ", ".join(map("{}".format, target_values)))
 
-        bool_msk = np.isin(source_raster.data, np.array(target_values)).astype("uint8")
+        bool_msk = np.isin(nanarray, np.array(target_values)).astype("uint8")
 
     # Mask all valid values
     elif target_values == "all":
         # Using getmaskarray is necessary in case .data.mask is nomask (False)
-        bool_msk = (~np.ma.getmaskarray(source_raster.data)).astype("uint8")
+        bool_msk = (~np.ma.getmaskarray(nanarray)).astype("uint8")
 
     else:
         raise ValueError("in_value must be a number, a tuple or a sequence")
@@ -99,7 +101,7 @@ def _polygonize(
     results = (
         {"properties": {"raster_value": v}, "geometry": s}
         for i, (s, v) in enumerate(
-            shapes(source_raster.data.astype(final_dtype), mask=bool_msk, transform=source_raster.transform)
+            shapes(nanarray.astype(final_dtype), mask=bool_msk, transform=source_raster.transform)
         )
     )
 
@@ -115,7 +117,7 @@ def _polygonize(
 
 def _rasterize(
     gdf: gpd.GeoDataFrame,
-    raster: gu.Raster | None = None,
+    raster: gu.Raster | xr.DataArray | None = None,
     crs: CRS | int | None = None,
     xres: float | None = None,
     yres: float | None = None,
@@ -132,7 +134,10 @@ def _rasterize(
         crs = gdf.crs
 
     if raster is not None:
-        crs = raster.crs  # type: ignore
+        if isinstance(raster, gu.Raster):
+            crs = raster.crs  # type: ignore
+        else:
+            crs = raster.rst.crs
 
     vect = gdf.to_crs(crs)
 
@@ -163,10 +168,14 @@ def _rasterize(
         # Calculate raster transform
         transform = rio.transform.from_bounds(left, bottom, right, top, width, height)
 
-    # otherwise use directly raster's dimensions
+    # Otherwise use directly raster's dimensions
     else:
-        out_shape = raster.shape  # type: ignore
-        transform = raster.transform  # type: ignore
+        if isinstance(raster, gu.Raster):
+            out_shape = raster.shape  # type: ignore
+            transform = raster.transform  # type: ignore
+        else:
+            out_shape = raster.rst.shape
+            transform = raster.rst.transform
 
     # Set default burn value, index from 1 to len(self.ds)
     if in_value is None:
@@ -264,9 +273,14 @@ def _create_mask(
 
         # For a reference, extract transform and CRS
         if ref is not None:
-            transform = ref.transform
-            out_shape = ref.shape
-            crs = ref.crs
+            if isinstance(ref, xr.DataArray):
+                transform = ref.rst.transform
+                out_shape = ref.rst.shape
+                crs = ref.rst.crs
+            else:
+                transform = ref.transform
+                out_shape = ref.shape
+                crs = ref.crs
 
         # For a user-input res
         else:

geoutils/pointcloud/pointcloud.py

@@ -28,6 +28,7 @@
 import geopandas as gpd
 import numpy as np
 import pandas as pd
+import xarray as xr
 from pyproj import CRS
 from rasterio.coords import BoundingBox
 from rasterio.transform import from_origin
@@ -772,13 +773,11 @@ def __setitem__(self, index: Any, assign: Any) -> None:
             ind = ind.astype(bool)  # In case the 3D Z column was used, it can only be stored as floating
             # Assign
             if self._has_z:
-                print(ind)
                 new_geo = gpd.points_from_xy(
                     x=self.geometry.x.values[ind], y=self.geometry.y.values[ind], z=assign, crs=self.crs
                 )
                 self.ds.loc[ind, "geometry"] = new_geo
             else:
-                print(self.data_column)
                 self.ds.loc[ind, [self.data_column]] = assign
 
         else:
@@ -1517,13 +1516,16 @@ def grid(
         :return: Raster from gridded point cloud.
         """
 
-        if isinstance(ref, gu.Raster):
+        if isinstance(ref, (gu.Raster, xr.DataArray)):
             if grid_coords is not None:
                 warnings.warn(
                     "Both reference point cloud and grid coordinates were passed for gridding, "
                     "using only the reference point cloud."
                 )
-            grid_coords = ref.coords(grid=False)
+            if isinstance(ref, gu.Raster):
+                grid_coords = ref.coords(grid=False)
+            else:
+                grid_coords = ref.rst.coords(grid=False)
         else:
             if res is not None:
                 xsize = (self.bounds.right - self.bounds.left) / res

geoutils/raster/_geotransformations.py

@@ -30,13 +30,14 @@
 import affine
 import numpy as np
 import rasterio as rio
+import xarray as xr
 from packaging.version import Version
 from rasterio.crs import CRS
 from rasterio.enums import Resampling
 
 import geoutils as gu
 from geoutils._misc import silence_rasterio_message
-from geoutils._typing import DTypeLike, NDArrayBool, NDArrayNum
+from geoutils._typing import DTypeLike, NDArrayNum
 from geoutils.raster.georeferencing import (
     _cast_pixel_interpretation,
     _default_nodata,
@@ -144,18 +145,18 @@ def _user_input_reproject(
         if isinstance(ref, gu.Raster):
             # Raise a warning if the reference is a raster that has a different pixel interpretation
             _cast_pixel_interpretation(source_raster.area_or_point, ref.area_or_point)
-            ds_ref = ref
-        elif isinstance(ref, str):
-            if not os.path.exists(ref):
-                raise ValueError("Reference raster does not exist.")
-            ds_ref = gu.Raster(ref, load_data=False)
+            # Read reprojecting params from ref raster
+            crs = ref.crs
+            res = ref.res
+            bounds = ref.bounds
+        elif isinstance(ref, xr.DataArray):
+            _cast_pixel_interpretation(source_raster.area_or_point, ref.rst.area_or_point)
+            crs = ref.rst.crs
+            res = ref.rst.res
+            bounds = ref.rst.bounds
         else:
             raise TypeError("Type of ref not understood, must be path to file (str), Raster.")
 
-        # Read reprojecting params from ref raster
-        crs = ds_ref.crs
-        res = ds_ref.res
-        bounds = ds_ref.bounds
     else:
         # Determine target CRS
         crs = CRS.from_user_input(crs)
@@ -331,16 +332,29 @@ def _is_reproj_needed(src_shape: tuple[int, int], reproj_kwargs: dict[str, Any])
     )
 
 
-def _rio_reproject(
-    src_arr: NDArrayNum | NDArrayBool, src_mask: NDArrayBool, reproj_kwargs: dict[str, Any]
-) -> tuple[NDArrayNum | NDArrayBool, NDArrayBool]:
+def _rio_reproject(src_arr: NDArrayNum, reproj_kwargs: dict[str, Any]) -> NDArrayNum:
     """Rasterio reprojection wrapper.
 
     :param src_arr: Source array for data.
-    :param src_mask: Source array for mask, only required if array is not float.
     :param reproj_kwargs: Reprojection parameter dictionary.
     """
 
+    # All masked values must be set to a nodata value for rasterio's reproject to work properly
+    if np.ma.isMaskedArray(src_arr):
+        is_input_masked = True
+        src_mask = np.ma.getmaskarray(src_arr)
+        src_arr = src_arr.data  # type: ignore
+    else:
+        is_input_masked = False
+        src_mask = ~np.isfinite(src_arr)
+
+    # Check reprojection is possible with nodata (boolean raster will be converted, so no need to check)
+    if np.dtype(src_arr.dtype) != bool and (reproj_kwargs["src_nodata"] is None and np.sum(src_mask) > 0):
+        raise ValueError(
+            "No nodata set, set one for the raster with self.set_nodata() or use a temporary one "
+            "with `force_source_nodata`."
+        )
+
     # For a boolean type
     convert_bool = False
     if np.dtype(src_arr.dtype) == np.bool_:
@@ -395,13 +409,16 @@ def _rio_reproject(
     )
     # If Rasterio is recent enough version, force tolerance to 0 to avoid deformations on chunks
     # See: https://github.com/rasterio/rasterio/issues/2433#issuecomment-2786157846
-    if Version(rio.__version__) > Version("1.5"):
+    if Version(rio.__version__) >= Version("1.5.0"):
         reproj_kwargs.update({"tolerance": 0})
 
     # Pop dtype and dst_shape arguments that don't exist in Rasterio, and are only used above
     reproj_kwargs.pop("dtype")
     reproj_kwargs.pop("dst_shape")
 
+    # TODO: Figure out why those warning are raised for _dask_reproject with perfectly fine src_transforms
+    warnings.filterwarnings("ignore", category=rio.errors.NotGeoreferencedWarning)
+
     # XSCALE/YSCALE have been supported for a while, but not officially exposed in the API until Rasterio 1.5,
     # so we need to silence them in warnings to avoid noise for users
     with silence_rasterio_message(param_name="SCALE"):
@@ -418,4 +435,11 @@ def _rio_reproject(
     if convert_bool:
         dst_arr = dst_arr.astype(bool)
 
-    return dst_arr, dst_mask
+    # Set mask
+    if is_input_masked:
+        dst_arr = np.ma.masked_array(data=dst_arr, mask=dst_mask, fill_value=reproj_kwargs["dst_nodata"])
+    else:
+        dst_arr = dst_arr
+        dst_arr[dst_mask] = np.nan
+
+    return dst_arr