updated ndvi loading and dataset merge

darts-preprocessing/src/darts_preprocessing/preprocess.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 
+import xarray as xr
 from utils.data_pre_processing import calculate_ndvi, load_auxiliary, load_planet_scene
 
 """
@@ -26,16 +27,17 @@
 # load planet scene
 ds_planet = load_planet_scene(planet_scene_path)
 
-# calculate ndvi
+# calculate xr.dataset ndvi
 ds_ndvi = calculate_ndvi(ds_planet)
 
-# get dataarray for elevation
+# get xr.dataset for elevation
 ds_elevation = load_auxiliary(planet_scene_path, elevation_path, xr_dataset_name="relative_elevation")
 
-# # get dataarray for slope
-da_slope = load_auxiliary(planet_scene_path, slope_path, xr_dataset_name="slope")
+# get xr.dataset for slope
+ds_slope = load_auxiliary(planet_scene_path, slope_path, xr_dataset_name="slope")
 
-# # get dataarray for tcvis
-# da_elevation = load_auxiliary(planet_scene_path, elevation_path)
+# # get xr.dataset for tcvis
+# ds_tcvis = load_auxiliary(planet_scene_path, tcvis_path)
 
-# merge all datasets
+# merge to final dataset
+ds_merged = xr.merge([ds_planet, ds_ndvi, ds_elevation, ds_slope])

darts-preprocessing/src/darts_preprocessing/utils/data_pre_processing.py

@@ -29,6 +29,7 @@ def load_planet_scene(planet_scene_path: str | Path) -> xr.Dataset:
         planet_scene_path = Path(planet_scene_path)
 
     # Find the appropriate TIFF file
+    # TODO use get_planet_image_from_planet_scene_path function
     ps_image = list(planet_scene_path.glob(f"{planet_scene_path.name}_*_SR.tif"))
 
     if not ps_image:
@@ -38,25 +39,30 @@ def load_planet_scene(planet_scene_path: str | Path) -> xr.Dataset:
     return rxr.open_rasterio(ps_image[0]).to_dataset(name="planet")
 
 
-def calculate_ndvi(planet_scene_dataarray: xr.DataArray, nir_band: int = 4, red_band: int = 3) -> xr.Dataset:
+def calculate_ndvi(planet_scene_dataarray: xr.DataArray, nir_band: int = 4, red_band: int = 3) -> xr.DataArray:
     """Calculate NDVI from an xarray DataArray containing spectral bands.
 
     Parameters
     ----------
     planet_scene_dataarray : xr.DataArray
         The xarray DataArray containing the spectral bands, where the bands are
-        indexed along a dimension (e.g., 'band').
+        indexed along a dimension (e.g., 'band'). The DataArray should have
+        dimensions including 'band', 'y', and 'x'.
 
     nir_band : int, optional
-        The index of the NIR band in the DataArray (default is 4).
+        The index of the NIR band in the DataArray (default is 4). This index
+        should correspond to the position of the NIR band in the 'band' dimension.
 
     red_band : int, optional
-        The index of the Red band in the DataArray (default is 3).
+        The index of the Red band in the DataArray (default is 3). This index
+        should correspond to the position of the Red band in the 'band' dimension.
 
     Returns
     -------
     xr.DataArray
-        A new DataArray containing the calculated NDVI values.
+        A new DataArray containing the calculated NDVI values. The resulting
+        DataArray will have dimensions (band: 1, y: ..., x: ...) and will be
+        named "ndvi".
 
     Raises
     ------
@@ -65,16 +71,23 @@ def calculate_ndvi(planet_scene_dataarray: xr.DataArray, nir_band: int = 4, red_
 
     Notes
     -----
-    NDVI is calculated using the formula:
+    NDVI (Normalized Difference Vegetation Index) is calculated using the formula:
         NDVI = (NIR - Red) / (NIR + Red)
 
+    This index is commonly used in remote sensing to assess vegetation health
+    and density.
+
+    Example
+    -------
+    >>> ndvi_data = calculate_ndvi(planet_scene_dataarray)
+
     """
     # Calculate NDVI using the formula
     nir = planet_scene_dataarray.sel(band=nir_band).astype("float32")
     r = planet_scene_dataarray.sel(band=red_band).astype("float32")
     ndvi = (nir - r) / (nir + r)
 
-    return ndvi.to_dataarray().to_dataset(name="ndvi")
+    return ndvi.expand_dims(dim={"band": [1]}, axis=0).rename_vars({"planet": "ndvi"})
 
 
 def geom_from_image_bounds(image_path):