Merge pull request #39 from manaakiwhenua/a5

--band

Author: alpha-beta-soup

README.md

@@ -64,8 +64,17 @@ Options:
   -r, --resolution [0|1|2|3|4|5|6|7|8|9|10|11|12|13|14|15]
                                   H3 resolution to index  [required]
   -pr, --parent_res [0|1|2|3|4|5|6|7|8|9|10|11|12|13|14|15]
-                                  H3 Parent resolution to index and aggregate
+                                  H3 parent resolution to index and aggregate
                                   to. Defaults to resolution - 6
+  -b, --band TEXT                 Band(s) to include in the output. Can
+                                  specify multiple, e.g. `-b 1 -b 2 -b 4` for
+                                  bands 1, 2, and 4 (all unspecified bands are
+                                  ignored). If unused, all bands are included
+                                  in the output (this is the default
+                                  behaviour). Bands can be specified as
+                                  numeric indices (1-based indexing) or string
+                                  band labels (if present in the input), e.g.
+                                  -b B02 -b B07 -b B12.
   -u, --upscale INTEGER           Upscaling factor, used to upsample input
                                   data on the fly; useful when the raster
                                   resolution is lower than the target DGGS
@@ -101,7 +110,7 @@ Options:
                                   is a need to resample. This setting
                                   specifies this resampling algorithm.
                                   [default: average]
-  -co, --compact                  Compact the H3 cells up to the parent
+  -co, --compact                  Compact the cells up to the parent
                                   resolution. Compaction is not applied for
                                   cells without identical values across all
                                   bands.

raster2dggs/a5.py

@@ -3,7 +3,7 @@
 import tempfile
 
 from pathlib import Path
-from typing import Union
+from typing import Optional, Sequence, Union
 from rasterio.enums import Resampling
 
 import raster2dggs.constants as const
@@ -20,82 +20,88 @@
     "--resolution",
     required=True,
     type=click.Choice(list(map(str, range(const.MIN_A5, const.MAX_A5 + 1)))),
