diff --git a/.github/workflows/lint_and_test.yml b/.github/workflows/lint_and_test.yml index 2699179..a27ff47 100644 --- a/.github/workflows/lint_and_test.yml +++ b/.github/workflows/lint_and_test.yml @@ -46,10 +46,6 @@ jobs: run: ruff check src/ tests/ shell: bash -el {0} - - name: Run mypy type checking - run: mypy src/ tests/ - shell: bash -el {0} - - name: Run pytest run: pytest shell: bash -el {0} \ No newline at end of file diff --git a/boathouse.json b/boathouse.json index 5263db7..4d360b4 100644 --- a/boathouse.json +++ b/boathouse.json @@ -146,6 +146,101 @@ } ] }, + { + "name": "Burn", + "path": "overflow", + "cmd": [ + "burn" + ], + "readme": "https://fema-ffrd.github.io/overflow/", + "description": "Burn elevations into a DEM inside the contiguous regions of a binary or classified mask raster, either as a constant, relative to the existing surface, or as a per region statistic.", + "category": "Hydrology", + "keywords": [ + "dem", + "burn", + "mask", + "region", + "waterbody", + "terrain", + "preprocessing", + "hydrology", + "raster" + ], + "args": [ + { + "displayName": "DEM File", + "name": "dem_file", + "type": "file", + "description": "path to the GDAL supported raster dataset for the input DEM" + }, + { + "displayName": "Mask File", + "name": "mask_file", + "type": "file", + "description": "path to a binary or classified raster, co-registered with the DEM, whose regions are burned into the DEM" + }, + { + "displayName": "Output File", + "name": "output_file", + "type": "file", + "description": "path to the output file (must be GeoTiff)", + "direction": "output" + }, + { + "displayName": "Method", + "name": "method", + "type": "string", + "description": "how to choose the elevation to burn: constant, relative, or statistic", + "defaultValue": "statistic" + }, + { + "displayName": "Burn Values", + "name": "burn_values", + "type": "string", + "description": "required for the constant and relative methods: a mapping of mask value to burn value such as '1:225.5,3:210.0', or a single number applied to every selected mask value" + }, + { + "displayName": "Mask Values", + "name": "mask_values", + "type": "string", + "description": "comma separated mask values identifying regions, such as '1,3'. Defaults to the keys of burn values, or to every non zero, non nodata mask value" + }, + { + "displayName": "Statistic", + "name": "statistic", + "type": "string", + "description": "only used with the statistic method: the statistic computed per region, one of min, max, mean", + "defaultValue": "min" + }, + { + "displayName": "Burn Offset", + "name": "burn_offset", + "type": "float", + "description": "only used with the statistic method: subtracted from each region's statistic", + "defaultValue": 0.0 + }, + { + "displayName": "Connectivity", + "name": "connectivity", + "type": "string", + "description": "8 treats diagonally touching cells as one region, 4 requires a shared edge", + "defaultValue": "8" + }, + { + "displayName": "Chunk Size", + "name": "chunk_size", + "type": "integer", + "description": "chunk size for tiled operations", + "defaultValue": 2048 + }, + { + "displayName": "Working Directory", + "name": "working_dir", + "type": "directory", + "description": "path to the working directory used for temporary files" + } + ] + }, { "name": "Fill", "path": "overflow", diff --git a/docs/algorithm-details/burn.md b/docs/algorithm-details/burn.md new file mode 100644 index 0000000..6fe4abb --- /dev/null +++ b/docs/algorithm-details/burn.md @@ -0,0 +1,149 @@ +# Burn Algorithm + +## Overview + +The burn algorithm rewrites a DEM inside regions defined by a co-registered mask raster. It reads the DEM and the mask and writes a modified DEM in which the selected regions carry an imposed elevation. + +Two of the three methods are trivial. `constant` and `relative` compute each output cell from the DEM value and mask value at that cell alone, so they need one pass over the data and nothing else. The `statistic` method is the interesting one: each contiguous region takes a statistic of the DEM beneath it, which no single cell and no single tile can determine on its own. For large rasters processed tile by tile, a region that crosses a tile boundary appears as two or more independent fragments, and the algorithm has to recognize them as one region before it can compute anything. + +## Definitions + +Let $M$ be the mask raster, $Z$ the DEM, and $S$ the set of selected mask values. + +A cell $c$ is **selected** when $M_c \in S$ and $M_c$ is neither the mask's nodata value nor NaN. + +A **region** is an equivalence class of selected cells under the transitive closure of the adjacency relation + +$$ +c \sim d \iff d \in N(c) \land M_c = M_d +$$ + +where $N(c)$ is the 8-connected neighborhood by default, or the 4-connected neighborhood when `connectivity=4`. Requiring $M_c = M_d$ rather than merely requiring both to be selected is what keeps two touching classes in separate regions. + +For a region $R$ the burned elevation is + +$$ +b(R) = \operatorname{stat}\left(\{Z_c : c \in R,\ Z_c \neq \text{nodata}\}\right) - \delta +$$ + +with $\operatorname{stat} \in \{\min, \max, \operatorname{mean}\}$ and $\delta$ the burn offset. When the set is empty the region has no defined value and is left untouched. + +## Mergeable Statistics + +The tiled algorithm works because $\min$, $\max$ and $\operatorname{mean}$ can all be computed from four quantities that combine associatively. For a region split into fragments $R_1, \dots, R_k$, accumulate per fragment + +$$ +\left(\min_i,\ \max_i,\ \Sigma_i = \sum_{c \in R_i} Z_c,\ n_i = |R_i|\right) +$$ + +and recover the whole region's statistics as + +$$ +\min = \min_i \min_i, \qquad +\max = \max_i \max_i, \qquad +\operatorname{mean} = \frac{\sum_i \Sigma_i}{\sum_i n_i} +$$ + +No fragment needs to know about any other while it is being processed, and the fragments never need to be re-read. This is the property that decides which statistics the tool can offer: a median or an arbitrary percentile is **not** recoverable this way without retaining every value in the region, which is why they are not supported. + +## Data Types and Structures + +### Input/Output Types + +| Parameter | Type | Description | No-Data Value | +|-----------|------|-------------|---------------| +| dem | any float or integer | Input elevation raster | required | +| mask | any numeric | Binary or classified region raster | optional | +| output | same as input dem | Burned elevation raster | inherited from dem | +| labels | int64 | Temporary per-cell region label | 0 | + +### Internal Structures + +| Structure | Type | Purpose | +|-----------|------|---------| +| `MaskSelection` | jitclass | Which mask values are selected and, for the direct methods, their burn values | +| `labels` | int64 2D array | Region label per cell, 0 outside any region | +| `stats` | float64 array, shape $(n, 4)$ | Per-region $(\min, \max, \Sigma, n)$ for one tile | +| `UnionFind` | jitclass over `Dict[int64, int64]` | Disjoint set joining labels that turn out to be one region | +| `label_perimeters` | int64 2D array | Region labels around each tile's perimeter | +| `mask_perimeters` | float64 2D array | Mask values around each tile's perimeter | +| `burn_lookup` | `Dict[int64, float64]` | Final label to burned elevation map | + +## Algorithm + +### Region labeling + +Within a tile, regions are found by flood fill in the shape of Algorithm 4 *LabelFlats* from [Barnes, Lehman & Mulla (2014)](https://rbarnes.org/sci/2014_flats.pdf), with equality of mask value standing in for equality of elevation. Cells are scanned in row-major order; the first unlabeled selected cell seeds a new label, and a FIFO queue drives the fill outward across neighbors sharing that cell's mask value. + +A single queue is allocated once and reused across regions rather than per region, which matters when a mask holds many small blobs. + +### Local statistics + +A second scan over the labeled tile accumulates $(\min, \max, \Sigma, n)$ per label, skipping DEM nodata and NaN. This is a single sequential pass; it is not parallelized within the tile because the accumulators would race. + +## Tiled Processing + +### Globally unique labels without renumbering + +Each tile draws its labels from a private range, + +$$ +\text{label\_offset}(t) = s^2 \cdot t + 1 +$$ + +for tile index $t$ and chunk size $s$. Since a tile of $s \times s$ cells can hold at most $s^2$ regions, the ranges cannot overlap, so labels from different tiles are distinct without any communication or renumbering pass. Label 0 is reserved for "not in a region". This is the same trick the tiled depression fill uses for watershed labels. + +### Joining regions across seams + +Only tile perimeters can participate in a cross-tile adjacency, so each tile stashes its perimeter labels and perimeter mask values, and nothing else survives the first pass. The solve then walks the tile grid, and for each tile joins it to its eastern and southern neighbors plus the two diagonal pairs meeting at its southeastern corner: + +``` ++ - - + - - + +| A | B | ++ - - * - - + +| C | D | ++ - - + - - + +``` + +A joins B along a vertical seam and C along a horizontal one; the corner joins A-D and B-C cover the diagonal adjacencies at the point marked `*`, which no edge scan reaches. Every adjacency in the grid is covered exactly once, and the bounds guards keep it correct for a tile grid that is a single tile wide or tall. + +Along a seam, index $i$ in one tile's perimeter sits directly against index $i$ in the other's. With 8-connectivity a cell also touches its two diagonal neighbors across the seam, so the scan compares $i$ against $i-1$, $i$ and $i+1$; with 4-connectivity it compares only $i$. Two labels are joined when both cells are labeled and hold the same mask value, exactly the rule the within-tile flood fill applies. + +Joins go into a union-find structure with union by size and path-compressed find, giving near constant amortized time per join. + +### Edge padding + +Tiles hanging off the right or bottom edge of the raster are padded out to a full square by the reader, and the labeler is told how many rows and columns hold real data so padded cells are never labeled. Without that bound, padding sharing a selected mask value could flood-fill into a real region and then join two genuinely disconnected regions across a seam, giving both the wrong statistic. + +### Solve and apply + +Every label is resolved to its representative, the four statistics are merged per representative, and the result is expanded back into a flat label-to-value lookup so the final pass is one dict read per cell. The apply pass re-reads the **original** DEM, not the partially written output, so statistics can never be contaminated by values the tool itself wrote. + +## Complexity + +Let $N$ be the number of cells, $R$ the number of region fragments and $T$ the number of tiles with chunk size $s$. + +| Phase | Time | Memory | +|-------|------|--------| +| `constant` / `relative` | $O(N)$ | $O(s^2)$ per worker | +| Label and accumulate | $O(N)$ | $O(s^2)$ per worker | +| Seam joining | $O(T s\,\alpha(R))$ | $O(T s)$ perimeters | +| Solve | $O(R\,\alpha(R))$ | $O(R)$ | +| Apply | $O(N)$ | $O(s^2)$ per worker | + +$\alpha$ is the inverse Ackermann function, effectively constant. The solve is proportional to the number of region fragments rather than to raster size, so a mask covering a small fraction of a very large DEM stays cheap. + +Tiles are processed concurrently by numba kernels compiled with `nogil`, using the same thread pool and back-pressure pattern as the other tiled tools. Raster writes and the merge of per-tile statistics into the global dicts are serialized under a lock, since numba typed dicts are not safe to insert into concurrently. + +## Input Requirements + +- The DEM must declare a nodata value. +- The mask must have the same shape and geotransform as the DEM. A nodata value is recommended but not required. +- Mask values are compared by exact equality, so a floating point mask must hold exact class values. +- Naming the mask's own nodata value as a selected value is rejected. + +## See Also + +- [Burn User Guide](../user-guide/terrain-conditioning/burn.md) - Parameters and usage +- [Fill](fill.md) - The tiled depression fill this design borrows its tiling structure from +- [Breach](breach.md) - The other terrain conditioning operation diff --git a/docs/algorithm-details/index.md b/docs/algorithm-details/index.md index ffdd8f5..f52ac47 100644 --- a/docs/algorithm-details/index.md +++ b/docs/algorithm-details/index.md @@ -7,6 +7,7 @@ This section provides in-depth documentation of the algorithms implemented in Ov Prepare raw DEMs for hydrological analysis by removing artificial depressions: - [Breach Algorithm](breach.md) - Remove depressions by carving least-cost flow paths +- [Burn Algorithm](burn.md) - Impose elevations on contiguous regions of a mask raster - [Fill Algorithm](fill.md) - Fill depressions using priority-flood ## Flow Routing diff --git a/docs/api/index.md b/docs/api/index.md index a579248..03e1086 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -8,6 +8,10 @@ This reference documents the Python API for the `overflow` library. options: heading_level: 3 +::: overflow.burn + options: + heading_level: 3 + ::: overflow.fill options: heading_level: 3 diff --git a/docs/cli/index.md b/docs/cli/index.md index 78c891e..e49caf2 100644 --- a/docs/cli/index.md +++ b/docs/cli/index.md @@ -72,6 +72,36 @@ overflow [COMMAND] [OPTIONS] ----- +### `burn` + +**Description:** Burn elevations into a DEM inside the regions of a mask raster. This command rewrites the DEM inside regions identified by a binary or classified mask raster, either to a constant elevation, to the DEM lowered by a fixed amount, or to a statistic of the DEM beneath each contiguous region. Regions are recognized as whole even where they straddle processing tiles. + +**Options:** + +| Option | Type | Default | Required | Description | +| :--- | :--- | :--- | :--- | :--- | +| `--dem_file` | TEXT | - | **Yes** | Path to the GDAL supported raster dataset for the input DEM. | +| `--mask_file` | TEXT | - | **Yes** | Path to a binary or classified raster, co-registered with the DEM, whose regions are burned into the DEM. | +| `--output_file` | TEXT | - | **Yes** | Path to the output file (must be GeoTiff). | +| `--method` | CHOICE | - | **Yes** | One of `constant`, `relative`, `statistic`. See below. | +| `--burn_values` | TEXT | None | No | Required for `constant` and `relative`: a mapping of mask value to burn value such as `"1:225.5,3:210.0"`, or a single number applied to every selected mask value. | +| `--mask_values` | TEXT | None | No | Comma separated mask values identifying regions, such as `"1,3"`. Defaults to the keys of `--burn_values`, or to every non zero, non nodata mask value. | +| `--statistic` | CHOICE | min | No | Only used with `--method statistic`: one of `min`, `max`, `mean`. | +| `--burn_offset` | FLOAT | 0.0 | No | Only used with `--method statistic`: subtracted from each region's statistic. | +| `--connectivity` | CHOICE | 8 | No | `8` treats diagonally touching cells as one region, `4` requires a shared edge. | +| `--chunk_size` | INTEGER | 2048 | No | Chunk size (use \<= 1 for in-memory processing). | +| `--working_dir` | TEXT | None | No | Working directory for temporary files. | + +**Methods:** + +| `--method` | Cell value written | +| :--- | :--- | +| `constant` | The burn value configured for the cell's mask value | +| `relative` | `dem - burn_value`, applied per cell so relief inside the region is preserved | +| `statistic` | The region's own `min`, `max` or `mean` of the DEM beneath it, less `--burn_offset` | + +----- + ### `flow-direction` **Description:** Compute D8 flow directions from a DEM and resolve flat areas. This command calculates the steepest descent direction for each cell using the D8 algorithm, then resolves flat areas to ensure continuous flow paths. diff --git a/docs/user-guide/terrain-conditioning/burn.md b/docs/user-guide/terrain-conditioning/burn.md new file mode 100644 index 0000000..f218800 --- /dev/null +++ b/docs/user-guide/terrain-conditioning/burn.md @@ -0,0 +1,157 @@ +# Burn + +## Purpose + +Burn imposes known elevations on a DEM inside regions taken from a second raster. The mask raster is typically binary (one class of interest) or classified (several), and each contiguous region within it is rewritten to a constant elevation, to the DEM lowered by a fixed amount, or to a statistic of the DEM beneath that region. + +## When to Use + +Burn applies elevation information the DEM itself does not carry. Use burn when: + +- Waterbodies from a land cover or hydrography raster should be flat rather than carrying survey noise +- Bridge decks and culvert footprints obstruct flow and need to be dropped to the channel elevation beneath them +- A levee breach, excavation or fill footprint has a known design elevation +- A region's elevation should be defined relative to what is already there, such as lowering a channel by a fixed depth + +Burn is a conditioning step and does not by itself remove every depression. Run [breach](breach.md) and [fill](fill.md) afterwards to produce a fully conditioned surface. + +!!! note + Burn only touches cells inside the selected regions. Every other cell, including DEM nodata cells, is copied through unchanged. + +## Regions + +A region is a maximal connected set of cells that share the **same** mask value and whose value is one of the selected values. Two selected classes that touch stay in separate regions, so a lake class and a wetland class that share a border are never merged. + +Regions are recognized as whole even when they straddle processing tiles: a lake that spans four tiles gets one minimum, not four. See [Burn Algorithm Details](../../algorithm-details/burn.md) for how that is done. + +## Parameters + +### dem_path +Path to the input DEM raster. Must be GDAL-readable, single band, and must declare a nodata value. + +### mask_path +Path to a binary or classified mask raster, co-registered with the DEM: same shape and same geotransform. Any GDAL-readable data type. A nodata value is recommended but not required. + +### output_path +Path for the output burned DEM, written as GeoTIFF. Data type, nodata value, projection and geotransform are inherited from the input DEM. + +### method +How to choose the elevation to burn. + +| Method | Cell value written | +|---|---| +| `constant` | the burn value configured for the cell's mask value | +| `relative` | `dem - burn_value`, applied per cell so relief inside the region is preserved | +| `statistic` | the region's own statistic of the DEM beneath it, less `burn_offset` | + +### mask_values +Which mask values identify regions, given as a comma separated string such as `"1,3"` or as a sequence. If omitted, the keys of `burn_values` are used. If neither names specific values, every non-zero, non-nodata mask value is selected, and each distinct value still forms its own regions. + +### burn_values +Required by `constant` and `relative`, ignored by `statistic`. Either a mapping of mask value to burn value, written as `"1:225.5,3:210.0"` or given as a dict, or a single number applied to every selected mask value. + +### statistic +Which statistic the `statistic` method computes per region: `min`, `max` or `mean`. Default `min`. + +### burn_offset +Subtracted from each computed statistic, so `statistic="min"` with `burn_offset=1.0` writes each region's minimum elevation less one. Used only by the `statistic` method. Default `0.0`. + +### connectivity +`8` (default) treats diagonally touching cells as one region. `4` requires a shared edge, so a diagonal chain of cells becomes several regions. + +### chunk_size +Tile dimension in pixels for processing. Default 2048. Set to 1 or less for in-memory processing when the DEM fits in available RAM. + +### working_dir +Directory for temporary files during tiled processing with the `statistic` method. If omitted, a system temp directory is created and removed afterwards. The `constant` and `relative` methods write no temporary files. + +### progress_callback +Optional callback for monitoring long operations. See [ProgressCallback API](../../api/index.md#overflow.ProgressCallback). + +## CLI Usage + +Flatten every waterbody to its own minimum elevation, one foot lower, so water drains through: + +```bash +overflow burn \ + --dem_file dem.tif \ + --mask_file waterbodies.tif \ + --output_file dem_burned.tif \ + --method statistic \ + --mask_values 1 \ + --statistic min \ + --burn_offset 1.0 +``` + +Set two classes to known design elevations: + +```bash +overflow burn \ + --dem_file dem.tif \ + --mask_file structures.tif \ + --output_file dem_burned.tif \ + --method constant \ + --burn_values "1:225.5,3:210.0" +``` + +Lower a channel class by 2 feet while preserving its shape: + +```bash +overflow burn \ + --dem_file dem.tif \ + --mask_file channels.tif \ + --output_file dem_burned.tif \ + --method relative \ + --burn_values "1:2.0" +``` + +## Python API Usage + +```python +import overflow + +overflow.burn( + dem_path="dem.tif", + mask_path="waterbodies.tif", + output_path="dem_burned.tif", + method="statistic", + mask_values="1", + statistic="min", + burn_offset=1.0, +) +``` + +`burn_values` also accepts a dict, which is often easier to build programmatically: + +```python +overflow.burn( + dem_path="dem.tif", + mask_path="structures.tif", + output_path="dem_burned.tif", + method="constant", + burn_values={1: 225.5, 3: 210.0}, +) +``` + +## Nodata Handling + +- DEM nodata cells are never written, so nodata holes survive the burn intact. +- DEM nodata cells are excluded from region statistics, so a lake with a data gap in the middle still gets the minimum of its valid cells. +- A region lying entirely over DEM nodata has no statistic and is left untouched. +- Mask nodata cells are never selected. Naming the mask's own nodata value in `mask_values` is an error. + +## Performance Considerations + +The `constant` and `relative` methods are per-cell functions of the DEM and mask values, so they run in a single tiled pass with no temporary raster. They are roughly as fast as copying the DEM. + +The `statistic` method needs three passes and writes a temporary Int64 label raster the size of the input, so budget disk space in `working_dir` accordingly. Memory during the solve is proportional to the number of region fragments, not to raster size, so masks covering a small fraction of the DEM stay cheap even on very large rasters. + +!!! warning "Integer DEMs truncate fractional burn values" + The output carries the input DEM's data type. Burning a `mean` statistic, a `relative` amount or an offset into an integer DEM truncates the result. Convert the DEM to a floating point type first if that matters. + +## See Also + +- [Breach](breach.md) - Carving flow paths through elevation barriers +- [Fill](fill.md) - Depression removal using priority flood +- [Flow Direction](../flow-routing/flow-direction.md) - Computing flow routing after conditioning +- [Burn Algorithm Details](../../algorithm-details/burn.md) - Implementation and theory diff --git a/mkdocs.yml b/mkdocs.yml index ba44626..36c4837 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -88,6 +88,7 @@ nav: - Pipeline: user-guide/pipeline.md - Terrain Conditioning: - Breach: user-guide/terrain-conditioning/breach.md + - Burn: user-guide/terrain-conditioning/burn.md - Fill: user-guide/terrain-conditioning/fill.md - Flow Routing: - Flow Direction: user-guide/flow-routing/flow-direction.md @@ -100,6 +101,7 @@ nav: - algorithm-details/index.md - Terrain Conditioning: - Breach: algorithm-details/breach.md + - Burn: algorithm-details/burn.md - Fill: algorithm-details/fill.md - Flow Routing: - Flow Direction: algorithm-details/flow-direction.md diff --git a/pyproject.toml b/pyproject.toml index 379d04e..28aa681 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "overflow-hydro" -version = "0.4.0" +version = "0.5.0" description = "High-performance Python library for hydrological terrain analysis with parallel, tiled algorithms" readme = "README.md" requires-python = ">=3.11" diff --git a/src/overflow/__init__.py b/src/overflow/__init__.py index 564b23f..dc5476e 100644 --- a/src/overflow/__init__.py +++ b/src/overflow/__init__.py @@ -2,6 +2,7 @@ Overflow - High-performance Python library for hydrological terrain analysis. - breach: Breach pits using least-cost paths +- burn: Burn mask regions into a DEM - fill: Fill depressions in a DEM - flow_direction: Compute D8 flow directions and resolve flats - accumulation: Calculate flow accumulation @@ -19,6 +20,11 @@ ) from overflow._basins.tiled import _label_watersheds_tiled from overflow._breach_paths_least_cost import _breach_paths_least_cost +from overflow._burn_mask.core import _burn_mask_core +from overflow._burn_mask.core import ( + resolve_burn_arguments as _resolve_burn_arguments, +) +from overflow._burn_mask.tiled import _burn_mask_tiled from overflow._extract_streams.core import _extract_streams_core from overflow._extract_streams.tiled import _extract_streams_tiled from overflow._fill_depressions.core import _fill_depressions @@ -43,7 +49,7 @@ ) from overflow.codes import FlowDirection -__version__ = "0.4.0" +__version__ = "0.5.0" def breach( @@ -76,6 +82,102 @@ def breach( ) +def burn( + dem_path: str, + mask_path: str, + output_path: str, + method: str, + mask_values: str | list | tuple | None = None, + burn_values: str | float | dict | None = None, + statistic: str = "min", + burn_offset: float = 0.0, + connectivity: int = 8, + chunk_size: int = DEFAULT_CHUNK_SIZE, + working_dir: str | None = None, + progress_callback: ProgressCallback | None = None, +) -> None: + """ + Burn elevations into a DEM inside the regions of a mask raster. + + A region is a maximal connected set of cells that share the same mask value and + whose value is selected. Cells of two different selected values stay in separate + regions even where they touch. The statistic method gives each contiguous region + its own statistic, computed over whole regions even when they straddle tiles. + + Args: + dem_path: Path to the input DEM raster file. + mask_path: Path to a binary or classified mask raster, co-registered with + the DEM (same shape and geotransform). + output_path: Path for the output burned DEM raster file. Its data type and + nodata value are carried over from the input DEM, so fractional burn + values are truncated when the DEM has an integer type. + method: How to choose the elevation to burn. "constant" writes the burn + value for the cell's mask value. "relative" subtracts the burn value + from the DEM, preserving relief inside the region. "statistic" writes + each region's own statistic of the DEM beneath it. + mask_values: Which mask values identify regions, as a comma separated + string such as "1,3" or a sequence. If None, the keys of burn_values + are used; if those are absent too, every non zero, non nodata mask + value is selected and each distinct value still forms its own regions. + burn_values: Required for the constant and relative methods, ignored by the + statistic method. Either a mapping of mask value to burn value, given + as a string such as "1:225.5,3:210.0" or as a dict, or a single number + applied to every selected mask value. + statistic: Which statistic the statistic method computes per region, one of + "min", "max" or "mean". Default is "min". + burn_offset: Subtracted from each computed statistic, so + statistic="min" with burn_offset=1.0 writes each region's minimum + elevation less one. Only used by the statistic method. Default is 0.0. + connectivity: 8 (default) to treat diagonally touching cells as one region, + or 4 to require a shared edge. + chunk_size: Size of processing chunks in pixels. Use chunk_size <= 1 for + in-memory processing. Default is 2048. + working_dir: Directory for temporary files during tiled processing with the + statistic method. If None, uses system temp directory. + progress_callback: Optional callback function for progress reporting. + Receives a float value between 0 and 1. + + Raises: + ValueError: If the DEM and mask are not co-registered, if the DEM has no + nodata value, or if any of the method, statistic, connectivity, mask + value or burn value arguments are unrecognized or inconsistent. + + Note: + DEM nodata cells are never written and are excluded from region statistics. + A region lying entirely over DEM nodata is left untouched. + """ + _validate_raster_compatibility(dem_path, mask_path) + method_flag, selection, statistic_flag, neighbor_offsets = _resolve_burn_arguments( + mask_path, method, mask_values, burn_values, statistic, connectivity + ) + if chunk_size <= 1: + _burn_mask_core( + dem_path, + mask_path, + output_path, + method_flag, + selection, + statistic_flag, + burn_offset, + neighbor_offsets, + ) + else: + _burn_mask_tiled( + dem_path, + mask_path, + output_path, + method_flag, + selection, + statistic_flag, + burn_offset, + neighbor_offsets, + connectivity, + chunk_size, + working_dir, + progress_callback, + ) + + def fill( input_path: str, output_path: str | None = None, @@ -363,6 +465,7 @@ def flow_length( __all__ = [ # Core functions "breach", + "burn", "fill", "flow_direction", "accumulation", diff --git a/src/overflow/_basins/tiled/global_state.py b/src/overflow/_basins/tiled/global_state.py index 7ca38a0..fd6b8bc 100644 --- a/src/overflow/_basins/tiled/global_state.py +++ b/src/overflow/_basins/tiled/global_state.py @@ -72,9 +72,12 @@ class GlobalState: Handles the flow direction and watershed connections at a corner between two tiles. - _handle_edge(self, tile_a: TileEdgeData, tile_b: TileEdgeData, side_a: Side): Handles the flow direction and watershed connections at an edge between two tiles. - - _join_adjacent_tiles(self, tile_a: TileEdgeData, tile_b: TileEdgeData, tile_c: TileEdgeData, - tile_d: TileEdgeData): Joins the adjacent tiles by handling the flow direction and watershed connections at - their edges and corners. + - _join_east_neighbor(self, tile_a: TileEdgeData, tile_b: TileEdgeData): + Joins a tile to the tile on its east across their shared vertical edge. + - _join_south_neighbor(self, tile_a: TileEdgeData, tile_c: TileEdgeData): + Joins a tile to the tile on its south across their shared horizontal edge. + - _join_diagonal_neighbors(self, tile_a: TileEdgeData, tile_b: TileEdgeData, tile_c: TileEdgeData, + tile_d: TileEdgeData): Joins the two diagonal pairs meeting where four tiles touch. - complete_graph(self): Completes the graph of watershed connections by iterating over all tiles and joining adjacent tiles. """ @@ -194,7 +197,24 @@ def _handle_edge( tile_b.watershed.get_side(side_b)[i - 1] ) - def _join_adjacent_tiles( + def _join_east_neighbor(self, tile_a: TileEdgeData, tile_b: TileEdgeData): + """ + Connect a tile to the tile on its east across their shared vertical edge. + + Handled in both directions because a watershed link is directional: it is + recorded only where a cell's flow direction actually crosses the seam. + """ + self._handle_edge(tile_a, tile_b, Side.RIGHT) + self._handle_edge(tile_b, tile_a, Side.LEFT) + + def _join_south_neighbor(self, tile_a: TileEdgeData, tile_c: TileEdgeData): + """ + Connect a tile to the tile on its south across their shared horizontal edge. + """ + self._handle_edge(tile_a, tile_c, Side.BOTTOM) + self._handle_edge(tile_c, tile_a, Side.TOP) + + def _join_diagonal_neighbors( self, tile_a: TileEdgeData, tile_b: TileEdgeData, @@ -202,41 +222,53 @@ def _join_adjacent_tiles( tile_d: TileEdgeData, ): """ - Connect the edges and corners of adjacent tiles. - Between the four adjacent tiles, the connections are as follows: + Connect the two diagonal pairs meeting where four tiles touch. + + - - + - - + | A | B | + - - * - - + | C | D | + - - + - - + - A and B are connected at their adjacent edge. - B and D are connected at their adjacent edge. - D and C are connected at their adjacent edge. - C and A are connected at their adjacent edge. - A and D are connected at their corners. - C and B are connected at their corners. + + A and D are connected at their corners, as are C and B. Every other + diagonal adjacency is already covered by the edge joins. """ self._handle_corner(tile_a, tile_d, Corner.BOTTOM_RIGHT) self._handle_corner(tile_d, tile_a, Corner.TOP_LEFT) self._handle_corner(tile_b, tile_c, Corner.BOTTOM_LEFT) self._handle_corner(tile_c, tile_b, Corner.TOP_RIGHT) - self._handle_edge(tile_a, tile_b, Side.RIGHT) - self._handle_edge(tile_b, tile_a, Side.LEFT) - self._handle_edge(tile_c, tile_d, Side.RIGHT) - self._handle_edge(tile_d, tile_c, Side.LEFT) - self._handle_edge(tile_a, tile_c, Side.BOTTOM) - self._handle_edge(tile_c, tile_a, Side.TOP) - self._handle_edge(tile_b, tile_d, Side.BOTTOM) - self._handle_edge(tile_d, tile_b, Side.TOP) def complete_graph(self): """ Complete the global graph by joining adjacent tiles. + + Walks every tile and joins it to its eastern and southern neighbors plus + the two diagonal pairs at its southeastern corner, so every adjacency in + the tile grid is covered exactly once. + + The bounds guards are load bearing. Iterating a 2x2 block over + range(tile_rows - 1) x range(tile_cols - 1) instead visits nothing at all + when the tile grid is a single tile row or column, silently leaving every + seam unjoined so basins are cut off at tile boundaries. """ - for tile_row in range(self.tile_rows - 1): - for tile_col in range(self.tile_cols - 1): + for tile_row in range(self.tile_rows): + for tile_col in range(self.tile_cols): tile_a = self._get_tile_edge_data(tile_row, tile_col) - tile_b = self._get_tile_edge_data(tile_row, tile_col + 1) - tile_c = self._get_tile_edge_data(tile_row + 1, tile_col) - tile_d = self._get_tile_edge_data(tile_row + 1, tile_col + 1) - self._join_adjacent_tiles(tile_a, tile_b, tile_c, tile_d) + has_east = tile_col + 1 < self.tile_cols + has_south = tile_row + 1 < self.tile_rows + + if has_east: + self._join_east_neighbor( + tile_a, self._get_tile_edge_data(tile_row, tile_col + 1) + ) + if has_south: + self._join_south_neighbor( + tile_a, self._get_tile_edge_data(tile_row + 1, tile_col) + ) + if has_east and has_south: + self._join_diagonal_neighbors( + tile_a, + self._get_tile_edge_data(tile_row, tile_col + 1), + self._get_tile_edge_data(tile_row + 1, tile_col), + self._get_tile_edge_data(tile_row + 1, tile_col + 1), + ) diff --git a/src/overflow/_burn_mask/__init__.py b/src/overflow/_burn_mask/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/overflow/_burn_mask/core/__init__.py b/src/overflow/_burn_mask/core/__init__.py new file mode 100644 index 0000000..bb0908d --- /dev/null +++ b/src/overflow/_burn_mask/core/__init__.py @@ -0,0 +1,47 @@ +from .burn_mask import ( + CORE_LABEL_OFFSET, + REGION_STAT_COLUMNS, + REGION_STAT_COUNT, + REGION_STAT_MAX, + REGION_STAT_MIN, + REGION_STAT_SUM, + MaskSelection, + _burn_mask_core, + accumulate_region_stats, + apply_burn, + build_burn_lookup, + build_mask_selection, + burn_direct, + label_regions, + parse_burn_method, + parse_burn_values, + parse_connectivity, + parse_mask_values, + parse_statistic, + region_value, + resolve_burn_arguments, +) + +__all__ = [ + "CORE_LABEL_OFFSET", + "REGION_STAT_COLUMNS", + "REGION_STAT_COUNT", + "REGION_STAT_MAX", + "REGION_STAT_MIN", + "REGION_STAT_SUM", + "MaskSelection", + "_burn_mask_core", + "accumulate_region_stats", + "apply_burn", + "build_burn_lookup", + "build_mask_selection", + "burn_direct", + "label_regions", + "parse_burn_method", + "parse_burn_values", + "parse_connectivity", + "parse_mask_values", + "parse_statistic", + "region_value", + "resolve_burn_arguments", +] diff --git a/src/overflow/_burn_mask/core/burn_mask.py b/src/overflow/_burn_mask/core/burn_mask.py new file mode 100644 index 0000000..476b2d2 --- /dev/null +++ b/src/overflow/_burn_mask/core/burn_mask.py @@ -0,0 +1,770 @@ +import numpy as np +from numba import boolean, float64, njit, prange # type: ignore[attr-defined] +from numba.experimental import jitclass +from numba.typed import Dict # type: ignore[attr-defined] +from numba.types import int64 as nb_int64 + +from overflow._util.constants import ( + BURN_METHOD_CONSTANT, + BURN_METHOD_RELATIVE, + BURN_METHOD_STATISTIC, + BURN_NO_REGION_LABEL, + BURN_STAT_MAX, + BURN_STAT_MEAN, + BURN_STAT_MIN, + NEIGHBOR_OFFSETS_4, + NEIGHBOR_OFFSETS_8, +) +from overflow._util.queue import Int64PairQueue +from overflow._util.raster import create_dataset, open_dataset + +# number of columns in the per region statistics array: min, max, sum, count +REGION_STAT_COLUMNS = 4 +# first label used by the in memory path; BURN_NO_REGION_LABEL is reserved +CORE_LABEL_OFFSET = 1 +REGION_STAT_MIN = 0 +REGION_STAT_MAX = 1 +REGION_STAT_SUM = 2 +REGION_STAT_COUNT = 3 + +_BURN_METHODS = { + "constant": BURN_METHOD_CONSTANT, + "relative": BURN_METHOD_RELATIVE, + "statistic": BURN_METHOD_STATISTIC, +} +_BURN_STATISTICS = { + "min": BURN_STAT_MIN, + "max": BURN_STAT_MAX, + "mean": BURN_STAT_MEAN, +} + + +@jitclass( + [ + ("values", float64[:]), + ("burns", float64[:]), + ("select_all", boolean), + ("nodata", float64), + ("has_nodata", boolean), + ("broadcast_burn", float64), + ("has_broadcast", boolean), + ] +) +class MaskSelection: + """ + Describes which mask raster values identify regions to burn, and for the + constant and relative methods how much to burn for each of those values. + + Selection is by exact equality against the mask value, so classified masks + should hold exact class values. Mask nodata and NaN are never selected. + + Attributes: + values (float64[:]): The mask values that identify regions. Empty when + select_all is True. + burns (float64[:]): The burn value for each entry of values, in the same + order. Empty for the statistic method or when has_broadcast is True. + select_all (boolean): If True, every non zero, non nodata mask value + identifies regions. Each distinct value still forms its own regions. + nodata (float64): The mask raster's nodata value. + has_nodata (boolean): Whether the mask raster declares a nodata value. + broadcast_burn (float64): A single burn value applied to every selected + mask value. + has_broadcast (boolean): Whether broadcast_burn should be used in place of + the per value burns array. + """ + + def __init__( + self, + values: np.ndarray, + burns: np.ndarray, + select_all: bool, + nodata: float, + has_nodata: bool, + broadcast_burn: float, + has_broadcast: bool, + ): + self.values = values + self.burns = burns + self.select_all = select_all + self.nodata = nodata + self.has_nodata = has_nodata + self.broadcast_burn = broadcast_burn + self.has_broadcast = has_broadcast + + def is_selected(self, value) -> bool: + """Return True if a mask value identifies a region to burn. + + Args: + value: A cell value from the mask raster. + + Returns: + bool: True if the value identifies a region. + """ + # value != value is a dtype agnostic NaN check; numba rejects np.isnan on ints + if value != value: + return False + if self.has_nodata and value == self.nodata: + return False + if self.select_all: + return bool(value != 0) + # an explicit loop rather than any(): numba does not compile generator + # expressions inside jitclass methods + for i in range(self.values.shape[0]): # noqa: SIM110 + if value == self.values[i]: + return True + return False + + def burn_for(self, value) -> float: + """Return the burn value configured for a selected mask value. + + Only meaningful for the constant and relative methods. Callers must have + checked is_selected first. + + Args: + value: A cell value from the mask raster. + + Returns: + float: The configured burn value, or NaN if the value has none. + """ + if self.has_broadcast: + return self.broadcast_burn + for i in range(self.values.shape[0]): + if value == self.values[i]: + return float(self.burns[i]) + return np.nan + + +@njit +def label_regions( + mask: np.ndarray, + selection: MaskSelection, + label_offset: int, + neighbor_offsets, + valid_rows: int, + valid_cols: int, +) -> tuple: + """Label each contiguous region of the mask with a unique integer. + + A region is a maximal connected set of cells that share the same mask value and + whose value is selected. Cells of two different selected values never join into + one region even when they touch. This is a flood fill in the shape of + Algorithm 4 LabelFlats (https://rbarnes.org/sci/2014_flats.pdf), with equality + of mask value standing in for equality of elevation. + + Labels are assigned from label_offset upwards. Callers processing tiles pass a + per tile offset so that labels are globally unique without a renumbering pass. + + Tiles that hang off the right or bottom edge of the raster are padded out to a + full square by the reader. Those padded cells must not be labeled: if they were, + two genuinely disconnected regions on either side of a tile seam could be joined + through padding that shares their mask value. valid_rows and valid_cols bound + the real data. + + Args: + mask (np.ndarray): The mask raster tile. + selection (MaskSelection): Which mask values identify regions. + label_offset (int): The first label to assign. + neighbor_offsets (np.ndarray): Row/column offsets defining connectivity, + either NEIGHBOR_OFFSETS_4 or NEIGHBOR_OFFSETS_8. + valid_rows (int): Number of rows of the tile that hold real raster data. + valid_cols (int): Number of columns of the tile that hold real raster data. + + Returns: + tuple: (labels, region_count) where labels is an int64 array the same shape + as mask holding BURN_NO_REGION_LABEL outside regions, and region_count + is the number of regions labeled. + """ + n_row, n_col = mask.shape + n_row = min(n_row, valid_rows) + n_col = min(n_col, valid_cols) + labels = np.full(mask.shape, BURN_NO_REGION_LABEL, dtype=np.int64) + region_count = 0 + # one queue reused across regions; allocating per region is wasteful when the + # mask holds many small regions + to_fill = Int64PairQueue([(np.int64(0), np.int64(0))]) + to_fill.pop() + for row in range(n_row): + for col in range(n_col): + if labels[row, col] != BURN_NO_REGION_LABEL: + continue + value = mask[row, col] + if not selection.is_selected(value): + continue + new_label = np.int64(label_offset + region_count) + region_count += 1 + to_fill.push((np.int64(row), np.int64(col))) + while len(to_fill) > 0: + fill_row, fill_col = to_fill.pop() + not_in_bounds = ( + fill_row < 0 + or fill_row >= n_row + or fill_col < 0 + or fill_col >= n_col + ) + if not_in_bounds: + continue + if labels[fill_row, fill_col] != BURN_NO_REGION_LABEL: + continue + if mask[fill_row, fill_col] != value: + continue + labels[fill_row, fill_col] = new_label + for k in range(neighbor_offsets.shape[0]): + to_fill.push( + ( + np.int64(fill_row + neighbor_offsets[k, 0]), + np.int64(fill_col + neighbor_offsets[k, 1]), + ) + ) + return labels, region_count + + +@njit +def accumulate_region_stats( + dem: np.ndarray, + labels: np.ndarray, + label_offset: int, + region_count: int, + dem_nodata: float, +) -> np.ndarray: + """Accumulate min, max, sum and count of the DEM under each labeled region. + + These four quantities are all mergeable: the statistics of a region split + across tiles are recovered by taking the min of the mins, the max of the maxes + and the sums of the sums and counts. That is what lets min, max and mean be + computed exactly for regions that straddle tile boundaries. + + DEM nodata and NaN cells are excluded, so a region lying entirely over DEM + nodata ends up with a count of zero. + + Args: + dem (np.ndarray): The DEM tile. + labels (np.ndarray): Region labels for the same tile. + label_offset (int): The first label used by this tile. + region_count (int): The number of regions labeled in this tile. + dem_nodata (float): The DEM's nodata value. + + Returns: + np.ndarray: A float64 array of shape (region_count, 4) holding + (min, max, sum, count) for each region, ordered by label. + """ + stats = np.zeros((region_count, REGION_STAT_COLUMNS), dtype=np.float64) + stats[:, REGION_STAT_MIN] = np.inf + stats[:, REGION_STAT_MAX] = -np.inf + n_row, n_col = dem.shape + for row in range(n_row): + for col in range(n_col): + label = labels[row, col] + if label == BURN_NO_REGION_LABEL: + continue + elevation = dem[row, col] + if elevation == dem_nodata or np.isnan(elevation): + continue + index = label - label_offset + value = np.float64(elevation) + if value < stats[index, REGION_STAT_MIN]: + stats[index, REGION_STAT_MIN] = value + if value > stats[index, REGION_STAT_MAX]: + stats[index, REGION_STAT_MAX] = value + stats[index, REGION_STAT_SUM] += value + stats[index, REGION_STAT_COUNT] += 1.0 + return stats + + +@njit +def region_value( + minimum: float, + maximum: float, + total: float, + count: float, + statistic: int, + burn_offset: float, +) -> float: + """Reduce a region's accumulated statistics to the elevation to burn. + + Args: + minimum (float): The region's minimum DEM value. + maximum (float): The region's maximum DEM value. + total (float): The sum of the region's DEM values. + count (float): The number of valid DEM cells in the region. + statistic (int): One of BURN_STAT_MIN, BURN_STAT_MAX, BURN_STAT_MEAN. + burn_offset (float): Subtracted from the computed statistic. + + Returns: + float: The elevation to burn, or NaN if the region has no valid DEM + cells and should be left untouched. + """ + if count == 0.0: + return np.nan + if statistic == BURN_STAT_MIN: + value = minimum + elif statistic == BURN_STAT_MAX: + value = maximum + else: + value = total / count + return value - burn_offset + + +@njit +def build_burn_lookup( + stats: np.ndarray, label_offset: int, statistic: int, burn_offset: float +): + """Build a label to burn value lookup from a single array of region statistics. + + Used by the in memory path, where every region is already whole. The tiled path + builds its lookup in GlobalState.solve_regions instead, after merging regions + that span tiles. + + Args: + stats (np.ndarray): Region statistics from accumulate_region_stats. + label_offset (int): The first label the statistics correspond to. + statistic (int): One of BURN_STAT_MIN, BURN_STAT_MAX, BURN_STAT_MEAN. + burn_offset (float): Subtracted from each computed statistic. + + Returns: + Dict: A numba typed dict mapping each label to its burn value. Regions with + no valid DEM cells map to NaN. + """ + lookup = Dict.empty(key_type=nb_int64, value_type=float64) + for index in range(stats.shape[0]): + lookup[np.int64(label_offset + index)] = region_value( + stats[index, REGION_STAT_MIN], + stats[index, REGION_STAT_MAX], + stats[index, REGION_STAT_SUM], + stats[index, REGION_STAT_COUNT], + statistic, + burn_offset, + ) + return lookup + + +@njit(nogil=True, parallel=True) +def burn_direct( + dem: np.ndarray, + mask: np.ndarray, + selection: MaskSelection, + method: int, + dem_nodata: float, + tile_row: int, + tile_col: int, +) -> tuple: + """Burn constant or relative values into a DEM tile, in place. + + Both methods are pure per cell functions of the DEM value and the mask value, + so no region labeling is needed and the tiled driver can run them in a single + pass with no cross tile reconciliation. + + Args: + dem (np.ndarray): The DEM tile, modified in place. + mask (np.ndarray): The mask tile. + selection (MaskSelection): Which mask values to burn and by how much. + method (int): BURN_METHOD_CONSTANT or BURN_METHOD_RELATIVE. + dem_nodata (float): The DEM's nodata value. + tile_row (int): The tile's row in the tile grid, passed through to the + caller so results can be routed without tracking futures. + tile_col (int): The tile's column in the tile grid. + + Returns: + tuple: (dem, tile_row, tile_col) with dem modified in place. + """ + n_row, n_col = dem.shape + for row in prange(n_row): + for col in range(n_col): + elevation = dem[row, col] + if elevation == dem_nodata or np.isnan(elevation): + continue + value = mask[row, col] + if not selection.is_selected(value): + continue + burn = selection.burn_for(value) + if np.isnan(burn): + continue + if method == BURN_METHOD_CONSTANT: + dem[row, col] = burn + else: + dem[row, col] = elevation - burn + return dem, tile_row, tile_col + + +@njit(nogil=True, parallel=True) +def apply_burn( + dem: np.ndarray, + labels: np.ndarray, + burn_lookup, + dem_nodata: float, + tile_row: int, + tile_col: int, +) -> tuple: + """Write each region's solved burn value into a DEM tile, in place. + + Args: + dem (np.ndarray): The DEM tile, modified in place. + labels (np.ndarray): Region labels for the same tile. + burn_lookup: Numba typed dict mapping every label present in labels to its + burn value. Labels mapping to NaN are left untouched. + dem_nodata (float): The DEM's nodata value. + tile_row (int): The tile's row in the tile grid, passed through to the + caller so results can be routed without tracking futures. + tile_col (int): The tile's column in the tile grid. + + Returns: + tuple: (dem, tile_row, tile_col) with dem modified in place. + """ + n_row, n_col = dem.shape + for row in prange(n_row): + for col in range(n_col): + label = labels[row, col] + if label == BURN_NO_REGION_LABEL: + continue + elevation = dem[row, col] + if elevation == dem_nodata or np.isnan(elevation): + continue + # every label written by label_regions is present in the lookup by + # construction, so this cannot miss + burn = burn_lookup[label] + if np.isnan(burn): + continue + dem[row, col] = burn + return dem, tile_row, tile_col + + +def parse_burn_method(method: str) -> int: + """Translate the public burn method string into its integer flag. + + Args: + method (str): One of "constant", "relative", "statistic". + + Returns: + int: The matching BURN_METHOD_* constant. + + Raises: + ValueError: If the method is not recognized. + """ + try: + return _BURN_METHODS[method] + except KeyError: + raise ValueError( + f"Unknown burn method '{method}', expected one of: " + f"{', '.join(_BURN_METHODS)}" + ) from None + + +def parse_statistic(statistic: str) -> int: + """Translate the public statistic string into its integer flag. + + Args: + statistic (str): One of "min", "max", "mean". + + Returns: + int: The matching BURN_STAT_* constant. + + Raises: + ValueError: If the statistic is not recognized. + """ + try: + return _BURN_STATISTICS[statistic] + except KeyError: + raise ValueError( + f"Unknown statistic '{statistic}', expected one of: " + f"{', '.join(_BURN_STATISTICS)}" + ) from None + + +def parse_connectivity(connectivity: int) -> np.ndarray: + """Translate a connectivity setting into its neighbor offsets. + + Args: + connectivity (int): Either 4 or 8. + + Returns: + np.ndarray: NEIGHBOR_OFFSETS_4 or NEIGHBOR_OFFSETS_8. + + Raises: + ValueError: If the connectivity is neither 4 nor 8. + """ + if connectivity == 8: + return NEIGHBOR_OFFSETS_8 + if connectivity == 4: + return NEIGHBOR_OFFSETS_4 + raise ValueError(f"Connectivity must be 4 or 8, got {connectivity}") + + +def parse_burn_values(burn_values: str | float | dict) -> dict | float: + """Parse the burn values argument into a mapping or a single broadcast value. + + Accepts the CLI's paired mapping string ("1:225.5,3:210.0"), a bare number + string applied to every selected mask value ("1.5"), or the already parsed + Python forms of either. + + Args: + burn_values (str | float | dict): The value to parse. + + Returns: + dict | float: A mask value to burn value mapping, or a single float to + broadcast to every selected mask value. + + Raises: + ValueError: If a mapping entry is malformed or holds a non numeric value. + """ + if isinstance(burn_values, dict): + return {float(key): float(value) for key, value in burn_values.items()} + if not isinstance(burn_values, str): + return float(burn_values) + text = burn_values.strip() + if ":" not in text: + try: + return float(text) + except ValueError: + raise ValueError( + f"Could not parse burn values '{burn_values}', expected a number " + "or a mapping like '1:225.5,3:210.0'" + ) from None + parsed = {} + for entry in text.split(","): + if not entry.strip(): + continue + parts = entry.split(":") + if len(parts) != 2: + raise ValueError( + f"Could not parse burn value entry '{entry}', expected 'value:burn'" + ) + try: + parsed[float(parts[0])] = float(parts[1]) + except ValueError: + raise ValueError( + f"Burn value entry '{entry}' has a non numeric part" + ) from None + if not parsed: + raise ValueError(f"Burn values '{burn_values}' contained no entries") + return parsed + + +def parse_mask_values(mask_values: str | list | tuple | None) -> list | None: + """Parse the mask values argument into a list of floats. + + Args: + mask_values (str | list | tuple | None): A comma separated string such as + "1,3", an already parsed sequence, or None. + + Returns: + list | None: The parsed mask values, or None if none were given. + + Raises: + ValueError: If an entry is not numeric or the string holds no entries. + """ + if mask_values is None: + return None + if not isinstance(mask_values, str): + return [float(value) for value in mask_values] + entries = [entry.strip() for entry in mask_values.split(",") if entry.strip()] + if not entries: + raise ValueError(f"Mask values '{mask_values}' contained no entries") + try: + return [float(entry) for entry in entries] + except ValueError: + raise ValueError( + f"Could not parse mask values '{mask_values}', expected a comma " + "separated list of numbers like '1,3'" + ) from None + + +def build_mask_selection( + method: int, + mask_values: str | list | tuple | None, + burn_values: str | float | dict | None, + mask_nodata: float | None, +) -> MaskSelection: + """Resolve the mask value and burn value arguments into a MaskSelection. + + When mask_values is omitted, the mapping keys of burn_values supply it. When + neither names specific values, every non zero, non nodata mask value is + selected and each distinct value still forms its own regions. + + Args: + method (int): One of the BURN_METHOD_* constants. + mask_values (str | list | tuple | None): Mask values identifying regions. + burn_values (str | float | dict | None): Burn values for the constant and + relative methods. + mask_nodata (float | None): The mask raster's nodata value, if it has one. + + Returns: + MaskSelection: The resolved selection. + + Raises: + ValueError: If burn values are missing or incomplete for the constant and + relative methods, or if a selected mask value is the mask's nodata. + """ + values = parse_mask_values(mask_values) + needs_burn_values = method in (BURN_METHOD_CONSTANT, BURN_METHOD_RELATIVE) + + parsed_burns: dict | float | None = None + if burn_values is not None: + parsed_burns = parse_burn_values(burn_values) + elif needs_burn_values: + raise ValueError( + "burn_values is required for the constant and relative methods" + ) + + if values is None and isinstance(parsed_burns, dict): + values = sorted(parsed_burns) + + if values is not None and mask_nodata is not None: + conflicting = [value for value in values if value == mask_nodata] + if conflicting: + raise ValueError( + f"Mask value {conflicting[0]} is the mask raster's nodata value and " + "cannot identify a region" + ) + + burns_array = np.empty(0, dtype=np.float64) + broadcast_burn = np.nan + has_broadcast = False + if needs_burn_values: + if isinstance(parsed_burns, dict): + assert values is not None + missing = [value for value in values if value not in parsed_burns] + if missing: + raise ValueError( + f"No burn value given for mask value(s) " + f"{', '.join(str(value) for value in missing)}" + ) + burns_array = np.array( + [parsed_burns[value] for value in values], dtype=np.float64 + ) + else: + broadcast_burn = float(parsed_burns) # type: ignore[arg-type] + has_broadcast = True + + return MaskSelection( + np.array(values if values is not None else [], dtype=np.float64), + burns_array, + values is None, + float(mask_nodata) if mask_nodata is not None else np.nan, + mask_nodata is not None, + broadcast_burn, + has_broadcast, + ) + + +def resolve_burn_arguments( + mask_path: str, + method: str, + mask_values: str | list | tuple | None, + burn_values: str | float | dict | None, + statistic: str, + connectivity: int, +) -> tuple: + """Validate and resolve the public burn arguments into their internal forms. + + Shared by the in memory and tiled paths so both agree on what the arguments + mean, and so every argument is checked before any raster work starts. + + Args: + mask_path (str): Path to the mask raster, opened to read its nodata value. + method (str): One of "constant", "relative", "statistic". + mask_values (str | list | tuple | None): Mask values identifying regions. + burn_values (str | float | dict | None): Burn values for the constant and + relative methods. + statistic (str): One of "min", "max", "mean". + connectivity (int): Either 4 or 8. + + Returns: + tuple: (method_flag, selection, statistic_flag, neighbor_offsets). + + Raises: + ValueError: If any argument is unrecognized, malformed or inconsistent. + """ + method_flag = parse_burn_method(method) + statistic_flag = parse_statistic(statistic) + neighbor_offsets = parse_connectivity(connectivity) + mask_ds = open_dataset(mask_path) + mask_nodata = mask_ds.GetRasterBand(1).GetNoDataValue() + mask_ds = None + selection = build_mask_selection(method_flag, mask_values, burn_values, mask_nodata) + return method_flag, selection, statistic_flag, neighbor_offsets + + +def _burn_mask_core( + dem_path: str, + mask_path: str, + output_path: str, + method: int, + selection: MaskSelection, + statistic: int, + burn_offset: float, + neighbor_offsets: np.ndarray, +) -> None: + """ + Burn mask regions into a DEM entirely in memory. + + Suitable for DEMs that fit in RAM. Region labeling sees the whole raster at + once, so no cross tile reconciliation is needed. + + Parameters + ---------- + dem_path : str + Path to the input DEM raster. + mask_path : str + Path to the co-registered mask raster. + output_path : str + Path to the output burned DEM raster. + method : int + One of the BURN_METHOD_* constants. + selection : MaskSelection + Which mask values identify regions and, for the constant and relative + methods, how much to burn for each. + statistic : int + One of the BURN_STAT_* constants. Only used by the statistic method. + burn_offset : float + Subtracted from each computed statistic. Only used by the statistic method. + neighbor_offsets : np.ndarray + Row/column offsets defining region connectivity. + """ + dem_ds = open_dataset(dem_path) + mask_ds = open_dataset(mask_path) + dem_band = dem_ds.GetRasterBand(1) + mask_band = mask_ds.GetRasterBand(1) + dem_nodata = dem_band.GetNoDataValue() + if dem_nodata is None: + raise ValueError("Input DEM must have a no data value") + + dem = dem_band.ReadAsArray() + mask = mask_band.ReadAsArray() + + if method == BURN_METHOD_STATISTIC: + labels, region_count = label_regions( + mask, + selection, + CORE_LABEL_OFFSET, + neighbor_offsets, + mask.shape[0], + mask.shape[1], + ) + stats = accumulate_region_stats( + dem, labels, CORE_LABEL_OFFSET, region_count, dem_nodata + ) + burn_lookup = build_burn_lookup( + stats, CORE_LABEL_OFFSET, statistic, burn_offset + ) + dem, _, _ = apply_burn(dem, labels, burn_lookup, dem_nodata, 0, 0) + else: + dem, _, _ = burn_direct(dem, mask, selection, method, dem_nodata, 0, 0) + + output_ds = create_dataset( + output_path, + dem_nodata, + dem_band.DataType, + dem_band.XSize, + dem_band.YSize, + dem_ds.GetGeoTransform(), + dem_ds.GetProjection(), + ) + output_band = output_ds.GetRasterBand(1) + output_band.WriteArray(dem) + output_band.FlushCache() + output_ds.FlushCache() + output_band = None + output_ds = None + dem_band = None + mask_band = None + dem_ds = None + mask_ds = None diff --git a/src/overflow/_burn_mask/tiled/__init__.py b/src/overflow/_burn_mask/tiled/__init__.py new file mode 100644 index 0000000..aec5ce6 --- /dev/null +++ b/src/overflow/_burn_mask/tiled/__init__.py @@ -0,0 +1,11 @@ +from .burn_mask_tiled import _burn_mask_tiled, label_tile, tile_valid_extent +from .global_state import GlobalState, handle_corner, handle_edge + +__all__ = [ + "GlobalState", + "_burn_mask_tiled", + "handle_corner", + "handle_edge", + "label_tile", + "tile_valid_extent", +] diff --git a/src/overflow/_burn_mask/tiled/burn_mask_tiled.py b/src/overflow/_burn_mask/tiled/burn_mask_tiled.py new file mode 100644 index 0000000..165ae62 --- /dev/null +++ b/src/overflow/_burn_mask/tiled/burn_mask_tiled.py @@ -0,0 +1,458 @@ +import concurrent.futures +import math +import os +import queue +import shutil +import tempfile +import time +from threading import Lock + +import numba +import numpy as np +from numba import njit # type: ignore[attr-defined] +from osgeo import gdal + +from overflow._burn_mask.core import ( + MaskSelection, + accumulate_region_stats, + apply_burn, + burn_direct, + label_regions, +) +from overflow._util.constants import BURN_METHOD_STATISTIC, BURN_NO_REGION_LABEL +from overflow._util.perimeter import get_tile_perimeter +from overflow._util.progress import ProgressCallback, ProgressTracker, silent_callback +from overflow._util.raster import ( + RasterChunk, + create_dataset, + open_dataset, + raster_chunker, +) + +from .global_state import GlobalState + +# temporary file for storing region labels +LABELS_FILENAME = "burn_labels.tif" + + +def setup_working_dir(working_dir: str | None) -> tuple[str, bool]: + """ + Setup working directory for burning mask regions into a DEM. + """ + cleanup_working_dir = False + if working_dir is None: + working_dir = tempfile.mkdtemp() + cleanup_working_dir = True + return working_dir, cleanup_working_dir + + +def setup_datasets( + dem_path: str, mask_path: str, output_path: str +) -> tuple[gdal.Dataset, gdal.Dataset, gdal.Dataset, float, float | None]: + """ + Open the input DEM and mask and create the output DEM. + + Returns the three datasets, the DEM's nodata value and the mask's nodata value + (which may be None, since binary masks commonly do not declare one). + """ + dem_ds = open_dataset(dem_path) + mask_ds = open_dataset(mask_path) + dem_band = dem_ds.GetRasterBand(1) + dem_nodata = dem_band.GetNoDataValue() + if dem_nodata is None: + raise ValueError("Input DEM must have a no data value") + mask_nodata = mask_ds.GetRasterBand(1).GetNoDataValue() + output_ds = create_dataset( + output_path, + dem_nodata, + dem_band.DataType, + dem_band.XSize, + dem_band.YSize, + dem_ds.GetGeoTransform(), + dem_ds.GetProjection(), + ) + return dem_ds, mask_ds, output_ds, dem_nodata, mask_nodata + + +def tile_valid_extent( + tile_row: int, tile_col: int, chunk_size: int, band: gdal.Band +) -> tuple[int, int]: + """ + Return how many rows and columns of a tile hold real raster data. + + Tiles hanging off the right or bottom edge of the raster are padded out to a + full square by the reader, and those padded cells must be excluded from region + labeling. + """ + valid_rows = min(chunk_size, max(band.YSize - tile_row * chunk_size, 0)) + valid_cols = min(chunk_size, max(band.XSize - tile_col * chunk_size, 0)) + return valid_rows, valid_cols + + +@njit(nogil=True) +def label_tile( + dem: np.ndarray, + mask: np.ndarray, + selection: MaskSelection, + neighbor_offsets: np.ndarray, + dem_nodata: float, + valid_rows: int, + valid_cols: int, + tile_row: int, + tile_col: int, + global_state: GlobalState, +) -> tuple: + """ + Label the mask regions in one tile and accumulate their DEM statistics. + + Labels are drawn from a range private to this tile, so they are globally unique + without a renumbering pass. The tile's label and mask perimeters are stashed in + the global state; they are all the later solve needs to discover which of these + labels are really the same region. + + This is called in parallel for each tile. Each tile writes only its own row of + the perimeter arrays, so no lock is needed here; the statistics are merged into + the global dicts by the caller under a lock instead. + """ + tile_index = np.int64(tile_row) * np.int64(global_state.num_cols) + np.int64( + tile_col + ) + chunk_size_squared = np.int64(global_state.chunk_size) * np.int64( + global_state.chunk_size + ) + label_offset = chunk_size_squared * tile_index + np.int64(BURN_NO_REGION_LABEL + 1) + labels, region_count = label_regions( + mask, selection, label_offset, neighbor_offsets, valid_rows, valid_cols + ) + stats = accumulate_region_stats(dem, labels, label_offset, region_count, dem_nodata) + global_state.label_perimeters[tile_index] = get_tile_perimeter(labels) + global_state.mask_perimeters[tile_index] = get_tile_perimeter(mask).astype( + np.float64 + ) + return labels, stats, label_offset, tile_row, tile_col + + +def _burn_direct_tiled( + dem_ds: gdal.Dataset, + mask_ds: gdal.Dataset, + output_ds: gdal.Dataset, + method: int, + selection: MaskSelection, + dem_nodata: float, + mask_fill: float | None, + chunk_size: int, + tracker: ProgressTracker, +) -> None: + """ + Burn constant or relative values in a single tiled pass. + + Neither method depends on anything outside the cell being written, so there is + no labeling, no temporary raster and no cross tile reconciliation. + """ + dem_band = dem_ds.GetRasterBand(1) + mask_band = mask_ds.GetRasterBand(1) + output_band = output_ds.GetRasterBand(1) + + max_workers = numba.config.NUMBA_NUM_THREADS # type: ignore[attr-defined] + task_queue: queue.Queue[int] = queue.Queue(max_workers) + lock = Lock() + chunk_counter = [0] # Use list for mutability in closure + total_chunks = math.ceil(dem_band.YSize / chunk_size) * math.ceil( + dem_band.XSize / chunk_size + ) + + def handle_burn_tile_result(future): + dem, tile_row, tile_col = future.result() + with lock: + dem_tile = RasterChunk(tile_row, tile_col, chunk_size, 0) + dem_tile.from_numpy(dem) + dem_tile.write(output_band) + task_queue.get() + chunk_counter[0] += 1 + tracker.callback(message=f"Chunk {chunk_counter[0]}/{total_chunks}") + + tracker.update(1, step_name="Burn regions") + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + for dem_tile in raster_chunker(dem_band, chunk_size): + while task_queue.full(): + time.sleep(0.1) + task_queue.put(0) + mask_tile = RasterChunk(dem_tile.row, dem_tile.col, chunk_size, 0) + mask_tile.read(mask_band, mask_fill) + future = executor.submit( + burn_direct, + dem_tile.data, + mask_tile.data, + selection, + method, + dem_nodata, + dem_tile.row, + dem_tile.col, + ) + future.add_done_callback(handle_burn_tile_result) + + while not task_queue.empty(): + time.sleep(0.1) + + output_band.FlushCache() + output_ds.FlushCache() + + +def _burn_statistic_tiled( + dem_ds: gdal.Dataset, + mask_ds: gdal.Dataset, + output_ds: gdal.Dataset, + selection: MaskSelection, + statistic: int, + burn_offset: float, + neighbor_offsets: np.ndarray, + connectivity: int, + dem_nodata: float, + mask_fill: float | None, + chunk_size: int, + working_dir: str, + tracker: ProgressTracker, +) -> None: + """ + Burn a per region statistic in three passes. + + 1. Label each tile's regions independently and accumulate min, max, sum and + count of the DEM under each label. + 2. Walk the tile seams, join labels that turn out to belong to the same region, + and merge their statistics. This is what makes a region that straddles two + or more tiles come out with one statistic rather than one per fragment. + 3. Re-read the original DEM and write each region's solved value. + """ + dem_band = dem_ds.GetRasterBand(1) + mask_band = mask_ds.GetRasterBand(1) + output_band = output_ds.GetRasterBand(1) + + labels_ds = create_dataset( + os.path.join(working_dir, LABELS_FILENAME), + BURN_NO_REGION_LABEL, + gdal.GDT_Int64, + dem_band.XSize, + dem_band.YSize, + dem_ds.GetGeoTransform(), + dem_ds.GetProjection(), + ) + labels_band = labels_ds.GetRasterBand(1) + + n_chunks_row = math.ceil(dem_band.YSize / chunk_size) + n_chunks_col = math.ceil(dem_band.XSize / chunk_size) + global_state = GlobalState(n_chunks_row, n_chunks_col, chunk_size, connectivity) + + max_workers = numba.config.NUMBA_NUM_THREADS # type: ignore[attr-defined] + task_queue: queue.Queue[int] = queue.Queue(max_workers) + lock = Lock() + total_chunks = n_chunks_row * n_chunks_col + + # pass 1: label regions and accumulate statistics per tile + chunk_counter = [0] # Use list for mutability in closure + + def handle_label_tile_result(future): + labels, stats, label_offset, tile_row, tile_col = future.result() + with lock: + global_state.merge_tile_stats(label_offset, stats) + labels_tile = RasterChunk(tile_row, tile_col, chunk_size, 0) + labels_tile.from_numpy(labels) + labels_tile.write(labels_band) + task_queue.get() + chunk_counter[0] += 1 + tracker.callback(message=f"Chunk {chunk_counter[0]}/{total_chunks}") + + tracker.update(1, step_name="Label regions") + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + for dem_tile in raster_chunker(dem_band, chunk_size): + while task_queue.full(): + time.sleep(0.1) + task_queue.put(0) + mask_tile = RasterChunk(dem_tile.row, dem_tile.col, chunk_size, 0) + mask_tile.read(mask_band, mask_fill) + valid_rows, valid_cols = tile_valid_extent( + dem_tile.row, dem_tile.col, chunk_size, dem_band + ) + future = executor.submit( + label_tile, + dem_tile.data, + mask_tile.data, + selection, + neighbor_offsets, + dem_nodata, + valid_rows, + valid_cols, + dem_tile.row, + dem_tile.col, + global_state, + ) + future.add_done_callback(handle_label_tile_result) + + while not task_queue.empty(): + time.sleep(0.1) + + # flush cache between writing and reading + labels_band.FlushCache() + labels_ds.FlushCache() + + # pass 2: join regions that span tiles and reduce them to burn values + tracker.update(2, step_name="Solve regions") + global_state.connect_tile_edges_and_corners() + burn_lookup = global_state.solve_regions(statistic, burn_offset) + + # pass 3: write the solved value for each region + chunk_counter_apply = [0] # Use list for mutability in closure + + def handle_apply_tile_result(future): + dem, tile_row, tile_col = future.result() + with lock: + dem_tile = RasterChunk(tile_row, tile_col, chunk_size, 0) + dem_tile.from_numpy(dem) + dem_tile.write(output_band) + task_queue.get() + chunk_counter_apply[0] += 1 + tracker.callback(message=f"Chunk {chunk_counter_apply[0]}/{total_chunks}") + + tracker.update(3, step_name="Apply burn values") + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + for dem_tile in raster_chunker(dem_band, chunk_size): + while task_queue.full(): + time.sleep(0.1) + task_queue.put(0) + labels_tile = RasterChunk(dem_tile.row, dem_tile.col, chunk_size, 0) + labels_tile.read(labels_band) + future = executor.submit( + apply_burn, + dem_tile.data, + labels_tile.data, + burn_lookup, + dem_nodata, + dem_tile.row, + dem_tile.col, + ) + future.add_done_callback(handle_apply_tile_result) + + while not task_queue.empty(): + time.sleep(0.1) + + output_band.FlushCache() + output_ds.FlushCache() + labels_ds.Close() + + +def _burn_mask_tiled( + dem_path: str, + mask_path: str, + output_path: str, + method: int, + selection: MaskSelection, + statistic: int, + burn_offset: float, + neighbor_offsets: np.ndarray, + connectivity: int, + chunk_size: int, + working_dir: str | None = None, + progress_callback: ProgressCallback | None = None, +) -> None: + """ + Burn mask regions into a DEM using a parallel tiled approach. + + The input DEM is divided into tiles that are processed concurrently by numba + kernels compiled with nogil, following the same threading pattern as the other + tiled tools in this package. + + The constant and relative methods are per cell functions of the DEM and mask + values, so they run in a single pass. The statistic method needs each contiguous + region's statistic, which a single tile cannot know for a region that continues + into its neighbors, so it runs in three passes: label and accumulate per tile, + join labels across tile seams and merge their statistics, then apply. + + Parameters + ---------- + dem_path : str + Path to the input DEM raster. + mask_path : str + Path to the co-registered mask raster. + output_path : str + Path to the output burned DEM raster. + method : int + One of the BURN_METHOD_* constants. + selection : MaskSelection + Which mask values identify regions and, for the constant and relative + methods, how much to burn for each. + statistic : int + One of the BURN_STAT_* constants. Only used by the statistic method. + burn_offset : float + Subtracted from each computed statistic. Only used by the statistic method. + neighbor_offsets : np.ndarray + Row/column offsets defining region connectivity. + connectivity : int + Either 4 or 8, matching neighbor_offsets. + chunk_size : int + Size of each tile in cells. + working_dir : str | None + Directory for the temporary labels raster used by the statistic method. + If None, a system temp directory is created and removed afterwards. + progress_callback : ProgressCallback | None + Optional callback for progress updates. If None, the operation runs + silently. + + Returns + ------- + None + """ + if progress_callback is None: + progress_callback = silent_callback + is_statistic = method == BURN_METHOD_STATISTIC + tracker = ProgressTracker( + progress_callback, "Burning mask regions", total_steps=3 if is_statistic else 1 + ) + + dem_ds, mask_ds, output_ds, dem_nodata, mask_nodata = setup_datasets( + dem_path, mask_path, output_path + ) + # tiles at the right and bottom raster edges are padded out to a full square. + # The pad value is irrelevant to the result because label_regions ignores cells + # outside the raster, but the reader still needs something to fill with when the + # mask declares no nodata value. + mask_fill = None if mask_nodata is not None else 0.0 + + cleanup_working_dir = False + try: + if is_statistic: + working_dir, cleanup_working_dir = setup_working_dir(working_dir) + _burn_statistic_tiled( + dem_ds, + mask_ds, + output_ds, + selection, + statistic, + burn_offset, + neighbor_offsets, + connectivity, + dem_nodata, + mask_fill, + chunk_size, + working_dir, + tracker, + ) + else: + _burn_direct_tiled( + dem_ds, + mask_ds, + output_ds, + method, + selection, + dem_nodata, + mask_fill, + chunk_size, + tracker, + ) + finally: + output_ds.Close() + mask_ds.Close() + dem_ds.Close() + if cleanup_working_dir and working_dir is not None: + shutil.rmtree(working_dir) diff --git a/src/overflow/_burn_mask/tiled/global_state.py b/src/overflow/_burn_mask/tiled/global_state.py new file mode 100644 index 0000000..fe85998 --- /dev/null +++ b/src/overflow/_burn_mask/tiled/global_state.py @@ -0,0 +1,332 @@ +import numpy as np +from numba import njit # type: ignore[attr-defined] +from numba.experimental import jitclass +from numba.typed import Dict # type: ignore[attr-defined] +from numba.types import DictType, float64, int64 + +from overflow._burn_mask.core import ( + REGION_STAT_COUNT, + REGION_STAT_MAX, + REGION_STAT_MIN, + REGION_STAT_SUM, + region_value, +) +from overflow._util.constants import BURN_NO_REGION_LABEL +from overflow._util.perimeter import Float64Perimeter, Int64Perimeter +from overflow._util.raster import Corner, Side +from overflow._util.union_find import UnionFind + + +@njit +def handle_edge( + labels_a: np.ndarray, + masks_a: np.ndarray, + labels_b: np.ndarray, + masks_b: np.ndarray, + union_find: UnionFind, + connectivity: int, +) -> None: + """ + Join regions that continue across the shared edge of two adjacent tiles. + + Both label vectors run in the same direction along the seam, so index i in + tile A sits directly against index i in tile B. With 8-connectivity a cell also + touches its two diagonal neighbors across the seam, which is why the inner loop + scans i-1, i and i+1; with 4-connectivity only i is adjacent. + + Two labels are joined when both cells are part of a region and hold the same + mask value, which is the same rule the within-tile flood fill applies. + + Args: + labels_a (np.ndarray): Region labels from tile A along the shared edge. + masks_a (np.ndarray): Mask values from tile A along the shared edge. + labels_b (np.ndarray): Region labels from tile B along the shared edge. + masks_b (np.ndarray): Mask values from tile B along the shared edge. + union_find (UnionFind): The global disjoint set of region labels, modified + in place. + connectivity (int): Either 4 or 8. + + Returns: + None. The function modifies union_find in place. + """ + for i in range(labels_a.shape[0]): + label_a = labels_a[i] + if label_a == BURN_NO_REGION_LABEL: + continue + value_a = masks_a[i] + first = i - 1 if connectivity == 8 else i + last = i + 1 if connectivity == 8 else i + for ni in range(first, last + 1): + if ni < 0 or ni >= labels_b.shape[0]: + continue + label_b = labels_b[ni] + if label_b == BURN_NO_REGION_LABEL: + continue + if masks_b[ni] != value_a: + continue + union_find.union(label_a, label_b) + + +@njit +def handle_corner( + label_a: int64, + value_a: float64, + label_b: int64, + value_b: float64, + union_find: UnionFind, +) -> None: + """ + Join regions that continue across the diagonal touch point of two tiles. + + Only reachable with 8-connectivity, and only for the two diagonal pairs meeting + at the point where four tiles join. Every other diagonal adjacency is already + covered by handle_edge. + + Args: + label_a (int64): Region label at tile A's corner. + value_a (float64): Mask value at tile A's corner. + label_b (int64): Region label at tile B's corner. + value_b (float64): Mask value at tile B's corner. + union_find (UnionFind): The global disjoint set of region labels, modified + in place. + + Returns: + None. The function modifies union_find in place. + """ + if label_a == BURN_NO_REGION_LABEL or label_b == BURN_NO_REGION_LABEL: + return + if value_a != value_b: + return + union_find.union(label_a, label_b) + + +@jitclass( + [ + ("label_perimeters", int64[:, :]), + ("mask_perimeters", float64[:, :]), + ("region_min", DictType(int64, float64)), + ("region_max", DictType(int64, float64)), + ("region_sum", DictType(int64, float64)), + ("region_count", DictType(int64, float64)), + ] +) +class GlobalState: + """ + Global state for burning per region statistics across a tiled raster. + + Each tile labels its own regions from a private label range, so a region that + straddles a tile boundary starts out as two or more independent labels with + independent statistics. This class holds what is needed to put them back + together: the label and mask values around each tile's perimeter, a disjoint set + over labels, and the per label statistics. + + The statistics are kept as four parallel dicts rather than one dict of records + because a numba typed dict has a single value type. + + Attributes: + union_find (UnionFind): Disjoint set joining labels of the same region. + label_perimeters (np.ndarray): Per tile region labels around the perimeter. + mask_perimeters (np.ndarray): Per tile mask values around the perimeter, + held as float64 so one class covers every mask raster data type. + region_min (DictType(int64, float64)): Per label minimum DEM value. + region_max (DictType(int64, float64)): Per label maximum DEM value. + region_sum (DictType(int64, float64)): Per label sum of DEM values. + region_count (DictType(int64, float64)): Per label count of valid DEM cells. + num_rows (int64): Number of tile rows. + num_cols (int64): Number of tile columns. + chunk_size (int64): Size of each tile in cells. + connectivity (int64): Either 4 or 8. + """ + + union_find: UnionFind + label_perimeters: np.ndarray + mask_perimeters: np.ndarray + region_min: DictType(int64, float64) # type: ignore[valid-type] + region_max: DictType(int64, float64) # type: ignore[valid-type] + region_sum: DictType(int64, float64) # type: ignore[valid-type] + region_count: DictType(int64, float64) # type: ignore[valid-type] + num_rows: int64 + num_cols: int64 + chunk_size: int64 + connectivity: int64 + + def __init__( + self, num_rows: int64, num_cols: int64, chunk_size: int64, connectivity: int64 + ): + self.union_find = UnionFind() + self.num_rows = num_rows + self.num_cols = num_cols + self.chunk_size = chunk_size + self.connectivity = connectivity + self.region_min = Dict.empty(key_type=int64, value_type=float64) + self.region_max = Dict.empty(key_type=int64, value_type=float64) + self.region_sum = Dict.empty(key_type=int64, value_type=float64) + self.region_count = Dict.empty(key_type=int64, value_type=float64) + tile_count = num_rows * num_cols + perimeter_cell_count = 4 * chunk_size - 4 + self.label_perimeters = np.full( + (tile_count, perimeter_cell_count), BURN_NO_REGION_LABEL, dtype=np.int64 + ) + self.mask_perimeters = np.zeros( + (tile_count, perimeter_cell_count), dtype=np.float64 + ) + + def row_col_to_tile_index(self, row: int64, col: int64) -> int64: + """Return the flat tile index of the tile at the given tile row and column.""" + return row * self.num_cols + col + + def _get_label_perimeter(self, tile_index: int64) -> Int64Perimeter: + return Int64Perimeter( + self.label_perimeters[tile_index], + self.chunk_size, + self.chunk_size, + tile_index, + ) + + def _get_mask_perimeter(self, tile_index: int64) -> Float64Perimeter: + return Float64Perimeter( + self.mask_perimeters[tile_index], + self.chunk_size, + self.chunk_size, + tile_index, + ) + + def merge_tile_stats(self, label_offset: int64, stats: np.ndarray) -> None: + """Record one tile's per region statistics. + + Labels are unique to the tile that created them, so this is a plain insert + with no merging. Callers must serialize this, since numba typed dicts are + not safe to insert into concurrently. + + Args: + label_offset (int64): The first label used by the tile. + stats (np.ndarray): The tile's (region_count, 4) statistics array. + """ + for index in range(stats.shape[0]): + label = int64(label_offset + index) + self.region_min[label] = stats[index, REGION_STAT_MIN] + self.region_max[label] = stats[index, REGION_STAT_MAX] + self.region_sum[label] = stats[index, REGION_STAT_SUM] + self.region_count[label] = stats[index, REGION_STAT_COUNT] + + def connect_tile_edges_and_corners(self) -> None: + """Join every pair of labels that meet across a tile boundary. + + Walks each tile once and joins it to its eastern and southern neighbors, + plus the two diagonal pairs meeting at its southeastern corner. Every + adjacency in the tile grid is covered exactly once, and the bounds guards + keep it correct for degenerate grids that are a single tile wide or tall. + """ + for row_index in range(self.num_rows): + for col_index in range(self.num_cols): + # + - - + - - + + # | A | B | + # + - - * - - + + # | C | D | + # + - - + - - + + tile_index_a = self.row_col_to_tile_index(row_index, col_index) + has_east = col_index + 1 < self.num_cols + has_south = row_index + 1 < self.num_rows + if has_east: + tile_index_b = self.row_col_to_tile_index(row_index, col_index + 1) + self._combine_edge( + tile_index_a, Side.RIGHT, tile_index_b, Side.LEFT + ) + if has_south: + tile_index_c = self.row_col_to_tile_index(row_index + 1, col_index) + self._combine_edge( + tile_index_a, Side.BOTTOM, tile_index_c, Side.TOP + ) + if has_east and has_south and self.connectivity == 8: + tile_index_b = self.row_col_to_tile_index(row_index, col_index + 1) + tile_index_c = self.row_col_to_tile_index(row_index + 1, col_index) + tile_index_d = self.row_col_to_tile_index( + row_index + 1, col_index + 1 + ) + self._combine_corner( + tile_index_a, + Corner.BOTTOM_RIGHT, + tile_index_d, + Corner.TOP_LEFT, + ) + self._combine_corner( + tile_index_b, + Corner.BOTTOM_LEFT, + tile_index_c, + Corner.TOP_RIGHT, + ) + + def _combine_edge( + self, tile_index_a: int64, side_a: Side, tile_index_b: int64, side_b: Side + ) -> None: + """Join regions across the shared edge of two adjacent tiles.""" + handle_edge( + self._get_label_perimeter(tile_index_a).get_side(side_a), + self._get_mask_perimeter(tile_index_a).get_side(side_a), + self._get_label_perimeter(tile_index_b).get_side(side_b), + self._get_mask_perimeter(tile_index_b).get_side(side_b), + self.union_find, + self.connectivity, + ) + + def _combine_corner( + self, + tile_index_a: int64, + corner_a: Corner, + tile_index_b: int64, + corner_b: Corner, + ) -> None: + """Join regions across the diagonal touch point of two tiles.""" + handle_corner( + self._get_label_perimeter(tile_index_a).get_corner(corner_a), + self._get_mask_perimeter(tile_index_a).get_corner(corner_a), + self._get_label_perimeter(tile_index_b).get_corner(corner_b), + self._get_mask_perimeter(tile_index_b).get_corner(corner_b), + self.union_find, + ) + + def solve_regions(self, statistic: int64, burn_offset: float64): + """Merge per tile statistics into whole regions and reduce them to values. + + Every label is resolved to its representative, the four mergeable + statistics are combined per representative, and the result is expanded back + into a flat lookup so the apply pass is a single dict read per cell. + + Args: + statistic (int64): One of the BURN_STAT_* constants. + burn_offset (float64): Subtracted from each computed statistic. + + Returns: + Dict: A numba typed dict mapping every label to its region's burn value. + Regions with no valid DEM cells map to NaN, meaning leave alone. + """ + root_min = Dict.empty(key_type=int64, value_type=float64) + root_max = Dict.empty(key_type=int64, value_type=float64) + root_sum = Dict.empty(key_type=int64, value_type=float64) + root_count = Dict.empty(key_type=int64, value_type=float64) + for label in self.region_min: + root = self.union_find.find(label) + if root in root_min: + if self.region_min[label] < root_min[root]: + root_min[root] = self.region_min[label] + if self.region_max[label] > root_max[root]: + root_max[root] = self.region_max[label] + root_sum[root] += self.region_sum[label] + root_count[root] += self.region_count[label] + else: + root_min[root] = self.region_min[label] + root_max[root] = self.region_max[label] + root_sum[root] = self.region_sum[label] + root_count[root] = self.region_count[label] + burn_lookup = Dict.empty(key_type=int64, value_type=float64) + for label in self.region_min: + root = self.union_find.find(label) + burn_lookup[label] = region_value( + root_min[root], + root_max[root], + root_sum[root], + root_count[root], + statistic, + burn_offset, + ) + return burn_lookup diff --git a/src/overflow/_fill_depressions/tiled/global_state.py b/src/overflow/_fill_depressions/tiled/global_state.py index be7365b..0e6976c 100644 --- a/src/overflow/_fill_depressions/tiled/global_state.py +++ b/src/overflow/_fill_depressions/tiled/global_state.py @@ -218,49 +218,70 @@ def _get_tile_perimeter(self, tile_index: int64) -> PerimeterData: return PerimeterData(labels, elevations) def connect_tile_edges_and_corners(self): - """Combine the spillover graphs of all tiles into a single global graph.""" - for row_index in range(self.num_rows - 1): - for col_index in range(self.num_cols - 1): - # + - - + - - + - # | A | B | - # + - - * - - + - # | C | D | - # + - - + - - + + """Combine the spillover graphs of all tiles into a single global graph. + + Walks every tile and joins it to the tile on its east and the tile on its + south, plus the two diagonal pairs meeting at its southeastern corner: + + + - - + - - + + | A | B | + + - - * - - + + | C | D | + + - - + - - + + + A joins B across a vertical seam and C across a horizontal one; the corner + joins A-D and B-C cover the diagonal adjacencies at the point marked *, + which no edge scan reaches. Every adjacency in the tile grid is covered + exactly once. + + The bounds guards are load bearing. Iterating a 2x2 block over + range(num_rows - 1) x range(num_cols - 1) instead visits nothing at all + when the tile grid is a single tile row or column, which silently leaves + every seam unjoined and fills straddling depressions to the wrong level. + """ + for row_index in range(self.num_rows): + for col_index in range(self.num_cols): tile_index_a = self._row_col_to_tile_index(row_index, col_index) - tile_index_b = self._row_col_to_tile_index(row_index, col_index + 1) - tile_index_c = self._row_col_to_tile_index(row_index + 1, col_index) - tile_index_d = self._row_col_to_tile_index(row_index + 1, col_index + 1) perimeter_data_a = self._get_tile_perimeter(tile_index_a) - perimeter_data_b = self._get_tile_perimeter(tile_index_b) - perimeter_data_c = self._get_tile_perimeter(tile_index_c) - perimeter_data_d = self._get_tile_perimeter(tile_index_d) - - # Connect A-B, A-C, D-B, D-C edges - self._combine_edge( - perimeter_data_a, Side.RIGHT, perimeter_data_b, Side.LEFT - ) - self._combine_edge( - perimeter_data_a, Side.BOTTOM, perimeter_data_c, Side.TOP - ) - self._combine_edge( - perimeter_data_d, Side.TOP, perimeter_data_b, Side.BOTTOM - ) - self._combine_edge( - perimeter_data_d, Side.LEFT, perimeter_data_c, Side.RIGHT - ) - # Connect A-D and B-C corners - self._combine_corner( - perimeter_data_a, - Corner.BOTTOM_RIGHT, - perimeter_data_d, - Corner.TOP_LEFT, - ) - self._combine_corner( - perimeter_data_b, - Corner.BOTTOM_LEFT, - perimeter_data_c, - Corner.TOP_RIGHT, - ) + has_east = col_index + 1 < self.num_cols + has_south = row_index + 1 < self.num_rows + + if has_east: + perimeter_data_b = self._get_tile_perimeter( + self._row_col_to_tile_index(row_index, col_index + 1) + ) + self._combine_edge( + perimeter_data_a, Side.RIGHT, perimeter_data_b, Side.LEFT + ) + if has_south: + perimeter_data_c = self._get_tile_perimeter( + self._row_col_to_tile_index(row_index + 1, col_index) + ) + self._combine_edge( + perimeter_data_a, Side.BOTTOM, perimeter_data_c, Side.TOP + ) + if has_east and has_south: + perimeter_data_b = self._get_tile_perimeter( + self._row_col_to_tile_index(row_index, col_index + 1) + ) + perimeter_data_c = self._get_tile_perimeter( + self._row_col_to_tile_index(row_index + 1, col_index) + ) + perimeter_data_d = self._get_tile_perimeter( + self._row_col_to_tile_index(row_index + 1, col_index + 1) + ) + self._combine_corner( + perimeter_data_a, + Corner.BOTTOM_RIGHT, + perimeter_data_d, + Corner.TOP_LEFT, + ) + self._combine_corner( + perimeter_data_b, + Corner.BOTTOM_LEFT, + perimeter_data_c, + Corner.TOP_RIGHT, + ) def _combine_edge( self, diff --git a/src/overflow/_resolve_flats/tiled/global_state.py b/src/overflow/_resolve_flats/tiled/global_state.py index 27ab32d..c6c6add 100644 --- a/src/overflow/_resolve_flats/tiled/global_state.py +++ b/src/overflow/_resolve_flats/tiled/global_state.py @@ -62,8 +62,9 @@ def complete_graph(self): """ Complete the global graph by joining adjacent tiles. - This method iterates over each pair of adjacent tiles in the DEM and joins their - edges and corners using the `join_adjacent_tiles` method of the global graph. + This method walks every tile and joins it to the tile on its east and the + tile on its south, plus the two diagonal pairs meeting at its southeastern + corner. Every adjacency in the tile grid is covered exactly once. The adjacent tiles are represented as follows: + - - + - - + @@ -71,22 +72,38 @@ def complete_graph(self): + - - * - - + | C | D | + - - + - - + + + The bounds guards are load bearing. Iterating a 2x2 block over + range(tile_rows - 1) x range(tile_cols - 1) instead visits nothing at all + when the tile grid is a single tile row or column, silently leaving every + seam unjoined so flats straddling a tile boundary are resolved as if they + ended there. """ - for tile_row in range(self.graph.tile_rows - 1): - for tile_col in range(self.graph.tile_cols - 1): + for tile_row in range(self.graph.tile_rows): + for tile_col in range(self.graph.tile_cols): elevations_a = self.elevations[self._get_tile_index(tile_row, tile_col)] - elevations_b = self.elevations[ - self._get_tile_index(tile_row, tile_col + 1) - ] - elevations_c = self.elevations[ - self._get_tile_index(tile_row + 1, tile_col) - ] - elevations_d = self.elevations[ - self._get_tile_index(tile_row + 1, tile_col + 1) - ] - self.graph.join_adjacent_tiles( - elevations_a, elevations_b, elevations_c, elevations_d - ) + has_east = tile_col + 1 < self.graph.tile_cols + has_south = tile_row + 1 < self.graph.tile_rows + + if has_east: + self.graph.join_east_neighbor( + elevations_a, + self.elevations[self._get_tile_index(tile_row, tile_col + 1)], + ) + if has_south: + self.graph.join_south_neighbor( + elevations_a, + self.elevations[self._get_tile_index(tile_row + 1, tile_col)], + ) + if has_east and has_south: + self.graph.join_diagonal_neighbors( + elevations_a, + self.elevations[self._get_tile_index(tile_row, tile_col + 1)], + self.elevations[self._get_tile_index(tile_row + 1, tile_col)], + self.elevations[ + self._get_tile_index(tile_row + 1, tile_col + 1) + ], + ) @njit(nogil=True) diff --git a/src/overflow/_resolve_flats/tiled/graphs.py b/src/overflow/_resolve_flats/tiled/graphs.py index e583071..6196952 100644 --- a/src/overflow/_resolve_flats/tiled/graphs.py +++ b/src/overflow/_resolve_flats/tiled/graphs.py @@ -242,7 +242,33 @@ def _handle_corner( ) self._join_neighbors(global_index_a, global_index_b) - def join_adjacent_tiles( + def join_east_neighbor( + self, elevations_a: Float32Perimeter, elevations_b: Float32Perimeter + ): + """ + Join a tile to the tile on its east across their shared vertical edge. + + Args: + elevations_a (Float32Perimeter): The elevation perimeter of the tile. + elevations_b (Float32Perimeter): The elevation perimeter of its eastern + neighbor. + """ + self._handle_edge(elevations_a, elevations_b, Side.RIGHT, Side.LEFT) + + def join_south_neighbor( + self, elevations_a: Float32Perimeter, elevations_c: Float32Perimeter + ): + """ + Join a tile to the tile on its south across their shared horizontal edge. + + Args: + elevations_a (Float32Perimeter): The elevation perimeter of the tile. + elevations_c (Float32Perimeter): The elevation perimeter of its southern + neighbor. + """ + self._handle_edge(elevations_a, elevations_c, Side.BOTTOM, Side.TOP) + + def join_diagonal_neighbors( self, elevations_a: Float32Perimeter, elevations_b: Float32Perimeter, @@ -250,7 +276,7 @@ def join_adjacent_tiles( elevations_d: Float32Perimeter, ): """ - Join adjacent tiles in the global graph. + Join the two diagonal pairs meeting at the point where four tiles touch. Args: elevations_a (Float32Perimeter): The elevation perimeter of tile A. @@ -264,15 +290,11 @@ def join_adjacent_tiles( + - - * - - + | C | D | + - - + - - + + + Only the A-D and B-C diagonals at the point marked * need handling here. + Every other diagonal adjacency along a seam is already covered by the edge + joins, which compare each cell against its neighbor's i-1, i and i+1. """ - # connect edge A-B - self._handle_edge(elevations_a, elevations_b, Side.RIGHT, Side.LEFT) - # connect edge B-D - self._handle_edge(elevations_b, elevations_d, Side.BOTTOM, Side.TOP) - # connect edge D-C - self._handle_edge(elevations_d, elevations_c, Side.LEFT, Side.RIGHT) - # connect edge C-A - self._handle_edge(elevations_c, elevations_a, Side.TOP, Side.BOTTOM) # connect corner A-D self._handle_corner( elevations_a, elevations_d, Corner.BOTTOM_RIGHT, Corner.TOP_LEFT diff --git a/src/overflow/_util/constants.py b/src/overflow/_util/constants.py index 0d701c6..bc91539 100644 --- a/src/overflow/_util/constants.py +++ b/src/overflow/_util/constants.py @@ -70,6 +70,16 @@ WEIGHTS_NODATA_MODE_ZERO = 0 WEIGHTS_NODATA_MODE_PROPAGATE = 1 EDGE_LABEL = np.int64(1) # edge label of fill depressions algorithm +# burn methods, see the burn tool +BURN_METHOD_CONSTANT = 0 +BURN_METHOD_RELATIVE = 1 +BURN_METHOD_STATISTIC = 2 +# per region statistics available to BURN_METHOD_STATISTIC +BURN_STAT_MIN = 0 +BURN_STAT_MAX = 1 +BURN_STAT_MEAN = 2 +# label used by the burn tool for cells that are not part of any masked region +BURN_NO_REGION_LABEL = 0 # Define flags for each side TOP = 0b0001 # Binary representation of 1 RIGHT = 0b0010 # Binary representation of 2 diff --git a/src/overflow/_util/raster.py b/src/overflow/_util/raster.py index e4a2bb9..ee24bd5 100644 --- a/src/overflow/_util/raster.py +++ b/src/overflow/_util/raster.py @@ -44,7 +44,12 @@ def gdal_data_type_to_numpy_data_type(gdal_dtype: int) -> np.dtype: def read_raster_with_bounds_handling( - x_offset: int, y_offset: int, x_size: int, y_size: int, raster_band: gdal.Band + x_offset: int, + y_offset: int, + x_size: int, + y_size: int, + raster_band: gdal.Band, + fill_value: float | None = None, ) -> np.ndarray: """Read a chunk of a raster band and return it as a numpy array. This function allows for reading out of bounds regions. If the window, or part of the window, extends beyond the edge of the raster, @@ -56,6 +61,10 @@ def read_raster_with_bounds_handling( x_size (int): The number of columns in the chunk. y_size (int): The number of rows in the chunk. raster_band (gdal.Band): The raster band to read from. + fill_value (float | None): Value used to pad out of bounds regions. If None + (the default) the band's nodata value is used and the band is required to + have one. Pass an explicit value to read bands without a nodata value, + such as binary mask rasters. Returns: np.ndarray: The chunk of the raster band as a numpy array. The shape of the array will be: (y_size, x_size). @@ -63,7 +72,7 @@ def read_raster_with_bounds_handling( assert x_size >= 0, "x_size must be positive" assert y_size >= 0, "y_size must be positive" # Get the no data value from the raster band - no_data_value = raster_band.GetNoDataValue() + no_data_value = raster_band.GetNoDataValue() if fill_value is None else fill_value assert no_data_value is not None, "The raster band has no no data value" # Get the GDAL data type from the raster band gdal_dtype = raster_band.DataType @@ -135,7 +144,7 @@ def from_numpy(self, data: np.ndarray): """ self.data = data - def read(self, band: gdal.Band): + def read(self, band: gdal.Band, fill_value: float | None = None): """Read a chunk of a raster band including an overlapping buffer region on all edges. If part of the chunk, including the buffer region, extends beyond the edge of the raster, the out of bounds region will be filled with nodata. The chunk will be stored as a numpy array @@ -143,13 +152,16 @@ def read(self, band: gdal.Band): Args: band (gdal.Band): The raster band to read from. + fill_value (float | None): Value used to pad out of bounds regions. If None + (the default) the band's nodata value is used and the band is required + to have one. """ x_offset = self.col * self.size - self.buffer_size y_offset = self.row * self.size - self.buffer_size x_size = self.size + 2 * self.buffer_size y_size = self.size + 2 * self.buffer_size self.data = read_raster_with_bounds_handling( - x_offset, y_offset, x_size, y_size, band + x_offset, y_offset, x_size, y_size, band, fill_value ) def _get_unbuffered_data(self, band: gdal.Band) -> np.ndarray: diff --git a/src/overflow/_util/union_find.py b/src/overflow/_util/union_find.py new file mode 100644 index 0000000..01ea9a5 --- /dev/null +++ b/src/overflow/_util/union_find.py @@ -0,0 +1,86 @@ +from numba.experimental import jitclass +from numba.typed import Dict # type: ignore[attr-defined] +from numba.types import DictType, int64 + + +@jitclass( + [ + ("parent", DictType(int64, int64)), + ("size", DictType(int64, int64)), + ] +) +class UnionFind: + """ + A disjoint-set (union-find) structure over int64 keys. + + Keys are added lazily, so the structure does not need to know the universe of + labels up front. This matters for tiled processing, where labels are handed out + per tile from a sparse global range and only the ones that actually appear on a + tile perimeter ever need to be joined. + + Uses union by size and iterative path compression, giving near constant time + amortized find and union. Path compression is written as an explicit two pass + loop because numba jitclass methods cannot recurse. + + Attributes: + parent (DictType(int64, int64)): Maps each key to its parent key. A key that + maps to itself is the representative (root) of its set. + size (DictType(int64, int64)): Maps each root key to the number of keys in + its set. Entries for non root keys are stale and must not be read. + """ + + def __init__(self): + self.parent = Dict.empty(key_type=int64, value_type=int64) + self.size = Dict.empty(key_type=int64, value_type=int64) + + def add(self, x: int64) -> None: + """Add a key as a singleton set if it is not already present. + + Args: + x (int64): The key to add. + """ + if x not in self.parent: + self.parent[x] = x + self.size[x] = int64(1) + + def find(self, x: int64) -> int64: + """Return the representative key of the set containing x. + + Adds x as a singleton if it is not already present. Compresses the path + from x to the root so subsequent lookups are direct. + + Args: + x (int64): The key to look up. + + Returns: + int64: The representative key of x's set. + """ + self.add(x) + # walk up to the root + root = x + while self.parent[root] != root: + root = self.parent[root] + # second pass: point every key on the path directly at the root + current = x + while self.parent[current] != root: + next_key = self.parent[current] + self.parent[current] = root + current = next_key + return root + + def union(self, a: int64, b: int64) -> None: + """Merge the sets containing a and b. + + Args: + a (int64): A key in the first set. + b (int64): A key in the second set. + """ + root_a = self.find(a) + root_b = self.find(b) + if root_a == root_b: + return + # attach the smaller tree under the larger one to keep paths short + if self.size[root_a] < self.size[root_b]: + root_a, root_b = root_b, root_a + self.parent[root_b] = root_a + self.size[root_a] += self.size[root_b] diff --git a/src/overflow/cli.py b/src/overflow/cli.py index 83a2a52..c108a2e 100644 --- a/src/overflow/cli.py +++ b/src/overflow/cli.py @@ -9,6 +9,7 @@ __version__, accumulation, breach, + burn, fill, flow_direction, streams, @@ -258,6 +259,141 @@ def flow_direction_cli( raise click.Abort() +@main.command(name="burn") +@click.option( + "--dem_file", + help="path to the GDAL supported raster dataset for the input DEM", + required=True, +) +@click.option( + "--mask_file", + help=( + "path to a binary or classified GDAL supported raster dataset, " + "co-registered with the DEM, whose regions are burned into the DEM" + ), + required=True, +) +@click.option( + "--output_file", + help="path to the output file (must be GeoTiff)", + required=True, +) +@click.option( + "--method", + help=( + "how to choose the elevation to burn: 'constant' writes the burn value " + "for the cell's mask value; 'relative' subtracts the burn value from the " + "DEM, preserving relief; 'statistic' writes each contiguous region's own " + "statistic of the DEM beneath it" + ), + type=click.Choice(["constant", "relative", "statistic"]), + required=True, +) +@click.option( + "--burn_values", + help=( + "required for --method constant and relative: a mapping of mask value to " + "burn value such as '1:225.5,3:210.0', or a single number applied to " + "every selected mask value" + ), + required=False, + default=None, +) +@click.option( + "--mask_values", + help=( + "comma separated mask values identifying regions, such as '1,3'. Defaults " + "to the keys of --burn_values, or to every non zero, non nodata mask value" + ), + required=False, + default=None, +) +@click.option( + "--statistic", + help="only used with --method statistic: the statistic computed per region", + type=click.Choice(["min", "max", "mean"]), + default="min", +) +@click.option( + "--burn_offset", + help=( + "only used with --method statistic: subtracted from each region's " + "statistic, so --statistic min --burn_offset 1.0 writes each region's " + "minimum elevation less one" + ), + type=float, + default=0.0, +) +@click.option( + "--connectivity", + help=("8 treats diagonally touching cells as one region, 4 requires a shared edge"), + type=click.Choice(["4", "8"]), + default="8", +) +@click.option( + "--chunk_size", + help="chunk size (use <= 1 for in-memory processing)", + default=DEFAULT_CHUNK_SIZE, +) +@click.option( + "--working_dir", + help="path to the working directory used for temporary files", + required=False, + default=None, +) +def burn_cli( + dem_file: str, + mask_file: str, + output_file: str, + method: str, + burn_values: str | None, + mask_values: str | None, + statistic: str, + burn_offset: float, + connectivity: str, + chunk_size: int, + working_dir: str | None, +): + """ + Burn elevations into a DEM inside the regions of a mask raster. + + This command rewrites the DEM inside regions identified by a binary or + classified mask raster, either to a constant elevation, to the DEM lowered by + a fixed amount, or to a statistic of the DEM beneath each contiguous region. + Regions are recognized as whole even where they straddle processing tiles. + """ + success = False + try: + progress_display = RichProgressDisplay() + with timer("Burn", spinner=False): + with progress_display.progress_context("Burn Mask Regions"): + burn( + dem_file, + mask_file, + output_file, + method, + mask_values, + burn_values, + statistic, + burn_offset, + int(connectivity), + chunk_size, + working_dir, + progress_display.callback, + ) + resource_stats.add_output_file("Burned DEM", output_file) + success = True + + console.print(resource_stats.get_summary_panel(success=success)) + except Exception as exc: + console.print( + f"[bold red]Error:[/bold red] burn failed with the following exception: {str(exc)}" + ) + if not success: + console.print(resource_stats.get_summary_panel(success=False)) + raise click.Abort() + + @main.command(name="accumulation") @click.option( "--input_file", diff --git a/tests/test_basins.py b/tests/test_basins.py index c5cc83e..9f283d8 100644 --- a/tests/test_basins.py +++ b/tests/test_basins.py @@ -171,6 +171,68 @@ def test_label_watersheds_tiled( gdal.Unlink(gpkg_path) +@pytest.mark.parametrize("transpose", [False, True], ids=["1xN grid", "Nx1 grid"]) +def test_label_watersheds_tiled_degenerate_tile_grid(transpose): + """Seams must be joined when the tile grid is a single tile row or column. + + Regression test. Joining seams by walking a 2x2 tile block over + range(rows - 1) x range(cols - 1) iterates zero times on a 1xN or Nx1 tile + grid, leaving every seam unjoined so the watershed graph never links a tile to + its neighbor. + + Every cell here drains west along its row to column 0, then along column 0 to a + single outlet at (4, 0), so the whole raster is one basin. With seams unjoined, + only the tile holding the outlet reaches it and the rest of the raster is left + unlabeled. + """ + fdr = np.full((8, 32), 4, dtype=np.ubyte) # west + fdr[:4, 0] = 6 # south, down column 0 toward the outlet + fdr[5:, 0] = 2 # north, up column 0 toward the outlet + fdr[4, 0] = 4 # west, off the raster edge: the outlet + outlet = (4, 0) + if transpose: + # rotate the whole flow field a quarter turn: west becomes north + fdr = np.full((32, 8), 2, dtype=np.ubyte) # north + fdr[0, :4] = 0 # east, along row 0 toward the outlet + fdr[0, 5:] = 4 # west, along row 0 toward the outlet + fdr[0, 4] = 2 # north, off the raster edge: the outlet + outlet = (0, 4) + + fdr_path = "/vsimem/degenerate_grid_fdr.tif" + output_path = "/vsimem/degenerate_grid_basins.tif" + gpkg_path = "/vsimem/degenerate_grid_basins.gpkg" + driver = gdal.GetDriverByName("GTiff") + rows, cols = fdr.shape + dataset = driver.Create(fdr_path, cols, rows, 1, gdal.GDT_Byte) + dataset.SetGeoTransform([0, 1, 0, 0, 0, 1]) + srs = osr.SpatialReference() + srs.ImportFromEPSG(4326) + dataset.SetProjection(srs.ExportToWkt()) + band = dataset.GetRasterBand(1) + band.SetNoDataValue(FLOW_DIRECTION_NODATA) + band.WriteArray(fdr) + band.FlushCache() + dataset.FlushCache() + band = None + dataset = None + + drainage_points = Dict.empty(UniTuple(int64, 2), int64) + drainage_points[outlet] = 0 + + try: + _label_watersheds_tiled(fdr_path, drainage_points, output_path, chunk_size=8) + ds = gdal.Open(output_path) + watersheds = ds.GetRasterBand(1).ReadAsArray() + ds = None + # the entire raster drains to the one outlet, so it is a single basin + assert np.all(watersheds != 0), "cells beyond the outlet's tile went unlabeled" + assert len(np.unique(watersheds)) == 1 + finally: + for path in (fdr_path, output_path, gpkg_path): + if gdal.VSIStatL(path) is not None: + gdal.Unlink(path) + + @pytest.fixture(name="test_drainage_points_gpkg") def fixture_test_drainage_points_gpkg(test_fdr_with_drainage_points_filepath): """Create a GeoPackage file with drainage points.""" diff --git a/tests/test_burn.py b/tests/test_burn.py new file mode 100644 index 0000000..33ec24f --- /dev/null +++ b/tests/test_burn.py @@ -0,0 +1,680 @@ +import click.testing +import numpy as np +import pytest +from osgeo import gdal + +from overflow import burn +from overflow._util.union_find import UnionFind +from overflow.cli import burn_cli + +NODATA = -9999.0 +# chunk sizes chosen so the same region is split differently by each one, and so +# that some of them leave a partial tile hanging off the raster edge +TILED_CHUNK_SIZES = [3, 4, 5, 7, 8, 13, 64] + + +def write_raster(path, array, nodata, data_type): + """Write a numpy array to an in-memory GeoTiff for testing.""" + driver = gdal.GetDriverByName("GTiff") + dataset = driver.Create(path, array.shape[1], array.shape[0], 1, data_type) + dataset.SetGeoTransform([0.0, 1.0, 0.0, 0.0, 0.0, -1.0]) + band = dataset.GetRasterBand(1) + band.WriteArray(array) + if nodata is not None: + band.SetNoDataValue(nodata) + dataset.FlushCache() + band = None + dataset = None + return path + + +def unlink_if_exists(path): + """Remove an in-memory raster, tolerating one that was never created.""" + if gdal.VSIStatL(path) is not None: + gdal.Unlink(path) + + +def read_raster(path): + """Read the first band of a raster into a numpy array.""" + dataset = gdal.Open(path) + array = dataset.GetRasterBand(1).ReadAsArray() + dataset = None + return array + + +@pytest.fixture(name="dem_path") +def fixture_dem_path(): + """Create a 20x20 DEM whose every cell holds a distinct elevation. + + Distinct values make each region's min, max and mean unambiguous. One cell + inside a region is nodata so it can be checked for exclusion. + + Yields: + str: Path to the DEM raster. + """ + dem = np.arange(400, dtype=np.float32).reshape(20, 20) + dem[2, 2] = NODATA + path = "/vsimem/test_burn_dem.tif" + write_raster(path, dem, NODATA, gdal.GDT_Float32) + yield path + gdal.Unlink(path) + + +@pytest.fixture(name="mask_path") +def fixture_mask_path(): + """Create a 20x20 classified mask with regions that straddle tile seams. + + The regions are, by class: + 1: a wide band across rows 1-3, split by every vertical seam under test + 1: a blob at rows 9-12, split by horizontal seams and disconnected from the band + 1: two diagonally touching cells, one region under 8-connectivity and two + under 4-connectivity + 2: an L whose two arms only meet by passing through other tiles + + Yields: + str: Path to the mask raster. + """ + mask = np.zeros((20, 20), dtype=np.uint8) + mask[1:4, 1:12] = 1 + mask[6:8, 2:16] = 2 + mask[6:16, 14:16] = 2 + mask[9:13, 16:18] = 1 + mask[7, 7] = 1 + mask[8, 8] = 1 + path = "/vsimem/test_burn_mask.tif" + write_raster(path, mask, 0.0, gdal.GDT_Byte) + yield path + gdal.Unlink(path) + + +@pytest.fixture(name="output_path") +def fixture_output_path(): + """Provide a scratch output path that is cleaned up after the test. + + Yields: + str: Path for the burned DEM. + """ + path = "/vsimem/test_burn_output.tif" + yield path + unlink_if_exists(path) + + +@pytest.fixture(name="small_dem_path") +def fixture_small_dem_path(): + """Create a 5x5 DEM with a simple gradient. + + Yields: + str: Path to the DEM raster. + """ + dem = np.array( + [ + [10, 20, 30, 40, 50], + [11, 21, 31, 41, 51], + [12, 22, 32, 42, 52], + [13, 23, 33, 43, 53], + [14, 24, 34, 44, 54], + ], + dtype=np.float32, + ) + path = "/vsimem/test_burn_small_dem.tif" + write_raster(path, dem, NODATA, gdal.GDT_Float32) + yield path + gdal.Unlink(path) + + +@pytest.fixture(name="small_mask_path") +def fixture_small_mask_path(): + """Create a 5x5 mask holding one class 1 region and one class 3 region. + + Yields: + str: Path to the mask raster. + """ + mask = np.array( + [ + [0, 1, 1, 0, 0], + [0, 1, 0, 0, 3], + [0, 0, 0, 3, 3], + [1, 0, 0, 0, 0], + [1, 1, 0, 0, 0], + ], + dtype=np.uint8, + ) + path = "/vsimem/test_burn_small_mask.tif" + write_raster(path, mask, 0.0, gdal.GDT_Byte) + yield path + gdal.Unlink(path) + + +def test_burn_constant(small_dem_path, small_mask_path, output_path): + """Each mask class is set to its own constant elevation.""" + burn( + small_dem_path, + small_mask_path, + output_path, + method="constant", + burn_values="1:100,3:200", + chunk_size=1, + ) + expected = np.array( + [ + [10, 100, 100, 40, 50], + [11, 100, 31, 41, 200], + [12, 22, 32, 200, 200], + [100, 23, 33, 43, 53], + [100, 100, 34, 44, 54], + ], + dtype=np.float32, + ) + assert np.allclose(read_raster(output_path), expected) + + +def test_burn_relative(small_dem_path, small_mask_path, output_path): + """Relative burning subtracts per cell, preserving relief inside a region.""" + burn( + small_dem_path, + small_mask_path, + output_path, + method="relative", + burn_values="1:10,3:20", + chunk_size=1, + ) + expected = np.array( + [ + [10, 10, 20, 40, 50], + [11, 11, 31, 41, 31], + [12, 22, 32, 22, 32], + [3, 23, 33, 43, 53], + [4, 14, 34, 44, 54], + ], + dtype=np.float32, + ) + assert np.allclose(read_raster(output_path), expected) + + +def test_burn_relative_broadcast(small_dem_path, small_mask_path, output_path): + """A bare burn value applies to every selected mask value.""" + burn( + small_dem_path, + small_mask_path, + output_path, + method="relative", + burn_values="2.5", + chunk_size=1, + ) + result = read_raster(output_path) + mask = read_raster(small_mask_path) + dem = read_raster(small_dem_path) + assert np.allclose(result[mask != 0], dem[mask != 0] - 2.5) + assert np.allclose(result[mask == 0], dem[mask == 0]) + + +@pytest.mark.parametrize( + "statistic,expected_region_1,expected_region_3", + [("min", 20.0, 42.0), ("max", 30.0, 52.0), ("mean", 71.0 / 3, 145.0 / 3)], +) +def test_burn_statistic( + small_dem_path, + small_mask_path, + output_path, + statistic, + expected_region_1, + expected_region_3, +): + """Each contiguous region is flattened to its own statistic.""" + burn( + small_dem_path, + small_mask_path, + output_path, + method="statistic", + mask_values="1,3", + statistic=statistic, + chunk_size=1, + ) + result = read_raster(output_path) + # the upper class 1 region: DEM values 20, 30, 21 + assert np.allclose(result[0, 1], expected_region_1) + assert np.allclose(result[0, 2], expected_region_1) + assert np.allclose(result[1, 1], expected_region_1) + # the class 3 region: DEM values 51, 42, 52 + assert np.allclose(result[1, 4], expected_region_3) + assert np.allclose(result[2, 3], expected_region_3) + assert np.allclose(result[2, 4], expected_region_3) + # the lower class 1 region is a separate region and keeps its own statistic + assert not np.allclose(result[3, 0], expected_region_1) + + +def test_burn_statistic_offset(small_dem_path, small_mask_path, output_path): + """burn_offset is subtracted from each region's statistic.""" + burn( + small_dem_path, + small_mask_path, + output_path, + method="statistic", + mask_values="1,3", + statistic="min", + burn_offset=1.5, + chunk_size=1, + ) + result = read_raster(output_path) + assert np.allclose(result[0, 1], 20.0 - 1.5) + assert np.allclose(result[2, 3], 42.0 - 1.5) + + +def test_burn_statistic_leaves_unselected_cells_alone( + small_dem_path, small_mask_path, output_path +): + """Cells outside the selected mask values are copied through unchanged.""" + burn( + small_dem_path, + small_mask_path, + output_path, + method="statistic", + mask_values="3", + statistic="min", + chunk_size=1, + ) + result = read_raster(output_path) + dem = read_raster(small_dem_path) + mask = read_raster(small_mask_path) + assert np.allclose(result[mask != 3], dem[mask != 3]) + assert np.allclose(result[mask == 3], 42.0) + + +def test_burn_default_mask_values_keeps_classes_separate( + small_dem_path, small_mask_path, output_path +): + """Omitting mask values selects every non zero class, each still its own region.""" + burn( + small_dem_path, + small_mask_path, + output_path, + method="statistic", + statistic="min", + chunk_size=1, + ) + result = read_raster(output_path) + assert np.allclose(result[0, 1], 20.0) + assert np.allclose(result[2, 3], 42.0) + assert np.allclose(result[4, 0], 13.0) + + +@pytest.mark.parametrize("connectivity,expected", [(8, 1), (4, 2)]) +def test_burn_connectivity(dem_path, output_path, connectivity, expected): + """Diagonally touching cells are one region under 8-connectivity, two under 4.""" + mask = np.zeros((20, 20), dtype=np.uint8) + mask[5, 5] = 1 + mask[6, 6] = 1 + mask_path = "/vsimem/test_burn_diagonal_mask.tif" + write_raster(mask_path, mask, 0.0, gdal.GDT_Byte) + try: + burn( + dem_path, + mask_path, + output_path, + method="statistic", + mask_values="1", + statistic="min", + connectivity=connectivity, + chunk_size=1, + ) + result = read_raster(output_path) + distinct = {result[5, 5], result[6, 6]} + assert len(distinct) == expected + if expected == 1: + # joined: both take the minimum of the pair + assert np.allclose(result[5, 5], 105.0) + assert np.allclose(result[6, 6], 105.0) + else: + # separate: each keeps its own value + assert np.allclose(result[5, 5], 105.0) + assert np.allclose(result[6, 6], 126.0) + finally: + gdal.Unlink(mask_path) + + +@pytest.mark.parametrize("chunk_size", TILED_CHUNK_SIZES) +@pytest.mark.parametrize( + "kwargs", + [ + {"method": "statistic", "mask_values": "1,2", "statistic": "min"}, + {"method": "statistic", "mask_values": "1,2", "statistic": "max"}, + {"method": "statistic", "mask_values": "1,2", "statistic": "mean"}, + { + "method": "statistic", + "mask_values": "1,2", + "statistic": "min", + "burn_offset": 5.0, + }, + { + "method": "statistic", + "mask_values": "1,2", + "statistic": "min", + "connectivity": 4, + }, + {"method": "constant", "burn_values": "1:100,2:200"}, + {"method": "relative", "burn_values": "1:10,2:20"}, + ], + ids=["min", "max", "mean", "offset", "connectivity4", "constant", "relative"], +) +def test_compare_core_and_tiled(dem_path, mask_path, chunk_size, kwargs): + """Tiled processing matches in-memory processing at every tile size. + + The mask fixture holds regions that straddle vertical seams, horizontal seams + and the point where four tiles meet, so this is the check that the cross tile + region solve actually works. + """ + core_path = "/vsimem/test_burn_core.tif" + tiled_path = "/vsimem/test_burn_tiled.tif" + try: + burn(dem_path, mask_path, core_path, chunk_size=1, **kwargs) + burn(dem_path, mask_path, tiled_path, chunk_size=chunk_size, **kwargs) + assert np.allclose(read_raster(core_path), read_raster(tiled_path)) + finally: + gdal.Unlink(core_path) + gdal.Unlink(tiled_path) + + +@pytest.mark.parametrize("chunk_size", [4, 5, 8]) +def test_burn_tiled_region_spanning_seam(chunk_size, output_path): + """A region split across tiles gets one statistic, not one per fragment.""" + # a single 16 cell horizontal bar that every tested chunk size cuts in two + dem = np.arange(64, dtype=np.float32).reshape(8, 8) + mask = np.zeros((8, 8), dtype=np.uint8) + mask[3, :] = 1 + dem_path = "/vsimem/test_burn_seam_dem.tif" + mask_path = "/vsimem/test_burn_seam_mask.tif" + write_raster(dem_path, dem, NODATA, gdal.GDT_Float32) + write_raster(mask_path, mask, 0.0, gdal.GDT_Byte) + try: + burn( + dem_path, + mask_path, + output_path, + method="statistic", + mask_values="1", + statistic="min", + chunk_size=chunk_size, + ) + result = read_raster(output_path) + # the whole bar takes the minimum of the whole bar, dem[3, 0] == 24 + assert np.allclose(result[3, :], 24.0) + finally: + gdal.Unlink(dem_path) + gdal.Unlink(mask_path) + + +@pytest.mark.parametrize("shape", [(4, 32), (32, 4)]) +def test_burn_tiled_single_tile_row_or_column(shape, output_path): + """Rasters whose tile grid is a single row or column still join their seams. + + A tile grid of 1xN or Nx1 is the case a 2x2 tile stencil would silently skip. + """ + rows, cols = shape + dem = np.arange(rows * cols, dtype=np.float32).reshape(rows, cols) + mask = np.ones((rows, cols), dtype=np.uint8) + dem_path = "/vsimem/test_burn_strip_dem.tif" + mask_path = "/vsimem/test_burn_strip_mask.tif" + write_raster(dem_path, dem, NODATA, gdal.GDT_Float32) + write_raster(mask_path, mask, 0.0, gdal.GDT_Byte) + try: + burn( + dem_path, + mask_path, + output_path, + method="statistic", + mask_values="1", + statistic="max", + chunk_size=8, + ) + # the mask covers the whole raster, so it is one region with one maximum + assert np.allclose(read_raster(output_path), float(rows * cols - 1)) + finally: + gdal.Unlink(dem_path) + gdal.Unlink(mask_path) + + +@pytest.mark.parametrize("chunk_size", [1, 4]) +def test_burn_dem_nodata_excluded_and_preserved(chunk_size, output_path): + """DEM nodata cells are excluded from statistics and never written.""" + dem = np.array( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, NODATA, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + [13.0, 14.0, 15.0, 16.0], + ], + dtype=np.float32, + ) + mask = np.zeros((4, 4), dtype=np.uint8) + mask[0:2, 0:2] = 1 + dem_path = "/vsimem/test_burn_nodata_dem.tif" + mask_path = "/vsimem/test_burn_nodata_mask.tif" + write_raster(dem_path, dem, NODATA, gdal.GDT_Float32) + write_raster(mask_path, mask, 0.0, gdal.GDT_Byte) + try: + burn( + dem_path, + mask_path, + output_path, + method="statistic", + mask_values="1", + statistic="mean", + chunk_size=chunk_size, + ) + result = read_raster(output_path) + # mean over 1, 2, 5 only; the nodata cell contributes nothing + assert np.allclose(result[0, 0], 8.0 / 3) + assert np.allclose(result[0, 1], 8.0 / 3) + assert np.allclose(result[1, 0], 8.0 / 3) + assert result[1, 1] == NODATA + finally: + gdal.Unlink(dem_path) + gdal.Unlink(mask_path) + + +@pytest.mark.parametrize("chunk_size", [1, 4]) +def test_burn_region_entirely_nodata_is_unchanged(chunk_size, output_path): + """A region with no valid DEM cells is left alone rather than filled with junk.""" + dem = np.arange(16, dtype=np.float32).reshape(4, 4) + dem[3, 3] = NODATA + mask = np.zeros((4, 4), dtype=np.uint8) + mask[3, 3] = 1 + mask[0, 0] = 1 + dem_path = "/vsimem/test_burn_empty_dem.tif" + mask_path = "/vsimem/test_burn_empty_mask.tif" + write_raster(dem_path, dem, NODATA, gdal.GDT_Float32) + write_raster(mask_path, mask, 0.0, gdal.GDT_Byte) + try: + burn( + dem_path, + mask_path, + output_path, + method="statistic", + mask_values="1", + statistic="min", + chunk_size=chunk_size, + ) + assert np.allclose(read_raster(output_path), dem) + finally: + gdal.Unlink(dem_path) + gdal.Unlink(mask_path) + + +@pytest.mark.parametrize("chunk_size", [1, 4]) +def test_burn_mask_without_nodata_value(chunk_size, output_path): + """A mask raster that declares no nodata value still works.""" + dem = np.arange(36, dtype=np.float32).reshape(6, 6) + mask = np.zeros((6, 6), dtype=np.uint8) + mask[2:5, 2:5] = 1 + dem_path = "/vsimem/test_burn_nonodata_dem.tif" + mask_path = "/vsimem/test_burn_nonodata_mask.tif" + write_raster(dem_path, dem, NODATA, gdal.GDT_Float32) + write_raster(mask_path, mask, None, gdal.GDT_Byte) + try: + burn( + dem_path, + mask_path, + output_path, + method="statistic", + mask_values="1", + statistic="min", + chunk_size=chunk_size, + ) + result = read_raster(output_path) + assert np.allclose(result[2:5, 2:5], 14.0) + assert np.allclose(result[0, :], dem[0, :]) + finally: + gdal.Unlink(dem_path) + gdal.Unlink(mask_path) + + +def test_burn_output_preserves_geotransform_and_nodata( + small_dem_path, small_mask_path, output_path +): + """The output carries the input DEM's georeferencing, data type and nodata.""" + burn( + small_dem_path, + small_mask_path, + output_path, + method="constant", + burn_values="1:1,3:1", + chunk_size=1, + ) + source = gdal.Open(small_dem_path) + result = gdal.Open(output_path) + assert result.GetGeoTransform() == source.GetGeoTransform() + assert result.GetRasterBand(1).DataType == source.GetRasterBand(1).DataType + assert result.GetRasterBand(1).GetNoDataValue() == NODATA + source = None + result = None + + +def test_burn_rejects_incompatible_rasters(small_dem_path, output_path): + """A mask that is not co-registered with the DEM is rejected.""" + mask = np.zeros((6, 6), dtype=np.uint8) + mask_path = "/vsimem/test_burn_bad_shape_mask.tif" + write_raster(mask_path, mask, 0.0, gdal.GDT_Byte) + try: + with pytest.raises(ValueError): + burn( + small_dem_path, + mask_path, + output_path, + method="statistic", + mask_values="1", + chunk_size=1, + ) + finally: + gdal.Unlink(mask_path) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"method": "nonsense", "burn_values": "1:1"}, + {"method": "statistic", "statistic": "median"}, + {"method": "statistic", "connectivity": 6}, + {"method": "constant"}, + {"method": "relative"}, + {"method": "constant", "burn_values": "1:oops"}, + {"method": "constant", "burn_values": "1:2:3"}, + {"method": "constant", "burn_values": "1:2", "mask_values": "1,3"}, + {"method": "statistic", "mask_values": "1,oops"}, + {"method": "statistic", "mask_values": "0"}, + ], + ids=[ + "bad_method", + "bad_statistic", + "bad_connectivity", + "constant_without_burn_values", + "relative_without_burn_values", + "non_numeric_burn_value", + "malformed_burn_value", + "missing_burn_value_for_class", + "non_numeric_mask_value", + "mask_value_is_nodata", + ], +) +def test_burn_rejects_bad_arguments( + small_dem_path, small_mask_path, output_path, kwargs +): + """Malformed or inconsistent arguments raise before any raster work starts.""" + with pytest.raises(ValueError): + burn(small_dem_path, small_mask_path, output_path, chunk_size=1, **kwargs) + + +def test_burn_cli(small_dem_path, small_mask_path, output_path): + """The CLI runs the statistic method end to end.""" + runner = click.testing.CliRunner() + result = runner.invoke( + burn_cli, + [ + "--dem_file", + small_dem_path, + "--mask_file", + small_mask_path, + "--output_file", + output_path, + "--method", + "statistic", + "--mask_values", + "1,3", + "--statistic", + "min", + "--burn_offset", + "1.0", + "--chunk_size", + "1", + ], + ) + assert result.exit_code == 0 + burned = read_raster(output_path) + assert np.allclose(burned[2, 3], 41.0) + + +def test_burn_cli_constant(small_dem_path, small_mask_path, output_path): + """The CLI accepts the paired burn value mapping for the constant method.""" + runner = click.testing.CliRunner() + result = runner.invoke( + burn_cli, + [ + "--dem_file", + small_dem_path, + "--mask_file", + small_mask_path, + "--output_file", + output_path, + "--method", + "constant", + "--burn_values", + "1:100,3:200", + "--chunk_size", + "2", + ], + ) + assert result.exit_code == 0 + burned = read_raster(output_path) + assert np.allclose(burned[0, 1], 100.0) + assert np.allclose(burned[2, 3], 200.0) + + +def test_union_find_joins_and_separates(): + """UnionFind reports one representative per connected group.""" + union_find = UnionFind() + union_find.union(1, 2) + union_find.union(2, 3) + union_find.union(10, 11) + assert union_find.find(1) == union_find.find(3) + assert union_find.find(10) == union_find.find(11) + assert union_find.find(1) != union_find.find(10) + # an unseen key is its own representative + assert union_find.find(99) == 99 + + +def test_union_find_merges_long_chains(): + """Chained unions collapse to a single representative.""" + union_find = UnionFind() + for i in range(1, 200): + union_find.union(i, i + 1) + roots = {union_find.find(i) for i in range(1, 201)} + assert len(roots) == 1 diff --git a/tests/test_fill_depressions.py b/tests/test_fill_depressions.py index 524df3f..c36a265 100644 --- a/tests/test_fill_depressions.py +++ b/tests/test_fill_depressions.py @@ -227,3 +227,55 @@ def test_fill_depressions_tiled(dem_filepath, expected_filled_dem_values): band = None ds = None gdal.Unlink(output_filepath) + + +@pytest.mark.parametrize("transpose", [False, True], ids=["1xN grid", "Nx1 grid"]) +def test_fill_depressions_tiled_degenerate_tile_grid(transpose): + """Seams must be joined when the tile grid is a single tile row or column. + + Regression test. Joining seams by walking a 2x2 tile block over + range(rows - 1) x range(cols - 1) iterates zero times on a 1xN or Nx1 tile + grid, which leaves every seam unjoined. Nothing raises; the fill is simply + wrong, and only for rasters that are smaller than chunk_size in one dimension + and larger in the other. + + The DEM is a plateau at 10 holding a depression at elevation 1 that straddles + the tile seam. Its only outlet is a pass at elevation 5 lying wholly inside the + second tile, so filling the first tile's half correctly is impossible without + the seam graph. Correct output raises the whole depression to 5; an unjoined + seam raises the first tile's half to 10 instead. + """ + chunk_size = 8 + dem = np.full((8, 16), 10.0, dtype=np.float32) + dem[2:6, 5:11] = 1.0 # depression straddling the seam at column 8 + dem[3, 11:16] = 5.0 # spill channel reaching the right edge, inside tile 2 + if transpose: + dem = np.ascontiguousarray(dem.T) + + dem_path = "/vsimem/degenerate_grid_dem.tif" + output_path = "/vsimem/degenerate_grid_filled.tif" + driver = gdal.GetDriverByName("GTiff") + dataset = driver.Create(dem_path, dem.shape[1], dem.shape[0], 1, gdal.GDT_Float32) + band = dataset.GetRasterBand(1) + band.WriteArray(dem) + band.SetNoDataValue(-9999.0) + dataset.FlushCache() + band = None + dataset = None + + try: + _fill_depressions_tiled(dem_path, output_path, chunk_size, "/vsimem") + ds = gdal.Open(output_path) + filled = ds.GetRasterBand(1).ReadAsArray() + ds = None + + expected = dem.copy() + if transpose: + expected[5:11, 2:6] = 5.0 + else: + expected[2:6, 5:11] = 5.0 + assert np.allclose(filled, expected) + finally: + for path in (dem_path, output_path): + if gdal.VSIStatL(path) is not None: + gdal.Unlink(path) diff --git a/tests/test_fix_flats_tiled.py b/tests/test_fix_flats_tiled.py index 3bdc0fa..2eb7397 100644 --- a/tests/test_fix_flats_tiled.py +++ b/tests/test_fix_flats_tiled.py @@ -444,6 +444,34 @@ def test_resolve_flats_tiled_no_seam_artifacts(seed, chunk_size): ) +@pytest.mark.parametrize("seed", range(8)) +@pytest.mark.parametrize("shape", [(8, 48), (48, 8)], ids=["1xN grid", "Nx1 grid"]) +def test_resolve_flats_tiled_degenerate_tile_grid(seed, shape): + """Seams must be joined when the tile grid is a single tile row or column. + + Regression test. Joining seams by walking a 2x2 tile block over + range(rows - 1) x range(cols - 1) iterates zero times on a 1xN or Nx1 tile + grid, leaving every seam unjoined so flats straddling a tile boundary are + resolved as if they ended there. The shapes here are smaller than chunk_size + in one dimension and larger in the other, which is exactly the case the block + walk misses. + """ + rng = np.random.default_rng(seed) + dem = priority_flood_fill(rng.integers(0, 4, size=shape).astype(np.float32)) + initial_fdr = make_initial_fdr(dem) + + fixed_tiled = _run_resolve_flats_tiled(dem, initial_fdr, chunk_size=8) + + assert find_reciprocal_pairs(fixed_tiled) == [] + assert_valid_drainage(fixed_tiled) + + fixed_core = fix_flats(dem, initial_fdr.copy(), inplace=False) + np.testing.assert_array_equal( + fixed_tiled == FLOW_DIRECTION_UNDEFINED, + fixed_core == FLOW_DIRECTION_UNDEFINED, + ) + + def test_seam_scale_mismatch_regression(): """Regression for a real-data bug: per-tile flat masks are on non-comparable numeric scales across a tile seam (the neighboring tile