-    help="A5 resolution to index",
+    help=const.OPTION_HELP['resolution']('A5'),
 )
 @click.option(
     "-pr",
     "--parent_res",
     required=False,
     type=click.Choice(list(map(str, range(const.MIN_A5, const.MAX_A5 + 1)))),
-    help="A5 parent resolution to index and aggregate to. Defaults to resolution - 6",
+    help=const.OPTION_HELP['parent_res']('A5', 'resolution - 6'),
+)
+@click.option(
+    "-b",
+    "--band",
+    required=False,
+    multiple=True,
+    help=const.OPTION_HELP['band'],
 )
 @click.option(
     "-u",
     "--upscale",
     default=const.DEFAULTS["upscale"],
     type=int,
-    help="Upscaling factor, used to upsample input data on the fly; useful when the raster resolution is lower than the target DGGS resolution. Default (1) applies no upscaling. The resampling method controls interpolation.",
+    help=const.OPTION_HELP['upscale'],
 )
 @click.option(
     "-c",
     "--compression",
     default=const.DEFAULTS["compression"],
     type=str,
-    help="Compression method to use for the output Parquet files. Options include 'snappy', 'gzip', 'brotli', 'lz4', 'zstd', etc. Use 'none' for no compression.",
-)
+    help=const.OPTION_HELP['compression'],)
 @click.option(
     "-t",
     "--threads",
     default=const.DEFAULTS["threads"],
-    help="Number of threads to use when running in parallel. The default is determined based dynamically as the total number of available cores, minus one.",
+    help=const.OPTION_HELP['threads'],
 )
 @click.option(
     "-a",
     "--aggfunc",
     default=const.DEFAULTS["aggfunc"],
     type=click.Choice(
-        ["count", "mean", "sum", "prod", "std", "var", "min", "max", "median", "mode"]
+        const.AGGFUNC_OPTIONS
     ),
-    help="Numpy aggregate function to apply when aggregating cell values after DGGS indexing, in case of multiple pixels mapping to the same DGGS cell.",
+    help=const.OPTION_HELP['aggfunc'],
 )
 @click.option(
     "-d",
     "--decimals",
     default=const.DEFAULTS["decimals"],
     type=int,
-    help="Number of decimal places to round values when aggregating. Use 0 for integer output.",
-)
+    help=const.OPTION_HELP['decimals'],)
 @click.option("-o", "--overwrite", is_flag=True)
 @click.option(
     "--warp_mem_limit",
     default=const.DEFAULTS["warp_mem_limit"],
     type=int,
-    help="Input raster may be warped to EPSG:4326 if it is not already in this CRS. This setting specifies the warp operation's memory limit in MB.",
+    help=const.OPTION_HELP['warp_mem_limit'],
 )
 @click.option(
     "--resampling",
     default=const.DEFAULTS["resampling"],
     type=click.Choice(Resampling._member_names_),
-    help="Input raster may be warped to EPSG:4326 if it is not already in this CRS. Or, if the upscale parameter is greater than 1, there is a need to resample. This setting specifies this resampling algorithm.",
+    help=const.OPTION_HELP['resampling'],
 )
 @click.option(
     "-co",
     "--compact",
     is_flag=True,
-    help="Compact the cells up to the parent resolution. Compaction is not applied for cells without identical values across all bands.",
+    help=const.OPTION_HELP['compact'],
 )
 @click.option(
     "--tempdir",
     default=const.DEFAULTS["tempdir"],
     type=click.Path(),
-    help="Temporary data is created during the execution of this program. This parameter allows you to control where this data will be written.",
+    help=const.OPTION_HELP['tempdir'],
 )
 @click.version_option(version=__version__)
 def a5(
     raster_input: Union[str, Path],
     output_directory: Union[str, Path],
     resolution: str,
     parent_res: str,
+    band: Optional[Sequence[Union[int, str]]],
     upscale: int,
     compression: str,
     threads: int,
@@ -139,5 +145,6 @@ def a5(
         int(resolution),
         parent_res,
         warp_args,
+        band,
         **kwargs,
     )

raster2dggs/cli.py

@@ -8,15 +8,6 @@
 from raster2dggs.s2 import s2
 from raster2dggs.a5 import a5
 
-#   If the program does terminal interaction, make it output a short
-# notice like this when it starts in an interactive mode:
-
-#     <program>  Copyright (C) <year>  <name of author>
-#     This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
-#     This is free software, and you are welcome to redistribute it
-#     under certain conditions; type `show c' for details.
-
-
 @click.group()
 @click.version_option(version=__version__)
 def cli():

raster2dggs/common.py

@@ -2,6 +2,7 @@
 import errno
 import tempfile
 import logging
+import numpy as np
 import threading
 import rioxarray
 import dask
@@ -11,7 +12,7 @@
 import pandas as pd
 import pyarrow.parquet as pq
 
-from typing import Union, Callable
+from typing import Union, Optional, Sequence, Callable
 from pathlib import Path
 from rasterio import crs
 from rasterio.vrt import WarpedVRT
@@ -67,7 +68,7 @@ def resolve_input_path(raster_input: Union[str, Path]) -> Union[str, Path]:
 
 def assemble_warp_args(resampling: str, warp_mem_limit: int) -> dict:
     warp_args: dict = {
-        "resampling": Resampling._member_map_[resampling],
+        "resampling": Resampling[resampling],
         "crs": crs.CRS.from_epsg(
             4326
         ),  # Input raster must be converted to WGS84 (4326) for DGGS indexing
@@ -76,13 +77,17 @@ def assemble_warp_args(resampling: str, warp_mem_limit: int) -> dict:
 
     return warp_args
 
+def first_mode(x):
+    m = pd.Series.mode(x, dropna=False)
+    # Result is empty if all x is nan
+    return m.iloc[0] if not m.empty else np.nan
 
 def create_aggfunc(aggfunc: str) -> str:
     if aggfunc == "mode":
         logging.warning(
             "Mode aggregation: arbitrary behaviour: if there is more than one mode when aggregating, only the first value will be recorded."
         )
-        aggfunc = lambda x: pd.Series.mode(x)[0]
+        aggfunc = first_mode
 
     return aggfunc
 
@@ -91,7 +96,7 @@ def assemble_kwargs(
     upscale: int,
     compression: str,
     threads: int,
-    aggfunc: str,
+    aggfunc: Union[str, Callable],
     decimals: int,
     warp_mem_limit: int,
     resampling: str,
@@ -158,13 +163,12 @@ def address_boundary_issues(
     LOGGER.debug(
         f"Reading Stage 1 output ({pq_input}) and setting index for parent-based partitioning"
     )
-    with TqdmCallback(desc="Reading window partitions"):
-        # Set index as parent cell
-        pad_width = const.zero_padding(indexer.dggs)
-        index_col = f"{indexer.dggs}_{parent_res:0{pad_width}d}"
-        ddf = dd.read_parquet(pq_input).set_index(index_col)
+    # Set index as parent cell
+    pad_width = const.zero_padding(indexer.dggs)
+    index_col = f"{indexer.dggs}_{parent_res:0{pad_width}d}"
+    ddf = dd.read_parquet(pq_input).set_index(index_col)
 
-    with TqdmCallback(desc="Counting parents"):
+    with TqdmCallback(desc="Reading window partitions and counting parents"):
         # Count parents, to get target number of partitions
         uniqueparents = sorted(list(ddf.index.unique().compute()))
 
@@ -209,6 +213,7 @@ def initial_index(
     resolution: int,
     parent_res: Union[None, int],
     warp_args: dict,
+    bands: Optional[Sequence[Union[int, str]]] = None,
     **kwargs,
 ) -> Path:
     """
@@ -237,10 +242,33 @@ def initial_index(
 
         # https://rasterio.readthedocs.io/en/latest/api/rasterio.warp.html#rasterio.warp.calculate_default_transform
         with rio.Env(CHECK_WITH_INVERT_PROJ=True):
-            with rio.open(raster_input) as src:
+            with rio.open(raster_input, mode='r', sharing=False) as src:
                 LOGGER.debug("Source CRS: %s", src.crs)
                 # VRT used to avoid additional disk use given the potential for reprojection to 4326 prior to DGGS indexing
-                band_names = src.descriptions
+                band_names = tuple(src.descriptions) if src.descriptions else tuple()
+                count = src.count # Bands
+                labels_by_index = {
+                    i: (band_names[i-1] if i-1 < len(band_names) and band_names[i-1] else f"band_{i}")
+                    for i in range(1, count + 1)
+                }
+                if not bands: # Covers None or empty tuple
+                    selected_indices = list(range(1, count + 1))
+                else:
+                    if all(isinstance(b, int) or str(b).isdigit() for b in bands):
+                        selected_indices = list(map(int, bands))
+                    else:
+                        name_to_index = {v: k for k, v in labels_by_index.items()}
+                        try:
+                            selected_indices = [name_to_index[str(b)] for b in bands]
+                        except KeyError as e:
+                            raise ValueError(f"Requested band name not found: {e.args[0]}")
+                    # Validate
+                    for i in selected_indices:
+                        if i < 1 or i > count:
+                            raise ValueError(f"Band index out of range: {i} (1..{count})")
+                    # De-duplicate, preserving order
+                    seen = set()
+                    selected_indices = [i for i in selected_indices if not (i in seen or seen.add(i))]
 
                 upscale_factor = kwargs["upscale"]
                 if upscale_factor > 1:
@@ -262,7 +290,7 @@ def initial_index(
                     upsample_args = dict({})
 
                 with WarpedVRT(
-                    src, src_crs=src.crs, **warp_args, **upsample_args
+                    src, src_crs=src.crs, **warp_args, **upsample_args,
                 ) as vrt:
                     LOGGER.debug("VRT CRS: %s", vrt.crs)
                     da: xr.Dataset = rioxarray.open_rasterio(
@@ -272,13 +300,21 @@ def initial_index(
                         default_name=const.DEFAULT_NAME,
                     ).chunk(**{"y": "auto", "x": "auto"})
 
+                    # Band selection
+                    if "band" in da.dims and (len(selected_indices) != count):
+                        if "band" in da.coords: # rioxarray commonly exposes 1..N as band coords
+                            da = da.sel(band=selected_indices)
+                        else:
+                            da = da.isel(band=[i - 1 for i in selected_indices])
+
                     windows = [window for _, window in vrt.block_windows()]
                     LOGGER.debug(
                         "%d windows (the same number of partitions will be created)",
                         len(windows),
                     )
 
-                    write_lock = threading.Lock()
+                    selected_labels = tuple([labels_by_index[i] for i in selected_indices])
+                    compression = kwargs["compression"]
 
                     def process(window):
                         sdf = da.rio.isel_window(window)
@@ -288,28 +324,22 @@ def process(window):
                             resolution,
                             parent_res,
                             vrt.nodata,
-                            band_labels=band_names,
+                            band_labels=selected_labels,
                         )
 
-                        with write_lock:
-                            pq.write_to_dataset(
-                                result,
-                                root_path=tmpdir,
-                                compression=kwargs["compression"],
-                            )
+                        pq.write_to_dataset(
+                            result,
+                            root_path=tmpdir,
+                            compression=compression,
+                        )
 
                         return None
 
-                    with tqdm(total=len(windows), desc="Raster windows") as pbar:
-                        with ThreadPoolExecutor(
+                    with ThreadPoolExecutor(
                             max_workers=kwargs["threads"]
-                        ) as executor:
-                            futures = [
-                                executor.submit(process, window) for window in windows
-                            ]
-                            for future in as_completed(futures):
-                                result = future.result()
-                                pbar.update(1)
+                        ) as executor, tqdm(total=len(windows), desc="Raster windows") as pbar:
+                        for _ in executor.map(process, windows, chunksize=1):
+                            pbar.update(1)
 
             LOGGER.debug("Stage 1 (primary indexing) complete")
             return address_boundary_issues(

raster2dggs/constants.py

@@ -48,3 +48,20 @@ def zero_padding(dggs: str) -> int:
     if max_res is None:
         raise ValueError(f"Unknown DGGS type: {dggs}")
     return len(str(max_res))
+
+OPTION_HELP = {
+    'resolution': lambda dggs: f"{dggs} resolution to index",
+    'parent_res': lambda dggs, default: f"{dggs} parent resolution to index and aggregate to. Defaults to {default}",
+    'band': "Band(s) to include in the output. Can specify multiple, e.g. `-b 1 -b 2 -b 4` for bands 1, 2, and 4 (all unspecified bands are ignored). If unused, all bands are included in the output (this is the default behaviour). Bands can be specified as numeric indices (1-based indexing) or string band labels (if present in the input), e.g. -b B02 -b B07 -b B12.",
+    'upscale': "Upscaling factor, used to upsample input data on the fly; useful when the raster resolution is lower than the target DGGS resolution. Default (1) applies no upscaling. The resampling method controls interpolation.",
+    'compression': "Compression method to use for the output Parquet files. Options include 'snappy', 'gzip', 'brotli', 'lz4', 'zstd', etc. Use 'none' for no compression.",
+    'threads': "Number of threads to use when running in parallel. The default is determined based dynamically as the total number of available cores, minus one.",
+    'aggfunc': "Numpy aggregate function to apply when aggregating cell values after DGGS indexing, in case of multiple pixels mapping to the same DGGS cell.",
+    'decimals': "Number of decimal places to round values when aggregating. Use 0 for integer output.",
+    'warp_mem_limit': "Input raster may be warped to EPSG:4326 if it is not already in this CRS. This setting specifies the warp operation's memory limit in MB.",
+    'resampling': "Input raster may be warped to EPSG:4326 if it is not already in this CRS. Or, if the upscale parameter is greater than 1, there is a need to resample. This setting specifies this resampling algorithm.",
+    'compact': "Compact the cells up to the parent resolution. Compaction is not applied for cells without identical values across all bands.",
+    'tempdir': "Temporary data is created during the execution of this program. This parameter allows you to control where this data will be written."
+}
+
+AGGFUNC_OPTIONS = ["count", "mean", "sum", "prod", "std", "var", "min", "max", "median", "mode"]