diff --git a/benchmarks/README.md b/benchmarks/README.md index b017defb7..3c54f318f 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -4,11 +4,13 @@ This directory contains repeatable performance measurements and pass/fail large ## Organization -- `workflows/` defines deterministic inputs, operation methods, calculation engines, chunk strategies, execution modes - and result computation shared by every suite (ASV benchmark + large data tests), +- `workflows/registry.py` lists operation methods, calculation engines, chunk strategies and supported execution modes, +- `workflows/runner.py` prepares files and computes complete outputs shared by ASV and large data tests; focused + workflows such as `grouped_reference.py` and `variography.py` prepare arrays for individual comparisons, - `asv_suite/operations.py` measures operations without a dedicated scaling comparison at one fixed configuration, - `asv_suite/comparisons.py` defines one-axis comparisons and generates their valid ASV cases and classes, with fixed - Numba worker checks and the GDAL CLI kept as a separate external reference, + Numba worker checks and GDAL CLI/Flox kept as separate external references, +- `asv_suite/variography.py` measures raster and point pair sampling, reduction of prepared pairs and complete variograms, - `asv_suite/render_results.py` renders the raw measurements into method, engine, strategy and execution-mode comparisons and the two concise graphics used by the documentation, - `gdal_comparison/` contains the GDAL CLI equivalent operations for performance comparison, @@ -28,6 +30,8 @@ results/ ## Performance benchmarks +### Run benchmarks quickly while developing + To run a benchmark while developing: ```bash @@ -36,6 +40,16 @@ asv run --quick --show-stderr -E existing --bench For example, `` can be `EagerIdwNumbaGriddingRasterSize.time_operation`. +Set `GEOUTILS_ASV_PR_CHECK=1` to use the reduced parameter ranges from the pull-request checks: + +```bash +GEOUTILS_ASV_PR_CHECK=1 asv run --quick --show-stderr -E existing --bench +``` + +Omit `GEOUTILS_ASV_PR_CHECK=1` to use the complete parameter ranges. + +### Compare revisions and view results + To compare a new implementation with the `main` branch: commit current changes, then use: ```bash @@ -62,6 +76,100 @@ asv preview --browser --html-dir benchmarks/results/asv/preview Pass `--baseline-commit ` to the renderer to add `comparisons/performance-change.md`, a compact before/after table for eager, Dask and Multiprocessing end-to-end time normalized to the GDAL CLI on the same revision. +For a local preview of the documentation graphics, run: + +```bash +python -m benchmarks.asv_suite.render_results --doc-only --doc-dir benchmarks/results/documentation +``` + +### Benchmark structure and parameters + +The shared operation benchmarks use two structures. `OperationBenchmarks` measures operations without a dedicated scaling +comparison at one fixed configuration. Its `case` parameter combines the execution mode and operation, such as +`dask-reproject`. The generated classes in `comparisons.py` vary one input at a time and identify the fixed execution +mode, method and calculation engine. For example, `EagerIdwNumbaGriddingRasterSize` varies `raster_size` while keeping eager execution, +IDW and Numba fixed. + +#### Input parameters + +- `raster_size` is the number of pixels along each side of a square raster. +- `chunk_size` is the number of pixels along each side of a square chunk. +- `interpolated_points` is the number of raster locations queried by interpolation. +- `subsample_size` is the requested number of valid observations. +- `points_per_axis` defines a square point grid, so the total point count is its square. +- `groups_per_axis` defines groups along both raster axes, so the total group count is its square. + +Some comparisons also vary input layout. For grouped statistics, local groups occupy contiguous areas, while +interleaved groups are spread across chunks. The benchmarks use two value arrays with different missing cells. + +#### Implementation options + +Comparison series vary one implementation choice while keeping the others fixed: + +- `method` selects the algorithm used by an operation. For grouped statistics, `moments` computes count, mean, standard + deviation and extrema, while `robust` computes exact median and NMAD. +- `calculation_engine` selects the numerical library that performs the calculation, such as SciPy or Numba. +- `strategy` selects how a chunked operation coordinates or combines results between chunks. Grouped statistics use + `dense` to store every declared group in each chunk, `sparse` to store only groups present in a chunk, and `groupwise` + to gather complete groups for exact statistics. `auto` selects `groupwise` for exact statistics and switches from + `dense` to `sparse` above 4096 groups for mergeable statistics. +- `execution_mode` selects eager, Dask or Multiprocessing execution. + +The GDAL CLI remains a separate external reference. Quick runs validate execution and reporting rather than repeatable +performance differences. The user statistics guide explains the memory limits of the grouped-statistics strategies. + +### Grouped statistics and Flox + +The optional Flox comparison uses prepared arrays and measures finite counts, mean and population standard deviation +through complete dataframe output. Both libraries receive the same float64 values, missing observations, boolean mask +and declared categories. Dask runs with one threaded worker and 256 × 256 spatial chunks. The GeoUtils multiprocessing +series uses the same tile size and one real process. Its pool starts before input construction so workers receive only +serialized tiles. Worker recycling is disabled for this comparison: startup stays outside the measurements, while +tiling, transfers, merging and complete output remain inside. Existing file-based multiprocessing comparisons continue +to measure worker initialization separately. Flox uses its default engine and map-reduce for lazy labels. Install +`flox` separately to include these references; ASV skips them when it is absent. + +```bash +asv run --quick --show-stderr -E existing --bench 'GroupedFlox' +``` + +These comparisons vary raster size from 256 × 256 to 4096 × 4096 with 256 local groups, and vary interleaved group count +from 16 to 4225 on a fixed 1024 × 1024 raster. Input construction stays outside the measurements. The results appear +beside the existing grouped statistics plots in the generated ASV report. + +### Pair sampling and variography + +The dedicated `asv_suite/variography.py` module separates pair generation from reduction, then measures their combined +cost through the public `variogram()` function. Every case records elapsed time and peak process memory. Inputs are +prepared before measurement, while spatial search construction, pair sampling and complete output construction are +included where applicable. + +| Benchmark class | Changing input | Fixed workload | +| --- | --- | --- | +| `VariogramPairCount` | 10,000–1,000,000 pairs; Matheron or Dowd estimator | 24 distance bins | +| `VariogramLagCount` | 24–256 distance bins; Matheron or Dowd estimator | 100,000 prepared pairs | +| `RasterPairSampling` | 1,000–100,000 pairs; five sampling methods; eager or Dask | 1024 × 1024 raster | +| `RasterPairSamplingSize` | 256–4096 pixels per side; eager or Dask | 10,000 pairs; chunk anchors | +| `PointPairSamplingSize` | 1,000–100,000 source points; three search strategies | 2,000 pairs; constant point density | +| `RasterVariogramSize` | 256–4096 pixels per side; eager or Dask | 10,000 pairs; 24 Dowd estimates | + +Raster fixtures contain scattered missing cells and smoothly varying values. Dask uses one thread and 256 × 256 chunks +from prepared arrays, so these cases measure scheduling and selected value reads without disk throughput. Pair sampling +does not currently accept a multiprocessing configuration. Point searches load all coordinates and therefore use eager +inputs here. The prepared pair reduction cases isolate the numerical work from these spatial access costs. + +Install `geoutils[geostat]` to include the named estimators; ASV skips those cases if SciKit-GStat is absent. Estimator +imports and initial compilation stay outside timing. Pair sampling itself needs no geostatistics package. + +```bash +GEOUTILS_ASV_PR_CHECK=1 asv run --quick --show-stderr -E existing --bench 'asv_suite.variography' +``` + +Omit the pull-request flag to measure the full ranges. These standalone measurements appear in native ASV history, +accessible from the combined report's history link. They do not enter the GDAL comparison graphics. + +### Continuous integration and published reports + In CI, `benchmark-asv-check` verifies changed benchmarks on every pull request (using `GEOUTILS_ASV_PR_CHECK=1` to use reduced parameters). The weekly or manually triggered `benchmark-asv` workflow records measurements on new `main` commits and stores their raw history on @@ -70,17 +178,16 @@ rebuilds the latest saved history and deploys the website and documentation grap Trigger it manually only to rebuild these outputs without new measurements. The user documentation links to the latest complete graphics published there. +The options and scaling pages include every comparison, including grouped statistics. The compact GDAL documentation +graphics remain focused on operations with GDAL equivalents. + The benchmark dependencies include Pytest because ASV imports every Python module under `benchmarks/` during discovery, including the large data test, without executing its tests. -For a local preview of the documentation graphics, run: - -```bash -python -m benchmarks.asv_suite.render_results --doc-only --doc-dir benchmarks/results/documentation -``` - ## Large data tests +### Run large data tests + Normal Pytest skips these intentionally expensive checks, while pull-request CI always runs them once on Ubuntu with Python 3.12. Run the complete suite locally with: @@ -90,3 +197,8 @@ python -m pytest --large-data -m large_data -ra Select one parameter with `-k` while developing and add `--lf` to repeat only failed cases. The practical instructions and environment variables are documented at the top of `test_large_data.py`. + +Grouped statistics have Dask memory contracts for dense and sparse summaries and exact localized medians/NMAD, +each at two chunk sizes. Multiprocessing grouping currently loads source arrays in the client before worker tiling, +so it is benchmarked but does not claim the same larger-than-memory contract. Run only grouped memory checks with +`python -m pytest --large-data -m large_data -k grouped_stats -ra`. diff --git a/benchmarks/asv_suite/comparisons.py b/benchmarks/asv_suite/comparisons.py index 2b1c91558..258ba7feb 100644 --- a/benchmarks/asv_suite/comparisons.py +++ b/benchmarks/asv_suite/comparisons.py @@ -10,6 +10,10 @@ from benchmarks.asv_suite import asv_parameter_values, asv_pr_check_enabled from benchmarks.gdal_comparison.commands import ComparisonOperation from benchmarks.gdal_comparison.runner import GdalRunner +from benchmarks.workflows.grouped_reference import ( + compute_grouped_reference, + prepare_grouped_reference, +) from benchmarks.workflows.registry import ( OPERATION_METHODS, OPERATION_STRATEGIES, @@ -19,11 +23,19 @@ OperationStrategyName, ) from benchmarks.workflows.runner import BenchmarkConfig, BenchmarkRunner +from geoutils._misc import import_optional +from geoutils.multiproc import MultiprocConfig +from geoutils.multiproc.cluster import MpCluster +from geoutils.profiler import profile_call + +######################################### +# Comparison dimensions and case helpers # +######################################### # Comparisons vary one GeoUtils choice at a time: method, calculation engine, chunk strategy or execution mode # The label dictionaries give the stored values readable names in plots ComparisonDimension = Literal["method", "calculation_engine", "strategy", "execution_mode"] -ExternalReference = Literal["gdal_cli"] +ExternalReference = Literal["gdal_cli", "flox"] GDAL_CLI_LABEL = "GDAL CLI" EXECUTION_MODE_LABELS: dict[ExecutionMode, str] = { @@ -35,6 +47,7 @@ "scipy": "SciPy", "numba": "Numba", "rasterio": "Rasterio/GDAL", + "numpy": "NumPy", } METHOD_LABELS = { "nearest": "Nearest", @@ -48,6 +61,10 @@ "label_union": "Label union", "label_stitch": "Label stitch", "geometry_stitch": "Geometry stitch", + "auto": "Automatic", + "dense": "Dense summaries", + "sparse": "Sparse summaries", + "groupwise": "Complete groups", } @@ -94,12 +111,13 @@ class ExternalReferenceCase: external_reference: ExternalReference pr_check: bool = False strategy: None = None + execution_mode: ExecutionMode | None = None @property def benchmark_class(self) -> str: """Return the generated public ASV class name for this reference.""" - values = (self.external_reference, self.method, self.comparison_group) + values = (self.external_reference, self.execution_mode, self.method, self.comparison_group) return "".join(_class_token(value) for value in values if value is not None) @@ -241,6 +259,10 @@ def _external_case( return ExternalReferenceCase(comparison_group, operation, method, "gdal_cli", pr_check=pr_check) +############################## +# Registered operation cases # +############################## + # Define the cases needed to compare each operation across execution modes, calculation engines, methods or strategies # Each helper changes only that choice and keeps the other operation settings fixed _INTERPOLATION_MODES = _execution_cases("interpolation-point-count", "interp_points", "linear", "scipy") @@ -261,6 +283,76 @@ def _external_case( _RASTERIZATION_MODES = _execution_cases("rasterization-raster-size", "rasterize", None, "rasterio") _SUBSAMPLE_STRATEGIES = _strategy_cases("subsample-size", "subsample", None, None, execution_mode="dask") +# Isolate input size, chunk size, membership layout and group count for shared grouped-statistic kernels +_GROUPED_MODES = _execution_cases( + "grouped-stats-raster-size", + "grouped_stats", + "moments", + "numpy", + strategy="dense", + pr_modes=("eager", "dask", "multiprocessing"), +) +_GROUPED_STRATEGIES = { + scenario: _strategy_cases(scenario, "grouped_stats", "moments", "numpy", execution_mode="dask") + for scenario in ( + "grouped-stats-raster-size", + "grouped-stats-chunk-size", + "grouped-stats-interleaved-chunks", + "grouped-stats-group-count", + ) +} +_GROUPED_ROBUST_MODES = _execution_cases( + "grouped-stats-robust-size", + "grouped_stats", + "robust", + "numpy", + strategy="groupwise", + pr_modes=("dask", "multiprocessing"), +) + +# Compare the same prepared arrays with an optional external library, across input size and group count +_GROUPED_FLOX_MODES = { + scenario: _execution_cases( + scenario, + "grouped_stats", + "moments", + "numpy", + strategy="auto", + execution_modes=("eager", "dask", "multiprocessing"), + pr_modes=("eager", "dask", "multiprocessing"), + ) + for scenario in ("grouped-flox-raster-size", "grouped-flox-group-count") +} +_GROUPED_FLOX_REFERENCES = tuple( + ExternalReferenceCase( + scenario, + "grouped_stats", + "moments", + "flox", + pr_check=True, + execution_mode=cast(ExecutionMode, execution_mode), + ) + for scenario in _GROUPED_FLOX_MODES + for execution_mode in ("eager", "dask") +) + +# Check each distinct layout and the automatic sparse threshold with a bounded pull-request workload +for _scenario, _strategies in _GROUPED_STRATEGIES.items(): + _GROUPED_STRATEGIES[_scenario] = tuple( + ( + replace(case, pr_check=True) + if (_scenario, case.strategy) + in { + ("grouped-stats-chunk-size", "dense"), + ("grouped-stats-interleaved-chunks", "groupwise"), + ("grouped-stats-group-count", "sparse"), + ("grouped-stats-group-count", "auto"), + } + else case + ) + for case in _strategies + ) + # Compare all four gridding methods across execution modes while keeping SciPy as the calculation engine _GRID_METHODS = ("nearest", "linear", "idw", "mean") _GRID_MODE_CASES = { @@ -310,6 +402,10 @@ def _external_case( # Combine every GeoUtils case and remove duplicates when the same combination appears in several comparisons BENCHMARK_CASES = _merge_cases( + _GROUPED_MODES, + *tuple(_GROUPED_STRATEGIES.values()), + _GROUPED_ROBUST_MODES, + *tuple(_GROUPED_FLOX_MODES.values()), _INTERPOLATION_MODES, _REPROJECTION_MODES, _FILTER_MODES, @@ -346,6 +442,7 @@ def _external_case( _RASTERIZATION_REFERENCE, *_GRID_REFERENCES.values(), _GRID_POINT_REFERENCE, + *_GROUPED_FLOX_REFERENCES, ) # Map each generated ASV class name back to the operation settings needed during setup @@ -353,6 +450,11 @@ def _external_case( EXTERNAL_REFERENCE_CASE_BY_CLASS = {case.benchmark_class: case for case in EXTERNAL_REFERENCE_CASES} +############################## +# Report labels and plots # +############################## + + def _series_label(case: BenchmarkCase, dimension: ComparisonDimension) -> str: """Return the plot label for the dimension varied by one GeoUtils case.""" @@ -395,6 +497,7 @@ class Comparison: workload_template: str logarithmic_x: bool = False documentation: bool = True + summary: bool = True series_dimension: ComparisonDimension = "execution_mode" calculation_engine: CalculationEngine | None = None strategy: OperationStrategyName | None = None @@ -413,6 +516,128 @@ class Comparison: # Define the report plots, including their displayed series and the operation settings held fixed COMPARISONS: tuple[Comparison, ...] = ( + *tuple( + Comparison( + slug=scenario, + title=title, + description=( + "Compares GeoUtils stats() with Flox on two prebuilt float64 arrays, the same boolean mask and " + "declared categories. Both return finite count, mean and population standard deviation (ddof=0), " + "including completed Dask results and dataframe construction. Dask uses one threaded worker; " + "GeoUtils multiprocessing uses one persistent process initialized before the arrays. Both use " + "256 × 256 tiles. Worker startup is excluded, while tile serialization and merging are timed. " + "Flox uses its default engine and map-reduce for lazy group labels." + ), + parameter_label=parameter_label, + series=( + *_comparison_series(_GROUPED_FLOX_MODES[scenario], "execution_mode"), + *tuple( + (f"Flox ({EXECUTION_MODE_LABELS[case.execution_mode]})", case.benchmark_class) + for case in _GROUPED_FLOX_REFERENCES + if case.comparison_group == scenario and case.execution_mode is not None + ), + ), + operation="grouped_stats", + method="moments", + calculation_engine="numpy", + strategy="auto", + workload_template=workload, + documentation=False, + summary=False, + ) + for scenario, title, parameter_label, workload in ( + ( + "grouped-flox-raster-size", + "GeoUtils and Flox grouped statistics by raster size", + "Size of raster (pixels per side)", + "{parameter} × {parameter} raster; 256 local groups; two masked float64 values", + ), + ( + "grouped-flox-group-count", + "GeoUtils and Flox grouped statistics by group count", + "Number of groups per axis", + "1,024 × 1,024 raster; {parameter} × {parameter} interleaved groups; two masked float64 values", + ), + ) + ), + Comparison( + slug="grouped-stats-execution-size", + title="Grouped moments by raster size and execution mode", + description=( + "Computes count, mean, standard deviation and extrema for two values with independent missing data " + "in 64 rectangular groups. Dask reads chunks lazily; multiprocessing loads the input in the client." + ), + parameter_label="Size of raster (pixels per side)", + series=_comparison_series(_GROUPED_MODES, "execution_mode"), + operation="grouped_stats", + method="moments", + calculation_engine="numpy", + strategy="dense", + workload_template="{parameter} × {parameter} raster; 256 × 256 chunks; 64 local groups; two values", + documentation=False, + ), + *tuple( + Comparison( + slug=scenario, + title=title, + description=( + "Computes count, mean, standard deviation and extrema for two values with independent gaps. " + "Dense summaries allocate every declared group per chunk; sparse summaries retain encountered " + "groups; groupwise gathers complete observations. Automatic uses the declared group count." + ), + parameter_label=parameter_label, + series=_comparison_series(_GROUPED_STRATEGIES[scenario], "strategy"), + operation="grouped_stats", + method="moments", + calculation_engine="numpy", + execution_mode="dask", + series_dimension="strategy", + documentation=False, + workload_template=workload, + ) + for scenario, title, parameter_label, workload in ( + ( + "grouped-stats-raster-size", + "Grouped reduction strategies by raster size", + "Size of raster (pixels per side)", + "{parameter} × {parameter} raster; 256 × 256 chunks; 64 local groups; two values", + ), + ( + "grouped-stats-chunk-size", + "Grouped reduction strategies by chunk size (local groups)", + "Size of chunks (pixels per side)", + "1,024 × 1,024 raster; {parameter} × {parameter} chunks; 64 local groups; two values", + ), + ( + "grouped-stats-interleaved-chunks", + "Grouped reduction strategies by chunk size (interleaved groups)", + "Size of chunks (pixels per side)", + "1,024 × 1,024 raster; {parameter} × {parameter} chunks; 64 interleaved groups; two values", + ), + ( + "grouped-stats-group-count", + "Grouped reduction strategies by declared group count", + "Number of groups per axis", + "1,024 × 1,024 raster; 128 × 128 chunks; {parameter} × {parameter} local groups; two values", + ), + ) + ), + Comparison( + slug="grouped-stats-robust-size", + title="Exact grouped median and NMAD by raster size", + description=( + "Gathers complete observations in 64 rectangular groups to calculate exact medians and NMAD for " + "two values with independent missing data. Memory depends on the largest complete group." + ), + parameter_label="Size of raster (pixels per side)", + series=_comparison_series(_GROUPED_ROBUST_MODES, "execution_mode"), + operation="grouped_stats", + method="robust", + calculation_engine="numpy", + strategy="groupwise", + workload_template="{parameter} × {parameter} raster; 256 × 256 chunks; 64 local groups; two values", + documentation=False, + ), Comparison( slug="interpolation-point-count", title="Linear interpolation by number of points (SciPy engine)", @@ -586,6 +811,11 @@ class Comparison: ) +##################################### +# ASV measurements and input sizes # +##################################### + + # The classes below define which numeric input changes, such as raster size, chunk size or point count # Generated subclasses later combine that input axis with one concrete operation configuration class _ComparisonBenchmark: @@ -822,14 +1052,167 @@ def make_config(self, parameter: int) -> BenchmarkConfig: ) +class _GroupedStatsRasterSize(_ComparisonBenchmark): + """Vary raster size around fixed chunks and localized groups.""" + + param_names = ["raster_size"] + params = [asv_parameter_values([512, 1024, 2048], pr_check_value=256)] + + def make_config(self, parameter: int) -> BenchmarkConfig: + """Prepare two values and 64 spatial groups on the selected raster size.""" + + return BenchmarkConfig(shape=(parameter, parameter), chunks=(256, 256)) + + +class _GroupedStatsChunkSize(_ComparisonBenchmark): + """Vary chunk size while keeping the raster and group boundaries fixed.""" + + param_names = ["chunk_size"] + params = [asv_parameter_values([64, 193, 512], pr_check_value=97)] + + def make_config(self, parameter: int) -> BenchmarkConfig: + """Include uneven edge chunks and groups crossing partition boundaries.""" + + size = 256 if asv_pr_check_enabled() else 1024 + return BenchmarkConfig(shape=(size, size), chunks=(parameter, parameter)) + + +class _GroupedStatsInterleavedChunks(_GroupedStatsChunkSize): + """Repeat every group throughout the raster while varying chunk size.""" + + def make_config(self, parameter: int) -> BenchmarkConfig: + """Keep observations interleaved across all chunks for each tested partition size.""" + + return replace(super().make_config(parameter), grouped_layout="interleaved") + + +class _GroupedStatsGroupCount(_ComparisonBenchmark): + """Vary declared groups across the automatic dense-to-sparse selection threshold.""" + + param_names = ["groups_per_axis"] + params = [asv_parameter_values([4, 16, 65], pr_check_value=65)] + + def make_config(self, parameter: int) -> BenchmarkConfig: + """Include 4225 groups so automatic reduction exercises its sparse branch.""" + + size = 256 if asv_pr_check_enabled() else 1024 + return BenchmarkConfig(shape=(size, size), chunks=(128, 128), grouped_regions_per_axis=parameter) + + +class _GroupedFloxRasterSize(_ComparisonBenchmark): + """Compare complete grouped results after preparing identical in-memory NumPy or Dask inputs.""" + + param_names = ["raster_size"] + params = [asv_parameter_values([256, 1024, 4096], pr_check_value=256)] + + def make_config(self, parameter: int) -> BenchmarkConfig: + """Vary raster size around 256 local groups and fixed spatial chunks.""" + + return BenchmarkConfig(shape=(parameter, parameter), chunks=(256, 256), grouped_regions_per_axis=16) + + def setup(self, parameter: int) -> None: + """Prepare common arrays and load optional libraries outside the measurement.""" + + # Select the same execution mode for GeoUtils and its corresponding Flox reference + benchmark_class = type(self).__name__ + case = BENCHMARK_CASE_BY_CLASS.get(benchmark_class) + reference = EXTERNAL_REFERENCE_CASE_BY_CLASS.get(benchmark_class) + selected = case or reference + assert selected is not None and selected.execution_mode is not None + self.implementation: Literal["geoutils", "flox"] = "geoutils" if reference is None else "flox" + if reference is not None: + try: + import_optional("flox", extra_name="benchmark") + except ImportError as exc: + raise NotImplementedError("Install optional flox to run this comparison") from exc + + # Fix the scheduler for both libraries and construct all observations before timing starts + self.dask = import_optional("dask", extra_name="benchmark") + config = self.make_config(parameter) + + # Start one persistent worker before building arrays so it receives only serialized tiles + self.mp_cluster: MpCluster | None = None + self.mp_config: MultiprocConfig | None = None + if selected.execution_mode == "multiprocessing": + self.mp_cluster = MpCluster({"nb_workers": 1, "max_tasks_per_child": None}) + self.mp_config = MultiprocConfig(chunks=config.chunks, cluster=self.mp_cluster) + + # Keep complete prepared inputs in the client, with identical logical tiles for both worker backends + self.inputs = prepare_grouped_reference( + config.shape[0], + config.grouped_regions_per_axis, + config.grouped_layout, + selected.execution_mode, + ) + + def teardown(self, parameter: int) -> None: + """Stop worker processes and release arrays after each independent ASV measurement.""" + + cluster = getattr(self, "mp_cluster", None) + if cluster is not None: + cluster.close() + if hasattr(self, "inputs"): + del self.inputs + + def time_operation(self, parameter: int) -> None: + """Compute every requested result with one threaded or multiprocessing worker.""" + + with self.dask.config.set(scheduler="threads", num_workers=1): + compute_grouped_reference(*self.inputs, implementation=self.implementation, mp_config=self.mp_config) + + def track_end_to_end_time_s(self, parameter: int) -> float: + """Measure masking, grouping and complete output construction from prepared inputs.""" + + start = time.perf_counter() + self.time_operation(parameter) + return time.perf_counter() - start + + def track_peak_process_tree_mem_mb(self, parameter: int) -> float: + """Measure peak process memory while the complete grouped result is calculated.""" + + _, metrics = profile_call(self.time_operation, parameter, dask=False, include_children=True) + assert metrics.peak_process_tree_mem_mb is not None + return metrics.peak_process_tree_mem_mb + + +class _GroupedFloxGroupCount(_GroupedFloxRasterSize): + """Compare grouped reductions as interleaved group count crosses the sparse threshold.""" + + param_names = ["groups_per_axis"] + params = [asv_parameter_values([4, 16, 65], pr_check_value=65)] + + def make_config(self, parameter: int) -> BenchmarkConfig: + """Vary declared groups on a fixed raster with membership repeated across chunks.""" + + size = 256 if asv_pr_check_enabled() else 1024 + return BenchmarkConfig( + shape=(size, size), chunks=(256, 256), grouped_regions_per_axis=parameter, grouped_layout="interleaved" + ) + + +setattr(_GroupedFloxRasterSize.track_end_to_end_time_s, "unit", "seconds") +setattr(_GroupedFloxRasterSize.track_peak_process_tree_mem_mb, "unit", "MB") + + class _NumbaWorkerIntegration(_GriddingRasterSize): """Exercise each Numba kernel once in Dask and multiprocessing workers.""" params = [asv_parameter_values([1024], pr_check_value=512)] +##################################### +# Public ASV class registration # +##################################### + # Select the input axis and fixture configuration used by each named comparison group _SCENARIO_BASES: dict[str, type[_ComparisonBenchmark]] = { + "grouped-flox-raster-size": _GroupedFloxRasterSize, + "grouped-flox-group-count": _GroupedFloxGroupCount, + "grouped-stats-raster-size": _GroupedStatsRasterSize, + "grouped-stats-chunk-size": _GroupedStatsChunkSize, + "grouped-stats-interleaved-chunks": _GroupedStatsInterleavedChunks, + "grouped-stats-group-count": _GroupedStatsGroupCount, + "grouped-stats-robust-size": _GroupedStatsRasterSize, "interpolation-point-count": _InterpolationPointCount, "reprojection-raster-size": _ReprojectionRasterSize, "filter-chunk-size": _FilterChunkSize, diff --git a/benchmarks/asv_suite/global_stats.py b/benchmarks/asv_suite/global_stats.py new file mode 100644 index 000000000..d85b809e4 --- /dev/null +++ b/benchmarks/asv_suite/global_stats.py @@ -0,0 +1,54 @@ +"""Compare GeoUtils global Dask reduction with equivalent native Dask statistics.""" + +from __future__ import annotations + +from typing import Literal + +import numpy as np + +from benchmarks.asv_suite import asv_parameter_values +from geoutils._misc import import_optional +from geoutils.stats.reduction import ( + _normalize_statistics, + _reduce_values, + _statistics_dask, +) + + +class GlobalDaskReduction: + """Measure the shared GeoUtils reducer and native Dask reductions on the same prepared array.""" + + number = 1 + repeat = 3 + rounds = 1 + warmup_time = 0 + timeout = 300 + param_names = ["implementation", "raster_size"] + params = [["geoutils", "native_dask"], asv_parameter_values([256, 1024, 4096], 256)] + + def setup(self, implementation: Literal["geoutils", "native_dask"], raster_size: int) -> None: + """Prepare one finite Dask raster and the same mergeable statistic request for both implementations.""" + + import_optional("dask", extra_name="benchmark") + import dask.array as da + + del implementation + generator = np.random.default_rng(42) + values = generator.normal(size=(raster_size, raster_size)) + values[::97, ::89] = np.nan + self.values = da.from_array(values, chunks=(256, 256)) + self.aliases = {"mean", "min", "max", "sum", "sumofsquares", "rmse", "std"} + self.statistics = _normalize_statistics(sorted(self.aliases), grouped=False) + + def time_reduction(self, implementation: Literal["geoutils", "native_dask"], raster_size: int) -> None: + """Compute the requested global estimates with the selected Dask reduction implementation.""" + + import dask + + del raster_size + with dask.config.set(scheduler="threads", num_workers=1): + if implementation == "geoutils": + _reduce_values([self.values], self.statistics) + else: + estimates, count = _statistics_dask(self.values, self.aliases) + dask.compute(estimates, count) diff --git a/benchmarks/asv_suite/operations.py b/benchmarks/asv_suite/operations.py index 15386a039..4cddf38ec 100644 --- a/benchmarks/asv_suite/operations.py +++ b/benchmarks/asv_suite/operations.py @@ -10,7 +10,7 @@ from benchmarks.workflows.runner import BenchmarkConfig, BenchmarkRunner # Scaling comparisons already cover these operations at three input values -_SCALING_OPERATIONS = {"filter", "reproject", "interp_points", "polygonize", "rasterize", "grid"} +_SCALING_OPERATIONS = {"filter", "reproject", "interp_points", "polygonize", "rasterize", "grid", "grouped_stats"} # Remove operations already measured while varying raster, chunk or point count # The remaining operation/execution-mode pairs are measured once with a fixed input configuration diff --git a/benchmarks/asv_suite/render_results.py b/benchmarks/asv_suite/render_results.py index 2421d3a05..fcd0f0a2a 100644 --- a/benchmarks/asv_suite/render_results.py +++ b/benchmarks/asv_suite/render_results.py @@ -96,6 +96,10 @@ "Scaling with the number of sampled values", "The returned sample size varies while the source raster remains fixed.", ), + "Number of groups per axis": ( + "Scaling with the number of groups", + "The number of rectangular groups varies while raster and chunk dimensions remain fixed.", + ), } OPERATION_LABELS: dict[OperationName, str] = { @@ -105,6 +109,7 @@ "filter": "Filtering", "reproject": "Reprojection", "statistics": "Statistics", + "grouped_stats": "Grouped statistics", "subsample": "Subsampling", "interp_points": "Point interpolation", "polygonize": "Polygonization", @@ -137,6 +142,7 @@ "rasterize": "Vector ⟶ Raster", "create_mask": "Vector ⟶ Raster", "statistics": "Raster ⟶ Other", + "grouped_stats": "Raster ⟶ Other", "write": "Raster ⟶ Other", } @@ -183,9 +189,13 @@ def __init__(self, commit_hash: str = "preview-current", geoutils_scale: float = parameter_values = (3, 9, 33) elif comparison.parameter_label == "Number of sampled values": parameter_values = (256, 2048, 16384) + elif comparison.parameter_label == "Number of groups per axis": + parameter_values = (4, 16, 65) elif comparison.parameter_label == "Size of chunks (pixels per side)": - parameter_values = (256, 512, 1024) - elif comparison.operation == "grid": + parameter_values = (64, 193, 512) if comparison.operation == "grouped_stats" else (256, 512, 1024) + elif comparison.slug == "grouped-flox-raster-size": + parameter_values = (256, 1024, 4096) + elif comparison.operation in {"grid", "grouped_stats"}: parameter_values = (512, 1024, 2048) else: parameter_values = (1024, 2048, 4096) @@ -400,7 +410,7 @@ def collect_comparison_measurements( method=selected_case.method, calculation_engine=(benchmark_case.calculation_engine if benchmark_case is not None else None), strategy=benchmark_case.strategy if benchmark_case is not None else None, - execution_mode=benchmark_case.execution_mode if benchmark_case is not None else None, + execution_mode=selected_case.execution_mode, external_reference=(reference_case.external_reference if reference_case is not None else None), series_dimension=comparison.series_dimension, parameter=parameter, @@ -488,14 +498,17 @@ def _largest_shared_parameter( # Shared parameters keep the normalized time bars based on exactly the same input size parameters_by_series = [] - for series_label, _ in comparison.series: - parameters_by_series.append( - { - record.parameter - for record in records - if record.comparison == comparison.slug and record.series_label == series_label - } - ) + for series_label, class_name in comparison.series: + parameters = { + record.parameter + for record in records + if record.comparison == comparison.slug and record.series_label == series_label + } + # Optional Flox references may be skipped while all GeoUtils measurements remain available + reference = EXTERNAL_REFERENCE_CASE_BY_CLASS.get(class_name) + if not parameters and reference is not None and reference.external_reference == "flox": + continue + parameters_by_series.append(parameters) shared_parameters = set.intersection(*parameters_by_series) if not shared_parameters: raise ValueError(f"No shared parameter found for {comparison.slug}") @@ -947,7 +960,7 @@ def _execution_mode_measurements( ) -> tuple[int, dict[ExecutionMode, ComparisonMeasurement]]: """Return GeoUtils execution modes at the largest workload shared by every mode.""" - selected_labels = [label for label, _ in comparison.series if label != GDAL_CLI_LABEL] + selected_labels = [label for label, class_name in comparison.series if class_name in BENCHMARK_CASE_BY_CLASS] if parameter is None: parameter = _largest_parameter_for_series(comparison, records, selected_labels) measurements = ( @@ -971,6 +984,8 @@ def _engine_summary_comparisons() -> tuple[Comparison, ...]: # Prefer a direct engine comparison, then fall back to an execution comparison containing Eager selected: dict[tuple[OperationName, str | None], tuple[int, Comparison]] = {} for comparison in COMPARISONS: + if not comparison.summary: + continue priority = 0 if comparison.series_dimension == "calculation_engine": priority = 2 if comparison.parameter_label == "Size of raster (pixels per side)" else 1 @@ -991,7 +1006,7 @@ def _engine_summary_table(records: list[ComparisonMeasurement]) -> str: """Compare eager GeoUtils calculation engines with the external GDAL CLI.""" rows: list[tuple[Comparison, str]] = [] - columns: tuple[CalculationEngine | ExternalReference, ...] = ("rasterio", "scipy", "numba", "gdal_cli") + columns: tuple[CalculationEngine | ExternalReference, ...] = ("rasterio", "scipy", "numba", "numpy", "gdal_cli") for comparison in _engine_summary_comparisons(): parameter = _summary_reference_parameter(comparison, records) _, by_implementation = _eager_implementation_measurements(comparison, records, parameter) @@ -1023,11 +1038,12 @@ def _engine_summary_table(records: list[ComparisonMeasurement]) -> str: '
', '', '', - '', + '', '', '', + '', '', - _group_operation_rows(rows, column_count=6), + _group_operation_rows(rows, column_count=7), "
Operation and methodReference workloadGeoUtils calculation engineGeoUtils calculation engineExternal reference
Rasterio/GDALSciPyNumbaNumPyGDAL CLI
", ] ) @@ -1036,7 +1052,11 @@ def _engine_summary_table(records: list[ComparisonMeasurement]) -> str: def _execution_summary_comparisons() -> tuple[Comparison, ...]: """Return every plot that directly compares GeoUtils execution modes.""" - return tuple(comparison for comparison in COMPARISONS if comparison.series_dimension == "execution_mode") + return tuple( + comparison + for comparison in COMPARISONS + if comparison.summary and comparison.series_dimension == "execution_mode" + ) def _execution_summary_table(records: list[ComparisonMeasurement]) -> str: @@ -1108,7 +1128,7 @@ def _summary_reference_parameter(comparison: Comparison, records: list[Compariso def _headline_summary_table(records: list[ComparisonMeasurement]) -> str: """List engines and execution modes in separate columns beside the external GDAL CLI.""" - engine_order: tuple[CalculationEngine, ...] = ("rasterio", "scipy", "numba") + engine_order: tuple[CalculationEngine, ...] = ("rasterio", "scipy", "numba", "numpy") mode_order: tuple[ExecutionMode, ...] = ("eager", "dask", "multiprocessing") engine_comparisons = { (comparison.operation, comparison.method): comparison for comparison in _engine_summary_comparisons() @@ -1186,13 +1206,14 @@ def _headline_summary_table(records: list[ComparisonMeasurement]) -> str: '
', '', '', - '', + '', '', '', '', + '', '', '', - _group_operation_rows(rows, column_count=9), + _group_operation_rows(rows, column_count=10), "
Operation and methodReference workloadGeoUtils calculation engineGeoUtils calculation engineGeoUtils execution modeExternal reference
Rasterio/GDALSciPyNumbaNumPyEagerDaskMultiprocessingGDAL CLI
", ] ) @@ -1513,7 +1534,7 @@ def _concept_map() -> str: 'Reprojection — Bilinear', '
Calculation engineThe library GeoUtils uses to ' 'carry out the calculation.
Rasterio / GDAL' - 'SciPyNumba
', + 'SciPyNumbaNumPy', '
Execution modeHow an operation runs: ' "in-memory or chunked out-of-memory." '
EagerDask' @@ -1521,7 +1542,8 @@ def _concept_map() -> str: '
Chunk strategyHow a chunked operation reconciles ' 'separate partial results.
Subsampling — Sequential' 'Subsampling — Top-k' - 'Polygonization — Label stitch
', + 'Polygonization — Label stitch' + 'Grouped statistics — Dense / Sparse / Complete groups
', "", '
External reference: GDAL CLI', "Standalone GDAL file-to-file command run outside GeoUtils. It is distinct from Rasterio/GDAL, " diff --git a/benchmarks/asv_suite/variography.py b/benchmarks/asv_suite/variography.py new file mode 100644 index 000000000..d233db5d2 --- /dev/null +++ b/benchmarks/asv_suite/variography.py @@ -0,0 +1,206 @@ +"""Measure pair sampling and variogram reduction separately, then through the complete public workflow.""" + +from __future__ import annotations + +from typing import Any, Literal + +from benchmarks.asv_suite import asv_parameter_values, asv_pr_check_enabled +from benchmarks.workflows.variography import ( + prepare_pair_pointcloud, + prepare_pair_raster, + prepare_variogram_pairs, +) +from geoutils._misc import import_optional +from geoutils.profiler import profile_call +from geoutils.stats import Variogram, variogram + +############################ +# Shared measurement setup # +############################ + + +class _VariographyBenchmark: + """Measure complete results after constructing inputs and importing optional estimators.""" + + number = 1 + repeat = 3 + rounds = 1 + warmup_time = 0 + timeout = 300 + + def time_operation(self, *parameters: Any) -> None: + """Compute the complete pair dataset or reduced variogram from prepared inputs.""" + + self._execute() + + def track_peak_process_tree_mem_mb(self, *parameters: Any) -> float: + """Measure process memory including prepared inputs and the completed result.""" + + _, metrics = profile_call(self._execute, dask=False, include_children=True) + assert metrics.peak_process_tree_mem_mb is not None + return metrics.peak_process_tree_mem_mb + + def _execute(self) -> Any: + """Compute the operation supplied by each concrete benchmark.""" + + raise NotImplementedError + + +setattr(_VariographyBenchmark.track_peak_process_tree_mem_mb, "unit", "MB") + + +def _prepare_estimator(estimator: str) -> None: + """Load and warm the optional estimator before timing its repeated application to distance bins.""" + + # ASV records an unavailable optional package as a skipped case, and compilation stays outside timing + try: + import_optional("skgstat", package_name="scikit-gstat", extra_name="geostat") + except ImportError as exc: + raise NotImplementedError("Install geoutils[geostat] to measure variogram estimators") from exc + Variogram.from_pairs(prepare_variogram_pairs(64), estimator=estimator, n_lags=4) + + +######################################### +# Reduction of already sampled pairs # +######################################### + + +class VariogramPairCount(_VariographyBenchmark): + """Vary pair count for the mean-square and robust median estimators at 24 fixed distance bins.""" + + param_names = ["n_pairs", "estimator"] + params = [asv_parameter_values([10_000, 100_000, 1_000_000], 1_000), ["matheron", "dowd"]] + + def setup(self, n_pairs: int, estimator: str) -> None: + """Prepare the same finite pairs for both estimators outside the measured call.""" + + _prepare_estimator(estimator) + self.pairs = prepare_variogram_pairs(n_pairs) + self.estimator = estimator + self.n_lags = 24 + + def _execute(self) -> Variogram: + """Reduce all prepared pairs to their distance-bin estimates and counts.""" + + return Variogram.from_pairs(self.pairs, estimator=self.estimator, n_lags=self.n_lags) + + +class VariogramLagCount(VariogramPairCount): + """Vary distance-bin count at 100,000 pairs to expose repeated full-input scans.""" + + param_names = ["n_lags", "estimator"] + params = [asv_parameter_values([24, 100, 256], 24), ["matheron", "dowd"]] + + def setup(self, n_lags: int, estimator: str) -> None: + """Keep the pair sample fixed while changing only its number of distance bins.""" + + super().setup(1_000 if asv_pr_check_enabled() else 100_000, estimator) + self.n_lags = n_lags + + +######################################### +# Spatial pair sampling # +######################################### + + +class RasterPairSampling(_VariographyBenchmark): + """Compare all regular-grid sampling methods on prepared eager and Dask rasters.""" + + param_names = ["n_pairs", "execution_mode", "sampling_method"] + sampling_methods = ["independent", "anchors", "chunk_anchors", "anchor_batched", "random_xy"] + params = [ + asv_parameter_values([1_000, 10_000, 100_000], 1_000), + ["eager", "dask"], + ["chunk_anchors", "random_xy"] if asv_pr_check_enabled() else sampling_methods, + ] + + def setup(self, n_pairs: int, execution_mode: Literal["eager", "dask"], sampling_method: str) -> None: + """Prepare one raster and fix batching, distance limits and the random seed for every method.""" + + size = 256 if asv_pr_check_enabled() else 1024 + self._prepare(size, n_pairs, execution_mode, sampling_method) + + def _prepare(self, size: int, n_pairs: int, execution_mode: Literal["eager", "dask"], sampling_method: str) -> None: + """Share the source and sampling settings between pair-count and raster-size comparisons.""" + + # Keep the same source values and worker count for every sampling method + self.source = prepare_pair_raster(size, execution_mode) + self.dask = import_optional("dask", extra_name="benchmark") + + # Bound candidate batches and reuse the same map-distance range and random seed + self.pair_kwargs = { + "n_pairs": n_pairs, + "sampling": "random_xy" if sampling_method == "random_xy" else "loglag", + "strategy": "chunk_anchors" if sampling_method == "random_xy" else sampling_method, + "min_distance": 1, + "max_distance": size / 2, + "batch_pairs": 100_000, + "anchors_per_round": 2_000, + "random_state": 42, + } + + def _execute(self) -> Any: + """Draw finite pairs and construct all endpoint values and coordinates with one Dask thread.""" + + with self.dask.config.set(scheduler="threads", num_workers=1): + return self.source.pairsample(**self.pair_kwargs) + + +class RasterPairSamplingSize(RasterPairSampling): + """Vary source raster size around a fixed pair count and the default chunk-anchor strategy.""" + + param_names = ["raster_size", "execution_mode"] + params = [asv_parameter_values([256, 1024, 4096], 256), ["eager", "dask"]] + + def setup(self, raster_size: int, execution_mode: Literal["eager", "dask"]) -> None: + """Keep 10,000 requested pairs while increasing the number of source chunks.""" + + n_pairs = 1_000 if asv_pr_check_enabled() else 10_000 + self._prepare(raster_size, n_pairs, execution_mode, "chunk_anchors") + + +class PointPairSamplingSize(_VariographyBenchmark): + """Compare exact ring searches and nearest-vector sampling as the irregular point set grows.""" + + param_names = ["n_points", "strategy"] + params = [asv_parameter_values([1_000, 10_000, 100_000], 1_000), ["kdtree", "hashgrid", "nn_logvector"]] + + def setup(self, n_points: int, strategy: str) -> None: + """Prepare a constant-density point cloud and enough nearby candidates for each search method.""" + + self.source = prepare_pair_pointcloud(n_points) + self.pair_kwargs = { + "n_pairs": 200 if asv_pr_check_enabled() else 2_000, + "strategy": strategy, + "min_distance": 1, + "max_distance": n_points**0.5 / 2, + "anchors_per_round": 2_000, + "nn_tolerance": 0.5, + "random_state": 42, + } + + def _execute(self) -> Any: + """Build the spatial search, sample finite pairs and construct their complete labelled dataset.""" + + return self.source.pairsample(**self.pair_kwargs) + + +######################################### +# Complete variogram workflow # +######################################### + + +class RasterVariogramSize(RasterPairSamplingSize): + """Measure public variogram() from pair generation through 24 robust distance-bin estimates.""" + + def setup(self, raster_size: int, execution_mode: Literal["eager", "dask"]) -> None: + """Prepare the raster and estimator while leaving sampling and reduction inside timing.""" + + super().setup(raster_size, execution_mode) + _prepare_estimator("dowd") + + def _execute(self) -> Variogram: + """Sample and reduce pairs through the public stats module with one Dask thread.""" + + with self.dask.config.set(scheduler="threads", num_workers=1): + return variogram(self.source, estimator="dowd", n_lags=24, **self.pair_kwargs) diff --git a/benchmarks/test_large_data.py b/benchmarks/test_large_data.py index 8b2dba0f9..0b499fd85 100644 --- a/benchmarks/test_large_data.py +++ b/benchmarks/test_large_data.py @@ -36,6 +36,7 @@ from benchmarks.workflows.registry import ( OPERATION_BENCHMARK_CASES, OPERATION_BY_NAME, + OperationStrategyName, split_operation_case, ) from benchmarks.workflows.runner import ( @@ -254,3 +255,24 @@ def test_neighborhood_gridding_stays_out_of_core( # IDW has its own reduction, while mean represents the shared circular-statistic neighborhood path self._check_case(case_name=case_name, large_data_config=config) + + @pytest.mark.parametrize("strategy", ["dense", "sparse", "groupwise"]) + @pytest.mark.parametrize("chunk_scale", [1, 2]) + def test_grouped_stats_stays_out_of_core( + self, strategy: OperationStrategyName, chunk_scale: int, large_data_config: BenchmarkConfig + ) -> None: + """Checks that grouped summaries and exact local medians finish below full-raster worker memory.""" + + # Keep regions fixed when changing chunks so group boundaries cross at least one tested partition layout + # Sixty-four regions per axis bound each complete group needed for exact median and NMAD + chunks = tuple(max(16, size // chunk_scale) for size in large_data_config.chunks) + config = replace( + large_data_config, + chunks=chunks, + grouped_regions_per_axis=64, + operation_strategy=strategy, + operation_method="robust" if strategy == "groupwise" else "moments", + ) + + # Reuse the isolated-process contract for finite counts, worker health and measured memory growth + self._check_case(case_name="dask-grouped_stats", large_data_config=config) diff --git a/benchmarks/workflows/grouped_reference.py b/benchmarks/workflows/grouped_reference.py new file mode 100644 index 000000000..a31a79442 --- /dev/null +++ b/benchmarks/workflows/grouped_reference.py @@ -0,0 +1,89 @@ +"""Prepare equivalent grouped statistics workloads for GeoUtils and optional Flox comparisons.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any, Literal + +import numpy as np +import pandas as pd + +from benchmarks.workflows.registry import ExecutionMode +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.multiproc import MultiprocConfig + + +def prepare_grouped_reference( + size: int, + groups_per_axis: int, + layout: Literal["local", "interleaved"], + execution_mode: ExecutionMode, +) -> tuple[Any, Any, Any, NDArrayNum]: + """Create two float64 value arrays with distinct gaps, declared integer groups and a shared mask. + + Local groups occupy rectangles; interleaved groups repeat throughout the raster. Values have shape + (2, size, size), groups and mask have shape (size, size), and Dask uses 256 by 256 spatial chunks. + Prepare these inputs before timing so both libraries measure only masking, reduction and result construction. + """ + + # 1/ Prepare two bounded signals with different missing observations + positions = np.arange(size * size).reshape(size, size) + first = (positions % 997).astype(np.float64) * 0.125 + values = np.stack((first, first * 2 + 10)) + values[0, positions % 17 == 0] = np.nan + values[1, positions % 29 == 0] = np.nan + + # 2/ Keep group membership and the common selection independent of missing values + rows, columns = np.arange(size)[:, None], np.arange(size)[None, :] + if layout == "local": + groups = (rows * groups_per_axis // size) * groups_per_axis + columns * groups_per_axis // size + else: + groups = (rows % groups_per_axis) * groups_per_axis + columns % groups_per_axis + groups = groups.astype(np.int32) + mask = positions % 13 != 0 + categories = np.arange(groups_per_axis**2) + + # 3/ Give both implementations the same lazy arrays and spatial partitions + if execution_mode == "dask": + import_optional("dask", extra_name="benchmark") + import dask.array as da + + values = da.from_array(values, chunks=(1, 256, 256)) + groups = da.from_array(groups, chunks=(256, 256)) + mask = da.from_array(mask, chunks=(256, 256)) + return values, groups, mask, categories + + +def compute_grouped_reference( + values: Any, + groups: Any, + mask: Any, + categories: NDArrayNum, + *, + implementation: Literal["geoutils", "flox"], + statistics: Sequence[str] = ("mean", "std"), + mp_config: MultiprocConfig | None = None, +) -> pd.DataFrame: + """Compute the same finite counts and selected statistics with the GeoUtils or Flox backend. + + Both backends return all declared groups, ignore missing and masked observations, and use population standard + deviation (ddof=0). Flox uses its default eager engine and map-reduce for lazy labels. All lazy results are + computed together, and both paths include construction of the same typed dataframe in the measured call. + An optional multiprocessing configuration sends GeoUtils tiles to an already initialized worker pool; tile + serialization, dispatch and result merging stay inside this call. + """ + + # Measure the selected reduction backend through the same complete GeoUtils public API + from geoutils.stats import stats + + return stats( + {"first": values[0], "second": values[1]}, + statistics, + by={"zone": groups}, + categories={"zone": categories}, + mask=mask, + observed=False, + backend=implementation, + mp_config=mp_config, + ) diff --git a/benchmarks/workflows/registry.py b/benchmarks/workflows/registry.py index 0b300b746..85315f7ee 100644 --- a/benchmarks/workflows/registry.py +++ b/benchmarks/workflows/registry.py @@ -25,8 +25,10 @@ # Define the dimensions supported by GeoUtils and tested in benchmarks: execution modes, # calculation engines, operation names, operation methods and chunk strategies ExecutionMode = Literal["eager", "dask", "multiprocessing"] -CalculationEngine = Literal["scipy", "numba", "rasterio"] -OperationStrategyName = Literal["sequential", "topk", "label_union", "label_stitch", "geometry_stitch"] +CalculationEngine = Literal["scipy", "numba", "rasterio", "numpy"] +OperationStrategyName = Literal[ + "sequential", "topk", "label_union", "label_stitch", "geometry_stitch", "auto", "dense", "sparse", "groupwise" +] OperationName = Literal[ "crop", "translate", @@ -34,6 +36,7 @@ "filter", "reproject", "statistics", + "grouped_stats", "subsample", "interp_points", "polygonize", @@ -77,6 +80,8 @@ class OperationCase: # List the supported method and calculation-engine combinations for each numerical operation # Single-method operations stay explicit so the engine is always recorded in benchmark results OPERATION_METHODS: tuple[OperationMethod, ...] = ( + OperationMethod("grouped_stats", "moments", ("numpy",), default=True), + OperationMethod("grouped_stats", "robust", ("numpy",)), OperationMethod("interp_points", "linear", ("scipy",), default=True), OperationMethod("reproject", "nearest", ("rasterio",), default=True), OperationMethod("filter", "mean", ("scipy",), default=True), @@ -90,6 +95,10 @@ class OperationCase: # List the alternative ways chunked operations select or reconcile results; eager execution has no strategy OPERATION_STRATEGIES: tuple[OperationStrategy, ...] = ( + OperationStrategy("grouped_stats", "auto", default=True), + OperationStrategy("grouped_stats", "dense"), + OperationStrategy("grouped_stats", "sparse"), + OperationStrategy("grouped_stats", "groupwise"), OperationStrategy("subsample", "sequential"), OperationStrategy("subsample", "topk", default=True), OperationStrategy("polygonize", "label_union"), @@ -107,6 +116,8 @@ class OperationCase: OperationCase("filter", ("dask", "multiprocessing"), 1), OperationCase("reproject", ("dask", "multiprocessing"), 1), OperationCase("statistics", ("dask",), 1), + # Multiprocessing currently tiles arrays already resident in the client, so only Dask is out of core + OperationCase("grouped_stats", ("dask",), 1), OperationCase("subsample", ("dask", "multiprocessing"), 1), OperationCase("interp_points", ("dask", "multiprocessing"), 1), OperationCase("polygonize", ("dask", "multiprocessing"), 1), diff --git a/benchmarks/workflows/runner.py b/benchmarks/workflows/runner.py index bb76256ed..ea7b53626 100644 --- a/benchmarks/workflows/runner.py +++ b/benchmarks/workflows/runner.py @@ -45,6 +45,10 @@ from geoutils.interface.gridding import GriddingEngine, GriddingMethod from geoutils.profiler import ProfileMetrics, profile_call +################################### +# Configuration and measurements # +################################### + # Keep input sizes, worker settings and measured results consistent across ASV, GDAL and large-data tests @dataclass @@ -65,6 +69,8 @@ class BenchmarkConfig: polygon_regions_per_axis: int = 1 vector_features_per_axis: int = 1 point_features_per_axis: int = 5 + grouped_regions_per_axis: int = 8 + grouped_layout: Literal["local", "interleaved"] = "local" operation_method: str | None = None calculation_engine: CalculationEngine | None = None operation_strategy: OperationStrategyName | None = None @@ -116,6 +122,11 @@ def worker_restarted(self) -> bool: return self.worker_pids_before != self.worker_pids_after +############################## +# Size and output helpers # +############################## + + # Calculate memory limits and read one output pixel without loading a complete raster def logical_raster_size_mb(config: BenchmarkConfig) -> float: """Return the uncompressed float32 raster size in decimal megabytes.""" @@ -156,6 +167,11 @@ def read_raster_center(filename: str) -> float: return float(dataset.read(1, window=rio.windows.Window(col, row, 1, 1))[0, 0]) +############################## +# Deterministic source files # +############################## + + # Write deterministic test rasters, polygons and points without allocating the complete raster in memory def _write_constant_raster(filename: str, config: BenchmarkConfig) -> None: """Write a deterministic constant raster one storage block at a time.""" @@ -287,6 +303,11 @@ def _write_point_source(filename: str, points_per_axis: int = 5) -> None: points.to_file(filename, driver="GPKG") +############################################ +# Worker lifecycle and complete operations # +############################################ + + # Prepare the shared inputs, start the selected execution mode and force each operation to produce a complete output class BenchmarkRunner: """Prepare deterministic files and execute one GeoUtils implementation.""" @@ -659,6 +680,66 @@ def _interpolation_points(self) -> tuple[np.ndarray[Any, Any], np.ndarray[Any, A y = rng.uniform(45.01, 45.99, size=self.config.ninterp) return x, y + def _grouped_statistics(self, raster: Any, method: str, strategy: str | None) -> float: + """Compute grouped moments or exact robust estimates on deterministic values with independent gaps. + + Local groups occupy rectangular regions; interleaved groups span the entire input. Dask builds all value + and membership arrays lazily. Multiprocessing benchmarks the current array interface, which loads values + in the client before tiling them for workers. The returned fingerprint checks complete finite counts. + """ + + # Generate coordinates with the same execution backend as the input raster + height, width = self.config.shape + if self.backend == "dask": + import_optional("dask", extra_name="benchmark") + import dask.array as da + + rows = da.arange(height, chunks=self.config.chunks[0])[:, None] + columns = da.arange(width, chunks=self.config.chunks[1])[None, :] + else: + rows = np.arange(height)[:, None] + columns = np.arange(width)[None, :] + regions = self.config.grouped_regions_per_axis + if regions < 1 or regions > min(height, width): + raise ValueError("Grouped regions per axis must fit within the raster dimensions.") + + # Separate localized membership from groups repeated through every chunk + if self.config.grouped_layout == "local": + groups = (rows * regions // height) * regions + columns * regions // width + else: + groups = (rows % regions) * regions + columns % regions + positions = rows * width + columns + base = raster.data.squeeze() + signal = base + (rows % 97) * 0.125 + (columns % 53) * 0.25 + values = { + "signal": np.where(positions % 17 != 0, signal, np.nan), + "offset": np.where(positions % 29 != 0, 2 * signal + 10, np.nan), + } + + # Measure the complete public calculation, including exact group gathering when requested + from geoutils.stats import stats + + statistics = ["mean", "std", "min", "max"] if method == "moments" else ["median", "nmad"] + config = self._multiproc_config("grouped_stats") if self.backend == "multiprocessing" else None + result = stats( + values, + by={"zone": groups}, + categories={"zone": range(regions**2)}, + statistics=statistics, + strategy=cast(Literal["auto", "dense", "sparse", "groupwise"], strategy or "auto"), + mp_config=config, + ) + + # Every pixel belongs to a group; missing values follow independent, analytically known periods + count = height * width + for name, period in (("signal", 17), ("offset", 29)): + expected = count - (count + period - 1) // period + if result[(name, "count")].sum() != expected: + raise AssertionError(f"Grouped benchmark lost finite observations in {name!r}.") + if not np.isfinite(result[name].to_numpy()).all(): + raise AssertionError(f"Grouped benchmark returned invalid estimates in {name!r}.") + return 1.0 + def _execute(self, operation: OperationName) -> float: """Build and fully compute one named benchmark operation.""" @@ -742,10 +823,14 @@ def _execute(self, operation: OperationName) -> float: import_optional("dask", extra_name="benchmark") import dask - statistics = raster.rst.get_stats(["mean", "std", "valid count"]) + statistics = raster.rst.stats(["mean", "std", "valid count"]) mean, _, _ = dask.compute(*statistics.values()) return float(mean) + if operation == "grouped_stats": + assert operation_method is not None + return self._grouped_statistics(raster, operation_method, operation_strategy) + if operation == "subsample": # Return only a fixed-size selection from the much larger raster mp_config = self._multiproc_config(operation) if self.backend == "multiprocessing" else None diff --git a/benchmarks/workflows/variography.py b/benchmarks/workflows/variography.py new file mode 100644 index 000000000..e805a2a6a --- /dev/null +++ b/benchmarks/workflows/variography.py @@ -0,0 +1,69 @@ +"""Prepare deterministic raster, point and pair inputs for variography benchmarks.""" + +from __future__ import annotations + +from typing import Any, Literal + +import numpy as np +import xarray as xr +from rasterio.transform import from_origin + +import geoutils as gu +from geoutils._misc import import_optional + + +def prepare_variogram_pairs(n_pairs: int) -> xr.Dataset: + """Create complete endpoint values at log-uniform distances, independently of spatial pair sampling. + + A fixed seed gives every estimator identical float64 inputs. The increasing difference amplitude creates a + nonconstant variogram, while random endpoint offsets avoid a special case with one constant endpoint. + """ + + # Spread observations across short and long distances without enumerating a spatial distance matrix + rng = np.random.default_rng(42) + distances = np.exp(rng.uniform(0, np.log(1024), n_pairs)) + first = rng.normal(size=n_pairs) + differences = rng.normal(size=n_pairs) * np.sqrt(1 - np.exp(-distances / 100)) + + # Keep only the public pair layout consumed by Variogram.from_pairs() + return xr.Dataset( + { + "distance": ("pair", distances), + "value": (("pair", "endpoint"), np.column_stack((first, first + differences))), + }, + attrs={"min_distance": 1.0, "max_distance": 1024.0}, + ) + + +def prepare_pair_raster(size: int, execution_mode: Literal["eager", "dask"]) -> Any: + """Create a smooth projected raster with scattered missing cells and 256 by 256 Dask chunks. + + Both modes start from the same prepared float32 values. Dask measures selected chunk reads and task scheduling + from memory; this fixture does not measure disk throughput or claim a larger-than-memory contract. + """ + + # Vary values in both directions and remove a known fraction of cells to require finite endpoint checks + rows, columns = np.arange(size)[:, None], np.arange(size)[None, :] + values = (np.sin(columns / 31) + np.cos(rows / 53)).astype(np.float32) + values[(rows * size + columns) % 17 == 0] = np.nan + transform = from_origin(0, size, 1, 1) + + # Expose the public object method in both modes while leaving lazy values uncomputed + if execution_mode == "dask": + import_optional("dask", extra_name="benchmark") + import dask.array as da + + array = da.from_array(values, chunks=(256, 256)) + return gu.RasterAccessor.from_array(array, transform, 32633, nodata=-99999).rst + return gu.Raster.from_array(values, transform, 32633, nodata=-99999) + + +def prepare_pair_pointcloud(n_points: int) -> gu.PointCloud: + """Create irregular projected points at roughly unit spacing with finite, smoothly varying values.""" + + # Keep average point density constant so increasing point count increases the search extent + rng = np.random.default_rng(42) + coordinates = rng.uniform(0, np.sqrt(n_points), size=(n_points, 2)) + x, y = coordinates.T + values = np.sin(x / 31) + np.cos(y / 53) + return gu.PointCloud.from_xyz(x, y, values, crs=32633) diff --git a/dev-environment.yml b/dev-environment.yml index 7d1e32164..b57281316 100644 --- a/dev-environment.yml +++ b/dev-environment.yml @@ -31,6 +31,7 @@ dependencies: - numba # For filters and point cloud gridding - dask # For out-of-memory operations - dask-geopandas # For out-of-memory vector/point operations + - flox # For grouped statistics # Test dependencies - gdal<3.13.3 # To test functionalities against GDAL @@ -63,3 +64,7 @@ dependencies: - pip: - -e ./ + - scikit-gstat>=1.0.23 # For variogram fitting + - gstools>=1.3 # For variogram conversion, kriging and random fields + - gpytorch>=1.11 # For Gaussian-process covariance conversion + - torch>=2 # GPyTorch backend (repeated here as not enforced in GPyTorch) diff --git a/doc/source/api.md b/doc/source/api.md index 177323588..02e4466c8 100644 --- a/doc/source/api.md +++ b/doc/source/api.md @@ -134,14 +134,33 @@ Use {meth}`~geoutils.open_raster` for an {class}`xarray.DataArray`, or instantia ~raster.base.RasterBase.plot ``` +(api-raster-statistics)= ### Statistics +See {ref}`stats` for estimators, grouping by intervals or categories, and variograms. + +```{eval-rst} +.. autosummary:: + :toctree: gen_modules/ + :template: raster_method.rst + + ~raster.base.RasterBase.stats + ~raster.base.RasterBase.variogram +``` + +(api-raster-sampling)= +### Sampling + +See {ref}`sampling` for selecting valid observations, common locations and spatial pairs. + ```{eval-rst} .. autosummary:: :toctree: gen_modules/ :template: raster_method.rst - ~raster.base.RasterBase.get_stats + ~raster.base.RasterBase.subsample + ~raster.base.RasterBase.cosample + ~raster.base.RasterBase.pairsample ``` ### Data manipulation @@ -157,7 +176,6 @@ Use {meth}`~geoutils.open_raster` for an {class}`xarray.DataArray`, or instantia ~raster.base.RasterBase.set_nodata ~raster.base.RasterBase.get_nanarray ~raster.base.RasterBase.get_mask - ~raster.base.RasterBase.subsample ``` ### Loading, writing and converting @@ -584,14 +602,31 @@ documentation](https://shapely.readthedocs.io/en/stable/properties.html). PointCloud.grid ``` +(api-point-statistics)= ### Statistics +See {ref}`stats` for the same statistical workflows on point cloud values. + +```{eval-rst} +.. autosummary:: + :toctree: gen_modules/ + + PointCloud.stats + PointCloud.variogram +``` + +(api-point-sampling)= +### Sampling + +See {ref}`sampling` for sampling values and locations from point clouds. + ```{eval-rst} .. autosummary:: :toctree: gen_modules/ - PointCloud.get_stats PointCloud.subsample + PointCloud.cosample + PointCloud.pairsample ``` ### Testing methods @@ -604,6 +639,99 @@ documentation](https://shapely.readthedocs.io/en/stable/properties.html). PointCloud.georeferenced_coords_equal ``` +(api-statistics)= +## Statistics + +The {ref}`Statistics feature page` introduces these functions and result objects. The corresponding spatial +methods are listed under {ref}`api-raster-statistics` and {ref}`api-point-statistics`. + +### Array estimators + +```{eval-rst} +.. autosummary:: + :toctree: gen_modules/ + + stats.nmad + stats.linear_error + stats.rmse + stats.sum_square +``` + +### Grouped statistics and plotting + +**Zonal statistics are grouped statistics with bins defined by vector features.** Use +`raster.stats(by={"zone": (zones, "id")})`, or the same point cloud method, to calculate statistics by +feature ID. See {ref}`stats-zonal` for examples. The array function below handles already aligned values and groupers. + +```{eval-rst} +.. autosummary:: + :toctree: gen_modules/ + + stats.stats + stats.plot_grouped_stats +``` + +### Variograms and covariance models + +```{eval-rst} +.. autosummary:: + :toctree: gen_modules/ + + stats.variogram + Variogram +``` + +```{eval-rst} +.. autosummary:: + :toctree: gen_modules/ + + Variogram.estimate + Variogram.from_pairs + Variogram.from_model + Variogram.fit + Variogram.plot + Variogram.variogram + Variogram.covariance + Variogram.correlation + Variogram.combine +``` + +### Export and backend conversion + +```{eval-rst} +.. autosummary:: + :toctree: gen_modules/ + + Variogram.to_dataframe + Variogram.to_xarray + Variogram.to_dict + Variogram.from_dict + Variogram.from_skgstat + Variogram.without_backend + Variogram.to_gstools + Variogram.to_gpytorch +``` + +(api-sampling)= +## Sampling + +The {ref}`Sampling feature page` describes selection from one dataset, matching locations between datasets, +and spatial pairs. Their object methods are listed under {ref}`api-raster-sampling` and {ref}`api-point-sampling`. + +### Common-location samples + +{meth}`~geoutils.Raster.cosample` and {meth}`~geoutils.PointCloud.cosample` return a raster or point cloud on the +support selected by `at`. Accessor calls return an {class}`xarray.DataArray` or {class}`geopandas.GeoDataFrame`. +Bands or columns contain `"self"`, `"other"`, then named auxiliaries in mapping order. Raster band names are stored +in `tags["long_name"]` (Xarray `attrs["long_name"]`); point outputs retain the support's selected index labels and +use `"self"` as their active data column. See {ref}`sampling-cosample` for examples. + +### Pair samples + +{meth}`~geoutils.Raster.pairsample` and {meth}`~geoutils.PointCloud.pairsample` return an {class}`xarray.Dataset` +with dimensions `pair` and `endpoint`, containing values, source indexes, coordinates and distances. See +{ref}`sampling-pairs` for its use and the available distance sampling schemes. + ## Multiprocessing configuration To perform **chunked execution** on GeoUtils objects, pass this Multiprocessing configuration to function that support it. diff --git a/doc/source/benchmarking_performance.md b/doc/source/benchmarking_performance.md index b0eaafe67..2851d66da 100644 --- a/doc/source/benchmarking_performance.md +++ b/doc/source/benchmarking_performance.md @@ -1,7 +1,7 @@ (benchmarking-performance)= # Performance -GeoUtils benchmarks its functionalities to measure **RAM and wall time**, check **lazy and out-of-core behaviour** with +GeoUtils benchmarks its functionalities to measure **execution time and memory usage**, check **lazy and out-of-core behaviour** with Dask and Multiprocessing, track **improvements or regressions over time**, and compare different **execution modes** and **calculation engines** against the [**GDAL CLI**](https://gdal.org/en/stable/programs/index.html). Chunked operations also compare strategies used to coordinate selections or reconcile results across chunks. @@ -11,7 +11,7 @@ The benchmarks use GeoUtils' profiling tool described in {ref}`profiling`, which ## Improvements or regressions over time The [**GeoUtils benchmark webpage**](https://glaciohack.github.io/geoutils/) provides performance of core functionalities and their changes with commit history. -It relies on [Airspeed Velocity (ASV)](https://asv.readthedocs.io/) to record fixed performance measurements for each commit and publish them. +It relies on [Airspeed Velocity (ASV)](https://asv.readthedocs.io/) to record reproducible performance measurements for each commit and publish them. ## Comparison across execution modes, engines and GDAL CLI @@ -25,19 +25,19 @@ GeoUtils remains distinct from the external GDAL CLI reference. Two reference graphics summarize the core results: -- **End-to-end time relative to GDAL** for each comparable operation -- **Peak RAM as raster size increases**, including the full process and its workers +- **Execution time relative to GDAL** for each comparable operation (using end-to-end time with file reading/writing for comparison), +- **Peak memory usage as raster size increases**, including the full process and its workers. :::{figure} https://glaciohack.github.io/geoutils/documentation/time_relative_to_gdal.svg -:alt: End-to-end GeoUtils execution-mode time relative to GDAL for four raster operations +:alt: GeoUtils execution mode time relative to GDAL for four raster operations -End-to-end time on the largest raster size shared by every execution mode and the GDAL CLI. GDAL is the reference at one. +Execution time on the largest raster size shared by every execution mode and the GDAL CLI. GDAL is the reference. ::: :::{figure} https://glaciohack.github.io/geoutils/documentation/peak_ram_by_raster_size.svg -:alt: Peak process-tree RAM by raster size for GeoUtils execution modes and GDAL +:alt: Peak memory usage by raster size for GeoUtils execution modes and GDAL -Peak RAM for the benchmark process and all execution-mode workers as raster dimensions increase. +Peak memory usage for the benchmark process and all execution-mode workers as raster dimensions increase. ::: These graphics show the latest complete CI benchmark and may be newer than this documentation version. @@ -45,9 +45,9 @@ These graphics show the latest complete CI benchmark and may be newer than this ## Test suite for scalable execution and large datasets -Every GeoUtils operation advertised as {ref}`chunked or lazy ` is tested -on all supported Python versions and operating systems for its respect of Dask laziness, deferred I/O, and loading behaviour, +Every GeoUtils operation listed as {ref}`chunked or lazy ` is tested +on all supported Python versions and operating systems to ensure its respect of Dask laziness, deferred I/O, and loading behaviour, while yielding **exactly** the same output as in-memory operations. -In addition, GeoUtils also includes large data tests running on the latest Python and Ubuntu, +In addition, GeoUtils also includes large data tests running on latest Python and Ubuntu releases, verifying that these operations use less memory than loading the full raster would require, while yielding a correct result. diff --git a/doc/source/benchmarking_profiling.md b/doc/source/benchmarking_profiling.md index a604446a5..858a4aaa9 100644 --- a/doc/source/benchmarking_profiling.md +++ b/doc/source/benchmarking_profiling.md @@ -13,8 +13,8 @@ kernelspec: (profiling)= # Profiling -GeoUtils has a **built-in profiling tool to measure time and memory** used by a function on your own data and hardware. -The same measurements support the controlled comparisons presented in {ref}`benchmarking-performance`. +GeoUtils has a **built-in profiling tool to measure execution time and memory usage** of a function with your own data and hardware. +The same measurements support the comparisons presented in {ref}`benchmarking-performance`. ```{note} The profiling functionalities rely on [psutil](https://psutil.readthedocs.io/en/latest/) and [plotly](https://plotly.com/) as optional dependencies. You can install them manually or with ``pip install geoutils[opt]`` @@ -69,7 +69,7 @@ Every decorated function called after this is recorded by the profiler. ### The profiled functions GeoUtils profiles the shared implementations of its core numerical operations, including `reproject`, `crop`, -`polygonize`, `rasterize`, `grid`, `get_stats`, `subsample`, `filter`, `interp_points`, `sieve` and `fill_nodata`. +`polygonize`, `rasterize`, `grid`, `stats`, `subsample`, `filter`, `interp_points`, `sieve` and `fill_nodata`. Object methods and Pandas or Xarray accessors therefore record the same underlying operation where supported, without also recording their lightweight wrappers. Memory is sampled at the interval configured by each decorator, which defaults to 0.005 seconds. @@ -86,7 +86,7 @@ def raster_workflow(source_raster): cropped = source_raster.crop((left, bottom, (left + right) / 2, top)) reprojected = cropped.reproject(crs=4326) filtered = reprojected.filter("mean", size=3) - return filtered.get_stats(["mean", "std", "nmad"]) + return filtered.stats(["mean", "std", "nmad"]) # The workflow and its four numerical operations are added to the collection statistics = raster_workflow(raster) diff --git a/doc/source/diagrams/diagram_chunked_subsample.py b/doc/source/diagrams/diagram_chunked_subsample.py index 12bfd5db6..164d6da92 100644 --- a/doc/source/diagrams/diagram_chunked_subsample.py +++ b/doc/source/diagrams/diagram_chunked_subsample.py @@ -9,7 +9,7 @@ from matplotlib.patches import Rectangle from geoutils.raster.array import get_mask_from_array -from geoutils.stats.sampling import ( +from geoutils.sampling.subsampling import ( _get_subsample_size_from_user_input, _splitmix64, _subsample_numpy, diff --git a/doc/source/feature_overview.md b/doc/source/feature_overview.md index a4e2f298f..4cba3c7e7 100644 --- a/doc/source/feature_overview.md +++ b/doc/source/feature_overview.md @@ -8,9 +8,9 @@ As many of our numerical operations rely on **NumPy, SciPy or Numba**, those are The **{ref}`summary tables` directly below** lists the core features of GeoUtils, their scalability and available backends. ```{seealso} -If you are interested in porting from GDAL/OGR, see our {ref}`cheatsheet-osgeo` page. -While tables below provide a scalability summary, the detailed **input/output behaviour of all operations** is available on the {ref}`scalability-support` page. -For measured backend comparisons and guidance on interpreting performance, see {ref}`benchmarking-performance`. +If you are interested in **porting from GDAL/OGR**, see our {ref}`cheatsheet-osgeo` page. +While tables below provide a scalability summary, the detailed **scalable execution behaviour of all operations** is available on the {ref}`scalability-support` page. +For **performance comparisons**, see the {ref}`benchmarking-performance` page. ``` ## Summary @@ -84,16 +84,36 @@ We first describe GeoUtils' core **data operations**, which operate on underlyin - ✅ - Rasterio / GeoPandas -* - {meth}`~geoutils.Raster.get_stats()` +* - {meth}`~geoutils.Raster.stats()` - Compute statistics of valid values over a valid mask. - ❌ - NumPy / SciPy +* - {meth}`~geoutils.Raster.stats()` with ``by`` + - Compute statistics by continuous bins, discrete categories or vector zones (zonal statistics). + - ✅ + - Pandas / NumPy / Dask + * - {meth}`~geoutils.Raster.subsample()` - - Randomly sample valid values. Chunk-invariant seed ensures reproducibility. + - Randomly sample valid values. Choose `strategy="topk"` for the same cells across chunk layouts. - ✅ - NumPy +* - {meth}`~geoutils.Raster.cosample()` + - Select matching finite values from two datasets. Returns a raster or point cloud on the chosen support. + - ✅ (rasters) + - NumPy / Dask + +* - {meth}`~geoutils.Raster.pairsample()` + - Select finite pairs across spatial distances. Returns a compact pair dataset. + - ✅ (rasters) + - NumPy / SciPy / Dask + +* - {meth}`~geoutils.Raster.variogram()` + - Estimate and fit semivariance by distance from sampled pairs. Returns lag statistics and a model. + - ✅ (rasters) + - NumPy / SciKit-GStat + * - {meth}`~geoutils.Raster.filter()` - Filter over window. Fast vectorized logic with NaN support. - ✅ diff --git a/doc/source/history.md b/doc/source/history.md index 79a8bb251..b0218bcf3 100644 --- a/doc/source/history.md +++ b/doc/source/history.md @@ -7,10 +7,10 @@ More information on how the package was created, who are the people behind it, a ## Creation GeoUtils was created during the [GlacioHack](https://github.com/GlacioHack) hackaton event, that took place online on November 8, 2020 and was initiated by -Amaury Dehecq2. +Amaury Dehecq1. ```{margin} -2More on our GlacioHack founder at [adehecq.github.io](https://adehecq.github.io/)! +1More on our GlacioHack founder at [adehecq.github.io](https://adehecq.github.io/)! ``` GeoUtils is inspired by previous efforts that were built directly on top of GDAL and OGR, namely: @@ -21,11 +21,11 @@ GeoUtils is inspired by previous efforts that were built directly on top of GDAL - the package [salem](https://github.com/fmaussion/salem). The initial core development of GeoUtils was mainly performed by members of the Glaciology group of the _Laboratory of Hydraulics, Hydrology and -Glaciology (VAW)_ at ETH Zürich3 and of the _University of Fribourg_, both in Switzerland. The package also received contributions by members of +Glaciology (VAW)_ at ETH Zürich2 and of the _University of Fribourg_, both in Switzerland. The package also received contributions by members of the _University of Oslo_, Norway, the _University of Washington_, US and _Université Grenobles Alpes_, France. ```{margin} -3Check-out [glaciology.ch](https://glaciology.ch) on our founding group of VAW glaciology! +2Check-out [glaciology.ch](https://glaciology.ch) on our founding group of VAW glaciology! ``` ## Joining effort with **demcompare** diff --git a/doc/source/index.md b/doc/source/index.md index f06ed2218..48731e78a 100644 --- a/doc/source/index.md +++ b/doc/source/index.md @@ -114,6 +114,7 @@ transformations raster_vector_point distance_ops stats +sampling filters ``` diff --git a/doc/source/multiprocessing.md b/doc/source/multiprocessing.md index 2ae0fdfef..ed6a0a82a 100644 --- a/doc/source/multiprocessing.md +++ b/doc/source/multiprocessing.md @@ -105,7 +105,7 @@ from typing import Any # Compute mean def compute_statistics(raster: gu.Raster) -> dict[str, np.floating[Any]]: - return raster.get_stats(stats_name=["mean", "valid_count"]) + return raster.stats(["mean", "valid_count"]) stats_results = map_blocks(compute_statistics, filename_rast, config_basic) total_count = sum([stats["valid_count"] for stats in stats_results]) @@ -119,6 +119,33 @@ To include block location in the results, set `return_block_info=True`. --- +## Reprojecting point clouds + +{meth}`~geoutils.PointCloud.reproject` accepts `mp_config` to read, transform and write point rows in chunks. +LAS, LAZ and GeoPackage files can remain unloaded throughout processing. For this operation, `chunks` is an integer +number of points per task, rather than raster dimensions. + +```python +import geoutils as gu +from geoutils.multiproc import ClusterGenerator, MultiprocConfig + +points = gu.PointCloud("observations.laz") + +with ClusterGenerator("multi", nb_workers=4) as cluster: + config = MultiprocConfig(chunks=100_000, outfile="projected.gpkg", cluster=cluster) + projected = points.reproject(crs=32633, mp_config=config) + +assert not points.is_loaded +assert not projected.is_loaded +``` + +The output format follows the filename extension, or an explicit `driver="GPKG"`, `"LAS"` or `"LAZ"`. +GeoPackage is the default when no extension is supplied. Point order and value columns are preserved; LAS/LAZ +coordinates use the file's stored precision. Dataframe accessor calls return an eager GeoDataFrame. Dask point +clouds keep using the lazy `reproject()` path without `mp_config`. + +--- + ## Choosing the right function | Use case | Function | diff --git a/doc/source/pointcloud_class.md b/doc/source/pointcloud_class.md index 54a38f777..283ab3cfa 100644 --- a/doc/source/pointcloud_class.md +++ b/doc/source/pointcloud_class.md @@ -157,11 +157,11 @@ See {ref}`core-array-funcs` for more details. ## Statistics -Statistics of a point cloud, optionally subsetting to an inlier mask, can be computed using {func}`~geoutils.PointCloud.get_stats`. +Statistics of a point cloud can be computed using {func}`~geoutils.PointCloud.stats`. ```{code-cell} ipython3 # Get mean, max and STD of the point cloud -pc.get_stats(["mean", "max", "std"]) +pc.stats(["mean", "max", "std"]) ``` A point cloud can also be quickly subsampled using {func}`~geoutils.PointCloud.subsample`, which considers only valid values, and returns either a point diff --git a/doc/source/raster_class.md b/doc/source/raster_class.md index 8774de6f6..bfbcde717 100644 --- a/doc/source/raster_class.md +++ b/doc/source/raster_class.md @@ -419,11 +419,11 @@ rast_reproj.to_xarray() ## Statistics -Statistics of a raster, optionally subsetting to an inlier mask, can be computed using {func}`~geoutils.Raster.get_stats`. +Statistics of a raster, optionally subsetting to an inlier mask, can be computed using {func}`~geoutils.Raster.stats`. ```{code-cell} ipython3 # Get mean, max and STD of the raster -rast.get_stats(["mean", "max", "std"]) +rast.stats(["mean", "max", "std"]) ``` A raster can also be quickly subsampled using {func}`~geoutils.Raster.subsample`, which can consider only valid values, and return either a point cloud or an diff --git a/doc/source/sampling.md b/doc/source/sampling.md new file mode 100644 index 000000000..533352d74 --- /dev/null +++ b/doc/source/sampling.md @@ -0,0 +1,363 @@ +--- +file_format: mystnb +jupytext: + formats: md:myst + text_representation: + extension: .md + format_name: myst +kernelspec: + display_name: geoutils-env + language: python + name: geoutils +--- +(sampling)= +# Sampling + +GeoUtils contains functionalities to **sample subsets of valid data from rasters and point clouds**, all supporting **chunked execution** to scale efficiently on large datasets. + +Three types of sampling operations are supported: + +- **Subsampling** selects a **random subset of valid values** in a single dataset, +- **Co-sampling** selects the **co-located sample of valid values** common to two primary datasets, +- **Pair-sampling** selects **pairs of values within a single dataset**. + +Co-sampling is especially useful for analyses requiring multiple large datasets (such as image registration), +while pair-sampling is at the core of geostatistics (variography and kriging). Subsampling is widely used across applications, +starting within co-sampling and pair-sampling themselves. + +```{note} +Sampling can be run directly through statistical operations in {ref}`stats` (e.g., for summary statistics or variogram estimation). +Use the API below for custom usage. +``` + +## Summary and quick use + +| Operation | Selection | Result | +| --- |-----------------------------------------------| --- | +| {meth}`~geoutils.Raster.subsample` | Valid observations in one dataset | Values or indexes | +| {meth}`~geoutils.Raster.cosample` | Locations valid in both datasets | Raster or point cloud on the selected support | +| {meth}`~geoutils.Raster.pairsample` | Pairs of locations within one dataset | {class}`xarray.Dataset` with paired values and distances | + + +```{code-cell} ipython3 +:tags: [remove-cell] + +# Match the figure resolution and text size used by the other feature pages +from matplotlib import pyplot as plt +plt.rcParams["figure.dpi"] = 600 +plt.rcParams["savefig.dpi"] = 600 +plt.rcParams["font.size"] = 9 +``` + +```{code-cell} ipython3 +:tags: [hide-cell] +:mystnb: +: code_prompt_show: "Show the code for opening example files" +: code_prompt_hide: "Hide the code for opening example files" + +import geoutils as gu +import numpy as np + +# Open a projected elevation raster and glacier outlines for the sampling examples +rast = gu.Raster(gu.examples.get_path("exploradores_aster_dem")) +glaciers = gu.Vector(gu.examples.get_path("exploradores_rgi_outlines")) +``` + +(sampling-subsample)= +## Subsampling + +{meth}`~geoutils.Raster.subsample` or {meth}`~geoutils.PointCloud.subsample`. + +Subsampling selects a **random subset of valid values, without replacement**. + +Use `subsample` to set the subsample size: a fraction between 0 and 1 selects that proportion of valid data, while a number above one sets the requested count. +For example, `subsample=0.1` selects 10%, `subsample=1` keeps all valid values, and `subsample=1000` selects 1000 values. + +```{code-cell} ipython3 +# Select a reproducible valid subsample of 2000 elevations +sample = rast.subsample(subsample=2000, random_state=42, strategy="topk") +sample[:5] +``` + +Pass `mask` to restrict eligible locations **before calculating the sample size**. Boolean arrays keep True values, +while vector outlines keep locations inside their geometries. Missing mask entries are excluded. For raster +sampling, mask rasters must share the source grid. For point sampling, point masks must follow the source points' +ordered coordinates and CRS; raster masks must share their CRS and are read with nearest interpolation. + +```{code-cell} ipython3 +# Select 10% of valid elevations inside glacier outlines +glacier_sample = rast.subsample(0.1, mask=glaciers, random_state=42) +glacier_sample[:5] +``` + +Use `return_indices=True` to get **sample locations instead of values**, for example to select the same cells in +several aligned arrays. Raster indexes are rows and columns; point cloud indexes are positions in the original table. +Masks do not change these index positions, and sampled values keep the source dtype. + +```{code-cell} ipython3 +# Recover the same raster cells with the same sample size, strategy and seed +rows, columns = rast.subsample(2000, return_indices=True, random_state=42, strategy="topk") +np.array_equal(rast.data[rows, columns], sample) +``` + +To keep coordinates and georeferencing together with the sample, use {meth}`~geoutils.Raster.to_pointcloud`: + +```{code-cell} ipython3 +# Keep sampled elevations together with their coordinates and CRS +points = rast.to_pointcloud(subsample=2000, random_state=42) +``` + +```{code-cell} ipython3 +:tags: [hide-input] +:mystnb: +: code_prompt_show: "Show the code for plotting the figure" +: code_prompt_hide: "Hide the code for plotting the figure" + +fig, axes = plt.subplots(1, 2, figsize=(8, 3)) +rast.plot(ax=axes[0], cmap="terrain", vmin=300, vmax=4000, cbar_title="Elevation (m)") +axes[0].set_title("Complete raster") +points.plot(ax=axes[1], cmap="terrain", vmin=300, vmax=4000, cbar_title="Elevation (m)", markersize=2) +axes[1].set_title("2000 sampled locations") +axes[1].set_yticklabels([]) +plt.tight_layout() +``` + +(sampling-cosample)= +## Co-sampling + +{meth}`geoutils.Raster.cosample` or {meth}`geoutils.PointCloud.cosample`. + +Co-sampling selects **values from two datasets at the same (co-located) valid locations**. It accounts for georeferencing, nodata +and optional masks, so the selected values can be compared directly. Use `subsample` to additionally select a common valid subset. + +For example, compare a DEM with a coarser version outside glacier outlines. `align="reproject"` allows the coarse +raster to be resampled onto the calling raster's grid: + +```{code-cell} ipython3 +# Compare fine and coarse elevations on the common grid outside glacier outlines +coarse = rast.reproject(res=90) +paired = rast.cosample( + coarse, + mask=glaciers, + mask_mode="outside", + align="reproject", +) +fine_elevation, coarse_elevation = paired.split_bands() +paired +``` + +The output has **two bands with a common grid and mask**: `"self"` for the calling raster and `"other"` for the second +raster. Standard raster operations can then calculate their difference: + +```{code-cell} ipython3 +# Subtract the two bands on their common finite support +differences = fine_elevation - coarse_elevation +``` + +```{code-cell} ipython3 +:tags: [hide-input] +:mystnb: +: code_prompt_show: "Show the code for plotting the figure" +: code_prompt_hide: "Hide the code for plotting the figure" + +fig, ax = plt.subplots(figsize=(5, 3)) +differences.plot(ax=ax, cmap="RdBu", vmin=-50, vmax=50, cbar_title="Fine − coarse elevation (m)") +glaciers.plot(ref_crs=rast, ax=ax, ec="k", fc="none", linewidth=0.5) + +# Keep the view on the raster extent even though the outlines extend beyond it +ax.set_xlim(rast.bounds.left, rast.bounds.right) +ax.set_ylim(rast.bounds.bottom, rast.bounds.top) +ax.set_title("Common sample outside glaciers") +plt.tight_layout() +``` + +### Common locations and output + +Use `at` to choose the **output grid or point locations**. It accepts `"self"`, `"other"` or a spatial object. +By default, two rasters use the calling raster's grid, while a raster–point comparison uses the point locations. +Use `raster_point_mode` to choose the direction of a raster–point comparison independently of the calling object: + +| Mode | Output locations | Value calculation | +| --- | --- | --- | +| `"resample_raster"` | Point coordinates | Evaluate rasters with `resample_method`, passed to {meth}`~geoutils.Raster.interp_points` | +| `"grid_points"` | Raster grid | Grid point values with `grid_method`, passed to {meth}`~geoutils.PointCloud.grid` | + +Both methods default to `"linear"`. `resample_method="nearest"` selects the nearest raster value; +`resample_method="linear"` interpolates neighboring raster values at the target point coordinates. +For point output, all point cloud inputs must share the chosen ordered coordinates. + +```{code-cell} ipython3 +# Interpolate the coarse raster at selected point observations +at_points = coarse.cosample( + points, + raster_point_mode="resample_raster", + resample_method="linear", + subsample=500, + random_state=42, +) +at_points.ds[["self", "other", "geometry"]].head() +``` + +**The output follows the selected locations and calling interface:** + +| Locations | GeoUtils object call | Accessor call | +| --- | --- | --- | +| Raster grid | {class}`~geoutils.Raster` | {class}`xarray.DataArray` | +| Point coordinates | {class}`~geoutils.PointCloud` | {class}`geopandas.GeoDataFrame` | + +A raster keeps its grid, with unsampled cells masked. A point cloud contains only selected points and keeps their +original index labels and order. To compare two rasters on a point set, pass that point cloud as `at`. + +Use one object family throughout a call: Raster/PointCloud objects, or DataArray/GeoDataFrame objects accessed +through `.rst` and `.pc`. This also applies to spatial auxiliaries, explicit `at`, and raster or point masks. +DataArrays and GeoDataFrames may mix eager and Dask storage. Plain arrays and vector outlines work with either +family. For point output, all point inputs must share the same ordered horizontal coordinates; this is checked +before raster alignment or interpolation. Point inputs may have different locations when gridded onto a raster. + +An explicit `at` selects exact locations and determines the conversion direction when the mode is omitted. +When both are specified, they must agree. With an explicit mode and no `at`, exactly one primary input must supply +the requested spatial type; otherwise choose `at` explicitly. Omitting both retains the defaults described above. +Grid or CRS mismatches raise unless `align="reproject"` is set. + +For example, use circular means of nearby point observations on a raster grid: + +```python +on_grid = coarse.cosample( + points, + raster_point_mode="grid_points", + grid_method="mean", + grid_kwargs={"dist_nodata_pixel": 2, "min_points": 3}, +) +``` + +Here the radius is two output pixels and each estimate requires three finite points. Gridding estimates values at +grid locations and can change the number and spatial distribution of observations used in a comparison. +Additional gridding options belong in `grid_kwargs`; raster interpolation options such as `nodata_propagation` +belong in `resample_kwargs`. Target locations and method names use the explicit co-sampling arguments. + +`resample_method="reduce"` is reserved for reducing raster windows around point coordinates. It currently raises +`NotImplementedError` because its integration requires revision of {meth}`~geoutils.Raster.reduce_points`. +That existing method remains available separately. Co-sampling now uses `resample_method` in place of its previous +`interpolation` argument; the `stats` interpolation argument for grouped calculations is unchanged. + +### Auxiliary variables + +Auxiliary variables carry **additional values at the same common locations**, such as terrain attributes or +measurement weights. Locations must be valid in both primary datasets and every auxiliary variable. + +Spatial objects supply their own georeferencing and use the first raster band or active point values by default. +Select another band with `auxiliary={"slope": (slope_raster, 2)}` or a point column with +`auxiliary={"intensity": (points, "intensity")}`. For plain arrays, `auxiliary_at` identifies the input whose grid +or point ordering they follow. Raster arrays must match one input band's shape; point arrays must be one-dimensional +with one value per input point: + +```{code-cell} ipython3 +# Carry an aligned elevation predictor through the same selection +with_auxiliary = rast.cosample( + coarse, + auxiliary={"elevation": rast.data}, + auxiliary_at="self", + align="reproject", +) +first, second, elevation = with_auxiliary.split_bands() +with_auxiliary.tags["long_name"] +``` + +Bands or columns contain `"self"`, `"other"`, then auxiliaries in mapping order. Raster band names are stored in +`tags["long_name"]`, or `attrs["long_name"]` for Xarray. Point clouds use these names as columns, with `"self"` +as the active data column. Use the usual `.data`, `.ds` or `.split_bands()` methods to access the values. + +```{code-cell} ipython3 +# Use native Xarray bands for the same grid comparison +native = coarse.to_xarray().rst.cosample(coarse.to_xarray()) +native.isel(band=0) - native.isel(band=1) +``` + +(sampling-pairs)= +## Pair-sampling + +{meth}`geoutils.Raster.pairsample` or {meth}`geoutils.PointCloud.pairsample`. + +Pair-sampling selects **pairs of valid values and their spatial separation within one dataset**. It provides the +observations used in variography without calculating distances between every possible pair of locations. + +```{code-cell} ipython3 +# Sample pairs at short and long separations within 5 km +pairs = rast.pairsample(n_pairs=20_000, max_distance=5000, random_state=42) +pairs.isel(pair=slice(0, 5)) +``` + +The output is an {class}`xarray.Dataset` indexed by `pair` and `endpoint`. It contains both values, their coordinates +and source indexes, and their distance. **Distances use the coordinate units**: use a projected CRS in metres for +distances in metres. + +```{code-cell} ipython3 +# Calculate signed differences while keeping the pair labels +pair_differences = pairs["value"].sel(endpoint="second") - pairs["value"].sel(endpoint="first") +pair_differences.isel(pair=slice(0, 5)) +``` + +### Sampling across distances + +**Logarithmic distance sampling** (`sampling="loglag"`, the default) represents both short and long separations. +**Uniform endpoint sampling** (`sampling="random_xy"`) selects locations uniformly, which usually gives fewer pairs +at short distances. Use `min_distance`, `max_distance` and `mask` to restrict the selection. + +```{code-cell} ipython3 +# Compare logarithmic distance sampling with uniformly selected endpoints +random_pairs = rast.pairsample( + n_pairs=20_000, sampling="random_xy", max_distance=5000, random_state=42 +) +``` + +```{code-cell} ipython3 +:tags: [hide-input] +:mystnb: +: code_prompt_show: "Show the code for plotting the figure" +: code_prompt_hide: "Hide the code for plotting the figure" + +fig, ax = plt.subplots(figsize=(6, 3)) +distance_edges = np.geomspace(rast.res[0], 5000, 13) +ax.hist(pairs["distance"], bins=distance_edges, histtype="step", label="Logarithmic distances") +ax.hist(random_pairs["distance"], bins=distance_edges, histtype="step", label="Uniform endpoints") +ax.set(xscale="log", xlabel="Distance (m)", ylabel="Number of pairs") +ax.legend() +plt.tight_layout() +``` + +`n_pairs` sets a target count. Sparse data, masks or distance limits can leave fewer accepted pairs; +`pairs.sizes["pair"]` gives the returned count. See {ref}`api-sampling` for the available sampling strategies. + +To estimate a variogram from these pairs, use {meth}`~geoutils.Variogram.from_pairs`. Alternatively, +{meth}`~geoutils.Raster.variogram` or {meth}`~geoutils.PointCloud.variogram` performs sampling and estimation in one +call, as described in {ref}`stats-variograms`. + +(sampling-reproducibility)= +## Reproducibility and chunked execution + +**Set `random_state` to reproduce a sample** with the same inputs and sampling settings. Raster subsampling offers +two strategies: + +- `"sequential"` draws from the sequence of valid values. It is the default for `subsample()` and can depend on chunk layout. +- `"topk"` selects the same cells regardless of chunk layout. It is the default for raster `cosample()` and grouped statistics. + +For grouped `stats()`, pass this choice as **`subsampling_strategy`**. Its separate `strategy` argument controls +aggregation across chunks; see {ref}`stats-grouped`. + +This guarantee applies to the `topk` raster sampler. Pair-sampling has its own strategies; for example, +`"chunk_anchors"` uses the chunk layout to select pairs. + +**Co-sampling supports Dask and multiprocessing through its spatial operations.** Dask accessor calls keep raster +bands or point partitions lazy. Counts and sample positions are computed during preparation, including a check for +an empty common selection. Final point interpolation and removal of missing values stay lazy; computing the result +can therefore leave fewer points, or none, when interpolation spreads nodata into the selected locations. Raster +output keeps the complete grid, even when `subsample` limits the selected cells; point output preserves the selected +row order and index labels. + +Pass `mp_config` directly to `cosample()` to use multiprocessing with eager or unloaded inputs. Raster output is +written by tiles to `mp_config.outfile`, and temporary intermediate files are cleaned automatically. Point output is +collected after interpolating the selected rows. Dask inputs and `mp_config` cannot be combined. + +Pair-sampling returns an in-memory pair dataset and, for point clouds, currently reads all source coordinates and +values. Use an absolute sample count to limit the returned rows or pairs. See {ref}`scalability-logic` for the chunked +algorithms. diff --git a/doc/source/scalability_concept.md b/doc/source/scalability_concept.md index a44a2bf22..30d13ae9a 100644 --- a/doc/source/scalability_concept.md +++ b/doc/source/scalability_concept.md @@ -64,7 +64,7 @@ Accessing {attr}`~geoutils.Raster.data`, or calling operations that require the print(f"Is raster loaded before data operation? {ds.rst.is_loaded}") # We compute statistics, which loads the array -ds.rst.get_stats() +ds.rst.stats() # The raster is now loaded print(f"Is raster loaded after data operation? {ds.rst.is_loaded}") diff --git a/doc/source/scalability_support.md b/doc/source/scalability_support.md index 190470117..7f7834465 100644 --- a/doc/source/scalability_support.md +++ b/doc/source/scalability_support.md @@ -102,7 +102,7 @@ The **memory usage** column lists the number of input chunks loaded in memory fo - {bdg-secondary}`In-memory` - {bdg-secondary}`In-memory` - — -* - {meth}`~geoutils.Raster.get_stats` +* - {meth}`~geoutils.Raster.stats` - {bdg-secondary}`In-memory` - {bdg-secondary}`In-memory` - — @@ -168,7 +168,7 @@ The **memory usage** column lists the number of input chunks loaded in memory fo - - -* - {meth}`~geoutils.PointCloud.get_stats` +* - {meth}`~geoutils.PointCloud.stats` - {bdg-secondary}`In-memory` - {bdg-secondary}`In-memory` - — diff --git a/doc/source/stats.md b/doc/source/stats.md index 339423336..bd983a665 100644 --- a/doc/source/stats.md +++ b/doc/source/stats.md @@ -11,103 +11,488 @@ kernelspec: name: geoutils --- (stats)= - # Statistics -GeoUtils supports statistical analysis tailored to geospatial objects. +GeoUtils provides **summary statistics, grouped statistics and variography for rasters and point clouds**, accounting +for **georeferencing and nodata**. The same methods are available on GeoUtils objects and their Xarray or GeoPandas +accessors. -For a {class}`~geoutils.Raster` or a {class}`~geoutils.PointCloud`, the statistics are naturally performed on the {attr}`~geoutils.Raster.data` attribute -which is clearly defined. +Three types of statistical operations are supported: -[//]: # (For a {class}`~geoutils.Vector`, statistics have to be performed on a specific column.) +- **Summary statistics** describe the **distribution of valid values** in a dataset, +- **Grouped statistics** describe values **within bins or categories** defined by other variables, +- **Variography** describes **spatial variability as a function of distance**. -```{warning} -The API for statistical features is preliminary and might change with the release of zonal and grouped statistics. -``` +**Zonal statistics are a special case of grouped statistics: vector features define the bins.** For example, +`stats()` can calculate the mean elevation of each glacier or catchment using its outline. -## Estimators +```{note} +Statistical operations can select samples directly, for example when estimating a variogram or grouped statistics. +See {ref}`sampling` to select observations for other analyses, and {ref}`api-statistics` for the API reference. +``` -The {func}`~geoutils.Raster.get_stats` method allows to extract key statistical estimators from a raster or a point cloud, optionally subsetting to an -inlier mask. +## Summary and quick use -Supported statistics are : -- **Mean:** arithmetic mean of the data, ignoring masked values. -- **Median:** middle value when the valid data points are sorted in increasing order, ignoring masked values. -- **Max:** maximum value among the data, ignoring masked values. -- **Min:** minimum value among the data, ignoring masked values. -- **Sum:** sum of all data, ignoring masked values. -- **Sum of squares:** sum of the squares of all data, ignoring masked values. -- **90th percentile:** point below which 90% of the data falls, ignoring masked values. -- **IQR (Interquartile Range):** difference between the 75th and 25th percentile of a dataset, ignoring masked values. -- **LE90 (Linear Error with 90% confidence):** difference between the 95th and 5th percentiles of a dataset, representing the range within which 90% of the data points lie. Ignore masked values. -- **NMAD (Normalized Median Absolute Deviation):** robust measure of variability in the data, less sensitive to outliers compared to standard deviation. Ignore masked values. -- **RMSE (Root Mean Square Error):** commonly used to express the magnitude of errors or variability and can give insight into the spread of the data. Only relevant when the raster represents a difference of two objects. Ignore masked values. -- **Std (Standard deviation):** measures the spread or dispersion of the data around the mean, ignoring masked values. -- **Valid count:** number of finite data points in the array. It counts the non-masked elements. -- **Total count:** total size of the raster. -- **Percentage valid points:** ratio between **Valid count** and **Total count**. +| Operation | Calculation | Result | +| --- | --- | --- | +| {meth}`~geoutils.Raster.stats` | Statistics of all valid values or by bins, categories, or vector zones | Number, dictionary, or {class}`pandas.DataFrame` | +| {meth}`~geoutils.Raster.variogram` | Spatial variability across distances | {class}`~geoutils.Variogram` | -If an inlier mask is passed: -- **Total inlier count:** number of data points in the inlier mask. -- **Valid inlier count:** number of unmasked data points in the array after applying the inlier mask. -- **Percentage inlier points:** ratio between **Valid inlier count** and **Valid count**. Useful for classification statistics. -- **Percentage valid inlier points:** ratio between **Valid inlier count** and **Total inlier count**. +```{code-cell} ipython3 +:tags: [remove-cell] -Callable functions are supported as well. +# Match the figure resolution and text size used by the other feature pages +from matplotlib import pyplot as plt +plt.rcParams["figure.dpi"] = 600 +plt.rcParams["savefig.dpi"] = 600 +plt.rcParams["font.size"] = 9 +``` ```{code-cell} ipython3 +:tags: [hide-cell] +:mystnb: +: code_prompt_show: "Show the code for opening example files" +: code_prompt_hide: "Hide the code for opening example files" + import geoutils as gu import numpy as np +import pandas as pd + +# Open an elevation raster and glacier outlines covering the same region +rast = gu.Raster(gu.examples.get_path("exploradores_aster_dem")) +glaciers = gu.Vector(gu.examples.get_path("exploradores_rgi_outlines")) +``` + +(stats-estimators)= +## Summary statistics + +{meth}`geoutils.Raster.stats` or {meth}`geoutils.PointCloud.stats`. + +Summary statistics describe **central values, spread and valid counts**. For rasters, they use the selected band; +for point clouds, they use the active {attr}`~geoutils.PointCloud.data_column` or the geometry's Z coordinate. +Built-in estimators exclude nodata. + +```{code-cell} ipython3 +# Compute the default summary statistics +rast.stats() +``` + +Request **one statistic for a number**, or **several statistics for a dictionary**. Use `"all"` to request every +available estimator and count. + +```{code-cell} ipython3 +rast.stats("mean") +``` + +```{code-cell} ipython3 +rast.stats(["mean", "median", "std", "nmad"]) +``` + +```{dropdown} Available estimators and counts +| Statistic | Meaning | +| --- | --- | +| Mean, median | Arithmetic average or middle value | +| Min, max | Smallest or largest value | +| Sum, sum of squares | Sum of values or of their squares | +| 90th percentile | Value below which 90% of observations fall | +| IQR | Difference between the 75th and 25th percentiles | +| LE90 | Difference between the 95th and 5th percentiles | +| NMAD | Median absolute deviation from the median, scaled by 1.4826 | +| RMSE | Square root of the mean squared value, useful for error data | +| Standard deviation | Spread around the mean | +| Valid count, total count | Number of finite observations and total number of locations | +| Percentage valid points | Valid count divided by total count, as a percentage | + +With an inlier mask, additional counts describe the selected locations, their finite values and their proportions +of the mask and dataset. +``` + +**Robust estimators**, such as the median and NMAD, are less sensitive to outliers than the mean and standard deviation. +The array functions {func}`~geoutils.stats.nmad`, {func}`~geoutils.stats.linear_error` and {func}`~geoutils.stats.rmse` +are also available directly. + +Custom estimators receive the data array and must handle its mask or NaNs. For example, count valid elevations +above a threshold: + +```{code-cell} ipython3 +def count_high_elevations(data: np.ndarray) -> int: + """Count observations above 1500 m in the selected elevation data.""" + + # Exclude masked and non-finite observations before counting high elevations + values = np.ma.asarray(data).compressed() + return int(np.count_nonzero(np.isfinite(values) & (values > 1500))) + +rast.stats(count_high_elevations) +``` + +Use `mask` to **restrict the statistics to selected locations**: + +```{code-cell} ipython3 +# Summarize elevations inside all glacier outlines together +glacier_mask = glaciers.create_mask(rast) +rast.stats(["mean", "std", "valid count"], mask=glacier_mask) +``` + +To calculate a separate statistic for each glacier, use {ref}`zonal statistics` below. + +(stats-grouped)= +## Grouped statistics + +{meth}`geoutils.Raster.stats`, {meth}`geoutils.PointCloud.stats` or {func}`geoutils.stats.stats` with ``by``. + +Grouped statistics describe **values within bins or categories of one or more variables**. Bins can be continuous +intervals, such as elevation bands, or discrete categories, such as land-cover classes or vector zones. + +Use `by` to name the grouping variables and `values` to select the data to summarize. Raster values use band numbers +starting at one; point cloud values use column names. A mapping names those values in the output and can also contain +**external rasters, point clouds or `(object, band_or_column)` selections**. Vector attributes can provide numeric +values or grouping variables; supply `bins` to treat an attribute as continuous. + +### Continuous bins + +Use `bins` to define **intervals of a continuous variable**. For example, summarize elevations within elevation bands: + +```{code-cell} ipython3 +# Divide the first raster band into explicit elevation intervals +elevation_edges = [0, 1000, 2000, 3000, 4000] +elevation_stats, elevation_masks = rast.stats( + ("mean", "min", "max"), + by={"elevation": 1}, + values={"elevation": 1}, + bins={"elevation": elevation_edges}, + return_masks=True, +) +elevation_stats +``` + +The output table has **one row per group and one column per value–statistic combination**. Each value also has a +finite `count`. With `return_masks=True`, the same row key retrieves the group's geographic mask: + +```{code-cell} ipython3 +# Retrieve the same elevation interval in the table and on the raster grid +first_interval = elevation_stats.index[0] +first_mask = elevation_masks[first_interval] +elevation_stats.loc[first_interval] +``` + +```{code-cell} ipython3 +:tags: [hide-input] +:mystnb: +: code_prompt_show: "Show the code for plotting the figure" +: code_prompt_hide: "Hide the code for plotting the figure" + +fig, axes = plt.subplots(1, 2, figsize=(8, 3)) +rast.plot(ax=axes[0], cmap="terrain", cbar_title="Elevation (m)") +axes[0].set_title("Elevation raster") +first_mask.plot(ax=axes[1], cbar_title="Inside interval (1=yes, 0=no)") +axes[1].set_title(f"Elevation interval {first_interval} m") +axes[1].set_yticklabels([]) +plt.tight_layout() +``` + +An integer, such as `bins={"elevation": 10}`, creates ten equal-width bins. Numeric edges include each lower edge +and the final upper edge. Pass a {class}`pandas.IntervalIndex` to choose which edges are included: + +```{code-cell} ipython3 +# Assign boundary elevations to the interval ending at that elevation +right_closed = pd.IntervalIndex.from_breaks(elevation_edges, closed="right") +rast.stats(by={"elevation": 1}, bins={"elevation": right_closed}) +``` + +### Discrete categories + +Use `categories` to define **discrete classes and their order**. A boolean glacier mask, for example, separates terrain +inside and outside the outlines: -# Instantiate a raster from a filename on disk -filename_rast = gu.examples.get_path("exploradores_aster_dem") -rast = gu.Raster(filename_rast) -rast +```{code-cell} ipython3 +# Compare elevation distributions inside and outside glaciers +glacier_stats = rast.stats( + ("mean", "min", "max", "nmad"), + by={"glacier": glacier_mask}, + values={"elevation": 1}, + categories={"glacier": [False, True]}, +) +glacier_stats ``` -By default and without any specification, this function computes the following main statistics: -minimum, maximum, mean, standard deviation, normalized median absolute deviation, total count, and percentage of valid points. +The same operation applies to land-cover codes or other categorical variables. Values outside the declared categories +are excluded. Boolean and Pandas categorical inputs can supply their categories without an explicit declaration. + +(stats-zonal)= +### Zonal statistics + +**Zonal statistics use vector features as bins.** Pass `(vector, "column")` in `by` to group raster cells or points +by a vector attribute. **A unique feature ID gives one group per feature**; repeated IDs combine features into a +group. Categories are inferred from the column, so no `bins` or `categories` argument is needed. + +For example, the glacier inventory's `RGIId` column identifies each glacier: + ```{code-cell} ipython3 -rast.get_stats() +# Calculate separate elevation statistics for each glacier outline +glacier_zonal_stats = rast.stats( + ("mean", "std", "min", "max"), + by={"glacier": (glaciers, "RGIId")}, + values={"elevation": 1}, + observed=False, +) +glacier_zonal_stats.sort_values(("elevation", "count"), ascending=False).head() ``` -To compute all available statistics, set `stats_name` to `all`. +Locations outside all zones are excluded, and each statistic uses the finite values in its zone. `observed=False` +retains zones without observations, with zero counts and undefined statistics. Use `return_masks=True` to retrieve +zone masks by the same feature IDs. + +The same grouping applies to point measurements: + ```{code-cell} ipython3 -rast.get_stats("all") +# Summarize a sample of point elevations within each glacier +points = rast.to_pointcloud(data_column_name="elevation", subsample=2000, random_state=42) +points.stats(("mean", "std"), by={"glacier": (glaciers, "RGIId")}).head() ``` -Get a single statistic (e.g., 'mean') as a float: +Raster cells are assigned by their centres, and points by their coordinates. Grouping assigns each location to one +category; use non-overlapping zones for independent feature statistics. Vector and GeoDataFrame inputs are both +accepted, including when the raster values are backed by Dask. + +```{note} +`by={"zone": (zones, "id")}` groups by feature IDs. A bare vector, `by={"inside": zones}`, groups the union of all +features into inside/outside categories. `mask=zones` instead restricts all groups to locations inside the features. +``` + +### Several variables and plotting + +**Several grouping variables define their joint bins or categories**. For example, combine elevation bands with +glacier membership to compare elevation distributions inside and outside glaciers at similar elevations: + ```{code-cell} ipython3 -rast.get_stats("mean") +# Retain empty combinations so the complete comparison grid remains visible +joint_stats = rast.stats( + ("median", "nmad"), + by={"elevation": 1, "glacier": glacier_mask}, + values={"elevation": 1}, + bins={"elevation": elevation_edges}, + categories={"glacier": [False, True]}, + observed=False, +) +joint_stats ``` -Get multiple statistics: +The rows form a {class}`pandas.MultiIndex`. Vector zones can be combined with continuous bins in the same way, +by including both variables in `by`. `observed=False` retains empty combinations; the default, `observed=True`, +omits combinations without eligible locations. + +Use {func}`~geoutils.stats.plot_grouped_stats` to **plot a statistic and its sample counts** for one or two grouping +variables. `min_count` hides estimates based on too few observations. + ```{code-cell} ipython3 -rast.get_stats(["mean", "max", "std"]) +:tags: [hide-input] +:mystnb: +: code_prompt_show: "Show the code for plotting the figure" +: code_prompt_hide: "Hide the code for plotting the figure" + +# Display within-group elevation spread and the number of observations supporting it +fig, ax = plt.subplots(figsize=(7, 4)) +fig.subplots_adjust(right=0.78) +axes = gu.stats.plot_grouped_stats(joint_stats, value="elevation", statistic="nmad", min_count=100, ax=ax) +axes["statistic"].set(xlabel="Elevation (m)", ylabel="Inside glacier") + +# Place the color scale beyond the count panels so every elevation interval stays visible +axes["colorbar"].set_position((0.84, 0.11, 0.025, 0.5)) +_ = axes["colorbar"].set_ylabel("Elevation NMAD (m)") ``` -Using a custom callable statistic: +To summarize only one variable, make a separate grouping call. Statistics such as medians cannot generally be +combined from subgroup results. + +### Choosing a grouped reduction strategy + +`strategy` controls **how values belonging to a group are combined across chunks**. The names `dense` and `sparse` +describe intermediate group summaries. They do not describe the number of finite input pixels, select different +groups, or change the requested statistic. `observed` independently controls whether empty groups appear in the result. + +| Strategy | What each chunk contributes | Suitable calculations and limits | +| --- | --- | --- | +| `"dense"` | A summary slot for every declared group combination, including groups absent from the chunk | Counts, means, standard deviations, sums, RMSE and extrema. Avoids matching group IDs during merges, but allocation grows with the complete group space | +| `"sparse"` | Summaries and IDs for groups encountered in that chunk | The same mergeable statistics. Useful when each chunk contains a small fraction of the declared groups; sorting and merging IDs adds work | +| `"groupwise"` | The actual observations for complete groups | Exact medians, quantiles, NMAD and custom functions. Reads intersecting chunks and batches small groups sharing those chunks; a large group still requires memory proportional to its full membership | +| `"auto"` | Chooses one of the above | Uses groupwise for exact quantiles, NMAD or custom functions; otherwise dense for at most 4096 declared group combinations, sparse above that | + +For example, 100 elevation bins combined with 100 land-cover categories declare **10,000 combinations**, even if +only a few occur. Dense keeps slots for all combinations in each chunk. Sparse initially stores only those present; +its summaries can grow as chunks are merged. This distinction is most useful when groups occupy localized regions. +If every chunk contains every group, sparse has little allocation advantage. + +The automatic threshold is a heuristic based on declared group count. It does not measure group occupancy or +available memory. Requesting an exact median alongside a mean selects groupwise for the entire calculation. +For chunked input, requesting those exact statistics with dense or sparse raises an error. Eager input can compute +exact statistics directly because its complete values are already resident in memory. + +Dask and multiprocessing share the numerical kernels. Dask can read arrays in chunks; multiprocessing currently +tiles arrays already available in the client. Neither the choice of strategy nor smaller chunks bounds the memory +needed by one very large exact group. The result table is returned in memory, and `observed=False` can itself produce +a large table of every declared group combination. + +Groupwise execution may reread a chunk when several batches need it. Smaller chunks can reduce temporary array +sizes while increasing the number of gathering tasks, so they do not necessarily make exact statistics faster. + +`subsampling_strategy` is a separate control: `"topk"` keeps the selected raster cells independent of chunk layout +for a fixed seed, while `"sequential"` follows the ordinary sampling workflow. Subsampling changes which observations +enter the estimates; the reduction strategy controls how those selected observations are combined. Small floating +point differences can arise from different summation orders across chunk layouts. + +Set **`subsample_per_group=True`** to sample within each group. For example, this uses at most 1,000 eligible +locations per land-cover category: + +```python +rast.stats( + ["mean", "std", "nmad"], + by={"landcover": landcover}, + categories={"landcover": classes}, + subsample=1000, + subsample_per_group=True, + random_state=42, +) +``` + +With multiple grouping variables, the limit applies to each combined group. A fraction such as `subsample=0.1` +keeps ten percent of each group's eligible locations, rounded down; very small groups may receive no sample. +`subsample=1` keeps all eligible locations. Smaller groups keep all their locations when the requested maximum +exceeds their size. Every selected value column uses the same sampled locations, so missing values can lower its +finite count. Returned masks and observed group rows still describe membership before sampling. Without `by`, +the option uses the ordinary global sample. + +The ASV grouped-statistics comparisons vary raster size, chunk size, group count and local versus interleaved +membership. They report moments separately from exact median/NMAD calculations, with complete result computation +inside the measured operation. Larger-than-memory tests additionally check Dask worker memory and health. + +### Masks, alignment and chunked execution + +**Group masks retain all eligible locations in each group**, including locations with missing values in a selected +band or column. Subsampling affects the statistics, not these masks. Masks retain their spatial representation and +can be plotted or saved, for example with `first_mask.to_file("elevation_bin.tif")`. + +**Spatial preparation is shared with co-sampling.** Combining rasters uses the calling grid; combining rasters and +points uses the first point dataset's locations. Set `at` explicitly to choose another support. Raster values are +interpolated at points, and `align="reproject"` allows mismatched grids or CRSs to be aligned. Separate point datasets +must share the same ordered locations. Each selected value keeps its own finite count after alignment. + +For arrays already aligned to one another, use the array function directly: + ```{code-cell} ipython3 -def custom_stat(data): - return np.nansum(data > 100) # Count the number of pixels above 100 -rast.get_stats(custom_stat) +# Apply the same elevation grouping without geospatial objects +gu.stats.stats( + {"elevation": rast.data}, + ("mean", "min", "max"), + by={"elevation": rast.data}, + bins={"elevation": elevation_edges}, +) ``` -Passing an inlier mask: +**Dask and multiprocessing use the same aggregation kernels** and return an in-memory table. Dask follows the input +chunks and scheduler. For eager arrays or GeoUtils objects, pass a {class}`~geoutils.multiproc.MultiprocConfig` configuration +from `geoutils.multiproc` to distribute array tiles across workers. Spatial preparation precedes aggregation. + +Use `strategy` to control how groups are combined across chunks: + +| Strategy | How it computes statistics | Useful for | +| --- | --- | --- | +| `"dense"` | Reduce each chunk into an accumulator for every declared group, then combine summaries | Moderate numbers of bins or categories | +| `"sparse"` | Reduce and combine only group IDs present in each chunk | Many zones or sparsely populated group combinations | +| `"groupwise"` | Gather complete groups from intersecting chunks, batching small groups that share chunks | Exact medians, NMAD, quantiles and custom functions | + +**`strategy="auto"` chooses `groupwise` for exact statistics**, including median and NMAD in the default set. For +mergeable statistics such as count, mean, standard deviation or sum, it uses `dense` up to 4096 declared group +combinations and `sparse` above that. These defaults favor fast dense reductions while limiting the size of +intermediate summaries. +The `sparse` strategy uses ordinary NumPy arrays of observed groups; no sparse array dependency is required. + +Exact statistics require a complete group's values to fit in memory. Use `subsample` and `random_state` for a +reproducible sampled estimate. **`subsampling_strategy` controls sample selection separately from aggregation**; +see {ref}`sampling-reproducibility` for its `"topk"` and `"sequential"` options. + +(stats-variograms)= +## Variography + +Use {func}`geoutils.stats.variogram`, {meth}`geoutils.Raster.variogram`, or +{meth}`geoutils.PointCloud.variogram`. + +Variography describes **how differences between values change with spatial separation**. An empirical variogram +groups sampled pairs by distance and estimates their semivariance. A fitted model describes this spatial variability +with a correlation range, a structured variance (partial sill) and an optional nugget. + +Install `geoutils[geostat]` for the optional geostatistical backends. GeoUtils uses SciKit-GStat estimators and models +and returns a {class}`~geoutils.Variogram` with the empirical bins and fitted parameters. + ```{code-cell} ipython3 -inlier_mask = rast > 1500 -rast.get_stats(inlier_mask=inlier_mask) +# Estimate elevation variability over distances up to 5 km using one pair sample +variogram = rast.variogram( + n_pairs=20_000, + n_lags=12, + max_lag=5000, + model="spherical", + random_state=42, +) +variogram.to_dataframe().head() ``` -## Subsampling +```{code-cell} ipython3 +:tags: [hide-input] +:mystnb: +: code_prompt_show: "Show the code for plotting the figure" +: code_prompt_hide: "Hide the code for plotting the figure" -The {func}`~geoutils.Raster.subsample` method allows to efficiently extract a valid random subsample from a raster or a point cloud. It can conveniently -return the output as a point cloud, or as an array. +fig, ax = plt.subplots(figsize=(6, 3)) +variogram.plot(ax=ax) +ax.set(xlabel="Distance (m)", ylabel="Elevation semivariance (m²)") +plt.tight_layout() +``` -The subsample size can be defined either as a fraction of valid values (floating value strictly between 0 and 1), or as a number of samples (integer value -above 1). +**Distance uses coordinate units; semivariance uses squared value units.** This example describes variability in +terrain elevation. To estimate an error variogram, use error measurements or elevation differences on stable terrain. + +### Sampling and fitting + +**Logarithmic distance sampling represents short and long separations efficiently.** `n_pairs` sets the sample size +and `n_lags` sets the number of distance bins. The default uses one sample, with chunk scheduling handled by the +sampling backend. See {ref}`sampling-pairs` for pair selection. + +For an advanced estimate of sampling variability, set `n_runs=3` or more to repeat sampling with shared distance +bins. The result includes the standard error across those independent estimates. This describes sampling +variability; it does not give confidence intervals on fitted parameters. Repetitions run sequentially, while each +sample retains its backend's chunk scheduling. + +Fitting is optional. Use `.fit()` to fit or refit the empirical bins without repeating the sampling: + +```{code-cell} ipython3 +# Fit another model to the retained empirical variogram +refitted = variogram.fit("gaussian") +refitted.model.to_dict() +``` + +Use `model=["spherical", "gaussian"]` to fit a **sum of models at different spatial scales**. +{meth}`~geoutils.Variogram.from_pairs` estimates an empirical variogram from an existing pair sample. The result +stores bin statistics and model parameters, so the pair arrays can be discarded. + +### Reusing models + +A {class}`~geoutils.Variogram` stores model parameters and evaluates **semivariance, covariance and correlation**. +Export it to a dataframe, Xarray dataset or dictionary to retain its bins, counts and fitted model: ```{code-cell} ipython3 -# Subsample 10% of the raster valid values -rast.subsample(subsample=0.1) +# Keep a portable result and evaluate correlation at selected distances +dataset = variogram.to_xarray() +restored = gu.Variogram.from_dict(variogram.to_dict()) +restored.correlation(np.array([0, 100, 1000])) +``` + +Use {meth}`~geoutils.Variogram.from_model` to define a model from known parameters. Convert it to GSTools or GPyTorch +to use the same spatial model in **kriging, random field simulation or Gaussian process calculations**: + +```{code-block} python +known = gu.Variogram.from_model("gaussian", effective_range=500, partial_sill=4, nugget=0.2) +gstools_model = known.to_gstools(dim=2).model + +# GPyTorch receives the nugget separately as observation noise +conversion = known.to_gpytorch() +kernel, noise = conversion.kernel, conversion.noise ``` diff --git a/geoutils/__init__.py b/geoutils/__init__.py index a0d3a60b4..0f920e790 100644 --- a/geoutils/__init__.py +++ b/geoutils/__init__.py @@ -23,10 +23,10 @@ from geoutils import examples, projtools # noqa from geoutils._config import config # noqa -from geoutils.raster import Raster, xr_accessor # noqa isort:skip -from geoutils.raster.xr_accessor import open_raster, RasterAccessor # noqa isort:skip -from geoutils.vector import Vector, open_vector, VectorAccessor # noqa isort:skip -from geoutils.pointcloud import PointCloud, open_pointcloud, PointCloudAccessor # noqa isort:skip +from geoutils.raster import Raster, RasterAccessor, open_raster # noqa isort:skip +from geoutils.vector import Vector, VectorAccessor, open_vector # noqa isort:skip +from geoutils.pointcloud import PointCloud, PointCloudAccessor, open_pointcloud # noqa isort:skip +from geoutils.stats.variography import Variogram # noqa isort:skip # To-be-deprecated from geoutils.raster import Mask # noqa isort:skip diff --git a/geoutils/_dispatch.py b/geoutils/_dispatch.py index 15b827a0e..747eacfb6 100644 --- a/geoutils/_dispatch.py +++ b/geoutils/_dispatch.py @@ -84,6 +84,23 @@ def is_dask_geodataframe(obj: Any) -> bool: ################################################################################# +def get_geo_interface(obj: Any, attr_name: str, accessors: Sequence[str] = ("rst", "vct", "pc")) -> Any: + """ + Return the object or accessor providing a requested geospatial operation, or None if absent. + + Arguments follow get_geo_attr(). Use this when several operations and metadata fields must come from the + same interface, avoiding collisions with native Xarray or Pandas names such as count and shape. + """ + + if hasattr(obj, attr_name): + return obj + for accessor_name in accessors: + accessor = getattr(obj, accessor_name, None) + if accessor is not None and hasattr(accessor, attr_name): + return accessor + return None + + def get_geo_attr(obj: Any, attr_name: str, accessors: Sequence[str] = ("rst", "vct", "pc")) -> Any: """Retrieve an attribute from an object, or one of its accessors.""" @@ -119,6 +136,40 @@ def has_geo_attr(obj: Any, attr_name: str, accessors: Sequence[str] = ("rst", "v return False +# Helpers for recognizing spatial inputs and selecting their GeoUtils interfaces +################################################################################# + + +def _get_raster_interface(obj: Any) -> Any: + """Return the object or rst accessor providing raster coordinate conversion, or None if absent.""" + + return get_geo_interface(obj, "ij2xy", accessors=("rst",)) + + +def _get_pointcloud_interface(obj: Any) -> Any: + """Return the object or pc accessor providing ordered point coordinate comparison, or None if absent.""" + + return get_geo_interface(obj, "georeferenced_coords_equal", accessors=("pc",)) + + +def _is_raster(obj: Any) -> bool: + """Check for raster operations on an object or its rst accessor, without requiring a specific class.""" + + return _get_raster_interface(obj) is not None + + +def _is_pointcloud(obj: Any) -> bool: + """Check for point cloud operations on an object or its pc accessor, without requiring a specific class.""" + + return _get_pointcloud_interface(obj) is not None + + +def _is_vector(obj: Any) -> bool: + """Check for rasterization on an object or its vct accessor, including point clouds with vector operations.""" + + return has_geo_attr(obj, "rasterize", accessors=("vct",)) + + # Level 0 checks: directly on user input ######################################## @@ -661,7 +712,8 @@ def _check_match_grid( and has_geo_attr(ref, "crs") and isinstance(get_geo_attr(ref, "transform"), rio.Affine) ): - dst_shape = get_geo_attr(ref, "shape") + # Match the spatial grid even when a native Xarray reference includes a band dimension + dst_shape = get_geo_attr(ref, "shape")[-2:] dst_transform = get_geo_attr(ref, "transform") dst_crs = get_geo_attr(ref, "crs") diff --git a/geoutils/filters.py b/geoutils/filters.py index 36f43dd48..30613cbe9 100644 --- a/geoutils/filters.py +++ b/geoutils/filters.py @@ -17,15 +17,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Filters class to remove outliers and reduce noise in rasters. -""" +"""Filters to remove outliers and reduce noise in rasters.""" from __future__ import annotations import math from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, overload import numpy as np import scipy @@ -162,40 +160,22 @@ def _filter_base( if np.ma.isMaskedArray(array): array = array.filled(np.nan) - # With new SciPy, just use vectorized version - filter_map: dict[str, Callable[..., NDArrayNum]] - if Version(scipy.__version__) > Version("1.16.0"): - filter_map = { - "gaussian": gaussian_filter, - "median": lambda arr, size=size, **_: generic_filter_scipy( - arr, np.nanmedian, size=size, mode="constant", cval=np.nan - ), - "mean": lambda arr, size=size, **_: generic_filter_scipy( - arr, np.nanmean, size=size, mode="constant", cval=np.nan - ), - "max": lambda arr, size=size, **_: generic_filter_scipy( - arr, np.nanmax, size=size, mode="constant", cval=np.nan - ), - "min": lambda arr, size=size, **_: generic_filter_scipy( - arr, np.nanmin, size=size, mode="constant", cval=np.nan - ), - "distance": distance_filter, - } - # With old SciPy, maintain speed with tricks from older custom filters - else: - filter_map = { - "gaussian": gaussian_filter, - "median": lambda arr, size=size, **_: median_filter(arr, size=size), - "mean": lambda arr, size=size, **_: mean_filter(arr, size=size), - "max": lambda arr, size=size, **_: max_filter(arr, size=size), - "min": lambda arr, size=size, **_: min_filter(arr, size=size), - "distance": distance_filter, - } + # Use the named implementations so optimized filters and their engines remain available on every SciPy version + filter_map: dict[str, Callable[..., Any]] = { + "gaussian": gaussian_filter, + "median": median_filter, + "mean": mean_filter, + "max": max_filter, + "min": min_filter, + "distance": distance_filter, + } if isinstance(method, str): if method not in filter_map: raise ValueError(f"Unsupported filter method '{method}'. Available: {list(filter_map)}") func = filter_map[method] + if method in {"median", "mean", "max", "min"}: + kwargs["size"] = size elif callable(method): func = method else: @@ -242,7 +222,6 @@ def _multiproc_filter( size: int = 3, **kwargs: Any, ) -> Raster: - # Get depth of overlap depth = _overlap_depth_for_filter(method, size=size, **kwargs) @@ -290,7 +269,7 @@ def _filter( kwargs = {} if method == "gaussian": kwargs.update({"sigma": sigma}) - if method == "median": + if isinstance(method, str) and method in {"gaussian", "median", "mean", "max", "min", "distance"}: kwargs.update({"engine": engine}) if method == "distance": kwargs.update({"outlier_threshold": outlier_threshold}) @@ -306,30 +285,80 @@ def _filter( return source_raster.copy(new_array=array) -def gaussian_filter(array: NDArrayNum, sigma: float = 1, **kwargs: Any) -> NDArrayNum: +def _validate_filter_engine(engine: str) -> None: + """Check that a filter engine is supported and load Numba only when requested.""" + + if engine not in {"scipy", "numba"}: + raise ValueError('Engine must be "scipy" or "numba".') + if engine == "numba": + import_optional("numba") + + +def gaussian_filter( + array: NDArrayNum, sigma: float = 1, engine: Literal["scipy", "numba"] = "scipy", **kwargs: Any +) -> NDArrayNum: """ Apply a Gaussian filter to a raster that may contain NaNs. N.B: kernel_size is set automatically based on sigma. + For 3D arrays, each 2D band is filtered independently. + :param array: The input array to be filtered. :param sigma: The sigma of the Gaussian kernel + :param engine: Filtering engine to use, either "scipy" or "numba". :returns: The filtered array (same shape as input) """ - if array.ndim == 1: + if array.ndim not in [2, 3]: raise ValueError("Gaussian filter can't be applied to 1D arrays.") + _validate_filter_engine(engine) + if sigma < 0: + raise ValueError("Sigma must be non-negative.") + + # Add a band dimension to 2D inputs so both shapes follow the same filtering path + squeeze = array.ndim == 2 + bands = array[np.newaxis] if squeeze else array # Boolean mask: True where NaN - mask = np.isnan(array) + mask = np.isnan(bands) mask_f = (~mask).astype(float) # Replace NaNs with 0 - arr_filled = np.where(mask, 0, array) + arr_filled = np.where(mask, 0, bands) # Apply gaussian filter to values and mask - filtered = scipy.ndimage.gaussian_filter(arr_filled, sigma, mode="constant", cval=0, **kwargs) - normalization = scipy.ndimage.gaussian_filter(mask_f, sigma, mode="constant", cval=0, **kwargs) + if engine == "scipy": + if squeeze: + filtered = scipy.ndimage.gaussian_filter(arr_filled[0], sigma, mode="constant", cval=0, **kwargs)[ + np.newaxis + ] + normalization = scipy.ndimage.gaussian_filter(mask_f[0], sigma, mode="constant", cval=0, **kwargs)[ + np.newaxis + ] + else: + filtered = np.stack( + [scipy.ndimage.gaussian_filter(band, sigma, mode="constant", cval=0, **kwargs) for band in arr_filled] + ) + normalization = np.stack( + [scipy.ndimage.gaussian_filter(band, sigma, mode="constant", cval=0, **kwargs) for band in mask_f] + ) + else: + truncate = float(kwargs.pop("truncate", 4.0)) + if kwargs: + raise ValueError("The Numba engine only supports the 'truncate' Gaussian filter option.") + + # Build SciPy's one-dimensional Gaussian kernel and convolve along each spatial axis + radius = int(truncate * sigma + 0.5) + coordinates = np.arange(-radius, radius + 1, dtype=float) + kernel_1d = np.exp(-0.5 / sigma**2 * coordinates**2) if sigma > 0 else np.ones(1) + kernel_1d /= kernel_1d.sum() + horizontal_kernel = kernel_1d.reshape(1, 1, -1) + vertical_kernel = kernel_1d.reshape(1, -1, 1) + filtered = convolution(arr_filled, horizontal_kernel, engine="numba", cval=0)[:, 0] + filtered = convolution(filtered, vertical_kernel, engine="numba", cval=0)[:, 0] + normalization = convolution(mask_f, horizontal_kernel, engine="numba", cval=0)[:, 0] + normalization = convolution(normalization, vertical_kernel, engine="numba", cval=0)[:, 0] # Avoid division by zero with np.errstate(invalid="ignore", divide="ignore"): @@ -338,7 +367,7 @@ def gaussian_filter(array: NDArrayNum, sigma: float = 1, **kwargs: Any) -> NDArr # Where normalization is zero, set result to NaN filtered[normalization == 0] = np.nan - return filtered + return filtered[0] if squeeze else filtered @jit(nopython=True, parallel=True) @@ -393,6 +422,7 @@ def median_filter(array: NDArrayNum, size: int, engine: Literal["scipy", "numba" if size % 2 == 0: raise ValueError("`size` must be odd.") + _validate_filter_engine(engine) if array.ndim == 2: return _apply_median_filter_2d(array, size, engine) @@ -420,87 +450,246 @@ def _apply_median_filter_2d( return np.where(nans, array, median_vals) else: - import_optional("numba") median_vals = median_filter_numba(array, size) return np.where(nans, array, median_vals) -def mean_filter(array: NDArrayNum, size: int = 5) -> NDArrayNum: +@overload +def mean_filter( + array: NDArrayNum, + size: int = 5, + *, + kernel_shape: Literal["square", "circular"] = "square", + engine: Literal["scipy", "numba"] = "scipy", + preserve_nodata: bool = True, + boundless: bool = True, + return_counts: Literal[False] = False, +) -> NDArrayNum: ... + + +@overload +def mean_filter( + array: NDArrayNum, + size: int = 5, + *, + kernel_shape: Literal["square", "circular"] = "square", + engine: Literal["scipy", "numba"] = "scipy", + preserve_nodata: bool = True, + boundless: bool = True, + return_counts: Literal[True], +) -> tuple[NDArrayNum, NDArrayNum, int]: ... + + +def mean_filter( + array: NDArrayNum, + size: int = 5, + *, + kernel_shape: Literal["square", "circular"] = "square", + engine: Literal["scipy", "numba"] = "scipy", + preserve_nodata: bool = True, + boundless: bool = True, + return_counts: bool = False, +) -> NDArrayNum | tuple[NDArrayNum, NDArrayNum, int]: """ Apply a mean filter to a 2D array that may contain NaNs. - :param array: 2D input array - :param size: size of the square kernel - :no_data: no data value - :return: filtered array with same shape + :param array: 2D input array. + :param size: Size of the square or circular kernel. + :param kernel_shape: Shape of the kernel, either "square" or "circular". + :param engine: Filtering engine to use, either "scipy" or "numba". + :param preserve_nodata: Whether missing input cells remain missing in the output. + :param boundless: Whether to compute partial windows along array edges. + :param return_counts: Whether to also return finite cell counts and the total kernel cell count. + :returns: Filtered array, optionally with the finite counts and total kernel cell count. """ if array.ndim != 2: raise ValueError(f"Invalid array shape {array.shape}, expected 2D.") + _validate_filter_engine(engine) # Mask nodata values - nans = np.isnan(array) - mask = ~np.isnan(array) - array_filled = np.where(mask, array, 0) - # Compute sum over the kernel - sum_vals = scipy.ndimage.uniform_filter(array_filled, size=size, mode="constant", cval=0.0) - # Count of valid (non-nodata) pixels in the kernel - count_vals = scipy.ndimage.uniform_filter(mask.astype(float), size=size, mode="constant", cval=0.0) + valid = np.isfinite(array) + array_filled = np.where(valid, array, 0) + + # Define the cells included in the requested kernel + if kernel_shape == "square": + kernel = None + kernel_count = size**2 + elif kernel_shape == "circular": + kernel = _create_circular_mask((size, size)).astype(float) + kernel_count = int(np.count_nonzero(kernel)) + else: + raise ValueError('Kernel shape should be "square" or "circular".') + + # Keep the optimized SciPy implementation used by the existing square mean filter + if engine == "scipy" and kernel_shape == "square": + sum_vals = scipy.ndimage.uniform_filter(array_filled, size=size, mode="constant", cval=0.0) + count_vals = scipy.ndimage.uniform_filter(valid.astype(float), size=size, mode="constant", cval=0.0) + finite_counts = np.rint(count_vals * kernel_count) if return_counts else None + else: + if kernel is None: + kernel = np.ones((size, size), dtype=float) + sum_vals = convolution(array_filled[np.newaxis], kernel[np.newaxis], engine=engine, cval=0)[0, 0] + finite_counts = convolution(valid.astype(float)[np.newaxis], kernel[np.newaxis], engine=engine, cval=0)[0, 0] + count_vals = finite_counts with np.errstate(invalid="ignore", divide="ignore"): mean_vals = sum_vals / count_vals - return np.where(nans, array, mean_vals) - - -def min_filter(array: NDArrayNum, size: int = 5, **kwargs: Any) -> NDArrayNum: + # Exclude windows that extend beyond the image when complete patches are required + if not boundless: + before = (size - 1) // 2 + after = size // 2 + if before > 0: + mean_vals[:before, :] = np.nan + mean_vals[:, :before] = np.nan + if finite_counts is not None: + finite_counts[:before, :] = np.nan + finite_counts[:, :before] = np.nan + if after > 0: + mean_vals[-after:, :] = np.nan + mean_vals[:, -after:] = np.nan + if finite_counts is not None: + finite_counts[-after:, :] = np.nan + finite_counts[:, -after:] = np.nan + + if preserve_nodata: + mean_vals = np.where(valid, mean_vals, array) + + if return_counts: + assert finite_counts is not None + return mean_vals, finite_counts, kernel_count + return mean_vals + + +@jit(nopython=True, parallel=True, cache=True) +def _minmax_filter_numba(array: NDArrayNum, size: int, fill_value: float, find_maximum: bool) -> NDArrayNum: + """Apply a compiled minimum or maximum filter independently to stacked 2D arrays.""" + + before = size // 2 + padded = np.full( + (array.shape[0], array.shape[1] + size - 1, array.shape[2] + size - 1), + fill_value, + dtype=array.dtype, + ) + padded[ + :, + before : before + array.shape[1], + before : before + array.shape[2], + ] = array + output = np.empty_like(array) + + for band in prange(array.shape[0]): + for row in range(array.shape[1]): + for col in range(array.shape[2]): + result = fill_value + for window_row in range(size): + for window_col in range(size): + value = padded[band, row + window_row, col + window_col] + if (find_maximum and value > result) or (not find_maximum and value < result): + result = value + output[band, row, col] = result + + return output + + +def min_filter( + array: NDArrayNum, size: int = 5, engine: Literal["scipy", "numba"] = "scipy", **kwargs: Any +) -> NDArrayNum: """ - Apply a minimum filter to a raster that may contain NaNs, using scipy's implementation. + Apply a minimum filter to a raster that may contain NaNs. + + For 3D arrays, each 2D band is filtered independently. :param array: The input array to be filtered. :param size: the shape that is taken from the input array, at every element position, to define the input to the filter function + :param engine: Filtering engine to use, either "scipy" or "numba". :returns: The filtered array (same shape as input). """ # Check that array dimension is 2 or 3 if array.ndim not in [2, 3]: raise ValueError(f"Invalid array shape given: {array.shape}. Expected 2D or 3D array.") + _validate_filter_engine(engine) nans = np.isnan(array) # We replace temporarily NaNs by infinite values during filtering to avoid spreading NaNs array_nans_replaced = np.where(nans, np.inf, array) - array_nans_replaced_f = scipy.ndimage.minimum_filter( - array_nans_replaced, size=size, mode="constant", cval=np.inf, **kwargs - ) + if engine == "scipy": + if array.ndim == 2: + array_nans_replaced_f = scipy.ndimage.minimum_filter( + array_nans_replaced, size=size, mode="constant", cval=np.inf, **kwargs + ) + else: + array_nans_replaced_f = np.stack( + [ + scipy.ndimage.minimum_filter(band, size=size, mode="constant", cval=np.inf, **kwargs) + for band in array_nans_replaced + ] + ) + else: + if kwargs: + raise ValueError("The Numba engine does not support additional minimum filter options.") + bands = array_nans_replaced[np.newaxis] if array.ndim == 2 else array_nans_replaced + array_nans_replaced_f = _minmax_filter_numba(bands, size, np.inf, False) + if array.ndim == 2: + array_nans_replaced_f = array_nans_replaced_f[0] # In the end, we want the filtered array without infinite values, so we put back NaNs return np.where(nans, array, array_nans_replaced_f) -def max_filter(array: NDArrayNum, size: int = 5, **kwargs: Any) -> NDArrayNum: +def max_filter( + array: NDArrayNum, size: int = 5, engine: Literal["scipy", "numba"] = "scipy", **kwargs: Any +) -> NDArrayNum: """ - Apply a maximum filter to a raster that may contain NaNs, using scipy's implementation. + Apply a maximum filter to a raster that may contain NaNs. + + For 3D arrays, each 2D band is filtered independently. :param array: the input array to be filtered. :param size: the shape that is taken from the input array, at every element position, to define the input to the filter function + :param engine: Filtering engine to use, either "scipy" or "numba". :returns: the filtered array (same shape as input). """ # Check that array dimension is 2 or 3 if array.ndim not in [2, 3]: raise ValueError(f"Invalid array shape given: {array.shape}. Expected 2D or 3D array.") + _validate_filter_engine(engine) nans = np.isnan(array) # We replace temporarily NaNs by negative infinite values during filtering to avoid spreading NaNs array_nans_replaced = np.where(nans, -np.inf, array) - array_nans_replaced_f = scipy.ndimage.maximum_filter( - array_nans_replaced, size=size, mode="constant", cval=-np.inf, **kwargs - ) + if engine == "scipy": + if array.ndim == 2: + array_nans_replaced_f = scipy.ndimage.maximum_filter( + array_nans_replaced, size=size, mode="constant", cval=-np.inf, **kwargs + ) + else: + array_nans_replaced_f = np.stack( + [ + scipy.ndimage.maximum_filter(band, size=size, mode="constant", cval=-np.inf, **kwargs) + for band in array_nans_replaced + ] + ) + else: + if kwargs: + raise ValueError("The Numba engine does not support additional maximum filter options.") + bands = array_nans_replaced[np.newaxis] if array.ndim == 2 else array_nans_replaced + array_nans_replaced_f = _minmax_filter_numba(bands, size, -np.inf, True) + if array.ndim == 2: + array_nans_replaced_f = array_nans_replaced_f[0] # In the end we want the filtered array without negative infinite values, so we put back NaNs return np.where(nans, array, array_nans_replaced_f) -def distance_filter(array: NDArrayNum, sigma: float = 5, outlier_threshold: float = 2) -> NDArrayNum: +def distance_filter( + array: NDArrayNum, + sigma: float = 5, + outlier_threshold: float = 2, + engine: Literal["scipy", "numba"] = "scipy", +) -> NDArrayNum: """ Filter out pixels whose value is distant more than a set threshold from the average value of all neighbor \ pixels within a given radius. @@ -510,6 +699,7 @@ def distance_filter(array: NDArrayNum, sigma: float = 5, outlier_threshold: floa :param array: Input array to be filtered. :param sigma: Radius in which the average value is calculated (for Gaussian filter, this is sigma). :param outlier_threshold: the minimum difference abs(array - mean) for a pixel to be considered an outlier. + :param engine: Filtering engine to use, either "scipy" or "numba". :returns: the filtered array (same shape as input) """ @@ -517,8 +707,8 @@ def distance_filter(array: NDArrayNum, sigma: float = 5, outlier_threshold: floa valid_mask = np.isfinite(array) # Smooth both the data and the valid mask - smoothed = gaussian_filter(np.nan_to_num(array, nan=0.0), sigma=sigma) - normalization = gaussian_filter(valid_mask.astype(float), sigma=sigma) + smoothed = gaussian_filter(np.nan_to_num(array, nan=0.0), sigma=sigma, engine=engine) + normalization = gaussian_filter(valid_mask.astype(float), sigma=sigma, engine=engine) # Avoid division by zero with np.errstate(invalid="ignore", divide="ignore"): @@ -552,3 +742,114 @@ def generic_filter( if array.ndim not in [2, 3]: raise ValueError(f"Invalid array shape given: {array.shape}. Expected 2D or 3D array.") return filter_function(array, **kwargs) + + +######################################### +# STACKED CONVOLUTION +######################################### + + +def _create_circular_mask( + shape: tuple[int, int], center: tuple[int, int] | None = None, radius: float | None = None +) -> NDArrayBool: + """Create a circular kernel, using the array centre and half width by default.""" + + # Use the array center and its nearest edge to choose a fully contained default circle + w, h = shape + + if center is None: + center = (int(w / 2), int(h / 2)) + if radius is None: + radius = min(center[0], center[1], w - center[0], h - center[1]) + + # Select cells strictly inside the radius to preserve the patch kernel boundary convention + Y, X = np.ogrid[:w, :h] + dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2) + mask = dist_from_center < radius + + return mask + + +@jit(nopython=True, parallel=True, cache=True) +def _convolution_numba(imgs: NDArrayNum, filters: NDArrayNum, output: NDArrayNum) -> NDArrayNum: + """Accumulate convolution over image and kernel stacks using compiled loops.""" + + # Read image and kernel dimensions for the output loops + n_N, N1, N2 = imgs.shape + n_M, M1, M2 = filters.shape + + # Restrict windows to complete footprints within the padded input + row_range = N1 - M1 + 1 + col_range = N2 - M2 + 1 + + # Accumulate each output pixel from its complete input window + for ii in range(n_N): + for rr in prange(row_range): + for cc in range(col_range): + for m1 in range(M1): + for m2 in range(M2): + for ff in range(n_M): + imgval = imgs[ii, rr + m1, cc + m2] + + # Reverse both kernel axes to compute convolution + filterval = filters[ff, M1 - 1 - m1, M2 - 1 - m2] + output[ii, ff, rr, cc] += imgval * filterval + + return output + + +def convolution( + imgs: NDArrayNum, + filters: NDArrayNum, + engine: Literal["scipy", "numba"] = "scipy", + cval: float = np.nan, +) -> NDArrayNum: + """ + Convolution on a number n_N of 2D images of size N1 x N2 using a number of kernels n_M of sizes M1 x M2, using + either scipy.ndimage.convolve or accelerated numba loops. + Note that the indexes on n_M and n_N correspond to first axes on the array to speed up computations (prefetching). + Inspired by: https://laurentperrinet.github.io/sciblog/posts/2017-09-20-the-fastest-2d-convolution-in-the-world.html + + :param imgs: Input array of size (n_N, N1, N2) with n_N images of size N1 x N2 + :param filters: Input array of filters of size (n_M, M1, M2) with n_M filters of size M1 x M2 + :param engine: Filtering engine to use, either "scipy" or "numba". + :param cval: Value used outside the image boundaries. + + :return: Filled array of outputs of size (n_N, n_M, N1, N2) + """ + + # Validate image and kernel stacks before allocating the output + imgs = np.asarray(imgs, dtype=float) + filters = np.asarray(filters, dtype=float) + if imgs.ndim != 3 or filters.ndim != 3 or any(size < 1 for size in filters.shape): + raise ValueError("Images and filters must be 3D stacks with non-empty kernels.") + _validate_filter_engine(engine) + + # Initialize output array according to input shapes + n_N, N1, N2 = imgs.shape + n_M, M1, M2 = filters.shape + output = np.zeros((n_N, n_M, N1, N2)) + + # Apply each kernel to each image, preserving the existing NaN padding outside the image + if engine == "scipy": + for image_index in range(n_N): + for filter_index in range(n_M): + output[image_index, filter_index] = scipy.ndimage.convolve( + imgs[image_index], filters[filter_index], mode="constant", cval=cval + ) + else: + # Pad asymmetrically for even kernel widths so compiled loops match SciPy's kernel origin + half_M1 = int((M1 - 1) / 2) + half_M2 = int((M2 - 1) / 2) + imgs_pad = np.pad( + imgs, + pad_width=((0, 0), (half_M1, M1 // 2), (half_M2, M2 // 2)), + constant_values=cval, + ) + output = _convolution_numba( + imgs=imgs_pad, + filters=filters, + output=output, + ) + + return output diff --git a/geoutils/interface/gridding.py b/geoutils/interface/gridding.py index c5927a1a5..d57a48f2c 100644 --- a/geoutils/interface/gridding.py +++ b/geoutils/interface/gridding.py @@ -1186,6 +1186,7 @@ def _grid_pointcloud_to_raster( dist_nodata_pixel: float = 1.0, nodata: int | float = -9999, *, + data_column: str | None = None, distance_power: float = 2.0, min_points: int = 1, chunksizes: tuple[int, int] | None = None, @@ -1196,13 +1197,32 @@ def _grid_pointcloud_to_raster( nodata_propagation: NodataPropagation = "gdal", gridding_func: GridPointCloudCallable = _grid_pointcloud, ) -> Any: - """Grid a point cloud to a raster with eager, Dask, or Multiprocessing backends.""" + """ + Grid a point cloud to a raster with eager, Dask, or Multiprocessing backends. + + A Dask reference selects lazy output even when the point source is eager. Its spatial chunks are reused unless + chunksizes is supplied, following the same reference-grid behavior as rasterization. An explicit data_column + selects values without copying the source or changing its active column. + """ + + # Resolve the value column from metadata so file-backed sources stay available for bounded worker reads + if data_column is not None and ( + not isinstance(data_column, str) or data_column not in get_geo_attr(source_pointcloud, "columns") + ): + raise ValueError("Argument ``data_column`` must name an existing point column.") + data_column_name = get_geo_attr(source_pointcloud, "data_column") if data_column is None else data_column + + # Follow a lazy output reference without converting an eager point cloud to a Dask dataframe + ref_chunks = get_geo_attr(ref, "_chunks") if ref is not None and has_geo_attr(ref, "_chunks") else None + if ref_chunks is not None: + ref_chunks = ref_chunks[-2:] + dask = True # A single operation must have one owner for scheduling and memory management if dask and mp_config is not None: raise ValueError( "Cannot use Multiprocessing and Dask simultaneously. To use Dask, remove mp_config. " - "To use Multiprocessing, use an eager PointCloud object." + "To use Multiprocessing, use an eager PointCloud and an unchunked raster reference." ) if is_dask_dataframe(_source_dataframe(source_pointcloud)) and mp_config is not None: @@ -1247,7 +1267,7 @@ def _grid_pointcloud_to_raster( array = _grid_pointcloud_block_from_source( source_pointcloud=source_pointcloud, geogrid=dst_geogrid, - data_column_name=get_geo_attr(source_pointcloud, "data_column"), + data_column_name=data_column_name, gridding_func=gridding_func, **kwargs, ) @@ -1258,7 +1278,6 @@ def _grid_pointcloud_to_raster( if mp_config is not None: chunksizes = _split_chunk_size(mp_config.chunks) else: - ref_chunks = get_geo_attr(ref, "_chunks") if ref is not None and has_geo_attr(ref, "_chunks") else None chunksizes = ref_chunks if ref_chunks is not None else (1024, 1024) assert chunksizes is not None @@ -1273,7 +1292,7 @@ def _grid_pointcloud_to_raster( source_pointcloud=source_pointcloud, dst_geotiling=dst_geotiling, dst_block_geogrids=dst_block_geogrids, - data_column_name=get_geo_attr(source_pointcloud, "data_column"), + data_column_name=data_column_name, gridding_func=gridding_func, **kwargs, ) @@ -1294,7 +1313,7 @@ def _grid_pointcloud_to_raster( source_pointcloud=source_pointcloud, dst_geotiling=dst_geotiling, dst_block_geogrids=dst_block_geogrids, - data_column_name=get_geo_attr(source_pointcloud, "data_column"), + data_column_name=data_column_name, mp_config=mp_config, file_metadata=file_metadata, gridding_func=gridding_func, diff --git a/geoutils/interface/interpolation.py b/geoutils/interface/interpolation.py index 233f909b9..e245cbf5e 100644 --- a/geoutils/interface/interpolation.py +++ b/geoutils/interface/interpolation.py @@ -36,11 +36,14 @@ from geoutils._misc import import_optional from geoutils._typing import DTypeLike, NDArrayBool, NDArrayNum, Number from geoutils.interface._nodata import NodataPropagation, _validate_nodata_propagation -from geoutils.multiproc import MultiprocConfig, compute_tiling +from geoutils.multiproc import MultiprocConfig from geoutils.multiproc.chunked import cached_cumsum, normalize_chunks +from geoutils.multiproc.mparray import block_bounds_from_chunks from geoutils.projtools import reproject_from_latlon from geoutils.raster.referencing import _bounds, _coords, _outside_bounds, _res, _xy2ij +InterpolationMethod = Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] + method_to_order = {"nearest": 0, "linear": 1, "cubic": 3, "quintic": 5, "slinear": 1, "pchip": 3, "splinef2d": 3} if TYPE_CHECKING: @@ -49,10 +52,15 @@ from geoutils.raster.raster import Raster -def _interp_output_dtype(dtype: DTypeLike) -> DTypeLike: - """Return an interpolation dtype that can represent NaNs.""" +def _interp_output_dtype(dtype: DTypeLike, *, validity_only: bool = False) -> DTypeLike: + """ + Return an interpolation dtype that can represent NaNs. + + Validity-only interpolation uses float32 regardless of the original values, following _interp_points(). + """ - return np.float32 if np.issubdtype(dtype, np.integer) else dtype + # Promote booleans too so missing interpolated values do not become True + return np.float32 if validity_only or np.issubdtype(dtype, np.integer) or np.issubdtype(dtype, np.bool_) else dtype def _destination_pixel_indices( @@ -252,7 +260,7 @@ def _interpn_interpolator( fill_value: Number = np.nan, bounds_error: bool = False, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] = None, + method: InterpolationMethod = None, ) -> Callable[[tuple[NDArrayNum, NDArrayNum]], NDArrayNum]: """ Create SciPy interpolator with nodata spreading. Default method is linear and default spreading is at distance of @@ -305,7 +313,6 @@ def _interpn_interpolator( # For the RegularGridInterpolator if method in RegularGridInterpolator._ALL_METHODS: - # We create the classic interpolator interp = RegularGridInterpolator( points, values, method=method, bounds_error=bounds_error, fill_value=fill_value @@ -313,7 +320,6 @@ def _interpn_interpolator( # We create a new interpolator callable that propagates nodata as defined above def regulargrid_interpolator_with_nan(xi: tuple[NDArrayNum, NDArrayNum]) -> NDArrayNum: - # Get results results = interp(xi) # Get invalids @@ -326,13 +332,11 @@ def regulargrid_interpolator_with_nan(xi: tuple[NDArrayNum, NDArrayNum]) -> NDAr # For the RectBivariateSpline else: - # The coordinates must be in ascending order, which requires flipping the array too (more costly) interp = RectBivariateSpline(np.flip(points[0]), points[1], np.flip(values[:], axis=0)) # We create a new interpolator callable that propagates nodata as defined above, and supports fill_value def rectbivariate_interpolator_with_fillvalue(xi: tuple[NDArrayNum, NDArrayNum]) -> NDArrayNum: - # Get invalids invalids = interp_mask(xi) @@ -421,13 +425,15 @@ def _interp_points_base( transform: rio.transform.Affine, points: tuple[Number, Number] | tuple[NDArrayNum, NDArrayNum], area_or_point: Literal["Area", "Point"] | None = None, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] | None = None, + method: InterpolationMethod | None = None, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, shift_area_or_point: bool | None = None, force_scipy_function: Literal["map_coordinates", "interpn"] | None = None, nodata_propagation: NodataPropagation = "gdal", *, return_interpolator: Literal[False] = False, + array_indices: tuple[NDArrayNum, NDArrayNum] | None = None, + _validity_only: bool = False, **kwargs: Any, ) -> NDArrayNum: ... @@ -438,13 +444,15 @@ def _interp_points_base( transform: rio.transform.Affine, points: tuple[Number, Number] | tuple[NDArrayNum, NDArrayNum], area_or_point: Literal["Area", "Point"] | None = None, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] | None = None, + method: InterpolationMethod | None = None, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, shift_area_or_point: bool | None = None, force_scipy_function: Literal["map_coordinates", "interpn"] | None = None, nodata_propagation: NodataPropagation = "gdal", *, return_interpolator: Literal[True], + array_indices: tuple[NDArrayNum, NDArrayNum] | None = None, + _validity_only: bool = False, **kwargs: Any, ) -> Callable[[tuple[NDArrayNum, NDArrayNum]], NDArrayNum]: ... @@ -455,13 +463,15 @@ def _interp_points_base( transform: rio.transform.Affine, points: tuple[Number, Number] | tuple[NDArrayNum, NDArrayNum], area_or_point: Literal["Area", "Point"] | None = None, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] | None = None, + method: InterpolationMethod | None = None, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, shift_area_or_point: bool | None = None, force_scipy_function: Literal["map_coordinates", "interpn"] | None = None, nodata_propagation: NodataPropagation = "gdal", *, return_interpolator: bool = False, + array_indices: tuple[NDArrayNum, NDArrayNum] | None = None, + _validity_only: bool = False, **kwargs: Any, ) -> NDArrayNum | Callable[[tuple[NDArrayNum, NDArrayNum]], NDArrayNum]: ... @@ -471,19 +481,34 @@ def _interp_points_base( transform: rio.transform.Affine, points: tuple[Number, Number] | tuple[NDArrayNum, NDArrayNum] | None, area_or_point: Literal["Area", "Point"] | None = None, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] | None = None, + method: InterpolationMethod | None = None, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, shift_area_or_point: bool | None = None, force_scipy_function: Literal["map_coordinates", "interpn"] | None = None, nodata_propagation: NodataPropagation = "gdal", return_interpolator: bool = False, + array_indices: tuple[NDArrayNum, NDArrayNum] | None = None, + _validity_only: bool = False, **kwargs: Any, ) -> NDArrayNum | Callable[[tuple[NDArrayNum, NDArrayNum]], NDArrayNum]: + """ + Interpolate a raster at point coordinates. + + This internal function can optionally reuse global pixel indices to work on chunks. + The private validity-only option follows _interp_points(); its conversion happens within this loaded array. + """ # If interpolation method undefined, default to the global system config if method is None: method = config["interpolation_method"] + # Convert availability within the current block, preserving masked cells without copying the source values + if _validity_only: + finite = np.isfinite(array) + if np.ma.isMaskedArray(finite): + finite = finite.filled(False) + array = np.where(finite, np.float32(1), np.float32(np.nan)) + # If array is not a floating dtype (to support NaNs), convert dtype if not np.issubdtype(array.dtype, np.floating): array = array.astype(np.float32) @@ -500,13 +525,16 @@ def interpolate_nearest_or_linear(x: NDArrayNum, y: NDArrayNum) -> NDArrayNum: """Interpolate point coordinates with the shared nearest or linear policy.""" # Convert georeferenced coordinates to array indices before applying the common numeric kernel - i, j = _xy2ij( - x, - y, - transform=transform, - area_or_point=area_or_point, - shift_area_or_point=shift_area_or_point, - ) + if array_indices is None: + i, j = _xy2ij( + x, + y, + transform=transform, + area_or_point=area_or_point, + shift_area_or_point=shift_area_or_point, + ) + else: + i, j = array_indices return _interpolate_array_band( array=array, src_rows=i, @@ -595,7 +623,6 @@ def point_interpolator(xi: tuple[NDArrayNum, NDArrayNum]) -> NDArrayNum: return scipy_interpolator else: rpoints = scipy_interpolator((y, x)) # type: ignore - return rpoints @@ -615,7 +642,6 @@ def _get_interp_indices_per_block( interp_y: NDArrayNum, starts: list[tuple[int, ...]], num_chunks: tuple[int, int], - chunksize: tuple[int, int], xres: float, yres: float, left: float, @@ -625,15 +651,20 @@ def _get_interp_indices_per_block( # The argument "starts" contains the list of chunk first X/Y index for the full array, plus the last index ny, nx = num_chunks - y_chunksize, x_chunksize = chunksize y_starts, x_starts = starts # We use one bucket per block, assuming a flattened blocks shape ind_per_block = [[] for _ in range(ny * nx)] for i, (x, y) in enumerate(zip(interp_x, interp_y)): - # Because it is a regular grid, we know exactly in which block ID the coordinate will fall - xb = int(np.floor((x - left) / (xres * x_chunksize))) - yb = int(np.floor((top - y) / (yres * y_chunksize))) + # Use actual chunk boundaries because overlap can merge small edge chunks + xb = int(np.searchsorted(x_starts, (x - left) / xres, side="right") - 1) + yb = int(np.searchsorted(y_starts, (top - y) / yres, side="right") - 1) + + # Assign outer half pixels to the first block, matching the interpolation kernel's finite support + if left - xres / 2 <= x < left: + xb = 0 + if top < y <= top + yres / 2: + yb = 0 if 0 <= xb < nx and 0 <= yb < ny: ind_per_block[yb * nx + xb].append(i) @@ -696,6 +727,13 @@ def _dask_interp_points( # Convert input to 2D array points_arr = np.vstack((points[0], points[1])) + src_rows, src_cols = _xy2ij( + points[0], + points[1], + transform=transform, + area_or_point=kwargs["area_or_point"], + shift_area_or_point=kwargs["shift_area_or_point"], + ) # Map depth of overlap required for each interpolation method depth = method_to_order[kwargs["method"]] + 1 # The overlap size is the order + 1 @@ -704,12 +742,11 @@ def _dask_interp_points( left, top = bounds.left, bounds.top # Expand dask array for overlapping computations - chunksize = darr.chunksize expanded = da.overlap.overlap(darr, depth=depth, boundary="nearest") - # Get starting 2D index for each chunk of the full array - # (mirroring what is done in block_id of dask.array.map_blocks) - starts = [cached_cumsum(c, initial_zero=True) for c in darr.chunks] + # Recover core chunk boundaries after any automatic merging required for overlap + core_chunks = [tuple(size - 2 * depth for size in axis) for axis in expanded.chunks] + starts = [cached_cumsum(axis, initial_zero=True) for axis in core_chunks] num_chunks = expanded.numblocks # Get samples indices per blocks @@ -718,7 +755,6 @@ def _dask_interp_points( points_arr[1, :], starts, num_chunks, - chunksize, res[0], res[1], left, @@ -742,9 +778,20 @@ def _dask_interp_points( # Compute values delayed used = [i for i in range(len(blocks)) if len(ind_per_block[i]) > 0] - list_interp = [ - _delayed_interp_points_block(blocks[i], block_ids[i], points_arr[:, ind_per_block[i]], **kwargs) for i in used - ] + list_interp = [] + for i in used: + # Translate global indices by integer offsets to keep interpolation weights identical across chunks + block_kwargs = kwargs.copy() + if kwargs["method"] in ("nearest", "linear"): + row_offset = starts[0][indexes_yi[i]] - depth + col_offset = starts[1][indexes_xi[i]] - depth + block_kwargs["array_indices"] = ( + src_rows[ind_per_block[i]] - row_offset, + src_cols[ind_per_block[i]] - col_offset, + ) + list_interp.append( + _delayed_interp_points_block(blocks[i], block_ids[i], points_arr[:, ind_per_block[i]], **block_kwargs) + ) # We concatenate and re-order in a delayed manner def _concat_reorder(list_vals, list_inds): # type: ignore @@ -766,17 +813,10 @@ def _concat_reorder(list_vals, list_inds): # type: ignore joined = dask.delayed(_concat_reorder)(list_interp, list_inds_used) # Join into one array using a floating type whenever source values cannot represent NaN - output_dtype = _interp_output_dtype(darr.dtype) + output_dtype = _interp_output_dtype(darr.dtype, validity_only=kwargs.get("_validity_only", False)) interp_points = da.from_delayed(joined, shape=(len(points[0]),), dtype=output_dtype) # Padded edge chunks repeat their outer cells, so restore the bounds of the complete source raster - src_rows, src_cols = _xy2ij( - points[0], - points[1], - transform=transform, - area_or_point=kwargs["area_or_point"], - shift_area_or_point=kwargs["shift_area_or_point"], - ) inside = (src_rows >= -0.5) & (src_rows < darr.shape[0] - 0.5) inside &= (src_cols >= -0.5) & (src_cols < darr.shape[1] - 0.5) interp_points = da.where(inside, interp_points, np.nan) @@ -806,7 +846,7 @@ def _interp_points_partition( """Interpolate one point partition and return a point-cloud partition.""" # Preserve the planned output structure even when Dask sends an empty partition - out_dtype = _interp_output_dtype(source_raster.dtype) + out_dtype = _interp_output_dtype(source_raster.dtype, validity_only=extra_kwargs.get("_validity_only", False)) if len(part) == 0: return _empty_pointcloud_meta(data_column=data_column, crs=out_crs, dtype=out_dtype) @@ -821,7 +861,10 @@ def _interp_points_partition( shift_area_or_point=interp_options["shift_area_or_point"], ) # Detect partitions with no raster overlap before constructing interpolation work - ind_outofbounds: NDArrayBool = (i < 0) | (j < 0) | (i >= source_raster.shape[0]) | (j >= source_raster.shape[1]) + # Include outer half pixels accepted by nearest and linear interpolation when selecting partitions + margin = 0.5 if interp_options["method"] in {"nearest", "linear"} else 0 + ind_outofbounds: NDArrayBool = (i < -margin) | (j < -margin) + ind_outofbounds |= (i >= source_raster.shape[0] - margin) | (j >= source_raster.shape[1] - margin) if np.count_nonzero(~ind_outofbounds) == 0: z = np.full(len(part), np.nan, dtype=out_dtype) @@ -850,7 +893,7 @@ def _interp_points_partition( def _interp_points_dask_pointcloud( source_raster: RasterBase, points: Any, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"], + method: InterpolationMethod, band: int, input_latlon: bool, as_array: bool, @@ -876,7 +919,7 @@ def _interp_points_dask_pointcloud( out_crs = source_raster.crs points_in_crs = points if points.crs == out_crs else points.to_crs(out_crs) data_column = "z" - out_dtype = _interp_output_dtype(source_raster.dtype) + out_dtype = _interp_output_dtype(source_raster.dtype, validity_only=extra_kwargs.get("_validity_only", False)) meta = _empty_pointcloud_meta(data_column=data_column, crs=out_crs, dtype=out_dtype) # Package stable interpolation options once for each partition task @@ -900,25 +943,17 @@ def _interp_points_dask_pointcloud( out_crs, meta=meta, ) - # Import at runtime to avoid the point-cloud base importing this interpolation module in return - from geoutils.pointcloud.base import _set_dataframe_attrs - - # Restore the metadata expected by the GeoUtils ``pc`` accessor - _set_dataframe_attrs( - out, - { - "crs": out_crs, - "bounds": None, - "point_count": None, - "data_column": data_column, - "geometry_type": "Point", - }, - ) if as_array: - # Expose values as a Dask array without computing point partitions - return out[data_column].to_dask_array(lengths=True) - return out + # Read lengths from input points so array sizing does not execute the interpolation tasks + lengths = tuple(points.map_partitions(len).compute()) + return out[data_column].to_dask_array(lengths=lengths) + + # Import after package initialization because the point cloud package also imports interpolation + from geoutils.pointcloud.dataframe import _build_pointcloud_output + + # Set point output metadata without computing the interpolation tasks + return _build_pointcloud_output(out, data_column=data_column, as_dataframe=True) # SAME WITH MULTIPROCESSING @@ -928,9 +963,15 @@ def _wrapper_multiproc_interp_per_block( rst: Raster, block_id: dict[str, Any], interp_coords: NDArrayNum, + band: int = 1, **kwargs: Any, ) -> NDArrayNum: - """Wrapper to use interpolation per block.""" + """ + Read one raster tile and interpolate its selected band at the assigned points. + + Band selection and interpolation options follow _interp_points(); validity conversion stays inside its + shared array kernel so workers never need a complete source validity raster. + """ # Extract information out of block_id dictionary tile_idx = block_id["tile_idx"] @@ -938,9 +979,14 @@ def _wrapper_multiproc_interp_per_block( # Crop input raster for the given block rst_block = rst.icrop((tile_idx[2], tile_idx[0], tile_idx[3], tile_idx[1])) + # Loaded rasters keep every band when cropped; unloaded rasters already read only the requested band + array = rst_block.data + if array.ndim == 3: + array = array[band - 1] + # Interpolate to points by dispatching to base function interp_chunk = _interp_points_base( - array=rst_block.data, + array=array, transform=rst_block.transform, points=(interp_coords[0, :], interp_coords[1, :]), **kwargs, @@ -954,15 +1000,29 @@ def _multiproc_interp_points( rst: RasterBase, points: tuple[NDArrayNum, NDArrayNum], config: MultiprocConfig, + band: int = 1, **kwargs: Any, ) -> NDArrayNum: """ Interpolate raster at point coordinates on out-of-memory chunks. + + Band selection and interpolation options follow _interp_points(); config supplies tile sizes and the cluster. """ # Convert input to 2D array points_arr = np.vstack((points[0], points[1])) + # Compute global indices only for methods that pass them directly to blocks + src_indices = None + if kwargs["method"] in ("nearest", "linear"): + src_indices = _xy2ij( + points[0], + points[1], + transform=rst.transform, + area_or_point=kwargs["area_or_point"], + shift_area_or_point=kwargs["shift_area_or_point"], + ) + # Map depth of overlap required for each interpolation method depth = method_to_order[kwargs["method"]] + 1 # The overlap size is the order + 1 res = _res(rst.transform) @@ -971,11 +1031,10 @@ def _multiproc_interp_points( # Get multiprocessing chunk sizes chunks = normalize_chunks(chunks=config.chunks, shape=rst.shape) - chunksize = (chunks[0][0], chunks[1][0]) # Get starting 2D index for each chunk of the full array # (mirroring what is done in block_id of dask.array.map_blocks) - tiling = compute_tiling(tile_size=config.chunks, raster_shape=rst.shape, overlap=depth) + tiling = block_bounds_from_chunks(chunks=chunks, shape=rst.shape, overlap=depth) starts = [ cached_cumsum(chunks[0], initial_zero=True), cached_cumsum(chunks[1], initial_zero=True), @@ -989,7 +1048,6 @@ def _multiproc_interp_points( points_arr[1, :], starts, # type: ignore num_chunks, - chunksize, res[0], res[1], left, @@ -1003,6 +1061,16 @@ def _multiproc_interp_points( # Create tasks for multiprocessing tasks = [] for i in range(len(block_ids)): + # Reuse the full raster fractional indices instead of recalculating from each tile transform + block_kwargs = kwargs.copy() + if src_indices is not None: + src_rows, src_cols = src_indices + row_offset, _, col_offset, _ = block_ids[i]["tile_idx"] + block_kwargs["array_indices"] = ( + src_rows[ind_per_block[i]] - row_offset, + src_cols[ind_per_block[i]] - col_offset, + ) + # Launch the task on the cluster to process each tile tasks.append( config.cluster.submit( @@ -1010,7 +1078,8 @@ def _multiproc_interp_points( rst, block_ids[i], points_arr[:, ind_per_block[i]], - **kwargs, + band=band, + **block_kwargs, ) ) @@ -1041,7 +1110,7 @@ def _multiproc_interp_points( def _interp_points( source_raster: RasterBase, points: tuple[NDArrayNum, NDArrayNum] | tuple[Number, Number] | PointCloudLike, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] = None, + method: InterpolationMethod = None, band: int = 1, input_latlon: bool = False, as_array: bool = False, @@ -1051,9 +1120,19 @@ def _interp_points( return_interpolator: bool = False, mp_config: MultiprocConfig | None = None, nodata_propagation: NodataPropagation = "gdal", + _validity_only: bool = False, **kwargs: Any, ) -> Any: - """See description of Raster.interp_points.""" + """ + Check point interpolation inputs and dispatch to eager, Dask or multiprocessing calculation. + + Public options follow Raster.interp_points(). The shared _interp_points_base() kernel can also interpolate + source availability, allowing cosampling to select common finite points before reading their actual values. + + :param _validity_only: Replace finite source values by float32 one and missing values by NaN inside each loaded + array or worker block. Apply the same interpolation and nodata options to those values; the returned + numeric values describe availability and remain NaN where the interpolation cannot supply a value. + """ # If interpolation method undefined, default to the global system config if method is None: @@ -1062,7 +1141,6 @@ def _interp_points( propagation = _validate_nodata_propagation(nodata_propagation) # 1/ Input checks - if is_dask_geodataframe(points): if mp_config is not None: raise ValueError("Dask point-cloud inputs cannot be combined with Multiprocessing interpolation.") @@ -1078,7 +1156,7 @@ def _interp_points( shift_area_or_point=shift_area_or_point, force_scipy_function=force_scipy_function, return_interpolator=return_interpolator, - extra_kwargs=kwargs, + extra_kwargs={**kwargs, "_validity_only": _validity_only}, ) # Check and normalize input points @@ -1107,13 +1185,15 @@ def _interp_points( i, j = _xy2ij(x, y, transform=transform, area_or_point=area_or_point, shift_area_or_point=shift_area_or_point) - # Get index of points outside of bounds (i = row index vs shape[0], j = column index vs shape[1]) - ind_outofbounds: NDArrayBool = (i < 0) | (j < 0) | (i >= shape[0]) | (j >= shape[1]) + # Retain the outer half pixels accepted by nearest and linear array interpolation + margin = 0.5 if method in {"nearest", "linear"} else 0 + ind_outofbounds: NDArrayBool = (i < -margin) | (j < -margin) + ind_outofbounds |= (i >= shape[0] - margin) | (j >= shape[1] - margin) # If all points fell outside of bounds if np.count_nonzero(~ind_outofbounds) == 0: warnings.warn("All provided points were outside of raster bounds, returning only NaNs.") - output = np.full(x.shape[0], np.nan) + output = np.full(x.shape[0], np.nan, dtype=np.float32 if _validity_only else float) if as_array: return output else: @@ -1122,7 +1202,7 @@ def _interp_points( PointCloud, # Runtime import to avoid circular issues ) - return PointCloud.from_xyz(x=points[0], y=points[1], z=output, crs=source_raster.crs) + return PointCloud.from_xyz(x=x, y=y, z=output, crs=source_raster.crs) # Only work on points inside bounds pts_inbounds = x[~ind_outofbounds], y[~ind_outofbounds] @@ -1132,12 +1212,13 @@ def _interp_points( # 2/ Dispatch to either base (in-memory) function, Dask function, or Multiprocessing function class _InterpKwargs(TypedDict): area_or_point: Literal["Area", "Point"] | None - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] + method: InterpolationMethod dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None nodata_propagation: NodataPropagation shift_area_or_point: bool | None force_scipy_function: Literal["map_coordinates", "interpn"] | None return_interpolator: bool + _validity_only: bool interp_kwargs: _InterpKwargs = { "area_or_point": area_or_point, @@ -1147,6 +1228,7 @@ class _InterpKwargs(TypedDict): "shift_area_or_point": shift_area_or_point, "force_scipy_function": force_scipy_function, "return_interpolator": return_interpolator, + "_validity_only": _validity_only, } # Cannot use Multiprocessing backend and Dask backend simultaneously @@ -1174,7 +1256,7 @@ class _InterpKwargs(TypedDict): orig_bands = source_raster.bands source_raster._bands = (band,) z_inbounds = _multiproc_interp_points( - rst=source_raster, points=pts_inbounds, config=mp_config, **interp_kwargs, **kwargs + rst=source_raster, points=pts_inbounds, config=mp_config, band=band, **interp_kwargs, **kwargs ) # Rewrite original bands source_raster._bands = orig_bands @@ -1203,10 +1285,9 @@ class _InterpKwargs(TypedDict): return z_inbounds # Otherwise, return array of input length with NaNs for outside-bound points else: - # Get output length and dtype n = len(x) - dtype = source_raster.dtype + dtype = _interp_output_dtype(source_raster.dtype, validity_only=_validity_only) # Rebuild array (delayed if Dask, normal if NumPy) def _rebuild_with_nans(z_inbounds: NDArrayNum, mask_out: NDArrayBool, n: int, dtype: DTypeLike) -> NDArrayNum: @@ -1229,7 +1310,7 @@ def _rebuild_with_nans(z_inbounds: NDArrayNum, mask_out: NDArrayBool, n: int, dt PointCloud, # Runtime import to avoid circular issues ) - return PointCloud.from_xyz(x=points[0], y=points[1], z=z, crs=source_raster.crs) + return PointCloud.from_xyz(x=x, y=y, z=z, crs=source_raster.crs) ############################################################## @@ -1249,7 +1330,6 @@ def _reduce_points( as_array: bool = False, boundless: bool = True, ) -> NDArrayNum | tuple[NDArrayNum, NDArrayNum]: - # Check and normalize input points pts, input_scalar = _check_match_points(source_raster, points) @@ -1338,7 +1418,6 @@ def format_value(value: Any) -> Any: win: NDArrayNum | dict[int, NDArrayNum] = data else: - # Create rasterio's window for reading rio_window = rio.windows.Window(col, row, width, height) @@ -1374,7 +1453,7 @@ def format_value(value: Any) -> Any: ) if not as_array: - output_val = PointCloud.from_xyz(x=points[0], y=points[1], z=output_val, crs=source_raster.crs) + output_val = PointCloud.from_xyz(x=x, y=y, z=output_val, crs=source_raster.crs) if return_window: return (output_val, output_win) diff --git a/geoutils/interface/raster_point.py b/geoutils/interface/raster_point.py index 9e65611b1..d25ec36e5 100644 --- a/geoutils/interface/raster_point.py +++ b/geoutils/interface/raster_point.py @@ -16,7 +16,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Functionalities at the interface of rasters and point clouds.""" +"""Exact conversions between rasters and point clouds, without gridding or interpolation.""" from __future__ import annotations @@ -32,13 +32,17 @@ from geoutils._typing import NDArrayNum from geoutils.raster.array import get_mask_from_array from geoutils.raster.referencing import _default_nodata, _xy2ij -from geoutils.stats.sampling import _subsample_numpy if TYPE_CHECKING: from geoutils.pointcloud.pointcloud import PointCloud, PointCloudLike from geoutils.raster.base import RasterType +################################## +# 1/ REGULAR POINT CLOUD TO RASTER +################################## + + def _regular_pointcloud_to_raster( pointcloud: PointCloudLike, grid_coords: tuple[NDArrayNum, NDArrayNum] = None, @@ -120,6 +124,11 @@ def _regular_pointcloud_to_raster( return raster_arr, out_transform, gdf_pc.crs, out_nodata, area_or_point +######################### +# 2/ RASTER TO POINT CLOUD +######################### + + def _raster_to_pointcloud( source_raster: RasterType, data_column_name: str = "b1", @@ -188,6 +197,12 @@ def _raster_to_pointcloud( all_bands = [data_band] all_column_names = [data_column_name] + # Point sampling returns a compact eager result, so compute a separate copy of a Dask source + # Loading the caller's DataArray here would replace its lazy graph with an in-memory array + if source_raster._chunks is not None: + source_raster = source_raster.copy(deep=False).rst + source_raster.load() + # If subsample is the entire array, load it to optimize speed if subsample == 1 and not source_raster.is_loaded: source_raster.load() @@ -219,6 +234,8 @@ def _raster_to_pointcloud( valid_mask = np.ones(source_raster.data[0, :].shape, dtype=bool) # Get subsample on valid mask + from geoutils.sampling.subsampling import _subsample_numpy + # Build a low memory boolean masked array with invalid values masked to pass to subsampling ma_valid = np.ma.masked_array(data=np.ones(np.shape(valid_mask), dtype=bool), mask=~valid_mask) # Take a subsample within the valid values diff --git a/geoutils/interface/rasterization.py b/geoutils/interface/rasterization.py index 2201d28ff..fdd4761be 100644 --- a/geoutils/interface/rasterization.py +++ b/geoutils/interface/rasterization.py @@ -527,6 +527,9 @@ def _rasterize( mp_backend = mp_config is not None # A Dask reference keeps its chunked representation unless Multiprocessing is requested ref_chunks = get_geo_attr(ref, "_chunks") if ref is not None and has_geo_attr(ref, "_chunks") else None + if ref_chunks is not None: + # Match only spatial chunks when the reference also has a band dimension + ref_chunks = ref_chunks[-2:] dask_backend = bool(dask) or (da is not None and ref_chunks is not None) if mp_backend and dask_backend: @@ -695,25 +698,15 @@ def _create_mask_pointcloud_dask(source_vector: Vector, points: Any, as_array: b # Each point partition becomes an equally partitioned boolean point cloud out = points_in_crs.map_partitions(_mask_pointcloud_partition, source_geom, source_vector.crs, meta=meta) - # Import at runtime because the point-cloud base also uses rasterization through its vector parent - from geoutils.pointcloud.base import _set_dataframe_attrs - - # Restore the metadata expected by the GeoUtils ``pc`` accessor - _set_dataframe_attrs( - out, - { - "crs": source_vector.crs, - "bounds": None, - "point_count": None, - "data_column": "z", - "geometry_type": "Point", - }, - ) - if as_array: - # Return a lazy value array while keeping point partitions uncomputed + # Return mask values as a Dask array return out["z"].to_dask_array(lengths=True) - return out + + # Import after package initialization because the point cloud package also imports rasterization + from geoutils.pointcloud.dataframe import _build_pointcloud_output + + # Set point output metadata without computing the mask values + return _build_pointcloud_output(out, data_column="z", as_dataframe=True) def _create_mask_raster( diff --git a/geoutils/multiproc/__init__.py b/geoutils/multiproc/__init__.py index f9e221c92..b8f7913b4 100644 --- a/geoutils/multiproc/__init__.py +++ b/geoutils/multiproc/__init__.py @@ -16,6 +16,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Run geospatial work in multiprocessing blocks and reuse unloaded raster and point cloud inputs across passes.""" + from geoutils.multiproc.chunked import * # noqa from geoutils.multiproc.cluster import * # noqa from geoutils.multiproc.mparray import * # noqa diff --git a/geoutils/multiproc/chunked.py b/geoutils/multiproc/chunked.py index 5d14310cd..9da5a97c4 100644 --- a/geoutils/multiproc/chunked.py +++ b/geoutils/multiproc/chunked.py @@ -16,8 +16,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + """Module defining array and configuration routines for chunked operations with Multiprocessing.""" +from collections.abc import Iterator +from itertools import product from typing import Any, Literal, TypeVar, cast import geopandas as gpd @@ -250,10 +253,10 @@ def _chunks2d_from_chunksizes_shape( def normalize_chunks(chunks: ChunkSpec, shape: tuple[int, int]) -> tuple[tuple[int, ...], tuple[int, ...]]: """ - Normalize a Dask-like chunk specification into explicit 2D chunks. + Normalize the chunk user input into explicit 2D chunks, as in Dask. - Supports a single integer, a ``(y, x)`` chunk-size tuple, or an already-normalized - ``((y0, y1, ...), (x0, x1, ...))`` tuple. + Supports a single integer, a chunksize tuple for X/Y dimension, or a tuple containing chunksize in X/Y along the + full array ((x1, x2, ...), (y1, y2, ...)). """ if isinstance(chunks, int): @@ -282,6 +285,20 @@ def normalize_chunks(chunks: ChunkSpec, shape: tuple[int, int]) -> tuple[tuple[i return normalized_chunks +def iter_chunk_slices(shape: tuple[int, ...], chunks: int | tuple[int, ...]) -> Iterator[tuple[slice, ...]]: + """Divide an array shape into contiguous slices with principle row order.""" + + lengths = (chunks,) * len(shape) if isinstance(chunks, int) else chunks + slices = [ + [ + slice(start, min(start + lengths[axis % len(lengths)], size)) + for start in range(0, size, lengths[axis % len(lengths)]) + ] + for axis, size in enumerate(shape) + ] + yield from product(*slices) + + def cached_cumsum(chunks: tuple[int, ...], initial_zero: bool = True) -> tuple[int, ...]: """Like dask's cumulative chunk starts. For (3,3,1) -> [0,3,6,7] if initial_zero.""" out = [0] if initial_zero else [] diff --git a/geoutils/multiproc/cluster.py b/geoutils/multiproc/cluster.py index a04c89e74..117090e09 100644 --- a/geoutils/multiproc/cluster.py +++ b/geoutils/multiproc/cluster.py @@ -20,10 +20,52 @@ """This module defines the cluster configurations.""" import multiprocessing +from collections.abc import Iterable, Iterator from multiprocessing.pool import Pool from typing import Any, Callable, Dict, Optional +def _map_bounded( + cluster: "AbstractCluster", + function: Callable[..., Any], + arguments: Iterable[tuple[Any, ...]], + max_pending: int = 8, +) -> Iterator[tuple[int, Any]]: + """ + Run a function for many inputs without submitting all the work at once. + + Submit up to ``max_pending`` calls, wait for their results, and then start the next batch. This limits how many + jobs and input values the cluster must keep at one time. Return results in the same order as the inputs, together + with the position of each input. + + :param cluster: Cluster used to run the function calls. + :param function: Function to call for each set of arguments. + :param arguments: Sets of positional arguments passed to the function in order. + :param max_pending: Largest number of calls submitted before waiting for their results. + + :returns: Input positions and their function results, in input order. + """ + + # Check the batch size before reading any inputs + if max_pending <= 0: + raise ValueError("Argument ``max_pending`` must be a positive integer.") + + # Submit one small batch at a time + pending: list[tuple[int, Any]] = [] + for index, args in enumerate(arguments): + pending.append((index, cluster.submit(function, *args))) + if len(pending) == max_pending: + # Wait for the batch and return its results in input order + results = cluster.gather([future for _, future in pending]) + yield from ((task_index, result) for (task_index, _), result in zip(pending, results)) + pending = [] + + # Finish the last, smaller batch + if pending: + results = cluster.gather([future for _, future in pending]) + yield from ((task_index, result) for (task_index, _), result in zip(pending, results)) + + class ClusterGenerator: def __new__(cls, name: str, nb_workers: int = 2) -> "AbstractCluster": # type: ignore """ diff --git a/geoutils/multiproc/mparray.py b/geoutils/multiproc/mparray.py index 153fd51b0..48f70b599 100644 --- a/geoutils/multiproc/mparray.py +++ b/geoutils/multiproc/mparray.py @@ -16,13 +16,17 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Module defining array and configuration routines for chunked operations with Multiprocessing.""" + +"""Module defining chunked array operations with Multiprocessing.""" from __future__ import annotations import logging import tempfile import warnings +from collections.abc import Iterator +from contextlib import contextmanager from numbers import Integral +from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Literal, overload import numpy as np @@ -35,6 +39,7 @@ from geoutils.multiproc.cluster import AbstractCluster, ClusterGenerator if TYPE_CHECKING: + from geoutils.raster.base import RasterBase from geoutils.raster.raster import Raster @@ -74,7 +79,7 @@ def _split_chunk_size(chunks: ChunkSize) -> tuple[int, int]: class MultiprocConfig: """ - Configuration class for handling multiprocessing parameters in raster processing. + Configuration class for handling multiprocessing parameters in raster and point cloud processing. This class encapsulates settings related to multiprocessing, allowing users to specify chunks, output file, and an optional cluster for parallel processing. @@ -85,16 +90,17 @@ def __init__( self, chunks: ChunkSize, outfile: str | None = None, - driver: str = "GTiff", + driver: str | None = None, cluster: AbstractCluster | None = None, ): """ Initialize the MultiprocConfig instance with multiprocessing settings. :param chunks: The size of the chunks for splitting raster data. Pass an integer for square chunks, or a - ``(rows, cols)`` tuple for rectangular chunks. + ``(rows, cols)`` tuple for rectangular chunks. Point cloud operations use an integer number of points. :param outfile: The file path where the output will be written. - :param driver: Driver to write file with. + :param driver: Output format. None uses GeoTIFF for rasters; point reprojection infers LAS/LAZ or GeoPackage + from the output filename, defaulting to GeoPackage when no extension is given. :param cluster: A cluster object for distributed computing, or None for sequential processing. """ self.chunks = _validate_chunk_size(chunks) @@ -113,6 +119,23 @@ def __init__( def copy(self) -> MultiprocConfig: return MultiprocConfig(chunks=self.chunks, outfile=self.outfile, driver=self.driver, cluster=self.cluster) + @contextmanager + def temporary(self) -> Iterator[MultiprocConfig]: + """ + Use a temporary output file while reusing this configuration's chunk sizes, driver and worker cluster. + + Each context owns a separate directory for its output and any driver sidecar files. Keep the context open + while intermediate files are being read; leaving it removes the files without closing the shared cluster. + + :returns: Context manager yielding a configuration with a unique output filename. The original + configuration remains unchanged, and temporary files are removed even when processing raises an error. + """ + + with tempfile.TemporaryDirectory() as directory: + yield MultiprocConfig( + chunks=self.chunks, outfile=str(Path(directory) / "output"), driver=self.driver, cluster=self.cluster + ) + def _generate_tiling_grid( row_min: int, @@ -344,7 +367,7 @@ def _apply_func_block( def map_overlap( func: Callable[..., Raster], - raster_path: str | Raster, + raster_path: str | RasterBase, mp_config: MultiprocConfig, *args: Any, depth: int = 0, @@ -359,7 +382,8 @@ def map_overlap( Use this function when `func` returns a :class:`geoutils.Raster`, as it ensures the processed data is written to `config.outfile`. - :param func: A function to apply to each raster tile. It must return a :class:`geoutils.Raster` object. + :param func: Function returning a Raster on the tile's grid. It may change the number of bands; every tile must + use the same output band count, dtype and nodata. Metadata comes from the first input tile's result. :param raster_path: Path to the input raster file or an existing :class:`geoutils.Raster` object. :param mp_config: Configuration object containing chunks, output file path, and an optional cluster. The `outfile` parameter in `config` must be provided. @@ -367,6 +391,8 @@ def map_overlap( :param depth: The overlap size between blocks to avoid edge effects, default is 0. :param kwargs: Additional keyword arguments to pass to `func`. + :returns: File-backed raster with the output bands and metadata returned by func(). + :raises ValueError: If `config.outfile` is not provided. :raises RuntimeError: If an error occurs while processing the raster blocks. """ @@ -391,19 +417,19 @@ def map_overlap( # Submit the block processing task on the cluster tasks.append(mp_config.cluster.submit(_apply_func_block, func, raster, tile, depth, *args, **kwargs)) - # get first tile to retrieve dtype and nodata + # Read the first result's band count and metadata because the function may change them result_tile0, _ = mp_config.cluster.compute(tasks[0]) file_metadata = { "width": raster.width, "height": raster.height, - "count": raster.count, + "count": result_tile0.count, "crs": raster.crs, "transform": raster.transform, "dtype": result_tile0.dtype, "nodata": result_tile0.nodata, } - raster_output = _write_multiproc_result(tasks, mp_config, file_metadata) + raster_output = _write_multiproc_result(tasks, mp_config, file_metadata, tags=dict(result_tile0.tags)) # Warns user if output file is a BigTIFF if raster_output._is_bigtiff(): @@ -419,13 +445,31 @@ def _write_multiproc_result( tasks: list[Any], mp_config: MultiprocConfig, file_metadata: dict[str, Any], + *, + tags: dict[str, Any] | None = None, ) -> Raster: + """ + Write completed worker tiles to their output windows without collecting the full raster in memory. + + :param tasks: Worker results containing a Raster or NumPy tile and its row/column window boundaries. + :param mp_config: Worker cluster, output filename and driver, as described by map_overlap(). + :param file_metadata: Rasterio output dimensions, band count, dtype, georeferencing and nodata. + :param tags: Optional output metadata, including AREA_OR_POINT when pixel interpretation is known. Rasterio + stores these as strings; the returned Raster keeps the supplied Python values. + + :returns: File-backed Raster or mask with every completed tile written and its output metadata attached. + """ # To avoid circular import, runtime here from geoutils.raster import Raster # Create a new raster file to save the processed results - with rio.open(mp_config.outfile, "w", driver=mp_config.driver, **file_metadata, BIGTIFF="IF_NEEDED") as dst: + with rio.open( + mp_config.outfile, "w", driver=mp_config.driver or "GTiff", **file_metadata, BIGTIFF="IF_NEEDED" + ) as dst: + # Preserve result metadata, including pixel interpretation, when the caller provides it + if tags is not None: + dst.update_tags(**tags) try: # Retrieve completed blocks promptly so worker results can be released after writing for task_index, completed in mp_config.cluster.iter_completed(tasks): @@ -458,15 +502,17 @@ def _write_multiproc_result( except Exception as e: raise RuntimeError(f"Error retrieving raster blocks from multiprocessing tasks: {e}") - if is_mask: - return Raster(mp_config.outfile, as_mask=True) - return Raster(mp_config.outfile) + # Keep metadata types such as band-name tuples on the returned object after the normal file metadata read + output = Raster(mp_config.outfile, is_mask=is_mask) + if tags is not None: + output.tags.update(tags) + return output @overload def map_blocks( func: Callable[..., Any], - raster_path: str | Any, + raster_path: str | RasterBase, mp_config: MultiprocConfig, *args: Any, depth: int = 0, @@ -478,7 +524,7 @@ def map_blocks( @overload def map_blocks( func: Callable[..., Any], - raster_path: str | Any, + raster_path: str | RasterBase, mp_config: MultiprocConfig, *args: Any, depth: int = 0, @@ -490,7 +536,7 @@ def map_blocks( @overload def map_blocks( func: Callable[..., Any], - raster_path: str | Any, + raster_path: str | RasterBase, mp_config: MultiprocConfig, *args: Any, depth: int = 0, @@ -501,7 +547,7 @@ def map_blocks( def map_blocks( func: Callable[..., Any], - raster_path: str | Raster, + raster_path: str | RasterBase, mp_config: MultiprocConfig, *args: Any, depth: int = 0, @@ -571,7 +617,7 @@ def map_blocks( @deprecate(details="Use map_overlap() instead.") def map_overlap_multiproc_save( func: Callable[..., Raster], - raster_path: str | Raster, + raster_path: str | RasterBase, mp_config: MultiprocConfig, *args: Any, depth: int = 0, @@ -585,7 +631,7 @@ def map_overlap_multiproc_save( @deprecate(details="Use map_blocks() with return_block_info instead.") def map_multiproc_collect( func: Callable[..., Any], - raster_path: str | Raster, + raster_path: str | RasterBase, mp_config: MultiprocConfig, *args: Any, depth: int = 0, diff --git a/geoutils/multiproc/readers.py b/geoutils/multiproc/readers.py new file mode 100644 index 000000000..d82a8e220 --- /dev/null +++ b/geoutils/multiproc/readers.py @@ -0,0 +1,473 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Read raster or point cloud values in reusable blocks across one or more processing passes.""" + +from __future__ import annotations + +import copy +import math +from contextlib import ExitStack +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import pyogrio + +from geoutils._dispatch import _get_pointcloud_interface, _get_raster_interface +from geoutils._misc import import_optional +from geoutils.raster.array import get_mask_from_array + +if TYPE_CHECKING: + import geopandas as gpd + + from geoutils.multiproc.mparray import MultiprocConfig + from geoutils.pointcloud.base import PointCloudBase + from geoutils.raster.base import RasterBase + + +######################################## +# 1/ RASTER AND POINT CLOUD READERS +######################################## + + +@dataclass(frozen=True) +class _ValueReader: + """ + Describe raster or point cloud values that workers can read one block at a time. + + Values can come directly from a raster band or point cloud column. They can also be calculated from a raster or + vector on the requested raster cells or point rows. block() limits the reader to one such part before it is sent + to a worker. + + For rasters, ``selector`` is a band number. For point clouds, it is a column name, and None uses the main data + column or geometry height. The reader gets the value shape and data type from the file. + """ + + source: Any + selector: int | str | None = None + mask: Any | None = None + support: Any | None = None + interpolation: str = "linear" + chunks: int | tuple[int, int] = 512 + vector_values: Any | None = None + mask_mode: Literal["inside", "outside"] | None = None + coverage: bool = False + shape: tuple[int, ...] = field(init=False) + dtype: np.dtype[Any] = field(init=False) + kind: Literal["raster", "point", "interpolated", "vector"] = field(init=False) + + def __post_init__(self) -> None: + """Check the raster or point cloud source without loading its values.""" + + # Use the support shape because vector values are calculated when each block is read + if self.vector_values is not None: + assert self.support is not None + raster = _get_raster_interface(self.support) + shape = tuple(raster.shape) if raster is not None else (int(self.support.point_count),) + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "dtype", np.dtype(bool if self.mask_mode is not None or self.coverage else float)) + object.__setattr__(self, "kind", "vector") + return + + # Copy the source information so worker reads cannot load or change the source object + raster = _get_raster_interface(self.source) + pointcloud = _get_pointcloud_interface(self.source) if raster is None else None + source = raster if raster is not None else pointcloud + if source is None or source.name is None or source.is_loaded: + raise ValueError("Value readers require an unloaded raster or point cloud.") + source = copy.copy(source) + + # Record the selected raster band and the shape of its output + if raster is not None: + selector = 1 if self.selector is None else self.selector + if not isinstance(selector, (int, np.integer)) or not 1 <= selector <= source.count: + raise ValueError("Raster bands must be integers between one and the raster band count.") + selector = int(selector) + shape = tuple(source.shape) + dtype = np.dtype(bool if source.is_mask else source.dtype) + kind = "raster" + if self.support is not None: + shape = (int(self.support.point_count),) + dtype = np.dtype(bool if source.is_mask else float) + kind = "interpolated" + else: + # Record the point count and selected column type from the file + selector = source.data_column if self.selector is None else self.selector + if selector is not None and (not isinstance(selector, str) or selector not in source.columns): + raise ValueError(f"Point column {selector!r} does not exist.") + from geoutils.pointcloud.las import _is_laspy_supported + + if _is_laspy_supported(source.name): + laspy = import_optional("laspy") + with laspy.open(source.name) as reader: + count = reader.header.point_count + dimension = reader.header.point_format.dimension_by_name(selector or "Z") + if dimension.num_elements != 1: + raise ValueError("Point values require one scalar value per selected column.") + if selector in (None, "Z") or dimension.scales is not None or dimension.offsets is not None: + dtype = np.dtype(float) + else: + dtype = np.dtype("uint8") if dimension.dtype is None else dimension.dtype + else: + info = pyogrio.read_info(source.name, force_feature_count=True) + count = info["features"] + dtypes = dict(zip(info["fields"], info["dtypes"])) + dtype = np.dtype(float if selector is None else dtypes[selector]) + shape = (int(count),) + kind = "point" + + # Store the checked source and its basic information on the frozen reader + object.__setattr__(self, "source", source) + object.__setattr__(self, "selector", selector) + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "dtype", dtype) + object.__setattr__(self, "kind", kind) + + @property + def ndim(self) -> int: + """Return the number of dimensions in the raster or point cloud layout.""" + + return len(self.shape) + + @property + def size(self) -> int: + """Return the total number of raster cells or points.""" + + return math.prod(self.shape) + + def block(self, slices: slice | tuple[slice, ...]) -> _BlockReader: + """Create a reader for contiguous slices and the matching part of any in-memory mask.""" + + # Make slice bounds explicit before they are sent to workers + slices = (slices,) if isinstance(slices, slice) else slices + if not isinstance(slices, tuple) or len(slices) != self.ndim or any(not isinstance(s, slice) for s in slices): + raise TypeError("Argument ``slices`` must contain one slice per input dimension.") + ranges = [part.indices(length) for part, length in zip(slices, self.shape)] + if any(step != 1 for _, _, step in ranges): + raise ValueError("Argument ``slices`` must select contiguous increasing slices.") + slices = tuple(slice(start, max(start, stop)) for start, stop, _ in ranges) + shape = tuple(part.stop - part.start for part in slices) + + # Send only the matching part of an in-memory mask to each worker + mask = self.mask + if isinstance(mask, _ValueReader): + mask = mask.block(slices) + elif mask is not None: + mask = mask[slices] + return _BlockReader( + source=self.source, + selector=self.selector, + kind=self.kind, + slices=slices, + shape=shape, + dtype=self.dtype, + mask=mask, + support=self.support, + interpolation=self.interpolation, + chunks=self.chunks, + vector_values=self.vector_values, + mask_mode=self.mask_mode, + coverage=self.coverage, + ) + + def read(self, slices: slice | tuple[slice, ...]) -> Any: + """Read only the requested raster window or point rows and apply the optional mask.""" + + return self.block(slices).read() + + def read_points(self, rows: slice) -> gpd.GeoDataFrame: + """Read point rows with their geometry and selected column, ignoring the value mask.""" + + if self.kind != "point": + raise TypeError("Point rows require a point cloud reader.") + block = self.block(rows) + return _read_point_rows(block.source, block.slices[0], block.selector) + + +def _normalize_reader_mask(mask: Any, shape: tuple[int, ...]) -> Any: + """Validate a reader mask from file information and ordinary masks from their values.""" + + if isinstance(mask, _ValueReader): + if mask.shape != shape or not np.issubdtype(mask.dtype, np.bool_): + raise ValueError("Argument ``mask`` must be boolean and contain one value per input location.") + return mask + + from geoutils.sampling.support import _normalize_mask_array + + return _normalize_mask_array(mask, shape) + + +@dataclass(frozen=True) +class _BlockReader: + """Store the source and slices needed to read one raster or point cloud block inside a worker.""" + + source: Any + selector: int | str | None + kind: Literal["raster", "point", "interpolated", "vector"] + slices: tuple[slice, ...] + shape: tuple[int, ...] + dtype: np.dtype[Any] + mask: Any | None = None + support: Any | None = None + interpolation: str = "linear" + chunks: int | tuple[int, int] = 512 + vector_values: Any | None = None + mask_mode: Literal["inside", "outside"] | None = None + coverage: bool = False + + def read(self) -> Any: + """Read this raster window or set of point rows and apply its optional mask.""" + + # Return an empty result without giving file readers an invalid zero-size request + values: Any + if math.prod(self.shape) == 0: + values = np.empty(self.shape, dtype=self.dtype) + elif self.kind == "raster": + from geoutils.raster.transformation import _crop + + # Read the selected raster band and window + assert isinstance(self.selector, (int, np.integer)) + raster = copy.copy(self.source) + raster._bands = self.source.bands[int(self.selector) - 1] + raster._out_count = 1 + rows, columns = self.slices + bounds = (columns.start, rows.start, columns.stop, rows.stop) + values, _ = _crop(raster, bounds, distance_unit="pixel") + elif self.kind == "vector": + from affine import Affine + + from geoutils.raster import Raster + from geoutils.sampling.support import ( + _mask_at_support, + _sample_vector_values, + ) + + # Build only the requested part of the raster or point support + raster = _get_raster_interface(self.support) + if raster is not None: + rows, columns = self.slices + transform = raster.transform * Affine.translation(columns.start, rows.start) + support = Raster.from_array( + np.zeros(self.shape, dtype=bool), transform, raster.crs, area_or_point=raster.area_or_point + ) + points = None + else: + support = self.support + assert support is not None + points = _read_point_rows(support, self.slices[0], None) + + # Calculate a vector mask or values on this part of the support + if self.mask_mode is not None: + values = _mask_at_support(self.source, support, support_dataframe=points, mask_mode=self.mask_mode) + else: + assert self.vector_values is not None + values = _sample_vector_values(self.source, self.vector_values, support, points) + if self.coverage: + values = np.isfinite(values) + elif self.kind == "interpolated": + from geoutils.multiproc.mparray import MultiprocConfig + + # Read raster values at the point locations in this block + assert self.support is not None + points = _read_point_rows(self.support, self.slices[0], None) + coordinates = (points.geometry.x.to_numpy(), points.geometry.y.to_numpy()) + config = MultiprocConfig(chunks=self.chunks) + values = self.source.interp_points( + coordinates, method=self.interpolation, band=self.selector, as_array=True, mp_config=config + ) + if self.source.is_mask: + values = np.isfinite(values) & (values != 0) + else: + # Read point rows and check that a separate support file uses the same order + dataframe = _read_point_rows(self.source, self.slices[0], self.selector) + if self.support is not None: + from geoutils.pointcloud.testing import _point_coords_equal_eager + + reference = _read_point_rows(self.support, self.slices[0], None) + if not _point_coords_equal_eager(dataframe, reference): + raise ValueError("Point values do not share the ordered support coordinates.") + values = np.asarray(dataframe.geometry.z if self.selector is None else dataframe[self.selector]) + + # Apply the common mask while preserving integer values + if self.mask is not None: + mask = _read_values(self.mask) + keep = np.ma.filled(mask, False) + invalid = get_mask_from_array(values).reshape(values.shape) + values = np.ma.masked_where(~keep | invalid, values) + return values + + +######################################## +# 2/ READ RASTER CELLS AND POINT ROWS +######################################## + + +def _read_point_rows(source: PointCloudBase, rows: slice, selector: int | str | None) -> gpd.GeoDataFrame: + """Read selected point rows and columns without loading the complete point cloud.""" + + count = rows.stop - rows.start + assert selector is None or isinstance(selector, str) + columns = [] if selector is None else [selector] + + # Read in-memory rows directly and otherwise use the reader for the file format + if source.is_loaded or source._is_pd: + dataframe = source.ds.iloc[rows] + return dataframe[[*columns, dataframe.geometry.name]] + assert source.name is not None + from geoutils.pointcloud.las import _is_laspy_supported, _load_laspy_data_slice + + if _is_laspy_supported(source.name): + return _load_laspy_data_slice(source.name, columns=columns, start=rows.start, count=count) + + # A zero max_features means unlimited rows in Pyogrio, so empty requests read at most one row for their schema + dataframe = pyogrio.read_dataframe( + source.name, columns=columns, skip_features=rows.start, max_features=max(1, count) + ) + return dataframe.iloc[:0] if count == 0 else dataframe + + +def _read_values(value: Any) -> Any: + """Read a raster or point cloud block and pass ordinary in-memory values through unchanged.""" + + return value.read() if isinstance(value, _BlockReader) else value + + +def _read_block_sample(block: Any, indexes: tuple[Any, ...]) -> Any: + """Read one raster or point cloud block and return only its requested positions.""" + + return _read_values(block)[indexes] + + +def _read_selected_values(reader: _ValueReader, indexes: Any, mp_config: MultiprocConfig | None) -> Any: + """Read selected raster cells or points from only the blocks that contain them, preserving their order.""" + + from geoutils.multiproc.chunked import iter_chunk_slices + from geoutils.multiproc.cluster import _map_bounded + + assert mp_config is not None + arguments = [] + positions_in_sample = [] + coordinates = np.unravel_index(indexes, reader.shape) + + # Find which requested positions belong to each block + for slices in iter_chunk_slices(reader.shape, mp_config.chunks): + selected = np.ones(len(indexes), dtype=bool) + for positions, part in zip(coordinates, slices): + selected &= (positions >= part.start) & (positions < part.stop) + if not np.any(selected): + continue + + # Record local positions and their place in the requested order + local = tuple(positions[selected] - part.start for positions, part in zip(coordinates, slices)) + arguments.append((reader.block(slices), local)) + positions_in_sample.append(np.flatnonzero(selected)) + + # Read only blocks that contain selected positions + pieces = [result for _, result in _map_bounded(mp_config.cluster, _read_block_sample, arguments)] + if not pieces: + return np.empty(0, dtype=reader.dtype) + + # Restore the requested order and use the data type returned by the file reader + concatenate = np.ma.concatenate if any(np.ma.isMaskedArray(piece) for piece in pieces) else np.concatenate + sampled = concatenate(pieces) + order = np.argsort(np.concatenate(positions_in_sample)) + return sampled[order] + + +######################################## +# 3/ CREATE RASTER AND POINT CLOUD READERS +######################################## + + +def _reader_from_vector( + dataframe: Any, + values: Any, + support: Any, + mp_config: MultiprocConfig | None, + *, + mask_mode: Literal["inside", "outside"] | None = None, + coverage: bool = False, +) -> _ValueReader | None: + """Create a reader that places vector values or a vector mask on an unloaded raster or point cloud support.""" + + # Keep loaded and Dask supports on their existing processing paths + if mp_config is None or support.is_loaded: + return None + raster = _get_raster_interface(support) + if raster is not None and raster._is_xr: + return None + pointcloud = _get_pointcloud_interface(support) + if pointcloud is not None and pointcloud._is_pd: + return None + return _ValueReader(dataframe, support=support, vector_values=values, mask_mode=mask_mode, coverage=coverage) + + +def _reader_from_source( + value_source: Any, + selector: int | str | None, + support: RasterBase | PointCloudBase, + mp_config: MultiprocConfig | None, + *, + align: str = "raise", + interpolation: str = "linear", + stack: ExitStack | None = None, +) -> _ValueReader | None: + """Create a reader when an unloaded raster or point cloud can stay on disk.""" + + if mp_config is None: + return None + raster = _get_raster_interface(value_source) + if raster is not None: + # Use a raster reader only when its values can stay on disk + if raster._is_xr or raster.is_loaded or raster.name is None: + return None + support_raster = _get_raster_interface(support) + if support_raster is not None and raster.georeferenced_grid_equal(support_raster): + return _ValueReader(raster, selector) + support_points = _get_pointcloud_interface(support) + if support_points is not None and support_points._is_dask: + raise ValueError("Dask inputs cannot be combined with multiprocessing reads from files.") + needs_alignment = support_raster is not None or raster.crs != support.crs + if needs_alignment: + # Align the raster once because workers cannot read it directly on the requested support + if align != "reproject" or stack is None: + return None + from geoutils.sampling.support import _aligned_raster + + config = stack.enter_context(mp_config.temporary()) + config.driver = "GTiff" + raster = _aligned_raster(raster, raster, support, "values", align, mp_config=config) + if support_raster is not None: + return _ValueReader(raster, selector) + return _ValueReader( + raster, selector, support=support_points, interpolation=interpolation, chunks=mp_config.chunks + ) + + # Use a point reader only when values and support have the same ordered locations + pointcloud = _get_pointcloud_interface(value_source) + support_points = _get_pointcloud_interface(support) + if pointcloud is not None and support_points is not None: + if pointcloud._is_pd or pointcloud.is_loaded or pointcloud.name is None: + return None + if support_points._is_dask: + raise ValueError("Dask inputs cannot be combined with multiprocessing reads from files.") + if pointcloud.crs != support_points.crs: + if align != "reproject" or stack is None: + return None + config = stack.enter_context(mp_config.temporary()) + config.driver = "GPKG" + config.outfile += ".gpkg" + if isinstance(config.chunks, tuple): + config.chunks = math.prod(config.chunks) + pointcloud = pointcloud.reproject(crs=support_points.crs, mp_config=config) + if pointcloud.point_count != support_points.point_count: + raise ValueError("Point values do not share the ordered support coordinates.") + reference = None if pointcloud is support_points else support_points + return _ValueReader(pointcloud, selector, support=reference) + return None diff --git a/geoutils/pointcloud/base.py b/geoutils/pointcloud/base.py index babbf85e4..150f7aeed 100644 --- a/geoutils/pointcloud/base.py +++ b/geoutils/pointcloud/base.py @@ -21,6 +21,7 @@ from __future__ import annotations import warnings +from collections.abc import Mapping from typing import ( TYPE_CHECKING, Any, @@ -28,6 +29,7 @@ Iterable, Literal, TypeVar, + cast, overload, ) @@ -39,48 +41,39 @@ from geoutils import profiler from geoutils._dispatch import get_geo_attr, is_dask_dataframe from geoutils._misc import import_optional -from geoutils._typing import ArrayLike, NDArrayBool, NDArrayNum, Number +from geoutils._typing import ArrayLike, DTypeLike, NDArrayBool, NDArrayNum, Number from geoutils.interface._nodata import NodataPropagation from geoutils.interface.gridding import ( GriddingEngine, GriddingMethod, _grid_pointcloud_to_raster, ) +from geoutils.pointcloud.dataframe import ( + _build_pointcloud_output, + _get_dataframe_attrs, + _set_dataframe_attrs, +) from geoutils.pointcloud.testing import _georeferenced_coords_equal -from geoutils.stats.sampling import _subsample_pointcloud -from geoutils.stats.stats import _statistics +from geoutils.sampling.subsampling import _subsample_pointcloud +from geoutils.stats.stats import stats as _stats +from geoutils.stats.stats import variogram as _variogram from geoutils.vector.base import VectorBase +from geoutils.vector.transformation import _get_reproject_crs if TYPE_CHECKING: + import xarray as xr + + from geoutils.interface.interpolation import InterpolationMethod from geoutils.multiproc import MultiprocConfig + from geoutils.pointcloud.pointcloud import PointCloudLike from geoutils.raster.base import RasterLike + from geoutils.stats.variography import Variogram + from geoutils.vector.base import VectorLike PointCloudBaseType = TypeVar("PointCloudBaseType", bound="PointCloudBase") -def _get_dataframe_attrs(ds: Any) -> dict[str, Any]: - """Get GeoUtils metadata from Pandas or Dask dataframes.""" - - # Dask does not carry Pandas ``attrs`` reliably through graph operations - if is_dask_dataframe(ds): - try: - return object.__getattribute__(ds, "_geoutils_attrs") - except AttributeError: - return {} - return getattr(ds, "attrs", {}) - - -def _set_dataframe_attrs(ds: Any, attrs: dict[str, Any]) -> None: - """Set GeoUtils metadata on Pandas or Dask dataframes.""" - - # Keep a private copy on Dask collections and use the public mapping for Pandas - if is_dask_dataframe(ds): - object.__setattr__(ds, "_geoutils_attrs", attrs.copy()) - elif hasattr(ds, "attrs"): - ds.attrs.update(attrs) - - class PointCloudBase(VectorBase): """ Shared implementation for :class:`geoutils.PointCloud` and the ``pc`` Pandas accessor. @@ -98,10 +91,12 @@ def _is_dask(self) -> bool: @property def _has_z(self) -> bool: - """Whether the point geometries all have a Z coordinate or not.""" + """Whether all point geometries have a Z coordinate.""" if self._is_dask: return False + if not self.is_loaded: + return getattr(self, "_geometry_type", None) in ("Point Z", "3D Point") return all(p.has_z for p in self.ds.geometry) if len(self.ds.geometry) > 0 else False @property @@ -149,6 +144,11 @@ def data_column(self) -> str | None: Can be None if point geometries are 3D. """ + if self._is_pd: + # Multiple accessors can share a dataframe, so its metadata owns the selected column + attrs = _get_dataframe_attrs(self.ds) + if "data_column" in attrs: + return attrs["data_column"] return getattr(self, "_data_column", None) @data_column.setter @@ -164,9 +164,13 @@ def set_data_column(self, new_data_column: str | None) -> None: :param new_data_column: Column to use, or None to use Z coordinates stored in 3D point geometry. """ - if self.is_loaded and not self._is_dask and self._has_z: + if self._has_z: if new_data_column is None: self._data_column = None + if self._is_pd or self.is_loaded: + attrs = _get_dataframe_attrs(self.ds) + attrs["data_column"] = None + _set_dataframe_attrs(self.ds, attrs) return warnings.warn( f"Overriding 3D points with with data column '{new_data_column}'. Set data_column " @@ -221,22 +225,14 @@ def is_mask(self) -> bool: def _cast_pointcloud_output(self, new_ds: Any) -> Any: """Cast a GeoDataFrame-like point cloud output to the proper public type.""" - attrs = _get_dataframe_attrs(self.ds) - new_crs = getattr(new_ds, "crs", None) - if new_crs is not None and new_crs != attrs.get("crs"): - attrs["crs"] = new_crs - attrs["bounds"] = None - attrs["data_column"] = self.data_column - attrs["geometry_type"] = "Point" - _set_dataframe_attrs(new_ds, attrs) - - # Accessors expose dataframe-like outputs while PointCloud wraps eager outputs - if self._is_pd or self._is_dask: - return new_ds - - from geoutils.pointcloud.pointcloud import PointCloud - - return PointCloud(new_ds, data_column=self.data_column) + # Copy source metadata before updating the result so its cached values remain independent + attrs = _get_dataframe_attrs(self.ds).copy() + return _build_pointcloud_output( + new_ds, + data_column=self.data_column, + as_dataframe=self._is_pd or self._is_dask, + attrs=attrs, + ) def _override_gdf_output(self, other: Any) -> Any: """Keep point-preserving GeoDataFrame outputs as point clouds.""" @@ -264,29 +260,35 @@ def copy(self, new_array: NDArrayNum | NDArrayBool | Any | None = None) -> Any: if self.data_column is None: raise ValueError("Dask-backed point clouds require an explicit data column.") new_ds = new_ds.assign(**{self.data_column: new_array}) - return self._cast_pointcloud_output(new_ds) - - new_ds = self.ds.copy() - if new_array is not None: - if not isinstance(new_array, np.ndarray): - new_array = np.asarray(new_array) - new_array = new_array.squeeze() - if not (new_array.ndim == 1 and new_array.shape[0] == self.point_count): - raise ValueError( - "New data array must be 1-dimensional with the same number of points as the point " - "cloud being copied." - ) - if self.data_column is not None: - new_ds[self.data_column] = new_array - else: - new_ds.geometry = gpd.points_from_xy( - x=self.geometry.x.to_numpy(), - y=self.geometry.y.to_numpy(), - z=new_array, - crs=self.crs, - ) - - return self._cast_pointcloud_output(new_ds) + else: + new_ds = self.ds.copy() + if new_array is not None: + if not isinstance(new_array, np.ndarray): + new_array = np.asarray(new_array) + new_array = new_array.squeeze() + if not (new_array.ndim == 1 and new_array.shape[0] == self.point_count): + raise ValueError( + "New data array must be 1-dimensional with the same number of points as the point " + "cloud being copied." + ) + if self.data_column is not None: + new_ds[self.data_column] = new_array + else: + new_ds.geometry = gpd.points_from_xy( + x=self.geometry.x.to_numpy(), + y=self.geometry.y.to_numpy(), + z=new_array, + crs=self.crs, + ) + + output = self._cast_pointcloud_output(new_ds) + if self._is_dask: + # A lazy copy has the same point locations, so it can reuse the source's known count and bounds + source_attrs = _get_dataframe_attrs(self.ds) + output_attrs = _get_dataframe_attrs(output).copy() + output_attrs.update(point_count=source_attrs.get("point_count"), bounds=source_attrs.get("bounds")) + _set_dataframe_attrs(output, output_attrs) + return output @classmethod def from_xyz( @@ -454,44 +456,114 @@ def to_geoutils(self) -> Any: ds = self.ds.compute() if self._is_dask else self.ds return PointCloud(ds, data_column=self.data_column) - @overload - def get_stats( + def stats( self, - stats_name: str | Callable[[NDArrayNum], np.floating[Any]], - ) -> np.floating[Any]: ... + statistics: str | Callable[[Any], Any] | Iterable[str | Callable[[Any], Any]] | None = None, + *, + by: Mapping[str, Any] | None = None, + values: str | Iterable[str] | Mapping[str, Any] | None = None, + bins: Mapping[str, Any] | None = None, + categories: Mapping[str, Iterable[Any]] | None = None, + at: Literal["self"] | RasterLike | PointCloudLike | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, + mask_mode: Literal["inside", "outside"] = "inside", + subsample: int | float = 1, + subsample_per_group: bool = False, + random_state: int | np.random.Generator | None = None, + strategy: Literal["auto", "dense", "sparse", "groupwise"] = "auto", + backend: Literal["geoutils", "flox"] = "geoutils", + subsampling_strategy: Literal["sequential", "topk"] = "sequential", + interpolation: InterpolationMethod = "linear", + align: Literal["raise", "reproject"] = "raise", + observed: bool = True, + return_masks: bool = False, + mp_config: MultiprocConfig | None = None, + ) -> Any: + """Calculate summary statistics or statistics grouped by categories, bins, or vector zones. + + Omit ``by`` to summarize the active point values. Grouped inputs follow the same ``by``, ``categories``, + ``bins``, and vector-zone interface as Raster.stats(). + + :param statistics: Statistics to calculate (e.g. "mean", ["mean", "nmad"], or np.nanmedian). None returns + "min", "max", "mean", "median", "std", "nmad", "validcount", "totalcount" and "percentagevalidpoints". + "all" also includes "sum", "sumofsquares", "90thpercentile", "iqr", "le90" and "rmse", plus inlier counts + for masked global statistics. Grouped defaults replace "validcount" with "count"; every grouped result + includes "count". + :param by: Named variables to group by (e.g. {"elevation": elevation}); use {"glacier": (outlines, "id")} + for vector zones. Arrays must match the selected locations. Omit for global statistics. + :param values: Point columns to summarize (e.g. "height" or ["height", "intensity"]); defaults to the main + data column. Use a mapping for named inputs (e.g. {"elevation": (dem, 1)}). + :param bins: Continuous bins keyed by grouping name (e.g. {"elevation": 10}). Each definition is a count of + equal-width bins, increasing edges (e.g. [0, 2, 5]), or a Pandas IntervalIndex to choose open/closed sides. + :param categories: Ordered categories keyed by grouping name (e.g. {"landcover": [100, 110, 120]}). + Values outside these categories are excluded. + :param at: Grid or ordered point locations on which to calculate statistics (e.g. at=reference or at="self"). + Defaults to this point cloud's locations. + :param mask: Locations to include (True in a boolean mask, e.g. mask=points.data > 1000, or vector features). + Global counts describe values before this mask; "all" adds counts for values kept by the mask. + :param mask_mode: Keep locations "inside" or "outside" vector features; ignored for boolean masks. + :param subsample: Fraction (e.g. 0.1 for 10%) or maximum count (e.g. 10000) of eligible locations to use. + A value of 1 keeps all locations. Counts describe the sampled locations. + :param subsample_per_group: Apply subsample within each combined group (True, stratified sampling) or once + across all groups (False). Without by, both use one global sample. + :param random_state: Seed to reproduce subsampling (e.g. 42), or an existing random generator. + :param strategy: Combine chunk statistics for all groups ("dense"), only groups present in each chunk + ("sparse"), or gather each complete group ("groupwise"). "auto" chooses from the statistics and group count; + exact quantiles and custom functions require "auto" or "groupwise" for chunked data. + :param backend: Use the GeoUtils reducer ("geoutils") or optional Flox reducer ("flox") for grouped statistics; + see stats() for the Flox restrictions. + :param subsampling_strategy: "topk" keeps the same sampled locations across chunk layouts for a fixed seed; + "sequential" draws random locations using traversal order and can depend on the chunks. + :param interpolation: Raster values at point locations use interp_points() with SciPy methods "nearest", + "linear", "slinear", "cubic", "quintic", "pchip" or "splinef2d". + Raster groupers listed in categories use "nearest". + :param align: "raise" rejects different grids or coordinate systems; "reproject" aligns them to the output + locations. Point inputs must still share the same ordered coordinates. + :param observed: Omit declared group combinations with no eligible locations (True), or include them (False). + :param return_masks: Also return masks keyed by group labels (e.g. table, masks = points.stats(...)). + Masks cover complete groups before subsampling. Requires by. + :param mp_config: Worker and tile settings for multiprocessing, e.g. MultiprocConfig(chunks=512). + Cannot be combined with Dask inputs. + :returns: A statistic, summary dictionary, grouped dataframe, or grouped dataframe and mask mapping. + """ - @overload - def get_stats( - self, - stats_name: list[str | Callable[[NDArrayNum], np.floating[Any]]] | None = None, - ) -> dict[str, np.floating[Any]]: ... + return _stats( + self, + statistics, + by=by, + values=values, + bins=bins, + categories=categories, + at=at, + mask=mask, + mask_mode=mask_mode, + subsample=subsample, + subsample_per_group=subsample_per_group, + random_state=random_state, + strategy=strategy, + backend=backend, + subsampling_strategy=subsampling_strategy, + interpolation=interpolation, + align=align, + observed=observed, + return_masks=return_masks, + mp_config=mp_config, + ) @profiler.profile("geoutils.pointcloud.base.get_stats", memprof=True) def get_stats( self, stats_name: ( - str | Callable[[NDArrayNum], np.floating[Any]] | list[str | Callable[[NDArrayNum], np.floating[Any]]] | None + str + | Callable[[NDArrayNum], np.floating[Any]] + | Iterable[str | Callable[[NDArrayNum], np.floating[Any]]] + | None ) = None, - ) -> np.floating[Any] | dict[str, np.floating[Any]] | None: - """ - Retrieve statistics for the point-cloud values. - - :param stats_name: Statistic name, custom callable, or list of either. None returns the main statistics and - ``all`` returns every available statistic. - :returns: One value for a single statistic, or a dictionary for multiple, default or all statistics. - """ - - # Statistics return small eager values, so reduce the lazy data column here - data = self.data.compute().values if self._is_dask else np.asarray(self.data) + ) -> Any: + """Call stats() with the legacy argument names; deprecated in favor of stats().""" - if isinstance(stats_name, list) or stats_name is None or stats_name == "all": - return _statistics(data, stats_name) # type: ignore - if isinstance(stats_name, str): - return _statistics(data, [stats_name])[stats_name] # type: ignore - if callable(stats_name): - return stats_name(data) # type: ignore - warnings.warn(f"Statistic name {stats_name} is a not recognized string", category=UserWarning) - return None + warnings.warn("get_stats() is deprecated; use stats() instead.", DeprecationWarning, stacklevel=2) + return _stats(self, statistics=stats_name) @overload def subsample( @@ -500,6 +572,7 @@ def subsample( return_indices: Literal[False] = False, *, random_state: int | np.random.Generator | None = None, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum: ... @overload @@ -509,6 +582,7 @@ def subsample( return_indices: Literal[True], *, random_state: int | np.random.Generator | None = None, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None = None, ) -> tuple[NDArrayNum, ...]: ... @overload @@ -517,6 +591,8 @@ def subsample( subsample: float | int, return_indices: bool = False, random_state: int | np.random.Generator | None = None, + *, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum | tuple[NDArrayNum, ...]: ... @profiler.profile("geoutils.pointcloud.base.subsample", memprof=True) @@ -525,14 +601,22 @@ def subsample( subsample: float | int, return_indices: bool = False, random_state: int | np.random.Generator | None = None, + *, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum | tuple[NDArrayNum, ...]: """ - Randomly sample finite point-cloud values. + Randomly sample finite point cloud values allowed by mask, without replacement. - :param subsample: Fraction of values to sample when at most 1, otherwise the number of values. - :param return_indices: Whether to return sampled indexes instead of values. + :param subsample: Fraction of eligible finite values to sample when at most 1, otherwise the maximum number + of values. The mask is applied before calculating this size. + :param return_indices: Whether to return sampled row positions instead of values. :param random_state: Random generator or seed used to make sampling reproducible. - :returns: Sampled values, or sampled indexes when ``return_indices`` is True. + :param mask: Eligible points: True in a boolean array or spatial mask, or inside vector geometries. + Arrays must have one entry per point. Point masks must follow the same ordered coordinates; + raster masks use nearest interpolation. Point and raster masks must share this point cloud's CRS. + Missing mask entries are excluded (e.g. mask=points.data > 0). + :returns: One-dimensional NumPy values with the source dtype, or a one-element tuple of indices into the + original row order. These indices are positions, independent of any dataframe index labels. """ return _subsample_pointcloud( @@ -540,8 +624,331 @@ def subsample( subsample=subsample, return_indices=return_indices, random_state=random_state, + mask=mask, + ) + + def cosample( + self, + other: RasterLike | PointCloudLike | ArrayLike, + *, + other_band: int = 1, + auxiliary: Mapping[str, Any] | None = None, + auxiliary_at: Literal["self", "other"] | Mapping[str, Literal["self", "other"]] | None = None, + at: Literal["self", "other"] | RasterLike | PointCloudLike | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, + mask_mode: Literal["inside", "outside"] = "inside", + subsample: int | float = 1, + random_state: int | np.random.Generator | None = None, + strategy: Literal["sequential", "topk"] = "topk", + raster_point_mode: Literal["grid_points", "resample_raster"] | None = None, + grid_method: GriddingMethod = "linear", + resample_method: InterpolationMethod | Literal["reduce"] = "linear", + grid_kwargs: Mapping[str, Any] | None = None, + resample_kwargs: Mapping[str, Any] | None = None, + align: Literal["raise", "reproject"] = "raise", + mp_config: MultiprocConfig | None = None, + ) -> RasterLike | PointCloudLike: + """ + Sample this point cloud and another dataset at common finite locations. + + This point cloud provides the default spatial support. Use ``raster_point_mode="grid_points"`` to grid + point values onto a raster input instead. An explicit ``at`` chooses the exact output locations and must + agree with any explicit mode. Raw auxiliary arrays must identify their primary input's grid or point ordering. + + Spatial inputs, explicit output support and raster or point masks must use one family: Raster/PointCloud + objects, or DataArray/GeoDataFrame objects. The latter may mix eager and Dask storage. Plain arrays and + vector outlines are accepted with either family. + + :param other: Dataset to sample alongside this point cloud. A plain array follows this point cloud's order. + :param other_band: Band selected from other if it is a raster, counting from one. + :param auxiliary: Additional values by output name (e.g. {"slope": slope_raster}). Select a raster band with + {"slope": (slope_raster, 2)} or a point column with {"intensity": (points, "intensity")}. Spatial inputs + without a selector use the first raster band or active point values. + :param auxiliary_at: Input locations followed by plain auxiliary arrays: "self", "other", or a choice per name + (e.g. {"slope": "other"}). Spatial auxiliaries use their own coordinates. + :param at: Output locations: "self", "other", or a reference raster/point cloud. Defaults to this point cloud. + Point inputs must share the selected point order; "grid_points" selects a raster grid instead. + :param mask: Locations eligible for sampling, defined by a boolean array, spatial mask, or vector outlines. + :param mask_mode: Whether a vector mask keeps locations "inside" or "outside" its geometries. + :param subsample: Fraction of common finite locations (e.g. 0.1), or maximum count (e.g. 1000); 1 keeps all. + :param random_state: Seed or random generator for reproducible sampling (e.g. 42). + :param strategy: Raster sampling with "topk" or "sequential"; "topk" keeps the same seeded sample across chunk + sizes. Point output always uses "sequential". + :param raster_point_mode: Conversion direction: "grid_points" places points on a raster, "resample_raster" + reads rasters at points. Defaults to at's locations, or this point cloud's locations. Must agree with at. + :param grid_method: Point gridding by SciPy interpolation ("nearest", "linear", "cubic"), or circular "idw", + "mean", "minimum", "maximum", "range", "count", "stdev", "average_distance", "average_distance_pts". + The aliases "average", "min" and "max" select "mean", "minimum" and "maximum". + :param resample_method: Raster interpolation using the SciPy methods "nearest", "linear", "cubic", "quintic", + "slinear", "pchip" or "splinef2d". Window reduction ("reduce") is not implemented. + :param grid_kwargs: Options for PointCloud.grid(), e.g. {"dist_nodata_pixel": 2, "min_points": 3} sets a + two-pixel radius and minimum of three finite points for circular methods. Other options include + "distance_power" for IDW and "engine" ("scipy" or "numba"). + Set locations and method with at and grid_method. + :param resample_kwargs: Options for Raster.interp_points(), e.g. {"nodata_propagation": "ignore"}. The nodata + policies are "gdal", "ignore" and "propagate"; "dist_nodata_spread" controls extra spreading in pixels. + Set locations, band and method with the corresponding cosample() arguments. + :param align: Handling of mismatched grids or coordinate systems: "raise" an error, or "reproject" to match at. + Point inputs must still share the same ordered coordinates when sampled at points. + :param mp_config: Worker and tile settings for multiprocessing. Raster output uses its outfile; cannot be + combined with Dask inputs. + :returns: Raster or point cloud on the selected support; Xarray DataArray or eager/lazy GeoDataFrame for + accessor calls. Bands or columns contain "self", "other", then auxiliaries in mapping order. + Raster outputs retain the target grid with a common mask; point outputs retain selected geometries + and index labels, with "self" as the active data column. + """ + + from geoutils.sampling.cosampling import _cosample + + return _cosample( + self, + other, + band=1, + other_band=other_band, + auxiliary=auxiliary, + auxiliary_at=auxiliary_at, + at=at, + mask=mask, + mask_mode=mask_mode, + subsample=subsample, + random_state=random_state, + strategy=strategy, + raster_point_mode=raster_point_mode, + grid_method=grid_method, + resample_method=resample_method, + grid_kwargs=grid_kwargs, + resample_kwargs=resample_kwargs, + align=align, + mp_config=mp_config, + ) + + def pairsample( + self, + *, + n_pairs: int = 1_000_000, + sampling: Literal["loglag", "random_xy"] = "loglag", + min_distance: float | None = None, + max_distance: float | None = None, + random_state: int | np.random.Generator | None = None, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None = None, + strategy: Literal["kdtree", "hashgrid", "nn_logvector"] = "nn_logvector", + n_bins: int = 24, + anchors_per_round: int = 50_000, + attempts_per_anchor: int = 1, + max_rounds: int = 50, + cell_size: float | None = None, + nn_tolerance: float = 0.1, + nn_batch_size: int = 250_000, + nn_oversample: float = 2.0, + nn_max_batches: int = 200, + index_dtype: DTypeLike = np.int32, + distance_dtype: DTypeLike = np.float32, + ) -> xr.Dataset: + """Sample finite point pairs for statistics by distance. + + Exact ring strategies use a KD-tree or hash grid. ``"nn_logvector"`` proposes isotropic log-spaced vectors + and accepts a nearby observed endpoint, which is generally faster for large point clouds. + + Strategy controls apply to ``"loglag"``. ``"random_xy"`` uses ``max_rounds`` and ``nn_batch_size``. + Dask point tables are loaded because the search requires all coordinates. + + :param n_pairs: Requested number of pairs with two finite values; fewer may be returned if sampling stops early. + :param sampling: ``"loglag"`` balances short and long distances on a log scale; ``"random_xy"`` draws + endpoints uniformly. + :param min_distance: Smallest distance in CRS units (e.g. meters). Defaults to half the spacing estimated + from the eligible point density. + :param max_distance: Largest distance in CRS units. Defaults to the bounding box diagonal of eligible points. + :param random_state: Seed for reproducible sampling (e.g. 42). + :param mask: Eligible points: True in a boolean array or spatial mask, or inside vector geometries. + Arrays must have one entry per point. Point masks must follow the same ordered coordinates; + raster masks use nearest interpolation. Point and raster masks must share this point cloud's CRS. + Missing mask entries are excluded. + :param strategy: GeoUtils log-lag strategy: ``"kdtree"`` uses SciPy to search distance rings, ``"hashgrid"`` + searches rings using a spatial grid, and ``"nn_logvector"`` uses SciPy to match proposed endpoints + to nearby points. + :param n_bins: Log-spaced distance rings used by ``"kdtree"`` and ``"hashgrid"`` (e.g. 24). + :param anchors_per_round: First endpoints tested per round by ``"kdtree"`` and ``"hashgrid"``. + :param attempts_per_anchor: Distance rings tried per first endpoint by ``"kdtree"`` and ``"hashgrid"``. + :param max_rounds: Maximum rounds to fill the sample with ``"kdtree"``, ``"hashgrid"``, or ``"random_xy"``. + :param cell_size: Grid cell width in CRS units for ``"hashgrid"``. Defaults to one eighth of max_distance. + :param nn_tolerance: Allowed endpoint snap distance for ``"nn_logvector"``, as a fraction of the proposed + pair distance (e.g. 0.1 allows a 10% offset). + :param nn_batch_size: Maximum candidate pairs per batch with ``"nn_logvector"`` or ``"random_xy"``; + smaller batches use less temporary memory. + :param nn_oversample: Candidate count as a multiple of the remaining pairs with ``"nn_logvector"`` (e.g. 2). + :param nn_max_batches: Maximum batches to fill the sample with ``"nn_logvector"``. + :param index_dtype: Integer NumPy dtype for returned row indexes (e.g. ``"int64"`` for very large point clouds). + :param distance_dtype: Floating NumPy dtype for returned distances (e.g. ``"float64"`` for greater precision). + :returns: Xarray Dataset with pair and endpoint dimensions, containing original row indexes, values, + coordinates, and distances. + """ + + from geoutils.sampling.pairsampling import _sample_point_pairs + + return _sample_point_pairs( + self, + n_pairs=n_pairs, + sampling=sampling, + min_distance=min_distance, + max_distance=max_distance, + random_state=random_state, + mask=mask, + strategy=strategy, + n_bins=n_bins, + anchors_per_round=anchors_per_round, + attempts_per_anchor=attempts_per_anchor, + max_rounds=max_rounds, + cell_size=cell_size, + nn_tolerance=nn_tolerance, + nn_batch_size=nn_batch_size, + nn_oversample=nn_oversample, + nn_max_batches=nn_max_batches, + index_dtype=index_dtype, + distance_dtype=distance_dtype, + ) + + def variogram( + self, + *, + n_pairs: int = 1_000_000, + sampling: Literal["loglag", "random_xy"] = "loglag", + estimator: str | Callable[[NDArrayNum], float] = "dowd", + bins: Literal["log", "uniform"] | Iterable[float] = "log", + n_lags: int = 24, + min_lag: float | None = None, + max_lag: float | None = None, + n_runs: int = 1, + model: str | Callable[..., Any] | list[str | Callable[..., Any]] | None = None, + fit_kwargs: dict[str, Any] | None = None, + random_state: int | np.random.Generator | None = None, + mask: VectorLike | ArrayLike | None = None, + **pair_sampling_kwargs: Any, + ) -> Variogram: + """Estimate a lightweight empirical variogram from point pairs. + + :param n_pairs: Number of finite pairs targeted in each run (e.g. 100_000). + :param sampling: How to select pairs: ``"loglag"`` balances short and long distances, while ``"random_xy"`` + selects endpoints independently. + :param estimator: Semivariance estimator from SciKit-GStat: ``"dowd"``, ``"matheron"``, ``"cressie"``, + ``"genton"``, ``"minmax"``, ``"entropy"`` or ``"percentile"``. A function can instead map absolute pair + differences to one value per distance bin. + :param bins: Distance bins: ``"log"`` for logarithmic spacing, ``"uniform"`` for equal widths, or explicit + edges (e.g. [1, 10, 100]). + :param n_lags: Number of distance bins when bins is ``"log"`` or ``"uniform"``. + :param min_lag: Minimum sampled distance in CRS units; defaults to half the spacing estimated from density. + :param max_lag: Maximum sampled distance in CRS units; defaults to the extent diagonal of eligible points. + :param n_runs: Independent samples to average; repeat sampling to estimate each distance bin's standard error. + :param model: SciKit-GStat model to fit: ``"spherical"``, ``"exponential"``, ``"gaussian"``, ``"cubic"``, + ``"stable"`` or ``"matern"``, or the corresponding model function. Sum a list of models ordered from short + to long range (e.g. ["spherical", "exponential"]). ``None`` keeps only the empirical variogram. + :param fit_kwargs: Options for Variogram.fit(): ``use_nugget``, ``bounds``, ``p0`` or ``maxfev`` + (e.g. {"use_nugget": True}); optimization uses SciPy curve_fit(). + :param random_state: Seed or NumPy generator for reproducible sampling across runs (e.g. 42). + :param mask: Points to keep: True values in a boolean mask or points inside vector geometries. + :param pair_sampling_kwargs: Extra pairsample() options (e.g. ``strategy`` or ``max_rounds``). + :returns: Variogram with distance bins, pair counts and semivariance, plus sampling errors and a fitted model + when requested. + """ + + return _variogram( + self, + n_pairs=n_pairs, + sampling=sampling, + n_runs=n_runs, + estimator=estimator, + bins=bins, + n_lags=n_lags, + min_lag=min_lag, + max_lag=max_lag, + model=model, + fit_kwargs=fit_kwargs, + random_state=random_state, + mask=mask, + **pair_sampling_kwargs, ) + @overload + def reproject( + self: PointCloudBaseType, + ref: RasterLike | VectorLike | None = None, + crs: CRS | str | int | None = None, + *, + inplace: Literal[False] = False, + mp_config: MultiprocConfig | None = None, + ) -> PointCloudBaseType | gpd.GeoDataFrame: ... + + @overload + def reproject( + self: PointCloudBaseType, + ref: RasterLike | VectorLike | None = None, + crs: CRS | str | int | None = None, + *, + inplace: Literal[True], + mp_config: MultiprocConfig | None = None, + ) -> None: ... + + @overload + def reproject( + self: PointCloudBaseType, + ref: RasterLike | VectorLike | None = None, + crs: CRS | str | int | None = None, + *, + inplace: bool = False, + mp_config: MultiprocConfig | None = None, + ) -> PointCloudBaseType | gpd.GeoDataFrame | None: ... + + @profiler.profile("geoutils.pointcloud.base.reproject", memprof=True) + def reproject( + self: PointCloudBaseType, + ref: RasterLike | VectorLike | None = None, + crs: CRS | str | int | None = None, + inplace: bool = False, + *, + mp_config: MultiprocConfig | None = None, + ) -> PointCloudBaseType | gpd.GeoDataFrame | None: + """ + Reproject point coordinates, preserving their order and value columns. + + Without multiprocessing, eager inputs return eager results and Dask inputs remain lazy. Multiprocessing + reads and writes row partitions, keeping file-backed PointCloud inputs and results unloaded. LAS/LAZ + output rounds coordinates to its stored precision; GeoPackage preserves floating-point coordinates. + Reopened indices follow the file format. LAS attributes must fit their dimension types; GeoPackage + requires millisecond timestamps and nullable integers that remain exact when read as float64. + + :param ref: Raster or vector whose CRS should be matched; mutually exclusive with ``crs``. + :param crs: Target coordinate reference system; mutually exclusive with ``ref``. + :param inplace: Update this object for eager execution. Unsupported with Dask or multiprocessing. + :param mp_config: Worker configuration with an integer number of points per chunk. The output format is + inferred from ``outfile`` or selected by ``driver`` (``GPKG``, ``LAS`` or ``LAZ``), defaulting to + GeoPackage. Cannot be combined with Dask input. + :returns: Reprojected PointCloud or GeoDataFrame matching the input interface, or None when in place. + Multiprocessing PointCloud results are unloaded; dataframe accessor results are eager. + """ + + # Keep the shared vector implementation for eager and lazy dataframe transformations + if mp_config is None: + return super().reproject(ref=ref, crs=crs, inplace=inplace) + if self._is_dask: + raise ValueError("Argument ``mp_config`` cannot be combined with a Dask point cloud.") + if inplace: + raise ValueError("Argument ``inplace`` is not supported with ``mp_config``; use the returned point cloud.") + + # Resolve the target without reading point data, then let workers build the output file + from geoutils.pointcloud.transformation import _reproject_pointcloud + + target_crs = _get_reproject_crs(ref=ref, crs=crs) + projected = _reproject_pointcloud(self, crs=target_crs, mp_config=mp_config) + if self._is_pd: + # Read every output attribute and use native LAS Z when the file represents heights as a column + projected.load(columns="all") + return _build_pointcloud_output( + projected.ds, + data_column=projected.data_column, + as_dataframe=True, + attrs=_get_dataframe_attrs(self.ds), + ) + return cast(PointCloudBaseType, projected) + @profiler.profile("geoutils.pointcloud.base.grid", memprof=True) def grid( self, @@ -554,6 +961,7 @@ def grid( dist_nodata_pixel: float = 1.0, nodata: int | float = -9999, *, + data_column: str | None = None, distance_power: float = 2.0, min_points: int = 1, engine: GriddingEngine = "scipy", @@ -568,7 +976,7 @@ def grid( Define the output grid with a reference raster, regular X/Y coordinates, or a combination of resolution or shape and optional bounds. - :param ref: Reference raster whose grid should be matched. + :param ref: Reference raster whose grid should be matched. A Dask reference also selects lazy output. :param grid_coords: Regular X and Y coordinates defining the output grid. :param res: Output resolution in X and Y, mutually exclusive with ``shape``. :param shape: Output shape as ``(height, width)``, mutually exclusive with ``res``. @@ -577,6 +985,7 @@ def grid( aliases for ``mean``, ``minimum`` and ``maximum``. :param dist_nodata_pixel: Maximum point distance or circular neighborhood radius in output pixels. :param nodata: Nodata value of the output raster. + :param data_column: Point value column to grid. None uses the active point values. :param distance_power: Distance exponent used for inverse-distance weighting. :param min_points: Minimum number of finite points required inside a circular neighborhood. :param engine: Calculation engine, either ``scipy`` or ``numba``. @@ -598,6 +1007,7 @@ def grid( resampling=resampling, dist_nodata_pixel=dist_nodata_pixel, nodata=nodata, + data_column=data_column, distance_power=distance_power, min_points=min_points, engine=engine, diff --git a/geoutils/pointcloud/dataframe.py b/geoutils/pointcloud/dataframe.py new file mode 100644 index 000000000..6473a0c5e --- /dev/null +++ b/geoutils/pointcloud/dataframe.py @@ -0,0 +1,346 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Module for dataframe operations: managing values, row selection, metadata and output construction.""" + +from __future__ import annotations + +import warnings +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload + +import geopandas as gpd +import numpy as np +import pandas as pd + +from geoutils._dispatch import is_dask_array, is_dask_dataframe +from geoutils._misc import import_optional + +if TYPE_CHECKING: + from geoutils.pointcloud.pointcloud import PointCloudLike + + +DataFrameType = TypeVar("DataFrameType") + + +############################################ +# 1/ METADATA AND POINT CLOUD OUTPUTS +############################################ + + +def _get_dataframe_attrs(ds: Any) -> dict[str, Any]: + """Get GeoUtils metadata from Pandas or Dask dataframes.""" + + # Dask does not carry Pandas ``attrs`` reliably through graph operations + if is_dask_dataframe(ds): + try: + return object.__getattribute__(ds, "_geoutils_attrs") + except AttributeError: + return {} + return getattr(ds, "attrs", {}) + + +def _set_dataframe_attrs(ds: Any, attrs: dict[str, Any]) -> None: + """Set GeoUtils metadata on Pandas or Dask dataframes.""" + + # Keep a private copy on Dask collections and use the public mapping for Pandas + if is_dask_dataframe(ds): + object.__setattr__(ds, "_geoutils_attrs", attrs.copy()) + elif hasattr(ds, "attrs"): + ds.attrs.update(attrs) + + +@overload +def _build_pointcloud_output( + dataframe: DataFrameType, + *, + data_column: str | None, + as_dataframe: Literal[True], + attrs: Mapping[str, Any] | None = None, + preserve_locations: bool = False, +) -> DataFrameType: ... + + +@overload +def _build_pointcloud_output( + dataframe: Any, + *, + data_column: str | None, + as_dataframe: bool, + attrs: Mapping[str, Any] | None = None, + preserve_locations: bool = False, +) -> PointCloudLike: ... + + +def _build_pointcloud_output( + dataframe: Any, + *, + data_column: str | None, + as_dataframe: bool, + attrs: Mapping[str, Any] | None = None, + preserve_locations: bool = False, +) -> PointCloudLike: + """ + Build a point cloud result with metadata matching its current rows and coordinate system. + + _set_dataframe_attrs() records the active value column, CRS and point geometry type without changing supplied + metadata. Row selections may change counts and bounds, so bounds and lazy counts are cleared unless + preserve_locations is True, meaning the result has the same ordered points and unchanged X/Y coordinates. + Eager results always receive a fresh count. Return the eager or Dask dataframe when as_dataframe is True; + otherwise compute Dask rows before constructing a PointCloud. + """ + + # Copy supplied metadata and load Dask rows only when the caller needs a PointCloud object + metadata = {} if attrs is None else dict(attrs) + if not as_dataframe and is_dask_dataframe(dataframe): + dataframe = dataframe.compute() + + # Reuse spatial metadata only when the caller guarantees unchanged point rows and X/Y coordinates + point_count = metadata.get("point_count") if preserve_locations else None + if not is_dask_dataframe(dataframe): + point_count = len(dataframe) + bounds = metadata.get("bounds") if preserve_locations else None + metadata.update( + data_column=data_column, + geometry_type="Point", + crs=dataframe.crs, + point_count=point_count, + bounds=bounds, + ) + _set_dataframe_attrs(dataframe, metadata) + + # Accessors return dataframes; object callers receive an eager PointCloud with the selected value column + if as_dataframe: + return dataframe + + from geoutils.pointcloud.pointcloud import PointCloud + + return PointCloud(dataframe, data_column=data_column) + + +############################################ +# 2/ POINT DATAFRAME PARTITIONS +############################################ + + +def _import_dask_dataframe() -> Any: + """Import Dask DataFrame while suppressing optional dask-expr warnings from older environments.""" + + # Delay the optional import until a lazy dataframe operation is requested + import_optional("dask") + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=FutureWarning, module="dask.dataframe") + import dask.dataframe as dd + + return dd + + +def _point_partition_lengths(dataframe: Any) -> tuple[int, ...]: + """ + Compute the row count of each point partition without collecting its rows or attached values. + + Dataframe follows _assign_point_values(). The caller can reuse these counts while the row layout stays unchanged. + + :returns: One integer row count per partition, in partition order. + """ + + return tuple(int(length) for length in dataframe.geometry.map_partitions(len).compute()) + + +def _point_array_partitions( + dataframe: Any, arrays: Sequence[Any], *, partition_lengths: Sequence[int] | None = None +) -> list[Any]: + """ + Align one-dimensional value arrays with the row partitions of a lazy point dataframe. + + Dataframe, arrays and partition_lengths follow _assign_point_values(). Only partition lengths are collected here; + the point rows and corresponding values remain lazy. Dask Series keep each partition's index, including duplicates. + """ + + import_optional("dask") + import dask.array as da + + # Count rows per partition so array chunks can be matched by position, independently of index labels + lengths = _point_partition_lengths(dataframe) if partition_lengths is None else tuple(partition_lengths) + if len(lengths) != dataframe.npartitions or any(length < 0 for length in lengths): + raise ValueError("Point partition lengths must give one nonnegative row count per dataframe partition.") + total = sum(lengths) + dd = _import_dask_dataframe() + columns = [] + for values in arrays: + array = da.asarray(values) + if any(np.isnan(length) for length in array.shape): + array = array.compute_chunk_sizes() + if array.ndim != 1 or array.shape[0] != total: + raise ValueError("Point values must contain one value per dataframe row.") + + # Keep every empty point partition; Dask otherwise collapses a zero-length array to one block + if total == 0: + columns.append(dataframe.index.to_series().astype(array.dtype)) + continue + + # Give each value partition the exact index of its corresponding point partition + aligned = array.rechunk((lengths,)) + columns.append(dd.from_dask_array(aligned, index=dataframe.index)) + return columns + + +############################################ +# 3/ ASSIGN VALUES BY POINT POSITION +############################################ + + +def _assign_point_partition(dataframe: gpd.GeoDataFrame, *columns: Any, names: Sequence[str]) -> gpd.GeoDataFrame: + """Assign the matching value partitions prepared by _assign_point_values(), ignoring their index labels.""" + + # Assign directly so user column names such as self do not collide with dataframe.assign() arguments + result = dataframe.copy() + for name, column in zip(names, columns): + array = np.asarray(column) + if array.ndim != 1 or len(array) != len(dataframe): + raise ValueError("Point value partitions must contain one value per dataframe row.") + result[name] = array + return result + + +def _assign_point_values( + dataframe: Any, values: Mapping[str, Any], *, partition_lengths: Sequence[int] | None = None +) -> Any: + """ + Attach named values to point rows by position while keeping the dataframe's eager or lazy backend. + + _point_array_partitions() matches lazy values to point partitions. _assign_point_partition() assigns each + partition by position, avoiding index alignment when several points have the same row label. + Dask Series already following the point partitions pass through without computing their lengths. + + :param dataframe: GeoDataFrame or Dask-GeoPandas dataframe whose ordered point rows define the value locations. + :param values: Mapping of column names to one-dimensional NumPy or Dask arrays, with one value per point row. + Dask Series may also be passed when their partitions contain the same ordered rows as the dataframe. + :param partition_lengths: Optional known row counts for the current dataframe partitions, reused when aligning + array inputs. These counts must be updated after selecting rows. Matching Dask Series do not need them. + :returns: A dataframe with the named columns assigned and its existing columns, point order and labels preserved. + """ + + # Keep the original table when there are no additional values to attach + if not values: + return dataframe + + # Preserve eager output even when a selected value array was computed lazily + if not is_dask_dataframe(dataframe): + columns = list(values.values()) + if any(is_dask_array(array) or is_dask_dataframe(array) for array in columns): + import_optional("dask") + import dask + + columns = list(dask.compute(*columns)) + return _assign_point_partition(dataframe, *columns, names=list(values)) + # Use point Series directly and align only array inputs, sharing one partition count for every array + array_names = [name for name, value in values.items() if not is_dask_dataframe(value)] + aligned = {} + if array_names: + arrays = _point_array_partitions( + dataframe, [values[name] for name in array_names], partition_lengths=partition_lengths + ) + aligned = dict(zip(array_names, arrays)) + columns = [] + for name, value in values.items(): + if is_dask_dataframe(value): + if value.ndim != 1 or value.npartitions != dataframe.npartitions: + raise ValueError("Point value Series must follow the dataframe's partition layout.") + columns.append(value) + else: + columns.append(aligned[name]) + + # Attach columns positionally within each corresponding partition and preserve their numeric metadata + meta = dataframe._meta.copy() + for name, column in zip(values, columns): + meta[name] = pd.Series([], dtype=column.dtype) + return dataframe.map_partitions(_assign_point_partition, *columns, names=list(values), meta=meta) + + +############################################ +# 4/ SELECT POINT ROWS +############################################ + + +def _select_point_partition( + dataframe: gpd.GeoDataFrame, + indices: Any, + *, + starts: Any = None, + partition_info: dict[str, Any] | None = None, +) -> gpd.GeoDataFrame: + """ + Select point rows from one partition using the mask or global positions supplied to _select_point_rows(). + + :param starts: Cumulative partition row counts for integer selection, or None for a matching boolean mask. + :param partition_info: Dask partition metadata identifying which global row offset to use. + """ + + # Boolean masks already match this partition; integer indices refer to the full ordered dataframe + if starts is None: + mask = np.asarray(indices) + if mask.ndim != 1 or len(mask) != len(dataframe): + raise ValueError("Point mask partitions must contain one value per dataframe row.") + return dataframe.iloc[mask] + assert partition_info is not None + start, stop = starts[partition_info["number"] : partition_info["number"] + 2] + lower, upper = np.searchsorted(indices, (start, stop)) + return dataframe.iloc[indices[lower:upper] - start] + + +def _select_point_rows(dataframe: Any, indices: Any, *, partition_lengths: Sequence[int] | None = None) -> Any: + """ + Select point rows by position without collecting a complete lazy point dataframe. + + _point_array_partitions() aligns boolean masks with point partitions. For integer positions, partition lengths + locate the requested rows and _select_point_partition() selects them with iloc(), independently of row labels. + + Dataframe and partition_lengths follow _assign_point_values(). + + :param indices: One-dimensional boolean NumPy or Dask mask, or sorted integer NumPy positions to keep. + A boolean Dask Series must follow the dataframe's partition layout and avoids computing partition lengths. + :returns: Selected rows in their original order, preserving the dataframe backend, geometry and index labels. + """ + + # Eager point supports keep eager results, including when their eligibility mask was generated lazily + if not is_dask_dataframe(dataframe): + selected = indices.compute() if is_dask_array(indices) or is_dask_dataframe(indices) else indices + return dataframe.iloc[np.asarray(selected)] + + # Match boolean mask partitions without collecting the full eligibility array + if indices.dtype == np.bool_: + if is_dask_dataframe(indices): + if indices.ndim != 1 or indices.npartitions != dataframe.npartitions: + raise ValueError("Point mask Series must follow the dataframe's partition layout.") + mask = indices + else: + (mask,) = _point_array_partitions(dataframe, [indices], partition_lengths=partition_lengths) + return dataframe.map_partitions(_select_point_partition, mask, meta=dataframe._meta) + + # Translate global positions into local rows using only small partition-length summaries + positions = np.asarray(indices, dtype=np.int64) + if positions.ndim != 1 or np.any(positions[1:] < positions[:-1]): + raise ValueError("Point row positions must be a sorted one-dimensional integer array.") + lengths = _point_partition_lengths(dataframe) if partition_lengths is None else tuple(partition_lengths) + if len(lengths) != dataframe.npartitions or any(length < 0 for length in lengths): + raise ValueError("Point partition lengths must give one nonnegative row count per dataframe partition.") + starts = np.concatenate(([0], np.cumsum(lengths))) + if len(positions) and (positions[0] < 0 or positions[-1] >= starts[-1]): + raise IndexError("Point row position is outside the dataframe.") + return dataframe.map_partitions(_select_point_partition, positions, starts=starts, meta=dataframe._meta) diff --git a/geoutils/pointcloud/las.py b/geoutils/pointcloud/las.py index 1b5d792ce..192d6c4bc 100644 --- a/geoutils/pointcloud/las.py +++ b/geoutils/pointcloud/las.py @@ -288,7 +288,8 @@ def _load_laspy_data_slice( # Seek directly to the requested row range instead of reading earlier points with laspy.open(filename) as reader: crs = reader.header.parse_crs(prefer_wkt=False) - reader.seek(start) + if count > 0: + reader.seek(start) points = reader.read_points(count) return _laspy_points_to_geodataframe(points=points, crs=crs, columns=columns) diff --git a/geoutils/pointcloud/pd_accessor.py b/geoutils/pointcloud/pd_accessor.py index a91df1808..97fcef098 100644 --- a/geoutils/pointcloud/pd_accessor.py +++ b/geoutils/pointcloud/pd_accessor.py @@ -33,13 +33,13 @@ from geoutils._dispatch import is_dask_dataframe, is_dask_geodataframe from geoutils._misc import import_optional -from geoutils.pointcloud.base import ( - PointCloudBase, +from geoutils.pointcloud.base import PointCloudBase +from geoutils.pointcloud.dataframe import ( _get_dataframe_attrs, + _import_dask_dataframe, _set_dataframe_attrs, ) from geoutils.pointcloud.las import ( - _empty_las_geodataframe, _is_laspy_supported, _load_laspy_data_slice, _load_laspy_metadata, @@ -56,18 +56,6 @@ _DASK_ACCESSOR_REGISTERED = False -def _import_dask_dataframe() -> Any: - """Import Dask DataFrame while suppressing optional dask-expr warnings from older environments.""" - - # Delay the optional import until lazy LAS partitions are requested - import_optional("dask") - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=FutureWarning, module="dask.dataframe") - import dask.dataframe as dd - - return dd - - def _register_dask_pointcloud_accessor() -> None: """Register the ``pc`` accessor on Dask DataFrames lazily.""" @@ -78,6 +66,7 @@ def _register_dask_pointcloud_accessor() -> None: return # Register only after Dask is available so normal imports remain lightweight + # https://docs.dask.org/en/stable/dataframe-extend.html#accessors import_optional("dask") with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=FutureWarning, module="dask.dataframe") @@ -92,9 +81,10 @@ def _register_dask_pointcloud_accessor() -> None: def _infer_data_column(ds: Any) -> str | None: """Infer a point cloud data column from dataframe metadata and columns.""" - data_column = _get_dataframe_attrs(ds).get("data_column") - if data_column is not None: - return data_column + attrs = _get_dataframe_attrs(ds) + if "data_column" in attrs: + # An explicit None selects elevation from 3D geometry, even when auxiliary columns exist + return attrs["data_column"] nongeo_columns = [c for c in ds.columns if c != "geometry"] if "Z" in nongeo_columns: @@ -238,9 +228,14 @@ def open_pointcloud( ) for start in starts ] + # Read zero records to preserve each native LAS dtype in Dask metadata and empty results + empty = _load_laspy_data_slice_dataframe(filename, columns_to_load, start=0, count=0) + if not parts: + parts = [delayed(_load_laspy_data_slice_dataframe)(filename, columns_to_load, start=0, count=0)] + # Keep LAS numeric dimensions unchanged while assembling the lazy dataframe with dask.config.set({"dataframe.convert-string": False}): - ddf = dd.from_delayed(parts, meta=_empty_las_geodataframe(columns_to_load, crs=metadata.crs)) + ddf = dd.from_delayed(parts, meta=empty) # Add geospatial behavior and cache header metadata for accessor properties ddf = dgpd.from_dask_dataframe(ddf, geometry="geometry") @@ -332,7 +327,8 @@ def crs(self) -> CRS: """Coordinate reference system of the point cloud.""" if self._is_dask: - return _get_dataframe_attrs(self.ds).get("crs") + # Direct Dask-GeoPandas construction retains CRS in geometry metadata without a GeoUtils cache + return _get_dataframe_attrs(self.ds).get("crs", self.ds.crs) return self.ds.crs @property diff --git a/geoutils/pointcloud/pointcloud.py b/geoutils/pointcloud/pointcloud.py index ef100e26e..d30978db1 100644 --- a/geoutils/pointcloud/pointcloud.py +++ b/geoutils/pointcloud/pointcloud.py @@ -67,7 +67,7 @@ # This is a generic Vector-type (if subclasses are made, this will change appropriately) PointCloudType = TypeVar("PointCloudType", bound="PointCloud") -PointCloudLike = Union["PointCloud", gpd.GeoDataFrame] +PointCloudLike = Union[PointCloudBase, gpd.GeoDataFrame] # List of NumPy "array" functions that are handled. # Note: all universal function are supported: https://numpy.org/doc/stable/reference/ufuncs.html @@ -1095,7 +1095,7 @@ def info(self, verbose: bool = True, stats: bool = False) -> None | str: if stats: as_str_split.append("\nStatistics:") - statistics = self.get_stats() + statistics = self.stats() # Determine the maximum length of the stat names for alignment max_len = max(len(name) for name in statistics.keys()) diff --git a/geoutils/pointcloud/transformation.py b/geoutils/pointcloud/transformation.py new file mode 100644 index 000000000..a36c62ada --- /dev/null +++ b/geoutils/pointcloud/transformation.py @@ -0,0 +1,333 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reproject point clouds in independent row partitions with file outputs.""" + +from __future__ import annotations + +import os +import pathlib +import tempfile +from typing import TYPE_CHECKING, Any + +import geopandas as gpd +import numpy as np +import pandas as pd +import pyogrio +from pyproj import CRS + +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.pointcloud.las import ( + _as_geodataframe, + _build_laspy_header, + _dataframe_to_lasdata, + _is_laspy_supported, + _load_laspy_data_slice, + _point_partition_size, + _stitch_laspy_files, +) + +if TYPE_CHECKING: + from geoutils.multiproc import MultiprocConfig + from geoutils.pointcloud.base import PointCloudBase + from geoutils.pointcloud.pointcloud import PointCloud + + +############################################ +# 1/ ROW PARTITION TRANSFORMATION +############################################ + + +def _reproject_pointcloud_partition( + source: gpd.GeoDataFrame | pathlib.Path, + columns: list[str], + start: int, + count: int, + crs: CRS, + filename: pathlib.Path, + las_output: bool, +) -> tuple[pathlib.Path, NDArrayNum | None]: + """Read and reproject one row partition, staging its exact dataframe and LAS coordinate bounds.""" + + # Read independent row ranges so unloaded sources stay outside the parent process + if isinstance(source, pathlib.Path): + if _is_laspy_supported(source): + dataframe = _load_laspy_data_slice(source, columns=columns, start=start, count=count) + else: + dataframe = pyogrio.read_dataframe(source, skip_features=start, max_features=count) + else: + dataframe = source + projected = dataframe.to_crs(crs) + + # Use actual elevations for LAS, independently of the selected point value column + bounds = None + if las_output and len(projected) > 0: + if projected.geometry.has_z.all(): + elevation = projected.geometry.z.to_numpy() + if "Z" in projected.columns: + if not np.array_equal(projected["Z"].to_numpy(), elevation): + raise ValueError("LAS output cannot store different values in geometry Z and column 'Z'.") + projected = projected.drop(columns="Z") + elif "Z" in projected.columns: + elevation = projected["Z"].to_numpy() + else: + raise ValueError("LAS output requires 3D point geometry or a native elevation column named 'Z'.") + + # Derive one coordinate encoding from every projected partition before writing LAS records + coordinates = np.column_stack((projected.geometry.x, projected.geometry.y, elevation)) + if not np.isfinite(coordinates).all(): + raise ValueError("LAS output requires finite X, Y and Z coordinates.") + bounds = np.stack((coordinates.min(axis=0), coordinates.max(axis=0))) + + # Preserve all dataframe dtypes until the final format is selected; only paths return to the parent + projected.to_pickle(filename) + return filename, bounds + + +############################################ +# 2/ ORDERED OUTPUT CONSTRUCTION +############################################ + + +def _check_gpkg_attributes(dataframe: gpd.GeoDataFrame) -> None: + """Reject attributes that GeoPackage storage or its dataframe reader would round.""" + + for name in dataframe.columns: + if name == dataframe.geometry.name: + continue + values = dataframe[name] + + # GeoPackage timestamps store milliseconds, so finer source times cannot round-trip exactly + if pd.api.types.is_datetime64_any_dtype(values.dtype): + submillisecond = (values.dt.microsecond % 1000 != 0) | (values.dt.nanosecond != 0) + if (values.notna() & submillisecond).any(): + raise ValueError(f"GeoPackage cannot preserve submillisecond timestamps in column {name!r}.") + + # GeoPandas reads integer columns containing nulls as float64, even when nulls occur in another partition + if pd.api.types.is_integer_dtype(values.dtype) and values.hasnans: + valid = values.dropna() + try: + restored = valid.astype(np.float64).astype(valid.dtype) + exact = restored.equals(valid) + except (TypeError, ValueError, OverflowError): + exact = False + if not exact: + raise ValueError(f"GeoPackage cannot preserve nullable integer values in column {name!r}.") + + +def _reproject_las_header( + dataframe: gpd.GeoDataFrame, + bounds: list[NDArrayNum], + crs: CRS, + source_filename: pathlib.Path | None, +) -> Any: + """Build a shared LAS schema and coordinate encoding from projected bounds and source dimensions.""" + + # Keep native source dimensions and elevation precision when a LAS header is available + source_header = None + if source_filename is not None and _is_laspy_supported(source_filename) and source_filename.exists(): + laspy = import_optional("laspy") + with laspy.open(source_filename) as reader: + source_header = reader.header.copy() + scales = np.array([1e-8, 1e-8, 1e-3] if crs.is_geographic else [1e-3, 1e-3, 1e-3]) + unchanged_axes = [2] + if source_header is not None: + scales[2] = source_header.scales[2] + if source_header.parse_crs() == crs: + scales = source_header.scales.copy() + unchanged_axes = [0, 1, 2] + + # Center integer coordinates on the complete output extent and enlarge scales only to avoid overflow + offsets = np.zeros(3) + if bounds: + minimum = np.min([part[0] for part in bounds], axis=0) + maximum = np.max([part[1] for part in bounds], axis=0) + offsets = minimum + (maximum - minimum) / 2 + required_scales = (maximum - minimum) / (2 * (np.iinfo(np.int32).max - 1)) + scales = np.maximum(scales, required_scales) + if source_header is not None: + # Keep unchanged coordinates on their original integer lattice whenever their current extent fits + for axis in unchanged_axes: + original_offset = source_header.offsets[axis] + original_scale = source_header.scales[axis] + if not bounds: + offsets[axis] = original_offset + scales[axis] = original_scale + continue + extrema = np.array([minimum[axis], maximum[axis]]) + encoded_extrema = np.rint((extrema - original_offset) / original_scale) + if np.all(encoded_extrema >= np.iinfo(np.int32).min) and np.all(encoded_extrema <= np.iinfo(np.int32).max): + offsets[axis] = original_offset + scales[axis] = original_scale + + # Reuse the common writer schema, with native Z chosen from geometry or the LAS elevation column + elevation_column = "Z" if "Z" in dataframe.columns else None + return _build_laspy_header( + dataframe, + data_column=elevation_column, + version=None if source_header is None else source_header.version, + point_format=None if source_header is None else source_header.point_format, + offsets=tuple(offsets), + scales=tuple(scales), + crs=crs, + ) + + +def _write_reprojected_las_partition(filename: pathlib.Path, output_filename: pathlib.Path, header: Any) -> str: + """Encode a projected partition and reject attribute values changed by its LAS dimension types.""" + + # Encode one partition with the common writer's conversion before publishing any point records + dataframe = pd.read_pickle(filename) + elevation_column = "Z" if "Z" in dataframe.columns else None + try: + encoded = _dataframe_to_lasdata(dataframe, data_column=elevation_column, header=header) + except OverflowError as error: + raise ValueError("LAS output cannot preserve point attributes with the selected dimension types.") from error + + # Check the encoded attributes for integer truncation, overflow and scaled dimension rounding + columns = [column for column in dataframe.columns if column not in (dataframe.geometry.name, "Z")] + for column in columns: + expected_values = dataframe[column].to_numpy() + encoded_values = np.asarray(encoded[column]) + + # Compare Python scalars so NumPy cannot round large integer values while promoting mixed numeric dtypes + equal_values = expected_values.astype(object) == encoded_values.astype(object) + equal_values |= pd.isna(expected_values) & pd.isna(encoded_values) + if not np.all(equal_values): + raise ValueError(f"LAS output cannot preserve the values in column {column!r} with its dimension type.") + encoded.write(output_filename) + return os.fspath(output_filename) + + +############################################ +# 3/ MULTIPROCESSING REPROJECTION +############################################ + + +def _reproject_pointcloud(source: PointCloudBase, crs: CRS, mp_config: MultiprocConfig) -> PointCloud: + """ + Reproject independent row partitions and return an unopened point cloud at the configured output path. + + _reproject_pointcloud_partition() reads source slices or receives eager rows, applies GeoPandas to_crs(), and + stages exact projected dataframes. GPKG output appends these partitions in source order. LAS output first uses + _reproject_las_header() to choose common scales and offsets, then _write_reprojected_las_partition() and + _stitch_laspy_files() encode and stream the rows. Only paths and coordinate bounds are gathered in the parent. + Output row order and attribute columns follow the source; reopened indices follow the destination format. + """ + + from geoutils.pointcloud.pointcloud import PointCloud + + # Validate configuration before inspecting point records or creating output files + chunks = _point_partition_size(mp_config) + if chunks <= 0: + raise ValueError("Argument ``chunks`` must be a strictly positive integer.") + if source._is_dask: + raise ValueError("Argument ``mp_config`` cannot be combined with a Dask point cloud.") + output_filename = pathlib.Path(mp_config.outfile) + suffix = output_filename.suffix.lower() + formats = {".las": "LAS", ".laz": "LAZ", ".gpkg": "GPKG"} + driver = mp_config.driver.upper() if mp_config.driver is not None else formats.get(suffix, "GPKG") + if driver not in formats.values(): + raise ValueError("Argument ``driver`` must be 'GPKG', 'LAS' or 'LAZ' for point cloud reprojection.") + if (suffix and (suffix not in formats or formats[suffix] != driver)) or (not suffix and driver != "GPKG"): + raise ValueError("Arguments ``driver`` and ``outfile`` must select the same supported point cloud format.") + target_crs = CRS.from_user_input(crs) + + # Plan slices from source metadata without loading an unopened point cloud + source_filename = pathlib.Path(source.name) if not source._is_pd and source.name is not None else None + if not source.is_loaded: + if source_filename is None or ( + not _is_laspy_supported(source_filename) and pyogrio.read_info(source_filename)["driver"] != "GPKG" + ): + raise ValueError("Unloaded point cloud reprojection supports LAS, LAZ and GPKG sources.") + dataframe = None + else: + dataframe = _as_geodataframe(source.ds, crs=source.crs) + if driver == "GPKG": + _check_gpkg_attributes(dataframe) + columns = list(source._nongeo_columns) + point_count = source.point_count + + # Complete all source reads in temporary files before replacing any existing destination + with tempfile.TemporaryDirectory(prefix=".geoutils-reproject-", dir=output_filename.parent) as directory: + temporary_directory = pathlib.Path(directory) + futures = [] + for index, start in enumerate(range(0, max(point_count, 1), chunks)): + count = min(chunks, point_count - start) + partition_source = source_filename if dataframe is None else dataframe.iloc[start : start + count] + futures.append( + mp_config.cluster.submit( + _reproject_pointcloud_partition, + partition_source, + columns, + start, + count, + target_crs, + temporary_directory / f"projected_{index}.pkl", + driver != "GPKG", + ) + ) + projected_parts = mp_config.cluster.gather(futures) + temporary_output = temporary_directory / f"output.{driver.lower()}" + + # Stream the exact GPKG geometry and attributes in the original point order + if driver == "GPKG": + # Reserve internal fields independently of user columns such as 'fid' and 'geom' + column_names = {column.lower() for column in columns} + layer_options = {"FID": "fid", "GEOMETRY_NAME": "geom"} + for option, field_name in layer_options.items(): + while field_name.lower() in column_names: + field_name = "_" + field_name + layer_options[option] = field_name + for index, (filename, _) in enumerate(projected_parts): + projected = pd.read_pickle(filename) + geometry_type = None + if len(projected) == 0: + geometry_type = "Point Z" if source._has_z else "Point" + pyogrio.write_dataframe( + projected, + temporary_output, + layer="points", + driver="GPKG", + append=index > 0, + geometry_type=geometry_type, + layer_options=layer_options, + ) + else: + # Encode every LAS worker file using the same global bounds and native dimension schema + projected = pd.read_pickle(projected_parts[0][0]) + bounds = [bounds for _, bounds in projected_parts if bounds is not None] + header = _reproject_las_header(projected, bounds, target_crs, source_filename) + del projected + futures = [] + for index, (filename, _) in enumerate(projected_parts): + futures.append( + mp_config.cluster.submit( + _write_reprojected_las_partition, + filename, + temporary_directory / f"encoded_{index}.las", + header, + ) + ) + written_paths = mp_config.cluster.gather(futures) + _stitch_laspy_files(temporary_output, written_paths, header=header, chunk_size=chunks) + + # Publish only the completed file so a source and destination may safely refer to the same path + os.replace(temporary_output, output_filename) + + return PointCloud(output_filename, data_column=source.data_column) diff --git a/geoutils/raster/__init__.py b/geoutils/raster/__init__.py index 52d8af52c..4020799fd 100644 --- a/geoutils/raster/__init__.py +++ b/geoutils/raster/__init__.py @@ -28,6 +28,7 @@ from geoutils.raster.transformation import * # noqa from geoutils.raster.raster import Raster, RasterType, handled_array_funcs # noqa isort:skip +from geoutils.raster.xr_accessor import RasterAccessor, open_raster # noqa isort:skip -__all__ = ["RasterType", "Raster"] +__all__ = ["Raster", "RasterAccessor", "RasterType", "open_raster"] diff --git a/geoutils/raster/array.py b/geoutils/raster/array.py index c6f736f6e..068b2e4ff 100644 --- a/geoutils/raster/array.py +++ b/geoutils/raster/array.py @@ -26,13 +26,50 @@ import numpy as np import xarray as xr -from geoutils._dispatch import has_geo_attr +from geoutils._dispatch import get_geo_attr, has_geo_attr from geoutils._typing import MArrayNum, NDArrayBool, NDArrayNum if TYPE_CHECKING: from geoutils.raster.base import RasterLike, RasterType +################################# +# 1/ INTERNAL ARRAY NORMALIZATION +################################# + + +def _selected_raster_data(raster: Any, band: int = 1, *, fill_value: float | bool = np.nan) -> Any: + """Select one raster band without loading lazy data.""" + + # Unwrap Xarray while preserving the laziness of its underlying array + data = get_geo_attr(raster, "data") + if isinstance(data, xr.DataArray): + data = data.data + + # Validate the public one based index before selecting one data plane + if band < 1: + raise ValueError("band numbers start at 1.") + if data.ndim == 2: + if band != 1: + raise ValueError("A single band raster only accepts band=1.") + selected = data + elif data.ndim == 3: + if band > data.shape[0]: + raise ValueError("band exceeds the number of raster bands.") + selected = data[band - 1] + else: + raise ValueError("Raster values must have two spatial dimensions and an optional band dimension.") + + # Preserve unmasked integer data and promote only when missing values require NaN + if np.ma.isMaskedArray(selected): + if not np.ma.is_masked(selected): + return np.ma.getdata(selected) + if np.isnan(fill_value): + selected = selected.astype(np.result_type(selected.dtype, np.float32)) + selected = selected.filled(fill_value) + return selected + + def _masked_raster_data(source_raster: RasterLike) -> MArrayNum: """Return raster values and their mask as an in-memory masked array.""" @@ -76,6 +113,11 @@ def _as_bands(array: MArrayNum) -> tuple[MArrayNum, bool]: raise ValueError("Raster processing expects a two-dimensional or multiband array.") +############################## +# 2/ MASK AND EXTENT FUNCTIONS +############################## + + def get_mask_from_array(array: NDArrayNum | NDArrayBool | MArrayNum) -> NDArrayBool: """ Return the mask of invalid values, whether array is a ndarray with NaNs or a np.ma.masked_array. @@ -146,6 +188,11 @@ def get_valid_extent(array: NDArrayNum | NDArrayBool | MArrayNum) -> tuple[int, return rows_nonzero[0], rows_nonzero[-1], cols_nonzero[0], cols_nonzero[-1] +########################### +# 3/ COORDINATE TRANSFORMS +########################### + + def get_xy_rotated(raster: RasterType, along_track_angle: float) -> tuple[NDArrayNum, NDArrayNum]: """ Rotate x, y axes of image to get along- and cross-track distances. diff --git a/geoutils/raster/base.py b/geoutils/raster/base.py index d5d2414cf..86e4de505 100644 --- a/geoutils/raster/base.py +++ b/geoutils/raster/base.py @@ -25,6 +25,7 @@ import struct import warnings from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import ( TYPE_CHECKING, Any, @@ -61,7 +62,11 @@ from geoutils.gapfill import _fill_nodata from geoutils.interface._nodata import NodataPropagation from geoutils.interface.distance import _proximity_from_vector_or_raster -from geoutils.interface.interpolation import _interp_points, _reduce_points +from geoutils.interface.interpolation import ( + InterpolationMethod, + _interp_points, + _reduce_points, +) from geoutils.interface.raster_point import ( _raster_to_pointcloud, _regular_pointcloud_to_raster, @@ -86,8 +91,9 @@ ) from geoutils.raster.testing import _array_equal_or_close from geoutils.raster.transformation import _crop, _reproject, _translate -from geoutils.stats.sampling import _subsample -from geoutils.stats.stats import _statistics +from geoutils.sampling.subsampling import _subsample +from geoutils.stats.stats import stats as _stats +from geoutils.stats.stats import variogram as _variogram # Input/output is a RasterType (= Raster or RasterAccessor subclass) RasterType = TypeVar("RasterType", bound="RasterBase") @@ -96,7 +102,10 @@ _UNSET = object() if TYPE_CHECKING: + from geoutils.interface.gridding import GriddingMethod from geoutils.pointcloud.pointcloud import PointCloud, PointCloudLike + from geoutils.stats.variography import Variogram + from geoutils.vector.base import VectorLike from geoutils.vector.vector import Vector, VectorType @@ -739,7 +748,7 @@ def info(self, stats: bool = False, verbose: bool = True) -> None | str: self.load() if self.count == 1: - statistics = self.get_stats() + statistics = self.stats() # Determine the maximum length of the stat names for alignment max_len = max(len(name) for name in statistics.keys()) @@ -751,7 +760,7 @@ def info(self, stats: bool = False, verbose: bool = True) -> None | str: for b in range(self.count): # try to keep with rasterio convention. as_str.append(f"Band {b + 1}:") - statistics = self.get_stats(band=b + 1) + statistics = self.stats(values=b + 1) if isinstance(statistics, dict): max_len = max(len(name) for name in statistics.keys()) for name, value in statistics.items(): @@ -763,147 +772,137 @@ def info(self, stats: bool = False, verbose: bool = True) -> None | str: else: return "\n".join(as_str) - @overload - def get_stats( - self, - stats_name: str | Callable[[NDArrayNum], np.floating[Any]], - inlier_mask: RasterType | NDArrayBool | None = None, - band: int = None, - counts: tuple[int, int] | None = None, - ) -> np.floating[Any] | dict[str, np.floating[Any]] | dict[str, dict[str, np.floating[Any]]]: ... - - @overload - def get_stats( + def stats( self, - stats_name: list[str | Callable[[NDArrayNum], np.floating[Any]]] | None = None, - inlier_mask: RasterType | NDArrayBool | None = None, - band: int = None, - counts: tuple[int, int] | None = None, - ) -> dict[str, np.floating[Any]] | dict[str, np.floating[Any]] | dict[str, dict[str, np.floating[Any]]]: ... + statistics: str | Callable[[Any], Any] | Iterable[str | Callable[[Any], Any]] | None = None, + *, + by: Mapping[str, Any] | None = None, + values: int | Iterable[int] | Mapping[str, Any] | None = None, + bins: Mapping[str, Any] | None = None, + categories: Mapping[str, Iterable[Any]] | None = None, + at: Literal["self"] | RasterLike | PointCloudLike | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, + mask_mode: Literal["inside", "outside"] = "inside", + subsample: int | float = 1, + subsample_per_group: bool = False, + random_state: int | np.random.Generator | None = None, + strategy: Literal["auto", "dense", "sparse", "groupwise"] = "auto", + backend: Literal["geoutils", "flox"] = "geoutils", + subsampling_strategy: Literal["sequential", "topk"] = "topk", + interpolation: InterpolationMethod = "linear", + align: Literal["raise", "reproject"] = "raise", + observed: bool = True, + return_masks: bool = False, + mp_config: MultiprocConfig | None = None, + ) -> Any: + """Calculate summary statistics or statistics grouped by categories, bins, or vector zones. + + Omit ``by`` to summarize raster values. With ``by``, provide ``categories`` for discrete groups, ``bins`` for + continuous groups, or a vector feature column for zonal statistics:: + + # Global statistics + raster.stats() + raster.stats("mean") + raster.stats(["mean", "std", "nmad"]) + + # Categorical stats + raster.stats(["mean", "std"], by={"landcover": landcover}, categories={"landcover": classes}) + # Binned stats + raster.stats(["mean", "std"], by={"elevation": elevation}, bins={"elevation": elevation_bins}) + # Zonal stats + raster.stats(["mean", "std"], by={"glacier": (glacier_outlines, "id")}) + + # Multiple grouping + raster.stats(["mean", "std"], by={"elevation": elevation, "glacier": (glacier_outlines, "id")}, + bins={"elevation": elevation_bins}) + + + :param statistics: Statistics to calculate (e.g. "mean", ["mean", "nmad"], or np.nanmedian). None returns + "min", "max", "mean", "median", "std", "nmad", "validcount", "totalcount" and "percentagevalidpoints". + "all" also includes "sum", "sumofsquares", "90thpercentile", "iqr", "le90" and "rmse", plus inlier counts + for masked global statistics. Grouped defaults replace "validcount" with "count"; every grouped result + includes "count". + :param by: Named variables to group by (e.g. {"elevation": elevation}); use {"glacier": (outlines, "id")} + for vector zones. Arrays must match the selected locations. Omit for global statistics. + :param values: Bands to summarize, counting from one (e.g. [1, 3]); defaults to all bands. + Use a mapping for named inputs (e.g. {"elevation": (dem, 1)}). + :param bins: Continuous bins keyed by grouping name (e.g. {"elevation": 10}). Each definition is a count of + equal-width bins, increasing edges (e.g. [0, 2, 5]), or a Pandas IntervalIndex to choose open/closed sides. + :param categories: Ordered categories keyed by grouping name (e.g. {"landcover": [100, 110, 120]}). + Values outside these categories are excluded. + :param at: Grid or ordered point locations on which to calculate statistics (e.g. at=reference or at="self"). + Defaults to the first point input, if present, otherwise the source locations. + :param mask: Locations to include (True in a boolean mask, e.g. mask=dem > 1000, or features in a vector mask). + Global counts describe values before this mask; "all" adds counts for values kept by the mask. + :param mask_mode: Keep locations "inside" or "outside" vector features; ignored for boolean masks. + :param subsample: Fraction (e.g. 0.1 for 10%) or maximum count (e.g. 10000) of eligible locations to use. + A value of 1 keeps all locations. Counts describe the sampled locations. + :param subsample_per_group: Apply subsample within each combined group (True, stratified sampling) or once + across all groups (False). Without by, both use one global sample. + :param random_state: Seed to reproduce subsampling (e.g. 42), or an existing random generator. + :param strategy: Combine chunk statistics for all groups ("dense"), only groups present in each chunk + ("sparse"), or gather each complete group ("groupwise"). "auto" chooses from the statistics and group count; + exact quantiles and custom functions require "auto" or "groupwise" for chunked data. + :param backend: Use the GeoUtils reducer ("geoutils") or optional Flox reducer ("flox") for grouped statistics; + see stats() for the Flox restrictions. + :param subsampling_strategy: "topk" keeps the same sampled locations across chunk layouts for a fixed seed; + "sequential" draws random locations using traversal order and can depend on the chunks. + :param interpolation: Raster values at point locations use interp_points() with SciPy methods "nearest", + "linear", "slinear", "cubic", "quintic", "pchip" or "splinef2d". + Raster groupers listed in categories use "nearest". + :param align: "raise" rejects different grids or coordinate systems; "reproject" aligns them to the output + locations. Point inputs must still share the same ordered coordinates. + :param observed: Omit declared group combinations with no eligible locations (True), or include them (False). + :param return_masks: Also return masks keyed by group labels (e.g. table, masks = raster.stats(...)). + Masks cover complete groups before subsampling. Requires by. + :param mp_config: Worker and tile settings for multiprocessing, e.g. MultiprocConfig(chunks=512). + Cannot be combined with Dask inputs. + :returns: A statistic, summary dictionary, grouped dataframe, or grouped dataframe and mask mapping. + """ + + return _stats( + self, + statistics, + by=by, + values=values, + bins=bins, + categories=categories, + at=at, + mask=mask, + mask_mode=mask_mode, + subsample=subsample, + subsample_per_group=subsample_per_group, + random_state=random_state, + strategy=strategy, + backend=backend, + subsampling_strategy=subsampling_strategy, + interpolation=interpolation, + align=align, + observed=observed, + return_masks=return_masks, + mp_config=mp_config, + ) @profiler.profile("geoutils.raster.base.get_stats", memprof=True) def get_stats( self, stats_name: ( - str | Callable[[NDArrayNum], np.floating[Any]] | list[str | Callable[[NDArrayNum], np.floating[Any]]] | None + str + | Callable[[NDArrayNum], np.floating[Any]] + | Iterable[str | Callable[[NDArrayNum], np.floating[Any]]] + | None ) = None, inlier_mask: RasterType | NDArrayBool | None = None, - band: int = None, - counts: tuple[int, int] | None = None, - ) -> np.floating[Any] | dict[str, np.floating[Any]] | dict[str, dict[str, np.floating[Any]]]: - """ - Retrieve specified statistics or all available statistics for the raster data. Allows passing custom callables - to calculate custom stats. - - Common statistics are : - - - Mean: arithmetic mean of the data, ignoring masked values. - - Median: middle value when the valid data points are sorted in increasing order, ignoring masked values. - - Max: maximum value among the data, ignoring masked values. - - Min: minimum value among the data, ignoring masked values. - - Sum: sum of all data, ignoring masked values. - - Sum of squares: sum of the squares of all data, ignoring masked values. - - 90th percentile: point below which 90% of the data falls, ignoring masked values. - - IQR (Interquartile Range): difference between the 75th and 25th percentile of a dataset, \ - ignoring masked values. - - LE90 (Linear Error with 90% confidence): difference between the 95th and 5th percentiles of a dataset, \ - representing the range within which 90% of the data points lie. Ignore masked values. - - NMAD (Normalized Median Absolute Deviation): robust measure of variability in the data, \ - less sensitive to outliers compared to standard deviation. Ignore masked values. - - RMSE (Root Mean Square Error): commonly used to express the magnitude of errors or variability and can give \ - insight into the spread of the data. Only relevant when the raster represents a difference of two objects. \ - Ignore masked values. - - Std (Standard deviation): measures the spread or dispersion of the data around the mean, \ - ignoring masked values. - - Valid count: number of finite data points in the array. It counts the non-masked elements. - - Total count: total size (width x height) of the raster. - - Percentage valid points: ratio between Valid count and Total count. - - For all statistics up to and including "Std", NumPy Masked functions are used (directly or in the calculation) - in case of a masked array, NumPy module otherwise. - - "Valid count" represents all non zero and not masked pixels in the input data (final_count_nonzero), - calculated before the mask application in case of an inlier_mask. NumPy Masked functions are used is this case - or if the Raster was already a masked array. "Percentage valid points" is calculated accordingly. - - If an inlier mask is passed: - - - Total inlier count: number of data points in the inlier mask. - - Valid inlier count: number of unmasked data points in the array after applying the inlier mask. - - Percentage inlier points: ratio between Valid inlier count and Valid count. Useful for classification \ - statistics. - - Percentage valid inlier points: ratio between Valid inlier count and Total inlier count. - - They are all computed based on the previously stated final_count_nonzero. - - Callable functions are supported as well. - - By default and without any specification, this function computes the following main statistics: minimum, - maximum, mean, standard deviation, NMAD, total count, and percentage of valid points. - To compute all available statistics, set `stats_name` to `all`. - - :param stats_name: Name or list of names of the statistics to retrieve. If None, main statistics are returned. - Accepted names include: - `mean`, `median`, `max`, `min`, `sum`, `sum of squares`, `90th percentile`, `iqr`, `LE90`, `nmad`, `rmse`, - `std`, `valid count`, `total count`, `percentage valid points` and if an inlier mask is passed : - `valid inlier count`, `total inlier count`, `percentage inlier point`, `percentage valid inlier points`. - Custom callables can also be provided. - To compute all available statistics, set `stats_name` to `all`. - :param inlier_mask: Mask or boolean array of areas to include (inliers=True). - :param band: The index of the band for which to compute statistics. Default is 1. - :param counts: (number of finite data points in the array, number of valid points (=True, to keep) - in inlier_mask), initialize in case of an inlier_mask. DO NOT USE. - :returns: The requested statistic or a dictionary of statistics if multiple or all are requested. - """ - - # Case mono-band - if self.count == 1 and band is None: - band = 1 - - if band is not None: - # Get data band - data = self.data[band - 1, :, :] if self.count > 1 else self.data - - # Derive inlier mask - if inlier_mask is not None: - valid_points = np.count_nonzero(np.logical_and(np.isfinite(data), ~data.mask)) - if isinstance(inlier_mask, RasterBase) and inlier_mask.is_mask: - mask = inlier_mask.data - else: - mask = inlier_mask - inlier_points = np.count_nonzero(mask) - - rast = self.copy() - - # Mask pixels from the inlier_mask - if not np.ma.isMaskedArray(rast.data): - rast[~mask] = np.nan # type: ignore - else: - rast.set_mask(~mask) # type: ignore - return rast.get_stats(stats_name=stats_name, band=band, counts=(valid_points, inlier_points)) - - # Given list or all attributes to compute if None - if isinstance(stats_name, list) or stats_name is None or stats_name == "all": - return _statistics(data, stats_name, counts) # type: ignore - else: - # Single attribute to compute - if isinstance(stats_name, str): - return _statistics(data, [stats_name], counts)[stats_name] # type: ignore - elif callable(stats_name): - return stats_name(data) # type: ignore - else: - warnings.warn( - "Statistic name " + str(stats_name) + " is a not recognized string", category=UserWarning - ) - else: - # Case multi-band - stats = {} - for band in range(1, self.count + 1): - stats["band " + str(band)] = self.get_stats( - stats_name=stats_name, inlier_mask=inlier_mask, band=band, counts=counts - ) + band: int | None = None, + ) -> Any: + """Call stats() with the legacy argument names; deprecated in favor of stats().""" - return stats # type: ignore + warnings.warn( + "get_stats() is deprecated; use stats() with values and mask instead.", + DeprecationWarning, + stacklevel=2, + ) + return _stats(self, statistics=stats_name, mask=inlier_mask, values=band) def _raster_equal_allclose( self, @@ -948,11 +947,10 @@ def _raster_equal_allclose( names = names + ["fill_value", "dtype", "transform", "crs", "nodata"] equalities = equalities_data + [ - self.data.fill_value == other.data.fill_value, + np.array_equal(self.data.fill_value, other.data.fill_value, equal_nan=True), self.data.dtype == other.data.dtype, self.transform == other.transform, self.crs == other.crs, - self.nodata == other.nodata, ] # For Raster or DataArray else: @@ -960,7 +958,6 @@ def _raster_equal_allclose( dtype = other.rst.dtype if isinstance(other, xr.DataArray) else other.dtype transform = other.rst.transform if isinstance(other, xr.DataArray) else other.transform crs = other.rst.crs if isinstance(other, xr.DataArray) else other.crs - nodata = other.rst.nodata if isinstance(other, xr.DataArray) else other.nodata # Three cases: masked/NaN, NaN/masked or NaN/NaN if np.ma.isMaskedArray(self.data): @@ -987,9 +984,15 @@ def _raster_equal_allclose( self.dtype == dtype, self.transform == transform, self.crs == crs, - self.nodata == nodata, ] + # Compare nodata after data access, which can update metadata when a file is loaded as a boolean mask + other_nodata = other.rst.nodata if isinstance(other, xr.DataArray) else other.nodata + nodata_equal = self.nodata == other_nodata + if self.nodata is not None and other_nodata is not None: + nodata_equal = nodata_equal or bool(np.isnan(self.nodata) and np.isnan(other_nodata)) + equalities.append(nodata_equal) + complete_equality = all(equalities) if not complete_equality and warn_failure_reason: @@ -1420,6 +1423,8 @@ def reproject( result_raster = open_raster(mp_config.outfile, is_mask=self.is_mask) else: result_raster = self.__class__(mp_config.outfile) + # Restore logical interpretation without changing subclass constructor arguments + result_raster._is_mask = self.is_mask return self._cast_raster_output(result_raster) # type: ignore # Not in-place @@ -1620,7 +1625,7 @@ def outside_image(self, xi: ArrayLike, yj: ArrayLike, index: bool = True) -> boo def interp_points( self, points: tuple[Number, Number] | tuple[NDArrayNum, NDArrayNum] | PointCloudLike, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] = None, + method: InterpolationMethod = None, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, band: int = 1, input_latlon: bool = False, @@ -1636,7 +1641,7 @@ def interp_points( def interp_points( self, points: tuple[Number, Number] | tuple[NDArrayNum, NDArrayNum] | PointCloudLike, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] = None, + method: InterpolationMethod = None, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, band: int = 1, input_latlon: bool = False, @@ -1652,7 +1657,7 @@ def interp_points( def interp_points( self, points: tuple[Number, Number] | tuple[NDArrayNum, NDArrayNum] | PointCloudLike, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] = None, + method: InterpolationMethod = None, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, band: int = 1, input_latlon: bool = False, @@ -1668,7 +1673,7 @@ def interp_points( def interp_points( self, points: tuple[NDArrayNum, NDArrayNum] | tuple[Number, Number] | PointCloudLike, - method: Literal["nearest", "linear", "cubic", "quintic", "slinear", "pchip", "splinef2d"] = None, + method: InterpolationMethod = None, dist_nodata_spread: Literal["half_order_up", "half_order_down"] | int | None = None, band: int = 1, input_latlon: bool = False, @@ -1813,7 +1818,7 @@ def filter( :param sigma: Optional standard deviation for Gaussian filtering. Only used when `method="gaussian"`. :param engine: Optional engine to use for filtering, either "scipy" (default) or "numba". - Only used when `method="median"`. + Only used with built-in filters. :param outlier_threshold: The minimum difference abs(array - mean) for a pixel to be considered an outlier. Only used when `method="distance"`. :param kwargs : Additional keyword arguments passed to the underlying filter implementation. @@ -2150,6 +2155,7 @@ def subsample( random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", mp_config: MultiprocConfig | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum: ... @overload @@ -2162,6 +2168,7 @@ def subsample( random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", mp_config: MultiprocConfig | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, ) -> tuple[NDArrayNum, ...]: ... @overload @@ -2173,6 +2180,8 @@ def subsample( random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", mp_config: MultiprocConfig | None = None, + *, + mask: RasterLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum | tuple[NDArrayNum, ...]: ... @profiler.profile("geoutils.raster.base.subsample", memprof=True) @@ -2184,18 +2193,27 @@ def subsample( random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", mp_config: MultiprocConfig | None = None, + *, + mask: RasterLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum | tuple[NDArrayNum, ...]: """ - Randomly sample the raster. Only valid values are considered. + Randomly sample valid raster values allowed by mask, without replacement. - :param subsample: Subsample size. If <= 1, a fraction of the total pixels to extract. - If > 1, the number of pixels. + :param subsample: Subsample size. If <= 1, a fraction of eligible finite pixels to extract. + If > 1, the maximum number of pixels. The mask is applied before calculating this size. :param band: Band to subsample. Use return_indices=True and indexing to subsample the same points over several bands. :param return_indices: Whether to return the extracted indices only. :param random_state: Random state or seed number. + :param strategy: "sequential" draws using the traversal order and can depend on chunk layout; "topk" keeps + the same seeded sample across chunk layouts. + :param mp_config: Worker and tile settings for multiprocessing. Cannot be combined with a Dask source. + :param mask: Eligible cells: True in a boolean array or aligned mask raster, or inside vector geometries. + Arrays must match the raster shape; mask rasters must share its grid and CRS. Missing mask entries are + excluded (e.g. mask=raster.data > 0). - :return: Array of sampled valid values, or array of sampled indices. + :returns: One-dimensional sampled values with the source dtype, or a tuple of row and column index arrays + referring to the original grid, including when mask restricts the sample. """ return _subsample( @@ -2206,4 +2224,255 @@ def subsample( random_state=random_state, strategy=strategy, mp_config=mp_config, + mask=mask, + ) + + def cosample( + self, + other: RasterLike | PointCloudLike | ArrayLike, + *, + band: int = 1, + other_band: int = 1, + auxiliary: Mapping[str, Any] | None = None, + auxiliary_at: Literal["self", "other"] | Mapping[str, Literal["self", "other"]] | None = None, + at: Literal["self", "other"] | RasterLike | PointCloudLike | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, + mask_mode: Literal["inside", "outside"] = "inside", + subsample: int | float = 1, + random_state: int | np.random.Generator | None = None, + strategy: Literal["sequential", "topk"] = "topk", + raster_point_mode: Literal["grid_points", "resample_raster"] | None = None, + grid_method: GriddingMethod = "linear", + resample_method: InterpolationMethod | Literal["reduce"] = "linear", + grid_kwargs: Mapping[str, Any] | None = None, + resample_kwargs: Mapping[str, Any] | None = None, + align: Literal["raise", "reproject"] = "raise", + mp_config: MultiprocConfig | None = None, + ) -> RasterLike | PointCloudLike: + """ + Sample this raster and another dataset at common finite locations. + + By default, raster values are resampled at point locations when either primary input is a point cloud. + Use ``raster_point_mode="grid_points"`` to grid point values onto the raster instead. An explicit ``at`` + chooses the exact output locations and must agree with any explicit mode. Raw auxiliary arrays must + identify the primary input whose grid or point ordering they follow. + + Spatial inputs, explicit output support and raster or point masks must use one family: Raster/PointCloud + objects, or DataArray/GeoDataFrame objects. The latter may mix eager and Dask storage. Plain arrays and + vector outlines are accepted with either family. + + :param other: Dataset to sample alongside this raster. A plain array follows this raster's grid. + :param band: Band selected from this raster, counting from one. + :param other_band: Band selected from other if it is a raster, counting from one. + :param auxiliary: Additional values by output name (e.g. {"slope": slope_raster}). Select a raster band with + {"slope": (slope_raster, 2)} or a point column with {"intensity": (points, "intensity")}. Spatial inputs + without a selector use the first raster band or active point values. + :param auxiliary_at: Input locations followed by plain auxiliary arrays: "self", "other", or a choice per name + (e.g. {"slope": "other"}). Spatial auxiliaries use their own coordinates. + :param at: Output locations: "self", "other", or a reference raster/point cloud. Defaults to the other point + cloud's locations when present, or this raster's grid otherwise. + :param mask: Locations eligible for sampling, defined by a boolean array, spatial mask, or vector outlines. + :param mask_mode: Whether a vector mask keeps locations "inside" or "outside" its geometries. + :param subsample: Fraction of common finite locations (e.g. 0.1), or maximum count (e.g. 1000); 1 keeps all. + :param random_state: Seed or random generator for reproducible sampling (e.g. 42). + :param strategy: Raster sampling with "topk" or "sequential"; "topk" keeps the same seeded sample across chunk + sizes. Point output always uses "sequential". + :param raster_point_mode: Conversion direction: "grid_points" places points on a raster, "resample_raster" + reads rasters at points. Defaults to at's locations, or point locations when available. Must agree with at. + :param grid_method: Point gridding by SciPy interpolation ("nearest", "linear", "cubic"), or circular "idw", + "mean", "minimum", "maximum", "range", "count", "stdev", "average_distance", "average_distance_pts". + The aliases "average", "min" and "max" select "mean", "minimum" and "maximum". + :param resample_method: Raster interpolation using the SciPy methods "nearest", "linear", "cubic", "quintic", + "slinear", "pchip" or "splinef2d". Window reduction ("reduce") is not implemented. + :param grid_kwargs: Options for PointCloud.grid(), e.g. {"dist_nodata_pixel": 2, "min_points": 3} sets a + two-pixel radius and minimum of three finite points for circular methods. Other options include + "distance_power" for IDW and "engine" ("scipy" or "numba"). + Set locations and method with at and grid_method. + :param resample_kwargs: Options for Raster.interp_points(), e.g. {"nodata_propagation": "ignore"}. The nodata + policies are "gdal", "ignore" and "propagate"; "dist_nodata_spread" controls extra spreading in pixels. + Set locations, band and method with the corresponding cosample() arguments. + :param align: Handling of mismatched grids or coordinate systems: "raise" an error, or "reproject" to match at. + Point inputs must still share the same ordered coordinates when sampled at points. + :param mp_config: Worker and tile settings for multiprocessing. Raster output uses its outfile; cannot be + combined with Dask inputs. + :returns: Raster or point cloud on the selected support; Xarray DataArray or eager/lazy GeoDataFrame for + accessor calls. Bands or columns contain "self", "other", then auxiliaries in mapping order. Raster + band names are stored in ``tags["long_name"]`` (Xarray ``attrs["long_name"]``). Point outputs use + "self" as their active data column and preserve the support index and order. + """ + + from geoutils.sampling.cosampling import _cosample + + return _cosample( + self, + other, + band=band, + other_band=other_band, + auxiliary=auxiliary, + auxiliary_at=auxiliary_at, + at=at, + mask=mask, + mask_mode=mask_mode, + subsample=subsample, + random_state=random_state, + strategy=strategy, + raster_point_mode=raster_point_mode, + grid_method=grid_method, + resample_method=resample_method, + grid_kwargs=grid_kwargs, + resample_kwargs=resample_kwargs, + align=align, + mp_config=mp_config, + ) + + def pairsample( + self, + *, + band: int = 1, + n_pairs: int = 1_000_000, + sampling: Literal["loglag", "random_xy"] = "loglag", + min_distance: float | None = None, + max_distance: float | None = None, + random_state: int | np.random.Generator | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, + strategy: Literal["independent", "anchors", "chunk_anchors", "anchor_batched"] = "chunk_anchors", + deduplicate: Literal["none", "per_anchor", "global"] = "per_anchor", + batch_pairs: int = 2_000_000, + max_rounds: int = 50, + max_oversample: float = 8.0, + chunks_per_round: int = 8, + anchors_per_round: int = 20_000, + distances_per_anchor: int = 8, + angles_per_distance: int = 8, + hybrid_local_fraction: float = 0.0, + max_local_distance: float | None = None, + index_dtype: DTypeLike = np.int32, + distance_dtype: DTypeLike = np.float64, + ) -> xr.Dataset: + """Sample finite raster cell pairs for statistics by distance. + + Logarithmic lag sampling draws isotropic distances across short and long ranges. Anchor strategies reuse + raster cells and can confine part of the sample to source chunks, which limits reads from Dask-backed rasters. + + Strategy, duplicate, oversampling, anchor, and local distance controls apply to ``"loglag"``. + Both sampling schemes use ``batch_pairs`` and ``max_rounds``. + + :param band: Band to sample, counting from one. + :param n_pairs: Requested number of pairs with two finite values; fewer may be returned if sampling stops early. + :param sampling: ``"loglag"`` balances short and long distances on a log scale; ``"random_xy"`` draws + endpoints uniformly. + :param min_distance: Smallest distance in CRS units (e.g. meters). Defaults to the smaller pixel spacing. + :param max_distance: Largest distance in CRS units. Defaults to the diagonal between outermost cell centers. + :param random_state: Seed for reproducible sampling (e.g. 42). + :param mask: Eligible cells: True in a mask array or aligned mask raster, or inside vector geometries. + :param strategy: GeoUtils log-lag strategy: ``"independent"`` draws each pair separately, ``"anchors"`` + reuses first endpoints, ``"chunk_anchors"`` also limits source chunks, and ``"anchor_batched"`` draws + several distances and directions from each first endpoint. + :param deduplicate: ``"none"`` keeps repeats, ``"per_anchor"`` removes repeated targets within each anchor + batch, and ``"global"`` removes repeated pairs across all batches. ``"random_xy"`` always removes repeats. + :param batch_pairs: Maximum candidate pairs per batch; smaller batches use less temporary memory. + :param max_rounds: Maximum attempts to fill the sample after rejecting missing values or out-of-range pairs. + :param max_oversample: Maximum candidate count as a multiple of the target pair count (e.g. 8). + :param chunks_per_round: Maximum source chunks used when drawing anchors from selected chunks. + :param anchors_per_round: Maximum first endpoints reused per round by ``"anchors"``, ``"chunk_anchors"``, + or local ``"independent"`` sampling. + :param distances_per_anchor: Distances drawn per first endpoint with ``"anchor_batched"``. + :param angles_per_distance: Directions drawn per distance with ``"anchor_batched"``. + :param hybrid_local_fraction: Fraction of candidate pairs kept within the first endpoint's chunk (e.g. 0.5). + Zero samples across the full raster; one keeps all pairs local. + :param max_local_distance: Largest proposed local distance in CRS units. Defaults to the largest chunk diagonal. + :param index_dtype: Integer NumPy dtype for returned cell indexes (e.g. ``"int64"`` for very large rasters). + :param distance_dtype: Floating NumPy dtype for returned distances (e.g. ``"float32"`` to reduce memory). + :returns: Xarray Dataset with pair and endpoint dimensions, containing cell indexes, values, coordinates, + and distances. + """ + + from geoutils.sampling.pairsampling import _sample_raster_pairs + + return _sample_raster_pairs( + self, + band=band, + n_pairs=n_pairs, + sampling=sampling, + min_distance=min_distance, + max_distance=max_distance, + random_state=random_state, + mask=mask, + strategy=strategy, + deduplicate=deduplicate, + batch_pairs=batch_pairs, + max_rounds=max_rounds, + max_oversample=max_oversample, + chunks_per_round=chunks_per_round, + anchors_per_round=anchors_per_round, + distances_per_anchor=distances_per_anchor, + angles_per_distance=angles_per_distance, + hybrid_local_fraction=hybrid_local_fraction, + max_local_distance=max_local_distance, + index_dtype=index_dtype, + distance_dtype=distance_dtype, + ) + + def variogram( + self, + *, + band: int = 1, + n_pairs: int = 1_000_000, + sampling: Literal["loglag", "random_xy"] = "loglag", + estimator: str | Callable[[NDArrayNum], float] = "dowd", + bins: Literal["log", "uniform"] | Iterable[float] = "log", + n_lags: int = 24, + min_lag: float | None = None, + max_lag: float | None = None, + n_runs: int = 1, + model: str | Callable[..., Any] | list[str | Callable[..., Any]] | None = None, + fit_kwargs: dict[str, Any] | None = None, + random_state: int | np.random.Generator | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, + **pair_sampling_kwargs: Any, + ) -> Variogram: + """Estimate a lightweight empirical variogram from raster pairs. + + :param band: Raster band to sample, counting from 1. + :param n_pairs: Number of finite pairs targeted in each run (e.g. 100_000). + :param sampling: How to select pairs: ``"loglag"`` balances short and long distances, while ``"random_xy"`` + selects endpoints independently. + :param estimator: Semivariance estimator from SciKit-GStat: ``"dowd"``, ``"matheron"``, ``"cressie"``, + ``"genton"``, ``"minmax"``, ``"entropy"`` or ``"percentile"``. A function can instead map absolute pair + differences to one value per distance bin. + :param bins: Distance bins: ``"log"`` for logarithmic spacing, ``"uniform"`` for equal widths, or explicit + edges (e.g. [1, 10, 100]). + :param n_lags: Number of distance bins when bins is ``"log"`` or ``"uniform"``. + :param min_lag: Minimum sampled distance in CRS units; defaults to the smaller pixel dimension. + :param max_lag: Maximum sampled distance in CRS units; defaults to the diagonal between outermost cell centers. + :param n_runs: Independent samples to average; repeat sampling to estimate each distance bin's standard error. + :param model: SciKit-GStat model to fit: ``"spherical"``, ``"exponential"``, ``"gaussian"``, ``"cubic"``, + ``"stable"`` or ``"matern"``, or the corresponding model function. Sum a list of models ordered from short + to long range (e.g. ["spherical", "exponential"]). ``None`` keeps only the empirical variogram. + :param fit_kwargs: Options for Variogram.fit(): ``use_nugget``, ``bounds``, ``p0`` or ``maxfev`` + (e.g. {"use_nugget": True}); optimization uses SciPy curve_fit(). + :param random_state: Seed or NumPy generator for reproducible sampling across runs (e.g. 42). + :param mask: Cells to keep: True values in a boolean mask or aligned mask raster, or cells inside + vector geometries. + :param pair_sampling_kwargs: Extra pairsample() options (e.g. ``strategy`` or ``max_rounds``). + :returns: Variogram with distance bins, pair counts and semivariance, plus sampling errors and a fitted model + when requested. + """ + + return _variogram( + self, + band=band, + n_pairs=n_pairs, + sampling=sampling, + n_runs=n_runs, + estimator=estimator, + bins=bins, + n_lags=n_lags, + min_lag=min_lag, + max_lag=max_lag, + model=model, + fit_kwargs=fit_kwargs, + random_state=random_state, + mask=mask, + **pair_sampling_kwargs, ) diff --git a/geoutils/raster/raster.py b/geoutils/raster/raster.py index 67bb76996..b074c4251 100644 --- a/geoutils/raster/raster.py +++ b/geoutils/raster/raster.py @@ -992,6 +992,9 @@ def to_rio_dataset(self) -> rio.io.DatasetReader: else: ds.write(self.data) + # Preserve custom tags and pixel interpretation when exporting through an in-memory file + ds.update_tags(**self.tags) + # Then open as a DatasetReader return mfh.open() diff --git a/geoutils/raster/transformation.py b/geoutils/raster/transformation.py index 3a9140a98..04a233c99 100644 --- a/geoutils/raster/transformation.py +++ b/geoutils/raster/transformation.py @@ -686,6 +686,10 @@ def _wrapper_multiproc_reproject_per_block( # Call reproject per block dst_block_arr = _reproject_per_block(*src_arrs, block_ids=block_ids, combined_meta=combined_meta, **kwargs) + # Store logical masks as integers so the writer can fill missing cells with nodata rather than True + if dst_block_arr.dtype == np.bool_: + dst_block_arr = dst_block_arr.astype("uint8") + return dst_block_arr, (dst_block_id["ys"], dst_block_id["ye"], dst_block_id["xs"], dst_block_id["xe"]) @@ -759,7 +763,7 @@ def _multiproc_reproject( "count": rst.count, "crs": dst_crs, "transform": dst_transform, - "dtype": dtype, + "dtype": "uint8" if np.dtype(dtype) == np.bool_ else dtype, "nodata": dst_nodata, } @@ -944,6 +948,10 @@ def _crop( if crop_img.ndim == 3 and crop_img.shape[0] == 1: crop_img = crop_img.squeeze(axis=0) + # Restore logical mask values from their on-disk integer representation, keeping missing cells masked + if source_raster.is_mask: + crop_img = crop_img.astype(bool) + return crop_img, tfm diff --git a/geoutils/raster/xr_accessor.py b/geoutils/raster/xr_accessor.py index c0c8de778..f2ee04ce4 100644 --- a/geoutils/raster/xr_accessor.py +++ b/geoutils/raster/xr_accessor.py @@ -50,7 +50,8 @@ def open_raster(filename: str, is_mask: bool = False, **kwargs: Any) -> xr.DataA ds = rioxr.open_rasterio(filename, masked=True, **kwargs) # Remove the band dimension if there is only one - ds = ds.squeeze() # Delete band coordinate (only one dimension) + if ds.sizes.get("band") == 1: + ds = ds.squeeze("band") # Delete band coordinate (only one dimension) # If input needs to be interpreted as a boolean mask if is_mask: @@ -278,8 +279,9 @@ def from_array( else: data = np.ma.getdata(masked) - # Squeeze data - data = data.squeeze() + # Remove only a singleton band axis so one row or column still defines a raster grid + if data.ndim == 3 and data.shape[0] == 1: + data = data[0] # For a 2-d array if data.ndim == 2: @@ -310,15 +312,18 @@ def from_array( def to_geoutils(self) -> RasterBase: """ - Convert to Raster object from GeoUtils. + Convert the DataArray to an in-memory GeoUtils Raster. - :return: + :returns: A Raster with identical values and georeferencing. Dask inputs retain their lazy source graph; + ordinary file-backed DataArrays load their values during conversion. """ from geoutils.raster import Raster # Runtime import to avoid circularity issues + # Materialize a separate Dask result so conversion never replaces the source's lazy array + ds = self._obj.compute() if self._chunks is not None else self._obj return Raster.from_array( - data=self._obj.data, + data=ds.data, crs=self.crs, transform=self.transform, nodata=self.nodata, diff --git a/geoutils/sampling/__init__.py b/geoutils/sampling/__init__.py new file mode 100644 index 000000000..e3ff3fadd --- /dev/null +++ b/geoutils/sampling/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Sampling operations shared by raster and point cloud objects.""" diff --git a/geoutils/sampling/cosampling.py b/geoutils/sampling/cosampling.py new file mode 100644 index 000000000..328bdc914 --- /dev/null +++ b/geoutils/sampling/cosampling.py @@ -0,0 +1,1114 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +""" +Sample several geospatial datasets at the same locations. + +Note: This module is inspired from logic originally developed in xDEM for coregistration and uncertainty quantification. +""" + +from __future__ import annotations + +import warnings +from collections.abc import Iterable, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, cast + +import numpy as np +import xarray as xr + +from geoutils._dispatch import ( + _get_pointcloud_interface, + _get_raster_interface, + _is_raster, + get_geo_attr, + has_geo_attr, + is_dask_array, + is_dask_dataframe, +) +from geoutils._misc import import_optional +from geoutils._typing import ArrayLike, NDArrayBool, NDArrayNum +from geoutils.interface.gridding import GriddingMethod +from geoutils.raster.array import _selected_raster_data +from geoutils.sampling.subsampling import _sample_valid_indices +from geoutils.sampling.support import ( + _aligned_pointcloud, + _aligned_raster, + _as_array, + _mask_at_support, + _mask_on_raster, + _normalize_sampling_input, + _point_values_at_support, + _sampling_specification, + _sampling_support, +) + +if TYPE_CHECKING: + from geoutils.interface.interpolation import InterpolationMethod + from geoutils.multiproc import MultiprocConfig + from geoutils.pointcloud.base import PointCloudBase + from geoutils.pointcloud.pointcloud import PointCloudLike + from geoutils.raster.base import RasterBase, RasterLike + from geoutils.raster.raster import Raster + from geoutils.vector.base import VectorLike + + +################################# +# 1/ INPUT PREPARATION +################################# + + +@dataclass(frozen=True) +class _CosampleInput: + """ + Small dataclass to consistently store pointers to normalized values, selector, input support and kind. + + These are created in the input preparation step done in _prepare_cosample_input(), right below! + """ + + value: RasterBase | PointCloudBase | ArrayLike + selector: int | str | None + input_support: RasterBase | PointCloudBase + kind: Literal["raster", "point"] + + +def _prepare_cosample_input( + value: Any, + selector: int | str | None, + name: str, + native: _CosampleInput | None = None, +) -> _CosampleInput: + """ + Prepare a single input: differentiate raster/pointclouds and arrays, validate their relative shape if they are an + array input, and check their input band exists if provided. + + This function is used below in _prepare_cosample_all_inputs(). + """ + + # Normalize Xarray inputs without spatial coordinates as arrays, and resolve spatial interfaces once + value = _normalize_sampling_input(value) + raster = _get_raster_interface(value) + pointcloud = _get_pointcloud_interface(value) if raster is None else None + kind: Literal["raster", "point"] + if raster is not None: + value = input_support = raster + kind = "raster" + elif pointcloud is not None: + value = input_support = pointcloud + kind = "point" + else: + # Require every array input to indicate the primary input support it follows + if native is None: + raise ValueError(f"Argument ``auxiliary_at`` must identify the native support of array auxiliary {name!r}.") + input_support, kind = native.input_support, native.kind + value = _as_array(value) + + # Check the native shape without loading Dask arrays or spatial values + if kind == "raster": + if value.ndim == 3 and value.shape[0] == 1: + value = value[0] + if value.ndim != 2 or tuple(value.shape) != tuple(cast("RasterBase", input_support).shape): + raise ValueError(f"Array {name!r} must match the shape of its native raster input.") + elif value.ndim != 1: + raise ValueError(f"Raw point value {name!r} must contain one value per native point.") + + # Validate raster bands from metadata before alignment, interpolation or worker dispatch + if kind == "raster": + if selector is None and name not in {"self", "other"}: + selector = 1 + count = raster.count if raster is not None else 1 + if not isinstance(selector, (int, np.integer)): + raise TypeError(f"Raster selector for {name!r} must be a band number.") + if not 1 <= selector <= count: + raise ValueError(f"Band for {name!r} must be between one and the raster band count.") + selector = int(selector) + elif pointcloud is not None: + # Primary point inputs use their active values; auxiliary tuples can select another column + selector = pointcloud.data_column if name in {"self", "other"} or selector is None else selector + if selector is not None and (not isinstance(selector, str) or selector not in pointcloud.columns): + raise ValueError(f"Point column {selector!r} selected for {name!r} does not exist.") + else: + selector = None + + return _CosampleInput(value, selector, input_support, kind) + + +def _prepare_cosample_inputs( + first: RasterLike | PointCloudLike, + second: RasterLike | PointCloudLike | ArrayLike, + band: int, + other_band: int, + auxiliary: Mapping[str, Any] | None, + auxiliary_at: Literal["self", "other"] | Mapping[str, Literal["self", "other"]] | None, +) -> dict[str, _CosampleInput]: + """ + Prepare all inputs (storing kind, values, selector and native support) to later choose output support. + + See _cosample() for arguments. + + Raster/point cloud inputs use their own coordinates as input support. An array ``second`` follows the first + input support, and array auxiliaries define their input support through ``auxiliary_at``. + """ + + # Copy auxiliary dictionaries to avoid modifying the user input, and check their content + auxiliary = {} if auxiliary is None else dict(auxiliary) + if any(not isinstance(name, str) or not name for name in auxiliary): + raise ValueError("Auxiliary names must be non-empty strings.") + if {"self", "other", "geometry"}.intersection(auxiliary): + raise ValueError("Auxiliary names cannot be 'self', 'other' or 'geometry'.") + if not isinstance(auxiliary_at, Mapping): + if auxiliary_at not in (None, "self", "other"): + raise ValueError("Values in argument ``auxiliary_at`` must be 'self' or 'other'.") + auxiliary_at = {} if auxiliary_at is None else dict.fromkeys(auxiliary, auxiliary_at) + else: + auxiliary_at = dict(auxiliary_at) + if not set(auxiliary_at).issubset(auxiliary): + raise ValueError("Argument ``auxiliary_at`` contains a name that is not present in ``auxiliary``.") + if any(location not in ("self", "other") for location in auxiliary_at.values()): + raise ValueError("Values in argument ``auxiliary_at`` must be 'self' or 'other'.") + + # An array second input has to follow the first input's locations, including when explicitly selected as support + inputs = {"self": _prepare_cosample_input(first, band, "self")} + inputs["other"] = _prepare_cosample_input(second, other_band, "other", inputs["self"]) + + # We loop through every auxiliary input + for name, specification in auxiliary.items(): + # For a raster or point cloud, the input support is simply its coordinates + # For arrays, we use the input specified for that auxiliary + if isinstance(specification, tuple): + value, selector = _sampling_specification(first, specification) + else: + value, selector = specification, None + input_support_name = auxiliary_at.get(name) + native = inputs[input_support_name] if input_support_name is not None else None + inputs[name] = _prepare_cosample_input(value, selector, name, native) + + return inputs + + +def _check_cosample_input_types( + inputs: Mapping[str, _CosampleInput], + support: RasterBase | PointCloudBase, + mask: RasterLike | VectorLike | ArrayLike | None, +) -> None: + """ + Require inputs to be of the same object or accessor "family" (GeoUtils, or Xarray/Pandas). + + However, DataArrays and GeoDataFrames can mix eager and Dask. Plain arrays have to match their input locations, + and vector outlines only supply a mask, so both of them are accepted in every case. + """ + + # Include all input supports + values = [(name, input_data.input_support) for name, input_data in inputs.items()] + values.append(("at", support)) + + # On a grid, vector masks supply shapes rather than point values, including lazy geometry tables + if _is_raster(mask) or not _is_raster(support): + mask_interface = _get_raster_interface(mask) + if mask_interface is None: + mask_interface = _get_pointcloud_interface(mask) + if mask_interface is not None: + values.append(("mask", mask_interface)) + first = inputs["self"].input_support + use_accessors = getattr(first, "_is_xr", False) or getattr(first, "_is_pd", False) + + # Check the interface family (independently of whether its arrays or dataframe partitions are lazy) + for name, interface in values: + is_accessor = getattr(interface, "_is_xr", False) or getattr(interface, "_is_pd", False) + if is_accessor != use_accessors: + raise TypeError( + f"Cannot mix Raster/PointCloud objects with DataArray/GeoDataFrame inputs in cosample(): {name!r}. " + "Use one family for all geospatial inputs." + ) + + +def _choose_cosample_support( + inputs: Mapping[str, _CosampleInput], + at: Literal["self", "other"] | RasterLike | PointCloudLike | None, + raster_point_mode: Literal["grid_points", "resample_raster"] | None, +) -> RasterBase | PointCloudBase: + """ + Choose output locations based on the ``at`` input, and the raster-point mode. + + An ``at`` argument has precedence, otherwise the raster-point mode identifies the support. + With neither option, _sampling_support() chooses the first point cloud, or the first raster's grid. + + The arguments fulfill different roles: + - ``at`` can for instance decide which of 2 raster inputs is the reference, which raster_point_mode doesn't affect, + - ``raster_point_mode`` decides on the direction in case of a raster-point comparison (grid points to raster, + or resample raster at point coordinates), but can conflict with ``at`` if defined in the other direction. + """ + + # Choose explicit output locations before using the raster and point conversion direction + if isinstance(at, str): + if at not in {"self", "other"}: + raise ValueError("Argument ``at`` must be 'self', 'other' or a geospatial support object.") + at = inputs[at].input_support + if at is None and raster_point_mode is not None: + candidates = [] + kind = "raster" if raster_point_mode == "grid_points" else "point" + for name in ("self", "other"): + input_data = inputs[name] + if input_data.kind == kind and input_data.value is input_data.input_support: + candidates.append(input_data.input_support) + if len(candidates) != 1: + raise ValueError("The conversion mode requires one unambiguous input support; select ``at`` explicitly.") + at = candidates[0] + + # Main call to select common support + support = _sampling_support((inputs["self"].input_support, inputs["other"].input_support), at) + + # Reject a conversion direction that conflicts with the chosen output locations + support_is_raster = _is_raster(support) + if (support_is_raster and raster_point_mode == "resample_raster") or ( + not support_is_raster and raster_point_mode == "grid_points" + ): + raise ValueError( + "Argument ``raster_point_mode`` conflicts with the grid or point locations selected by ``at``." + ) + return support + + +######################################## +# 2/ COMMON ALIGNMENT AND VALIDITY +######################################## + + +def _align_cosample_inputs_for_raster_support( + inputs: Mapping[str, _CosampleInput], + support: RasterBase, + grid_method: GriddingMethod, + grid_kwargs: Mapping[str, Any], + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None, + temporary_files: ExitStack, +) -> tuple[dict[str, Any], dict[str, tuple[RasterBase, int]]]: + """ + Align every input with the output support, keeping Dask/MP/eager support. + + Point inputs are gridded, and rasters reprojected to the support grid if necessary. + NumPy and Dask use the selected arrays, while multiprocessing uses raster objects and band indexes that + workers can read by tile. + """ + + from geoutils.pointcloud.dataframe import ( + _assign_point_values, + _build_pointcloud_output, + _get_dataframe_attrs, + ) + + # Store rasters and band numbers for multiprocessing, or store arrays for NumPy and Dask + aligned_rasters = {} + arrays = {} + aligned_points = {} + for name, input_data in inputs.items(): + value = input_data.value + selected_band = cast(int, input_data.selector) if input_data.kind == "raster" else 1 + input_support = input_data.input_support + + # Give each gridded or reprojected input its own temporary file for multiprocessing + intermediate = temporary_files.enter_context(mp_config.temporary()) if mp_config is not None else None + + # Reuse each point projected coordinates for its spatial values and plain arrays + if input_data.kind == "point": + owner = cast("PointCloudBase", input_support) + if id(owner) not in aligned_points: + point_config = temporary_files.enter_context(mp_config.temporary()) if mp_config is not None else None + aligned_points[id(owner)] = _aligned_pointcloud(owner, support, name, align, mp_config=point_config) + pointcloud = aligned_points[id(owner)] + if value is not owner: + # Replace masked values with NaN and attach arrays by position, including single points and Dask chunks + raw = value + if np.ma.isMaskedArray(raw): + raw = np.where(np.ma.getmaskarray(raw), np.nan, np.ma.getdata(raw)) + geometry = pointcloud.ds[[pointcloud.ds.geometry.name]] + if geometry.geometry.name != "geometry": + geometry = geometry.rename_geometry("geometry") + dataframe = _assign_point_values(geometry, {name: raw}) + + # Select the array column while keeping the owner's coordinates and spatial metadata + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", message="Overriding 3D points with with data column", category=UserWarning + ) + copied = _build_pointcloud_output( + dataframe, + data_column=name, + as_dataframe=pointcloud._is_pd, + attrs=_get_dataframe_attrs(pointcloud.ds), + preserve_locations=True, + ) + pointcloud = _get_pointcloud_interface(copied) + + # Select point columns without copying or loading the source, then calculate one raster band + value = pointcloud.grid( + ref=support, + resampling=grid_method, + data_column=cast("str | None", input_data.selector), + mp_config=intermediate, + **grid_kwargs, + ) + + # Align every input once, then keep its raster for workers or read its selected NumPy/Dask band + raster = _aligned_raster(value, input_support, support, name, align, mp_config=intermediate) + if mp_config is not None: + aligned_rasters[name] = (raster, selected_band) + else: + arrays[name] = _selected_raster_data(raster, selected_band) + + return arrays, aligned_rasters + + +def _align_cosample_inputs_for_point_support( + inputs: Mapping[str, _CosampleInput], + support: PointCloudBase, + partition_lengths: tuple[int, ...] | None, + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None, + temporary_files: ExitStack, +) -> tuple[dict[str, Any], dict[str, tuple[RasterBase, int]]]: + """ + Align every input with the output points and keep the representation needed for later sampling. + + Point inputs are checked lazily at the output support (as inputs need to be aligned already), while rasters are + reprojected to the right CRS, and retained for later interpolation after the common validity mask is derived. + """ + + # Separate grid inputs from values already located at the output points + point_values = {} + grid_inputs = {} + aligned_points = {} + for name, input_data in inputs.items(): + if input_data.kind == "raster": + grid_inputs[name] = input_data + continue + + # Reject incompatible coordinates before any raster reprojection or interpolation starts + owner = cast("PointCloudBase", input_data.input_support) + if id(owner) not in aligned_points: + intermediate = temporary_files.enter_context(mp_config.temporary()) if mp_config is not None else None + aligned_points[id(owner)] = _aligned_pointcloud(owner, support, name, align, mp_config=intermediate) + value = aligned_points[id(owner)] if input_data.value is owner else input_data.value + point_values[name] = _point_values_at_support( + cast("PointCloudBase | ArrayLike", value), + input_data.selector, + support_dataframe=support.ds, + name=name, + point_partition_lengths=partition_lengths, + ) + + # Give each raster reprojection its own temporary file and keep selected bands for deferred interpolation + aligned_rasters = {} + for name, input_data in grid_inputs.items(): + intermediate = temporary_files.enter_context(mp_config.temporary()) if mp_config is not None else None + raster = _aligned_raster( + input_data.value, input_data.input_support, support, name, align, mp_config=intermediate + ) + aligned_rasters[name] = (raster, cast(int, input_data.selector)) + return point_values, aligned_rasters + + +def _intersect_validity(validity_layers: Iterable[Any]) -> Any: + """Intersect boolean validity raster/pointcloud layers without computing lazy arrays.""" + + # Combine finite coverage and mask eligibility while preserving the input array backend + common_validity = None + for validity in validity_layers: + if validity is None: + continue + common_validity = validity if common_validity is None else common_validity & validity + + # Every cosampling path supplies validity from at least one primary input + if common_validity is None: + raise RuntimeError("Cosampling requires at least one validity layer.") + return common_validity + + +############################### +# 3/ COSAMPLE ON RASTER SUPPORT +############################### + + +def _cosample_raster_eager( + arrays: Mapping[str, NDArrayNum], + common_validity: NDArrayBool, + subsample: int | float, + random_state: int | np.random.Generator | None, + strategy: Literal["sequential", "topk"], +) -> NDArrayNum: + """Cosample valid locations eagerly (in-memory) into a multi-band array output.""" + + # Use all valid cells when no smaller sample was requested + if subsample == 1: + selected = common_validity + if not np.any(common_validity): + raise ValueError("There is no finite data common to all cosampled values.") + else: + # Subsample pixel locations randomly for all output bands + rows, columns = _sample_valid_indices( + common_validity, subsample=subsample, random_state=random_state, strategy=strategy + ) + if rows.size == 0: + raise ValueError("There is no finite data common to all cosampled values.") + selected = np.zeros(common_validity.shape, dtype=bool) + selected[rows, columns] = True + + # Stack self, other and auxiliary arrays as output bands, setting unselected cells to NaN + return np.stack([np.where(selected, array, np.nan) for array in arrays.values()]) + + +def _cosample_raster_dask( + arrays: Mapping[str, Any], + common_validity: Any, + subsample: int | float, + random_state: int | np.random.Generator | None, + strategy: Literal["sequential", "topk"], +) -> Any: + """ + Cosample valid locations lazily across chunks into a multi-band Dask array output. + + Same logic as eager, but written in Dask. + """ + + import_optional("dask") + import dask.array as da + + # Use all valid cells when no smaller sample was requested + if subsample == 1: + selected = common_validity + + # Check that at least one cell is valid without computing the complete Dask output + if not bool(common_validity.any().compute()): + raise ValueError("There is no finite data common to all cosampled values.") + else: + # Randomly choose one set of cell positions for every output band + rows, columns = _sample_valid_indices( + common_validity, subsample=subsample, random_state=random_state, strategy=strategy + ) + if rows.size == 0: + raise ValueError("There is no finite data common to all cosampled values.") + + # Give each cell a unique number using row * width + column + # Mark sampled cell numbers within each Dask chunk instead of loading the full mask into memory + grid_rows = da.arange(common_validity.shape[0], chunks=common_validity.chunks[0])[:, None] + grid_columns = da.arange(common_validity.shape[1], chunks=common_validity.chunks[1])[None, :] + selected = da.isin( + grid_rows * common_validity.shape[1] + grid_columns, + rows * common_validity.shape[1] + columns, + ) + + # Stack self, other and auxiliary arrays as output bands, setting unselected cells to NaN + return np.stack([np.where(selected, array, np.nan) for array in arrays.values()]) + + +def _wrapper_cosample_raster_block_mp( + tile: RasterBase, + inputs: Mapping[str, tuple[RasterBase, int]], + support: RasterBase, + mask: RasterLike | VectorLike | ArrayLike | None, + mask_mode: str, + indices: tuple[NDArrayNum, NDArrayNum] | None = None, + validity_only: bool = False, +) -> Raster: + """ + Wrapper for Multiprocessing cosample in input blocks, used in _cosample_raster_mp() with map_overlap(). + + Same logic as eager above, but for a chunk. + """ + + from geoutils.raster.raster import Raster + + # Read the same geographic window from every input and store its requested band in arrays + tile = _get_raster_interface(tile) + arrays = {} + for name, (raster, band) in inputs.items(): + window = tile if name == "self" else raster.crop(tile.bounds) + arrays[name] = _selected_raster_data(window, band) + + # Find this tile's first row and column in the complete output grid + column, row = (~support.transform) * (tile.transform.c, tile.transform.f) + row, column = int(round(row)), int(round(column)) + + # Crop raster masks or slice array masks to this tile; vector masks are evaluated using coordinates + if mask is not None: + if _is_raster(mask): + mask = get_geo_attr(mask, "crop", accessors=("rst",))(tile.bounds) + elif not has_geo_attr(mask, "create_mask", accessors=("vct",)): + mask = _as_array(mask).reshape(support.shape)[row : row + tile.height, column : column + tile.width] + + # Intersect the user mask and finite coverage from every input before selecting any cells + validity_layers = [_mask_at_support(mask, tile, mask_mode=mask_mode)] + validity_layers.extend(np.isfinite(array) for array in arrays.values()) + common_validity = _intersect_validity(validity_layers) + + # Find sampled cells within this tile's rows, then check that their columns also fall inside the tile + # The sample is sorted by row so each tile can look up its cells without scanning the full sample + if indices is not None: + lower, upper = np.searchsorted(indices[0], (row, row + tile.height)) + rows, columns = indices[0][lower:upper] - row, indices[1][lower:upper] - column + inside = (columns >= 0) & (columns < tile.width) + selected = np.zeros(tile.shape, dtype=bool) + selected[rows[inside], columns[inside]] = True + common_validity &= selected + + # Return a band of 1 for valid cells and NaN elsewhere when choosing the sample + # For the final output, return one band per input with excluded cells set to NaN + if validity_only: + data = np.where(common_validity, np.float32(1), np.float32(np.nan)) + else: + data = np.stack([np.where(common_validity, array, np.nan) for array in arrays.values()]) + return Raster.from_array( + data, + tile.transform, + tile.crs, + nodata=np.nan, + area_or_point=support.area_or_point, + tags={} if validity_only else {"long_name": tuple(inputs)}, + ) + + +def _wrapper_has_finite_raster_block(tile: RasterBase) -> bool: + """ + Wrapper for Multiprocessing through map_blocks: check one validity tile for finite cells.""" + + return bool(np.any(np.isfinite(_selected_raster_data(tile)))) + + +def _cosample_raster_mp( + inputs: Mapping[str, tuple[RasterBase, int]], + support: RasterBase, + mask: RasterLike | VectorLike | ArrayLike | None, + mask_mode: str, + subsample: int | float, + random_state: int | np.random.Generator | None, + strategy: Literal["sequential", "topk"], + mp_config: MultiprocConfig, + temporary_files: ExitStack, +) -> Raster: + """ + Cosample valid values from a temporary "validity" file, then write their values by block. + + Same logic as eager above, but chunked using Multiprocessing as backend. + + _wrapper_cosample_raster_block_mp() applies the same mask in both passes. + We keep the temporary "validity" file alive until the values have been written, and sort + sampled rows so each worker can find the cells inside its block. + """ + + from geoutils.multiproc import map_blocks, map_overlap + + # Write 1 where all inputs are finite and the mask allows the cell, and NaN elsewhere + # Select cells from this temporary file before writing the output values + intermediate = temporary_files.enter_context(mp_config.temporary()) + reference = inputs["self"][0] + validity_raster = map_overlap( + _wrapper_cosample_raster_block_mp, reference, intermediate, inputs, support, mask, mask_mode, validity_only=True + ) + + # Check that valid cells exist, or randomly select a subset for all output bands + indices = None + if subsample == 1: + has_valid = any(map_blocks(_wrapper_has_finite_raster_block, validity_raster, intermediate)) + else: + indices = validity_raster.subsample( + subsample, return_indices=True, random_state=random_state, strategy=strategy, mp_config=intermediate + ) + has_valid = len(indices[0]) > 0 + + # Sort sampled cells by row so each tile can find its cells without scanning the full sample + order = np.argsort(indices[0], kind="stable") + indices = indices[0][order], indices[1][order] + if not has_valid: + raise ValueError("There is no finite data common to all cosampled values.") + + # Read each tile again to write the selected values from all inputs to the final output file + return map_overlap( + _wrapper_cosample_raster_block_mp, reference, mp_config, inputs, support, mask, mask_mode, indices=indices + ) + + +def _cosample_on_raster( + first: RasterLike | PointCloudLike, + inputs: Mapping[str, _CosampleInput], + *, + support: RasterBase, + mask: RasterLike | VectorLike | ArrayLike | None, + mask_mode: str, + subsample: int | float, + random_state: int | np.random.Generator | None, + strategy: Literal["sequential", "topk"], + grid_method: GriddingMethod, + grid_kwargs: Mapping[str, Any], + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None, + temporary_files: ExitStack, +) -> RasterLike: + """ + Cosample all inputs at the same raster locations into a multi-band raster output. + + See _cosample() for most input arguments. + + This function does in order: + + - _align_cosample_inputs_for_raster_support() places every input on the grid. + - _intersect_validity() then combines finite coverage with the optional mask into a common validity mask. + - Then an eager/Dask/MP _cosample_raster() function draws one common sample and constructs the multi-band output. + + :param inputs: Named values with their selected band and original locations, prepared by _cosample(). + :param support: Raster defining the output grid, resolved from ``at`` and ``raster_point_mode``. + :param temporary_files: Pass ExitStack to keep intermediate multiprocessing files alive until output is built. + + :returns: Raster with one band per input and a shared mask outside the selected cells. + """ + + # 1/ Align every input with the output grid without loading lazy values + arrays, aligned_rasters = _align_cosample_inputs_for_raster_support( + inputs, support, grid_method, grid_kwargs, align, mp_config, temporary_files + ) + + # 2/ Compute the common validity mask, then select the same sample of grid cells for all aligned inputs + + # 2a/ For Multiprocessing + if mp_config is not None: + # Align raster masks once before workers crop them, and validate array masks against the complete grid + if _is_raster(mask): + intermediate = temporary_files.enter_context(mp_config.temporary()) + mask = _aligned_raster(mask, mask, support, "mask", align, mp_config=intermediate) + elif mask is not None and not has_geo_attr(mask, "create_mask", accessors=("vct",)): + mask = _mask_on_raster(mask, support, mask_mode, align) + + result = _cosample_raster_mp( + aligned_rasters, + support, + mask, + mask_mode, + subsample, + random_state, + strategy, + mp_config, + temporary_files, + ) + + # Return the output file through the caller's raster interface + output = first._cast_raster_output(result) + if isinstance(output, xr.DataArray): + output.attrs["long_name"] = tuple(inputs) + return output + + # 2b/ For Dask or eager, intersect the user mask and finite coverage from every aligned value + validity_layers = [_mask_at_support(mask, support, mask_mode=mask_mode, align=align)] + validity_layers.extend(np.isfinite(array) for array in arrays.values()) + common_validity = _intersect_validity(validity_layers) + + # Subsample only the common valid cells, then stack their values into output bands + if is_dask_array(common_validity): + data = _cosample_raster_dask(arrays, common_validity, subsample, random_state, strategy) + else: + data = _cosample_raster_eager(arrays, common_validity, subsample, random_state, strategy) + + # 3/ Construct the raster output from the NumPy or Dask bands + tags = {"long_name": tuple(arrays)} + + # Return an Xarray for accessor calls without computing its Dask arrays + if getattr(first, "_is_xr", False) or getattr(first, "_is_pd", False): + from geoutils.raster.xr_accessor import RasterAccessor + + return RasterAccessor.from_array( + data, support.transform, support.crs, nodata=np.nan, area_or_point=support.area_or_point, tags=tags + ) + + # Load the result into a Raster + from geoutils.raster.raster import Raster + + data = data.compute() if is_dask_array(data) else data + return Raster.from_array( + data, support.transform, support.crs, nodata=np.nan, area_or_point=support.area_or_point, tags=tags + ) + + +############################## +# 4/ COSAMPLE ON POINT SUPPORT +############################## + + +def _raster_valid_at_points( + raster: RasterBase, + points: PointCloudLike | tuple[NDArrayNum, NDArrayNum], + resample_method: InterpolationMethod, + band: int, + resample_kwargs: Mapping[str, Any], + mp_config: MultiprocConfig | None = None, + point_partition_lengths: tuple[int, ...] | None = None, +) -> Any: + """ + Find valid raster points by interpolating raster block's validity, respecting Dask/MP. + + Arguments follow _cosample_on_points(), except for point_partition_lengths that follows _point_values_at_support(). + """ + + # Interpolate a mask of 1 for finite raster cells and NaN for missing cells to find points with data + # Use no extra nodata spreading unless the caller requests it + values = raster.interp_points( + points=points, + method=resample_method, + band=band, + as_array=not is_dask_dataframe(points), + mp_config=mp_config, + _validity_only=True, + **{"dist_nodata_spread": 0, **resample_kwargs}, + ) + + # Convert the interpolated Dask column to an array, reusing point counts per chunk when available + if is_dask_dataframe(values): + values = get_geo_attr(values, "data", ("pc",)).to_dask_array(lengths=point_partition_lengths) + return np.isfinite(values) + + +def _cosample_on_points( + first: RasterLike | PointCloudLike, + inputs: Mapping[str, _CosampleInput], + *, + support: PointCloudBase, + mask: RasterLike | VectorLike | ArrayLike | None, + mask_mode: str, + subsample: int | float, + random_state: int | np.random.Generator | None, + resample_method: InterpolationMethod, + resample_kwargs: Mapping[str, Any], + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None, + temporary_files: ExitStack, +) -> PointCloudLike: + """ + Cosample all inputs at the same point locations into a multi-column point cloud output. + + See _cosample() for input arguments. + + This function does in order: + - _align_cosample_inputs_for_point_support() checks point locations and separates their values from aligned rasters, + - _raster_valid_at_points() checks raster valid values before defining a common validity mask, + - Finally, raster values are interpolated at the selected points with interp_points(). + + :param support: Point cloud defining the ordered output coordinates, resolved from ``at`` and ``raster_point_mode``. + :param inputs: Named values, raster bands and original locations resolved by _prepare_cosample_inputs(). + :param temporary_files: Pass ExitStack to keep intermediate multiprocessing files alive until output is built. + + :returns: PointCloud with one column per input, retaining selected points in their original order. + """ + + from geoutils.pointcloud.dataframe import ( + _assign_point_values, + _build_pointcloud_output, + _point_partition_lengths, + _select_point_rows, + ) + + # 1/ Read the output point coordinates and align each input for sampling + dataframe = support.ds + is_dask_points = is_dask_dataframe(dataframe) + + # Count rows in each Dask chunk so values and masks can be matched to the same points + # Keep the Dask dataframe to not load all coordinates in memory + partition_lengths = _point_partition_lengths(dataframe) if is_dask_points else None + + # Check every point input before aligning rasters, whose values will be read only at the selected points + point_values, aligned_rasters = _align_cosample_inputs_for_point_support( + inputs, support, partition_lengths, align, mp_config, temporary_files + ) + + # 2/ Compute the common validity mask, then select the same sample of point locations for all aligned inputs + validity_layers = [np.isfinite(values) for values in point_values.values()] + + # Check where each raster has data and combine the result with the masks from the point values + if aligned_rasters: + points = dataframe if is_dask_points else (dataframe.geometry.x.to_numpy(), dataframe.geometry.y.to_numpy()) + for raster, selected_band in aligned_rasters.values(): + finite = _raster_valid_at_points( + raster, points, resample_method, selected_band, resample_kwargs, mp_config, partition_lengths + ) + validity_layers.append(finite) + + # Read the user mask at the output point coordinates and exclude points where it is False + intermediate = ( + temporary_files.enter_context(mp_config.temporary()) if mp_config is not None and mask is not None else None + ) + mask_values = _mask_at_support( + mask, + support, + support_dataframe=dataframe, + mask_mode=mask_mode, + align=align, + mp_config=intermediate, + point_partition_lengths=partition_lengths, + ) + validity_layers.append(mask_values) + + # Intersect all finite coverage and mask eligibility before drawing one common sample + common_validity = _intersect_validity(validity_layers) + + # Keep all remaining points, or randomly choose a smaller sample as requested + selected_rows = common_validity + if subsample != 1: + (indices,) = _sample_valid_indices( + common_validity, subsample=subsample, random_state=random_state, strategy="sequential" + ) + if indices.size == 0: + raise ValueError("There is no finite data common to all cosampled values.") + + # Sort the selected row numbers so the output follows the original point order + selected_rows = np.sort(indices) + elif is_dask_points and not bool(common_validity.any().compute()): + raise ValueError("There is no finite data common to all cosampled values.") + + # Create a table with only geometry and the requested point columns, then keep the selected rows + geometry = dataframe[[dataframe.geometry.name]] + if geometry.geometry.name != "geometry": + geometry = geometry.rename_geometry("geometry") + point_columns = _assign_point_values(geometry, point_values, partition_lengths=partition_lengths) + output = _select_point_rows(point_columns, selected_rows, partition_lengths=partition_lengths) + if not is_dask_dataframe(output) and output.empty: + raise ValueError("There is no finite data common to all cosampled values.") + + # 3/ Interpolate each raster at the selected coordinates and add its values as a new column + if aligned_rasters: + selected_points = ( + output if is_dask_dataframe(output) else (output.geometry.x.to_numpy(), output.geometry.y.to_numpy()) + ) + sampled = {} + for name, (raster, selected_band) in aligned_rasters.items(): + values = raster.interp_points( + points=selected_points, + method=resample_method, + band=selected_band, + as_array=not is_dask_dataframe(output), + mp_config=mp_config, + **resample_kwargs, + ) + + # Extract the interpolated column as a Dask Series so it keeps the same chunks as the selected points + sampled[name] = get_geo_attr(values, "data", ("pc",)) if is_dask_dataframe(values) else values + + # Final interpolation may exclude more points near nodata; remove rows with missing or infinite values + output = _assign_point_values(output, sampled) + output = output.replace([np.inf, -np.inf], np.nan).dropna(subset=list(inputs)) + + # Order all columns as self, other, auxiliaries, then geometry + output = output[[*inputs, "geometry"]] + + # Build the point output with self as its active column and metadata for the selected rows + # Accessor calls keep Dask data chunked; PointCloud calls load the result into memory + as_dataframe = getattr(first, "_is_xr", False) or getattr(first, "_is_pd", False) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", message="Overriding 3D points with with data column 'self'", category=UserWarning + ) + result = _build_pointcloud_output(output, data_column="self", as_dataframe=as_dataframe) + + # Check for empty results only when loaded, so Dask does not run the final interpolation yet + if not is_dask_dataframe(result) and get_geo_attr(result, "point_count", ("pc",)) == 0: + raise ValueError("There is no finite data common to all cosampled values.") + return result + + +########################### +# 5/ MAIN COSAMPLE FUNCTION +########################### + + +def _cosample( + first: RasterLike | PointCloudLike, + second: RasterLike | PointCloudLike | ArrayLike, + *, + band: int, + other_band: int, + auxiliary: Mapping[str, Any] | None, + auxiliary_at: Literal["self", "other"] | Mapping[str, Literal["self", "other"]] | None, + at: Literal["self", "other"] | RasterLike | PointCloudLike | None, + mask: RasterLike | VectorLike | ArrayLike | None, + mask_mode: Literal["inside", "outside"], + subsample: int | float, + random_state: int | np.random.Generator | None, + strategy: Literal["sequential", "topk"], + raster_point_mode: Literal["grid_points", "resample_raster"] | None, + grid_method: GriddingMethod, + resample_method: InterpolationMethod | Literal["reduce"], + grid_kwargs: Mapping[str, Any] | None, + resample_kwargs: Mapping[str, Any] | None, + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None = None, +) -> RasterLike | PointCloudLike: + """ + Cosample two datasets at the same locations, potentially with same-shape auxiliary data tied to them. + + Two primary inputs can be rasters or point clouds, and auxiliary outputs can also be arrays but need to match the + shape of one of the two main inputs. + The output is a multi-band raster or multi-column point cloud containing all data sampled at valid values of the + same locations. + + This function reuses eager, Dask and multiprocessing execution from other geospatial operations (grid, + interp_points). Additional steps include placing values on common support, determine validity (non-NaN/inf), + then drawing one shared sample. + Specifically: + - Dask retains array chunks or point partitions in a graph. + - Multiprocessing uses temporary raster files and map_overlap() to write the final bands without holding chunks + in-memory at once. + + This main function has the following steps: + + - First, _prepare_cosample_inputs() resolves each input's values, band and original locations. + - Then _choose_cosample_support() identifies the shared output locations and _check_cosample_input_types() + enforces consistent object types (GeoUtils or Xarray/Pandas). + - Finally, _cosample_on_raster() or _cosample_on_points() aligns the inputs, computes a common validity mask, then + selects inputs at the same location (with optional subsampling), and finally builds the output. + + :param first: First raster or point cloud, whose selected values become the "self" output. + :param second: Second raster or point cloud to sample alongside the first. An array has to match the first input + grid shape or point count. + :param band: Band selected from the first raster, counting from one. Point clouds use their main data column. + :param other_band: Band selected from the second input if it is a raster, counting from one. + :param auxiliary: Additional values by output name (e.g. {"slope": slope_raster}). Select a raster band or + point column with a pair, e.g. {"slope": (slope_raster, 2)} or {"intensity": (points, "intensity")}. + Spatial inputs default to the first raster band or active point values. + :param auxiliary_at: Input locations followed by plain auxiliary arrays: "self", "other", or a choice per name + (e.g. {"slope": "other"}). Spatial auxiliaries use their own coordinates. + :param at: Output locations: "self", "other", or a reference raster/point cloud. Defaults to the first point + cloud, or first raster's grid when neither input is a point cloud. + :param mask: Locations eligible for sampling, defined by a boolean array, spatial mask, or vector outlines. + :param mask_mode: Whether a vector mask keeps locations "inside" or "outside" its geometries. + :param subsample: Fraction of common finite locations (e.g. 0.1), or maximum count (e.g. 1000); 1 keeps all. + :param random_state: Seed or random generator for reproducible sampling (e.g. 42). + :param strategy: Raster sampling with "topk" or "sequential"; "topk" keeps the same seeded sample across chunk + sizes. Point output always uses "sequential". + :param raster_point_mode: Conversion direction: "grid_points" places points on a raster, "resample_raster" reads + rasters at points. Defaults to at's locations, or point locations when available. Must agree with at. + :param grid_method: Point gridding by SciPy interpolation ("nearest", "linear", "cubic"), or circular "idw", + "mean", "minimum", "maximum", "range", "count", "stdev", "average_distance", "average_distance_pts". + The aliases "average", "min" and "max" select "mean", "minimum" and "maximum". + :param resample_method: Raster interpolation using the SciPy methods "nearest", "linear", "cubic", "quintic", + "slinear", "pchip" or "splinef2d". Window reduction ("reduce") is not implemented. + :param grid_kwargs: Options for PointCloud.grid(), e.g. {"dist_nodata_pixel": 2, "min_points": 3} sets a two-pixel + radius and minimum of three finite points for circular methods. Other options include "distance_power" for + IDW and "engine" ("scipy" or "numba"). Set output locations and method with at and grid_method. + :param resample_kwargs: Options for Raster.interp_points(), e.g. {"nodata_propagation": "ignore"}. The nodata + policies are "gdal", "ignore" and "propagate"; "dist_nodata_spread" controls extra spreading in pixels. + Set locations, band and method with the corresponding cosample() arguments. + :param align: Handling of mismatched grids or coordinate systems: "raise" an error, or "reproject" to match at. + Point inputs must still share the same ordered coordinates when sampled at points. + :param mp_config: Worker and tile settings for multiprocessing. Raster output uses its outfile; cannot be + combined with Dask inputs. + + :returns: Raster bands or point cloud columns named 'self', 'other' and the auxiliaries. + All values share the same finite locations; raster cells outside the sample remain masked. + """ + + # 1/ Check input arguments and raise appropriate errors + # Basic type/value checks + if second is None: + raise TypeError("Argument ``other`` is required for cosample().") + if mask_mode not in {"inside", "outside"}: + raise ValueError("Argument ``mask_mode`` must be 'inside' or 'outside'.") + if strategy not in {"sequential", "topk"}: + raise ValueError("Argument ``strategy`` must be 'sequential' or 'topk'.") + if align not in {"raise", "reproject"}: + raise ValueError("Argument ``align`` must be 'raise' or 'reproject'.") + if not isinstance(subsample, (int, float)) or subsample <= 0: + raise ValueError("Argument ``subsample`` must be a positive number.") + if raster_point_mode not in {None, "grid_points", "resample_raster"}: + raise ValueError("Argument ``raster_point_mode`` must be 'grid_points', 'resample_raster' or None.") + + # Copy extra options, and reject arguments that should be passed to cosample() directly and not as kwargs + # (This check is required because interp_points() contains similarly-named inputs as cosample(), such as + # "mp_config" or "band", etc) + grid_kwargs = {} if grid_kwargs is None else dict(grid_kwargs) + resample_kwargs = {} if resample_kwargs is None else dict(resample_kwargs) + if {"ref", "grid_coords", "res", "shape", "bounds", "resampling", "data_column", "mp_config"}.intersection( + grid_kwargs + ): + raise ValueError( + "Use ``at``, ``grid_method``, point column selectors and ``mp_config`` outside ``grid_kwargs`` " + "to choose the grid, method, point column and backend." + ) + if { + "points", + "method", + "band", + "as_array", + "input_latlon", + "return_interpolator", + "mp_config", + "_validity_only", + }.intersection(resample_kwargs): + raise ValueError( + "Use ``at``, ``band``, ``resample_method`` and ``mp_config`` outside ``resample_kwargs``; " + "validity is managed internally." + ) + + # 2/ Resolve input locations and define the shared output support + inputs = _prepare_cosample_inputs(first, second, band, other_band, auxiliary, auxiliary_at) + mask = _normalize_sampling_input(mask) + support = _choose_cosample_support(inputs, at, raster_point_mode) + _check_cosample_input_types(inputs, support, mask) + + # Find the raster or point cloud interface that defines the output locations + support_is_raster = _is_raster(support) + if not support_is_raster and resample_method == "reduce": + raise NotImplementedError( + "Window reduction in cosample awaits revision of Raster.reduce_points(); " + "use reduce_points separately in the meantime." + ) + + # 3/ Select execution backend and, for Multiprocessing, keep temporary files alive until the output is complete + + # If any input is Dask but mp_config was passed, raise an error + if mp_config is not None: + input_values = [input_data.value for input_data in inputs.values()] + for value in [*input_values, support, mask]: + # Inspect raw collections and spatial metadata without reading file-backed DataArray values + lazy = is_dask_array(value) or is_dask_dataframe(value) + if has_geo_attr(value, "_chunks", accessors=("rst",)): + lazy |= get_geo_attr(value, "_chunks", accessors=("rst",)) is not None + if has_geo_attr(value, "_is_dask", accessors=("pc", "vct")): + lazy |= get_geo_attr(value, "_is_dask", accessors=("pc", "vct")) + if lazy: + raise ValueError("Cannot use Multiprocessing and Dask simultaneously in cosample().") + + # Keep temporary files alive until end of execution with ExitStack() + with ExitStack() as temporary_files: + # Dispatch only the kwargs relevant to the support operation (grid() or interp_points()) + if support_is_raster: + return _cosample_on_raster( + first, + inputs, + support=cast("RasterBase", support), + mask=mask, + mask_mode=mask_mode, + subsample=subsample, + random_state=random_state, + strategy=strategy, + grid_method=grid_method, + grid_kwargs=grid_kwargs, + align=align, + mp_config=mp_config, + temporary_files=temporary_files, + ) + + return _cosample_on_points( + first, + inputs, + support=cast("PointCloudBase", support), + mask=mask, + mask_mode=mask_mode, + subsample=subsample, + random_state=random_state, + resample_method=cast("InterpolationMethod", resample_method), + resample_kwargs=resample_kwargs, + align=align, + mp_config=mp_config, + temporary_files=temporary_files, + ) diff --git a/geoutils/sampling/pairsampling.py b/geoutils/sampling/pairsampling.py new file mode 100644 index 000000000..b16473e64 --- /dev/null +++ b/geoutils/sampling/pairsampling.py @@ -0,0 +1,1493 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +""" +Sample pairs of raster cells or point cloud rows. + +Note: This module is inspired from code originally developed in xDEM and SciKit-GStat for uncertainty quantification. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import xarray as xr +from scipy.spatial import cKDTree + +from geoutils._dispatch import ( + _get_pointcloud_interface, + get_geo_attr, + is_dask_array, + is_dask_dataframe, +) +from geoutils._misc import import_optional +from geoutils._typing import ArrayLike, NDArrayNum +from geoutils.raster.array import _selected_raster_data +from geoutils.sampling.support import _mask_at_support, _mask_on_raster + +if TYPE_CHECKING: + from geoutils.pointcloud.base import PointCloudBase + from geoutils.pointcloud.pointcloud import PointCloudLike + from geoutils.raster.base import RasterBase, RasterLike + from geoutils.vector.base import VectorLike + +############################# +# 1/ SHARED PAIR OPERATIONS +############################# + + +def _read_raster_pair_values(array: Any, first: NDArrayNum, second: NDArrayNum) -> tuple[NDArrayNum, NDArrayNum]: + """ + Read both endpoint vectors together so Dask shares source chunks between their selections. + + :param array: Two-dimensional NumPy or Dask array containing the selected raster band. + :param first: Flat cell indexes for the first endpoints, calculated as row * raster width + column. + :param second: Matching flat cell indexes for the second endpoints. + :returns: Two one-dimensional arrays of endpoint values, with missing values represented by NaN. + """ + + # Convert each endpoint's flat indexes to the corresponding raster row and column positions + use_dask = is_dask_array(array) + selections = [] + for indexes in (first, second): + rows, columns = np.divmod(np.asarray(indexes, dtype=np.int64), int(array.shape[1])) + selections.append(array.vindex[rows, columns] if use_dask else array[rows, columns]) + + # Compute both lazy selections in one graph without interleaving their endpoint buffers + if use_dask: + selections = list(import_optional("dask").compute(*selections)) + + # Express any masked endpoint as NaN before callers test whether its value is available + for index, values in enumerate(selections): + if np.ma.isMaskedArray(values): + values = values.filled(np.nan) + selections[index] = np.asarray(values) + return selections[0], selections[1] + + +def _deduplicate_pairs(first: NDArrayNum, second: NDArrayNum, *, n_observations: int) -> tuple[NDArrayNum, NDArrayNum]: + """ + Remove repeated pairs while treating A-B and B-A as the same pair. + + Endpoint arrays follow the layout described by _pair_dataset(). + + :param n_observations: Total number of source cells or rows, used to give each unordered pair a unique integer key. + :returns: First and second endpoint indexes with the smaller index first, preserving first occurrence order. + """ + + # Put the smaller row number first so A-B and B-A have the same key + low = np.minimum(first, second).astype(np.int64, copy=False) + high = np.maximum(first, second).astype(np.int64, copy=False) + keys = low * np.int64(n_observations) + high + + # Restore input order after selecting the first occurrence of each key + _, positions = np.unique(keys, return_index=True) + positions.sort() + return low[positions], high[positions] + + +############################ +# 2/ REGULAR RASTER SAMPLING +############################ + + +class _RegularPairSampler: + """ + Pair sampler for isotropic log-lag Monte Carlo sampling of a regular 2D grid. + + This method deals efficiently with large datasets by supporting Dask arrays for out-of-memory subsampling, and by + anticipating the probability of nodata occurring in pairs (with iterative top-up). To sample short to long lags + efficiently for variography, pairs are sampled by drawing separation vectors with log-uniform magnitude and + uniformly distributed orientation, corresponding to an isotropic Monte Carlo sampling of logarithmic spatial lags. + + References + ---------- + Sampling method is inspired from the log-lag sampling developed in Hugonnet et al. (2022), Section V-C and + Supplementary Section II-C. + + Few literature references exist that describe this algorithm specifically. However, it was + conceptualized fairly early, such as in Cressie (1993): + "In isotropic settings, all directions are equivalent, and lag distances may be grouped on a + logarithmic scale to ensure adequate sampling across short and long ranges." + + Or in fluid mechanics and turbulence, following Monin and Yaglom (1971): + "For isotropic turbulence, ensemble averages over separation vectors are replaced by averages over + uniformly distributed directions and logarithmically spaced magnitudes." + + Hugonnet et al. (2022): http://dx.doi.org/10.1109/JSTARS.2022.3188922 + Cressie (1993): http://dx.doi.org/10.1002/9781119115151 + Monin and Yaglom (1971). Statistical Fluid Mechanics: Mechanics of Turbulence (Vol. I). The MIT Press. + + The code expands an early version written in SciKit-GStat (as "RasterEquidistantMetricSpace"). + + Use in GeoUtils + --------------- + + sample() returns the first and second endpoint indices and their distances. _sample_raster_pairs() reads the + selected raster values and builds the labelled Xarray result; all other methods are internal. + + Summary of algo + --------------- + We want to sample a large number of point pairs from a regular 2D grid such that separation distances cover + short and long lags efficiently (approximately log-uniform in distance). We cannot enumerate all pairs due to + the size of the grid, so we also subsample. + + The core log-lag sampling is to: + 1) Subsample a distance r ~ Uniform(log(min_distance), log(max_distance)) (log-uniform in r) + 2) Subsample an angle θ ~ Uniform(0, 2π) (isotropic) + 3) Convert (r, θ) to integer pixel offsets (ix, iy) by rounding + 4) Choose origins and compute targets using the offsets + 5) Reject out-of-bounds pairs and (optionally) reject NaN endpoints + 6) Avoid pair duplication during sampling to circumvent costly duplicate removal + + Dask specifics + ------------- + - We never load the full array in memory. + - We count finite cells once at the start to cap the request at the number of distinct pairs that can exist. + - Then, iterating until top-up of valid values, for each candidate batch we read valid (finite) values at + sampled indices using `vindex` (out-of-memory). + + Strategies (strategy) + --------------------- + - "independent": + Each pair is generated independently (origin + offset). This is the basic method that is moderately efficient. + For 1M pairs, we have to sample 1M + 1M points (heavy graph with Dask.vindex). + + - "anchors": + We reuse a set of random anchor points for one endpoint of each pair. Targets are generated relative to these + anchors, so that anchor values (and their chunks) are reused across many pairs. For instance, for 1M pairs, + we index 1000 anchors points that each match 1000 random points, so in the end we index 1k + 1M points + (we thus use half the sample size of "independent"). + + - "chunk_anchors": + Like "anchors" but anchors are sampled from a small set of chunks per round to reduce chunk fan-out and + task overhead for Dask. This method seems to perform the best overall in both speed and memory (default). + + - "anchor_batched": + Structured generation: for each anchor sample multiple distances (log-uniform) and for each distance + sample multiple angles. This produces blocks of pairs with shared anchors and controlled lag coverage. + There might be some room for improvement in this method... which could make it more efficient to sample less + chunks for one vindex (mostly affects speed, while batch size limits temporary memory). + + Hybrid local/global (hybrid_local_fraction) + ------------------------------------------- + Optionally, one can require that a fraction (or all) of pairs remains within the same origin chunk. + This is situational (mostly for short-range variograms), but can massively reduce I/O overhead (both + endpoints are always in the same chunk). The other pairs are sampled "globally" to preserve long-range lag coverage. + + NaN handling + ------------ + To deal with NaNs without knowing their distribution ahead (Dask array), the following steps are applied: + 1. We estimate the global finite fraction f_valid by a single reduction (counting chunk per chunk), and deduce the + probability of a random pair containing at least 1 NaN: p_pair_valid ≈ f_valid^2. For instance, 10% of NaNs + in the array gives us an 81% chance of selecting a valid pair at random. + 2. We oversample so that random pairs will roughly match requested samples, then filter out pairs where + either endpoint is not finite (NaN/inf). + 3. We iterate (top-up) sampling until target count of valid pairs is reached. This is typically not + critical for variography (it rarely matters if sample count is slightly larger or smaller). + + Notes on scalability + -------------------- + Storing more than 1e8 endpoint pairs and distances is RAM-heavy regardless of strategy. + We use int32 indices when possible to reduce memory footprint. + """ + + ################# + # CONFIGURATION + ################# + + def __init__( + self, + array: Any, + *, + # Raster geometry and target sample size + dx: float, + dy: float, + n_pairs: int, + # Distance range + min_distance: float, + max_distance: float, + # Log-lag sampling strategies with various chunk-compatibility + strategy: Literal["independent", "anchors", "chunk_anchors", "anchor_batched"], + # Deduplication + deduplicate: Literal["none", "per_anchor", "global"], + # Random seed + random_state: int | np.random.Generator | None, + # Batching / Termination + batch_pairs: int, + max_rounds: int, + max_oversample: float, + # Chunk / Locality + chunks_per_round: int, + anchors_per_round: int, + # Parameters for anchor_batched + distances_per_anchor: int, + angles_per_distance: int, + # Hybrid local/global + hybrid_local_fraction: float, + max_local_distance: float | None, + # Dtypes to optimize memory usage + index_dtype: Any, + distance_dtype: Any, + ) -> None: + """ + Pair sampling on a regular raster grid. + + :param array: 2D NumPy or Dask array of shape (ny, nx). Values may include NaNs. For Dask arrays, value access + stays lazy until small vectors are computed internally for finiteness checks. + :param dx: Horizontal pixel spacing in coordinate units, such as meters. + :param dy: Vertical pixel spacing in coordinate units, such as meters. + :param n_pairs: Target number of valid pairs with two finite endpoints. + :param min_distance: Smallest distance included in log-distance sampling. + :param max_distance: Largest distance included in log-distance sampling. + :param strategy: Pair generation strategy: "independent", "anchors", "chunk_anchors", or "anchor_batched". + See the class docstring for details and performance trade-offs. + :param deduplicate: Duplicate handling: "global" sorts pairs at the end, "per_anchor" avoids duplicate targets + for each anchor, and "none" skips removal. Duplicate pairs should be avoided for variography because they + bias the distance distribution. + :param random_state: Seed or NumPy Generator used for reproducible random sampling. + :param batch_pairs: Maximum candidate pairs generated per round before NaN filtering. Larger batches reduce + Python and Dask scheduling overhead but require more memory for temporary arrays. + :param max_rounds: Maximum number of top-up rounds used to reach ``n_pairs`` valid pairs. Extra rounds help + when NaNs are clustered or local constraints lower the acceptance rate. + :param max_oversample: Maximum candidate multiplier relative to ``n_pairs``. This prevents very large temporary + arrays when the finite fraction is small. + :param chunks_per_round: Number of chunks selected for anchors in chunk-aligned strategies. Smaller values + improve I/O locality but reduce spatial coverage per round. + :param anchors_per_round: Number of first endpoints drawn per round by anchor-based strategies. + :param distances_per_anchor: Number of log-uniform radii drawn per anchor by "anchor_batched". + :param angles_per_distance: Number of directions drawn for each radius by "anchor_batched". + :param hybrid_local_fraction: Fraction of candidate pairs forced to remain in the first endpoint's chunk. Zero + gives a fully global sample; one keeps every pair local. + :param max_local_distance: Largest distance used for local pairs. Defaults to the chunk diagonal when omitted. + Larger values allow longer local lags but increase rejection at chunk boundaries. + :param index_dtype: Integer dtype for returned endpoint indices. int32 reduces memory when it can represent all + raster cells. + :param distance_dtype: Floating dtype for returned distances. float32 uses half the memory of float64. + """ + + # Store the grid and sampling options with consistent numeric types + self.array = array + self.shape = (int(array.shape[0]), int(array.shape[1])) + self.size = int(np.prod(self.shape)) + self.dx, self.dy = float(abs(dx)), float(abs(dy)) + self.n_pairs = int(n_pairs) + self.min_distance, self.max_distance = float(min_distance), float(max_distance) + self.strategy, self.deduplicate = strategy, deduplicate + self.rng = ( + random_state if isinstance(random_state, np.random.Generator) else np.random.default_rng(random_state) + ) + self.batch_pairs, self.max_rounds = int(batch_pairs), int(max_rounds) + self.max_oversample = float(max_oversample) + self.chunks_per_round, self.anchors_per_round = int(chunks_per_round), int(anchors_per_round) + self.distances_per_anchor, self.angles_per_distance = int(distances_per_anchor), int(angles_per_distance) + self.hybrid_local_fraction = float(hybrid_local_fraction) + self.index_dtype, self.distance_dtype = np.dtype(index_dtype), np.dtype(distance_dtype) + + # Check sampling options before creating any temporary arrays + if self.n_pairs < 1: + raise ValueError("Argument ``n_pairs`` must be a positive integer.") + if not 0 < self.min_distance < self.max_distance: + raise ValueError("Require 0 < ``min_distance`` < ``max_distance``.") + if strategy not in {"independent", "anchors", "chunk_anchors", "anchor_batched"}: + raise ValueError("Unknown regular grid pair sampling ``strategy``.") + if deduplicate not in {"none", "per_anchor", "global"}: + raise ValueError("Argument ``deduplicate`` must be 'none', 'per_anchor' or 'global'.") + if not 0 <= self.hybrid_local_fraction <= 1: + raise ValueError("Argument ``hybrid_local_fraction`` must be between 0 and 1.") + if min(self.batch_pairs, self.max_rounds, self.chunks_per_round, self.anchors_per_round) < 1: + raise ValueError("Batch, round, chunk, and anchor controls must be positive integers.") + if min(self.distances_per_anchor, self.angles_per_distance) < 1 or self.max_oversample <= 0: + raise ValueError("Distance, angle, and oversampling controls must be strictly positive.") + + # Use Dask chunks as local areas, or split an in-memory raster into similarly sized areas + if is_dask_array(array): + self.chunk_edges = tuple(np.r_[0, np.cumsum(chunks)] for chunks in array.chunks) + else: + self.chunk_edges = tuple(np.r_[np.arange(0, size, 2048), size] for size in self.shape) + chunk_rows, chunk_columns = (int(np.max(np.diff(edges))) for edges in self.chunk_edges) + self.max_local_distance = ( + float(np.hypot((chunk_columns - 1) * self.dx, (chunk_rows - 1) * self.dy)) + if max_local_distance is None + else float(max_local_distance) + ) + + ################### + # POSSIBLE PAIRS + ################### + + def _offsets(self, count: int, maximum: float) -> tuple[NDArrayNum, NDArrayNum]: + """ + Draw random directions with distance ranges represented evenly on a log scale. + + :param count: Number of proposed offsets before rounding and distance checks. + :param maximum: Requested upper distance limit in coordinate units, capped by max_distance. When it does not + exceed min_distance, use the sampler's full distance interval instead. + :returns: Matching row and column offsets, excluding zero offsets and distances outside the limits. + """ + + # Limit local distances to the configured nearby area + upper = min(self.max_distance, maximum) + if upper <= self.min_distance: + upper = self.max_distance + + # Draw distances and directions, then round them to row and column offsets + radius = np.exp(self.rng.uniform(np.log(self.min_distance), np.log(upper), count)) + angle = self.rng.uniform(0, 2 * np.pi, count) + column_offset = np.rint(radius * np.cos(angle) / self.dx).astype(np.int64) + row_offset = np.rint(radius * np.sin(angle) / self.dy).astype(np.int64) + + # Remove rounded offsets whose exact grid distance falls outside the requested range + exact_distance = np.hypot(column_offset * self.dx, row_offset * self.dy) + in_range = ( + ((row_offset != 0) | (column_offset != 0)) + & (exact_distance >= self.min_distance) + & (exact_distance <= upper) + ) + return row_offset[in_range], column_offset[in_range] + + def _sample_anchors(self, count: int, *, chunk_aligned: bool) -> NDArrayNum: + """ + Draw first endpoints across the grid or from a small set of chunks. + + :param count: Number of first endpoints to draw, allowing repeated cells. + :param chunk_aligned: Whether to draw only from chunks_per_round randomly selected chunks. + :returns: Flat raster cell indexes that can be reused as first endpoints. + """ + + # Draw directly from the full grid when pairs do not need to stay near selected chunks + if not chunk_aligned: + return self.rng.integers(0, self.size, count, dtype=np.int64) + + # Select only a few source chunks before drawing cell numbers + n_chunk_rows, n_chunk_columns = (len(edges) - 1 for edges in self.chunk_edges) + chunk_count = min(self.chunks_per_round, n_chunk_rows * n_chunk_columns) + chosen = self.rng.choice(n_chunk_rows * n_chunk_columns, chunk_count, replace=False) + + # Split first endpoints between those chunks to limit Dask reads + anchors: list[NDArrayNum] = [] + per_chunk = int(np.ceil(count / chunk_count)) + remaining = count + for flat_chunk in chosen: + # Use each edge chunk's true size so every drawn cell exists + chunk_row, chunk_column = divmod(int(flat_chunk), n_chunk_columns) + row_start, row_stop = self.chunk_edges[0][chunk_row : chunk_row + 2] + column_start, column_stop = self.chunk_edges[1][chunk_column : chunk_column + 2] + take = min(per_chunk, remaining) + rows = self.rng.integers(row_start, row_stop, take, dtype=np.int64) + columns = self.rng.integers(column_start, column_stop, take, dtype=np.int64) + anchors.append(rows * self.shape[1] + columns) + remaining -= take + if remaining == 0: + break + + # Join the first endpoints after all selected chunks have contributed + return np.concatenate(anchors) if anchors else np.empty(0, dtype=np.int64) + + def _same_chunk( + self, rows: NDArrayNum, columns: NDArrayNum, target_rows: NDArrayNum, target_columns: NDArrayNum + ) -> NDArrayNum: + """ + Check whether each pair stays within its actual Dask chunk or eager sampling area. + + :param rows: Raster row indexes of the first endpoints. + :param columns: Raster column indexes of the first endpoints. + :param target_rows: Matching row indexes of the second endpoints. + :param target_columns: Matching column indexes of the second endpoints. + :returns: One boolean per pair, true when both endpoints belong to the same chunk. + """ + + # Locate both endpoints against the real boundaries, including uneven interior chunks + row_edges, column_edges = self.chunk_edges + same_row = np.searchsorted(row_edges, rows, side="right") == np.searchsorted( + row_edges, target_rows, side="right" + ) + same_column = np.searchsorted(column_edges, columns, side="right") == np.searchsorted( + column_edges, target_columns, side="right" + ) + return same_row & same_column + + def _from_anchors(self, anchors: NDArrayNum, count: int, *, local: bool) -> tuple[NDArrayNum, NDArrayNum]: + """ + Reuse first endpoints and draw a separate offset for each pair. + + The candidate count is described by _candidates(). + + :param anchors: Flat raster cell indexes to reuse as first endpoints. + :param local: Whether to keep both endpoints in the same chunk and apply max_local_distance. + :returns: First and second endpoint indexes after checking grid boundaries and configured duplicate handling. + """ + + # Return empty integer arrays before NumPy can try to repeat an empty input + if anchors.size == 0 or count == 0: + return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) + + # Repeat first endpoints and draw an offset for every requested pair + repeated = np.resize(anchors, count) + rows, columns = np.divmod(repeated, self.shape[1]) + row_offset, column_offset = self._offsets(count, self.max_local_distance if local else self.max_distance) + length = min(count, row_offset.size) + repeated, rows, columns = repeated[:length], rows[:length], columns[:length] + target_rows, target_columns = rows + row_offset[:length], columns + column_offset[:length] + + # Remove offsets that leave the raster before converting endpoints to flat cell numbers + inside = ( + (target_rows >= 0) + & (target_rows < self.shape[0]) + & (target_columns >= 0) + & (target_columns < self.shape[1]) + ) + if local: + # Keep nearby pairs in the first endpoint's chunk to limit Dask reads + inside &= self._same_chunk(rows, columns, target_rows, target_columns) + first = repeated[inside] + second = target_rows[inside] * self.shape[1] + target_columns[inside] + + # Optionally keep each second endpoint only once for a given first endpoint + if self.deduplicate == "per_anchor" and first.size: + order = np.argsort(first, kind="stable") + keys = first[order].astype(np.int64) * np.int64(self.size) + second[order] + _, keep = np.unique(keys, return_index=True) + positions = order[np.sort(keep)] + first, second = first[positions], second[positions] + return first, second + + def _anchor_batched(self, anchors: NDArrayNum, *, local: bool) -> tuple[NDArrayNum, NDArrayNum]: + """ + Draw several distances and directions from every first endpoint. + + Anchors and local distance constraints follow _from_anchors(); the numbers of distances and directions + are configured by _RegularPairSampler.__init__(). + """ + + # Return empty integer arrays when no first endpoint was supplied + if anchors.size == 0: + return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) + + # Draw several directions for each distance so one first endpoint yields many pairs + upper = min(self.max_distance, self.max_local_distance) if local else self.max_distance + radii = np.exp( + self.rng.uniform(np.log(self.min_distance), np.log(upper), (anchors.size, self.distances_per_anchor)) + ) + angles = self.rng.uniform( + 0, + 2 * np.pi, + (anchors.size, self.distances_per_anchor, self.angles_per_distance), + ) + column_offset = np.rint(radii[:, :, None] * np.cos(angles) / self.dx).astype(np.int64).ravel() + row_offset = np.rint(radii[:, :, None] * np.sin(angles) / self.dy).astype(np.int64).ravel() + repeated = np.repeat(anchors, self.distances_per_anchor * self.angles_per_distance) + + # Remove zero and out-of-range offsets after rounding them to grid cells + exact_distance = np.hypot(column_offset * self.dx, row_offset * self.dy) + in_range = ( + ((row_offset != 0) | (column_offset != 0)) + & (exact_distance >= self.min_distance) + & (exact_distance <= upper) + ) + repeated, row_offset, column_offset = repeated[in_range], row_offset[in_range], column_offset[in_range] + rows, columns = np.divmod(repeated, self.shape[1]) + target_rows, target_columns = rows + row_offset, columns + column_offset + + # Exclude second endpoints outside the raster or, for nearby pairs, outside the selected chunk + inside = ( + (target_rows >= 0) + & (target_rows < self.shape[0]) + & (target_columns >= 0) + & (target_columns < self.shape[1]) + ) + if local: + inside &= self._same_chunk(rows, columns, target_rows, target_columns) + first = repeated[inside] + second = target_rows[inside] * self.shape[1] + target_columns[inside] + + # Optionally remove repeated second endpoints for each first endpoint + if self.deduplicate == "per_anchor" and first.size: + order = np.argsort(first, kind="stable") + keys = first[order].astype(np.int64) * np.int64(self.size) + second[order] + _, keep = np.unique(keys, return_index=True) + positions = order[np.sort(keep)] + first, second = first[positions], second[positions] + return first, second + + def _independent(self, count: int) -> tuple[NDArrayNum, NDArrayNum]: + """ + Draw each first endpoint and offset independently across the full grid. + + The candidate count and returned endpoint arrays are described by _candidates(). + """ + + # Draw every first endpoint independently + row_offset, column_offset = self._offsets(count, self.max_distance) + rows = self.rng.integers(0, self.shape[0], row_offset.size, dtype=np.int64) + columns = self.rng.integers(0, self.shape[1], column_offset.size, dtype=np.int64) + target_rows, target_columns = rows + row_offset, columns + column_offset + + # Keep only second endpoints that remain inside the raster after applying offsets + inside = ( + (target_rows >= 0) + & (target_rows < self.shape[0]) + & (target_columns >= 0) + & (target_columns < self.shape[1]) + ) + return rows[inside] * self.shape[1] + columns[inside], ( + target_rows[inside] * self.shape[1] + target_columns[inside] + ) + + #################### + # STRATEGY CHOICE + #################### + + def _candidates(self, count: int) -> tuple[NDArrayNum, NDArrayNum]: + """ + Draw one limited batch of possible pairs with the selected strategy. + + Sampling options are configured by _RegularPairSampler.__init__(). + + :param count: Maximum number of candidate pairs to return before checking endpoint values. + :returns: Matching flat cell indexes for the first and second endpoints, possibly fewer than count. + """ + + # Split the batch between nearby and full-range pairs in the requested proportion + local_count = int(round(count * self.hybrid_local_fraction)) + global_count = count - local_count + first_parts: list[NDArrayNum] = [] + second_parts: list[NDArrayNum] = [] + + # Draw full-range pairs independently and reuse first endpoints only for nearby pairs + if self.strategy == "independent": + if local_count: + anchors = self._sample_anchors(min(self.anchors_per_round, local_count), chunk_aligned=True) + first, second = self._from_anchors(anchors, local_count, local=True) + first_parts.append(first) + second_parts.append(second) + if global_count: + first, second = self._independent(global_count) + first_parts.append(first) + second_parts.append(second) + + # Reuse one set of first endpoints for both nearby and full-range pairs + elif self.strategy in {"anchors", "chunk_anchors"}: + anchors = self._sample_anchors( + min(self.anchors_per_round, count), chunk_aligned=self.strategy == "chunk_anchors" + ) + for part_count, local in ((local_count, True), (global_count, False)): + if part_count: + first, second = self._from_anchors(anchors, part_count, local=local) + first_parts.append(first) + second_parts.append(second) + + # Draw several distances and directions from each first endpoint + else: + pairs_per_anchor = self.distances_per_anchor * self.angles_per_distance + for part_count, local in ((local_count, True), (global_count, False)): + if part_count: + anchors = self._sample_anchors(int(np.ceil(part_count / pairs_per_anchor)), chunk_aligned=local) + first, second = self._anchor_batched(anchors, local=local) + first_parts.append(first[:part_count]) + second_parts.append(second[:part_count]) + + # Return empty integer arrays when neither part requested a pair + if not first_parts: + return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) + return np.concatenate(first_parts)[:count], np.concatenate(second_parts)[:count] + + ################# + # PAIR COLLECTION + ################# + + def sample(self) -> tuple[NDArrayNum, NDArrayNum, NDArrayNum]: + """ + Collect the requested number of pairs whose two raster values are available. + + Sampling options are configured by _RegularPairSampler.__init__(). + + :returns: Three one-dimensional arrays containing first endpoint indexes, second endpoint indexes, + and their distances in raster coordinate units. Indexes refer to flat raster cells in row order. + """ + + # Count available cells without loading a complete Dask mask + if is_dask_array(self.array): + dask_array = __import__("dask.array", fromlist=["array"]) + n_valid = int(dask_array.count_nonzero(dask_array.isfinite(self.array)).compute()) + else: + n_valid = int(np.count_nonzero(np.isfinite(self.array))) + + # Limit a unique sample to the number of different pairs that can exist + maximum_unique = n_valid * (n_valid - 1) // 2 + target = min(self.n_pairs, maximum_unique) + if target < self.n_pairs: + warnings.warn( + f"Argument ``n_pairs`` exceeds the {maximum_unique} possible finite pairs; using that maximum.", + UserWarning, + ) + if target == 0: + raise ValueError("At least two finite raster cells are required to sample pairs.") + + # Bound temporary arrays by both the batch limit and the allowed oversampling of the target + count = min(self.batch_pairs, max(1, int(np.ceil(target * self.max_oversample)))) + first_parts: list[NDArrayNum] = [] + second_parts: list[NDArrayNum] = [] + remaining, stalled = target, 0 + for _ in range(self.max_rounds): + if remaining == 0: + break + + # Draw another bounded batch when missing values or edge crossings leave too few pairs + first, second = self._candidates(count) + + # Read only proposed endpoints, then keep pairs where both values are available + if first.size: + first_values, second_values = _read_raster_pair_values(self.array, first, second) + finite = np.isfinite(first_values) & np.isfinite(second_values) + first, second = first[finite], second[finite] + + # Keep only the remaining number of pairs and detect rounds that find nothing + if first.size: + take = min(remaining, first.size) + first_parts.append(first[:take]) + second_parts.append(second[:take]) + remaining -= take + stalled = 0 + else: + stalled += 1 + if stalled >= 5: + break + + # Fail clearly when no round found one pair with two available values + if not first_parts: + raise ValueError("No finite raster pairs could be sampled.") + + # Remove duplicates across rounds when the caller requests global uniqueness + first, second = np.concatenate(first_parts), np.concatenate(second_parts) + if self.deduplicate == "global": + first, second = _deduplicate_pairs(first, second, n_observations=self.size) + if first.size < target: + warnings.warn( + f"Sampled {first.size} finite pairs out of {target} requested after {self.max_rounds} rounds.", + UserWarning, + ) + + # Use the requested integer type and calculate exact map distances for the result + first, second = first.astype(self.index_dtype, copy=False), second.astype(self.index_dtype, copy=False) + first_rows, first_columns = np.divmod(first.astype(np.int64), self.shape[1]) + second_rows, second_columns = np.divmod(second.astype(np.int64), self.shape[1]) + distances = np.hypot( + (second_columns - first_columns) * self.dx, + (second_rows - first_rows) * self.dy, + ).astype(self.distance_dtype, copy=False) + return first, second, distances + + +############################## +# 3/ IRREGULAR POINT SAMPLING +############################## + + +@dataclass(frozen=True) +class _GridSpec: + """Grid cells used to find nearby point rows without comparing every point.""" + + cell_size: float + x_min: float + y_min: float + n_columns: int + n_rows: int + + +class _IrregularPairSampler: + """ + Irregular 2D coordinate sampler returning endpoint indices and distances. + + Unlike regular grid sampling, irregular coordinates have no fixed row and column offsets, which makes it less + computationally efficient yet simplifies the approach a lot. The exact methods below ("kdtree" and "hashgrid") both + choose logarithmically spaced distance rings and sample observed point pairs within them. The approximate + "nn_logvector" method instead draws a log-uniform distance and random direction, then uses the observed point + nearest the proposed endpoint. + + Strategies: + - "kdtree" : exact annulus sampling via KDTree query_ball_point(r_out) + annulus filter + - "hashgrid" : exact annulus sampling via hash-grid + AABB culling + annulus filter + - "nn_logvector": approximate log-distance + random angle with vectorized KDTree NN queries + (no distance-bias correction; accepts that long distances may appear more often) + """ + + ################# + # CONFIGURATION + ################# + + def __init__( + self, + coordinates: NDArrayNum, + *, + # Target sample size and distance range + n_pairs: int, + min_distance: float, + max_distance: float, + n_bins: int, + # Search strategy + strategy: Literal["kdtree", "hashgrid", "nn_logvector"], + # Exact annulus controls + anchors_per_round: int, + attempts_per_anchor: int, + max_rounds: int, + # Hash-grid tuning + cell_size: float | None, + # nn_logvector tuning (vectorized) + nn_tolerance: float, + nn_batch_size: int, + nn_oversample: float, + nn_max_batches: int, + # Random seed + random_state: int | np.random.Generator | None, + # Dtypes to optimize memory usage + index_dtype: Any, + distance_dtype: Any, + ) -> None: + """ + Pair sampling on irregular point coordinates. + + :param coordinates: Finite X/Y coordinates with shape (n_points, 2). + :param n_pairs: Target number of point pairs. + :param min_distance: Smallest allowed pair distance. + :param max_distance: Largest allowed pair distance. + :param n_bins: Number of logarithmically spaced distance rings used by the exact strategies. + :param strategy: Pair search strategy: "kdtree", "hashgrid", or "nn_logvector". See the class docstring for + details. + :param anchors_per_round: Number of first points tested in each round by the exact strategies. + :param attempts_per_anchor: Number of distance rings tested for each first point by the exact strategies. + :param max_rounds: Maximum number of sampling rounds used by the exact strategies. + :param cell_size: Square cell width used by "hashgrid". Defaults to one eighth of ``max_distance``. + :param nn_tolerance: Largest nearest-point error accepted by "nn_logvector", as a fraction of the proposed + distance. + :param nn_batch_size: Maximum number of proposed endpoints checked together by "nn_logvector". + :param nn_oversample: Number of endpoints proposed by "nn_logvector" relative to the pairs still needed. + :param nn_max_batches: Maximum number of proposal batches used by "nn_logvector". + :param random_state: Seed or NumPy Generator used for reproducible random sampling. + :param index_dtype: Integer dtype for returned point indices. + :param distance_dtype: Floating dtype for returned distances. + """ + + # Store coordinates and sampling options with consistent numeric types + self.coordinates = np.asarray(coordinates, dtype=np.float64) + self.size = len(self.coordinates) + self.n_pairs = int(n_pairs) + self.min_distance, self.max_distance = float(min_distance), float(max_distance) + self.n_bins, self.strategy = int(n_bins), strategy + self.anchors_per_round, self.attempts_per_anchor = int(anchors_per_round), int(attempts_per_anchor) + self.max_rounds = int(max_rounds) + self.cell_size = self.max_distance / 8 if cell_size is None else float(cell_size) + self.nn_tolerance, self.nn_batch_size = float(nn_tolerance), int(nn_batch_size) + self.nn_oversample, self.nn_max_batches = float(nn_oversample), int(nn_max_batches) + self.rng = ( + random_state if isinstance(random_state, np.random.Generator) else np.random.default_rng(random_state) + ) + self.index_dtype, self.distance_dtype = np.dtype(index_dtype), np.dtype(distance_dtype) + + # Prepare distance rings now and build search helpers only if the chosen strategy needs them + self.edges = np.geomspace(self.min_distance, self.max_distance, self.n_bins + 1) + self._tree: cKDTree | None = None + self.grid: dict[tuple[int, int], NDArrayNum] | None = None + self.grid_spec: _GridSpec | None = None + + # Check coordinates and batch options before starting repeated searches + if self.coordinates.ndim != 2 or self.coordinates.shape[1] != 2 or self.size < 2: + raise ValueError("Argument ``coordinates`` must contain at least two X/Y points.") + if self.n_pairs < 1 or self.n_bins < 1: + raise ValueError("Arguments ``n_pairs`` and ``n_bins`` must be positive integers.") + if not 0 < self.min_distance < self.max_distance: + raise ValueError("Require 0 < ``min_distance`` < ``max_distance``.") + if strategy not in {"kdtree", "hashgrid", "nn_logvector"}: + raise ValueError("Unknown irregular point pair sampling ``strategy``.") + if self.anchors_per_round < 1 or self.attempts_per_anchor < 1 or self.max_rounds < 1: + raise ValueError("Exact search batch and round controls must be positive integers.") + if self.cell_size <= 0 or self.nn_tolerance <= 0: + raise ValueError("Arguments ``cell_size`` and ``nn_tolerance`` must be strictly positive.") + if self.nn_batch_size < 1 or self.nn_oversample <= 0 or self.nn_max_batches < 1: + raise ValueError("Nearest-neighbor batch controls must be strictly positive.") + + ######################## + # NEARBY POINT SEARCH + ######################## + + def _build_grid(self) -> None: + """Group point rows into grid cells used for nearby searches.""" + + # Convert coordinates to integer cells relative to the point cloud origin + x, y = self.coordinates.T + x_min, y_min = float(x.min()), float(y.min()) + columns = np.floor((x - x_min) / self.cell_size).astype(np.int32) + rows = np.floor((y - y_min) / self.cell_size).astype(np.int32) + n_columns, n_rows = int(columns.max()) + 1, int(rows.max()) + 1 + keys = columns.astype(np.int64) * np.int64(n_rows) + rows + + # Sort once so all point rows from one occupied cell sit together + order = np.argsort(keys, kind="stable") + boundaries = np.r_[0, np.flatnonzero(np.diff(keys[order])) + 1, self.size] + self.grid = {} + for start, stop in zip(boundaries[:-1], boundaries[1:]): + column, row = divmod(int(keys[order[start]]), n_rows) + self.grid[(column, row)] = order[start:stop].astype(np.int32, copy=False) + + # Keep the grid details needed to find cells around a distance ring + self.grid_spec = _GridSpec(self.cell_size, x_min, y_min, n_columns, n_rows) + + @property + def tree(self) -> cKDTree: + """Build SciPy's nearby point search tree only when a strategy needs it.""" + + # Avoid building the tree for grid-based strategies that never use it + if self._tree is None: + self._tree = cKDTree(self.coordinates) + return self._tree + + def _hash_candidates(self, anchor: int, inner: float, outer: float) -> NDArrayNum: + """ + Collect point rows from grid cells that may cross one distance ring. + + The anchor and distance boundaries are described by _one_in_ring(). This preliminary search can include + points outside the ring; _one_in_ring() checks their exact distances before selecting an endpoint. + """ + + # Build the search grid on first use so other strategies do not store it + if self.grid is None or self.grid_spec is None: + self._build_grid() + assert self.grid is not None and self.grid_spec is not None + x, y = self.coordinates[anchor] + center_column = int(np.floor((x - self.grid_spec.x_min) / self.cell_size)) + center_row = int(np.floor((y - self.grid_spec.y_min) / self.cell_size)) + radius = int(np.ceil(outer / self.cell_size)) + + # Visit only grid cells inside the square around the outer distance + parts: list[NDArrayNum] = [] + for column in range(max(0, center_column - radius), min(self.grid_spec.n_columns, center_column + radius + 1)): + for row in range(max(0, center_row - radius), min(self.grid_spec.n_rows, center_row + radius + 1)): + # Skip cells that cannot touch the requested distance ring + cell_x_min = self.grid_spec.x_min + column * self.cell_size + cell_y_min = self.grid_spec.y_min + row * self.cell_size + cell_x_max, cell_y_max = cell_x_min + self.cell_size, cell_y_min + self.cell_size + nearest_x = max(cell_x_min - x, 0.0, x - cell_x_max) + nearest_y = max(cell_y_min - y, 0.0, y - cell_y_max) + farthest_x = max(abs(x - cell_x_min), abs(x - cell_x_max)) + farthest_y = max(abs(y - cell_y_min), abs(y - cell_y_max)) + if nearest_x**2 + nearest_y**2 > outer**2 or farthest_x**2 + farthest_y**2 < inner**2: + continue + indexes = self.grid.get((column, row)) + if indexes is not None: + parts.append(indexes) + + # Join point rows from matching cells before checking their exact distances + return np.concatenate(parts) if parts else np.empty(0, dtype=np.int32) + + def _one_in_ring(self, anchor: int, inner: float, outer: float) -> tuple[int, float] | None: + """ + Select one second point within an exact distance range of the first. + + :param anchor: Row index of the first endpoint in the sampler's finite coordinate array. + :param inner: Inclusive lower distance boundary in coordinate units. + :param outer: Exclusive upper distance boundary, except that max_distance is included. + :returns: The selected second endpoint's row index and exact distance, or None when no point lies in the ring. + """ + + # Ask the selected search helper for points inside the outer distance + candidates = ( + np.asarray(self.tree.query_ball_point(self.coordinates[anchor], outer), dtype=np.int64) + if self.strategy == "kdtree" + else self._hash_candidates(anchor, inner, outer).astype(np.int64, copy=False) + ) + if candidates.size == 0: + return None + + # Check exact distances because both helpers may also return points outside the inner distance + differences = self.coordinates[candidates] - self.coordinates[anchor] + squared = np.sum(differences**2, axis=1) + # Include the final upper boundary so pairs exactly at max_distance remain available + below_outer = squared <= outer**2 if outer == self.max_distance else squared < outer**2 + in_ring = (squared >= inner**2) & below_outer & (candidates != anchor) + if not np.any(in_ring): + return None + + # Choose one matching endpoint at random so stored row order does not bias the sample + chosen = int(self.rng.choice(candidates[in_ring])) + return chosen, float(np.linalg.norm(self.coordinates[chosen] - self.coordinates[anchor])) + + ################### + # OFFSET MATCHING + ################### + + def _nearest_vector(self) -> tuple[NDArrayNum, NDArrayNum, NDArrayNum]: + """Draw offsets across log-spaced distances and match them to nearby points.""" + + # Collect limited batches until the requested number of pairs is reached + first_parts: list[NDArrayNum] = [] + second_parts: list[NDArrayNum] = [] + distance_parts: list[NDArrayNum] = [] + remaining = self.n_pairs + for _ in range(self.nn_max_batches): + if remaining == 0: + break + + # Draw random first points, log-spaced distances, and random directions + count = min(self.nn_batch_size, int(np.ceil(self.nn_oversample * remaining))) + anchors = self.rng.integers(0, self.size, count, dtype=np.int64) + radii = np.exp(self.rng.uniform(np.log(self.min_distance), np.log(self.max_distance), count)) + angles = self.rng.uniform(0, 2 * np.pi, count) + proposals = self.coordinates[anchors] + np.column_stack((radii * np.cos(angles), radii * np.sin(angles))) + + # Match each proposed endpoint to the nearest observed point within the allowed error + proposal_distance, neighbors = self.tree.query(proposals, k=1) + accepted = (neighbors != anchors) & (proposal_distance <= self.nn_tolerance * radii) + first, second = anchors[accepted], neighbors[accepted].astype(np.int64, copy=False) + distance = np.linalg.norm(self.coordinates[second] - self.coordinates[first], axis=1) + in_range = (distance >= self.min_distance) & (distance <= self.max_distance) + + # Check exact distances and keep only the number of pairs still needed + take = min(remaining, int(np.count_nonzero(in_range))) + if take: + first_parts.append(first[in_range][:take]) + second_parts.append(second[in_range][:take]) + distance_parts.append(distance[in_range][:take]) + remaining -= take + + # Distinguish finding no pair from finding fewer pairs than requested + if not first_parts: + raise ValueError("No point pairs could be sampled within the requested distances.") + return np.concatenate(first_parts), np.concatenate(second_parts), np.concatenate(distance_parts) + + ################# + # PAIR COLLECTION + ################# + + def sample(self) -> tuple[NDArrayNum, NDArrayNum, NDArrayNum]: + """ + Collect point pairs with the selected search strategy. + + Sampling options are configured by _IrregularPairSampler.__init__(). + + :returns: Three one-dimensional arrays containing first endpoint indexes, second endpoint indexes, + and distances in coordinate units. Indexes refer to rows in the sampler's finite coordinate array. + """ + + # Use proposed offsets directly for the nearest-point strategy + if self.strategy == "nn_logvector": + first, second, distances = self._nearest_vector() + else: + # Search one first point and distance range at a time for the exact strategies + first_values: list[int] = [] + second_values: list[int] = [] + distance_values: list[float] = [] + for _ in range(self.max_rounds): + if len(first_values) >= self.n_pairs: + break + + # Reuse first points when one round requests more pairs than there are points + replace = self.anchors_per_round > self.size + anchors = self.rng.choice(self.size, self.anchors_per_round, replace=replace) + for anchor in anchors: + for _ in range(self.attempts_per_anchor): + # Choose each log-spaced distance range equally often + bin_index = int(self.rng.integers(0, self.n_bins)) + selected = self._one_in_ring(int(anchor), self.edges[bin_index], self.edges[bin_index + 1]) + if selected is not None: + selected_index, distance = selected + first_values.append(int(anchor)) + second_values.append(selected_index) + distance_values.append(distance) + if len(first_values) == self.n_pairs: + break + if len(first_values) == self.n_pairs: + break + + # Convert the collected Python lists to numeric arrays + if not first_values: + raise ValueError("No point pairs could be sampled within the requested distances.") + first = np.asarray(first_values) + second = np.asarray(second_values) + distances = np.asarray(distance_values) + + # Warn when fewer pairs were found, but return every successful pair + if first.size < self.n_pairs: + warnings.warn(f"Sampled {first.size} point pairs out of {self.n_pairs} requested.", UserWarning) + return ( + first.astype(self.index_dtype, copy=False), + second.astype(self.index_dtype, copy=False), + distances.astype(self.distance_dtype, copy=False), + ) + + +################################## +# 4/ OBJECT METHOD IMPLEMENTATIONS +################################## + + +def _random_raster_pairs( + array: Any, + *, + dx: float, + dy: float, + n_pairs: int, + min_distance: float, + max_distance: float, + random_state: int | np.random.Generator | None, + max_rounds: int, + batch_pairs: int, +) -> tuple[NDArrayNum, NDArrayNum, NDArrayNum]: + """ + Draw independent raster endpoints and keep pairs in the requested distance range. + + Array layout, pixel spacing, and sampling controls are described by _RegularPairSampler.__init__(). + The three returned arrays follow _RegularPairSampler.sample(). + """ + + # Start empty result arrays that each round extends after removing duplicates + rng = random_state if isinstance(random_state, np.random.Generator) else np.random.default_rng(random_state) + size, n_columns = int(np.prod(array.shape)), int(array.shape[1]) + first: NDArrayNum = np.empty(0, dtype=np.int64) + second: NDArrayNum = np.empty(0, dtype=np.int64) + for _ in range(max_rounds): + remaining = n_pairs - first.size + if remaining <= 0: + break + + # Draw a limited batch of independent endpoints to fill the remaining sample + count = min(batch_pairs, max(10_000, remaining * 3)) + first_candidate = rng.integers(0, size, count, dtype=np.int64) + second_candidate = rng.integers(0, size, count, dtype=np.int64) + first_rows, first_columns = np.divmod(first_candidate, n_columns) + second_rows, second_columns = np.divmod(second_candidate, n_columns) + distances = np.hypot((second_columns - first_columns) * dx, (second_rows - first_rows) * dy) + + # Remove self-pairs, wrong distances, and pairs with a missing value + keep = (first_candidate != second_candidate) & (distances >= min_distance) & (distances <= max_distance) + # Read only pairs in range, sharing one Dask computation between both endpoints + first_candidate, second_candidate = first_candidate[keep], second_candidate[keep] + first_values, second_values = _read_raster_pair_values(array, first_candidate, second_candidate) + keep = np.isfinite(first_values) & np.isfinite(second_values) + if np.any(keep): + # Remove duplicates across all rounds before keeping the requested count + first, second = _deduplicate_pairs( + np.concatenate((first, first_candidate[keep])), + np.concatenate((second, second_candidate[keep])), + n_observations=size, + ) + first, second = first[:n_pairs], second[:n_pairs] + + # Distinguish finding no pair from finding fewer unique pairs than requested + if first.size == 0: + raise ValueError("No finite raster pairs could be sampled.") + if first.size < n_pairs: + warnings.warn(f"Sampled {first.size} unique raster pairs out of {n_pairs} requested.", UserWarning) + + # Calculate exact map distances after the final pair order is known + first_rows, first_columns = np.divmod(first, n_columns) + second_rows, second_columns = np.divmod(second, n_columns) + distances = np.hypot((second_columns - first_columns) * dx, (second_rows - first_rows) * dy) + return first, second, distances + + +def _pair_dataset( + *, + first: NDArrayNum, + second: NDArrayNum, + pair_values: NDArrayNum, + distances: NDArrayNum, + pair_coordinates: dict[str, NDArrayNum], + attrs: dict[str, Any], +) -> xr.Dataset: + """ + Build the Xarray Dataset returned by raster and point cloud pairsample() methods. + + :param first: One-dimensional original cell or row indexes for the first endpoint of each pair. + Raster cell indexes follow row order across the grid. + :param second: Matching original indexes for the second endpoints. + :param pair_values: Endpoint values with shape (n_pairs, 2), first endpoint before second. + :param distances: One spatial distance per pair, in the source coordinate units. + :param pair_coordinates: Named coordinate arrays with shape (n_pairs, 2), using the same endpoint order. + :param attrs: Sampling settings and source metadata to attach to the dataset. + :returns: Dataset with pair and endpoint dimensions, containing indexes, values, distances, and coordinates. + """ + + # Store the two endpoints along one labelled dimension shared by row numbers and values + indexes = np.column_stack((first, second)) + data_vars: dict[str, Any] = { + "index": (("pair", "endpoint"), indexes), + "value": (("pair", "endpoint"), pair_values), + "distance": ("pair", distances), + } + + # Add raster or point coordinates while keeping the same core dataset layout + for name, coordinate in pair_coordinates.items(): + data_vars[name] = (("pair", "endpoint"), coordinate) + + # Name the first and second endpoints so later code need not guess from array positions + return xr.Dataset( + data_vars=data_vars, + coords={"pair": np.arange(len(first)), "endpoint": ["first", "second"]}, + attrs=attrs, + ) + + +def _sample_raster_pairs( + raster: RasterBase, + *, + band: int, + n_pairs: int, + sampling: Literal["loglag", "random_xy"], + min_distance: float | None, + max_distance: float | None, + random_state: int | np.random.Generator | None, + mask: RasterLike | VectorLike | ArrayLike | None, + strategy: Literal["independent", "anchors", "chunk_anchors", "anchor_batched"], + deduplicate: Literal["none", "per_anchor", "global"], + batch_pairs: int, + max_rounds: int, + max_oversample: float, + chunks_per_round: int, + anchors_per_round: int, + distances_per_anchor: int, + angles_per_distance: int, + hybrid_local_fraction: float, + max_local_distance: float | None, + index_dtype: Any, + distance_dtype: Any, +) -> xr.Dataset: + """ + Check raster pairsample() inputs, draw pairs, and build its Xarray result. + + _RegularPairSampler.sample() draws log-spaced pairs, while _random_raster_pairs() draws independent + endpoint pairs. Then _pair_dataset() produces the Xarray labelled layout containing the pairs. + + Strategy, duplicate, oversampling, anchor, and local distance controls apply to ``"loglag"``. + Both sampling schemes use ``batch_pairs`` and ``max_rounds``. Dask raster values are read in chunks. + + :param raster: Raster to sample. + :param band: Band to sample, with start index of 1. + :param n_pairs: Requested number of pairs with two finite values, fewer may be returned if sampling stops early. + :param sampling: Pairwise sampling method, ``"loglag"`` balances short and long distances on a log scale, while + ``"random_xy"`` draws endpoints uniformly. + :param min_distance: Smallest distance in CRS units (e.g. meters). Defaults to the smaller pixel spacing. + :param max_distance: Largest distance in CRS units. Defaults to the diagonal between outermost cell centers. + :param random_state: Seed for reproducible sampling (e.g. 42). + :param mask: Eligible cells: True in a mask array or aligned mask raster, or inside vector geometries. + :param strategy: GeoUtils log-lag strategy: ``"independent"`` draws each pair separately, ``"anchors"`` + reuses first endpoints, ``"chunk_anchors"`` also limits source chunks, and ``"anchor_batched"`` draws + several distances and directions from each first endpoint. + :param deduplicate: ``"none"`` keeps repeats, ``"per_anchor"`` removes repeated targets within each anchor + batch, and ``"global"`` removes repeated pairs across all batches. ``"random_xy"`` always removes repeats. + :param batch_pairs: Maximum candidate pairs per batch; smaller batches use less temporary memory. + :param max_rounds: Maximum attempts to fill the sample after rejecting missing values or out-of-range pairs. + :param max_oversample: Maximum candidate count as a multiple of the target pair count (e.g. 8). + :param chunks_per_round: Maximum source chunks used when drawing anchors from selected chunks. + :param anchors_per_round: Maximum first endpoints reused per round by ``"anchors"``, ``"chunk_anchors"``, + or local ``"independent"`` sampling. + :param distances_per_anchor: Distances drawn per first endpoint with ``"anchor_batched"``. + :param angles_per_distance: Directions drawn per distance with ``"anchor_batched"``. + :param hybrid_local_fraction: Fraction of candidate pairs kept within the first endpoint's chunk (e.g. 0.5). + Zero samples across the full raster; one keeps all pairs local. + :param max_local_distance: Largest proposed local distance in CRS units. Defaults to the largest chunk diagonal. + :param index_dtype: Integer NumPy dtype for returned cell indexes (e.g. ``"int64"`` for very large rasters). + :param distance_dtype: Floating NumPy dtype for returned distances (e.g. ``"float32"`` to reduce memory). + :returns: Xarray Dataset with pair and endpoint dimensions, containing cell indexes, values, coordinates, + and distances. + """ + + # Select the raster band and check the requested output number types + array = _selected_raster_data(raster, band) + index_type, distance_type = np.dtype(index_dtype), np.dtype(distance_dtype) + if not np.issubdtype(index_type, np.integer) or not np.issubdtype(distance_type, np.floating): + raise TypeError("Arguments ``index_dtype`` and ``distance_dtype`` must be integer and floating, respectively.") + if int(np.prod(array.shape)) - 1 > np.iinfo(index_type).max: + raise ValueError("Argument ``index_dtype`` cannot represent every cell in this raster.") + + # Convert an array, raster, or vector mask to one boolean grid + if mask is not None: + mask_array = _mask_on_raster(mask, raster, "inside", "raise") + + # Apply the mask without loading a Dask source array + if is_dask_array(array): + dask_array = __import__("dask.array", fromlist=["array"]) + array = dask_array.where(mask_array, array, np.nan) + else: + array = np.where(mask_array, array, np.nan) + + # Choose default map distances from the cell size and raster extent + dx, dy = (float(abs(value)) for value in get_geo_attr(raster, "res")) + diagonal = float(np.hypot((array.shape[1] - 1) * dx, (array.shape[0] - 1) * dy)) + minimum = min(dx, dy) if min_distance is None else float(min_distance) + maximum = diagonal if max_distance is None else float(max_distance) + if not 0 < minimum < maximum: + raise ValueError("Require 0 < ``min_distance`` < ``max_distance``.") + + # Call the log-spaced or independent endpoint sampling workflow + if sampling == "loglag": + first, second, distances = _RegularPairSampler( + array, + dx=dx, + dy=dy, + n_pairs=n_pairs, + min_distance=minimum, + max_distance=maximum, + strategy=strategy, + deduplicate=deduplicate, + random_state=random_state, + batch_pairs=batch_pairs, + max_rounds=max_rounds, + max_oversample=max_oversample, + chunks_per_round=chunks_per_round, + anchors_per_round=anchors_per_round, + distances_per_anchor=distances_per_anchor, + angles_per_distance=angles_per_distance, + hybrid_local_fraction=hybrid_local_fraction, + max_local_distance=max_local_distance, + index_dtype=index_dtype, + distance_dtype=distance_dtype, + ).sample() + elif sampling == "random_xy": + first, second, distances = _random_raster_pairs( + array, + dx=dx, + dy=dy, + n_pairs=n_pairs, + min_distance=minimum, + max_distance=maximum, + random_state=random_state, + max_rounds=max_rounds, + batch_pairs=batch_pairs, + ) + else: + raise ValueError("Argument ``sampling`` must be 'loglag' or 'random_xy'.") + + # Use the requested number types and recover map coordinates for both endpoints + first = first.astype(index_type, copy=False) + second = second.astype(index_type, copy=False) + distances = distances.astype(distance_type, copy=False) + first_rows, first_columns = np.divmod(first.astype(np.int64), int(array.shape[1])) + second_rows, second_columns = np.divmod(second.astype(np.int64), int(array.shape[1])) + first_x, first_y = raster.ij2xy(first_rows, first_columns) + second_x, second_y = raster.ij2xy(second_rows, second_columns) + + # Load only the selected raster values when building the Xarray result + return _pair_dataset( + first=first, + second=second, + pair_values=np.column_stack(_read_raster_pair_values(array, first, second)), + distances=distances, + pair_coordinates={ + "row": np.column_stack((first_rows, second_rows)), + "column": np.column_stack((first_columns, second_columns)), + "x": np.column_stack((first_x, second_x)), + "y": np.column_stack((first_y, second_y)), + }, + attrs={ + "source": "raster", + "crs": str(get_geo_attr(raster, "crs")), + "sampling": sampling, + "strategy": strategy if sampling == "loglag" else "random_xy", + "deduplicate": deduplicate, + "requested_pairs": int(n_pairs), + "accepted_pairs": int(len(first)), + "min_distance": minimum, + "max_distance": maximum, + "band": int(band), + }, + ) + + +def _sample_point_pairs( + pointcloud: PointCloudBase, + *, + n_pairs: int, + sampling: Literal["loglag", "random_xy"], + min_distance: float | None, + max_distance: float | None, + random_state: int | np.random.Generator | None, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None, + strategy: Literal["kdtree", "hashgrid", "nn_logvector"], + n_bins: int, + anchors_per_round: int, + attempts_per_anchor: int, + max_rounds: int, + cell_size: float | None, + nn_tolerance: float, + nn_batch_size: int, + nn_oversample: float, + nn_max_batches: int, + index_dtype: Any, + distance_dtype: Any, +) -> xr.Dataset: + """ + Check point cloud pairsample() inputs, draw pairs, and build its Xarray result. + + _IrregularPairSampler.sample() searches for log-spaced pairs; the independent path draws and filters endpoints + directly. Then _pair_dataset() produces the Xarray labelled layout containing the pairs. + + Strategy controls apply to ``"loglag"``. ``"random_xy"`` uses ``max_rounds`` and ``nn_batch_size``. + Dask point tables are loaded because the search requires all coordinates. + + :param pointcloud: Point cloud to sample, using its main data column or geometry heights. + :param n_pairs: Requested number of pairs with two finite values; fewer may be returned if sampling stops early. + :param sampling: ``"loglag"`` balances short and long distances on a log scale; ``"random_xy"`` draws + endpoints uniformly. + :param min_distance: Smallest distance in CRS units (e.g. meters). Defaults to half the spacing estimated + from the eligible point density. + :param max_distance: Largest distance in CRS units. Defaults to the eligible point cloud's bounding box diagonal. + :param random_state: Seed for reproducible sampling (e.g. 42). + :param mask: Eligible points: True in a boolean array or spatial mask, or inside vector geometries. + Point masks must follow the same ordered coordinates; raster masks use nearest interpolation. + Missing mask entries are excluded. + :param strategy: GeoUtils log-lag strategy: ``"kdtree"`` uses SciPy to search distance rings, ``"hashgrid"`` + searches rings using a spatial grid, and ``"nn_logvector"`` uses SciPy to match proposed endpoints + to nearby points. + :param n_bins: Log-spaced distance rings used by ``"kdtree"`` and ``"hashgrid"`` (e.g. 24). + :param anchors_per_round: First endpoints tested per round by ``"kdtree"`` and ``"hashgrid"``. + :param attempts_per_anchor: Distance rings tried per first endpoint by ``"kdtree"`` and ``"hashgrid"``. + :param max_rounds: Maximum rounds to fill the sample with ``"kdtree"``, ``"hashgrid"``, or ``"random_xy"``. + :param cell_size: Grid cell width in CRS units for ``"hashgrid"``. Defaults to one eighth of max_distance. + :param nn_tolerance: Allowed endpoint snap distance for ``"nn_logvector"``, as a fraction of the proposed + pair distance (e.g. 0.1 allows a 10% offset). + :param nn_batch_size: Maximum candidate pairs per batch with ``"nn_logvector"`` or ``"random_xy"``; + smaller batches use less temporary memory. + :param nn_oversample: Candidate count as a multiple of the remaining pairs with ``"nn_logvector"`` (e.g. 2). + :param nn_max_batches: Maximum batches to fill the sample with ``"nn_logvector"``. + :param index_dtype: Integer NumPy dtype for returned row indexes (e.g. ``"int64"`` for very large point clouds). + :param distance_dtype: Floating NumPy dtype for returned distances (e.g. ``"float64"`` for greater precision). + :returns: Xarray Dataset with pair and endpoint dimensions, containing original row indexes, values, + coordinates, and distances. + """ + + # Load the point table because pair searches need all coordinates + dataframe = pointcloud.ds.compute() if is_dask_dataframe(pointcloud.ds) else pointcloud.ds + values = np.asarray( + dataframe[pointcloud.data_column] if pointcloud.data_column is not None else dataframe.geometry.z + ) + coordinates = np.column_stack((dataframe.geometry.x.to_numpy(), dataframe.geometry.y.to_numpy())) + + # Keep rows with available coordinates and values, then apply the optional mask + valid = np.isfinite(values) & np.all(np.isfinite(coordinates), axis=1) + if mask is not None: + # Reuse the loaded table for spatial masks and ordered-coordinate checks against point masks + support = _get_pointcloud_interface(dataframe) + mask_array = _mask_at_support(mask, support, support_dataframe=dataframe) + if mask_array is not None and is_dask_array(mask_array): + mask_array = mask_array.compute() + valid &= mask_array + + # Keep source row numbers so the result refers back to the original point table + original_indexes = np.flatnonzero(valid) + coordinates_valid, values_valid = coordinates[valid], values[valid] + if len(values_valid) < 2: + raise ValueError("At least two finite points are required to sample pairs.") + + # Check that the requested integer type can hold every original row number + index_type, distance_type = np.dtype(index_dtype), np.dtype(distance_dtype) + if not np.issubdtype(index_type, np.integer) or not np.issubdtype(distance_type, np.floating): + raise TypeError("Arguments ``index_dtype`` and ``distance_dtype`` must be integer and floating, respectively.") + if len(values) - 1 > np.iinfo(index_type).max: + raise ValueError("Argument ``index_dtype`` cannot represent every point in this point cloud.") + + # Choose default distances from the point extent and typical point spacing + bounds = np.ptp(coordinates_valid, axis=0) + diagonal = float(np.hypot(*bounds)) + density_spacing = float(np.sqrt(max(bounds[0] * bounds[1], 0) / len(values_valid))) + minimum = max(0.5 * density_spacing, float(np.finfo(float).eps)) if min_distance is None else float(min_distance) + maximum = diagonal if max_distance is None else float(max_distance) + if not 0 < minimum < maximum: + raise ValueError("Require 0 < ``min_distance`` < ``max_distance``.") + if n_pairs < 1 or max_rounds < 1: + raise ValueError("Arguments ``n_pairs`` and ``max_rounds`` must be positive integers.") + + # Use the selected point search for log-spaced distances + if sampling == "loglag": + first, second, distances = _IrregularPairSampler( + coordinates_valid, + n_pairs=n_pairs, + min_distance=minimum, + max_distance=maximum, + n_bins=n_bins, + strategy=strategy, + anchors_per_round=anchors_per_round, + attempts_per_anchor=attempts_per_anchor, + max_rounds=max_rounds, + cell_size=cell_size, + nn_tolerance=nn_tolerance, + nn_batch_size=nn_batch_size, + nn_oversample=nn_oversample, + nn_max_batches=nn_max_batches, + random_state=random_state, + index_dtype=index_dtype, + distance_dtype=distance_dtype, + ).sample() + elif sampling == "random_xy": + # Draw independent endpoints in limited rounds for uniform random sampling + rng = random_state if isinstance(random_state, np.random.Generator) else np.random.default_rng(random_state) + first = np.empty(0, dtype=np.int64) + second = np.empty(0, dtype=np.int64) + for _ in range(max_rounds): + remaining = n_pairs - first.size + if remaining <= 0: + break + + # Draw extra possible pairs before checking exact distances and removing duplicates + count = min(nn_batch_size, max(10_000, remaining * 3)) + first_candidate = rng.integers(0, len(values_valid), count, dtype=np.int64) + second_candidate = rng.integers(0, len(values_valid), count, dtype=np.int64) + candidate_distances = np.linalg.norm( + coordinates_valid[first_candidate] - coordinates_valid[second_candidate], axis=1 + ) + keep = ( + (first_candidate != second_candidate) + & (candidate_distances >= minimum) + & (candidate_distances <= maximum) + ) + if np.any(keep): + # Remove duplicates across rounds before keeping the requested number + first, second = _deduplicate_pairs( + np.concatenate((first, first_candidate[keep])), + np.concatenate((second, second_candidate[keep])), + n_observations=len(values_valid), + ) + first, second = first[:n_pairs], second[:n_pairs] + + # Return a smaller sample with a warning, but fail when no matching pair exists + if first.size == 0: + raise ValueError("No point pairs could be sampled within the requested distances.") + if first.size < n_pairs: + warnings.warn(f"Sampled {first.size} unique point pairs out of {n_pairs} requested.", UserWarning) + distances = np.linalg.norm(coordinates_valid[first] - coordinates_valid[second], axis=1) + else: + raise ValueError("Argument ``sampling`` must be 'loglag' or 'random_xy'.") + + # Map kept row numbers back to the original point table and requested number types + original_first = original_indexes[first].astype(index_type, copy=False) + original_second = original_indexes[second].astype(index_type, copy=False) + distances = distances.astype(distance_type, copy=False) + + # Build the same labelled Xarray layout used for raster pairs + return _pair_dataset( + first=original_first, + second=original_second, + pair_values=np.column_stack((values_valid[first], values_valid[second])), + distances=distances, + pair_coordinates={ + "x": np.column_stack((coordinates_valid[first, 0], coordinates_valid[second, 0])), + "y": np.column_stack((coordinates_valid[first, 1], coordinates_valid[second, 1])), + }, + attrs={ + "source": "pointcloud", + "crs": str(pointcloud.crs), + "sampling": sampling, + "strategy": strategy if sampling == "loglag" else "random_xy", + "requested_pairs": int(n_pairs), + "accepted_pairs": int(len(first)), + "min_distance": minimum, + "max_distance": maximum, + }, + ) diff --git a/geoutils/sampling/stratified.py b/geoutils/sampling/stratified.py new file mode 100644 index 000000000..0f6928959 --- /dev/null +++ b/geoutils/sampling/stratified.py @@ -0,0 +1,298 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Subsample positions independently within encoded integer strata.""" + +from __future__ import annotations + +from itertools import product +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np + +from geoutils._dispatch import is_dask_array +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.multiproc.chunked import iter_chunk_slices +from geoutils.multiproc.cluster import _map_bounded +from geoutils.sampling.subsampling import _splitmix64 + +if TYPE_CHECKING: + from geoutils.multiproc.mparray import MultiprocConfig + + +def _prune_stratified_candidates( + candidates: tuple[NDArrayNum, NDArrayNum, NDArrayNum], quotas: dict[int, int] | int | float +) -> tuple[NDArrayNum, NDArrayNum, NDArrayNum]: + """ + Keep each stratum's smallest random keys from one block or a batch of block candidates. + + Sort group IDs once so each selection inspects only its group's contiguous slice. Quotas follow + _sample_strata_block(): fractions are evaluated here only when the candidates contain the whole input; + chunked fractions require global counts computed before block selection. + + :param candidates: Three matching one-dimensional arrays of group IDs, random keys and original flat positions. + Their order need not be sorted until the final sample is returned. + :returns: The same three arrays restricted to each group's smallest keys, keeping their input dtypes. + """ + + # Place members of the same group together without scanning the full input separately for every group + labels, keys, positions = candidates + if labels.size == 0: + return candidates + order = np.argsort(labels, kind="stable") + labels, keys, positions = labels[order], keys[order], positions[order] + starts = np.r_[0, np.flatnonzero(labels[1:] != labels[:-1]) + 1] + groups = labels[starts] + counts = np.diff(np.r_[starts, labels.size]) + + # Keep whole small groups and partially select only slices that exceed their global quota + chosen = [] + for group, start, count in zip(groups, starts, counts): + # A shared fixed cap needs no separate count pass; fractions can also be resolved here for eager input + if isinstance(quotas, dict): + quota = min(quotas[int(group)], int(count)) + else: + quota = int(quotas * count) if quotas <= 1 else min(int(quotas), int(count)) + if quota == 0: + continue + if quota == count: + selected = np.arange(start, start + count) + else: + selected = start + np.argpartition(keys[start : start + count], quota - 1)[:quota] + chosen.append(selected) + + # Preserve the array dtypes even when rounding leaves every group with an empty sample + selection = np.concatenate(chosen) if chosen else np.empty(0, dtype=np.int64) + return labels[selection], keys[selection], positions[selection] + + +def _sample_strata_block( + group_ids: NDArrayNum, + origin: tuple[int, ...], + full_shape: tuple[int, ...], + quotas: dict[int, int] | int | float, + seed: int | None, + ranks: dict[int, NDArrayNum] | None, +) -> tuple[NDArrayNum, NDArrayNum, NDArrayNum]: + """ + Sample a block's strata while keeping positions relative to the full input. + + Topk keeps each group's best global-key candidates for later merging. Sequential sampling receives already + drawn member ranks relative to this block, so only those members need to leave the worker. + + :param group_ids: One NumPy block of the encoded group array passed to _stratified_subsample_indices(). + :param origin: Starting position of this block along each axis of the full input. + :param full_shape: Dimensions of the full input, used to convert block positions to original flat indices. + :param quotas: Mapping from group IDs to target sample counts across the full input, or one common cap above + one. A fraction at most one is accepted only when this block contains the whole input. + :param seed: Integer mixed with original flat indices to rank topk candidates. None selects by ranks instead. + :param ranks: For sequential sampling, a mapping from each group ID to zero-based ranks among that group's + members in this block's flat order. These are not indices into the block or the full input. None for topk. + :returns: Matching one-dimensional arrays of group IDs, random keys and flat indices into the full input. + Sequential results use zero keys because their members have already been selected. + """ + + # Ignore negative group IDs, which represent missing groupers and locations excluded by a user mask + flat = np.flatnonzero(group_ids.ravel() >= 0) + labels = group_ids.ravel()[flat] + if group_ids.shape == full_shape: + positions = flat + else: + # Convert tile-local positions only when the block covers part of the original input + coordinates = np.unravel_index(flat, group_ids.shape) + absolute = tuple(coordinate + start for coordinate, start in zip(coordinates, origin)) + positions = np.ravel_multi_index(absolute, full_shape) + + # Derive topk keys from original positions, never from a group's label or the block's local order + if seed is not None: + keys = _splitmix64(np.uint64(seed) ^ positions.astype(np.uint64)) + return _prune_stratified_candidates((labels, keys, positions), quotas) + + # Group sequential members once and translate their local ranks into original flat positions + assert ranks is not None + order = np.argsort(labels, kind="stable") + sorted_labels = labels[order] + starts = np.r_[0, np.flatnonzero(sorted_labels[1:] != sorted_labels[:-1]) + 1] if labels.size else [] + groups = sorted_labels[starts] + selected = [order[start + ranks[int(group)]] for group, start in zip(groups, starts)] + selection = np.concatenate(selected) if selected else np.empty(0, dtype=np.int64) + return labels[selection], np.zeros(selection.size, dtype=np.uint64), positions[selection] + + +def _merge_stratified_candidates( + blocks: list[tuple[NDArrayNum, NDArrayNum, NDArrayNum]], quotas: dict[int, int] | int | float +) -> tuple[NDArrayNum, NDArrayNum, NDArrayNum]: + """ + Combine a small batch of block candidates and prune each stratum back to its global quota. + + :param blocks: Candidate tuples returned by _sample_strata_block() or an earlier merge. + :param quotas: Target sample counts per group across the full input, or a common cap above one, as in + _sample_strata_block(). Fractions must already have been converted to full-group target counts. + :returns: One candidate tuple containing each group's smallest keys across the supplied blocks. + """ + + labels = np.concatenate([block[0] for block in blocks]) + keys = np.concatenate([block[1] for block in blocks]) + positions = np.concatenate([block[2] for block in blocks]) + return _prune_stratified_candidates((labels, keys, positions), quotas) + + +def _stratified_subsample_indices( + group_ids: Any, + subsample: int | float, + random_state: int | np.random.Generator | None = None, + strategy: Literal["sequential", "topk"] = "topk", + mp_config: MultiprocConfig | None = None, +) -> NDArrayNum: + """ + Choose original flat positions independently within each nonnegative integer group ID. + + Fractions select floor(fraction * group size); amounts above one cap the count per group without warnings for + smaller groups. One keeps all eligible locations. Values are not inspected, so callers can reuse the same + positions across several value columns with different missing observations. + + Count memberships first when chunked fractions or sequential ranks need group sizes, then let + _sample_strata_block() select from each block. Fixed-cap topk needs no separate count pass. + For topk, _merge_stratified_candidates() prunes bounded batches of candidates using keys from _splitmix64(). + This gives identical samples for NumPy, Dask and multiprocessing regardless of chunk boundaries. Sequential + sampling draws member ranks in group order, then maps them to blocks; its sample can depend on block layout. + + Dask builds an eight-way candidate merge tree. Multiprocessing handles up to eight tiles at a time and merges + their candidates in the caller. Both return computed NumPy indices; neither gathers the full group-ID array. + Count metadata and selected positions must fit in memory. NumPy inputs are already resident, including when + their tiles are sent to worker processes. + + :param group_ids: NumPy or Dask integer array with one encoded group per location. Nonnegative IDs identify + groups; negative IDs exclude locations. Grouping arrays and any user mask must already be combined here. + :param subsample: Positive finite fraction at most one, or maximum number of sampled locations per group above + one. Fractions round down separately in each group; one keeps all eligible locations. + :param random_state: Integer seed or NumPy Generator. Topk draws one seed from a supplied generator and uses zero + when omitted; sequential sampling uses the generator directly. A subsample of one consumes no draws. + :param strategy: Topk selects the smallest deterministic keys; sequential draws random member ranks per group. + :param mp_config: Tile sizes and worker cluster for a resident NumPy input. Omit for eager or Dask execution. + :returns: Computed one-dimensional NumPy array of flat indices into the original input, ordered by random key + for topk or by original position for sequential sampling and for a subsample of one. + """ + + # 1/ Check user input before starting any sampling tasks + if not np.issubdtype(group_ids.dtype, np.integer): + raise TypeError("Group IDs must be integers.") + if not np.isfinite(subsample) or subsample <= 0: + raise ValueError("Argument ``subsample`` must be a positive finite number.") + if strategy not in ("sequential", "topk"): + raise ValueError("Argument ``subsampling_strategy`` must be 'sequential' or 'topk'.") + use_dask = is_dask_array(group_ids) + if use_dask and mp_config is not None: + raise ValueError("Cannot use Multiprocessing and Dask simultaneously.") + if group_ids.size == 0: + return np.empty(0, dtype=np.int64) + + # Draw one seed for the whole topk call; sequential sampling uses the generator directly below + seed = None + if strategy == "topk" and subsample != 1: + if isinstance(random_state, np.random.Generator): + seed = int(random_state.integers(0, np.iinfo(np.uint32).max, dtype=np.uint32)) + else: + seed = 0 if random_state is None else int(random_state) + + # 2/ Divide the group IDs into blocks and record their original coordinate offsets + shape = tuple(int(length) for length in group_ids.shape) + if use_dask: + import_optional("dask") + import dask + + blocks = list(group_ids.to_delayed().ravel()) + starts = [np.r_[0, np.cumsum(lengths[:-1])] for lengths in group_ids.chunks] + origins = list(product(*starts)) + else: + # Reuse views of resident NumPy tiles; without multiprocessing, the input is one block + chunks = shape if mp_config is None else mp_config.chunks + tiles = list(iter_chunk_slices(shape, chunks)) + origins = [tuple(part.start for part in tile) for tile in tiles] + blocks = [group_ids[tile] for tile in tiles] + + # 3/ Count full group memberships only when fractions or sequential ranks require them in advance + quotas: dict[int, int] | int | float = subsample + if seed is None or (subsample < 1 and (use_dask or mp_config is not None)): + if use_dask: + memberships = list(dask.compute(*[dask.delayed(np.unique)(block, return_counts=True) for block in blocks])) + else: + memberships = [np.unique(block, return_counts=True) for block in blocks] + + # Round each group's fraction once, after counts from all its blocks have been combined + totals: dict[int, int] = {} + for groups, counts in memberships: + for group, count in zip(groups, counts): + if group >= 0: + totals[int(group)] = totals.get(int(group), 0) + int(count) + quotas = { + group: int(subsample * count) if subsample <= 1 else min(int(subsample), count) + for group, count in totals.items() + } + if not any(quotas.values()): + return np.empty(0, dtype=np.int64) + + # Draw sorted member ranks once for each sequential group, following group and then block order + ranks_per_block: list[dict[int, NDArrayNum] | None] = [None] * len(blocks) + if seed is None: + assert isinstance(quotas, dict) + rng = np.random.default_rng(random_state) + ranks = {} + for group in sorted(totals): + # Keep complete groups without consuming random numbers or constructing a discarded permutation + if quotas[group] == totals[group]: + ranks[group] = np.arange(totals[group]) + else: + ranks[group] = np.sort(rng.choice(totals[group], quotas[group], replace=False)) + + # Map each group's global member ranks onto its successive blocks without rescanning the input + offsets = dict.fromkeys(totals, 0) + for index, (groups, counts) in enumerate(memberships): + block_ranks: dict[int, NDArrayNum] = {} + for group, count in zip(groups, counts): + if group < 0: + continue + start = offsets[int(group)] + low, high = np.searchsorted(ranks[int(group)], [start, start + count]) + block_ranks[int(group)] = ranks[int(group)][low:high] - start + offsets[int(group)] += int(count) + ranks_per_block[index] = block_ranks + + # 4/ Select within blocks, then combine only their sampled positions or bounded topk candidates + if use_dask: + candidates = [ + dask.delayed(_sample_strata_block)(block, origin, shape, quotas, seed, ranks) + for block, origin, ranks in zip(blocks, origins, ranks_per_block) + ] + if seed is not None: + while len(candidates) > 1: + candidates = [ + dask.delayed(_merge_stratified_candidates)(candidates[start : start + 8], quotas) + for start in range(0, len(candidates), 8) + ] + results = list(dask.compute(*candidates)) + elif mp_config is not None: + # Bound pending tile payloads and intermediate candidates, including for process-based clusters + arguments = zip( + blocks, origins, [shape] * len(blocks), [quotas] * len(blocks), [seed] * len(blocks), ranks_per_block + ) + results = [] + for index, result in _map_bounded(mp_config.cluster, _sample_strata_block, arguments): + results.append(result) + if seed is not None and ((index + 1) % 8 == 0 or index + 1 == len(blocks)): + results = [_merge_stratified_candidates(results, quotas)] + else: + results = [_sample_strata_block(blocks[0], origins[0], shape, quotas, seed, ranks_per_block[0])] + + # Return a stable ordering by random key for topk and by original position for sequential or unsampled input + positions = np.concatenate([result[2] for result in results]) + if seed is None: + return np.sort(positions) + keys = np.concatenate([result[1] for result in results]) + return positions[np.lexsort((positions, keys))] diff --git a/geoutils/stats/sampling.py b/geoutils/sampling/subsampling.py similarity index 66% rename from geoutils/stats/sampling.py rename to geoutils/sampling/subsampling.py index 8c86655e8..0faa3bb1f 100644 --- a/geoutils/stats/sampling.py +++ b/geoutils/sampling/subsampling.py @@ -16,23 +16,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Module for array sampling statistics.""" +"""Module for subsampling: selecting a random subset of valid points in a N-D array.""" from __future__ import annotations import operator import warnings -from typing import TYPE_CHECKING, Any, Callable, Literal, TypedDict, overload +from typing import TYPE_CHECKING, Any, Callable, Literal, TypedDict, cast, overload import numpy as np +from geoutils._dispatch import is_dask_array from geoutils._misc import import_optional -from geoutils._typing import MArrayNum, NDArrayBool, NDArrayNum +from geoutils._typing import ArrayLike, DTypeLike, MArrayNum, NDArrayBool, NDArrayNum from geoutils.multiproc import MultiprocConfig, compute_tiling from geoutils.raster.array import get_mask_from_array if TYPE_CHECKING: - from geoutils.raster.raster import Raster, RasterBase + from geoutils.pointcloud.base import PointCloudBase + from geoutils.pointcloud.pointcloud import PointCloudLike + from geoutils.raster.base import RasterBase, RasterLike + from geoutils.vector.base import VectorLike # Import Dask as optional dependency try: @@ -77,12 +81,13 @@ def _get_subsample_size_from_user_input( npoints = min(int(subsample), total_nb_valids) if subsample > total_nb_valids: warnings.warn( - f"Subsample value of {subsample} is larger than the number of valid pixels of {total_nb_valids}," + f"Argument ``subsample`` with value {subsample} is larger than the number of valid pixels of " + f"{total_nb_valids}," f" using all valid pixels as a subsample.", category=UserWarning, ) else: - raise ValueError("Subsample must be > 0.") + raise ValueError("Argument ``subsample`` must be > 0.") return npoints @@ -141,6 +146,7 @@ def _subsample_numpy( *, random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", + mask: NDArrayBool | None = None, ) -> NDArrayNum: ... @@ -152,6 +158,7 @@ def _subsample_numpy( *, random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", + mask: NDArrayBool | None = None, ) -> tuple[NDArrayNum, ...]: ... @@ -162,6 +169,7 @@ def _subsample_numpy( *, random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", + mask: NDArrayBool | None = None, ) -> NDArrayNum | tuple[NDArrayNum, ...]: """ Subsample valid values of a 1D or 2D array. @@ -174,13 +182,16 @@ def _subsample_numpy( :param strategy: Sampling strategy: - "sequential": Random draw from valid indices (chunk-dependent, different output than chunked implementation). - "topk": Deterministic key-per-pixel draw (chunk-invariant, same output in chunked implementation). + :param mask: Prepared boolean eligibility mask with the same shape as array. :returns: The subsampled array (1D) or the indices to extract (same shape as input array). """ # Determine valid pixels and their global linear indices (row * nx + col) - mask = get_mask_from_array(array) - valids = np.flatnonzero(~mask.ravel()) # Robust 1D index list (global linear indices) + valid = ~get_mask_from_array(array).reshape(array.shape) + if mask is not None: + valid &= mask + valids = np.flatnonzero(valid.ravel()) # Robust 1D index list (global linear indices) total_nb_valids = int(valids.size) # If no valid values, early return @@ -238,7 +249,7 @@ def _subsample_numpy( return unraveled if return_indices else array[unraveled] else: - raise ValueError(f"Unknown strategy {strategy!r}. Choose 'sequential' or 'topk'.") + raise ValueError(f"Unknown ``strategy`` {strategy!r}. Choose 'sequential' or 'topk'.") ##################### @@ -295,12 +306,22 @@ def _get_indices_block_per_subsample( return relative_index_per_block +def _valid_subsample_mask(arr_chunk: NDArrayNum | NDArrayBool, mask_chunk: NDArrayBool | None = None) -> NDArrayBool: + """Find finite, unmasked values eligible for sampling, or eligible True values for boolean data.""" + + valid = ~get_mask_from_array(arr_chunk).reshape(arr_chunk.shape) + if arr_chunk.dtype == np.bool_: + valid &= np.ma.getdata(arr_chunk) + if mask_chunk is not None: + valid &= mask_chunk + return valid + + @delayed -def _delayed_nb_valids(arr_chunk: NDArrayNum | NDArrayBool) -> NDArrayNum: +def _delayed_nb_valids(arr_chunk: NDArrayNum | NDArrayBool, *, mask_chunk: NDArrayBool | None = None) -> NDArrayNum: """Count number of valid values per block.""" - if arr_chunk.dtype == np.bool_: - return np.array([np.count_nonzero(arr_chunk)]).reshape((1, 1)) - return np.array([np.count_nonzero(np.isfinite(arr_chunk))]).reshape((1, 1)) + valid = _valid_subsample_mask(arr_chunk, mask_chunk) + return np.array([np.count_nonzero(valid)]).reshape((1, 1)) @delayed @@ -312,6 +333,7 @@ def _delayed_topk_candidates_block( k: int, nx_full: int, # Width of full array return_indices_local: bool, + mask_chunk: NDArrayBool | None = None, ) -> tuple[NDArrayNum, NDArrayNum | NDArrayBool]: """ Return up to k valid samples from one block as (keys, payload). @@ -323,15 +345,15 @@ def _delayed_topk_candidates_block( * else: selected values from the array (dtype of arr_chunk, but typically float) """ + # Keep empty payloads in the same dtype as sampled values or original indices + payload_dtype: DTypeLike = np.int64 if return_indices_local else arr_chunk.dtype + # If no samples, return empty if k <= 0: - return np.empty((0,), dtype=np.uint64), np.empty((0,), dtype=np.int64) + return np.empty((0,), dtype=np.uint64), np.empty((0,), dtype=payload_dtype) # Only valid values are sampled (finite for numerical arrays, True for boolean arrays) - if np.issubdtype(arr_chunk.dtype, np.bool_): - valid = arr_chunk - else: - valid = np.isfinite(arr_chunk) + valid = _valid_subsample_mask(arr_chunk, mask_chunk) # Get nonzero indices for flattened array, and number of valid values flat = np.flatnonzero(valid.ravel()) @@ -339,7 +361,7 @@ def _delayed_topk_candidates_block( # If no valid if nvalid == 0: - return np.empty((0,), dtype=np.uint64), np.empty((0,), dtype=np.int64) + return np.empty((0,), dtype=np.uint64), np.empty((0,), dtype=payload_dtype) # Convert flat relative indices to local (row, col) within the chunk ncols = int(arr_chunk.shape[1]) @@ -426,28 +448,27 @@ def _delayed_gid_to_rc(gid: NDArrayNum, nx_full: int) -> tuple[NDArrayNum, NDArr @delayed def _delayed_subsample_block( - arr_chunk: NDArrayNum | NDArrayBool, subsample_indices: NDArrayNum + arr_chunk: NDArrayNum | NDArrayBool, subsample_indices: NDArrayNum, *, mask_chunk: NDArrayBool | None = None ) -> NDArrayNum | NDArrayBool: """Subsample the valid values at the corresponding 1D valid indices per block.""" - if arr_chunk.dtype == np.bool_: - return arr_chunk[arr_chunk][subsample_indices] - return arr_chunk[np.isfinite(arr_chunk)][subsample_indices] + valid = _valid_subsample_mask(arr_chunk, mask_chunk) + return arr_chunk[valid][subsample_indices] @delayed def _delayed_subsample_indices_block( - arr_chunk: NDArrayNum | NDArrayBool, subsample_indices: NDArrayNum, block_id: dict[str, Any] + arr_chunk: NDArrayNum | NDArrayBool, + subsample_indices: NDArrayNum, + block_id: dict[str, Any], + *, + mask_chunk: NDArrayBool | None = None, ) -> NDArrayNum: """Return 2D indices from the subsampled 1D valid indices per block.""" - if arr_chunk.dtype == np.bool_: - ix, iy = np.unravel_index(np.argwhere(arr_chunk.flatten())[subsample_indices], shape=arr_chunk.shape) - else: - # Unravel indices of valid data to the shape of the block - ix, iy = np.unravel_index( - np.argwhere(np.isfinite(arr_chunk.flatten()))[subsample_indices], shape=arr_chunk.shape - ) + # Unravel indices of valid data to the shape of the block + valid = _valid_subsample_mask(arr_chunk, mask_chunk) + ix, iy = np.unravel_index(np.argwhere(valid.flatten())[subsample_indices], shape=arr_chunk.shape) # Convert to full-array indexes by adding the row and column starting indexes for this block ix += block_id["row_start"] @@ -462,6 +483,9 @@ def _dask_subsample( return_indices: bool = False, random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", + preserve_order: bool = False, + *, + mask: da.Array | None = None, ) -> da.Array | tuple[da.Array, da.Array]: """ Subsample valid values out-of-memory from a 2D Dask array. @@ -470,6 +494,10 @@ def _dask_subsample( "sequential" is chunk-dependent but slightly faster. Returns a delayed subsampled Dask array of the output (either values or indices). + + :param preserve_order: Restore sequential random-draw order after collecting block results. The one-dimensional + sampling adapter uses this to match NumPy point sampling; existing raster calls keep block order. + :param mask: Boolean array marking values eligible for sampling, with the same shape as darr. """ # To raise appropriate error on missing optional dependency @@ -483,9 +511,15 @@ def _dask_subsample( # Create a delayed object for each block, and flatten the blocks into a 1d shape blocks = darr.to_delayed().ravel() + # Give each data block the matching mask cells without loading either array + mask_blocks = [None] * len(blocks) + if mask is not None: + mask_blocks = da.asarray(mask).rechunk(darr.chunks).to_delayed().ravel().tolist() + # Compute number of valid points for each block out-of-memory list_delayed_valids = [ - da.from_delayed(_delayed_nb_valids(b), shape=(1, 1), dtype=np.dtype("int32")) for b in blocks + da.from_delayed(_delayed_nb_valids(b, mask_chunk=m), shape=(1, 1), dtype=np.dtype("int32")) + for b, m in zip(blocks, mask_blocks) ] # Compute once, then flatten nb_valids_per_block = np.concatenate([x.ravel() for x in dask.compute(*list_delayed_valids)], axis=0).astype( @@ -521,10 +555,16 @@ def _dask_subsample( ] # STRATEGY 1: "sequential" (chunk-dependent) - if strategy == "sequential": + if strategy == "sequential" or (preserve_order and subsample == 1): # Get random 1D indexes for the subsample size - indices_1d = rng.choice(total_nb_valids, subsample_size, replace=False) + indices_1d = ( + np.arange(total_nb_valids) + if preserve_order and subsample == 1 + else rng.choice(total_nb_valids, subsample_size, replace=False) + ) + # Block selection sorts valid positions; recover the original draw order only when requested + draw_order = np.argsort(np.argsort(indices_1d)) if preserve_order else slice(None) # Sort which indexes belong to which chunk ind_per_block = _get_indices_block_per_subsample( @@ -536,7 +576,10 @@ def _dask_subsample( # Task a delayed subsample to be computed for each block, skipping blocks with no values to sample used = [i for i in range(len(blocks)) if len(ind_per_block[i]) > 0] list_subsamples = [ - _delayed_subsample_block(blocks[i], np.asarray(ind_per_block[i], dtype=np.int64)) for i in used + _delayed_subsample_block( + blocks[i], np.asarray(ind_per_block[i], dtype=np.int64), mask_chunk=mask_blocks[i] + ) + for i in used ] # Cast output to the right expected dtype and length, then compute and concatenate @@ -544,7 +587,7 @@ def _dask_subsample( da.from_delayed(s, shape=(len(ind_per_block[i]),), dtype=darr.dtype) for s, i in zip(list_subsamples, used) ] - return da.concatenate(list_subsamples_da, axis=0) + return da.concatenate(list_subsamples_da, axis=0)[draw_order] # To return indices else: @@ -552,7 +595,10 @@ def _dask_subsample( used = [i for i in range(len(blocks)) if len(ind_per_block[i]) > 0] list_subsample_indices = [ _delayed_subsample_indices_block( - blocks[i], np.asarray(ind_per_block[i], dtype=np.int64), block_id=block_ids[i] + blocks[i], + np.asarray(ind_per_block[i], dtype=np.int64), + block_id=block_ids[i], + mask_chunk=mask_blocks[i], ) for i in used ] @@ -562,7 +608,7 @@ def _dask_subsample( da.from_delayed(s, shape=(len(ind_per_block[i]), 2), dtype=np.int32) for s, i in zip(list_subsample_indices, used) ] - indices = da.concatenate(list_indices_da, axis=0) + indices = da.concatenate(list_indices_da, axis=0)[draw_order] return indices[:, 0], indices[:, 1] # STRATEGY 2: "topk" (chunk-invariant; deterministic by (seed, global linear index)) @@ -588,6 +634,7 @@ def _dask_subsample( k=subsample_size, nx_full=nx_full, return_indices_local=return_indices, + mask_chunk=mask_blocks[i], ) for i in range(len(blocks)) ] @@ -618,7 +665,7 @@ def _dask_subsample( return rr, cc else: - raise ValueError(f"Unknown strategy {strategy!r}, available strategies are 'sequential' or 'topk'.") + raise ValueError(f"Unknown ``strategy`` {strategy!r}, available strategies are 'sequential' or 'topk'.") ################################ @@ -626,39 +673,82 @@ def _dask_subsample( ################################ -def _wrapper_multiproc_nb_valids_per_block(rst: Raster, tile_idx: NDArrayNum) -> int: +def _read_subsample_raster_block( + rst: RasterBase, + tile_idx: NDArrayNum, + *, + band: int = 1, + mask: RasterLike | VectorLike | ArrayLike | None = None, +) -> tuple[NDArrayNum | NDArrayBool | MArrayNum, NDArrayBool]: + """ + Read one band and its eligible cells from a raster tile without changing the source. + + Crop raster masks and slice array masks to the same tile; evaluate vector masks from its coordinates. + Both sampling passes use this helper so their finite counts and selected values use the same mask. + Keep the original values and dtype separate from eligibility so integer data need no NaN conversion. + """ + + from geoutils._dispatch import _get_raster_interface, _is_raster, has_geo_attr + from geoutils.sampling.support import _as_array, _mask_at_support + + # Read the tile and select its requested band without replacing the source raster's band selection + window = (tile_idx[2], tile_idx[0], tile_idx[3], tile_idx[1]) + rst_block = _get_raster_interface(rst.icrop(window)) + data = _as_array(rst_block.data) + arr = data if data.ndim == 2 else data[band - 1] + + # Exclude nodata and nonfinite values; boolean rasters sample only their True cells + valid = _valid_subsample_mask(arr) + + # Read only the corresponding mask window, then combine it with the finite cells in this tile + if mask is not None: + if _is_raster(mask): + mask = _get_raster_interface(mask).icrop(window) + elif not has_geo_attr(mask, "create_mask", accessors=("vct",)): + mask = _as_array(mask)[tile_idx[0] : tile_idx[1], tile_idx[2] : tile_idx[3]] + valid &= cast(NDArrayBool, _mask_at_support(mask, rst_block)) + return arr, valid + + +def _wrapper_multiproc_nb_valids_per_block( + rst: RasterBase, + tile_idx: NDArrayNum, + *, + band: int = 1, + mask: RasterLike | VectorLike | ArrayLike | None = None, +) -> int: """Count valid values in one tile out-of-memory.""" - rst_block = rst.icrop((tile_idx[2], tile_idx[0], tile_idx[3], tile_idx[1])) - arr = rst_block.data - if np.issubdtype(arr.dtype, np.bool_): - return int(np.count_nonzero(arr)) - return int(np.count_nonzero(~get_mask_from_array(arr))) + _, valid = _read_subsample_raster_block(rst, tile_idx, band=band, mask=mask) + return int(np.count_nonzero(valid)) def _wrapper_multiproc_subsample_values_block( - rst: Raster, + rst: RasterBase, tile_idx: NDArrayNum, subsample_indices_rel: NDArrayNum, -) -> NDArrayNum: + *, + band: int = 1, + mask: RasterLike | VectorLike | ArrayLike | None = None, +) -> NDArrayNum | NDArrayBool | MArrayNum: """ Subsample values in one tile using 1D indices relative to the tile's valid-value list. """ # Get tile out-of-memory - rst_block = rst.icrop((tile_idx[2], tile_idx[0], tile_idx[3], tile_idx[1])) - arr = rst_block.data + arr, valid = _read_subsample_raster_block(rst, tile_idx, band=band, mask=mask) # Return subsample of finite values (or True values for boolean input) - if np.issubdtype(arr.dtype, np.bool_): - return arr[arr].ravel()[subsample_indices_rel] - return arr[np.isfinite(arr)].ravel()[subsample_indices_rel] + return arr[valid].ravel()[subsample_indices_rel] def _wrapper_multiproc_subsample_indices_block( - rst: Raster, + rst: RasterBase, tile_idx: NDArrayNum, subsample_indices_rel: NDArrayNum, + *, + band: int = 1, + mask: RasterLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum: """ Return indices of the sampled valid pixels in one tile. @@ -667,18 +757,14 @@ def _wrapper_multiproc_subsample_indices_block( """ # Get tile out-of-memory - rst_block = rst.icrop((tile_idx[2], tile_idx[0], tile_idx[3], tile_idx[1])) - arr = rst_block.data + arr, valid = _read_subsample_raster_block(rst, tile_idx, band=band, mask=mask) # Get starting row/col of the tile row0 = int(tile_idx[0]) col0 = int(tile_idx[2]) # Get relative indices of finite values (or True for boolean) - if np.issubdtype(arr.dtype, np.bool_): - flat_valid = np.flatnonzero(arr.ravel()) - else: - flat_valid = np.flatnonzero(np.isfinite(arr).ravel()) + flat_valid = np.flatnonzero(valid.ravel()) # Use input to draw them flat_sel = flat_valid[subsample_indices_rel.astype(np.int64)] @@ -692,14 +778,16 @@ def _wrapper_multiproc_subsample_indices_block( def _wrapper_multiproc_topk_candidates_block( - rst: Raster, + rst: RasterBase, tile_idx: NDArrayNum, *, seed: int, k: int, nx_full: int, return_indices: bool, -) -> tuple[NDArrayNum, NDArrayNum | NDArrayBool]: + band: int = 1, + mask: RasterLike | VectorLike | ArrayLike | None = None, +) -> tuple[NDArrayNum, NDArrayNum | NDArrayBool | MArrayNum]: """ Return up to k candidates from one tile as (keys, payload). @@ -710,27 +798,24 @@ def _wrapper_multiproc_topk_candidates_block( """ # If no subsample, early return if k <= 0: - return np.empty((0,), dtype=np.uint64), np.empty((0,), dtype=np.int64) + payload_dtype = np.int64 if return_indices else rst.dtype + return np.empty((0,), dtype=np.uint64), np.empty((0,), dtype=payload_dtype) # Get tile out-of-memory - rst_block = rst.icrop((tile_idx[2], tile_idx[0], tile_idx[3], tile_idx[1])) - arr = rst_block.data + arr, valid = _read_subsample_raster_block(rst, tile_idx, band=band, mask=mask) # Tile offsets in full-array indices row0 = int(tile_idx[0]) col0 = int(tile_idx[2]) # Get valids indices - if np.issubdtype(arr.dtype, np.bool_): - valid = arr - else: - valid = np.isfinite(arr) flat = np.flatnonzero(valid.ravel()) nvalid = int(flat.size) # If no valid, early return if nvalid == 0: - return np.empty((0,), dtype=np.uint64), np.empty((0,), dtype=np.int64) + payload_dtype = np.int64 if return_indices else arr.dtype + return np.empty((0,), dtype=np.uint64), np.empty((0,), dtype=payload_dtype) # Get relative row and columns ncols = int(arr.shape[1]) @@ -752,6 +837,7 @@ def _wrapper_multiproc_topk_candidates_block( return key_sel, gid[sel] # If we return values + vals: NDArrayNum | NDArrayBool | MArrayNum if np.issubdtype(arr.dtype, np.bool_): vals = np.ones(m, dtype=np.bool_) else: @@ -766,6 +852,9 @@ def _multiproc_subsample( return_indices: bool = False, random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", + *, + band: int = 1, + mask: RasterLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum | tuple[NDArrayNum, NDArrayNum]: """ Subsample valid values out-of-memory from a 2D raster array using Multiprocessing tasks. @@ -774,6 +863,7 @@ def _multiproc_subsample( "sequential" is chunk-dependent but slightly faster. Returns a concatenated subsampled NumPy array collected from all tasks (either values or indices). + The mask must already match the source grid; each worker reads only its corresponding window. """ # Get tiling @@ -788,7 +878,10 @@ def _multiproc_subsample( tile_ids = [tiling[indexes_row[i], indexes_col[i], :] for i in range(num_blocks)] # Count valid values per tile in parallel - tasks = [config.cluster.submit(_wrapper_multiproc_nb_valids_per_block, rst, tile_ids[i]) for i in range(num_blocks)] + tasks = [ + config.cluster.submit(_wrapper_multiproc_nb_valids_per_block, rst, tile_ids[i], band=band, mask=mask) + for i in range(num_blocks) + ] try: nb_valids_per_block = np.array(config.cluster.gather(tasks), dtype=np.int64) except Exception as e: @@ -829,6 +922,8 @@ def _multiproc_subsample( rst, tile_ids[i], np.asarray(ind_per_block[i], dtype=np.int64), + band=band, + mask=mask, ) for i in used ] @@ -848,6 +943,8 @@ def _multiproc_subsample( rst, tile_ids[i], np.asarray(ind_per_block[i], dtype=np.int64), + band=band, + mask=mask, ) for i in used ] @@ -885,6 +982,8 @@ def _multiproc_subsample( k=subsample_size, nx_full=nx_full, return_indices=return_indices, + band=band, + mask=mask, ) for i in range(num_blocks) ] @@ -921,7 +1020,7 @@ def _multiproc_subsample( return rows.astype(np.int64), cols.astype(np.int64) else: - raise ValueError(f"Unknown strategy {strategy!r}. Choose 'sequential' or 'topk'.") + raise ValueError(f"Unknown ``strategy`` {strategy!r}. Choose 'sequential' or 'topk'.") ###################################################### @@ -938,23 +1037,39 @@ def _subsample( random_state: int | np.random.Generator | None = None, strategy: Literal["sequential", "topk"] = "sequential", mp_config: MultiprocConfig | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, ) -> Any: """ Subsample an array at valid values, dispatching automatically to NumPy, Dask or Multiprocessing implementation. - :param source_raster: Input array (NumPy/masked or Dask). - :param subsample: Subsample size or fraction. - :param band: Band to subsample. + _mask_at_support() places masks on the source grid before the NumPy and Dask samplers count eligible values. + _multiproc_subsample() instead places masks within each tile to keep unloaded inputs out of memory. + All paths keep eligibility separate from data values so masking preserves the sampled dtype and indices. + + :param source_raster: Raster or raster accessor providing band values and their optional Dask chunks. + :param subsample: Positive fraction of finite values at most one, or maximum number of values above one. + :param band: Raster band to subsample, counting from one. :param return_indices: If True, return (rows, cols) indices instead of values. :param random_state: Seed or Generator. :param strategy: Either "sequential" (chunk/order dependent) or "topk" (chunk-invariant). + :param mp_config: Tile sizes and worker cluster for multiprocessing. Cannot be combined with a Dask source. + :param mask: Boolean array, aligned mask raster, or vector geometries restricting eligible cells. - :returns: - - values: 1D array of sampled values - - indices: (rows, cols) (axis order) - - for Dask input: returns lazy `da.Array` unless compute=True + :returns: One-dimensional sampled values, or a tuple of row and column index arrays. Dask generally returns + lazy arrays after computing finite counts; an empty Dask sample returns NumPy arrays. """ + from geoutils._dispatch import _get_raster_interface, _is_raster, has_geo_attr + from geoutils.sampling.support import ( + _as_array, + _mask_at_support, + _normalize_sampling_input, + ) + + # Check the selected band before reading data or submitting worker tasks + if not 1 <= band <= source_raster.count: + raise ValueError("Argument ``band`` must be between one and the raster band count.") + # Cannot use Multiprocessing backend and Dask backend simultaneously mp_backend = mp_config is not None # The check below can only run on Xarray @@ -962,8 +1077,8 @@ def _subsample( if mp_backend and dask_backend: raise ValueError( - "Cannot use Multiprocessing and Dask simultaneously. To use Dask, remove mp_config parameter " - "from subsample(). To use Multiprocessing, open the file without chunks." + "Cannot use Multiprocessing and Dask simultaneously. To use Dask, remove ``mp_config`` parameter " + "from subsample(). To use Multiprocessing, open the file without ``chunks``." ) class _SubsampleKwargs(TypedDict): @@ -979,48 +1094,146 @@ class _SubsampleKwargs(TypedDict): "strategy": strategy, } - # Multiprocessing (out-of-memory) + # Validate masks before worker dispatch, keeping spatial masks available for reading one tile at a time if mp_backend: assert mp_config is not None - # Temporary switch bands - orig_bands = source_raster.bands - source_raster._bands = (band,) - try: - return _multiproc_subsample(source_raster, config=mp_config, **subsample_kwargs) - finally: - source_raster._bands = orig_bands + mask = _normalize_sampling_input(mask) + if _is_raster(mask): + mask_raster = _get_raster_interface(mask) + if mask_raster._chunks is not None: + raise ValueError("Cannot use Multiprocessing and Dask masks simultaneously.") + if not source_raster.georeferenced_grid_equal(mask_raster): + raise ValueError("Raster value ``mask`` does not share the selected support grid.") + if not mask_raster.is_mask and not np.issubdtype(mask_raster.dtype, np.bool_): + raise ValueError("Argument ``mask`` must be boolean and contain one value per input location.") + mask = mask_raster + elif mask is not None and not has_geo_attr(mask, "create_mask", accessors=("vct",)): + mask = _mask_at_support(mask, source_raster) + if is_dask_array(mask): + raise ValueError("Cannot use Multiprocessing and Dask masks simultaneously.") + # Restrict disk reads to the selected band through a shallow view, leaving the caller's band selection intact + sampling_raster = source_raster + sampling_band = band + if not source_raster._is_xr and not source_raster.is_loaded: + sampling_raster = source_raster.copy(deep=False) + sampling_raster._bands = (source_raster.bands[band - 1],) + sampling_band = 1 + return _multiproc_subsample( + sampling_raster, config=mp_config, band=sampling_band, mask=mask, **subsample_kwargs + ) + + # Read one band without converting masked integer values to floating point + data = _as_array(source_raster.data) + arr = data if data.ndim == 2 else data[band - 1] + mask_array = _mask_at_support(mask, source_raster) + + # Keep Dask data and mask blocks lazy; an eager source collects only a supplied lazy mask + if dask_backend: + return _dask_subsample(arr, mask=mask_array, **subsample_kwargs) + if mask_array is not None and is_dask_array(mask_array): + mask_array = mask_array.compute() + return _subsample_numpy(arr, mask=mask_array, **subsample_kwargs) # type: ignore + + +def _sample_valid_indices( + valid: Any, + *, + subsample: int | float, + random_state: int | np.random.Generator | None, + strategy: Literal["sequential", "topk"], +) -> tuple[np.typing.NDArray[np.int64], ...]: + """ + Choose eligible array positions without collecting a complete lazy validity mask. + + Sampling options follow _subsample(). One-dimensional inputs use a single-column grid so Dask's existing block + sampler can work on point rows. This preserves global row positions, topk keys and NumPy's sequential draw order. + Two-dimensional sequential sampling keeps the existing Dask block order. + + :param valid: One- or two-dimensional boolean NumPy or Dask array marking eligible locations. + :returns: Computed integer position arrays, one per input dimension. + """ + + # Accept point rows or grid cells, the dimensions supported by the shared samplers + if valid.ndim not in (1, 2): + raise ValueError("Valid sampling locations must form a one- or two-dimensional array.") + + # Dask's sampler already selects True cells; avoid expanding boolean blocks into floating-point arrays + if is_dask_array(valid): + grid = valid[:, None] if valid.ndim == 1 else valid + indexes = _dask_subsample( + grid, + subsample=subsample, + return_indices=True, + random_state=random_state, + strategy=strategy, + preserve_order=valid.ndim == 1, + ) + # Compute both coordinate arrays together so they share the same sampling tasks + indexes = dask.compute(*indexes) + if valid.ndim == 1: + indexes = indexes[:1] else: - if source_raster.data.ndim != 2: - arr = source_raster.data[band - 1, :, :] - else: - arr = source_raster.data - # Dask (out-of-memory) - if dask_backend: - return _dask_subsample(arr, **subsample_kwargs) - # NumPy - else: - return _subsample_numpy(arr, **subsample_kwargs) # type: ignore + # NumPy's sampler selects finite values, so represent excluded cells by NaN + sampling_values = np.where(valid, 1.0, np.nan) + indexes = _subsample_numpy( + sampling_values, + subsample=subsample, + return_indices=True, + random_state=random_state, + strategy=strategy, + ) + return tuple(np.asarray(index, dtype=np.int64) for index in indexes) def _subsample_pointcloud( - source_pointcloud: Any, + source_pointcloud: PointCloudBase, subsample: float | int, return_indices: bool = False, random_state: int | np.random.Generator | None = None, + *, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None = None, ) -> NDArrayNum | tuple[NDArrayNum, ...]: - """Subsample point-cloud values after materializing any lazy data column.""" + """ + Subsample finite point cloud values, gathering only selected values from a lazy data column. + + _dask_subsample() treats point rows as a single-column grid and preserves NumPy's random-draw order. + _mask_at_support() places the optional mask on these rows before either sampler counts eligible values. + Eager columns use _subsample_numpy() directly. Both paths return computed results without changing the value dtype. + + :param source_pointcloud: Point cloud or accessor whose selected data column supplies the sampled values. + :param subsample: Positive fraction of finite values at most one, or maximum number of values above one. + :param return_indices: Return a one-element tuple of row-position indices instead of sampled values. + :param random_state: Integer seed or NumPy Generator used for the random draw. + :param mask: Boolean array or spatial mask restricting eligible point rows. + :returns: Computed one-dimensional NumPy values, or a one-element tuple of indices into the original row order. + """ - data = source_pointcloud.data.compute().values if source_pointcloud._is_dask else np.asarray(source_pointcloud.data) - if return_indices: - return _subsample_numpy( - array=data, + from geoutils.sampling.support import _as_array, _mask_at_support + + # Read native values and reuse any known partition lengths when placing the mask on the same point rows + data = _as_array(source_pointcloud.data) + partition_lengths = tuple(int(length) for length in data.chunks[0]) if source_pointcloud._is_dask else None + mask_array = _mask_at_support(mask, source_pointcloud, point_partition_lengths=partition_lengths) + + # Let the existing Dask sampler collect only the requested values or row positions from point partitions + if source_pointcloud._is_dask: + sampled = _dask_subsample( + data[:, None], subsample=subsample, - return_indices=True, + return_indices=return_indices, random_state=random_state, + preserve_order=True, + mask=None if mask_array is None else mask_array[:, None], ) - return _subsample_numpy( - array=data, - subsample=subsample, - return_indices=False, - random_state=random_state, - ) + if return_indices: + rows = sampled[0] + return (rows.compute() if is_dask_array(rows) else rows,) + assert not isinstance(sampled, tuple) + return sampled.compute() if is_dask_array(sampled) else sampled + + # Preserve the established NumPy sampling rules for eager point data + if mask_array is not None and is_dask_array(mask_array): + mask_array = mask_array.compute() + if return_indices: + return _subsample_numpy(data, subsample, return_indices=True, random_state=random_state, mask=mask_array) + return _subsample_numpy(data, subsample, return_indices=False, random_state=random_state, mask=mask_array) diff --git a/geoutils/sampling/support.py b/geoutils/sampling/support.py new file mode 100644 index 000000000..db3870009 --- /dev/null +++ b/geoutils/sampling/support.py @@ -0,0 +1,707 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +""" +Shared module to sample geospatial inputs on shared spatial supports: grids or point locations. + +Especially used by cosampling, but also pair/subsampling for masking. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterable +from contextlib import ExitStack +from typing import TYPE_CHECKING, Any, Literal, cast, overload + +import geopandas as gpd +import numpy as np +import pandas as pd +import xarray as xr + +from geoutils._dispatch import ( + _get_pointcloud_interface, + _get_raster_interface, + _is_pointcloud, + _is_raster, + _is_vector, + get_geo_attr, + has_geo_attr, + is_dask_array, + is_dask_dataframe, +) +from geoutils._misc import import_optional +from geoutils._typing import ArrayLike, NDArrayNum +from geoutils.raster.array import _selected_raster_data +from geoutils.vector.base import _as_geodataframe + +if TYPE_CHECKING: + from geoutils.interface.interpolation import InterpolationMethod + from geoutils.multiproc import MultiprocConfig + from geoutils.pointcloud.base import PointCloudBase + from geoutils.pointcloud.pointcloud import PointCloudLike + from geoutils.raster.base import RasterBase, RasterLike + from geoutils.vector.base import VectorLike + +################################# +# 1/ ARRAY VALUES AND INPUT GRIDS +################################# + + +def _normalize_sampling_input(value: Any) -> Any: + """Treat Xarray inputs without both x and y coordinates as plain arrays, preserving Dask storage.""" + + if isinstance(value, xr.DataArray) and not {"x", "y"}.issubset(value.coords): + return value.data + return value + + +def _as_array(value: Any) -> Any: + """ + Extract array values while preserving NumPy masks and Dask arrays. + + :param value: Array, Xarray DataArray, or Pandas/Dask Series, Index or DataFrame to unwrap. + + :returns: Array values with existing NumPy masks or Dask storage preserved. + """ + + # Extract values from Xarray and Pandas inputs as NumPy arrays, or Dask arrays for chunked data + if isinstance(value, xr.DataArray): + return value.data + if is_dask_dataframe(value): + return value.to_dask_array(lengths=True) + if isinstance(value, (pd.Series, pd.Index, pd.DataFrame)): + return value.to_numpy(copy=False) + return value if is_dask_array(value) else np.asanyarray(value) + + +def _normalize_mask_array(mask: Any | None, shape: tuple[int, ...]) -> Any: + """ + Return a boolean mask value per location, keeping missing mask entries ineligible. + + The mask comes from _mask_at_support(); this step checks array values after any spatial placement. + + :param shape: Required output shape. The mask must contain the same number of entries and is reshaped to match. + + :returns: Boolean NumPy or Dask array, or None when no mask was supplied. + """ + + # Return None for no mask, to avoid allocating extra memory + if mask is None: + return None + + # Ensure lazy arrays and proper missing values in eager masked arrays + values = _as_array(mask) + if np.ma.isMaskedArray(values): + values = values.filled(False) + elif is_dask_array(values): + # Apply the same missing-entry rule within lazy blocks without collecting the mask + import_optional("dask") + import dask.array as da + + values = da.ma.filled(values, False) + if not np.issubdtype(values.dtype, np.bool_) or values.size != math.prod(shape): + raise ValueError("Argument ``mask`` must be boolean and contain one value per input location.") + return values.reshape(shape) + + +def _raster_from_input( + value: RasterLike | ArrayLike, input_support: RasterLike | PointCloudLike, name: str +) -> RasterBase: + """ + Return a raster input or attach its original grid metadata to a plain array. + + ``value`` is the source passed to _values_at_support(); input_support and name follow that parent. + Spatial raster inputs use their own metadata. Plain arrays must match the input_support grid. + """ + + # Use the raster's own grid information when the input is already a raster or raster accessor + value = _normalize_sampling_input(value) + raster = _get_raster_interface(value) + if raster is not None: + return raster + + # Find the raster that supplies coordinates for a plain array; an array alone has no grid information + input_raster = _get_raster_interface(input_support) + if input_raster is None: + raise ValueError(f"Two-dimensional value {name!r} must be tied to a raster input.") + + # Read the array and drop a band dimension of length one, then require the original raster shape + array: Any = value if hasattr(value, "ndim") else np.asarray(value) + if array.ndim == 3 and array.shape[0] == 1: + array = array[0] + if array.ndim != 2 or tuple(array.shape) != tuple(input_raster.shape): + raise ValueError(f"Array {name!r} must match the shape of its native raster input.") + + # Replace masked entries with NaN using NumPy's dtype promotion so large integer values keep their precision + if np.ma.isMaskedArray(array): + array = np.where(np.ma.getmaskarray(array), np.nan, np.ma.getdata(array)) + + # Keep raw Dask arrays lazy even when their grid comes from a GeoUtils Raster + from_array = input_raster.from_array + if is_dask_array(array): + from geoutils.raster.xr_accessor import RasterAccessor + + from_array = RasterAccessor.from_array + + # Create a raster with the array values and the original grid's coordinates + # Mark missing array values without applying the original raster's nodata value + raster = from_array( + data=array if is_dask_array(array) else np.ma.masked_invalid(array), + transform=input_raster.transform, + crs=input_raster.crs, + nodata=None, + area_or_point=input_raster.area_or_point, + ) + return _get_raster_interface(raster) + + +def _aligned_raster( + value: RasterLike | ArrayLike, + input_support: RasterLike | PointCloudLike, + support: RasterLike | PointCloudLike, + name: str, + align: str, + mp_config: MultiprocConfig | None = None, +) -> RasterBase: + """ + Return a raster reprojected to the output support. + + Match the complete grid for raster output, or only the coordinate system for point output. + + See _values_at_support() for arguments. + """ + + # Turn plain arrays into rasters using their original grid before checking the output coordinates + raster = _raster_from_input(value, input_support, name) + + # Use nearest-neighbor resampling for masks so interpolation does not turn booleans into fractions + resampling = "nearest" if raster.is_mask else None + if has_geo_attr(support, "georeferenced_grid_equal", ("rst",)): + if cast("RasterBase", support).georeferenced_grid_equal(raster): + return raster + + # Reproject to the full output grid when allowed, or report the grid mismatch + if align == "reproject": + projected = raster.reproject(ref=support, resampling=resampling, silent=True, mp_config=mp_config) + return _get_raster_interface(projected) + raise ValueError(f"Raster value {name!r} does not share the selected support grid.") + + # For point output, match only the CRS because the output points do not define a raster grid + if raster.crs != support.crs: + if align != "reproject": + raise ValueError(f"Raster value {name!r} does not share the point support CRS.") + raster = raster.reproject(crs=support.crs, resampling=resampling, silent=True, mp_config=mp_config) + + # Return a raster or its accessor so later steps can call the same raster methods + return _get_raster_interface(raster) + + +def _mask_on_raster( + mask: RasterLike | VectorLike | ArrayLike, + support: RasterBase, + mask_mode: str, + align: str, + mp_config: MultiprocConfig | None = None, +) -> Any: + """ + Evaluate a user mask on raster support. + + See _mask_at_support() for arguments. + + Return one boolean value per support cell. Vector masks use mask_mode, while raster and plain array masks use + their boolean values directly. The caller handles an absent mask before calling this function. + """ + + # Read Xarray masks without x and y coordinates as plain arrays + mask = _normalize_sampling_input(mask) + is_raster = _is_raster(mask) + + # Mark cells inside vector shapes, then invert the mask when outside was requested + if not is_raster and has_geo_attr(mask, "create_mask", accessors=("vct",)): + create_mask = get_geo_attr(mask, "create_mask", accessors=("vct",)) + values = create_mask(ref=support, as_array=True, mp_config=mp_config) + values = _normalize_mask_array(values, tuple(support.shape)) + return values if mask_mode == "inside" else ~values + + # Read raster masks on the output grid; plain arrays already correspond to its cells + if is_raster: + mask_raster = _aligned_raster(mask, mask, support, "mask", align, mp_config=mp_config) + values = _selected_raster_data(mask_raster, fill_value=False) + else: + values = _as_array(mask) + + # Remove a band dimension of length one, but preserve rows and columns even when their size is one + if values.ndim == 3 and values.shape[0] == 1: + values = values[0] + + # Require the mask to match the output grid and contain boolean values + if tuple(values.shape) != tuple(support.shape): + raise ValueError("A raster support mask must be boolean and match the support grid.") + return _normalize_mask_array(values, tuple(support.shape)) + + +################################# +# 2/ SUPPORT AND VALUE SELECTION +################################# + + +@overload +def _sampling_support(inputs: Iterable[Any], at: RasterBase) -> RasterBase: ... + + +@overload +def _sampling_support(inputs: Iterable[Any], at: PointCloudBase) -> PointCloudBase: ... + + +@overload +def _sampling_support(inputs: Iterable[Any], at: RasterLike | PointCloudLike | None) -> RasterBase | PointCloudBase: ... + + +def _sampling_support(inputs: Iterable[Any], at: RasterLike | PointCloudLike | None) -> RasterBase | PointCloudBase: + """ + Choose the grid or point locations shared by all requested values. + + :param inputs: Ordered raster/point cloud inputs, optionally paired with a band or column selector. With no + explicit ``at`` passed, use the first point cloud if present, otherwise the first input. + :param at: Explicit raster or point cloud support, or None. The caller resolves any 'self'/'other' names first. + + :returns: Raster or point cloud interface supplying the chosen output locations. + """ + + at = _normalize_sampling_input(at) + + # Unless at is provided, use the first point cloud's locations, or the first input if no point cloud is present + if at is None: + # Take the object from each (object, selector) pair only when input locations determine the support + objects = [_normalize_sampling_input(value[0] if isinstance(value, tuple) else value) for value in inputs] + at = objects[0] + for value in objects: + pointcloud = _get_pointcloud_interface(value) + if pointcloud is not None: + at = pointcloud + break + + # Return the chosen raster or point cloud, using its accessor when needed + raster = _get_raster_interface(at) + pointcloud = _get_pointcloud_interface(at) + if raster is None and pointcloud is None: + raise TypeError("Argument ``at`` must select raster or point cloud support.") + return raster if raster is not None else pointcloud + + +def _sampling_specification(source: RasterLike | PointCloudLike, specification: Any) -> tuple[Any, Any]: + """ + Split a value request into its data object and optional band or column. + + :param source: Calling raster or point cloud used when the request contains only a selector. + :param specification: Band number, column name, separate data object, or (spatial object, selector) pair. + None selects the calling object's default value. + :returns: Data object and selector; a separate object with no explicit selection receives a None selector. + """ + + # Use the calling object when the request only specifies a band number or column name + if specification is None or isinstance(specification, (str, int, np.integer)): + return source, specification + + # Read (object, band or column) pairs only when the first item is a raster, point cloud or vector + if isinstance(specification, tuple) and len(specification) == 2: + value = specification[0] + if _is_raster(value) or _is_pointcloud(value) or _is_vector(value): + return specification + return specification, None + + +def _vector_values_at_points(points: gpd.GeoDataFrame, features: gpd.GeoDataFrame) -> pd.Series: + """ + Assign vector feature values to points, with later features winning at overlaps. + + Called by _sample_vector_values(): points is its support_dataframe (or one Dask partition), and features + contains geometries in the point CRS and their numeric value column. Return a Series in the original point + order, with NaN outside all features. + """ + + # Number points by row so duplicate index labels cannot mix up their matches + left = gpd.GeoDataFrame(geometry=points.geometry.reset_index(drop=True), crs=points.crs) + + # Match points to intersecting features in the same CRS + # Sort matches by feature row so the last overlapping feature supplies the value + matches = gpd.sjoin(left, features, how="inner", predicate="intersects").sort_values("index_right") + + # Write matched values back in point order and leave unmatched points as NaN + output = np.full(len(points), np.nan) + output[matches.index.to_numpy()] = matches["value"].to_numpy() + return pd.Series(output, index=points.index, name="value") + + +def _sample_vector_values( + dataframe: gpd.GeoDataFrame, + values: NDArrayNum, + support: RasterBase | PointCloudBase, + support_dataframe: Any | None, + mp_config: MultiprocConfig | None = None, +) -> Any: + """ + Place numeric feature values on the chosen grid or point locations. + + Rasterization handles grid alignment and Dask or multiprocessing blocks. Point values use a spatial join, + partitioned for Dask inputs, with later features winning at overlaps and missing values outside coverage. + + The output support, support_dataframe and mp_config follow _values_at_support(). + + :param dataframe: Vector features whose geometries define where each value applies. + :param values: One numeric value per dataframe row, in the same order. + :returns: Array on the support grid or ordered support points, with NaN outside all features. + """ + + # Rasterize feature row numbers starting at one, reserving zero for cells outside all features + if _is_raster(support): + indexes = np.arange(1, len(values) + 1) + rasterize = get_geo_attr(dataframe, "rasterize", accessors=("vct",)) + raster = rasterize(ref=support, in_value=indexes.tolist(), out_value=0, out_dtype=np.int32, mp_config=mp_config) + + # Replace feature numbers with their values, and replace zero with NaN + codes = _selected_raster_data(raster).astype(np.int64) + return np.take(np.concatenate(([np.nan], values)), codes) + + # Number features by row so duplicate index labels cannot mix up their matches + if support_dataframe is None: + raise RuntimeError("Point support coordinates were not prepared.") + features = gpd.GeoDataFrame( + {"value": values}, geometry=dataframe.geometry.reset_index(drop=True), crs=dataframe.crs + ) + + # Match feature coordinates to the output CRS once before sampling every point partition + if features.crs != support_dataframe.crs: + features = features.to_crs(support_dataframe.crs) + + # Assign feature values to output points, processing each Dask chunk separately when needed + if is_dask_dataframe(support_dataframe): + sampled = support_dataframe.map_partitions( + _vector_values_at_points, features, meta=pd.Series([], dtype=float, name="value") + ) + return sampled.to_dask_array(lengths=True) + return _vector_values_at_points(support_dataframe, features).to_numpy() + + +def _aligned_pointcloud( + pointcloud: PointCloudBase, + support: RasterBase | PointCloudBase, + name: str, + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None = None, +) -> PointCloudBase: + """ + Match point coordinates to the support CRS and, for point output, its ordered XY locations. + + Raster output only requires matching coordinate systems before gridding. Point output also requires identical + ordered coordinates after any reprojection, so later value selection can use positional arrays. + Multiprocessing callers provide a temporary configuration kept alive until the aligned points have been read. + """ + + # Match the output CRS before gridding points or comparing their ordered coordinates + if pointcloud.crs != support.crs: + if align != "reproject": + raise ValueError(f"Point value {name!r} does not share the support CRS.") + point_config = None + if mp_config is not None: + # Preserve exact coordinates in a separate point file and adapt raster tiles to row partitions + point_config = mp_config.copy() + point_config.driver = "GPKG" + point_config.outfile += ".gpkg" + if isinstance(point_config.chunks, tuple): + point_config.chunks = math.prod(point_config.chunks) + pointcloud = pointcloud.reproject(crs=support.crs, mp_config=point_config) + pointcloud = _get_pointcloud_interface(pointcloud) + + # Require the same point coordinates in the same order when output locations are points + if not _is_raster(support) and pointcloud is not support: + if not cast("PointCloudBase", support).georeferenced_coords_equal(pointcloud): + raise ValueError(f"Point value {name!r} does not share the ordered support coordinates.") + return pointcloud + + +def _point_values_at_support( + source: PointCloudBase | ArrayLike, + selector: int | str | None, + *, + support_dataframe: Any | None, + name: str, + point_partition_lengths: tuple[int, ...] | None = None, +) -> Any: + """ + Read an aligned point cloud column or plain array in output point order. + + Point cloud coordinates are prepared by _aligned_pointcloud() before this reader is called. Plain arrays must + contain one value per support point. Keep Dask storage and reuse known partition lengths when reading the + support dataframe. + """ + + # Read the requested column, or use Z coordinates when no column is selected + source_pointcloud = _get_pointcloud_interface(source) + if source_pointcloud is not None: + dataframe = source_pointcloud.ds + column = source_pointcloud.data_column if selector is None else selector + if column is not None and (not isinstance(column, str) or column not in dataframe.columns): + raise ValueError(f"Point column {column!r} selected for {name!r} does not exist.") + values = dataframe.geometry.z if column is None else dataframe[column] + + # Reuse output chunk lengths only when reading from that same dataframe + if is_dask_dataframe(dataframe): + lengths = ( + point_partition_lengths + if dataframe is support_dataframe and point_partition_lengths is not None + else True + ) + return values.to_dask_array(lengths=lengths) + return np.asarray(values) + + # Replace masked entries with NaN and remove extra dimensions around the point values + point_array = _as_array(source) + if np.ma.isMaskedArray(point_array): + point_array = np.where(np.ma.getmaskarray(point_array), np.nan, np.ma.getdata(point_array)) + array = np.atleast_1d(point_array.squeeze()) + + # Check the array length against the output point count, using saved Dask chunk lengths when available + point_count = ( + sum(point_partition_lengths) + if point_partition_lengths is not None + else (len(support_dataframe) if support_dataframe is not None else None) + ) + if array.ndim != 1 or len(array) != point_count: + raise ValueError(f"Raw point value {name!r} must contain one value per support point.") + return array + + +def _values_at_support( + source: RasterLike | PointCloudLike | VectorLike | ArrayLike, + selector: int | str | None, + *, + input_support: RasterLike | PointCloudLike, + support: RasterBase | PointCloudBase, + support_dataframe: Any | None, + name: str, + interpolation: InterpolationMethod, + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None, + point_partition_lengths: tuple[int, ...] | None = None, +) -> Any: + """ + Read one requested value (band of raster or column of point cloud) at every location chosen for the final result. + + _aligned_raster() prepares grids before band selection or interpolation. _aligned_pointcloud() checks point + locations before _point_values_at_support() reads their values, while _sample_vector_values() assigns feature + values by location. This parent defines the input and output supports used by the spatial helpers shared by + co-sampling and statistics. + + :param source: Raster, point cloud, vector, or plain array providing the values to place. + :param selector: Raster band number (counting from one) or point/vector column name. None uses the first raster + band or active point value; vectors require an explicit column. + :param input_support: Raster or point cloud providing the original grid or ordered coordinates associated with + plain input arrays. Spatial source objects already carry their own locations. + :param support: Raster or point cloud defining where the result is evaluated, which may differ from input_support. + :param support_dataframe: Prepared GeoDataFrame or Dask GeoDataFrame in output point order; None for a raster grid. + :param name: Value label used to identify this input in error messages. + :param interpolation: Raster interpolation method when evaluating values at points. + :param align: Whether mismatched grids or coordinate systems raise an error or are reprojected to the support. + :param mp_config: Optional multiprocessing configuration passed to spatial operations that support it. + :param point_partition_lengths: Optional known support_dataframe partition lengths, avoiding repeated row counts + when reading columns from that same table. + :returns: One selected value per support cell or point, as a NumPy or Dask array. + """ + + # Identify rasters and point clouds first so they are not mistaken for plain arrays or vectors + source = _normalize_sampling_input(source) + source_raster = _get_raster_interface(source) + source_pointcloud = _get_pointcloud_interface(source) + support_is_raster = _is_raster(support) + + # Read the requested numeric vector column at each output location + is_vector = source_raster is None and source_pointcloud is None and _is_vector(source) + if is_vector: + dataframe = _as_geodataframe(source) + if selector is None or selector not in dataframe.columns: + raise ValueError("Vector values require an explicit feature column.") + if not pd.api.types.is_numeric_dtype(dataframe[selector]): + raise TypeError("Selected vector values must be numeric.") + return _sample_vector_values( + dataframe, np.asarray(dataframe[selector], dtype=float), support, support_dataframe, mp_config=mp_config + ) + + # Inspect plain arrays only after excluding spatial objects, whose values may still be file-backed + if source_raster is None and source_pointcloud is None: + direct_values: Any = source + raw_ndim = np.ndim(direct_values) + input_raster = _get_raster_interface(input_support) + if raw_ndim >= 2 and input_raster is not None: + # Return arrays directly when their shape and original grid already match the output grid + support_shape = tuple(cast("RasterBase", support).shape) if support_is_raster else None + if raw_ndim == 3 and direct_values.shape[0] == 1: + direct_values = direct_values[0] + if ( + support_shape is not None + and tuple(direct_values.shape) == support_shape + and input_raster.georeferenced_grid_equal(support) + ): + if np.ma.isMaskedArray(direct_values): + direct_values = np.where(np.ma.getmaskarray(direct_values), np.nan, np.ma.getdata(direct_values)) + return direct_values + + # Give other plain arrays their original grid before moving them to the output locations + source_raster = _raster_from_input(source, input_support, name) + + # Select the requested raster band and match its grid or CRS to the output locations + if source_raster is not None: + if selector is not None and not isinstance(selector, (int, np.integer)): + raise TypeError(f"Raster selector for {name!r} must be a band number.") + band = 1 if selector is None else int(selector) + raster = _aligned_raster(source_raster, source_raster, support, name, align, mp_config=mp_config) + if support_is_raster: + return _selected_raster_data(raster, band) + + # Interpolate raster values at the output point coordinates + if support_dataframe is None: + raise RuntimeError("Point support coordinates were not prepared.") + points = ( + support_dataframe + if is_dask_dataframe(support_dataframe) + else (support_dataframe.geometry.x.to_numpy(), support_dataframe.geometry.y.to_numpy()) + ) + known_partitions = is_dask_dataframe(points) and point_partition_lengths is not None + values = raster.interp_points( + points=points, + method=interpolation, + band=band, + as_array=not known_partitions, + mp_config=mp_config, + ) + + # Reuse known chunk lengths when converting interpolated point columns back to arrays + if known_partitions: + return get_geo_attr(values, "data", ("pc",)).to_dask_array(lengths=point_partition_lengths) + return values + + # Reject point values on a raster grid here because only cosample() provides a gridding method + if support_is_raster: + if source_pointcloud is not None: + raise ValueError(f"Point value {name!r} cannot be evaluated on raster support without gridding.") + raise ValueError(f"Raw value {name!r} cannot be tied to the selected spatial support.") + + # Read point values after checking that their ordered coordinates match the output support + with ExitStack() as temporary_files: + if source_pointcloud is not None: + intermediate = temporary_files.enter_context(mp_config.temporary()) if mp_config is not None else None + source = _aligned_pointcloud(source_pointcloud, support, name, align, mp_config=intermediate) + return _point_values_at_support( + cast("PointCloudBase | ArrayLike", source), + selector, + support_dataframe=support_dataframe, + name=name, + point_partition_lengths=point_partition_lengths, + ) + + +def _mask_at_support( + mask: RasterLike | VectorLike | ArrayLike | None, + support: RasterBase | PointCloudBase, + *, + support_dataframe: Any | None = None, + mask_mode: str = "inside", + align: Literal["raise", "reproject"] = "raise", + mp_config: MultiprocConfig | None = None, + point_partition_lengths: tuple[int, ...] | None = None, +) -> Any | None: + """ + Reproject a boolean mask on raster or point support for sampling and statistics. + + _mask_on_raster() handles raster grids. For point support, create_mask() identifies points inside vector features + and _values_at_support() aligns raster or point masks. _normalize_mask_array() applies the same boolean + and missing-value rules to their results. + + A missing mask input stays None to avoid allocating an unused mask. + + :param mask: Boolean array or spatial raster, point cloud or vector mask. None keeps all locations eligible. + :param support: Raster or point cloud defining the grid or ordered points on which to evaluate the mask. + :param support_dataframe: Prepared output point table, optionally lazy. If omitted for point support, use its ds. + :param mask_mode: Keep locations inside or outside vector features. Boolean masks use their values directly. + :param align: Whether a raster or point mask with a different grid or coordinate system raises or is reprojected. + :param mp_config: Optional multiprocessing configuration for spatial mask placement. + :param point_partition_lengths: Optional known support_dataframe partition lengths, as in _values_at_support(). + :returns: Boolean array on the support, optionally lazy, or None when no mask was supplied. + """ + + # 1/ Check mask options and use the raster mask helper when the output is a grid + if mask is None: + return None + if mask_mode not in {"inside", "outside"}: + raise ValueError("Argument ``mask_mode`` must be 'inside' or 'outside'.") + if _is_raster(support): + return _mask_on_raster(mask, cast("RasterBase", support), mask_mode, align, mp_config=mp_config) + if support_dataframe is None: + support_dataframe = get_geo_attr(support, "ds", accessors=("pc",)) + mask = _normalize_sampling_input(mask) + + # Check for point clouds before vectors because point clouds also have vector methods + mask_is_raster = _is_raster(mask) + mask_is_pointcloud = _is_pointcloud(mask) + is_vector = not mask_is_raster and not mask_is_pointcloud and _is_vector(mask) + + # 2/ Calculate which output points the mask allows + + # If input is a vector + if is_vector: + # Use create_mask() to find points inside vector shapes, excluding their boundaries + create_mask = get_geo_attr(mask, "create_mask", accessors=("vct",)) + known_partitions = is_dask_dataframe(support_dataframe) and point_partition_lengths is not None + values = create_mask(ref=support_dataframe, as_array=not known_partitions, mp_config=mp_config) + + # Reuse known chunk lengths when reading the point mask from a Dask dataframe + if known_partitions: + values = get_geo_attr(values, "data", ("pc",)).to_dask_array(lengths=point_partition_lengths) + if mask_mode == "outside": + values = ~values + + # If input is raster or point cloud + elif mask_is_raster or mask_is_pointcloud: + # Require boolean raster or point values so numeric data cannot be mistaken for a mask + dtype = ( + get_geo_attr(mask, "dtype", accessors=("rst",)) + if mask_is_raster + else get_geo_attr(mask, "data", accessors=("pc",)).dtype + ) + is_boolean = np.issubdtype(dtype, np.bool_) + if mask_is_raster: + is_boolean |= get_geo_attr(mask, "is_mask", accessors=("rst",)) + if not is_boolean: + raise ValueError("A point support mask must contain boolean values.") + + # Read mask values at the output points; nearest interpolation preserves boolean raster values + sampled = _values_at_support( + mask, + 1 if mask_is_raster else None, + input_support=mask, + support=support, + support_dataframe=support_dataframe, + name="mask", + interpolation="nearest", + align=align, + mp_config=mp_config, + point_partition_lengths=point_partition_lengths, + ) + + # Exclude missing mask values as well as False values + values = np.isfinite(sampled) & (sampled != 0) + + # If input is a boolean array + else: + # Check that the plain array has one boolean value for each output point + count = sum(point_partition_lengths) if point_partition_lengths is not None else len(support_dataframe) + return _normalize_mask_array(mask, (count,)) + + # 3/ Return one boolean value per point without counting the Dask rows again + return _normalize_mask_array(values, (values.size,)) diff --git a/geoutils/stats/__init__.py b/geoutils/stats/__init__.py index 1cceb3d88..235ae01ce 100644 --- a/geoutils/stats/__init__.py +++ b/geoutils/stats/__init__.py @@ -16,6 +16,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Statistical functionalities: grouped/zonal/global stats, specific estimators, and variography/geostats. """ from geoutils.stats.estimators import * # noqa -from geoutils.stats.sampling import * # noqa +from geoutils.stats.grouping import * # noqa from geoutils.stats.stats import * # noqa +from geoutils.stats.variography import * # noqa diff --git a/geoutils/stats/estimators.py b/geoutils/stats/estimators.py index 68532d76d..3720fd2d2 100644 --- a/geoutils/stats/estimators.py +++ b/geoutils/stats/estimators.py @@ -46,17 +46,17 @@ def nmad(data: NDArrayNum, nfact: float = 1.4826) -> np.floating[Any]: def linear_error(data: NDArrayNum, interval: float = 90) -> np.floating[Any]: """ - Compute the linear error (LE) for a given dataset, representing the range of differences between the upper and - lower percentiles of the data. By default, this calculates the 90% confidence interval (LE90). + Compute the linear error (LE) for a given dataset, representing the difference between the upper and lower + percentiles of the data. By default, this calculates the central interval containing 90% of the values (LE90). :param data: A numpy array or masked array of data, typically representing the differences (errors) in elevation or - another quantity. - :param interval: The confidence interval to compute, specified as a percentage. For example, an interval of 90 will - compute the range between the 5th and 95th percentiles (LE90). This value must be between 0 and 100. + another quantity. + :param interval: The central interval to compute, specified as a percentage. For example, an interval of 90 will + compute the range between the 5th and 95th percentiles (LE90). This value must be greater than 0 and at most 100. - return: The computed linear error, which is the difference between the upper and lower percentiles. + :returns: The computed linear error, which is the difference between the upper and lower percentiles. - raises: ValueError if the `interval` is not between 0 and 100. + :raises ValueError: If interval is not greater than 0 and at most 100. """ # Validate the interval if not (0 < interval <= 100): @@ -77,9 +77,12 @@ def sum_square(data: NDArrayNum) -> np.floating[Any]: Calculate the sum of the square of a data array. :param data: A numpy array or masked array of data, typically representing the differences (errors) in elevation or - another quantity. + another quantity. :return: sum square """ + # Promote integer values before squaring so their original storage type cannot overflow + if np.issubdtype(data.dtype, np.integer): + data = data.astype(np.float64) if np.ma.isMaskedArray(data): return np.ma.sum(np.square(data)) else: @@ -91,9 +94,12 @@ def rmse(data: NDArrayNum) -> np.floating[Any]: Calculate the RMSE of a data array. :param data: A numpy array or masked array of data, typically representing the differences (errors) in elevation or - another quantity. + another quantity. :return: rmse """ + # Promote integer values before squaring so their original storage type cannot overflow + if np.issubdtype(data.dtype, np.integer): + data = data.astype(np.float64) if np.ma.isMaskedArray(data): return np.sqrt(np.ma.mean(np.square(data))) else: diff --git a/geoutils/stats/grouping.py b/geoutils/stats/grouping.py new file mode 100644 index 000000000..aad6a039b --- /dev/null +++ b/geoutils/stats/grouping.py @@ -0,0 +1,1919 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Group raster and point cloud values into bins, categories, or vector zones.""" + +from __future__ import annotations + +import copy +import math +import warnings +import weakref +from collections.abc import Hashable, Iterable, Iterator, Mapping, Sequence +from contextlib import ExitStack +from dataclasses import dataclass, replace +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Any, Literal, cast + +import geopandas as gpd +import numpy as np +import pandas as pd +import xarray as xr +from numpy.typing import NDArray + +from geoutils._dispatch import ( + _get_pointcloud_interface, + _get_raster_interface, + _is_pointcloud, + _is_raster, + _is_vector, + is_dask_array, + is_dask_dataframe, +) +from geoutils._misc import import_optional +from geoutils._typing import ArrayLike +from geoutils.multiproc.chunked import iter_chunk_slices +from geoutils.multiproc.cluster import _map_bounded +from geoutils.multiproc.readers import ( + _normalize_reader_mask, + _read_selected_values, + _read_values, + _reader_from_source, + _reader_from_vector, + _ValueReader, +) +from geoutils.raster.array import get_mask_from_array +from geoutils.sampling.stratified import _stratified_subsample_indices +from geoutils.sampling.support import ( + _aligned_pointcloud, + _as_array, + _normalize_mask_array, + _normalize_sampling_input, + _sample_vector_values, + _sampling_specification, + _values_at_support, +) +from geoutils.stats.reduction import _reduce_values +from geoutils.stats.selection import _sample_eligible_indices +from geoutils.vector.base import _as_geodataframe + +if TYPE_CHECKING: + from geoutils.interface.interpolation import InterpolationMethod + from geoutils.multiproc import MultiprocConfig + from geoutils.pointcloud.base import PointCloudBase + from geoutils.pointcloud.pointcloud import PointCloudLike + from geoutils.raster.base import RasterBase, RasterLike + from geoutils.stats.reduction import _Statistics + from geoutils.vector.base import VectorLike + + +__all__ = ["plot_grouped_stats"] + + +#################### +# 1/ DEFINE GROUPS +#################### + + +@dataclass(frozen=True) +class _GroupDefinition: + """Store ordered labels, histogram edges or equal-width bin count to define a _PreparedGrouper.""" + + groups: pd.Index | None = None + edges: NDArray[Any] | None = None + bin_count: int | None = None + + +@dataclass(frozen=True) +class _PreparedGrouper: + """Store grouping values with their definition (see above), including already encoded vector categories.""" + + values: Any + definition: _GroupDefinition + encoded: bool = False + + +def _resolve_group_definition( + name: str, + *, + values: Any = None, + bins: int | Iterable[float] | pd.IntervalIndex | None = None, + categories: Iterable[Hashable] | None = None, +) -> _GroupDefinition: + """Resolve bins or categories before inferring categories from the original value dtype.""" + + # Keep explicit numeric declarations independent of categorical or boolean storage + if bins is not None: + if isinstance(bins, (int, np.integer, np.bool_)): + if isinstance(bins, (bool, np.bool_)) or bins < 1: + raise ValueError(f"The bin count for {name!r} must be a positive integer.") + return _GroupDefinition(bin_count=int(bins)) + if isinstance(bins, pd.IntervalIndex): + intervals = bins.rename(name) + if intervals.empty or not intervals.is_non_overlapping_monotonic: + raise ValueError(f"Intervals for {name!r} must be non-empty, ordered and non-overlapping.") + if not all(np.isfinite(interval.left) and np.isfinite(interval.right) for interval in intervals): + raise ValueError(f"Intervals for {name!r} must have finite bounds.") + return _GroupDefinition(groups=intervals) + + # Treat a numeric sequence as histogram edges and include the final upper edge + edges = np.asarray(list(bins), dtype=float) + if edges.ndim != 1 or len(edges) < 2 or not np.all(np.isfinite(edges)) or not np.all(np.diff(edges) > 0): + raise ValueError(f"Bin edges for {name!r} must be finite and strictly increasing.") + return _GroupDefinition(groups=pd.IntervalIndex.from_breaks(edges, closed="left", name=name), edges=edges) + + # Use declared categories, Pandas categories, or the two boolean values + dtype = getattr(values, "dtype", None) + if categories is None and isinstance(dtype, pd.CategoricalDtype): + if not is_dask_dataframe(values) or cast(Any, values).cat.known: + categories = dtype.categories + if categories is None and pd.api.types.is_bool_dtype(dtype): + categories = (False, True) + if categories is None: + raise ValueError(f"Grouper {name!r} requires an entry in ``bins`` or ``categories``.") + category_index = pd.Index(list(categories), tupleize_cols=False) + if category_index.empty or category_index.has_duplicates or category_index.hasnans: + raise ValueError(f"Categories for {name!r} must be non-empty, unique and finite.") + + # Keep the declaration order in the result, including categories absent from the input + level = pd.CategoricalIndex(category_index, categories=category_index, ordered=True, name=name) + return _GroupDefinition(groups=level) + + +def _validate_group_declarations( + by: Mapping[str, Any], + bins: Mapping[str, int | Iterable[float] | pd.IntervalIndex] | None, + categories: Mapping[str, Iterable[Hashable]] | None, +) -> dict[str, _GroupDefinition]: + """Validate grouper names and explicit declarations before spatial placement or array conversion.""" + + # Check bin and category declarations before reading grouping values + if not by: + raise ValueError("Argument ``by`` must contain at least one named grouper.") + if any(not isinstance(name, str) or not name for name in by): + raise ValueError("Grouper names must be non-empty strings.") + bins = {} if bins is None else bins + categories = {} if categories is None else categories + + # Require one unambiguous interpretation for every declared grouper + unknown = (set(bins) | set(categories)).difference(by) + if unknown: + raise ValueError(f"Bin or category declarations do not match ``by``: {sorted(unknown)!r}.") + overlap = set(bins).intersection(categories) + if overlap: + raise ValueError(f"A grouper cannot define both ``bins`` and ``categories``: {sorted(overlap)!r}.") + return { + name: _resolve_group_definition(name, bins=bins.get(name), categories=categories.get(name)) + for name in by + if name in bins or name in categories + } + + +def _prepare_grouper(values: Any, definition: _GroupDefinition) -> _PreparedGrouper: + """Convert one grouper to array values while keeping its categories or bins.""" + + # Numeric bins read category values, while category groups can use exact integer category codes + if isinstance(getattr(values, "dtype", None), pd.CategoricalDtype): + if isinstance(definition.groups, pd.CategoricalIndex): + categories = definition.groups.categories + if isinstance(values, (pd.Categorical, pd.CategoricalIndex)): + return _PreparedGrouper(values.set_categories(categories).codes, definition, encoded=True) + categorical = values.cat.set_categories(categories) + codes = _as_array(categorical.cat.codes) + return _PreparedGrouper(codes, definition, encoded=True) + values = values.astype(float) + + # Unwrap Xarray and dataframe containers without changing NumPy masks or Dask storage + if isinstance(values, xr.DataArray): + values = values.data + elif is_dask_dataframe(values) or isinstance(values, (pd.Series, pd.Index)): + values = _as_array(values) + return _PreparedGrouper(values, definition) + + +##################################### +# 2/ NORMALIZE GROUPED INPUTS +##################################### + + +def _normalize_grouped_inputs_dask_eager( + values: Mapping[str, Any], + by: Mapping[str, _PreparedGrouper], + mask: Any | None, + mp_config: MultiprocConfig | None, +) -> tuple[dict[str, Any], dict[str, _PreparedGrouper], Any | None, tuple[int, ...], bool, Any | None]: + """Prepare NumPy or Dask inputs with matching shapes and chunks.""" + + # Unwrap value containers once, keeping native masks and Dask arrays + named_values = {name: _as_array(array) for name, array in values.items()} + first_value = next(iter(named_values.values())) + shape = tuple(first_value.shape) + if not shape: + raise ValueError("Argument ``values`` must contain at least one dimension.") + mask = _normalize_mask_array(mask, shape) + + # Use one Dask layout when any value, grouping variable or mask is lazy + all_inputs = [*named_values.values(), *(grouper.values for grouper in by.values()), mask] + raw_inputs = [value.data if isinstance(value, xr.DataArray) else value for value in all_inputs] + lazy_inputs = [value for value in raw_inputs if is_dask_array(value)] + use_dask = bool(lazy_inputs) + chunks = lazy_inputs[0].reshape(shape).chunks if use_dask else None + if use_dask and mp_config is not None: + raise ValueError("Dask inputs cannot be combined with Multiprocessing grouped statistics.") + if use_dask: + import_optional("dask") + import dask.array as da + + # Check every selected value against the shared shape and array type + arrays: dict[str, Any] = {} + for name, raw_values in named_values.items(): + array = da.asarray(raw_values) if use_dask else np.asanyarray(raw_values) + if tuple(array.shape) != shape: + raise ValueError(f"Value {name!r} must match the shape of the other selected values.") + if not np.issubdtype(array.dtype, np.number): + raise TypeError(f"Value {name!r} must contain numeric data.") + arrays[name] = array.reshape(shape) + if use_dask: + arrays[name] = arrays[name].rechunk(chunks) + + # Put every grouper on the same shape and Dask chunk layout as the selected values + groupers: dict[str, _PreparedGrouper] = {} + for name, prepared in by.items(): + raw_values = prepared.values + group_values: Any = raw_values if is_dask_array(raw_values) else np.asanyarray(raw_values) + flattened_category = prepared.encoded and group_values.ndim == 1 and group_values.size == math.prod(shape) + if tuple(group_values.shape) != shape and not flattened_category: + raise ValueError(f"Grouper {name!r} must contain one value per input location.") + group_values = group_values.reshape(shape) + if use_dask: + group_values = ( + group_values.rechunk(chunks) + if is_dask_array(group_values) + else da.from_array(group_values, chunks=chunks) + ) + if not isinstance(prepared.definition.groups, pd.CategoricalIndex) and not pd.api.types.is_numeric_dtype( + group_values.dtype + ): + raise TypeError(f"Continuous grouper {name!r} must contain numeric values.") + groupers[name] = _PreparedGrouper(group_values, prepared.definition, prepared.encoded) + return arrays, groupers, mask, shape, use_dask, chunks + + +def _normalize_grouped_inputs_mp( + values: Mapping[str, Any], + by: Mapping[str, _PreparedGrouper], + mask: Any | None, + mp_config: MultiprocConfig, +) -> tuple[ + dict[str, Any], + dict[str, _PreparedGrouper], + Any | None, + tuple[int, ...], + list[tuple[slice, ...]], +]: + """Validate file-backed inputs and define the tiles multiprocessing workers will read.""" + + # Preserve readers for worker access while unwrapping only values already held in memory + arrays = {name: value if isinstance(value, _ValueReader) else _as_array(value) for name, value in values.items()} + shape = tuple(next(iter(arrays.values())).shape) + if not shape: + raise ValueError("Argument ``values`` must contain at least one dimension.") + mask = _normalize_reader_mask(mask, shape) + + # Reject mixed Dask and multiprocessing execution before starting file reads + raw_groupers = [grouper.values for grouper in by.values()] + if any(is_dask_array(value) for value in [*arrays.values(), *raw_groupers, mask]): + raise ValueError("Dask inputs cannot be combined with Multiprocessing grouped statistics.") + if any(tuple(value.shape) != shape for value in arrays.values()): + raise ValueError("Selected values must have matching shapes.") + for name, value in arrays.items(): + if not np.issubdtype(value.dtype, np.number): + raise TypeError(f"Value {name!r} must contain numeric data.") + + # Check grouper shapes and types while keeping readers available for later worker passes + groupers: dict[str, _PreparedGrouper] = {} + for name, prepared in by.items(): + group_values = prepared.values + if not isinstance(group_values, _ValueReader): + group_values = _as_array(group_values) + if prepared.encoded and group_values.ndim == 1 and group_values.size == math.prod(shape): + group_values = group_values.reshape(shape) + if tuple(group_values.shape) != shape: + raise ValueError(f"Grouper {name!r} must match the shape of the selected values.") + if not isinstance(prepared.definition.groups, pd.CategoricalIndex) and not pd.api.types.is_numeric_dtype( + group_values.dtype + ): + raise TypeError(f"Continuous grouper {name!r} must contain numeric values.") + groupers[name] = _PreparedGrouper(group_values, prepared.definition, prepared.encoded) + + # Reuse one ordered set of slices for every multiprocessing pass + tiles = list(iter_chunk_slices(shape, mp_config.chunks)) + return arrays, groupers, mask, shape, tiles + + +###################################### +# 3/ DERIVE BIN EDGES FROM COUNT +###################################### + + +def _group_definition_from_limits(name: str, bin_count: int, lower: float, upper: float) -> _GroupDefinition: + """Build equal-width histogram bins from min/mas in the group, including constant inputs.""" + + # Give constant values a nonzero bin range so the edges stay strictly increasing + if lower == upper: + half_width = 0.5 * abs(float(lower)) if lower != 0 else 0.5 + lower, upper = lower - half_width, upper + half_width + edges = np.linspace(float(lower), float(upper), bin_count + 1) + return _GroupDefinition(groups=pd.IntervalIndex.from_breaks(edges, closed="left", name=name), edges=edges) + + +def _derive_bin_edges_from_count_dask_eager( + by: Mapping[str, _PreparedGrouper], + *, + mask: Any | None, + use_dask: bool, + chunks: Any | None, +) -> dict[str, _PreparedGrouper]: + """ + Derive equal-width bin edges for eager or Dask groupers declared by bin count. + + Only required for binning defined by a count requiring min/max knowledge of the variable, otherwise this function + skips the grouping variable. + + :param use_dask: Whether to derive limits lazily with Dask or eagerly with NumPy. + :param chunks: Common Dask chunk layout, or None for eager inputs. + :returns: Prepared groupers with bin counts replaced by their derived interval labels and edges. + """ + + # Use the common Dask layout when a lazy input requires calculating its range + if use_dask: + import_optional("dask") + import dask.array as da + + user_eligible = None if mask is None else da.asarray(mask).rechunk(chunks) if use_dask else mask + resolved: dict[str, _PreparedGrouper] = {} + for name, prepared in by.items(): + definition = prepared.definition + values = prepared.values + + # Find the common finite range only for automatically generated equal-width bins + if definition.bin_count is not None: + if values.size == 0: + raise ValueError(f"Grouper {name!r} has no finite values inside ``mask``.") + if use_dask: + import dask + + finite = da.ma.filled(da.isfinite(values), False) + if user_eligible is not None: + finite &= user_eligible + lower, upper, finite_count = dask.compute( + da.min(da.where(finite, values, np.inf)), + da.max(da.where(finite, values, -np.inf)), + finite.sum(), + ) + else: + finite = np.isfinite(np.ma.getdata(values)) & ~np.ma.getmaskarray(values) + if user_eligible is not None: + finite &= user_eligible + finite_values = np.asarray(values)[np.asarray(finite)] + finite_count = finite_values.size + lower = np.min(finite_values) if finite_count else np.nan + upper = np.max(finite_values) if finite_count else np.nan + if not finite_count: + raise ValueError(f"Grouper {name!r} has no finite values inside ``mask``.") + definition = _group_definition_from_limits(name, definition.bin_count, lower, upper) + resolved[name] = _PreparedGrouper(values, definition, prepared.encoded) + return resolved + + +def _wrapper_group_limits_block_mp(values: Any, mask: Any) -> tuple[float, float]: + """Find finite grouper bounds after the user mask in one multiprocessing worker block.""" + + array = _as_array(_read_values(values)) + eligible = ~get_mask_from_array(array).reshape(array.shape) + if mask is not None: + eligible &= _normalize_mask_array(_read_values(mask), array.shape) + finite = np.ma.getdata(array)[eligible] + if finite.size == 0: + return np.inf, -np.inf + return float(np.min(finite)), float(np.max(finite)) + + +def _derive_bin_edges_from_count_mp( + by: Mapping[str, _PreparedGrouper], + mask: Any | None, + tiles: Sequence[tuple[slice, ...]], + mp_config: MultiprocConfig, +) -> dict[str, _PreparedGrouper]: + """ + Derive equal-width bin edges for multiprocessing groupers declared by bin count. + + Same as eager/Dask logic. + Only required for binning defined by a count requiring min/max knowledge of the variable, otherwise this function + skips the grouping variable. + + _group_limits_mp() reads the finite minimum and maximum inside the user mask from each worker block. + """ + + resolved = {} + for name, prepared in by.items(): + definition = prepared.definition + group_values = prepared.values + + # If the bin count was used as a definition, we find the min/max to define the bin edges, otherwise we skip + if definition.bin_count is not None: + group_arguments = [] + for tile in tiles: + block = group_values.block(tile) if isinstance(group_values, _ValueReader) else group_values[tile] + block_mask = ( + mask.block(tile) if isinstance(mask, _ValueReader) else None if mask is None else mask[tile] + ) + group_arguments.append((block, block_mask)) + + # We start at infinity and recursively accumulate min/max to get the global group min/max + lower, upper = np.inf, -np.inf + for _, (block_lower, block_upper) in _map_bounded( + mp_config.cluster, _wrapper_group_limits_block_mp, group_arguments + ): + lower, upper = min(lower, block_lower), max(upper, block_upper) + if not np.isfinite(lower): + raise ValueError(f"Grouper {name!r} has no finite values inside ``mask``.") + + # Define bin edges + definition = _group_definition_from_limits(name, definition.bin_count, lower, upper) + resolved[name] = _PreparedGrouper(group_values, definition, prepared.encoded) + return resolved + + +################################# +# 4/ ASSIGN GROUP IDs +################################# + + +def _encode_grouper(values: NDArray[Any], *, groups: pd.Index, edges: NDArray[Any] | None = None) -> NDArray[Any]: + """ + Replace categories or intervals with their integer group IDs in one array block. + + Keep the input shape and mark missing or undeclared values with -1. + Numeric edges follow histogram conventions (left close, right open), users can pass an IntervalIndex instead to + have it explicitly defined. + + :param values: One block of a single grouping variable passed to _assign_group_ids_dask_eager(). + :param groups: Ordered category or interval labels. + :param edges: Explicit numeric bin boundaries, or None to match the labels or intervals in groups directly. + + :returns: Integer group IDs with the input shape and -1 for excluded values. + """ + + # Follow histogram edge rules and include values equal to the final upper edge + array = np.asarray(np.ma.getdata(values)) + missing = np.ma.getmaskarray(values) + if edges is not None: + codes = np.searchsorted(edges, array, side="right") - 1 + codes[array == edges[-1]] = len(edges) - 2 + invalid = ~np.isfinite(array) | (codes < 0) | (codes >= len(edges) - 1) + return np.where(invalid | missing, -1, codes).astype(np.int64) + + # Let Pandas apply explicit interval boundaries or match category labels + if isinstance(groups, pd.IntervalIndex): + codes = groups.get_indexer(array.ravel()) + else: + categories = groups.categories if isinstance(groups, pd.CategoricalIndex) else groups + codes = pd.Categorical(array.ravel(), categories=categories, ordered=True).codes + codes = np.asarray(codes, dtype=np.int64).reshape(array.shape) + return np.where(missing, -1, codes) + + +def _group_layout(by: Mapping[str, _PreparedGrouper]) -> tuple[list[pd.Index], int]: + """Return the ordered labels and total group count shared by every calculation backend.""" + + levels = [] + for grouper in by.values(): + assert grouper.definition.groups is not None + levels.append(grouper.definition.groups) + total_groups = math.prod(len(level) for level in levels) + if total_groups > np.iinfo(np.int64).max: + raise ValueError("The product of group counts exceeds the supported integer range.") + return levels, total_groups + + +@dataclass(frozen=True) +class _GroupAssignmentMP: + """Keep multiprocessing group IDs and their temporary storage available to the parent calculation.""" + + group_ids: Any + levels: list[pd.Index] + total_groups: int + observed_ids: set[int] + storage: ExitStack + directory: str + + +def _assign_group_ids_dask_eager( + by: Mapping[str, _PreparedGrouper], + *, + mask: Any | None, + shape: tuple[int, ...], + use_dask: bool, + chunks: Any | None, +) -> tuple[Any, list[pd.Index], int]: + """ + Assign combined group IDs to eager or Dask arrays. + + _encode_grouper() gives an ID to each variable's categories or intervals, following order of ``by``. + A location receives -1 if the user mask excludes it or any grouper is missing or undeclared. + + :param by: Groupers returned by _derive_bin_edges_from_count_dask_eager(). + :param shape: Common shape of the values, grouping variables and user mask. + :param use_dask: Whether to assign group IDs lazily with Dask or eagerly with NumPy. + :param chunks: Common Dask chunk layout, or None for eager inputs. + :returns: Combined group IDs with the input shape, ordered label indexes for each grouper, and the total number + of declared combinations. Negative IDs mark excluded locations. + """ + + # Start with the user mask because it applies to every grouping variable + if use_dask: + import_optional("dask") + import dask.array as da + if mask is None: + eligible = da.ones(shape, chunks=chunks, dtype=bool) if use_dask else np.ones(shape, dtype=bool) + else: + eligible = da.asarray(mask).rechunk(chunks) if use_dask else mask + + # Assign group IDs within each grouping variable in the order of "by" + encoded: list[Any] = [] + levels, total_groups = _group_layout(by) + for prepared, level in zip(by.values(), levels): + definition = prepared.definition + values = prepared.values.reshape(shape) + + # Number the intervals in each array block, use Pandas only for custom open or closed sides + if prepared.encoded: + codes = da.where(da.isfinite(values), values, -1) if use_dask else np.where(np.isfinite(values), values, -1) + codes = codes.astype(np.int64) + elif use_dask: + codes = values.map_blocks(_encode_grouper, groups=level, edges=definition.edges, dtype=np.int64) + else: + codes = _encode_grouper(values, groups=level, edges=definition.edges) + encoded.append(codes) + eligible = eligible & (codes >= 0) + + # Combine the separate IDs into one compact group ID per input location + group_ids = da.zeros(shape, chunks=chunks, dtype=np.int64) if use_dask else np.zeros(shape, dtype=np.int64) + for codes, level in zip(encoded, levels): + # Reserve one consecutive range for each earlier combination before adding this variable's number + group_ids = group_ids * len(level) + codes + group_ids = da.where(eligible, group_ids, -1) if use_dask else np.where(eligible, group_ids, -1) + + # Use the smallest signed integer type that can hold all groups plus the missing marker + for dtype in (np.int8, np.int16, np.int32, np.int64): + if total_groups - 1 <= np.iinfo(dtype).max: + group_ids = group_ids.astype(dtype) + break + return group_ids, levels, total_groups + + +def _wrapper_assign_group_ids_block_mp(groupers: Mapping[str, Any], mask: Any, shape: tuple[int, ...]) -> Any: + """Read and assign resolved group IDs within one multiprocessing worker block.""" + + prepared = { + name: _PreparedGrouper(_read_values(grouper.values), grouper.definition, grouper.encoded) + for name, grouper in groupers.items() + } + block_mask = _normalize_mask_array(_read_values(mask), shape) + ids, _, _ = _assign_group_ids_dask_eager(prepared, mask=block_mask, shape=shape, use_dask=False, chunks=None) + return ids + + +def _assign_group_ids_mp( + by: Mapping[str, _PreparedGrouper], + mask: Any | None, + shape: tuple[int, ...], + tiles: Sequence[tuple[slice, ...]], + mp_config: MultiprocConfig, +) -> _GroupAssignmentMP: + """ + Assign group IDs from unloaded readers into temporary storage with multiprocessing. + + Follows the same logic as the eager/Dask implementation above. + + _assign_group_ids_block_mp() handles one block at a time, we keep the combined numbers on disk so sampling, + reduction and optional returned masks can reuse them without loading all grouping arrays at once. + """ + + # Create the disk-backed array before workers begin returning group ID blocks + levels, total_groups = _group_layout(by) + storage = ExitStack() + try: + directory = storage.enter_context(TemporaryDirectory(prefix="geoutils-stats-")) + ids = ( + np.memmap(f"{directory}/groups.dat", mode="w+", dtype=np.int64, shape=shape) + if math.prod(shape) + else np.empty(shape, dtype=np.int64) + ) + if isinstance(ids, np.memmap): + storage.callback(ids._mmap.close) + observed_ids = set() + + # Send matching grouping and mask blocks to each worker in their original order + id_arguments = [] + for tile in tiles: + block_groupers = { + name: replace( + grouper, + values=( + grouper.values.block(tile) if isinstance(grouper.values, _ValueReader) else grouper.values[tile] + ), + ) + for name, grouper in by.items() + } + block_mask = mask.block(tile) if isinstance(mask, _ValueReader) else None if mask is None else mask[tile] + block_shape = tuple(part.stop - part.start for part in tile) + id_arguments.append((block_groupers, block_mask, block_shape)) + + for index, block_ids in _map_bounded(mp_config.cluster, _wrapper_assign_group_ids_block_mp, id_arguments): + ids[tiles[index]] = block_ids + observed_ids.update(int(value) for value in np.unique(block_ids) if value >= 0) + return _GroupAssignmentMP( + group_ids=ids, + levels=levels, + total_groups=total_groups, + observed_ids=observed_ids, + storage=storage, + directory=directory, + ) + except Exception: + storage.close() + raise + + +################################ +# 5/ SAMPLE GROUPED VALUES +################################ + + +def _sample_grouped_values_dask_eager( + values: Mapping[str, Any], + group_ids: Any, + *, + subsample: int | float, + subsample_per_group: bool, + random_state: int | np.random.Generator | None, + subsampling_strategy: Literal["sequential", "topk"], + use_dask: bool, +) -> tuple[dict[str, Any], NDArray[Any]]: + """Sample matching group IDs and selected values from eager or Dask arrays.""" + + # If we subsample per group, we call stratified subsampling + if subsample_per_group: + flat_indices = _stratified_subsample_indices( + group_ids, + subsample=subsample, + random_state=random_state, + strategy=subsampling_strategy, + ) + # Sample directly from all eligible locations when sampling is not per group + else: + flat_indices = _sample_eligible_indices( + group_ids >= 0, + subsample=subsample, + random_state=random_state, + strategy=subsampling_strategy, + ) + + # Read the same locations for every input value, computing Dask inputs together + selected = [ + group_ids.reshape(-1)[flat_indices], + *(array.reshape(-1)[flat_indices] for array in values.values()), + ] + if use_dask: + import dask + + selected = list(dask.compute(*selected)) + selected_ids = np.asarray(selected[0]) + selected_values = {name: np.asanyarray(array) for name, array in zip(values, selected[1:])} + return selected_values, selected_ids + + +def _sample_grouped_values_mp( + values: Mapping[str, Any], + group_ids: Any, + *, + shape: tuple[int, ...], + subsample: int | float, + subsample_per_group: bool, + random_state: int | np.random.Generator | None, + subsampling_strategy: Literal["sequential", "topk"], + mp_config: MultiprocConfig, + assignment: _GroupAssignmentMP | None, +) -> tuple[dict[str, Any], NDArray[Any]]: + """ + Sample values with matching group IDs using multiprocessing backend. + + It mirrors the eager/Dask logic above. + + NumPy values are indexed directly. For unloaded readers, we use _read_selected_values() so + workers read only the selected value locations. + """ + + # If we subsample per group, we call stratified subsampling + if subsample_per_group: + sample_ids = group_ids + indexes = _stratified_subsample_indices( + sample_ids, + subsample, + random_state=random_state, + strategy=subsampling_strategy, + mp_config=mp_config, + ) + # Sample directly when all group IDs are in memory + elif assignment is None: + indexes = _sample_eligible_indices( + group_ids >= 0, + subsample=subsample, + random_state=random_state, + strategy=subsampling_strategy, + ) + # Scan group IDs per chunk + else: + sample_ids = group_ids + if math.prod(shape): + sample_ids = np.memmap(f"{assignment.directory}/eligible.dat", mode="w+", dtype=np.int8, shape=shape) + assignment.storage.callback(sample_ids._mmap.close) + for tile in iter_chunk_slices(shape, mp_config.chunks): + sample_ids[tile] = np.where(group_ids[tile] >= 0, 0, -1) + indexes = _stratified_subsample_indices( + sample_ids, + subsample, + random_state=random_state, + strategy=subsampling_strategy, + mp_config=mp_config, + ) + + # Read the same locations for every input value, leaving unloaded inputs to worker readers + selected_ids = np.asarray(group_ids.reshape(-1)[indexes]) + selected_values = { + name: ( + _read_selected_values(array, indexes, mp_config) + if isinstance(array, _ValueReader) + else array.reshape(-1)[indexes] + ) + for name, array in values.items() + } + return selected_values, selected_ids + + +######################################## +# 6/ BUILD GROUP LABELS AND MASKS +######################################## + + +class _GroupMasks(Mapping[Hashable, Any]): + """ + Create boolean masks on the result grid or points from one shared group ID array. + + This functionality is a helper for the ``return_masks`` option, to get a boolean mask of each group + on the common support. + We build each mask only when its result key is requested, avoiding one stored array per group. + + :param group_ids: Complete combined memberships from _assign_group_ids_dask_eager() or + _assign_group_ids_mp(), before subsampling. + :param key_ids: Mapping from result row labels to the integer group IDs represented by those rows. + :param shape: Original value-array shape restored when a group mask is requested. + :param support: Raster or point cloud defining the result locations, or None to return array masks. + """ + + def __init__( + self, + group_ids: Any, + key_ids: Mapping[Hashable, int], + shape: tuple[int, ...], + support: RasterBase | PointCloudBase | None, + ) -> None: + """Keep the shared memberships and ordered result keys without creating individual masks.""" + + # Keep one group ID array and follow the result table's row order + self._group_ids = group_ids + self._key_ids = dict(key_ids) + self._shape = shape + self._support = support + + def __getitem__(self, key: Hashable) -> RasterLike | PointCloudLike | ArrayLike: + """Return one group's boolean mask as an array or an object on the original spatial support.""" + + # 1/ Select the requested group full membership + + # Raise the usual dictionary error before creating a mask + if key not in self._key_ids: + raise KeyError(key) + mask = (self._group_ids == self._key_ids[key]).reshape(self._shape) + + # Return a plain boolean array when the input had no raster or point locations + if self._support is None: + return mask + + # 2/ Build a raster mask with the same grid and coordinate system as the input + if _is_raster(self._support): + raster = cast("RasterBase", self._support) + return raster.from_array( + data=mask, + transform=raster.transform, + crs=raster.crs, + nodata=None, + area_or_point=raster.area_or_point, + tags=raster.tags.copy(), + ) + + # 3/ Keep point geometry and other columns while replacing the selected value with the mask + if _is_pointcloud(self._support): + from geoutils.pointcloud.dataframe import ( + _assign_point_values, + _build_pointcloud_output, + _get_dataframe_attrs, + ) + + pointcloud = cast("PointCloudBase", self._support) + if not pointcloud._is_pd and not pointcloud.is_loaded: + # Read attributes for the requested mask without loading the caller's file-backed point object + pointcloud = copy.copy(pointcloud) + pointcloud.load(columns="all") + dataframe = pointcloud.ds + column = pointcloud.data_column + if column is None: + # Add a boolean column when the point values were stored in geometry Z coordinates + column = "group_mask" + while column in dataframe.columns: + # Avoid replacing an existing point attribute with the new membership column + column = f"_{column}" + + # Preserve every coordinate and assign masks by row position, including single-point inputs + output = _assign_point_values(dataframe, {column: mask}) + with warnings.catch_warnings(): + # The mask intentionally selects its boolean column instead of the preserved geometry Z + warnings.filterwarnings("ignore", message="Overriding 3D points with with data column") + return _build_pointcloud_output( + output, + data_column=column, + as_dataframe=pointcloud._ACCESSOR_OUTPUT, + attrs=_get_dataframe_attrs(dataframe), + preserve_locations=True, + ) + raise TypeError("Group masks require array, raster or point cloud support.") + + def __iter__(self) -> Iterator[Hashable]: + """Iterate over group labels in the result table row order.""" + + return iter(self._key_ids) + + def __len__(self) -> int: + """Return the number of groups represented in the result table.""" + + return len(self._key_ids) + + +def _group_index( + levels: Sequence[pd.Index], names: Sequence[str], group_numbers: Sequence[int] | NDArray[Any] +) -> pd.Index: + """ + Build the ordered Pandas row index from combined group IDs. + + We use a categorical or interval index for one grouper, and a MultiIndex for several groupers. + + :param levels: Ordered label indexes returned by the eager, Dask or multiprocessing group assignment. + :param names: Grouper names in the same order as levels. + :param group_numbers: Combined integer group IDs to include, in the requested row order. + :returns: A Pandas index naming each requested group combination. + """ + + # Split each combined number back into one number per grouping variable + level_codes = [] + remainders = np.asarray(group_numbers, dtype=np.int64) + for level in reversed(levels): + remainders, codes = np.divmod(remainders, len(level)) + level_codes.append(codes) + level_codes.reverse() + + # Use a direct interval or category index when there is only one grouping variable + if len(levels) == 1: + selected = levels[0].take(level_codes[0]) + index = selected.rename(names[0]) + else: + index = pd.MultiIndex(levels=list(levels), codes=level_codes, names=list(names), verify_integrity=False) + return index + + +def _format_grouped_stats( + table: pd.DataFrame, + statistics: _Statistics, + resolved_strategy: str, + *, + value_names: Sequence[str], + by_names: Sequence[str], + levels: Sequence[pd.Index], + total_groups: int, + full_group_ids: Any, + shape: tuple[int, ...], + support: RasterBase | PointCloudBase | None, + observed: bool, + subsample: int | float, + subsample_per_group: bool, + subsampling_strategy: Literal["sequential", "topk"], + return_masks: bool, + group_numbers: NDArray[Any] | None = None, +) -> pd.DataFrame | tuple[pd.DataFrame, Mapping[Hashable, RasterLike | PointCloudLike | ArrayLike]]: + """Format final reduced values, adding back empty categories with NaNs.""" + + # Keep every group that was present before the optional subsampling + if group_numbers is not None: + group_numbers = np.asarray(group_numbers) + elif not observed: + group_numbers = np.arange(total_groups) + elif subsample != 1: + if is_dask_array(full_group_ids): + import dask.array as da + + group_numbers = da.unique(full_group_ids).compute() + else: + group_numbers = np.unique(full_group_ids) + group_numbers = group_numbers[group_numbers >= 0] + else: + group_numbers = table.index.to_numpy() + table = table.reindex(group_numbers) + + # Fill empty counts/groups with zero and leave estimates as NaN + columns = pd.MultiIndex.from_product([list(value_names), statistics.output_names], names=["value", "statistic"]) + table.columns = columns + for name in value_names: + table[(name, "count")] = table[(name, "count")].fillna(0).astype(np.int64) + for statistic, alias in zip(statistics.names, statistics.aliases): + if alias in {"validcount", "totalcount"}: + table[(name, statistic)] = table[(name, statistic)].fillna(0).astype(np.int64) + table.index = _group_index(levels, by_names, group_numbers) + + # Record sampling user inputs in the table attributes + table.attrs["grouped_stats"] = { + "observed": observed, + "subsample": subsample, + "subsample_per_group": subsample_per_group, + "strategy": resolved_strategy, + "subsampling_strategy": subsampling_strategy, + "mask_membership": "groupers", + } + + # Delay each boolean group mask until the caller reads it from the returned mapping + if return_masks: + key_ids = dict(zip(table.index, map(int, group_numbers))) + masks = _GroupMasks(full_group_ids, key_ids=key_ids, shape=shape, support=support) + return table, masks + return table + + +########################################## +# 7/ CALCULATE GROUPED STATISTICS +########################################## + + +def _group_values_for_flox( + by: Mapping[str, _PreparedGrouper], *, use_dask: bool +) -> tuple[list[Any], list[pd.Index | NDArray[Any]]]: + """Prepare grouping values and their declared order for the eager or Dask Flox backend.""" + + group_values = [] + expected_groups = [] + for prepared in by.values(): + values = prepared.values + groups = prepared.definition.groups + assert groups is not None + + # Keep ordinary category values for Flox to match, including their declared order + masked = np.ma.isMaskedArray(values._meta) if use_dask else np.ma.isMaskedArray(values) + if not prepared.encoded and not isinstance(groups, pd.IntervalIndex) and not masked: + group_values.append(values) + expected_groups.append(groups) + continue + + # Keep existing category codes, or encode bins whose boundary rules differ from Flox's native bins + if prepared.encoded: + codes = values.astype(np.int64) + elif use_dask: + codes = values.map_blocks( + _encode_grouper, + groups=groups, + edges=prepared.definition.edges, + dtype=np.int64, + ) + else: + codes = _encode_grouper(values, groups=groups, edges=prepared.definition.edges) + group_values.append(codes) + expected_groups.append(np.arange(len(groups), dtype=np.int64)) + return group_values, expected_groups + + +def _groupby_reduce_flox( + values: Any, + group_values: Sequence[Any], + expected_groups: Sequence[pd.Index | NDArray[Any]], + *, + function: str, + use_dask: bool, + fill_value: float | int, + dtype: Any, + finalize_kwargs: Mapping[str, Any] | None = None, +) -> Any: + """Run one eager or Dask Flox reduction over every grouping variable.""" + + flox = import_optional("flox", extra_name="flox") + expected: Any = expected_groups[0] if len(expected_groups) == 1 else tuple(expected_groups) + axes = tuple(range(-group_values[0].ndim, 0)) + return flox.groupby_reduce( + values, + *group_values, + func=function, + expected_groups=expected, + sort=False, + axis=axes, + fill_value=fill_value, + dtype=dtype, + method="map-reduce" if use_dask else None, + finalize_kwargs=finalize_kwargs, + )[0] + + +def _calculate_grouped_stats_flox( + values: Mapping[str, Any], + by: Mapping[str, _PreparedGrouper], + *, + statistics: _Statistics, + mask: Any | None, + full_group_ids: Any | None, + shape: tuple[int, ...], + support: RasterBase | PointCloudBase | None, + observed: bool, + subsample: int | float, + subsampling_strategy: Literal["sequential", "topk"], +) -> pd.DataFrame: + """ + Calculate grouped statistics with Flox from eager or Dask values and grouping variables. + + Flox matches the separate grouping values itself. GeoUtils only converts interval definitions to preserve their + boundary rules, applies the common mask and finite-data rules, and builds the established result dataframe. + """ + + # Check statistics that Flox cannot reproduce through its grouped reductions + unsupported = [ + name for name, alias in zip(statistics.names, statistics.aliases) if alias is None or alias == "nmad" + ] + if unsupported: + raise ValueError(f"The Flox backend does not support these statistics: {unsupported!r}.") + exact = {"median", "90thpercentile", "iqr", "le90"} + use_dask = any(is_dask_array(array) for array in values.values()) + if use_dask and any(alias in exact for alias in statistics.aliases): + raise ValueError("The Dask Flox backend does not support exact median or percentile statistics.") + + # Prepare Flox group inputs while keeping the categories and bins used for the final row index + group_values, expected_groups = _group_values_for_flox(by, use_dask=use_dask) + levels, total_groups = _group_layout(by) + if use_dask: + import_optional("dask") + import dask.array as da + + # Apply the user mask and each value's missing data independently before stacking the selected values + selected_values = [] + for array in values.values(): + if use_dask: + data = da.ma.getdata(array) + valid = ~da.ma.getmaskarray(array) & da.isfinite(data) + if mask is not None: + valid &= mask + selected_values.append(da.where(valid, data, np.nan)) + else: + data = np.ma.getdata(array) + valid = ~get_mask_from_array(array) + if mask is not None: + valid &= mask + selected_values.append(np.where(valid, data, np.nan)) + stacked = da.stack(selected_values) if use_dask else np.stack(selected_values) + + # Count all eligible group locations separately from each value's finite observations + first = next(iter(values.values())) + reduction_shape = tuple(first.shape) + if use_dask: + membership = da.ones(reduction_shape, chunks=first.chunks, dtype=float) + if mask is not None: + membership = da.where(mask, 1.0, np.nan) + else: + membership = np.ones(reduction_shape, dtype=float) + if mask is not None: + membership = np.where(mask, 1.0, np.nan) + scheduled = { + "membership": _groupby_reduce_flox( + membership, + group_values, + expected_groups, + function="count", + use_dask=use_dask, + fill_value=0, + dtype=np.int64, + ), + "count": _groupby_reduce_flox( + stacked, + group_values, + expected_groups, + function="count", + use_dask=use_dask, + fill_value=0, + dtype=np.int64, + ), + } + + # Ask Flox only for the reductions needed by the requested statistics + aliases = {alias for alias in statistics.aliases if alias is not None} + direct_functions = { + "mean": "nanmean", + "median": "nanmedian", + "max": "nanmax", + "min": "nanmin", + "sum": "nansum", + "std": "nanstd", + } + for alias, function in direct_functions.items(): + if alias in aliases: + scheduled[alias] = _groupby_reduce_flox( + stacked, + group_values, + expected_groups, + function=function, + use_dask=use_dask, + fill_value=np.nan, + dtype=np.float64, + finalize_kwargs={"ddof": 0} if alias == "std" else None, + ) + + # Reduce squared values for sum of squares and RMSE without assigning combined group IDs + squared = np.square(stacked) + if "sumofsquares" in aliases: + scheduled["sumofsquares"] = _groupby_reduce_flox( + squared, + group_values, + expected_groups, + function="nansum", + use_dask=use_dask, + fill_value=np.nan, + dtype=np.float64, + ) + if "rmse" in aliases: + scheduled["mean_square"] = _groupby_reduce_flox( + squared, + group_values, + expected_groups, + function="nanmean", + use_dask=use_dask, + fill_value=np.nan, + dtype=np.float64, + ) + + # Request the quantiles shared by the percentile, IQR and LE90 statistics + quantiles = {} + if "90thpercentile" in aliases: + quantiles["q90"] = 0.90 + if "iqr" in aliases: + quantiles.update(q75=0.75, q25=0.25) + if "le90" in aliases: + quantiles.update(q95=0.95, q05=0.05) + for name, quantile in quantiles.items(): + scheduled[name] = _groupby_reduce_flox( + stacked, + group_values, + expected_groups, + function="nanquantile", + use_dask=use_dask, + fill_value=np.nan, + dtype=np.float64, + finalize_kwargs={"q": quantile}, + ) + + # Compute all small Dask results together before constructing the Pandas output + if use_dask: + import dask + + keys = list(scheduled) + scheduled = dict(zip(keys, dask.compute(*(scheduled[key] for key in keys)))) + membership_count = np.asarray(scheduled["membership"], dtype=np.int64).reshape(total_groups) + finite_count = np.asarray(scheduled["count"], dtype=np.int64).reshape(len(values), total_groups) + + # Restore GeoUtils statistic names and count fields for each selected value + columns = {} + for value_index in range(len(values)): + columns[(value_index, "count")] = finite_count[value_index] + for name, statistic_alias in zip(statistics.names, statistics.aliases): + if statistic_alias == "validcount": + result = finite_count[value_index] + elif statistic_alias == "totalcount": + result = membership_count + elif statistic_alias == "percentagevalidpoints": + result = np.divide( + 100 * finite_count[value_index], + membership_count, + out=np.full(total_groups, np.nan), + where=membership_count > 0, + ) + elif statistic_alias == "rmse": + result = np.sqrt(np.asarray(scheduled["mean_square"])[value_index].reshape(total_groups)) + elif statistic_alias == "90thpercentile": + result = np.asarray(scheduled["q90"])[value_index].reshape(total_groups) + elif statistic_alias == "iqr": + result = np.asarray(scheduled["q75"])[value_index].reshape(total_groups) - np.asarray(scheduled["q25"])[ + value_index + ].reshape(total_groups) + elif statistic_alias == "le90": + result = np.asarray(scheduled["q95"])[value_index].reshape(total_groups) - np.asarray(scheduled["q05"])[ + value_index + ].reshape(total_groups) + else: + assert statistic_alias is not None + result = np.asarray(scheduled[statistic_alias])[value_index].reshape(total_groups) + columns[(value_index, name)] = result + table = pd.DataFrame(columns, index=np.arange(total_groups)) + table = table.loc[membership_count > 0] + + # Use the common formatter to restore declared group labels, empty groups and output metadata + return _format_grouped_stats( + table, + statistics, + "flox", + value_names=list(values), + by_names=list(by), + levels=levels, + total_groups=total_groups, + full_group_ids=np.empty(0, dtype=np.int64) if full_group_ids is None else full_group_ids, + shape=shape, + support=support, + observed=observed, + subsample=subsample, + subsample_per_group=False, + subsampling_strategy=subsampling_strategy, + return_masks=False, + ) + + +def _calculate_grouped_stats( + values: ArrayLike | Mapping[str, ArrayLike], + by: Mapping[str, _PreparedGrouper], + *, + statistics: _Statistics, + mask: Any | None, + subsample: int | float, + subsample_per_group: bool, + random_state: int | np.random.Generator | None, + strategy: Literal["auto", "dense", "sparse", "groupwise"], + backend: Literal["geoutils", "flox"], + subsampling_strategy: Literal["sequential", "topk"], + observed: bool, + return_masks: bool, + support: RasterBase | PointCloudBase | None, + mp_config: MultiprocConfig | None, +) -> pd.DataFrame | tuple[pd.DataFrame, Mapping[Hashable, RasterLike | PointCloudLike | ArrayLike]]: + """ + Calculate statistics for values with prepared group definitions and aligned grouping values. + + Both backends first normalize their array inputs and derive bin edges when only a bin count was given. The Flox + backend then lets Flox match group values and calculate the result. The GeoUtils backend assigns combined group + IDs, optionally samples them, and passes them to _reduce_values(). Both paths finish with _format_grouped_stats(). + """ + + # 1/ Normalize array inputs (eager vs Dask), chunk layout + named_values: dict[str, Any] = dict(values) if isinstance(values, Mapping) else {"value": values} + if not named_values or any(not isinstance(name, str) or not name for name in named_values): + raise ValueError("Argument ``values`` must contain at least one non-empty name.") + all_inputs = [*named_values.values(), *(grouper.values for grouper in by.values()), mask] + use_reader_mp = mp_config is not None and any(isinstance(value, _ValueReader) for value in all_inputs) + if use_reader_mp: + assert mp_config is not None + arrays, groupers, mask, shape, tiles = _normalize_grouped_inputs_mp(named_values, by, mask, mp_config) + use_dask, chunks = False, None + else: + arrays, groupers, mask, shape, use_dask, chunks = _normalize_grouped_inputs_dask_eager( + named_values, by, mask, mp_config + ) + tiles = None + + # 2/ Derive bin edges from each grouper min/max when only a bin count was passed (e.g., 20) for a grouper + if use_reader_mp: + assert mp_config is not None and tiles is not None + groupers = _derive_bin_edges_from_count_mp(groupers, mask, tiles, mp_config) + else: + groupers = _derive_bin_edges_from_count_dask_eager( + groupers, + mask=mask, + use_dask=use_dask, + chunks=chunks, + ) + + # Let Flox match the grouping values directly; combined group IDs are only needed for global subsampling + if backend == "flox": + full_group_ids = None + if subsample != 1: + full_group_ids, _, _ = _assign_group_ids_dask_eager( + groupers, + mask=mask, + shape=shape, + use_dask=use_dask, + chunks=chunks, + ) + sample_inputs = { + **{f"value_{index}": array for index, array in enumerate(arrays.values())}, + **{f"grouper_{index}": grouper.values for index, grouper in enumerate(groupers.values())}, + } + sampled, _ = _sample_grouped_values_dask_eager( + sample_inputs, + full_group_ids, + subsample=subsample, + subsample_per_group=False, + random_state=random_state, + subsampling_strategy=subsampling_strategy, + use_dask=use_dask, + ) + arrays = {name: sampled[f"value_{index}"] for index, name in enumerate(arrays)} + groupers = { + name: replace(grouper, values=sampled[f"grouper_{index}"]) + for index, (name, grouper) in enumerate(groupers.items()) + } + mask = None + return _calculate_grouped_stats_flox( + arrays, + groupers, + statistics=statistics, + mask=mask, + full_group_ids=full_group_ids, + shape=shape, + support=support, + observed=observed, + subsample=subsample, + subsampling_strategy=subsampling_strategy, + ) + + # 3/ Assign an integer ID to each group before sampling + assignment_mp = None + keep_mp_storage = False + try: + if use_reader_mp: + assert mp_config is not None and tiles is not None + assignment_mp = _assign_group_ids_mp(groupers, mask, shape, tiles, mp_config) + full_group_ids = assignment_mp.group_ids + levels = assignment_mp.levels + total_groups = assignment_mp.total_groups + else: + full_group_ids, levels, total_groups = _assign_group_ids_dask_eager( + groupers, + mask=mask, + shape=shape, + use_dask=use_dask, + chunks=chunks, + ) + + # 4/ Optionally subsample the common support, globally or per group + group_ids = full_group_ids + if subsample != 1: + if mp_config is not None: + arrays, group_ids = _sample_grouped_values_mp( + arrays, + full_group_ids, + shape=shape, + subsample=subsample, + subsample_per_group=subsample_per_group, + random_state=random_state, + subsampling_strategy=subsampling_strategy, + mp_config=mp_config, + assignment=assignment_mp, + ) + else: + arrays, group_ids = _sample_grouped_values_dask_eager( + arrays, + full_group_ids, + subsample=subsample, + subsample_per_group=subsample_per_group, + random_state=random_state, + subsampling_strategy=subsampling_strategy, + use_dask=use_dask, + ) + + # 5/ Reduce every grouped input for all the statistics defined by the user + table, resolved_strategy = _reduce_values( + list(arrays.values()), + statistics, + group_ids=group_ids, + total_groups=total_groups, + strategy=strategy, + mp_config=mp_config if subsample == 1 else None, + ) + result = _format_grouped_stats( + table, + statistics, + resolved_strategy, + value_names=list(arrays), + by_names=list(by), + levels=levels, + total_groups=total_groups, + full_group_ids=full_group_ids, + shape=shape, + support=support, + observed=observed, + subsample=subsample, + subsample_per_group=subsample_per_group, + subsampling_strategy=subsampling_strategy, + return_masks=return_masks, + group_numbers=( + np.array(sorted(assignment_mp.observed_ids), dtype=np.int64) + if observed and assignment_mp is not None + else None + ), + ) + + # Keep files available for masks returned after multiprocessing has finished + if return_masks and assignment_mp is not None: + weakref.finalize(result[1], assignment_mp.storage.close) + keep_mp_storage = True + return result + finally: + if assignment_mp is not None and not keep_mp_storage: + assignment_mp.storage.close() + + +######################################### +# 8/ SELECT SPATIAL GROUPING VALUES +######################################### + + +def _vector_group_values( + vector: VectorLike, + selector: str | None, + *, + support: RasterBase | PointCloudBase, + support_dataframe: gpd.GeoDataFrame | None, + name: str, + definition: _GroupDefinition | None, + mp_config: MultiprocConfig | None, +) -> _PreparedGrouper: + """ + Create grouping values on the common support from vector features or one selected attribute column. + + Without a selector (feature column), return a boolean inside/outside grouper. + With a feature column, sample integer category numbers directly, allowing nonnumeric labels on raster and + point support. + + The "support" argument follows _grouped_stats(), and "mp_config" follows stats(). + + :param vector: Vector or GeoDataFrame providing feature geometries and optional category attributes. + :param selector: Feature column containing category labels, or None to group by vector coverage. + :param support_dataframe: Point coordinates and attributes on the selected support, or None for raster support. + :param name: Name of this grouping variable in the result. + :param definition: Optional validated bin or category declaration; otherwise infer labels from the selected feature + column, excluding missing labels. + :returns: Grouping values or category codes on the selected support, together with their group definition. + """ + + # 1/ Read the vector and distinguish boolean inside/outside grouping from a feature attribute ID grouping + dataframe = _as_geodataframe(vector) + encoded = False + + # Treat a vector without a selected column as a grouping variable for inside and outside + if selector is None: + if definition is not None and not isinstance(definition.groups, pd.CategoricalIndex): + raise ValueError("Vector values require an explicit feature column.") + feature_values = np.ones(len(dataframe)) + definition = definition or _resolve_group_definition(name, categories=(False, True)) + + # 2/ Read feature labels once and keep the input order + else: + if selector not in dataframe.columns: + raise ValueError(f"Vector column {selector!r} does not exist.") + labels = dataframe[selector] + if definition is None: + if isinstance(labels.dtype, pd.CategoricalDtype): + definition = _resolve_group_definition(name, values=labels) + else: + category_values = pd.unique(labels[labels.notna()]) + if not len(category_values): + raise ValueError(f"Vector column {selector!r} has no categories.") + definition = _resolve_group_definition(name, categories=category_values) + if isinstance(definition.groups, pd.CategoricalIndex): + feature_values = pd.Categorical(labels, categories=definition.groups.categories, ordered=True).codes + encoded = True + else: + feature_values = _prepare_grouper(labels, definition).values + if not pd.api.types.is_numeric_dtype(feature_values.dtype): + raise TypeError("Selected vector values must be numeric.") + + # 3/ Use the same vector sampling as cosample() and keep Dask results lazy + if mp_config is not None: + values = _reader_from_vector(dataframe, feature_values, support, mp_config, coverage=selector is None) + if values is not None: + return _PreparedGrouper(values, definition, encoded=encoded) + if support_dataframe is None and _is_pointcloud(support): + support_dataframe = cast("PointCloudBase", support).ds + values = _sample_vector_values(dataframe, feature_values, support, support_dataframe, mp_config=mp_config) + if selector is None: + values = np.isfinite(values) + return _PreparedGrouper(values, definition, encoded=encoded) + + +def _select_groupers_at_support( + source: RasterLike | PointCloudLike | ArrayLike | Mapping[str, ArrayLike], + by: Mapping[str, Any], + *, + support: RasterBase | PointCloudBase | None, + definitions: Mapping[str, _GroupDefinition], + interpolation: InterpolationMethod, + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None, + stack: ExitStack | None = None, +) -> dict[str, _PreparedGrouper]: + """ + Reproject every grouping variable on the common spatial support. + + :param support: Raster or point cloud defining common support grid or point locations. None means the + source and grouping inputs are already aligned arrays without spatial support. + :param definitions: Validated bins or categories provided for grouping variables. + :param stack: Optional ExitStack for aligned temporary files from Multiprocessing chunked execution that must + remain available during reduction. + + :returns: Prepared grouping values on the selected support, in the input order. + """ + + # Use the locations already selected for values and masks by the shared statistics preparation + support_dataframe = None + groupers_at_support: dict[str, _PreparedGrouper] = {} + source_support = _get_raster_interface(source) + if source_support is None: + source_support = _get_pointcloud_interface(source) + + # Align each grouper to the values common grid, point order, or array shape + definition: _GroupDefinition | None + for name, specification in by.items(): + # Keep ordinary array groupers unchanged when the source has no spatial support + if support is None: + definition = definitions.get(name) or _resolve_group_definition(name, values=specification) + groupers_at_support[name] = _prepare_grouper(specification, definition) + continue + group_source, group_selector = _sampling_specification(source, specification) + + # Treat Xarrays without spatial coordinates as arrays + group_source = _normalize_sampling_input(group_source) + + # Detect type of grouping variable: raster, point cloud, or vector + is_raster = _is_raster(group_source) + is_pointcloud = _is_pointcloud(group_source) + is_vector = not is_raster and not is_pointcloud and _is_vector(group_source) + + # Native arrays already follow the source locations, so keep their masks and category metadata intact + if not (is_raster or is_pointcloud or is_vector) and source_support is support: + definition = definitions.get(name) or _resolve_group_definition(name, values=group_source) + prepared = _prepare_grouper(group_source, definition) + if _is_raster(support) and getattr(prepared.values, "ndim", 0) == 3 and prepared.values.shape[0] == 1: + prepared = _PreparedGrouper(prepared.values[0], definition, prepared.encoded) + groupers_at_support[name] = prepared + continue + + # Keep native file columns and matching raster bands unloaded for worker processes + if mp_config is not None: + definition = definitions.get(name) + if definition is None and is_raster and pd.api.types.is_bool_dtype(group_source.dtype): + definition = _resolve_group_definition(name, values=group_source) + categorical = definition is not None and isinstance(definition.groups, pd.CategoricalIndex) + file_values = _reader_from_source( + group_source, + group_selector, + support, + mp_config, + align=align, + interpolation="nearest" if categorical else interpolation, + stack=stack, + ) + if file_values is not None: + definition = definition or _resolve_group_definition(name, values=file_values) + groupers_at_support[name] = _PreparedGrouper(file_values, definition) + continue + + # Use feature labels for categorical zones + if is_vector: + if group_selector is not None and not isinstance(group_selector, str): + raise TypeError(f"Vector selector for {name!r} must be a column name.") + groupers_at_support[name] = _vector_group_values( + group_source, + group_selector, + support=support, + support_dataframe=support_dataframe, + name=name, + definition=definitions.get(name), + mp_config=mp_config, + ) + continue + + # Read point coordinates only when placing values actually requires them + if support_dataframe is None and _is_pointcloud(support): + support_dataframe = cast("PointCloudBase", support).ds + + # Preserve point cloud category labels and order before converting the column to an array + metadata_values = group_source + pointcloud = _get_pointcloud_interface(group_source) + if pointcloud is not None: + column = pointcloud.data_column if group_selector is None else group_selector + if column is not None: + if not isinstance(column, str) or column not in pointcloud.ds.columns: + raise ValueError(f"Point column {column!r} selected for {name!r} does not exist.") + metadata_values = pointcloud.ds[column] + definition = definitions.get(name) or _resolve_group_definition(name, values=metadata_values) + + if ( + pointcloud is not None + and _is_pointcloud(support) + and isinstance(getattr(metadata_values, "dtype", None), pd.CategoricalDtype) + ): + # Validate point locations before reading native category codes, avoiding conversion to object arrays + with ExitStack() as temporary_files: + intermediate = temporary_files.enter_context(mp_config.temporary()) if mp_config is not None else None + aligned = _aligned_pointcloud(pointcloud, support, name, align, mp_config=intermediate) + groupers_at_support[name] = _prepare_grouper(aligned.ds[column], definition) + continue + + # Always use nearest neighbor for categorical raster labels, otherwise interpolation from the user + grouped_values = _values_at_support( + group_source, + group_selector, + input_support=source, + support=support, + support_dataframe=support_dataframe, + name=name, + interpolation="nearest" if isinstance(definition.groups, pd.CategoricalIndex) else interpolation, + align=align, + mp_config=mp_config, + ) + groupers_at_support[name] = _prepare_grouper(grouped_values, definition) + + return groupers_at_support + + +##################### +# 9/ PARENT FUNCTION +##################### + + +def _grouped_stats( + source: RasterLike | PointCloudLike | ArrayLike | Mapping[str, ArrayLike], + by: Mapping[str, Any], + *, + values: ArrayLike | Mapping[str, ArrayLike], + support: RasterBase | PointCloudBase | None, + statistics: _Statistics, + mask: Any | None, + subsample: int | float, + subsample_per_group: bool, + random_state: int | np.random.Generator | None, + strategy: Literal["auto", "dense", "sparse", "groupwise"], + backend: Literal["geoutils", "flox"], + subsampling_strategy: Literal["sequential", "topk"], + interpolation: InterpolationMethod, + align: Literal["raise", "reproject"], + observed: bool, + return_masks: bool, + mp_config: MultiprocConfig | None, + definitions: Mapping[str, _GroupDefinition], + stack: ExitStack | None = None, +) -> pd.DataFrame | tuple[pd.DataFrame, Mapping[Hashable, RasterLike | PointCloudLike | ArrayLike]]: + """ + Calculate grouped statistics after selecting values (source input) and groupers (grouping variables from ``by``) on + one common spatial support. + + Logic: + _select_groupers_at_support() places groupers on the same common support (grid or points), then + _calculate_grouped_stats() derives missing bin edges, applies optional subsampling, calls the selected grouped + reduction backend, and formats the result table. + + See stats() for all argument descriptions. + + :returns: Grouped dataframe, and optionally group masks. + """ + + # This is the geospatial-specific step, in addition to the _select step already done before in stats(): + # We reproject every grouper variable on the common support (raster or point) + groupers_at_support = _select_groupers_at_support( + source, + by, + support=support, + definitions=definitions, + interpolation=interpolation, + align=align, + mp_config=mp_config, + stack=stack, + ) + + # Apply optional subsampling and calculate the result table with the selected grouped reduction backend + return _calculate_grouped_stats( + values, + groupers_at_support, + statistics=statistics, + mask=mask, + subsample=subsample, + subsample_per_group=subsample_per_group, + random_state=random_state, + strategy=strategy, + backend=backend, + subsampling_strategy=subsampling_strategy, + observed=observed, + return_masks=return_masks, + support=support, + mp_config=mp_config, + ) + + +########################### +# 10/ GROUPED STAT PLOTTING +########################### + + +def _plot_axis(index: pd.Index) -> tuple[NDArray[Any], NDArray[Any], list[str] | None]: + """ + Return plot edges, centers, and optional text labels for one grouping variable. + + Continuous binned intervals use their numeric bin bounds. + Other intervals (categorical, zonal, or non-continuous bins) use equally spaced positions with labels so that + gaps or nonnumeric values do not distort the plotted group order. + + :param index: Ordered labels for one grouping variable from the table passed to plot_grouped_stats(). + :returns: Bin edges, bin centers and optional tick labels, with one center per group. + """ + + # Draw continuous numeric intervals with their true widths + if isinstance(index, pd.IntervalIndex) and len(index) > 0: + adjacent = len(index) == 1 or np.all(np.asarray(index.right[:-1]) == np.asarray(index.left[1:])) + if adjacent: + edges = np.asarray([index[0].left, *index.right], dtype=float) + return edges, np.asarray(index.mid, dtype=float), None + + # Give categories and separated intervals equal widths + edges = np.arange(len(index) + 1, dtype=float) + centers = edges[:-1] + 0.5 + return edges, centers, [str(value) for value in index] + + +def plot_grouped_stats( + table: pd.DataFrame, + *, + value: str | None = None, + statistic: str = "nmad", + min_count: int = 0, + cmap: Any = "viridis", + vmin: float | None = None, + vmax: float | None = None, + ax: Any | None = None, + savefig_fname: str | None = None, +) -> Mapping[str, Any]: + """ + Plot grouped statistics for one or two grouping variables with their sample counts shown as side histograms. + + One grouping variable produces a curve below its histogram counts. + Two grouping variables produce a colored grid, with count histograms above and to the right. + + Passing ``min_count`` hides estimates without removing observations from the count panels. + + :param table: Grouped dataframe returned by stats() or an object stats() method. + :param value: Selected value column. It may be omitted when the table contains one value. + :param statistic: Statistic column to display. + :param min_count: Hide statistic cells with fewer finite observations. + :param cmap: Matplotlib colormap used for a two-dimensional statistic grid. + :param vmin: Lower color limit for a two-dimensional statistic grid. + :param vmax: Upper color limit for a two-dimensional statistic grid. + :param ax: Optional Matplotlib axes whose area is divided into the plot panels. + :param savefig_fname: Optional path used to save the completed figure. + :returns: Mapping naming the Matplotlib axes created for each panel. + """ + + # 1/ Validate the selected value, statistic, and group dimensions + + # Import Matplotlib only when the caller requests a plot + matplotlib = import_optional("matplotlib") + import matplotlib.pyplot as plt + + # Require the named column levels produced by stats() before looking up the selected values + if not isinstance(table, pd.DataFrame) or not isinstance(table.columns, pd.MultiIndex): + raise TypeError("Argument ``table`` must be a grouped statistics dataframe with MultiIndex columns.") + if list(table.columns.names) != ["value", "statistic"]: + raise ValueError("Argument ``table`` must have 'value' and 'statistic' column levels.") + + # Infer the value column only when exactly one is available + available_values = list(dict.fromkeys(table.columns.get_level_values("value"))) + if value is None: + if len(available_values) != 1: + raise ValueError("Argument ``value`` must be selected when ``table`` contains multiple values.") + value = available_values[0] + + # Every plot needs counts to show how many finite observations support each estimate + if (value, statistic) not in table.columns or (value, "count") not in table.columns: + raise ValueError(f"Value {value!r} must contain both {statistic!r} and 'count' statistics.") + if table.index.nlevels not in {1, 2}: + raise ValueError("plot_grouped_stats supports one or two group dimensions.") + if min_count < 0: + raise ValueError("Argument ``min_count`` cannot be negative.") + + # 2/ Prepare a common plotting area for the statistic and its count panels + + # Use the supplied axes as the plot area or create a new figure + if ax is None: + figure = plt.figure(figsize=(7, 6)) + frame = figure.add_axes((0.1, 0.1, 0.8, 0.8)) + elif isinstance(ax, matplotlib.axes.Axes): + frame = ax + figure = ax.figure + else: + raise TypeError("Argument ``ax`` must be a Matplotlib Axes or None.") + frame.set_axis_off() + + # 3/ Draw one-dimensional curves or a two-dimensional grid with matching counts + + # Draw counts above the statistic for one grouping variable + if table.index.nlevels == 1: + count_axis = frame.inset_axes((0.0, 0.72, 1.0, 0.28)) + statistic_axis = frame.inset_axes((0.0, 0.0, 1.0, 0.64)) + edges, centers, labels = _plot_axis(table.index) + + # Hide unsupported estimates while keeping their counts visible in the panel above + counts = table[(value, "count")].to_numpy(dtype=float) + values = table[(value, statistic)].where(table[(value, "count")] >= min_count).to_numpy(dtype=float) + + # Match count-bar widths to the numeric intervals or category positions + count_axis.bar(edges[:-1], counts, width=np.diff(edges), align="edge", color="0.7", edgecolor="white") + count_axis.set_xlim(edges[0], edges[-1]) + count_axis.set_ylabel("Count") + count_axis.tick_params(axis="x", labelbottom=False) + + # Align the statistic with the same group centers and display labels for nonnumeric positions + statistic_axis.plot(centers, values, marker="o") + statistic_axis.set(xlim=(edges[0], edges[-1]), xlabel=table.index.name, ylabel=statistic) + if labels is not None: + statistic_axis.set_xticks(centers, labels, rotation=45, ha="right") + axes = {"count": count_axis, "statistic": statistic_axis} + + # Draw a two-dimensional statistic grid with totals for each row and column + else: + statistic_axis = frame.inset_axes((0.0, 0.0, 0.68, 0.66)) + count_x_axis = frame.inset_axes((0.0, 0.72, 0.68, 0.28)) + count_y_axis = frame.inset_axes((0.74, 0.0, 0.26, 0.66)) + level_x, level_y = table.index.levels + edges_x, centers_x, labels_x = _plot_axis(level_x) + edges_y, centers_y, labels_y = _plot_axis(level_y) + + # Add every declared group combination so groups with no data appear as gaps + full_index = pd.MultiIndex.from_product([level_x, level_y], names=table.index.names) + counts = table[(value, "count")].reindex(full_index).unstack(level=1) + plotted = table[(value, statistic)].where(table[(value, "count")] >= min_count) + plotted = plotted.reindex(full_index).unstack(level=1) + mesh = statistic_axis.pcolormesh( + edges_x, + edges_y, + plotted.to_numpy(dtype=float).T, + cmap=cmap, + vmin=vmin, + vmax=vmax, + shading="flat", + ) + + # Draw row and column totals with the same widths as the statistic grid + counts_x = counts.sum(axis=1, skipna=True).to_numpy(dtype=float) + counts_y = counts.sum(axis=0, skipna=True).to_numpy(dtype=float) + count_x_axis.bar(edges_x[:-1], counts_x, width=np.diff(edges_x), align="edge", color="0.7", edgecolor="white") + count_y_axis.barh(edges_y[:-1], counts_y, height=np.diff(edges_y), align="edge", color="0.7", edgecolor="white") + count_x_axis.set(xlim=(edges_x[0], edges_x[-1]), ylabel="Count") + count_y_axis.set(ylim=(edges_y[0], edges_y[-1]), xlabel="Count") + count_x_axis.tick_params(axis="x", labelbottom=False) + count_y_axis.tick_params(axis="y", labelleft=False) + + # Give the statistic grid the same extent as both count panels + statistic_axis.set( + xlim=(edges_x[0], edges_x[-1]), + ylim=(edges_y[0], edges_y[-1]), + xlabel=table.index.names[0], + ylabel=table.index.names[1], + ) + if labels_x is not None: + statistic_axis.set_xticks(centers_x, labels_x, rotation=45, ha="right") + if labels_y is not None: + statistic_axis.set_yticks(centers_y, labels_y) + + # Expose all created axes so callers can adapt the finished figure + colorbar = figure.colorbar(mesh, ax=statistic_axis, label=statistic) + axes = { + "count_x": count_x_axis, + "count_y": count_y_axis, + "statistic": statistic_axis, + "colorbar": colorbar.ax, + } + + # 4/ Save after every panel and label has been added + if savefig_fname is not None: + figure.savefig(savefig_fname, bbox_inches="tight") + return axes diff --git a/geoutils/stats/reduction.py b/geoutils/stats/reduction.py new file mode 100644 index 000000000..f4f7d3a7c --- /dev/null +++ b/geoutils/stats/reduction.py @@ -0,0 +1,1270 @@ +# Copyright (c) 2025 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reduce arrays to descriptive statistics, globally or by group.""" + +from __future__ import annotations + +import math +import warnings +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, replace +from functools import partial +from typing import TYPE_CHECKING, Any + +import numpy as np +import pandas as pd +from numpy.typing import NDArray +from scipy.stats import iqr +from scipy.stats.mstats import mquantiles + +from geoutils._dispatch import is_dask_array +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.multiproc.readers import _ValueReader +from geoutils.raster.array import get_mask_from_array +from geoutils.stats.estimators import linear_error, nmad, rmse, sum_square + +if TYPE_CHECKING: + from geoutils.multiproc import MultiprocConfig + from geoutils.stats.selection import _SelectionCounts + +######################### +# 1/ HELPERS +######################### + +# We list the statistics that can be merged efficiently from aggregates computed in separate blocks +# For instance: a mean can be simply a sum divided by count, both that can be aggregated over many group-chunk +# intersections without requiring to load a whole group at once +# However, some statistics cannot aggregate like this and require complete groups, those are: median, percentiles, IQR, +# LE90, NMAD, and custom functions +MERGEABLE_STATISTICS = { + "validcount", + "totalcount", + "percentagevalidpoints", + "sum", + "mean", + "std", + "min", + "max", + "sumofsquares", + "rmse", +} + +# List aliases for statistics, total counts, and inlier counts +_STATS_ALIAS_CALLABLE = { + "mean": "Mean", + "median": "Median", + "max": "Max", + "min": "Min", + "sum": "Sum", + "sumofsquares": "Sum of squares", + "90thpercentile": "90th percentile", + "iqr": "IQR", + "le90": "LE90", + "nmad": "NMAD", + "rmse": "RMSE", + "std": "Standard deviation", +} +_STATS_ALIAS_COUNTS = { + "validcount": "Valid count", + "totalcount": "Total count", + "percentagevalidpoints": "Percentage valid points", +} +_STATS_ALIAS_GEN = _STATS_ALIAS_CALLABLE | _STATS_ALIAS_COUNTS +_STATS_ALIAS_MASK = { + "validinliercount": "Valid inlier count", + "totalinliercount": "Total inlier count", + "percentagevalidinlierpoints": "Percentage valid inlier points", + "percentageinlierpoints": "Percentage inlier points", +} +# List of all statistics when user input is "all" +_STATS_ALIAS_ALL = _STATS_ALIAS_GEN | _STATS_ALIAS_MASK + +# Allow synonyms +_SYNONYMS = { + "maximum": "max", + "minimum": "min", + "sum2": "sumofsquares", + "90percentile": "90thpercentile", + "rms": "rmse", + "standarddeviation": "std", +} + +# List of defaults statistics used when no user input is passed +_STATS_LIST_MIN = [ + "min", + "max", + "mean", + "std", + "validcount", + "totalcount", + "percentagevalidpoints", +] + + +@dataclass(frozen=True) +class _Statistics: + """ + Store requested statistics, output names and validated internal aliases. + + :param requested: Statistic names or callables. + :param names: Matching output column names. + :param aliases: Matching internal statistic names (None for callables). + """ + + requested: list[str | Callable[[Any], Any]] + names: list[str] + aliases: list[str | None] + grouped: bool = True + single: bool = False + + @property + def output_names(self) -> list[str]: + """Return requested labels with the mandatory count for grouped output.""" + + return ["count", *self.names] if self.grouped else self.names + + +def _get_stat_common_alias(stat_name: str) -> str | None: + """Return the internal statistic name for a user alias.""" + + # Ignore case, spaces and underscores in names such as "standard deviation" + normalized_name = "".join(stat_name.lower().replace("_", "").split()) + normalized_name = _SYNONYMS.get(normalized_name, normalized_name) + return normalized_name if normalized_name in _STATS_ALIAS_ALL else None + + +def _normalize_statistics( + statistics: str | Callable[[Any], Any] | Iterable[str | Callable[[Any], Any]] | None, + *, + grouped: bool = True, + masked: bool = False, +) -> _Statistics: + """ + Validate statistics and choose their result names and internal aliases. + + See stats() for description of ``statistics`` argument. + + :param grouped: Include the grouped count column, reject unknown names and omit summary-only mask counts. + :param masked: Include mask-specific counts in the complete global selection requested by ``"all"``. + :returns: A _Statistics object separating requested statistics, output names and internal aliases. + """ + + # 1/ Expand the user request into the statistics to calculate + # Keep the established default selection and add inlier counts only when a mask was used + default_names = statistics is None or isinstance(statistics, str) and statistics == "all" + if statistics is None: + requested: list[str | Callable[[Any], Any]] = list(_STATS_LIST_MIN) + elif isinstance(statistics, str) and statistics == "all": + requested = list(_STATS_ALIAS_GEN) + if masked and not grouped: + requested += list(_STATS_ALIAS_MASK) + elif isinstance(statistics, str) or callable(statistics): + requested = [statistics] + else: + requested = list(statistics) + + # Grouped results always include a leading count column + if grouped: + requested = [statistic for statistic in requested if statistic != "count"] + if default_names: + requested = [statistic for statistic in requested if statistic != "validcount"] + if requested == ["all"]: + requested = [*_STATS_ALIAS_CALLABLE, "totalcount", "percentagevalidpoints"] + elif "all" in requested: + raise ValueError("Statistic 'all' cannot be combined with other statistics.") + + # 2/ Separate output labels from the internal names used by the estimators + names = [] + for statistic in requested: + if isinstance(statistic, str): + names.append(statistic) + elif callable(statistic): + function = statistic + while isinstance(function, partial): + function = function.func + names.append(getattr(function, "__name__", type(function).__name__)) + else: + raise TypeError("Argument ``statistics`` must contain names or callable functions.") + if default_names and not grouped: + names = [_STATS_ALIAS_ALL[name] for name in names] + aliases = [_get_stat_common_alias(statistic) if isinstance(statistic, str) else None for statistic in requested] + + # 3/ Check user names before calculation, raising warning for unknown ones and skipping them + if len(set(names)) != len(names): + raise ValueError("Statistic names must be unique.") + if grouped and "count" in names: + raise ValueError("Statistic name 'count' is reserved for the mandatory grouped count.") + allowed = _STATS_ALIAS_GEN if grouped else _STATS_ALIAS_ALL + unknown = [ + statistic for statistic, alias in zip(requested, aliases) if isinstance(statistic, str) and alias not in allowed + ] + if unknown and grouped: + raise ValueError(f"Unknown statistic names: {unknown!r}.") + for statistic in unknown: + warnings.warn(f"Statistic name {statistic} is not recognized", category=UserWarning) + single = callable(statistics) or isinstance(statistics, str) and statistics != "all" + return _Statistics(requested=requested, names=names, aliases=aliases, grouped=grouped, single=single) + + +def _resolve_strategy( + aliases: Sequence[str | None], strategy: str, total_groups: int, chunked: bool +) -> tuple[str, bool]: + """ + Choose chunk strategy for the requested statistics. + + See _reduce_values() for ``strategy`` and ``total_groups`` arguments, and ``aliases`` are from _Statistics. + + :param chunked: Whether Dask or multiprocessing will calculate separate input blocks. + + :returns: The strategy and whether every requested statistic can be merged from block summaries. + """ + + # Mask-related counts use the same count reductions as ordinary validity statistics + mergeable = all(alias in MERGEABLE_STATISTICS or alias in _STATS_ALIAS_MASK for alias in aliases) + if strategy == "auto": + if not mergeable: + strategy = "groupwise" + else: + strategy = "dense" if total_groups <= 4096 else "sparse" + elif strategy not in {"dense", "sparse", "groupwise"}: + raise ValueError("Argument ``strategy`` must be 'auto', 'dense', 'sparse' or 'groupwise'.") + if chunked and not mergeable and strategy != "groupwise": + raise ValueError("Exact quantiles and custom statistics require ``strategy``='groupwise' or 'auto'.") + return strategy, mergeable + + +######################### +# 2/ EAGER REDUCTION +######################### + + +def _reduce_complete_values_eager(data: NDArrayNum, aliases: set[str]) -> tuple[dict[str, Any], int]: + """ + Evaluate whole-array statistics with NumPy's masked or NaN-aware functions. + + We preserve input type when calling estimators: masked arrays use NumPy Masked functions, while + ordinary arrays use functions that ignore NaNs. + + We also return the validity count, total count and percentages. + + See _reduce_values() for ``data`` argument description. + + :param aliases: Internal names of the numerical statistics to calculate, excluding count statistics. + + :returns: Estimates keyed by their internal names and the input validity count. + """ + + # Keep the existing distinction between masked values and NaNs in ordinary arrays + masked = np.ma.isMaskedArray(data) + final_count = int(np.count_nonzero(~np.ma.getmaskarray(data) if masked else np.isfinite(data))) + module = np.ma if masked else np + prefix = "" if masked else "nan" + functions = {name: getattr(module, prefix + name) for name in ("mean", "median", "max", "min", "sum", "std")} + + # Use the existing estimators and matching percentile definitions for masked arrays + functions.update( + sumofsquares=sum_square, + **{ + "90thpercentile": ( + (lambda array: mquantiles(array, prob=0.9, alphap=1, betap=1)[0]) + if masked + else partial(np.nanpercentile, q=90) + ) + }, + le90=partial(linear_error, interval=90), + iqr=partial(iqr, nan_policy="omit"), + nmad=nmad, + rmse=rmse, + ) + + # Report empty inputs consistently and avoid calling estimators that require at least one valid value + if final_count == 0: + warnings.warn("Empty raster, returns NaN for all stats", category=UserWarning) + result = {alias: functions[alias](data) if final_count else np.nan for alias in aliases} + return result, final_count + + +def _extrema_result(values: Sequence[Any] | NDArray[Any]) -> Any: + """Keep integer extrema beyond floating-point precision exact, including missing group results.""" + + present = [value for value in values if not pd.isna(value)] + if present and all(isinstance(value, (int, np.integer)) for value in present): + if any(abs(int(value)) > 2**53 for value in present): + dtype = "UInt64" if any(value > np.iinfo(np.int64).max for value in present) else "Int64" + return pd.array(values, dtype=dtype) + return np.asarray(values, dtype=float) + + +def _aggregate_eager( + values: Sequence[NDArray[Any]], group_ids: NDArray[Any] | None, statistics: _Statistics +) -> pd.DataFrame: + """ + Calculate every statistic from complete in-memory groups after one stable sort. + + Locations are sorted once by group ID, placing all values from each group in one slice. Statistics that need the + complete group, such as quantiles and user functions, can then use those slices without filtering the full array + separately for every group. The stable sort keeps the original order within each group. With no group IDs, + use the complete input directly and avoid allocating or sorting an artificial group array. + + See _reduce_values() for ``values`` and ``group_ids`` descriptions, and see _aggregate_chunked() for + ``statistics``. + + :returns: A table with integer group rows and (value position, statistic name) columns. + """ + + # 1/ Locate complete groups while keeping the original order within each group + + # Use one direct slice for a full-array reduction, or sort grouped locations once by their group ID + global_reduction = group_ids is None + if global_reduction: + order: slice | NDArray[Any] = slice(None) + labels = np.array([0], dtype=np.int64) + starts = np.array([0], dtype=np.int64) + sizes = np.array([values[0].size], dtype=np.int64) + else: + ids = np.asarray(group_ids).ravel() + selected = np.flatnonzero(ids >= 0) + order = selected[np.argsort(ids[selected], kind="stable")] + labels, starts, sizes = np.unique(ids[order], return_index=True, return_counts=True) + + # 2/ Calculate the requested estimators and counts for each selected value + + # Reuse the validated names for every group instead of parsing the same request again + aliases = {alias for alias in statistics.aliases if alias in _STATS_ALIAS_CALLABLE} + columns: dict[tuple[int, str], Any] = {} + for value_index, array in enumerate(values): + # Share the group ordering but handle missing observations independently for each selected value + ordered = np.asanyarray(array) if global_reduction else np.asanyarray(array).ravel()[order] + if not global_reduction: + invalid = get_mask_from_array(ordered).reshape(ordered.shape) + if np.any(invalid): + if np.issubdtype(ordered.dtype, np.integer): + ordered = np.ma.array(np.ma.getdata(ordered), mask=invalid) + else: + ordered = np.where(invalid, np.nan, np.ma.getdata(ordered)) + results: dict[str, list[Any]] = {"count": [], **{name: [] for name in statistics.names}} + + # Pass complete group values to statistics such as median and user functions + for start, size in zip(starts, sizes): + group = ordered if global_reduction else ordered[start : start + size] + count = int(np.count_nonzero(~get_mask_from_array(group))) + if global_reduction: + computed, count = _reduce_complete_values_eager(group, aliases) + else: + computed, _ = _reduce_complete_values_eager(group, aliases) if count and aliases else ({}, count) + results["count"].append(count) + + # Distinguish finite values from all group locations when constructing counts and percentages + for statistic, name, alias in zip(statistics.requested, statistics.names, statistics.aliases): + result: Any + if alias == "validcount": + result = count + elif alias == "totalcount": + result = int(size) + elif alias == "percentagevalidpoints": + result = 100 * count / size if size else np.nan + elif callable(statistic): + # Give user functions the complete group, with missing observations still present + callable_values = ( + np.where(np.ma.getmaskarray(group), np.nan, np.ma.getdata(group)) + if np.ma.isMaskedArray(group) and np.any(np.ma.getmaskarray(group)) + else group + ) + result = statistic(callable_values) if count else np.nan + else: + result = computed.get(alias or "", np.nan) + results[name].append(result) + + # Keep counts as integers and preserve extrema beyond floating-point precision when needed + extrema_names = {name for name, alias in zip(statistics.names, statistics.aliases) if alias in {"min", "max"}} + count_names = { + name for name, alias in zip(statistics.names, statistics.aliases) if alias in {"validcount", "totalcount"} + } + for name, result_values in results.items(): + if name in extrema_names: + columns[(value_index, name)] = _extrema_result(result_values) + else: + dtype = np.int64 if name == "count" or name in count_names else float + columns[(value_index, name)] = np.asarray(result_values, dtype=dtype) + + # 3/ Assemble the shared table format used by eager and chunked reductions + return pd.DataFrame(columns, index=labels) + + +############################################# +# 3/ CHUNKED REDUCTION (DASK AND MULTIPROCESSING) +############################################# + + +def _reader_block_requires_eager(block: Any) -> bool: + """Identify values for which global NumPy semantics differ from finite grouped reductions.""" + + from geoutils.multiproc.readers import _read_values + + values = _read_values(block) + nonfinite = ~np.isfinite(values) if np.ma.isMaskedArray(values) else np.isinf(values) + return bool(np.ma.filled(np.any(nonfinite), False)) + + +def _reader_requires_eager(array: _ValueReader, mp_config: MultiprocConfig) -> bool: + """Check exceptional nonfinite values in bounded reads before choosing global estimators.""" + + from geoutils.multiproc.chunked import iter_chunk_slices + from geoutils.multiproc.cluster import _map_bounded + + arguments = ((array.block(tile),) for tile in iter_chunk_slices(array.shape, mp_config.chunks)) + return any(result for _, result in _map_bounded(mp_config.cluster, _reader_block_requires_eager, arguments)) + + +def _statistics_dask(data: Any, aliases: set[str]) -> tuple[dict[str, Any], Any]: + """ + Build only the requested Dask reductions while returning lazy scalar results. + + Reuse the global median when both median and NMAD are requested. Counts and estimates remain lazy for the caller + to compute together; exact quantiles may still need to bring values from several chunks into one worker. + + This optional native Dask calculation is kept as a reference for tests and benchmarks. Public stats() uses + _reduce_values() so global and grouped reductions share the GeoUtils implementation. + + Data follows _reduce_values(), and aliases follow _reduce_complete_values_eager(). + + :returns: Lazy estimates keyed by their internal names and a lazy finite-value count. + """ + + import_optional("dask") + import dask.array as da + + # Dask requires explicit axes for full-array quantiles. This remains lazy, but exact global quantiles can still be + # memory-intensive at execution time because Dask has to combine data across chunks + axes = tuple(range(data.ndim)) + median = da.nanquantile(data, 0.50, axis=axes) if aliases & {"median", "nmad"} else None + functions: dict[str, Callable[[], Any]] = { + name: partial(getattr(da, "nan" + name), data) for name in ("mean", "max", "min", "sum", "std") + } + squared_values = data.astype(np.float64) if np.issubdtype(data.dtype, np.integer) else data + + # Defer graph construction for quantiles and squared values until their statistics are requested + functions.update( + median=lambda: median, + sumofsquares=lambda: da.nansum(da.square(squared_values)), + **{"90thpercentile": lambda: da.nanquantile(data, 0.90, axis=axes)}, + le90=lambda: da.nanquantile(data, 0.95, axis=axes) - da.nanquantile(data, 0.05, axis=axes), + iqr=lambda: da.nanquantile(data, 0.75, axis=axes) - da.nanquantile(data, 0.25, axis=axes), + nmad=lambda: 1.4826 * da.nanquantile(da.fabs(data - median), 0.50, axis=axes), + rmse=lambda: da.sqrt(da.nanmean(da.square(squared_values))), + ) + finite = da.isfinite(data) + if np.ma.isMaskedArray(data._meta): + finite = da.ma.filled(finite, False) + final_count = finite.sum() + results = {alias: functions[alias]() for alias in aliases} + for alias in aliases & {"sum", "sumofsquares"}: + results[alias] = da.where(final_count > 0, results[alias], np.nan) + return results, final_count + + +def _reduce_block( + values: Sequence[NDArray[Any]], + group_ids: NDArray[Any] | None, + total_groups: int, + dense: bool, + statistics: set[str], +) -> tuple[NDArray[Any], NDArray[Any], dict[str, NDArray[Any]]]: + """ + Summarize groups in one array block so summaries from many blocks can be combined. + + A dense summary reserves one position for every declared group; a sparse summary stores only groups present in the + block. Each selected value has its own finite count, while group sizes include missing selected values. For standard + deviation, the summary stores the mean and sum of squared deviations needed by the pairwise merge. Stored quantities + have shape (number of selected values, number of stored groups), and only requested quantities are allocated. + + Values and group_ids follow _reduce_values(), restricted to this block. Total_groups still includes every + declared group across the full input. + + :param dense: Reserve every declared group when True; otherwise store only groups present in the block. + :param statistics: Internal names of the requested mergeable statistics, including any count statistics. + :returns: Stored integer group labels, their total location counts, and named summary arrays with one row per + selected value and one column per stored group. These are the summaries consumed by _merge_blocks(). + """ + + from geoutils.multiproc.readers import _read_values + + # Read file descriptors inside the worker while leaving supplied arrays unchanged + values = [_read_values(value) for value in values] + + # 1/ Number the groups represented in this block + + # Use one group for a full-array block, or store the requested dense or sparse grouped labels + if group_ids is None: + eligible = None + labels = np.array([0], dtype=np.int64) + codes = None + size = np.array([values[0].size], dtype=np.int64) + else: + ids = np.asarray(group_ids).ravel() + eligible = ids >= 0 + + # Dense summaries use the declared numbers directly; sparse summaries give present groups consecutive slots + if dense: + labels = np.arange(total_groups, dtype=np.int64) + codes = ids[eligible].astype(np.int64) + else: + labels, codes = np.unique(ids[eligible], return_inverse=True) + size = np.bincount(codes, minlength=len(labels)) + shape = (len(values), len(labels)) + state = {"count": np.zeros(shape, dtype=np.int64)} + + # 2/ Create only the summary arrays needed by the requested statistics + needs_mean = bool(statistics & {"mean", "std"}) + if needs_mean: + state["mean"] = np.zeros(shape, dtype=float) + if "std" in statistics: + state["m2"] = np.zeros(shape, dtype=float) + summaries = statistics & {"sum", "sumofsquares", "min", "max"} + if "rmse" in statistics: + summaries.add("sumofsquares") + + # Empty groups start with neutral values so missing observations cannot change a sum, minimum, or maximum + for name in summaries: + fill = np.inf if name == "min" else -np.inf if name == "max" else 0.0 + integer_extrema = name in {"min", "max"} and any( + np.issubdtype(array.dtype, np.integer) and array.dtype.itemsize >= 8 for array in values + ) + state[name] = np.full(shape, fill, dtype=object if integer_extrema else float) + + # 3/ Reuse group IDs while handling missing data separately for each value array + for index, array in enumerate(values): + data = np.ma.getdata(array).ravel() + finite = ~get_mask_from_array(array).ravel() + if eligible is not None: + data = data[eligible] + finite = finite[eligible] + valid_codes = None if codes is None else codes[finite] + + # Count finite values in each group, using a direct count for whole-array summaries + count = ( + np.array([np.count_nonzero(finite)], dtype=np.int64) + if valid_codes is None + else np.bincount(valid_codes, minlength=len(labels)) + ) + state["count"][index] = count + + # Counts need only the finite-value mask; prepare numerical values only for requested estimates + if len(state) == 1: + continue + valid_data = data[finite] + needs_float = bool(state.keys() & {"mean", "m2", "sum", "sumofsquares"}) + data = valid_data.astype(float, copy=False) if needs_float else valid_data + + # Measure spread around each local mean to keep small variation accurate beside large values + if needs_mean or "sum" in state: + sums = ( + np.array([np.sum(data)], dtype=float) + if valid_codes is None + else np.bincount(valid_codes, weights=data, minlength=len(labels)) + ) + if "sum" in state: + state["sum"][index] = sums + if needs_mean: + mean = np.divide(sums, count, out=np.zeros(len(labels)), where=count > 0) + state["mean"][index] = mean + + # Store squared deviations from the local mean; the merge later corrects for differences between means + if "m2" in state: + deviations = data - (mean[0] if valid_codes is None else mean[valid_codes]) + state["m2"][index] = ( + np.array([np.sum(deviations**2)], dtype=float) + if valid_codes is None + else np.bincount(valid_codes, weights=deviations**2, minlength=len(labels)) + ) + + # Keep uncentered squared values for sum of squares and RMSE, which measure magnitude rather than spread + if "sumofsquares" in state: + state["sumofsquares"][index] = ( + np.array([np.sum(data**2)], dtype=float) + if valid_codes is None + else np.bincount(valid_codes, weights=data**2, minlength=len(labels)) + ) + + # Leave empty groups at their neutral extrema; finalization turns their estimates into NaN + if data.size: + if "min" in state: + if valid_codes is None: + state["min"][index, 0] = np.min(valid_data) + else: + np.minimum.at(state["min"][index], valid_codes, valid_data) + if "max" in state: + if valid_codes is None: + state["max"][index, 0] = np.max(valid_data) + else: + np.maximum.at(state["max"][index], valid_codes, valid_data) + return labels, size, state + + +def _merge_blocks( + summaries: Sequence[tuple[NDArray[Any], NDArray[Any], dict[str, NDArray[Any]]]], +) -> tuple[NDArray[Any], NDArray[Any], dict[str, NDArray[Any]]]: + """ + Combine dense or sparse block summaries with stable pairwise means and variances. + + Sparse group labels are aligned before counts, sums, minima, maxima, means and squared deviations are combined. The + mean and variance update follows Chan, Golub and LeVeque (1983), avoiding subtraction of large squared sums. + Group sizes include all eligible locations, while the counts used to weight means include only finite values. + + :param summaries: Nonempty sequence of block summaries in the format returned by _reduce_block(). + :returns: One summary in that same format, containing the union of the input groups. + """ + + # 1/ Align group IDs when each block stored only the groups it contained + first_labels, _, first_state = summaries[0] + same_labels = all(np.array_equal(summary[0], first_labels) for summary in summaries[1:]) + labels = first_labels if same_labels else np.unique(np.concatenate([summary[0] for summary in summaries])) + size = np.zeros(len(labels), dtype=np.int64) + shape = (first_state["count"].shape[0], len(labels)) + + # 2/ Allocate the same quantities as the incoming summaries, using neutral starting values + combined: dict[str, NDArray[Any]] = {} + for name in first_state: + fill = np.inf if name == "min" else -np.inf if name == "max" else 0 + combined[name] = np.full(shape, fill, dtype=first_state[name].dtype) + + # 3/ Merge each block's counts, means, spread, and other requested quantities + for block_labels, block_size, state in summaries: + # Avoid label indexing when all summaries already use the same group order + positions = slice(None) if same_labels else np.searchsorted(labels, block_labels) + size[positions] += block_size + old_count = combined["count"][:, positions] + new_count = old_count + state["count"] + + # Weight the mean shift by the fraction of finite values contributed by the incoming block + if "mean" in combined: + old_mean = combined["mean"][:, positions] + delta = state["mean"] - old_mean + fraction = np.divide(state["count"], new_count, out=np.zeros_like(delta), where=new_count > 0) + combined["mean"][:, positions] = old_mean + delta * fraction + + # Correct the summed squared deviations for the difference between the two block means + if "m2" in combined: + correction = delta**2 * old_count * fraction + combined["m2"][:, positions] += state["m2"] + correction + combined["count"][:, positions] = new_count + + # Combine sums and ranges without keeping the original values + for name in state.keys() - {"count", "mean", "m2"}: + current = combined[name][:, positions] + if name == "min": + current = np.minimum(current, state[name]) + elif name == "max": + current = np.maximum(current, state[name]) + else: + current = current + state[name] + combined[name][:, positions] = current + return labels, size, combined + + +def _finalize_blocks( + summary: tuple[NDArray[Any], NDArray[Any], dict[str, NDArray[Any]]], + statistics: _Statistics, +) -> pd.DataFrame: + """ + Turn combined block summaries into the same result columns as the in-memory path. + + Take the square root of the mean squared deviation for population standard deviation, or the mean squared value + for RMSE. Groups with locations but no finite selected values keep their counts and receive NaN estimates. + + Statistics is the normalized request described in _aggregate_chunked(). + + :param summary: Combined labels, total location counts and summary arrays returned by _merge_blocks(). + :returns: A table with integer group rows and (value position, statistic name) columns. + """ + + # Keep groups that contain locations even when every selected value is missing + labels, size, state = summary + observed = size > 0 + columns: dict[tuple[int, str], Any] = {} + + # Convert each selected value's summaries independently because their finite counts can differ + for index in range(state["count"].shape[0]): + count = state["count"][index] + columns[(index, "count")] = count[observed] + for name, alias in zip(statistics.names, statistics.aliases): + if alias == "validcount": + result = count + elif alias == "totalcount": + result = size + elif alias == "percentagevalidpoints": + result = np.divide(100 * count, size, out=np.full(len(size), np.nan), where=size > 0) + + # Normalize by the population count, keeping empty groups undefined instead of dividing by zero + elif alias in {"std", "rmse"}: + numerator = state["m2" if alias == "std" else "sumofsquares"][index] + result = np.sqrt(np.divide(numerator, count, out=np.full(len(size), np.nan), where=count > 0)) + else: + assert alias is not None + result = np.where(count > 0, state[alias][index], np.nan) + selected_result = result[observed] + columns[(index, name)] = _extrema_result(selected_result) if alias in {"min", "max"} else selected_result + + # Keep integer group IDs until the grouping parent restores category or interval labels + return pd.DataFrame(columns, index=labels[observed]) + + +def _collect_block( + values: Sequence[NDArray[Any]], group_ids: NDArray[Any] | None, labels: Sequence[int] +) -> tuple[NDArray[Any] | None, list[NDArray[Any]]]: + """ + Collect full group values from one block for statistics such as median or a user function. + + Apply the same group selection to every value array and retain missing values for total counts. With no group + IDs, return the complete flattened arrays without constructing a group ID array. + + Values and group_ids follow _reduce_values(), restricted to this block. + + :param labels: Integer group IDs to collect from this block; ignored when group_ids is None. + :returns: Selected flattened group IDs, or None without groups, followed by matching flattened value arrays. + """ + + from geoutils.multiproc.readers import _read_values + + # Read only this worker's file slices before selecting group members + values = [_read_values(value) for value in values] + + # Return only group members, including missing values needed for total counts + if group_ids is None: + selected: slice | NDArray[Any] = slice(None) + ids = None + else: + ids = np.asarray(group_ids).ravel() + selected = ids == labels[0] if len(labels) == 1 else np.isin(ids, labels) + ids = ids[selected] + return ids, [np.asanyarray(array).ravel()[selected] for array in values] + + +def _aggregate_collected( + blocks: Sequence[tuple[NDArray[Any] | None, list[NDArray[Any]]]], statistics: _Statistics +) -> pd.DataFrame: + """ + Calculate exact statistics after joining a group's values from every block it crosses. + + Concatenate corresponding value arrays and pass their shared group IDs to _aggregate_eager(). This keeps + the eager estimator definitions while allowing one complete group to span several chunks. + + Statistics is the normalized request described in _aggregate_chunked(). + + :param blocks: Nonempty sequence of matching group IDs and value arrays returned by _collect_block(). + :returns: The grouped table produced by _aggregate_eager() for the joined values. + """ + + # Join each group's values without padding it to the full input shape + ids = None if blocks[0][0] is None else np.concatenate([block[0] for block in blocks if block[0] is not None]) + arrays = [] + for index in range(len(blocks[0][1])): + values = [block[1][index] for block in blocks] + concatenate = np.ma.concatenate if any(np.ma.isMaskedArray(value) for value in values) else np.concatenate + arrays.append(concatenate(values)) + return _aggregate_eager(arrays, ids, statistics) + + +def _aggregate_chunked( + values: Sequence[Any], + group_ids: Any | None, + statistics: _Statistics, + total_groups: int, + strategy: str, + mp_config: MultiprocConfig | None, +) -> pd.DataFrame: + """ + Calculate grouped statistics across Dask chunks or NumPy tiles handled by worker processes. + + ``dense`` reduces each input chunk to fixed-size arrays containing every declared group. + ``sparse`` stores only groups present in each chunk and aligns their labels when chunks are combined. + + Both combine counts, sums, minima, maxima, means, standard deviations and RMSE without keeping the original values. + + Median, NMAD, other quantiles, and user functions cannot be aggregated across chunks from these summaries. For + these statistics, ``groupwise`` finds the chunks containing each group and collects all its values before + calculation. Groups found in the same chunks are read together in batches near one chunk's size. + + _collect_block() and _aggregate_collected() calculate complete groups. For mergeable statistics, _reduce_block() + summarizes each block, _merge_blocks() combines small batches of summaries, and _finalize_blocks() builds the + result columns. Both backends share these calculations, with the following execution differences: + + - Dask schedules block reductions and an eight-way merge tree. Multiprocessing submits up to eight reductions + at a time, then merges their summaries in the caller using a bounded set of binary accumulation levels. + Different merge trees and block layouts can produce small floating-point differences. + - Both backends find group memberships before collecting values and process group batches one at a time. + Dask schedules collection and complete-group calculation as tasks; multiprocessing collects members in + workers and calculates their statistics in the caller. A complete group is never split, so its memory use + can exceed the target batch size. Dask may recompute contributing blocks between successive batches. + - Collected values follow block-grid order, then flattened order within each block. Equal block layouts give + the same order across these backends, but tiling can change the order from an eager array's flattened order. + Consequently, grouped user functions that depend on observation order can depend on the block layout. + + Both paths finish computation before returning a Pandas table. + + Values, group_ids, total_groups and mp_config follow _reduce_values(). Strategy is already resolved from auto + to dense, sparse or groupwise by that parent. + + :param statistics: Validated _Statistics from _normalize_statistics(), reused by every block calculation. + :returns: A computed table with integer group rows and (value position, statistic name) columns. + """ + + # 1/ Divide values and group IDs into matching blocks for the selected backend + + # Return an empty table without starting tasks when the input is empty + if values[0].size == 0: + empty_ids = None if group_ids is None else np.empty(0, dtype=int) + return _aggregate_eager([np.empty(0) for _ in values], empty_ids, statistics) + + # Split values and group IDs into matching Dask chunks or NumPy tiles + use_dask = is_dask_array(group_ids) or any(is_dask_array(array) for array in values) + if use_dask: + import_optional("dask") + import dask + import dask.array as da + + chunk_source = next(array for array in [group_ids, *values] if is_dask_array(array)) + chunks = chunk_source.chunks + value_blocks = [list(da.asarray(array).rechunk(chunks).to_delayed().ravel()) for array in values] + id_blocks: list[Any] + + # A whole-array reduction needs no group ID blocks, but still follows the same block ordering + if group_ids is None: + id_blocks = [None] * len(value_blocks[0]) + else: + ids = da.asarray(group_ids).rechunk(chunks) + id_blocks = list(ids.to_delayed().ravel()) + submit = dask.delayed + block_size = math.prod(max(lengths) for lengths in chunks) + else: + # Follow the requested worker tile size in every array dimension + if mp_config is None: + raise ValueError("Chunked NumPy aggregation requires ``mp_config``.") + from geoutils.multiproc.chunked import iter_chunk_slices + + lengths = (mp_config.chunks, mp_config.chunks) if isinstance(mp_config.chunks, int) else mp_config.chunks + + # Reuse identical tile views for group memberships and every selected value + tiles = list(iter_chunk_slices(values[0].shape, mp_config.chunks)) + id_blocks = [None] * len(tiles) if group_ids is None else [group_ids[tile] for tile in tiles] + from geoutils.multiproc.readers import _ValueReader + + value_blocks = [ + [array.block(tile) if isinstance(array, _ValueReader) else array[tile] for tile in tiles] + for array in values + ] + block_size = math.prod(lengths[axis % 2] for axis in range(values[0].ndim)) + + # 2/ Collect complete groups when the requested statistics need all their values + if strategy == "groupwise": + # Find group memberships first, reading only group IDs until the required value blocks are known + if group_ids is None: + locations = {0: list(range(len(id_blocks)))} + sizes = {0: int(values[0].size)} + elif use_dask: + memberships = list( + dask.compute(*[dask.delayed(np.unique)(block, return_counts=True) for block in id_blocks]) + ) + else: + memberships = [np.unique(block, return_counts=True) for block in id_blocks] + if group_ids is not None: + locations = {} + sizes = {} + for block_index, (labels, counts) in enumerate(memberships): + # Ignore the missing marker and record every block contributing locations to each group + for label, count in zip(labels[labels >= 0], counts[labels >= 0]): + locations.setdefault(int(label), []).append(block_index) + sizes[int(label)] = sizes.get(int(label), 0) + int(count) + + # Read several small groups together when they cross the same blocks + cohorts: dict[tuple[int, ...], list[int]] = {} + for label, block_indexes in locations.items(): + cohorts.setdefault(tuple(block_indexes), []).append(label) + batches: list[tuple[list[int], tuple[int, ...]]] = [] + for cohort_blocks, labels in cohorts.items(): + batch: list[int] = [] + size = 0 + for label in labels: + # Start another batch before adding a group would exceed the target size; never split one group + if batch and size + sizes[label] > block_size: + batches.append((batch, cohort_blocks)) + batch, size = [], 0 + batch.append(label) + size += sizes[label] + if batch: + # Keep the last partly filled batch for this set of blocks + batches.append((batch, cohort_blocks)) + + # Keep each batch near one block's size unless one group is larger by itself + tables = [] + for labels, cohort_blocks in batches: + if use_dask: + members = [ + dask.delayed(_collect_block)([blocks[index] for blocks in value_blocks], id_blocks[index], labels) + for index in cohort_blocks + ] + table = dask.delayed(_aggregate_collected)(members, statistics).compute() + else: + assert mp_config is not None + + # Workers select only the needed group members; the parent joins them for the exact estimators + handles = [ + mp_config.cluster.submit( + _collect_block, [blocks[index] for blocks in value_blocks], id_blocks[index], labels + ) + for index in cohort_blocks + ] + members = mp_config.cluster.gather(handles) + table = _aggregate_collected(members, statistics) + tables.append(table) + + # Keep exact integer columns when another batch contains only missing extrema + if tables: + for column in tables[0].columns: + integer_dtypes = {str(table[column].dtype) for table in tables} & {"Int64", "UInt64"} + if integer_dtypes: + dtype = "UInt64" if "UInt64" in integer_dtypes else "Int64" + for table in tables: + table[column] = table[column].astype(dtype) + + # Restore integer group order across batches, keeping the expected columns if no groups were found + return ( + pd.concat(tables).sort_index() + if tables + else _aggregate_eager([np.empty(0) for _ in values], np.empty(0, dtype=int), statistics) + ) + + # 3/ Combine small summaries when the requested statistics can be merged across blocks + + # Ask each block for only the counts, sums, means, or ranges needed by the request + reduction_statistics = {alias for alias in statistics.aliases if alias is not None} + dense = strategy == "dense" + if use_dask: + tasks = [ + submit(_reduce_block)( + [blocks[index] for blocks in value_blocks], block, total_groups, dense, reduction_statistics + ) + for index, block in enumerate(id_blocks) + ] + + # Merge groups of eight tasks at each level so no single task receives every block summary + while len(tasks) > 1: + tasks = [submit(_merge_blocks)(tasks[start : start + 8]) for start in range(0, len(tasks), 8)] + summary = tasks[0].compute() + else: + assert mp_config is not None + # Combine worker results in small groups so finished summaries do not keep piling up in memory + levels: list[Any] = [] + for start in range(0, len(id_blocks), 8): + handles = [ + mp_config.cluster.submit( + _reduce_block, + [blocks[index] for blocks in value_blocks], + id_blocks[index], + total_groups, + dense, + reduction_statistics, + ) + for index in range(start, min(start + 8, len(id_blocks))) + ] + summary = _merge_blocks(mp_config.cluster.gather(handles)) + depth = 0 + + # Keep at most one accumulated summary per level, merging equally sized batches in submission order + while depth < len(levels) and levels[depth] is not None: + summary = _merge_blocks([levels[depth], summary]) + levels[depth] = None + depth += 1 + if depth == len(levels): + # Add a level when this batch has combined with every earlier level + levels.append(summary) + else: + levels[depth] = summary + + # Combine the remaining occupied levels once all worker batches have finished + summary = _merge_blocks([level for level in levels if level is not None]) + return _finalize_blocks(summary, statistics) + + +######################### +# 4/ SHARED PARENT REDUCTION +######################### + + +def _reduce_values( + values: Sequence[Any], + statistics: _Statistics, + *, + group_ids: Any | None = None, + total_groups: int = 1, + strategy: str = "auto", + mp_config: MultiprocConfig | None = None, +) -> tuple[pd.DataFrame, str]: + """ + Reduce one or more arrays globally or by integer group IDs. + + _resolve_strategy() chooses block summaries or complete groups. + _aggregate_chunked() schedules the shared calculations for Dask or multiprocessing. + Eager grouped inputs use the same _reduce_block() and _finalize_blocks() summaries, while one global group + or exact grouped estimates use _aggregate_eager(). + + The shared block helpers define the estimates, while _aggregate_chunked() documents the backend differences in + scheduling, memory use and observation order. This parent returns a computed table for every backend. + + Statistics, strategy and mp_config follow stats(). + + :param values: Nonempty sequence of numeric NumPy or Dask arrays with matching shapes and missing values + represented by NaN. Dask inputs must have known shapes and compatible chunks. + :param group_ids: Integer array with the same shape as the values. IDs from zero to total_groups - 1 identify + groups, and negative IDs exclude locations. None treats all locations as one group. + :param total_groups: Number of declared group combinations, including groups with no eligible locations. + :returns: A computed table with integer group rows and (value position, statistic name) columns, and the resolved + reduction strategy. + """ + + # Use matching block tasks for Dask and multiprocessing, with direct NumPy reduction otherwise + use_dask = is_dask_array(group_ids) or any(is_dask_array(array) for array in values) + if use_dask and mp_config is not None: + raise ValueError("Dask inputs cannot be combined with Multiprocessing statistics.") + chunked = use_dask or mp_config is not None + resolved_strategy, mergeable = _resolve_strategy(statistics.aliases, strategy, total_groups, chunked) + + # Schedule blocks only when requested; eager arrays can use the same summaries without task overhead + if chunked: + table = _aggregate_chunked(values, group_ids, statistics, total_groups, resolved_strategy, mp_config) + elif group_ids is None: + table = _aggregate_eager(values, group_ids, statistics) + elif mergeable: + summary = _reduce_block( + values, + group_ids, + total_groups, + resolved_strategy == "dense", + {alias for alias in statistics.aliases if alias is not None}, + ) + table = _finalize_blocks(summary, statistics) + else: + table = _aggregate_eager(values, group_ids, statistics) + return table, resolved_strategy + + +############################ +# 5/ GLOBAL OUTPUT +############################ + + +def _global_reduction_statistics(statistics: _Statistics) -> _Statistics: + """Keep numerical estimators in the shared reducer while leaving global-only output to its formatter.""" + + selected = [ + (statistic, name, alias) + for statistic, name, alias in zip(statistics.requested, statistics.names, statistics.aliases) + if alias in _STATS_ALIAS_CALLABLE + ] + return replace( + statistics, + requested=[statistic for statistic, _, _ in selected], + names=[name for _, name, _ in selected], + aliases=[alias for _, _, alias in selected], + single=False, + ) + + +def _global_values_require_eager( + values: Sequence[Any], + statistics: _Statistics, + strategy: str, + mp_config: MultiprocConfig | None, +) -> bool: + """Identify global inputs that need their complete array for established NumPy or callable behavior.""" + + # User functions receive the complete selected array, including its original dimensions + if any(callable(statistic) for statistic in statistics.requested): + return True + if strategy == "groupwise" or not any(alias in _STATS_ALIAS_CALLABLE for alias in statistics.aliases): + return False + + # NumPy's global estimators include infinities while the grouped block reducer selects finite observations + dask_checks = [] + for array in values: + if not np.issubdtype(array.dtype, np.inexact): + continue + if is_dask_array(array): + import_optional("dask") + import dask.array as da + + dask_checks.append(da.ma.filled(da.isinf(array), False).any()) + else: + if isinstance(array, _ValueReader): + assert mp_config is not None + if _reader_requires_eager(array, mp_config): + return True + elif bool(np.ma.filled(np.any(np.isinf(array)), False)): + return True + if dask_checks: + import dask + + return any(dask.compute(*dask_checks)) + return False + + +def _materialize_global_values(values: Sequence[Any]) -> list[Any]: + """Read complete raster or point cloud values and compute Dask arrays together for global-only behavior.""" + + materialized = list(values) + dask_positions = [] + dask_values = [] + for index, array in enumerate(materialized): + if isinstance(array, _ValueReader): + window = tuple(slice(0, length) for length in array.shape) + materialized[index] = array.read(window) + elif is_dask_array(array): + dask_positions.append(index) + dask_values.append(array) + if dask_values: + import_optional("dask") + import dask + + for index, array in zip(dask_positions, dask.compute(*dask_values)): + materialized[index] = array + return materialized + + +def _materialize_global_counts( + counts: Sequence[_SelectionCounts | None], +) -> list[_SelectionCounts | None]: + """Compute all small Dask mask counts together before formatting global summaries.""" + + materialized = list(counts) + positions = [] + values = [] + for index, selection_counts in enumerate(counts): + if selection_counts is None: + continue + positions.append(index) + values.extend((selection_counts.valid_before_mask, selection_counts.selected_locations)) + if any(is_dask_array(value) for value in values): + import_optional("dask") + import dask + + values = list(dask.compute(*values)) + for position, start in zip(positions, range(0, len(values), 2)): + selection_counts = counts[position] + assert selection_counts is not None + materialized[position] = replace( + selection_counts, + valid_before_mask=values[start], + selected_locations=values[start + 1], + ) + return materialized + + +def _format_global_stats( + table: pd.DataFrame, + value_names: Sequence[str], + values: Sequence[Any], + counts: Sequence[_SelectionCounts | None], + statistics: _Statistics, +) -> Any: + """Restore scalar or dictionary global output from the one-group reduction table and selection counts.""" + + results = {} + for value_index, (value_name, array, selection_counts) in enumerate(zip(value_names, values, counts)): + final_count = int(table.loc[0, (value_index, "count")]) + valid_count = final_count if selection_counts is None else int(selection_counts.valid_before_mask) + total_count = int(array.size) + selected_count = None if selection_counts is None else int(selection_counts.selected_locations) + value_result = {} + + # Map the shared numerical columns and the global mask counts back to the requested output labels + for statistic, name, alias in zip(statistics.requested, statistics.names, statistics.aliases): + if callable(statistic): + result = statistic(array) + elif alias in _STATS_ALIAS_CALLABLE: + result = table.loc[0, (value_index, name)] + elif alias == "validcount": + result = valid_count + elif alias == "totalcount": + result = total_count + elif alias == "percentagevalidpoints": + result = 100 * valid_count / total_count if total_count else np.nan + elif alias == "validinliercount" and selection_counts is not None: + result = final_count + elif alias == "totalinliercount" and selection_counts is not None: + result = selected_count + elif alias == "percentageinlierpoints" and selection_counts is not None: + result = 100 * final_count / valid_count if valid_count else np.nan + elif alias == "percentagevalidinlierpoints" and selection_counts is not None: + result = 100 * final_count / selected_count if selected_count else 0 + else: + result = np.nan + if not callable(statistic) and isinstance(result, np.generic) and not np.ma.is_masked(result): + result = result.item() + value_result[name] = result + results[value_name] = next(iter(value_result.values())) if statistics.single else value_result + return next(iter(results.values())) if len(results) == 1 else results + + +def _reduce_global_values( + selected_values: Mapping[str, tuple[Any, _SelectionCounts | None]], + statistics: _Statistics, + *, + strategy: str, + mp_config: MultiprocConfig | None, +) -> Any: + """ + Calculate global statistics through the same reducer as one complete grouped bin. + + _global_reduction_statistics() selects the numerical estimates shared with grouped statistics. _reduce_values() + calculates them with one implicit group, and _format_global_stats() restores the established global counts, + callable results and scalar or dictionary output. + + Selected_values is returned by _sample_and_mask_global_values(); statistics, strategy and mp_config follow stats(). + + :returns: Statistics for one selected value, or a mapping from value names to their statistics. + """ + + # Resolve the requested backend before removing global-only counts and callable output from the shared table + names = list(selected_values) + values = [selected_values[name][0] for name in names] + counts = [selected_values[name][1] for name in names] + chunked = any(is_dask_array(array) for array in values) or mp_config is not None + strategy_aliases = [ + alias + for statistic, alias in zip(statistics.requested, statistics.aliases) + if callable(statistic) or alias in _STATS_ALIAS_CALLABLE + ] + resolved_strategy, _ = _resolve_strategy(strategy_aliases, strategy, 1, chunked) + + # Read complete arrays only for custom functions or the exceptional global handling of infinities + use_eager = chunked and _global_values_require_eager(values, statistics, resolved_strategy, mp_config) + if use_eager: + values = _materialize_global_values(values) + reduction_statistics = _global_reduction_statistics(statistics) + table, _ = _reduce_values( + values, + reduction_statistics, + strategy=resolved_strategy, + mp_config=None if use_eager else mp_config, + ) + + # Finish the small global count fields and custom functions after every numerical backend returns its table + counts = _materialize_global_counts(counts) + return _format_global_stats(table, names, values, counts, statistics) diff --git a/geoutils/stats/selection.py b/geoutils/stats/selection.py new file mode 100644 index 000000000..b684c783e --- /dev/null +++ b/geoutils/stats/selection.py @@ -0,0 +1,509 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Select reprojected values, masks and samples for statistics.""" + +from __future__ import annotations + +import math +from collections.abc import Iterable, Mapping +from contextlib import ExitStack +from dataclasses import dataclass, replace +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Any, Literal, cast + +import numpy as np +import pandas as pd +from numpy.typing import NDArray + +from geoutils._dispatch import ( + _get_pointcloud_interface, + _get_raster_interface, + _is_pointcloud, + _is_raster, + is_dask_array, +) +from geoutils._misc import import_optional +from geoutils._typing import ArrayLike +from geoutils.multiproc.chunked import iter_chunk_slices +from geoutils.multiproc.cluster import _map_bounded +from geoutils.multiproc.readers import ( + _normalize_reader_mask, + _read_selected_values, + _read_values, + _reader_from_source, + _reader_from_vector, + _ValueReader, +) +from geoutils.raster.array import get_mask_from_array +from geoutils.sampling.stratified import _stratified_subsample_indices +from geoutils.sampling.subsampling import _sample_valid_indices, _subsample_numpy +from geoutils.sampling.support import ( + _as_array, + _mask_at_support, + _normalize_mask_array, + _normalize_sampling_input, + _sampling_specification, + _sampling_support, + _values_at_support, +) + +if TYPE_CHECKING: + from geoutils.interface.interpolation import InterpolationMethod + from geoutils.multiproc.mparray import MultiprocConfig + from geoutils.pointcloud.base import PointCloudBase + from geoutils.pointcloud.pointcloud import PointCloudLike + from geoutils.raster.base import RasterBase, RasterLike + from geoutils.vector.base import VectorLike + + +@dataclass(frozen=True) +class _SelectionCounts: + """Store value validity before masking and the number of locations selected by the mask.""" + + valid_before_mask: Any + selected_locations: Any + + +######################################## +# 1/ SELECT A COMMON SPATIAL SUPPORT +######################################## + + +def _select_values_and_mask_at_support( + source: RasterLike | PointCloudLike | ArrayLike | Mapping[str, ArrayLike], + *, + by: Mapping[str, Any] | None, + values: int | str | Iterable[int | str] | Mapping[str, Any] | None, + at: Literal["self"] | RasterLike | PointCloudLike | None, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None, + mask_mode: str, + interpolation: InterpolationMethod, + align: Literal["raise", "reproject"], + mp_config: MultiprocConfig | None, + stack: ExitStack | None = None, +) -> tuple[Any, Any | None, RasterBase | PointCloudBase | None]: + """ + Select values and an optional mask on the same support for summary and grouped statistics. + + _sampling_support() chooses the grid or point locations, and _values_at_support() aligns selected bands, columns, + or named inputs to them. Keep native arrays unchanged when no alignment is needed so reductions retain their + masked-array behavior and Dask laziness. + + All arguments follow stats(). + + :returns: Selected values, a boolean eligibility mask or None, and the raster or point cloud defining their + locations. Plain array inputs keep their input representation and have no spatial support. + """ + + # 1/ Check user input before reading any possibly lazy values + if mask_mode not in {"inside", "outside"}: + raise ValueError("Argument ``mask_mode`` must be 'inside' or 'outside'.") + if align not in {"raise", "reproject"}: + raise ValueError("Argument ``align`` must be 'raise' or 'reproject'.") + if isinstance(at, str) and at != "self": + raise ValueError("Argument ``at`` must be 'self' or a raster or point cloud support object.") + + # Use the shared dispatch checks, treating Xarrays without spatial coordinates as plain arrays + spatial_source = _normalize_sampling_input(source) + is_raster = _is_raster(spatial_source) + is_pointcloud = _is_pointcloud(spatial_source) + if not is_raster and not is_pointcloud: + if values is not None or at is not None: + raise ValueError("Arguments ``values`` and ``at`` require a raster or point cloud source.") + return source, mask, None + source = _get_raster_interface(spatial_source) if is_raster else _get_pointcloud_interface(spatial_source) + raster = cast("RasterBase", source) if is_raster else None + + # Choose the same output grid or point locations for selected values and groupers + inputs = [ + source, + *(by.values() if by is not None else []), + *(values.values() if isinstance(values, Mapping) else []), + ] + support = _sampling_support(inputs, source if isinstance(at, str) else at) + support_dataframe = None + + # 2/ Name the requested values and read them on the selected grid or points + # Preserve existing summary band labels and explicit names supplied in a mapping + if isinstance(values, Mapping): + value_specs = dict(values) + elif raster is not None: + # Interpret numeric selections as one-based bands; omitting values selects every band + bands: Iterable[Any] + if values is None: + bands = range(1, raster.count + 1) + elif isinstance(values, (int, np.integer)): + bands = [values] + else: + bands = values + if not isinstance(bands, Iterable) or isinstance(bands, (str, bytes)): + raise TypeError("Raster ``values`` must select one or more integer band numbers.") + value_specs = {} + for selected_band in bands: + if not isinstance(selected_band, (int, np.integer)) or not 1 <= selected_band <= raster.count: + raise ValueError("Raster bands must be integers between one and the raster band count.") + name = f"band_{selected_band}" if by is not None else f"band {selected_band}" + value_specs[name] = int(selected_band) + else: + # Use the active point column by default, or geometry heights when no column is selected + if values is None: + pointcloud = cast("PointCloudBase", source) + value_specs = {pointcloud.data_column or "z": pointcloud.data_column} + else: + columns = [values] if isinstance(values, str) else values + if not isinstance(columns, Iterable): + raise TypeError("Point cloud ``values`` must select one or more column names.") + value_specs = {} + for column in columns: + if not isinstance(column, str): + raise TypeError("Point cloud ``values`` must select column names.") + value_specs[column] = column + if not value_specs or any(not isinstance(name, str) or not name for name in value_specs): + raise ValueError("Selected value names must be non-empty strings.") + + # Read native raster bands directly; interpolate or align only when the selected support requires it + selected_values: dict[str, Any] = {} + for name, specification in value_specs.items(): + value_source, selector = _sampling_specification(source, specification) + # Plain Xarrays follow the selected locations; their dimensions alone do not define a raster + value_source = _normalize_sampling_input(value_source) + file_values = _reader_from_source( + value_source, selector, support, mp_config, align=align, interpolation=interpolation, stack=stack + ) + if file_values is None and mp_config is not None and not support.is_loaded: + from geoutils._dispatch import _is_vector + + if not _is_raster(value_source) and not _is_pointcloud(value_source) and _is_vector(value_source): + from geoutils.sampling.support import _as_geodataframe + + dataframe = _as_geodataframe(value_source) + if selector is None or selector not in dataframe.columns: + raise ValueError("Vector values require an explicit feature column.") + if not pd.api.types.is_numeric_dtype(dataframe[selector]): + raise TypeError("Selected vector values must be numeric.") + file_values = _reader_from_vector( + dataframe, np.asarray(dataframe[selector], dtype=float), support, mp_config + ) + if file_values is not None: + selected_values[name] = file_values + continue + if value_source is raster and support is raster: + # Read a native band directly to retain its NumPy mask and avoid unnecessary interpolation + selected_band = 1 if selector is None else selector + if not isinstance(selected_band, (int, np.integer)) or not 1 <= selected_band <= raster.count: + raise ValueError("Raster bands must be integers between one and the raster band count.") + data = raster.data + selected_values[name] = data[selected_band - 1] if data.ndim == 3 else data + else: + # Let the sampling helpers align, rasterize or interpolate spatial inputs as required + if support_dataframe is None and not _is_raster(support): + support_dataframe = cast("PointCloudBase", support).ds + selected_values[name] = _values_at_support( + value_source, + selector, + input_support=source, + support=support, + support_dataframe=support_dataframe, + name=name, + interpolation=interpolation, + align=align, + mp_config=mp_config, + ) + + # 3/ Place the optional mask through the same sampling helpers used by cosample() + # Keep value validity separate because stats() reports finite counts independently for each selected value + file_mask = ( + _reader_from_source(mask, None, support, mp_config, align=align, interpolation="nearest", stack=stack) + if mask is not None + else None + ) + if file_mask is not None: + return selected_values, file_mask, support + if mask is not None: + from geoutils._dispatch import _is_vector + + if not _is_raster(mask) and not _is_pointcloud(mask) and _is_vector(mask): + from geoutils.sampling.support import _as_geodataframe + + dataframe = _as_geodataframe(mask) + file_mask = _reader_from_vector( + dataframe, + np.ones(len(dataframe)), + support, + mp_config, + mask_mode=cast(Literal["inside", "outside"], mask_mode), + ) + if file_mask is not None: + return selected_values, file_mask, support + elif not _is_raster(support) and not _is_raster(mask) and not _is_pointcloud(mask): + return selected_values, _normalize_mask_array(mask, (cast("PointCloudBase", support).point_count,)), support + if mask is not None and support_dataframe is None and not _is_raster(support) and _is_pointcloud(mask): + support_dataframe = cast("PointCloudBase", support).ds + support_mask = _mask_at_support( + mask, + support, + support_dataframe=support_dataframe, + mask_mode=mask_mode, + align=align, + mp_config=mp_config, + ) + return selected_values, support_mask, support + + +#################################### +# 2/ SAMPLE ELIGIBLE LOCATIONS +#################################### + + +def _sample_eligible_indices( + eligible: Any, + *, + subsample: int | float, + random_state: int | np.random.Generator | None, + strategy: Literal["sequential", "topk"], +) -> NDArray[Any]: + """ + Select flat eligible positions with the same sampling rules for summary and grouped statistics. + + Grid sampling delegates to _sample_valid_indices(): ``topk`` selects the same locations across chunk layouts, + while ``sequential`` follows each backend's traversal order. Other array dimensions use the NumPy sampler after + computing a lazy eligibility mask. Sampling finishes here before either reduction backend receives the values. + + Subsample and random_state follow stats(). + + :param eligible: Boolean array marking locations available for sampling, independently of value validity. + :param strategy: The subsampling_strategy option from stats(), distinct from its reduction strategy. + :returns: Selected positions in the flattened input as an in-memory integer array. + """ + + # Delegate grid sampling to the shared sampler, including its Dask task handling + if eligible.ndim == 2: + rows, columns = _sample_valid_indices( + eligible, subsample=subsample, random_state=random_state, strategy=strategy + ) + return rows * eligible.shape[1] + columns + + # Keep the existing flat sampling order for point values and in-memory inputs + if is_dask_array(eligible): + eligible = eligible.compute() + (flat_indices,) = _subsample_numpy( + np.where(np.asarray(eligible).ravel(), 1.0, np.nan), + subsample=subsample, + return_indices=True, + random_state=random_state, + strategy=strategy, + ) + return flat_indices + + +def _sample_reader_indices( + shape: tuple[int, ...], + mask: Any, + subsample: int | float, + random_state: Any, + strategy: Literal["topk", "sequential"], + mp_config: MultiprocConfig, +) -> Any: + """Sample one global group while reading file masks and storing eligibility in bounded blocks.""" + + if math.prod(shape) == 0: + return np.empty(0, dtype=np.int64) + + # Sample an unmasked input as one group without allocating or writing a complete group array + if mask is None: + ids = np.broadcast_to(np.array(0, dtype=np.int8), shape) + return _stratified_subsample_indices( + ids, subsample, random_state=random_state, strategy=strategy, mp_config=mp_config + ) + + with ExitStack() as storage: + directory = storage.enter_context(TemporaryDirectory(prefix="geoutils-stats-sample-")) + ids = np.memmap(f"{directory}/eligible.dat", mode="w+", dtype=np.int8, shape=shape) + storage.callback(ids._mmap.close) + tiles = list(iter_chunk_slices(shape, mp_config.chunks)) + if isinstance(mask, _ValueReader): + arguments = ((mask.block(tile),) for tile in tiles) + for index, block_mask in _map_bounded(mp_config.cluster, _read_values, arguments): + # Read only the mask: value validity does not determine the common sample + ids[tiles[index]] = np.where(np.ma.filled(block_mask, False), 0, -1) + else: + for tile in tiles: + ids[tile] = np.where(np.ma.filled(mask[tile], False), 0, -1) + return _stratified_subsample_indices( + ids, subsample, random_state=random_state, strategy=strategy, mp_config=mp_config + ) + + +########################################### +# 3/ APPLY MASKS AND KEEP SUMMARY COUNTS +########################################### + + +def _mask_global_values( + array: Any, + mask: Any | None, + selected_locations: Any | None, + mp_config: MultiprocConfig | None = None, +) -> tuple[Any, _SelectionCounts | None]: + """ + Apply an eligible-location mask and retain counts before and after that selection. + + Boolean mask validation belongs to _normalize_mask_array(). Here, missing values remain excluded independently of + the user mask. An eager array stays eager unless its mask is lazy; in that case, selection must wait for Dask too. + The original array is never changed, and its dimensions are kept for user-defined statistics. + + :param array: One selected value array, already aligned to the mask. + :param mask: Validated boolean eligibility array of the same shape, or None to keep the input unchanged. + :returns: Values with excluded locations masked or set to NaN, plus their validity before masking and the common + number of selected locations. Without a mask, the counts are None. + """ + + # Keep unmasked inputs exactly as supplied, including NumPy masks and lazy Dask graphs + if mask is None: + return array, None + + if isinstance(array, _ValueReader): + assert mp_config is not None + assert selected_locations is not None + return _mask_reader_values(array, mask, selected_locations, mp_config) + if isinstance(mask, _ValueReader): + mask = _normalize_mask_array(mask.read(tuple(slice(0, length) for length in mask.shape)), array.shape) + assert selected_locations is not None + + # Count finite input values before masking so summary validity still describes the full selected array + valid = np.isfinite(array) + if np.ma.isMaskedArray(valid): + valid = valid.filled(False) + counts = _SelectionCounts(valid_before_mask=valid.sum(), selected_locations=selected_locations) + + # Defer selection only when the array or its own mask is lazy; other selected arrays do not affect this choice + if is_dask_array(array) or is_dask_array(mask): + import_optional("dask") + import dask.array as da + + raw_values = np.ma.getdata(array) if np.ma.isMaskedArray(array) else array + selected = da.where(mask & valid, raw_values, np.nan) + else: + # Preserve the masked-array estimators and the original shape for in-memory data + selected = np.ma.masked_where(~mask | ~valid, array) + return selected, counts + + +def _count_reader_value_block(values: Any) -> int: + """Count finite values in one raster or point cloud reader block.""" + + data = _read_values(values) + return int(np.count_nonzero(~get_mask_from_array(data))) + + +def _count_reader_mask_block(mask: Any) -> int: + """Count selected locations in one raster or point cloud mask block.""" + + selected = _read_values(mask) + return int(np.count_nonzero(np.ma.filled(selected, False))) + + +def _count_selected_locations(mask: Any, shape: tuple[int, ...], mp_config: MultiprocConfig | None) -> Any: + """Count one common mask once for all selected global values.""" + + if not isinstance(mask, _ValueReader): + return mask.sum() + assert mp_config is not None + arguments = ((mask.block(tile),) for tile in iter_chunk_slices(shape, mp_config.chunks)) + return sum(result for _, result in _map_bounded(mp_config.cluster, _count_reader_mask_block, arguments)) + + +def _mask_reader_values( + array: _ValueReader, mask: Any, selected_locations: int, mp_config: MultiprocConfig +) -> tuple[_ValueReader, _SelectionCounts]: + """Attach a mask to a reader after counting its original finite values by blocks.""" + + arguments = ((array.block(tile),) for tile in iter_chunk_slices(array.shape, mp_config.chunks)) + + # Bound pending reads and keep only the accumulated validity count in the caller + valid_before_mask = 0 + for _, count in _map_bounded(mp_config.cluster, _count_reader_value_block, arguments): + valid_before_mask += count + counts = _SelectionCounts(valid_before_mask=valid_before_mask, selected_locations=selected_locations) + return replace(array, mask=mask), counts + + +################################ +# 4/ GLOBAL SAMPLING AND MASKING +################################ + + +def _sample_and_mask_global_values( + values: ArrayLike | Mapping[str, ArrayLike], + *, + mask: Any | None, + subsample: int | float, + random_state: int | np.random.Generator | None, + subsampling_strategy: Literal["sequential", "topk"], + mp_config: MultiprocConfig | None, +) -> dict[str, tuple[Any, _SelectionCounts | None]]: + """ + Name global values, sample common locations, and apply the common mask. + + Ordinary containers are converted to NumPy or Dask arrays, while raster and point cloud readers remain unloaded. + Sampling first chooses the same eligible locations for every value. Masking then keeps each value's missing data + independent and records the counts needed for global output. + + Values and selection options follow _global_stats(). + + :returns: Selected values and their optional mask counts, keyed by output value name. + """ + + # 1/ Name values and check their common shape + named_values = dict(values) if isinstance(values, Mapping) else {"value": values} + if not named_values or any(not isinstance(name, str) or not name for name in named_values): + raise ValueError("Argument ``values`` must contain at least one non-empty name.") + arrays = { + name: array if isinstance(array, _ValueReader) else _as_array(array) for name, array in named_values.items() + } + shape = next(iter(arrays.values())).shape + if any(array.shape != shape for array in arrays.values()): + raise ValueError("Selected values must have matching shapes.") + + # Keep readers unloaded and reject two task schedulers in the same calculation + mask = _normalize_reader_mask(mask, shape) + use_dask = is_dask_array(mask) or any(is_dask_array(array) for array in arrays.values()) + if use_dask and mp_config is not None: + raise ValueError("Dask inputs cannot be combined with Multiprocessing statistics.") + + # 2/ Select the same eligible positions from every value when sampling is requested + had_mask = mask is not None + if subsample != 1: + first = next(iter(arrays.values())) + if isinstance(mask, _ValueReader) or any(isinstance(array, _ValueReader) for array in arrays.values()): + assert mp_config is not None + indexes = _sample_reader_indices(shape, mask, subsample, random_state, subsampling_strategy, mp_config) + else: + eligible = mask if mask is not None else np.ones_like(first, dtype=bool) + indexes = _sample_eligible_indices( + eligible, subsample=subsample, random_state=random_state, strategy=subsampling_strategy + ) + arrays = { + name: ( + _read_selected_values(array, indexes, mp_config) + if isinstance(array, _ValueReader) + else array.reshape(-1)[indexes] + ) + for name, array in arrays.items() + } + mask = np.ones(len(indexes), dtype=bool) if had_mask else None + + # 3/ Read a mask once when in-memory values need it, then share its selected count across every value + if isinstance(mask, _ValueReader) and any(not isinstance(array, _ValueReader) for array in arrays.values()): + mask = _normalize_mask_array(mask.read(tuple(slice(0, length) for length in shape)), shape) + selected_locations = None if mask is None else _count_selected_locations(mask, shape, mp_config) + return { + name: _mask_global_values(array, mask, selected_locations, mp_config=mp_config) + for name, array in arrays.items() + } diff --git a/geoutils/stats/stats.py b/geoutils/stats/stats.py index 3a63c9fb5..2a7aec0f8 100644 --- a/geoutils/stats/stats.py +++ b/geoutils/stats/stats.py @@ -16,369 +16,480 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Module for zonal statistics.""" +"""Apply global/grouped statistics and variography to rasters and point clouds.""" from __future__ import annotations import warnings -from collections.abc import Callable -from functools import partial -from typing import Any, Literal +from collections.abc import Callable, Hashable, Iterable, Mapping +from contextlib import ExitStack +from typing import TYPE_CHECKING, Any, Literal import numpy as np -from scipy.stats import iqr -from scipy.stats.mstats import mquantiles -from geoutils._dispatch import is_dask_array +from geoutils._dispatch import _is_pointcloud, _is_raster from geoutils._misc import import_optional -from geoutils._typing import NDArrayNum -from geoutils.stats.estimators import linear_error, nmad, rmse, sum_square - -_STATS_ALIAS_CALLABLE = { - "mean": "Mean", - "median": "Median", - "max": "Max", - "min": "Min", - "sum": "Sum", - "sumofsquares": "Sum of squares", - "90thpercentile": "90th percentile", - "iqr": "IQR", - "le90": "LE90", - "nmad": "NMAD", - "rmse": "RMSE", - "std": "Standard deviation", -} - -_STATS_ALIAS_COUNTS = { - "validcount": "Valid count", - "totalcount": "Total count", - "percentagevalidpoints": "Percentage valid points", -} - -_STATS_ALIAS_GEN = _STATS_ALIAS_CALLABLE | _STATS_ALIAS_COUNTS - -_SYNONYMES = { - "maximum": "max", - "minimum": "min", - "sum": "Sum", - "sum2": "sumofsquares", - "90percentile": "90thpercentile", - "rms": "rmse", - "standarddeviation": "std", -} - -_STATS_ALIAS_MASK = { - "validinliercount": "Valid inlier count", - "totalinliercount": "Total inlier count", - "percentagevalidinlierpoints": "Percentage valid inlier points", - "percentageinlierpoints": "Percentage inlier points", -} # type: ignore - - -_STATS_ALIAS_ALL = _STATS_ALIAS_GEN | _STATS_ALIAS_MASK -_ALIAS_STATS_GEN = {v: k for k, v in _STATS_ALIAS_GEN.items()} -_ALIAS_STATS_MASK = {v: k for k, v in _STATS_ALIAS_MASK.items()} -_ALIAS_STATS_ALL = _ALIAS_STATS_GEN | _ALIAS_STATS_MASK - - -_STATS_LIST_MIN = [ - "min", - "max", - "mean", - "median", - "std", - "nmad", - "validcount", - "totalcount", - "percentagevalidpoints", -] - - -def _get_stat_common_alias(stat_name: str, stats_dict: dict[str, Any]) -> str | None: - """Return the internal statistic name for a user-facing alias.""" - - if stat_name in stats_dict: - return stat_name - - # Spaces and underscores are optional in names such as "standard deviation" - for separator in (None, "_"): - normalized_name = "".join(stat_name.lower().split(separator)) - if normalized_name in _STATS_ALIAS_ALL: - return normalized_name - if normalized_name in _SYNONYMES: - return _SYNONYMES[normalized_name] - return None - - -def _get_default_stat_names(stats_name: Literal["all"] | None, counts: tuple[int, int] | None) -> list[str]: - """Return internal statistic names for the default or complete selection.""" - - if stats_name is None: - return _STATS_LIST_MIN - - stat_names = list(_STATS_ALIAS_GEN) - if counts is not None: - stat_names += list(_STATS_ALIAS_MASK) - return stat_names - - -def _statistics_dask( - data: Any, - stats_name: list[str | Callable[[NDArrayNum], np.floating[Any]]] | Literal["all"] | None = None, - counts: tuple[int, int] | None = None, -) -> dict[str, Any]: - """Calculate common statistics on a Dask array while returning lazy scalar results.""" - - import_optional("dask") - import dask.array as da - - finite = da.isfinite(data) - final_count_nonzero = finite.sum() - valid_count = final_count_nonzero if counts is None else counts[0] - - # Dask requires explicit axes for full-array quantiles. This remains lazy, but exact global quantiles can still be - # memory-intensive at execution time because Dask has to combine data across chunks - axes = tuple(range(data.ndim)) - - def _quantile(q: float) -> Any: - return da.nanquantile(data, q, axis=axes) - - median = _quantile(0.50) - q05 = _quantile(0.05) - q25 = _quantile(0.25) - q75 = _quantile(0.75) - q90 = _quantile(0.90) - q95 = _quantile(0.95) - - stats_dict: dict[str, Any] = { - "mean": da.nanmean(data), - "median": median, - "max": da.nanmax(data), - "min": da.nanmin(data), - "sum": da.nansum(data), - "sumofsquares": da.nansum(da.square(data)), - "90thpercentile": q90, - "le90": q95 - q05, - "iqr": q75 - q25, - "nmad": 1.4826 * da.nanquantile(da.fabs(data - median), 0.50, axis=axes), - "rmse": da.sqrt(da.nanmean(da.square(data))), - "std": da.nanstd(data), - "validcount": valid_count, - "totalcount": data.size, - "percentagevalidpoints": (valid_count / data.size) * 100 if data.size else np.nan, - } - - if counts is not None: - stats_dict.update( - { - "validinliercount": final_count_nonzero, - "totalinliercount": counts[1], - "percentageinlierpoints": (final_count_nonzero / counts[0]) * 100, - "percentagevalidinlierpoints": (final_count_nonzero / counts[1]) * 100 if counts[1] != 0 else 0, - } - ) +from geoutils._typing import ArrayLike, NDArrayNum +from geoutils.stats.grouping import _grouped_stats, _validate_group_declarations +from geoutils.stats.reduction import _normalize_statistics, _reduce_global_values +from geoutils.stats.selection import ( + _sample_and_mask_global_values, + _select_values_and_mask_at_support, +) +from geoutils.stats.variography import Variogram, _estimate_variogram + +if TYPE_CHECKING: + from geoutils.interface.interpolation import InterpolationMethod + from geoutils.multiproc import MultiprocConfig + from geoutils.pointcloud.base import PointCloudBase + from geoutils.pointcloud.pointcloud import PointCloudLike + from geoutils.raster.base import RasterBase, RasterLike + from geoutils.stats.reduction import _Statistics + from geoutils.vector.base import VectorLike + +__all__ = ["stats", "variogram"] + + +def _global_stats( + values: ArrayLike | Mapping[str, ArrayLike], + statistics: _Statistics, + *, + mask: Any | None, + subsample: int | float, + random_state: int | np.random.Generator | None, + strategy: Literal["auto", "dense", "sparse", "groupwise"], + subsampling_strategy: Literal["sequential", "topk"], + mp_config: MultiprocConfig | None, +) -> Any: + """ + Calculate global statistic. + + This is done the same way as grouped statistics, simply by creating one implicit group! + + _sample_and_mask_global_values() chooses a common sample and applies the mask. + _reduce_global_values() reduces every selected value through the same calculation as one complete grouped bin. + + One selected value returns its statistics directly, several values return a dictionary keyed by their names. + + See stats() for argument descriptions. + + :param values: One numeric array or container, or a mapping of output names to selected arrays with matching + shapes, as returned by _select_values_and_mask_at_support(). + :param mask: Boolean eligibility array on the selected locations, or None to keep every location. + + :returns: Statistics for one selected value, or a mapping from value names to their statistics. + """ - if stats_name is None or stats_name == "all": - stat_names = _get_default_stat_names(stats_name, counts) - return {_STATS_ALIAS_ALL[stat_name]: stats_dict[stat_name] for stat_name in stat_names} - - res_dict: dict[str, Any] = {} - for stat_name in stats_name: - if isinstance(stat_name, str): - stat_common_alias = _get_stat_common_alias(stat_name, stats_dict) - if stat_common_alias in stats_dict: - res_dict[stat_name] = stats_dict[stat_common_alias] - elif stat_common_alias is not None: - res_dict[stat_name] = np.nan - else: - warnings.warn("Statistic name " + stat_name + " is not recognized", category=UserWarning) - res_dict[stat_name] = np.float32(np.nan) - elif callable(stat_name): - res_dict[stat_name.__name__] = stat_name(data) - else: - warnings.warn("Statistic name " + stat_name + " is not recognized", category=UserWarning) - res_dict[stat_name] = np.float32(np.nan) - - return res_dict - - -def _statistics( - data: NDArrayNum, - stats_name: list[str | Callable[[NDArrayNum], np.floating[Any]]] | Literal["all"] | None = None, - counts: tuple[int, int] | None = None, -) -> dict[str, float]: + # Sample the global values accounting for the optional user-input mask and subsample size + values_to_reduce = _sample_and_mask_global_values( + values, + mask=mask, + subsample=subsample, + random_state=random_state, + subsampling_strategy=subsampling_strategy, + mp_config=mp_config, + ) + + # Reduce all values through one shared calculation and restore the established global output form + return _reduce_global_values( + values_to_reduce, + statistics, + strategy=strategy, + mp_config=mp_config if subsample == 1 else None, + ) + + +def stats( + source: RasterLike | PointCloudLike | ArrayLike | Mapping[str, ArrayLike], + statistics: str | Callable[[Any], Any] | Iterable[str | Callable[[Any], Any]] | None = None, + *, + by: Mapping[str, Any] | None = None, + values: int | str | Iterable[int | str] | Mapping[str, Any] | None = None, + bins: Mapping[str, Any] | None = None, + categories: Mapping[str, Iterable[Hashable]] | None = None, + at: Literal["self"] | RasterLike | PointCloudLike | None = None, + mask: RasterLike | PointCloudLike | VectorLike | ArrayLike | None = None, + mask_mode: Literal["inside", "outside"] = "inside", + subsample: int | float = 1, + subsample_per_group: bool = False, + random_state: int | np.random.Generator | None = None, + strategy: Literal["auto", "dense", "sparse", "groupwise"] = "auto", + backend: Literal["geoutils", "flox"] = "geoutils", + subsampling_strategy: Literal["sequential", "topk"] = "topk", + interpolation: InterpolationMethod = "linear", + align: Literal["raise", "reproject"] = "raise", + observed: bool = True, + return_masks: bool = False, + mp_config: MultiprocConfig | None = None, +) -> Any: """ - Calculate common statistics for an N-D array : - - - Mean: arithmetic mean of the data, ignoring masked values. - - Median: middle value when the valid data points are sorted in increasing order, ignoring masked values. - - Max: maximum value among the data, ignoring masked values. - - Min: minimum value among the data, ignoring masked values. - - Sum: sum of all data, ignoring masked values. - - Sum of squares: sum of the squares of all data, ignoring masked values. - - 90th percentile: point below which 90% of the data falls, ignoring masked values. - - IQR (Interquartile Range): difference between the 75th and 25th percentile of a dataset, ignoring masked values. + Calculate statistics, either global (whole array) or grouped with other geospatial objects (continuous binning or + categorical grouping, including zonal grouping). + + Omit ``by`` for global statistics. With ``by``, provide ``categories`` for discrete groups, ``bins`` for + continuous groups along the variable. For a vector, provide the feature column through a tuple directly in ``by``, + which performs geometric zonal statistics:: + + # Global statistics + raster.stats() + raster.stats("mean") + raster.stats(["mean", "std", "nmad"]) + + # Categorical stats + raster.stats(["mean", "std"], by={"landcover": lc}, categories={"landcover": lc_classes}) + # Binned stats + raster.stats(["mean", "std"], by={"elevation": dem}, bins={"elevation": elevation_bins}) + # Zonal stats + raster.stats(["mean", "std"], by={"feature": (outlines, "id")}) + + # Multiple grouping with vector and raster inputs + raster.stats(["mean", "std"], by={"elevation": dem, "feature": (outlines, "id")}, + bins={"elevation": elevation_bins}) + + Use ``values`` with index of band (for raster) or label of column (for point cloud) to select input, which + defaults to all bands for a raster, and the main data column for a point cloud. + + # Select only band 2 (defaults to all bands) + raster.stats(["mean", "std"], values=[1, 2]) + + # Select specific point cloud columns (defaults to main column) + point.stats(["mean", "std"], values=["z", "intensity"]) + + # Provide name mapping of bands for output dataframe + raster.stats(["mean", "std"], values={"red": 1, "blue": 2}) + + Raster and point clouds are cosampled at the spatial support of ``at``. Set it to "self" to use the source + locations; otherwise, point inputs take precedence when no support is specified. + + Source input can be masked with ``mask`` and ``mask_mode``, and subsampled using ``subsample`` and + ``subsample_per_group``. + + Supported statistics + -------------------- + + - Mean: arithmetic mean of the data, + - Median: middle value when the valid data points are sorted in increasing order, + - Max: maximum value among the data, + - Min: minimum value among the data, + - Sum: sum of all data, + - Sum of squares: sum of the squares of all data, + - 90th percentile: point below which 90% of the data falls, + - IQR (Interquartile Range): difference between the 75th and 25th percentile of a dataset - LE90 (Linear Error with 90% confidence): difference between the 95th and 5th percentiles of a dataset, \ - representing the range within which 90% of the data points lie. Ignore masked values. + representing the range within which 90% of the data points lie. - NMAD (Normalized Median Absolute Deviation): robust measure of variability in the data, less sensitive to \ - outliers compared to standard deviation. Ignore masked values. + outliers compared to standard deviation. - RMSE (Root Mean Square Error): commonly used to express the magnitude of errors or variability and can give \ insight into the spread of the data. Only relevant when the raster represents a difference of two objects. \ - Ignore masked values. - - Std (Standard deviation): measures the spread or dispersion of the data around the mean, ignoring masked values. - - Valid count: number of finite data points in the array. It counts the non-masked elements. + - Std (Standard deviation): measures the spread or dispersion of the data around the mean, + - Valid count: number of unmasked entries for masked arrays, or finite entries for ordinary arrays. - Total count: total size of the raster. - Percentage valid points: ratio between Valid count and Total count. - For all statistics up to and including "Std", NumPy Masked functions are used (directly or in the calculation) - in case of a masked array, NumPy module otherwise. - - "Valid count" represents all non zero and not masked pixels in the input data (final_count_nonzero), previously - calculated in case of a Raster.get_stats() called with an inlier_mask, before the mask application. NumPy Masked - functions is used is this case or if the input was already a masked array. - Percentage valid points is calculated accordingly. - If an inlier mask is passed: - Total inlier count: number of data points in the inlier mask. - Valid inlier count: number of unmasked data points in the array after applying the inlier mask. - Percentage inlier points: ratio between Valid inlier count and Valid count. Useful for classification statistics. - Percentage valid inlier points: ratio between Valid inlier count and Total inlier count. - They are all computed based on the previously stated final_count_nonzero. - Callable functions are supported as well. - :param data: Array on which to compute statistics. - :param stats_name: list of names of the statistics to retrieve. If None, all statistics are returned. - Accepted names include: - `mean`, `median`, `max`, `min`, `sum`, `sum of squares`, `90th percentile`, `iqr`, `LE90`, `nmad`, `rmse`, - `std`, `valid count`, `total count`, `percentage valid points` and if an inlier mask is passed : - `valid inlier count`, `total inlier count`, `percentage inlier points`, `percentage valid inlier points`. - Custom callables can also be provided. - :param counts: Tuple with number of finite data points in array and number of valid points in inlier_mask. - - :returns: A dictionary containing the calculated statistics for the selected band. + Grouping details + ---------------- + + Every grouper must have an entry in ``bins`` or ``categories`` unless it has a boolean or Pandas categorical + dtype. Numeric edge sequences use left-closed intervals and include the final right edge. Pass a + :class:`pandas.IntervalIndex` to control edge closure explicitly. The result index follows the order of ``by``; + columns have ``value`` and ``statistic`` levels, and a finite ``count`` is always included for each value. + + When ``return_masks`` is true, the second result behaves as a mapping from each dataframe index key to a boolean + array or spatial object. Its masks describe complete eligible group membership after ``mask`` and valid groupers, + before random subsampling and independently of missing selected values. + + Chunk strategies + ---------------- + Dask and multiprocessing divide the inputs into chunks and use the same NumPy calculation on each one. ``dense`` + stores a fixed-size summary containing every declared group per chunk. It is efficient when the number of group + combinations is moderate. ``sparse`` stores only groups present in each chunk, reducing memory when many possible + groups are absent, at the cost of aligning group labels when chunks are combined. Both aggregate counts, sums, + minima, maxima, means, standard deviations and RMSE across chunks without keeping the original values. + + Median, NMAD, other quantiles, and user functions cannot be aggregated across chunks from the ``dense`` or + ``sparse`` summaries because they need all values from a group at once. Requesting these combinations raises an + error. ``groupwise`` instead finds the chunks containing each group and collects all its values before calculation, + so its memory use depends on the largest group. ``auto`` selects ``groupwise`` for these statistics; otherwise, it + selects ``dense`` for up to 4096 group combinations and ``sparse`` above that. With ordinary in-memory inputs, all + group values are already available and these statistics are calculated directly. + + Subsampling strategies + ---------------------- + Subsampling is global by default. Set ``subsample_per_group=True`` to apply the fraction or maximum separately + to each combination of grouping variables. Fractions round down within each group; a maximum larger than a + group keeps every eligible location in that group. ``subsample=1`` always keeps all eligible locations. + + ``topk`` assigns a reproducible SplitMix64 key to each flat input position and keeps the smallest keys, giving + the same sample for NumPy, Dask and multiprocessing regardless of chunk layout. ``sequential`` draws without + replacement from the available positions; its result can depend on chunk layout. + + Per-group sampling delegates to _stratified_subsample_indices(). Fractions use full group counts before rounding + quotas; fixed topk caps select directly within chunks. Dask and multiprocessing share the selection calculation; + Dask combines topk candidates in a task tree, while multiprocessing combines bounded batches in the caller. + Group counts and the final sample must fit in memory. + + Sampling uses group membership and the user mask, then selects the same locations for every value array. + Missing observations are still handled separately for each value. Dask computes the selected arrays together; + the reduced sample is then calculated eagerly for both Dask and multiprocessing inputs. Without subsampling, + _reduce_values() uses their shared block calculations, with scheduling differences documented in + _aggregate_chunked(). + + References + ---------- + The overall calculation follows the split-apply-combine strategy described by Wickham (2011), + "The Split-Apply-Combine Strategy for Data Analysis": + https://doi.org/10.18637/jss.v040.i01 + + The pairwise mean and variance merge follows Chan, Golub and LeVeque (1983), + "Algorithms for Computing the Sample Variance: Analysis and Recommendations": + https://doi.org/10.1080/00031305.1983.10483115 + + The key mixer used by ``topk`` follows SplitMix64 from Steele, Lea and Flood (2014), + "Fast Splittable Pseudorandom Number Generators": + https://doi.org/10.1145/2660193.2660195 + + :param source: Input raster or point cloud to summarize or group. + :param statistics: Statistics to calculate (e.g. "mean", ["mean", "nmad"], or np.nanmedian). None returns + "min", "max", "mean", "median", "std", "nmad", "validcount", "totalcount" and "percentagevalidpoints". + "all" also includes "sum", "sumofsquares", "90thpercentile", "iqr", "le90" and "rmse", plus inlier counts + for masked global statistics. Grouped defaults replace "validcount" with "count"; every grouped result + includes "count". + :param by: Named variables to group by (e.g. {"elevation": dem}); use {"glacier": (outlines, "id")} + for vector zones. Arrays must match the source input shape. Omit for global statistics. + :param values: Bands (e.g. [1, 3]) or point columns (e.g. "height") to summarize; defaults to all raster bands + or the main point data column. Use a mapping to name your band inputs (e.g. {"elevation": (dem, 1)}). + :param bins: Continuous bins keyed by grouping name (e.g. {"elevation": 10}). Each definition is a count of + equal-width bins, increasing edges (e.g. [0, 2, 5]), or a Pandas IntervalIndex to choose open/closed sides. + :param categories: Ordered categories keyed by grouping name (e.g. {"landcover": [100, 110, 120]}). + Values outside these categories are excluded. + :param at: Grid or ordered point locations on which to calculate statistics (e.g. at=reference or at="self"). + Defaults to the first point input, if present, otherwise the source locations. + :param mask: Locations to include (True in a boolean mask, e.g. mask=dem > 1000, or features in a vector mask). + Global counts describe values before this mask; "all" adds counts for values kept by the mask. + :param mask_mode: Keep locations "inside" or "outside" vector features; ignored for boolean masks. + :param subsample: Fraction (e.g. 0.1 for 10%) or maximum count (e.g. 10000) of eligible locations to use. + A value of 1 keeps all locations. Counts describe the sampled locations. + :param subsample_per_group: Whether to apply subsample within each combined group (True, stratified sampling) or + once across all groups (False). Without ``by``, both use one global sample. + :param random_state: Seed to reproduce subsampling (e.g. 42), or an existing random generator. + :param strategy: Combine chunk statistics for all groups ("dense"), only groups present in each chunk + ("sparse"), or gather each complete group ("groupwise"). "auto" chooses from the statistics and group count; + exact quantiles and custom functions require "auto" or "groupwise" for chunked data. + :param backend: Use the GeoUtils reducer ("geoutils") or optional Flox reducer ("flox") for grouped statistics. + Flox cannot return group masks, sample within groups, use multiprocessing, or calculate custom functions and + NMAD. Its Dask path also excludes exact medians and percentiles. + :param subsampling_strategy: "topk" keeps the same sampled locations across chunk layouts for a fixed seed; + "sequential" draws random locations using traversal order and can depend on the chunks. + :param interpolation: Raster values at point locations use interp_points() with SciPy methods "nearest", + "linear", "slinear", "cubic", "quintic", "pchip" or "splinef2d". + Raster groupers listed in categories use "nearest". + :param align: "raise" rejects different grids or coordinate systems; "reproject" aligns them to the output + locations. Point inputs must still share the same ordered coordinates. + :param observed: Omit declared group combinations with no eligible locations (True), or include them (False). + :param return_masks: Also return masks keyed by group labels (e.g. table, masks = raster.stats(...)). + Masks cover complete groups before subsampling. Requires by. + :param mp_config: Worker and tile settings for multiprocessing, e.g. MultiprocConfig(chunks=512). + Cannot be combined with Dask inputs. + :returns: A statistic, summary dictionary, grouped dataframe, or grouped dataframe and mask mapping. """ - if is_dask_array(data): - return _statistics_dask(data=data, stats_name=stats_name, counts=counts) - - if np.ma.isMaskedArray(data): - - # Count non zero and not masked pixels in the input data - final_count_nonzero = np.count_nonzero(~np.ma.getmaskarray(data)) - - # Compute valid count from non zero and not masked pixels in the input data - # beforehand saved in counts[0] in case of a inler_mask parameter in get_stats() - valid_count = final_count_nonzero if counts is None else counts[0] - - stats_dict = { - "mean": np.ma.mean, - "median": np.ma.median, - "max": np.ma.max, - "min": np.ma.min, - "sum": np.ma.sum, - "sumofsquares": sum_square, - "90thpercentile": partial(lambda x: mquantiles(x, prob=0.9, alphap=1, betap=1)[0]), - "le90": partial(linear_error, interval=90), - "iqr": partial(iqr, nan_policy="omit"), # ignore masked value (nan), - "nmad": nmad, - "rmse": rmse, - "std": np.ma.std, - } # type: ignore - - else: - # Count non zero pixels in the input data - final_count_nonzero = np.count_nonzero(np.isfinite(data)) - - # Compute valid count from non zero and not masked pixels in the input data - # beforehand saved in counts[0] in case of a inler_mask parameter in get_stats() - valid_count = final_count_nonzero if counts is None else counts[0] - - stats_dict = { - "mean": np.nanmean, - "median": np.nanmedian, - "max": np.nanmax, - "min": np.nanmin, - "sum": np.nansum, - "sumofsquares": sum_square, - "90thpercentile": partial(np.nanpercentile, q=90), - "le90": partial(linear_error, interval=90), - "iqr": partial(iqr, nan_policy="omit"), # ignore masked value (nan), - "nmad": nmad, - "rmse": rmse, - "std": np.nanstd, - } # type: ignore - - # Pixels counts - stats_dict.update( - { - "validcount": valid_count, - "totalcount": data.size, - "percentagevalidpoints": (valid_count / data.size) * 100 if data.size else np.nan, - } + # 1/ Validate user inputs + if strategy not in {"auto", "dense", "sparse", "groupwise"}: + raise ValueError("Argument ``strategy`` must be 'auto', 'dense', 'sparse' or 'groupwise'.") + if backend not in {"geoutils", "flox"}: + raise ValueError("Argument ``backend`` must be 'geoutils' or 'flox'.") + if subsampling_strategy not in {"sequential", "topk"}: + raise ValueError("Argument ``subsampling_strategy`` must be 'sequential' or 'topk'.") + if not isinstance(subsample, (int, float)) or subsample <= 0: + raise ValueError("Argument ``subsample`` must be a positive number.") + if not isinstance(subsample_per_group, (bool, np.bool_)): + raise TypeError("Argument ``subsample_per_group`` must be a boolean.") + if by is None and (bins is not None or categories is not None or not observed or return_masks): + raise ValueError( + "Argument ``by`` is required for ``bins``, ``categories``, ``observed``=False or ``return_masks``=True." + ) + if backend == "flox": + if by is None: + raise ValueError("The Flox backend requires grouped statistics through argument ``by``.") + if subsample_per_group or return_masks or mp_config is not None or strategy != "auto": + raise ValueError( + "The Flox backend requires ``subsample_per_group``=False, ``return_masks``=False, " + "``mp_config``=None and ``strategy``='auto'." + ) + import_optional("flox", extra_name="flox") + + # Warn before spatial selection reads values from raster or point cloud containers + spatial_inputs = [source, at, mask, *by.values(), *(values.values() if isinstance(values, Mapping) else [])] + spatial_inputs = [value[0] if isinstance(value, tuple) and value else value for value in spatial_inputs] + if any(_is_raster(value) or _is_pointcloud(value) for value in spatial_inputs): + warnings.warn( + "The Flox backend loads Raster and PointCloud inputs before calculating grouped statistics.", + category=UserWarning, + stacklevel=2, + ) + if ( + by is None + and statistics is not None + and not isinstance(statistics, (str, Iterable)) + and not callable(statistics) + ): + warnings.warn(f"Statistic name {statistics} is a not recognized string", category=UserWarning) + return None + + # Validate statistics and group inputs + normalized_statistics = _normalize_statistics( + statistics, grouped=by is not None, masked=by is None and mask is not None ) - - if counts is not None: - stats_dict.update( - { - "validinliercount": final_count_nonzero, - "totalinliercount": counts[1], - "percentageinlierpoints": (final_count_nonzero / counts[0]) * 100, - "percentagevalidinlierpoints": (final_count_nonzero / counts[1]) * 100 if counts[1] != 0 else 0, - } + definitions = None if by is None else _validate_group_declarations(by, bins, categories) + + # 2/ Dispatch to input preparation and global/grouped statistics subfunctions + with ExitStack() as stack: + # Inspect all inputs to choose the common "support" (raster grid or point locations), then reproject + # the values from source + the mask on it, as those are used by both global/grouped stats + # (But we don't reproject grouping variables for now, this is done in _grouped_stats below) + values_at_support, support_mask, support = _select_values_and_mask_at_support( + source, + by=by, + values=values, + at=at, + mask=mask, + mask_mode=mask_mode, + interpolation=interpolation, + align=align, + mp_config=mp_config, + stack=stack, ) - # If there are no valid data points, raise a warning - if final_count_nonzero == 0: - warnings.warn("Empty raster, returns Nan for all stats", category=UserWarning) - - if stats_name is None or stats_name == "all": - stat_names_res = _get_default_stat_names(stats_name, counts) - - res_dict = { - _STATS_ALIAS_ALL[stat_name]: ( - stats_dict[stat_name](data) # type: ignore - if (callable(stats_dict[stat_name]) and final_count_nonzero != 0) - else ( - np.nan - # If there are no valid data points, set callable statistics to NaN - if (callable(stats_dict[stat_name]) and final_count_nonzero == 0) - else stats_dict[stat_name] - ) + # If no grouping, compute global stats + if by is None: + return _global_stats( + values_at_support, + normalized_statistics, + mask=support_mask, + subsample=subsample, + random_state=random_state, + strategy=strategy, + subsampling_strategy=subsampling_strategy, + mp_config=mp_config, ) - for stat_name in stat_names_res - } # type: ignore - - else: - res_dict = {} # type: ignore - for stat_name in stats_name: - - # Compute stat if in stats_dict keys - if isinstance(stat_name, str): - - # Get common alias - stat_common_alias = _get_stat_common_alias(stat_name, stats_dict) - if stat_common_alias is not None: - if stat_common_alias in stats_dict: - res_dict[stat_name] = stats_dict[stat_common_alias] - if callable(res_dict[stat_name]): - if final_count_nonzero == 0: - res_dict[stat_name] = np.nan - else: - res_dict[stat_name] = res_dict[stat_name](data) # type: ignore - else: - res_dict[stat_name] = np.nan - else: - warnings.warn("Statistic name " + stat_name + " is not recognized", category=UserWarning) - res_dict[stat_name] = np.float32(np.nan) # type: ignore - - # Compute stat if callable - elif callable(stat_name): - res_dict[stat_name.__name__] = stat_name(data) # type: ignore - - else: - warnings.warn("Statistic name " + stat_name + " is not recognized", category=UserWarning) - res_dict[stat_name] = np.float32(np.nan) # type: ignore - - return {k: (v.item() if isinstance(v, np.generic) else v) for k, v in res_dict.items()} # type: ignore + + # Otherwise perform grouping + assert definitions is not None + return _grouped_stats( + source, + by, + values=values_at_support, + support=support, + statistics=normalized_statistics, + mask=support_mask, + subsample=subsample, + subsample_per_group=subsample_per_group, + random_state=random_state, + strategy=strategy, + backend=backend, + subsampling_strategy=subsampling_strategy, + interpolation=interpolation, + align=align, + observed=observed, + return_masks=return_masks, + mp_config=mp_config, + definitions=definitions, + stack=stack, + ) + + +def variogram( + source: RasterBase | PointCloudBase, + *, + band: int | None = None, + n_pairs: int = 1_000_000, + sampling: Literal["loglag", "random_xy"] = "loglag", + estimator: str | Callable[[NDArrayNum], float] = "dowd", + bins: Literal["log", "uniform"] | Iterable[float] = "log", + n_lags: int = 24, + min_lag: float | None = None, + max_lag: float | None = None, + n_runs: int = 1, + model: str | Callable[..., Any] | list[str | Callable[..., Any]] | None = None, + fit_kwargs: dict[str, Any] | None = None, + random_state: int | np.random.Generator | None = None, + mask: RasterLike | VectorLike | ArrayLike | None = None, + **pair_sampling_kwargs: Any, +) -> Variogram: + """ + Estimate a variogram from spatial pairs efficiently sampled from a raster or point cloud. + + The output Variogram object can be easily re-fit with ``fit()``, plotted with ``plot()`` or exported + to various formats (e.g., GSTools, GPyTorch) with for example ``to_gstools()``. It can also be passed to + resampling and gridding function supporting kriging in GeoUtils. + + The empirical variogram's ``sampling`` uses pairwise logarithmic distance "loglag" by default to efficiently + capture both short distances and large distances on large datasets, which typically outperforms random + endpoints selection ("random_xy"). + Both options support out-of-memory execution through Dask for raster inputs. + + SciKit-GStat supplies the semivariance estimators and theoretical model formulas. GeoUtils fits the chosen + model to the measured distance bins with SciPy curve_fit(). + + :param source: Raster or point cloud whose spatial variability is estimated. + :param band: Raster band to sample, counting from 1. Omit for a point cloud. + :param n_pairs: Number of finite pairs targeted in each run (e.g. 100_000). + :param sampling: How to select pairs: ``"loglag"`` balances short and long distances, while ``"random_xy"`` + selects endpoints independently. + :param estimator: Semivariance estimator from SciKit-GStat: ``"dowd"``, ``"matheron"``, ``"cressie"``, + ``"genton"``, ``"minmax"``, ``"entropy"`` or ``"percentile"``. A function can instead map absolute pair + differences to one value per distance bin. + :param bins: Distance bins: ``"log"`` for logarithmic spacing, ``"uniform"`` for equal widths, or explicit + edges (e.g. [1, 10, 100]). + :param n_lags: Number of distance bins when bins is ``"log"`` or ``"uniform"``. + :param min_lag: Minimum sampled distance in CRS units; defaults to the smaller pixel spacing or half the estimated + point spacing. + :param max_lag: Maximum sampled distance in CRS units; defaults to the diagonal of the source locations. + :param n_runs: Independent samples to average; repeat sampling to estimate each distance bin's standard error. + :param model: SciKit-GStat model to fit: ``"spherical"``, ``"exponential"``, ``"gaussian"``, ``"cubic"``, + ``"stable"`` or ``"matern"``, or the corresponding model function. Sum a list of models ordered from short + to long range (e.g. ["spherical", "exponential"]). ``None`` keeps only the empirical variogram. + :param fit_kwargs: Options for Variogram.fit(): ``use_nugget``, ``bounds``, ``p0`` or ``maxfev`` + (e.g. {"use_nugget": True}); optimization uses SciPy curve_fit(). + :param random_state: Seed or NumPy generator for reproducible sampling across runs (e.g. 42). + :param mask: Locations to keep: True values in a boolean mask, an aligned mask raster for raster inputs, + or locations inside vector geometries. + :param pair_sampling_kwargs: Extra pairsample() options for the source (e.g. ``strategy`` or ``max_rounds``). + :returns: Variogram with distance bins, pair counts and semivariance, plus sampling errors and a fitted model + when requested. + """ + + # Keep raster band selection out of point pair sampling kwargs + pair_kwargs: dict[str, Any] = {} if band is None else {"band": band} + pair_kwargs.update( + { + "n_pairs": n_pairs, + "sampling": sampling, + "min_distance": min_lag, + "max_distance": max_lag, + "mask": mask, + } + ) + pair_kwargs.update(pair_sampling_kwargs) + + return _estimate_variogram( + source, + n_runs=n_runs, + estimator=estimator, + bins=bins, + n_lags=n_lags, + min_lag=min_lag, + max_lag=max_lag, + models=model, + fit_kwargs=fit_kwargs, + random_state=random_state, + pair_kwargs=pair_kwargs, + ) diff --git a/geoutils/stats/variography.py b/geoutils/stats/variography.py new file mode 100644 index 000000000..3dca3a393 --- /dev/null +++ b/geoutils/stats/variography.py @@ -0,0 +1,1467 @@ +# Copyright (c) 2026 GeoUtils developers +# +# This file is part of the GeoUtils project: +# https://github.com/glaciohack/geoutils +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Estimate, fit, and convert variograms across Python packages for inter-operability.""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import pandas as pd +import xarray as xr +from scipy.optimize import curve_fit + +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum + +if TYPE_CHECKING: + from geoutils.pointcloud.base import PointCloudBase + from geoutils.raster.base import RasterBase + +__all__ = ["Variogram"] + + +############################# +# 1/ VARIOGRAM MODEL METADATA +############################# + +_BASE_MODELS = {"spherical", "exponential", "gaussian", "cubic", "stable", "matern"} +_COMPOSITE_MODELS = {"sum", "product"} + + +@dataclass(frozen=True) +class VariogramModel: + """ + Parameters of a fitted theoretical variogram in a form shared by all supported packages: GSTools, GPyTorch, and + SciKit-GStat. + + The ``effective_range`` follows SciKit-GStat's convention (as numerical ranges are defined differently between + packages), and ``partial_sill`` excludes the nugget, making the conversion to covariance kernels unambiguous. + A composite model holds its independent structures in ``components`` and keeps their shared nugget on the parent + model. + + :param model_name: Base model name or ``"sum"``/``"product"`` for a composition. + :param effective_range: Distance at which a base model effectively reaches its sill. + :param partial_sill: Structured variance excluding the nugget. + :param nugget: Uncorrelated variance added at positive distances. + :param smoothness: Matérn smoothness parameter. + :param shape: Stable model shape parameter. + :param active_dims: Feature columns used by a converted covariance kernel. + :param components: Base structures contained by a composite model. + """ + + model_name: str + effective_range: float | None = None + partial_sill: float | None = None + nugget: float = 0.0 + smoothness: float | None = None + shape: float | None = None + active_dims: tuple[int, ...] | None = None + components: tuple[VariogramModel, ...] = () + + ##################### + # MODEL VALIDATION + ##################### + + def __post_init__(self) -> None: + # Replace accepted short names with the one standard name used by every conversion + aliases = { + "cub": "cubic", + "exp": "exponential", + "gau": "gaussian", + "mat": "matern", + "rbf": "gaussian", + "sph": "spherical", + "sta": "stable", + } + model_name = aliases.get(self.model_name.lower(), self.model_name.lower()) + object.__setattr__(self, "model_name", model_name) + object.__setattr__(self, "components", tuple(self.components)) + + # Convert numeric fields once so loaded and newly fitted models behave alike + for name in ("effective_range", "partial_sill", "nugget", "smoothness", "shape"): + value = getattr(self, name) + if value is not None: + object.__setattr__(self, name, float(value)) + if self.active_dims is not None: + object.__setattr__(self, "active_dims", tuple(int(value) for value in self.active_dims)) + if self.partial_sill is not None and (not np.isfinite(self.partial_sill) or self.partial_sill < 0): + raise ValueError("Variogram partial sill must be a finite, non-negative number.") + + # Check combined models and calculate their sill from their component models + if model_name in _COMPOSITE_MODELS: + if len(self.components) < 2: + raise ValueError("A composite variogram requires at least two components.") + if any(component.model_name == model_name for component in self.components): + raise ValueError(f"Nested {model_name} variograms are not supported; flatten the components first.") + if any(component.nugget != 0 for component in self.components): + raise ValueError("Composite components must omit nuggets; set one shared nugget on the parent model.") + if self.partial_sill is None: + partial_sills = [component.partial_sill or 0.0 for component in self.components] + combined_sill = sum(partial_sills) if model_name == "sum" else np.prod(partial_sills) + object.__setattr__(self, "partial_sill", float(combined_sill)) + elif model_name not in _BASE_MODELS: + raise ValueError(f"Unsupported variogram model {model_name!r}.") + else: + # Require the range, variance, and optional shape values needed by this model type + if self.components: + raise ValueError("Only a composite variogram can contain components.") + if self.effective_range is None or not np.isfinite(self.effective_range) or self.effective_range <= 0: + raise ValueError("Variogram effective range must be a finite, strictly positive number.") + if self.partial_sill is None: + raise ValueError("A base variogram requires a partial sill.") + if model_name == "matern" and self.smoothness is None: + raise ValueError("A Matérn variogram requires a smoothness parameter.") + if model_name == "stable" and self.shape is None: + raise ValueError("A stable variogram requires a shape parameter.") + + # Check the shared numeric fields after model-specific defaults are applied + if not np.isfinite(self.nugget) or self.nugget < 0: + raise ValueError("Variogram nugget must be a finite, non-negative number.") + if self.smoothness is not None and (not np.isfinite(self.smoothness) or self.smoothness <= 0): + raise ValueError("Variogram smoothness must be a finite, strictly positive number.") + if self.shape is not None and (not np.isfinite(self.shape) or self.shape <= 0): + raise ValueError("Variogram shape must be a finite, strictly positive number.") + if self.active_dims is not None and ( + len(self.active_dims) == 0 + or len(set(self.active_dims)) != len(self.active_dims) + or min(self.active_dims) < 0 + ): + raise ValueError("Argument ``active_dims`` must contain unique, non-negative dimensions.") + + #################### + # MODEL EVALUATION + #################### + + @property + def sill(self) -> float: + """Total sill, including the nugget.""" + + return float((self.partial_sill or 0.0) + self.nugget) + + def variogram(self, distance: NDArrayNum | float) -> NDArrayNum: + """ + Evaluate the theoretical variogram at one or more distances. + + :param distance: Spatial distance or array of distances. + :returns: Semivariance at each distance. + """ + + # Use one array calculation, then return a scalar when the input was scalar + scalar_input = np.ndim(distance) == 0 + distances = np.atleast_1d(np.asarray(distance, dtype=float)) + + # Add component semivariances, then add the shared nugget once + if self.model_name == "sum": + values = sum((component.variogram(distances) for component in self.components), np.zeros_like(distances)) + output = values + np.where(distances > 0, self.nugget, 0.0) + return output[0] if scalar_input else output + + # Multiply component covariances, then convert the result back to semivariance + if self.model_name == "product": + if self.partial_sill is None: + raise AssertionError("A product variogram model must define its partial sill.") + covariance = np.ones_like(distances) + for component in self.components: + covariance *= component.covariance(distances) + output = (float(self.partial_sill) - covariance) + np.where(distances > 0, self.nugget, 0.0) + return output[0] if scalar_input else output + + # Confirm that validation supplied the numeric fields required by SciKit-GStat + if self.effective_range is None or self.partial_sill is None: + raise AssertionError("A base variogram model must define its range and partial sill.") + + # Pass parameters in the order expected by the selected SciKit-GStat model function + skgstat = import_optional("skgstat", package_name="scikit-gstat", extra_name="geostat") + model_function = getattr(skgstat.models, self.model_name) + arguments: list[float] = [self.effective_range, self.partial_sill] + if self.model_name == "matern": + if self.smoothness is None: + raise AssertionError("A Matérn variogram model must define its smoothness.") + arguments.append(float(self.smoothness)) + elif self.model_name == "stable": + if self.shape is None: + raise AssertionError("A stable variogram model must define its shape.") + arguments.append(float(self.shape)) + + # Calculate the distance-dependent part and add the nugget above zero distance + values = np.asarray(model_function(distances.ravel(), *arguments), dtype=float).reshape(distances.shape) + output = values + np.where(distances > 0, self.nugget, 0.0) + return output[0] if scalar_input else output + + def covariance(self, distance: NDArrayNum | float) -> NDArrayNum: + """ + Evaluate covariance implied by this variogram. + + :param distance: Spatial distance or array of distances. + :returns: Covariance at each distance. + """ + + return self.sill - self.variogram(distance) + + def correlation(self, distance: NDArrayNum | float) -> NDArrayNum: + """ + Evaluate correlation implied by this variogram. + + :param distance: Spatial distance or array of distances. + :returns: Correlation at each distance. + """ + + if self.sill == 0: + raise ValueError("A variogram with zero sill does not define correlation.") + return self.covariance(distance) / self.sill + + ################### + # MODEL COMPOSITION + ################### + + @classmethod + def sum(cls, components: Sequence[VariogramModel], nugget: float = 0.0) -> VariogramModel: + """ + Combine independent nested structures into a summed variogram model. + + :param components: Fitted structures to add. + :param nugget: Shared uncorrelated variance. + :returns: Summed model with normalized components. + """ + + return cls.combine(components, combination="sum", nugget=nugget) + + @classmethod + def combine( + cls, + components: Sequence[VariogramModel], + *, + combination: str, + nugget: float = 0.0, + ) -> VariogramModel: + """ + Combine independently parameterized components by addition or multiplication. + + :param components: Fitted structures to combine. + :param combination: Either ``"sum"`` or ``"product"``. + :param nugget: Shared uncorrelated variance. + :returns: Composite model with normalized components. + """ + + if combination not in _COMPOSITE_MODELS: + raise ValueError("Argument ``combination`` must be 'sum' or 'product'.") + + # Flatten nested combinations of the same kind and store all nuggets once on the parent + flattened: list[VariogramModel] = [] + for component in components: + if component.model_name == combination: + flattened.extend(component.components) + else: + flattened.append(replace(component, nugget=0.0)) + return cls(model_name=combination, nugget=nugget, components=tuple(flattened)) + + ################# + # SERIALIZATION + ################# + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-compatible model description.""" + + # Convert nested components too so the complete model contains only plain Python data + return { + "model_name": self.model_name, + "effective_range": self.effective_range, + "partial_sill": self.partial_sill, + "nugget": self.nugget, + "smoothness": self.smoothness, + "shape": self.shape, + "active_dims": self.active_dims, + "components": [component.to_dict() for component in self.components], + } + + @classmethod + def from_dict(cls, values: Mapping[str, Any]) -> VariogramModel: + """ + Restore a model from VariogramModel.to_dict() output. + + :param values: Serialized model fields. + :returns: Restored fitted model. + """ + + # Restore nested components before the dataclass checks the complete parent model + components = tuple(cls.from_dict(component) for component in values.get("components", ())) + return cls( + model_name=str(values["model_name"]), + effective_range=values.get("effective_range"), + partial_sill=values.get("partial_sill"), + nugget=float(values.get("nugget", 0.0)), + smoothness=values.get("smoothness"), + shape=values.get("shape"), + active_dims=( + None if values.get("active_dims") is None else tuple(int(value) for value in values["active_dims"]) + ), + components=components, + ) + + +################################ +# 2/ RESULTS FROM OTHER PACKAGES +################################ + + +@dataclass(frozen=True) +class GPyTorchVariogram: + """ + A GPyTorch covariance kernel and its separate observation noise nugget. + + :param kernel: Converted GPyTorch covariance kernel. + :param noise: Nugget variance for the observation likelihood. + """ + + kernel: Any + noise: float + + +@dataclass(frozen=True) +class GSToolsVariogram: + """ + A GSTools covariance model and the source coordinate columns it uses. + + :param model: Converted GSTools covariance model. + :param active_dims: Feature columns to select before passing coordinates to GSTools. + """ + + model: Any + active_dims: tuple[int, ...] | None + + +############################ +# 3/ MEASURED VARIOGRAM RESULT +############################ + + +@dataclass(frozen=True) +class Variogram: + """ + Measured and fitted variogram values that can be saved without a fitting package object. + + The arrays have one value per distance bin. ``backend_object`` is absent by default because a SciKit-GStat Variogram + retains sampled coordinates, pairwise distances and pairwise differences. Pass ``keep_backend=True`` during + estimation or conversion only when its extra methods are worth the additional memory. + + :param lags: Mean sampled distance in each distance bin. + :param semivariance: Measured semivariance in each distance bin. + :param counts: Number of sampled pairs in each distance bin. + :param semivariance_error: Sampling error estimated across independent runs. + :param bin_lower_edges: Inclusive lower boundaries of distance bins. + :param bin_edges: Upper boundaries of distance bins. + :param fitted_semivariance: Fitted values evaluated at ``lags``. + :param model: Portable fitted model parameters. + :param estimator: Name of the formula used to estimate semivariance. + :param distance: Distance measure name. + :param binning: Method used to build distance bins. + :param backend: Package used for fitting or import. + :param backend_object: Optional retained object from that package. + :param fit_result: Small details returned by the fit. + :param attrs: Additional serializable metadata. + """ + + lags: NDArrayNum + semivariance: NDArrayNum + counts: NDArrayNum + semivariance_error: NDArrayNum | None = None + bin_lower_edges: NDArrayNum | None = None + bin_edges: NDArrayNum | None = None + fitted_semivariance: NDArrayNum | None = None + model: VariogramModel | None = None + estimator: str | None = None + distance: str | None = None + binning: str | None = None + backend: str | None = None + backend_object: Any = field(default=None, repr=False, compare=False) + fit_result: Mapping[str, Any] = field(default_factory=dict, compare=False) + attrs: Mapping[str, Any] = field(default_factory=dict, compare=False) + + ################### + # RESULT VALIDATION + ################### + + def __post_init__(self) -> None: + # Copy required distance bin arrays so the result owns its small stored data + lags = np.asarray(self.lags, dtype=float).copy() + semivariance = np.asarray(self.semivariance, dtype=float).copy() + counts = np.asarray(self.counts, dtype=np.int64).copy() + if lags.ndim != 1 or semivariance.ndim != 1 or counts.ndim != 1: + raise ValueError("Variogram lag statistics must be one-dimensional.") + if not (len(lags) == len(semivariance) == len(counts)): + raise ValueError("Variogram lags, semivariances and counts must have equal lengths.") + if np.any(counts < 0): + raise ValueError("Variogram pair counts cannot be negative.") + + # Check optional arrays against the same number of distance bins + optional_arrays: dict[str, NDArrayNum | None] = { + "semivariance_error": self.semivariance_error, + "bin_lower_edges": self.bin_lower_edges, + "bin_edges": self.bin_edges, + "fitted_semivariance": self.fitted_semivariance, + } + normalized: dict[str, NDArrayNum | None] = {} + for name, values in optional_arrays.items(): + if values is None: + normalized[name] = None + continue + array = np.asarray(values, dtype=float).copy() + if array.ndim != 1 or len(array) != len(lags): + raise ValueError(f"Argument ``{name}`` must be one-dimensional and aligned with ``lags``.") + normalized[name] = array + + # Make result arrays read-only so later edits cannot disagree with the fitted model + for array in (lags, semivariance, counts, *[value for value in normalized.values() if value is not None]): + array.setflags(write=False) + object.__setattr__(self, "lags", lags) + object.__setattr__(self, "semivariance", semivariance) + object.__setattr__(self, "counts", counts) + object.__setattr__(self, "semivariance_error", normalized["semivariance_error"]) + object.__setattr__(self, "bin_lower_edges", normalized["bin_lower_edges"]) + object.__setattr__(self, "bin_edges", normalized["bin_edges"]) + object.__setattr__(self, "fitted_semivariance", normalized["fitted_semivariance"]) + object.__setattr__(self, "fit_result", dict(self.fit_result)) + object.__setattr__(self, "attrs", dict(self.attrs)) + + ############################ + # CONSTRUCTION AND FITTING + ############################ + + @classmethod + def estimate( + cls, + coordinates: NDArrayNum, + values: NDArrayNum, + *, + model: str = "spherical", + active_dims: tuple[int, ...] | None = None, + keep_backend: bool = False, + **kwargs: Any, + ) -> Variogram: + """ + Estimate a variogram from available point coordinates using SciKit-GStat. + + :param coordinates: Observation coordinates arranged by row. + :param values: One value per observation. + :param model: SciKit-GStat theoretical model to fit. + :param active_dims: Feature columns used by later covariance conversion. + :param keep_backend: Whether to retain the SciKit-GStat object, which can use substantial memory. + :param kwargs: Additional SciKit-GStat Variogram options. + :returns: Per-distance measurements and fitted parameters without the source object by default. + """ + + # Import SciKit-GStat only for this direct coordinate-based estimate + skgstat = import_optional("skgstat", package_name="scikit-gstat", extra_name="geostat") + + # Put coordinates into one row per observation and check value length + coordinates_array = np.asarray(coordinates, dtype=float) + values_array = np.asarray(values, dtype=float).squeeze() + if coordinates_array.ndim == 1: + coordinates_array = coordinates_array[:, np.newaxis] + if coordinates_array.ndim != 2 or values_array.ndim != 1: + raise ValueError( + "Argument ``coordinates`` must be (observation, feature) and ``values`` must be one-dimensional." + ) + if len(coordinates_array) != len(values_array): + raise ValueError("Arguments ``coordinates`` and ``values`` must contain the same number of observations.") + + # Remove rows with missing coordinates or values before creating the SciKit-GStat object + valid = np.isfinite(values_array) & np.all(np.isfinite(coordinates_array), axis=1) + if np.count_nonzero(valid) < 2: + raise ValueError("At least two finite observations are required to estimate a variogram.") + backend_variogram = skgstat.Variogram(coordinates_array[valid], values_array[valid], model=model, **kwargs) + + # Copy the small distance bin result and discard the larger SciKit-GStat object by default + return cls.from_skgstat( + backend_variogram, + active_dims=active_dims, + keep_backend=keep_backend, + ) + + @classmethod + def from_model( + cls, + model_name: str, + effective_range: float, + partial_sill: float, + *, + nugget: float = 0.0, + smoothness: float | None = None, + shape: float | None = None, + active_dims: tuple[int, ...] | None = None, + ) -> Variogram: + """ + Create a variogram from known fitted parameters without measured values. + + :param model_name: Supported theoretical model name. + :param effective_range: Distance at which the model effectively reaches its sill. + :param partial_sill: Structured variance excluding the nugget. + :param nugget: Uncorrelated variance. + :param smoothness: Matérn smoothness parameter. + :param shape: Stable model shape parameter. + :param active_dims: Feature columns used by later covariance conversion. + :returns: Variogram containing only fitted parameters. + """ + + # Use the standard Matérn smoothness when the caller supplies no value + if model_name.lower() == "matern" and smoothness is None: + smoothness = 1.5 + + # Use empty measured arrays because this result contains only a supplied model + return cls( + lags=np.empty(0), + semivariance=np.empty(0), + counts=np.empty(0, dtype=np.int64), + model=VariogramModel( + model_name=model_name, + effective_range=effective_range, + partial_sill=partial_sill, + nugget=nugget, + smoothness=smoothness, + shape=shape, + active_dims=active_dims, + ), + ) + + @classmethod + def combine( + cls, + *variograms: Variogram, + combination: str = "sum", + nugget: float = 0.0, + ) -> Variogram: + """ + Combine fitted structures for covariance conversion. + + :param variograms: Two or more lightweight variograms with fitted models. + :param combination: Either ``"sum"`` or ``"product"``. + :param nugget: Shared uncorrelated variance. + :returns: Lightweight variogram containing the composite model. + """ + + # Require fitted models because measured distance bins alone do not define covariance at every distance + if len(variograms) < 2 or any(variogram.model is None for variogram in variograms): + raise ValueError("At least two variograms with fitted models are required.") + + # Combine only fitted model parameters and leave out unrelated measured bins + models = tuple(variogram.model for variogram in variograms if variogram.model is not None) + return cls( + lags=np.empty(0), + semivariance=np.empty(0), + counts=np.empty(0, dtype=np.int64), + model=VariogramModel.combine(models, combination=combination, nugget=nugget), + ) + + @classmethod + def from_pairs( + cls, + pairs: xr.Dataset, + *, + estimator: str | Callable[[NDArrayNum], float] = "dowd", + bins: Literal["log", "uniform"] | Iterable[float] = "log", + n_lags: int = 24, + min_lag: float | None = None, + max_lag: float | None = None, + ) -> Variogram: + """ + Calculate measured variogram values from a pair dataset. + + This method only reads the pair distances and endpoint values. The pair dataset can therefore be discarded as + soon as the per-distance statistics have been computed. + + :param pairs: Dataset returned by Raster.pairsample() or PointCloud.pairsample(). + :param estimator: SciKit-GStat estimator name or a function accepting absolute pair differences. + :param bins: ``"log"``, ``"uniform"`` or explicit distance boundaries. + :param n_lags: Number of distance bins used for named binning. + :param min_lag: Lower distance boundary. Defaults to the smallest positive pair distance. + :param max_lag: Upper distance boundary. Defaults to the largest pair distance. + :returns: Per-distance variogram values with no retained pair data. + """ + + # Check the expected Xarray pair layout before reading endpoint values + required = {"distance", "value"} + if not isinstance(pairs, xr.Dataset) or not required.issubset(pairs.data_vars): + raise TypeError("Argument ``pairs`` must be an Xarray Dataset containing 'distance' and 'value'.") + if pairs["value"].dims != ("pair", "endpoint") or pairs.sizes.get("endpoint") != 2: + raise ValueError( + "Variable 'value' in argument ``pairs`` must have dimensions ('pair', 'endpoint') of length two." + ) + if pairs["distance"].dims != ("pair",): + raise ValueError("Variable 'distance' in argument ``pairs`` must have dimensions ('pair',).") + + # Calculate the absolute value difference in each pair and remove missing pairs + distances = np.asarray(pairs["distance"], dtype=float) + endpoint_values = np.asarray(pairs["value"], dtype=float) + differences = np.abs(endpoint_values[:, 0] - endpoint_values[:, 1]) + valid = np.isfinite(distances) & np.isfinite(differences) & (distances > 0) + distances, differences = distances[valid], differences[valid] + if distances.size == 0: + raise ValueError("Argument ``pairs`` contains no finite observations with positive distance.") + + # Build log-spaced or equal-width bins, or check the caller's exact bin edges + binning: str + if isinstance(bins, str): + # Use the requested sampling limits when available so repeated runs share the same bins + minimum = float(pairs.attrs.get("min_distance", np.min(distances))) if min_lag is None else float(min_lag) + maximum = float(pairs.attrs.get("max_distance", np.max(distances))) if max_lag is None else float(max_lag) + if not 0 < minimum < maximum: + raise ValueError("Require 0 < ``min_lag`` < ``max_lag``.") + if n_lags < 1 or bins not in {"log", "uniform"}: + raise ValueError("Argument ``bins`` must be 'log' or 'uniform', with ``n_lags`` at least one.") + edges = ( + np.geomspace(minimum, maximum, n_lags + 1) + if bins == "log" + else np.linspace(minimum, maximum, n_lags + 1) + ) + binning = bins + else: + edges = np.asarray(tuple(bins), dtype=float) + if edges.ndim != 1 or len(edges) < 2 or not np.all(np.diff(edges) > 0): + raise ValueError("Argument ``bins`` must contain at least two increasing lag boundaries.") + binning = "explicit" + + # Choose the named semivariance formula or use the caller's function + if callable(estimator): + estimator_function = estimator + estimator_name = getattr(estimator, "__name__", "callable") + else: + skgstat = import_optional("skgstat", package_name="scikit-gstat", extra_name="geostat") + if not hasattr(skgstat.estimators, estimator): + raise ValueError(f"Unknown SciKit-GStat ``estimator`` {estimator!r}.") + estimator_function = getattr(skgstat.estimators, estimator) + estimator_name = estimator + + # Assign each pair to one distance bin, including both outer edges + membership = np.digitize(distances, edges, right=True) - 1 + membership[distances == edges[0]] = 0 + experimental = np.full(len(edges) - 1, np.nan, dtype=float) + counts = np.zeros(len(edges) - 1, dtype=np.int64) + lag_centers = np.full(len(edges) - 1, np.nan, dtype=float) + + # Sort once so each estimator receives a contiguous bin without scanning all pairs again + order = np.argsort(membership, kind="stable") + sorted_membership = membership[order] + boundaries = np.searchsorted(sorted_membership, np.arange(len(edges))) + sorted_distances, sorted_differences = distances[order], differences[order] + for index, (start, stop) in enumerate(zip(boundaries[:-1], boundaries[1:])): + counts[index] = stop - start + if stop > start: + # Keep the original order within each bin, including for user supplied estimators + experimental[index] = float(estimator_function(sorted_differences[start:stop])) + lag_centers[index] = float(np.mean(sorted_distances[start:stop])) + + # Return only per-bin arrays and plain source details that can be saved + return cls( + lags=lag_centers, + semivariance=experimental, + counts=counts, + semivariance_error=np.full(len(experimental), np.nan), + bin_lower_edges=edges[:-1], + bin_edges=edges[1:], + estimator=estimator_name, + distance="euclidean", + binning=binning, + attrs={**pairs.attrs, "pair_count": int(np.sum(counts))}, + ) + + def fit( + self, + models: str | Callable[..., Any] | Sequence[str | Callable[..., Any]] = "spherical", + *, + use_nugget: bool = False, + bounds: Sequence[tuple[float, float]] | None = None, + p0: Sequence[float] | None = None, + maxfev: int | None = None, + ) -> Variogram: + """ + Fit one or more summed theoretical models to the measured bins. + + Finite, positive sampling errors are used as weights. The returned copy retains only fitted parameters and + the small covariance matrix produced by the optimizer. + + :param models: Model name, SciKit-GStat model function or sequence ordered from short to long range. + :param use_nugget: Whether to fit a shared non-negative nugget. + :param bounds: Lower and upper bound for every fitted parameter. + :param p0: Initial parameter values in range/sill order, followed by optional model shape and nugget. + :param maxfev: Maximum number of model evaluations. + :returns: New variogram containing fitted model parameters in the shared form. + """ + + # Turn short names, model functions, or a sequence into standard model names + requested_models: list[str | Callable[..., Any]] = [] + if isinstance(models, str): + requested_models.extend(models.split("+")) + elif callable(models): + requested_models = [models] + else: + requested_models = list(models) + aliases = { + "cub": "cubic", + "exp": "exponential", + "gau": "gaussian", + "mat": "matern", + "sph": "spherical", + "sta": "stable", + } + model_names = [] + for requested in requested_models: + name = requested.strip().lower() if isinstance(requested, str) else getattr(requested, "__name__", "") + model_names.append(aliases.get(name, name)) + if not model_names or any(name not in _BASE_MODELS for name in model_names): + raise ValueError(f"Argument ``models`` must contain names from {sorted(_BASE_MODELS)}.") + + # Fit only bins with measured values and require more bins than fitted parameters + valid = np.isfinite(self.lags) & np.isfinite(self.semivariance) + if np.count_nonzero(valid) < 2: + raise ValueError("At least two finite empirical lag classes are required for fitting.") + + skgstat = import_optional("skgstat", package_name="scikit-gstat", extra_name="geostat") + parameter_counts = [3 if name in {"stable", "matern"} else 2 for name in model_names] + + def summed_model(distance: NDArrayNum, *parameters: float) -> NDArrayNum: + # Add each component's parameters in the order expected by SciPy's curve fit + values = np.zeros_like(np.asarray(distance, dtype=float)) + position = 0 + for name, count in zip(model_names, parameter_counts): + values += getattr(skgstat.models, name)(distance, *parameters[position : position + count]) + position += count + if use_nugget: + values += np.where(np.asarray(distance) > 0, parameters[position], 0.0) + return values + + # Choose starting range and variance values from the measured distances and semivariances + maximum_lag = float(np.nanmax(self.lags[valid])) + maximum_variance = float(np.nanmax(self.semivariance[valid])) + if maximum_variance <= 0: + maximum_variance = 1.0 + if p0 is None: + guesses: list[float] = [] + for index, count in enumerate(parameter_counts, start=1): + guesses.extend((index * maximum_lag / len(model_names), maximum_variance / len(model_names))) + if count == 3: + guesses.append(1.0) + if use_nugget: + guesses.append(maximum_variance * 0.05) + p0 = guesses + + # Check optional starting values and limits against the number of fitted parameters + expected = sum(parameter_counts) + int(use_nugget) + if len(p0) != expected: + raise ValueError(f"Argument ``p0`` must contain {expected} parameters for the selected models.") + + # Keep fitted parameters nonnegative unless the caller supplies other limits + if bounds is None: + model_bounds: list[tuple[float, float]] = [] + for count in parameter_counts: + model_bounds.extend(((float(np.finfo(float).eps), maximum_lag), (0.0, np.inf))) + if count == 3: + model_bounds.append((float(np.finfo(float).eps), np.inf)) + if use_nugget: + model_bounds.append((0.0, np.inf)) + bounds = model_bounds + if len(bounds) != expected: + raise ValueError(f"Argument ``bounds`` must contain {expected} lower/upper pairs for the selected models.") + lower, upper = np.asarray(bounds, dtype=float).T + + # Give bins with smaller measured errors more influence when usable errors exist + errors = None + if self.semivariance_error is not None: + candidate_errors = self.semivariance_error[valid] + if np.any(np.isfinite(candidate_errors) & (candidate_errors > 0)): + positive = np.isfinite(candidate_errors) & (candidate_errors > 0) + replacement = float(np.nanmedian(candidate_errors[positive])) + errors = np.where(positive, candidate_errors, replacement) + + # Fit all summed components together so their parameters can adjust to one another + coefficients, covariance = curve_fit( + summed_model, + self.lags[valid], + self.semivariance[valid], + p0=np.asarray(p0, dtype=float), + bounds=(lower, upper), + sigma=errors, + absolute_sigma=errors is not None, + method="trf", + maxfev=maxfev, + ) + + # Split SciPy's fitted numbers back into the shared component model objects + components: list[VariogramModel] = [] + position = 0 + for name, count in zip(model_names, parameter_counts): + parameters = coefficients[position : position + count] + position += count + components.append( + VariogramModel( + model_name=name, + effective_range=float(parameters[0]), + partial_sill=float(parameters[1]), + smoothness=float(parameters[2]) if name == "matern" else None, + shape=float(parameters[2]) if name == "stable" else None, + ) + ) + + # Store the shared nugget on the single model or the combined parent model + nugget = float(coefficients[position]) if use_nugget else 0.0 + fitted_model = ( + replace(components[0], nugget=nugget) + if len(components) == 1 + else VariogramModel.sum(components, nugget=nugget) + ) + + # Return a new read-only result with fitted values and small fit details + return replace( + self, + fitted_semivariance=fitted_model.variogram(self.lags), + model=fitted_model, + backend="scikit-gstat models/scipy fit", + fit_result={"coefficients": coefficients.tolist(), "covariance": covariance.tolist()}, + ) + + @classmethod + def from_skgstat( + cls, + variogram: Any, + *, + active_dims: tuple[int, ...] | None = None, + keep_backend: bool = False, + ) -> Variogram: + """ + Copy a small result from a fitted SciKit-GStat Variogram. + + :param variogram: Fitted SciKit-GStat Variogram object. + :param active_dims: Feature columns used by later covariance conversion. + :param keep_backend: Whether to retain the input object. + :returns: Per-distance measurements and fitted parameters. + """ + + # Read SciKit-GStat's description once to get its standard model names and settings + description = variogram.describe() + configured_model = str(description.get("params", {}).get("model", description["model"])).lower() + + # Copy fitted numbers into the shared `VariogramModel` form + model = _model_from_skgstat( + variogram, + configured_model, + description, + active_dims=active_dims, + ) + + # Copy per-bin measurements and fitted values from SciKit-GStat + lag_centers, experimental = variogram.get_empirical(bin_center=True) + lag_centers = np.asarray(lag_centers, dtype=float) + fitted = np.asarray(variogram.fitted_model(lag_centers), dtype=float) + fit_result = { + key: description.get(key) + for key in ("normalized_effective_range", "normalized_sill", "normalized_nugget") + if key in description + } + + # Keep the full SciKit-GStat object only when the caller requests it + return cls( + lags=lag_centers, + semivariance=np.asarray(experimental, dtype=float), + counts=np.asarray(variogram.bin_count, dtype=np.int64), + semivariance_error=np.full(len(lag_centers), np.nan), + bin_lower_edges=np.r_[0.0, np.asarray(variogram.bins, dtype=float)[:-1]], + bin_edges=np.asarray(variogram.bins, dtype=float), + fitted_semivariance=fitted, + model=model, + estimator=str(description["estimator"]), + distance=str(description["dist_func"]), + binning=str(description.get("params", {}).get("bin_func", "unknown")), + backend="skgstat", + backend_object=variogram if keep_backend else None, + fit_result=fit_result, + ) + + ################################ + # REPRESENTATION AND STORAGE + ################################ + + def without_backend(self) -> Variogram: + """Return a copy that releases any retained fitting package object.""" + + return replace(self, backend_object=None) + + def to_dataframe(self) -> pd.DataFrame: + """Return one row per distance bin in a Pandas DataFrame.""" + + # Add required distance bin columns first, then optional error, edge, and fitted columns + data: dict[str, Any] = { + "lag": self.lags, + "semivariance": self.semivariance, + "count": self.counts, + } + if self.semivariance_error is not None: + data["semivariance_error"] = self.semivariance_error + if self.bin_lower_edges is not None: + data["bin_lower_edge"] = self.bin_lower_edges + if self.bin_edges is not None: + data["bin_edge"] = self.bin_edges + if self.fitted_semivariance is not None: + data["fitted_semivariance"] = self.fitted_semivariance + return pd.DataFrame(data) + + def to_xarray(self) -> xr.Dataset: + """Return labelled lag statistics in an Xarray Dataset.""" + + # Store all measured arrays on one labelled distance bin dimension + data_vars: dict[str, Any] = { + "semivariance": ("lag", self.semivariance), + "count": ("lag", self.counts), + } + if self.semivariance_error is not None: + data_vars["semivariance_error"] = ("lag", self.semivariance_error) + if self.bin_lower_edges is not None: + data_vars["bin_lower_edge"] = ("lag", self.bin_lower_edges) + if self.bin_edges is not None: + data_vars["bin_edge"] = ("lag", self.bin_edges) + if self.fitted_semivariance is not None: + data_vars["fitted_semivariance"] = ("lag", self.fitted_semivariance) + + # Store model and method descriptions as plain attributes that survive Xarray save and load + attrs = { + **self.attrs, + "estimator": self.estimator or "", + "distance": self.distance or "", + "binning": self.binning or "", + "backend": self.backend or "", + } + if self.model is not None: + attrs["model"] = json.dumps(self.model.to_dict()) + return xr.Dataset(data_vars=data_vars, coords={"lag": self.lags}, attrs=attrs) + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-compatible representation without backend state.""" + + # Convert arrays and nested models to lists and dictionaries accepted by JSON + return { + "lags": self.lags.tolist(), + "semivariance": self.semivariance.tolist(), + "counts": self.counts.tolist(), + "semivariance_error": (None if self.semivariance_error is None else self.semivariance_error.tolist()), + "bin_lower_edges": None if self.bin_lower_edges is None else self.bin_lower_edges.tolist(), + "bin_edges": None if self.bin_edges is None else self.bin_edges.tolist(), + "fitted_semivariance": None if self.fitted_semivariance is None else self.fitted_semivariance.tolist(), + "model": None if self.model is None else self.model.to_dict(), + "estimator": self.estimator, + "distance": self.distance, + "binning": self.binning, + "backend": self.backend, + "fit_result": dict(self.fit_result), + "attrs": dict(self.attrs), + } + + @classmethod + def from_dict(cls, values: Mapping[str, Any]) -> Variogram: + """ + Restore a lightweight variogram from Variogram.to_dict() output. + + :param values: Serialized per-distance measurements, optional fitted model, and metadata. + :returns: Variogram containing the restored measurements and model, without an external backend object. + """ + + # Restore the optional model before the dataclass checks all per-bin arrays + model_values = values.get("model") + return cls( + lags=np.asarray(values["lags"], dtype=float), + semivariance=np.asarray(values["semivariance"], dtype=float), + counts=np.asarray(values["counts"], dtype=np.int64), + semivariance_error=( + None + if values.get("semivariance_error") is None + else np.asarray(values["semivariance_error"], dtype=float) + ), + bin_lower_edges=( + None if values.get("bin_lower_edges") is None else np.asarray(values["bin_lower_edges"], dtype=float) + ), + bin_edges=None if values.get("bin_edges") is None else np.asarray(values["bin_edges"], dtype=float), + fitted_semivariance=( + None + if values.get("fitted_semivariance") is None + else np.asarray(values["fitted_semivariance"], dtype=float) + ), + model=None if model_values is None else VariogramModel.from_dict(model_values), + estimator=values.get("estimator"), + distance=values.get("distance"), + binning=values.get("binning"), + backend=values.get("backend"), + fit_result=values.get("fit_result", {}), + attrs=values.get("attrs", {}), + ) + + ######################## + # EVALUATION AND PLOTTING + ######################## + + def variogram(self, distance: NDArrayNum | float) -> NDArrayNum: + """ + Evaluate the fitted theoretical variogram. + + :param distance: Spatial distance or array of distances. + :returns: Semivariance at each distance. + """ + + if self.model is None: + raise ValueError("A fitted variogram model is required for evaluation.") + return self.model.variogram(distance) + + __call__ = variogram + + def covariance(self, distance: NDArrayNum | float) -> NDArrayNum: + """ + Evaluate covariance implied by the fitted model. + + :param distance: Spatial distance or array of distances. + :returns: Covariance at each distance. + """ + + if self.model is None: + raise ValueError("A fitted variogram model is required for evaluation.") + return self.model.covariance(distance) + + def correlation(self, distance: NDArrayNum | float) -> NDArrayNum: + """ + Evaluate correlation implied by the fitted model. + + :param distance: Spatial distance or array of distances. + :returns: Correlation at each distance. + """ + + if self.model is None: + raise ValueError("A fitted variogram model is required for evaluation.") + return self.model.correlation(distance) + + def plot(self, ax: Any | None = None, *, show_error: bool = True, **kwargs: Any) -> Any: + """ + Plot measured bins and the fitted model when present. + + :param ax: Existing Matplotlib axes. A new figure and axes are created by default. + :param show_error: Whether to draw available sampling errors. + :param kwargs: Keyword arguments passed to the measured point plot. + :returns: Matplotlib axes containing the variogram. + """ + + # Import Matplotlib only when the caller requests a plot + pyplot = import_optional("matplotlib.pyplot", package_name="matplotlib") + if ax is None: + _, ax = pyplot.subplots() + + # Draw measured bins and optional error bars, then add the fitted curve + error = self.semivariance_error if show_error else None + ax.errorbar(self.lags, self.semivariance, yerr=error, fmt="o", **kwargs) + if self.model is not None and np.any(np.isfinite(self.lags)): + distances = np.linspace(0, float(np.nanmax(self.lags)), 500) + ax.plot(distances, self.variogram(distances)) + ax.set(xlabel="Lag distance", ylabel="Semivariance") + return ax + + ############################ + # CONVERSIONS TO OTHER PACKAGES + ############################ + + def to_gstools(self, *, dim: int = 2) -> GSToolsVariogram: + """ + Convert the fitted model to GSTools with its source feature dimensions. + + :param dim: Number of dimensions passed to the GSTools covariance model. + :returns: Native covariance model and dimensions selected by the source model. + """ + + # Import GSTools only for this conversion and require a fitted model + gstools = import_optional("gstools", extra_name="geostat") + if self.model is None: + raise ValueError("A fitted variogram model is required for conversion.") + if dim <= 0: + raise ValueError("GSTools model dimension must be strictly positive.") + + # Require all components to use the same coordinate columns expected by one GSTools model + active_dims = {component.active_dims for component in self.model.components} if self.model.components else set() + if len(active_dims) > 1: + raise NotImplementedError("GSTools cannot combine variogram components that select different dimensions.") + + # Return coordinate-column choices separately because GSTools does not store them + selected_dims = next(iter(active_dims)) if active_dims else self.model.active_dims + return GSToolsVariogram( + model=_model_to_gstools(self.model, gstools=gstools, dim=dim), + active_dims=selected_dims, + ) + + def gpytorch_parameters(self) -> dict[str, Any]: + """Describe the fitted model with plain parameters used by the GPyTorch conversion.""" + + if self.model is None: + raise ValueError("A fitted variogram model is required for conversion.") + return _model_to_gpytorch_parameters(self.model) + + def to_gpytorch(self, *, active_dims: tuple[int, ...] | None = None, trainable: bool = True) -> GPyTorchVariogram: + """ + Convert supported fitted structures to a GPyTorch covariance kernel. + + :param active_dims: Optional feature column override applied to every structure. + :param trainable: Whether converted kernel parameters may be optimized. + :returns: Native covariance kernel and separate likelihood noise. + """ + + # Import GPyTorch only for this conversion and require a fitted model + gpytorch = import_optional("gpytorch", extra_name="gp") + if self.model is None: + raise ValueError("A fitted variogram model is required for conversion.") + kernel, noise = _model_to_gpytorch( + self.model, + gpytorch=gpytorch, + active_dims=active_dims, + ) + + # Make kernel parameters fixed when the caller will use the model only for prediction + if not trainable: + for parameter in kernel.parameters(): + parameter.requires_grad_(False) + return GPyTorchVariogram(kernel=kernel, noise=noise) + + +############################ +# 4/ REPEATED PAIR ESTIMATION +############################ + + +def _estimate_variogram( + source: RasterBase | PointCloudBase, + *, + n_runs: int, + estimator: str | Callable[[NDArrayNum], float], + bins: Literal["log", "uniform"] | Iterable[float], + n_lags: int, + min_lag: float | None, + max_lag: float | None, + models: str | Callable[..., Any] | Sequence[str | Callable[..., Any]] | None, + fit_kwargs: Mapping[str, Any] | None, + random_state: int | np.random.Generator | None, + pair_kwargs: Mapping[str, Any], +) -> Variogram: + """ + Sample one or more pair sets and combine their per-distance variogram values. + + Each run calls the source pairsample() method, then Variogram.from_pairs() groups value differences by distance. + Repeated runs share the first run's bins and are combined before Variogram.fit() fits any requested model. + + Source, estimation, binning, repetition, fitting, and random-state options are documented by + geoutils.stats.variogram(). + + :param models: Public model argument forwarded to Variogram.fit(): one model or a sequence of models to sum, + or None to keep only the measured variogram. + :param pair_kwargs: Complete options for source.pairsample(), including sample size, mask, distance limits, + and any source-specific controls. This function supplies a separate random_state for each run. + :returns: Variogram with mean measurements and total pair counts across runs, sampling errors when repeated, + and any requested fitted model. + """ + + # Check the repeat count before creating one random seed per run + if not isinstance(n_runs, (int, np.integer)) or isinstance(n_runs, bool) or n_runs < 1: + raise ValueError("Argument ``n_runs`` must be a positive integer.") + rng = random_state if isinstance(random_state, np.random.Generator) else np.random.default_rng(random_state) + seeds = rng.integers(0, np.iinfo(np.int32).max, n_runs) + + # Materialize explicit boundaries once because a generator cannot be read by every repeated run + if not isinstance(bins, str): + bins = tuple(bins) + + def run(seed: np.integer[Any], bin_spec: Literal["log", "uniform"] | Iterable[float]) -> Variogram: + """ + Calculate one run with its own seed and either the initial or shared distance bins. + + :param seed: Integer seed for this run's pair sample. + :param bin_spec: Initial binning choice, or fixed boundaries copied from the first run. + """ + + # Convert one pair sample to per-distance values and then release the pairs + pairs = source.pairsample(random_state=int(seed), **pair_kwargs) + return Variogram.from_pairs( + pairs, + estimator=estimator, + bins=bin_spec, + n_lags=n_lags, + min_lag=min_lag, + max_lag=max_lag, + ) + + # Build distance bins from the first run so all later runs use the same edges + first = run(seeds[0], bins) + if n_runs == 1: + result = replace(first, attrs={**first.attrs, "n_runs": 1}) + return result if models is None else result.fit(models, **dict(fit_kwargs or {})) + + # Draw each remaining sample through the object's own raster or point implementation + shared_bins: Literal["log", "uniform"] | Iterable[float] = bins + if isinstance(bins, str): + assert first.bin_lower_edges is not None and first.bin_edges is not None + shared_bins = np.r_[first.bin_lower_edges[0], first.bin_edges] + + runs = [first, *[run(seed, shared_bins) for seed in seeds[1:]]] + + # Stack the small per-distance arrays from all runs + semivariances = np.vstack([run_result.semivariance for run_result in runs]) + lag_centers = np.vstack([run_result.lags for run_result in runs]) + counts = np.sum(np.vstack([run_result.counts for run_result in runs]), axis=0) + mean_semivariance = np.full(semivariances.shape[1], np.nan) + mean_lags = np.full(semivariances.shape[1], np.nan) + errors = np.full(semivariances.shape[1], np.nan) + + # Average each distance bin and calculate sampling error only when runs are repeated + for index in range(semivariances.shape[1]): + finite_semivariance = semivariances[np.isfinite(semivariances[:, index]), index] + finite_lags = lag_centers[np.isfinite(lag_centers[:, index]), index] + if finite_semivariance.size: + mean_semivariance[index] = np.mean(finite_semivariance) + if finite_semivariance.size > 1: + errors[index] = np.std(finite_semivariance, ddof=1) / np.sqrt(finite_semivariance.size) + if finite_lags.size: + mean_lags[index] = np.mean(finite_lags) + + # Reuse the first run's bin details with the combined values and counts + result = replace( + first, + lags=mean_lags, + semivariance=mean_semivariance, + semivariance_error=errors, + counts=counts, + attrs={**first.attrs, "n_runs": n_runs, "pair_count": int(np.sum(counts))}, + ) + return result if models is None else result.fit(models, **dict(fit_kwargs or {})) + + +################################ +# 5/ HELPERS FOR OTHER PACKAGES +################################ + + +def _model_from_skgstat( + variogram: Any, + configured_model: str, + description: Mapping[str, Any], + *, + active_dims: tuple[int, ...] | None, +) -> VariogramModel: + """ + Copy SciKit-GStat fitted numbers into the shared `VariogramModel` form. + + The fitted variogram and active_dims arguments are described by Variogram.from_skgstat(). + + :param configured_model: Lowercase SciKit-GStat model name, with summed components separated by plus signs. + :param description: Model parameters and settings returned by the fitted SciKit-GStat object's describe(). + :returns: Portable fitted model, preserving component order and the shared nugget for summed structures. + """ + + # Read a single model's range, sill, shape, and nugget from SciKit-GStat's description + if "+" not in configured_model: + return VariogramModel( + model_name=configured_model, + effective_range=float(description["effective_range"]), + partial_sill=float(description["sill"]), + nugget=float(description["nugget"]), + smoothness=None if description.get("smoothness") is None else float(description["smoothness"]), + shape=None if description.get("shape") is None else float(description["shape"]), + active_dims=active_dims, + ) + + # Read one shared nugget after all summed component parameters + names = [name.strip() for name in configured_model.split("+")] + coefficients = list(np.asarray(variogram.cof, dtype=float)) + use_nugget = bool(description.get("params", {}).get("use_nugget", False)) + nugget = float(coefficients.pop()) if use_nugget else 0.0 + components: list[VariogramModel] = [] + position = 0 + for name in names: + # Read each component's range and sill before its optional shape value + effective_range, partial_sill = coefficients[position : position + 2] + position += 2 + smoothness = shape = None + if name == "matern": + smoothness = float(coefficients[position]) + position += 1 + elif name == "stable": + shape = float(coefficients[position]) + position += 1 + components.append( + VariogramModel( + model_name=name, + effective_range=float(effective_range), + partial_sill=float(partial_sill), + smoothness=smoothness, + shape=shape, + active_dims=active_dims, + ) + ) + + # Build the shared summed model with one nugget on its parent + return VariogramModel.sum(components, nugget=nugget) + + +def _model_to_gstools(model: VariogramModel, *, gstools: Any, dim: int) -> Any: + """ + Build the equivalent GSTools model from one shared `VariogramModel`. + + The dimension argument is described by Variogram.to_gstools(). + + :param model: Portable fitted model to convert, including any summed components. + :param gstools: GSTools module already imported by Variogram.to_gstools(). + :returns: Native GSTools covariance model with the fitted range, variance, and nugget. + """ + + # Convert each summed component, then apply the shared nugget once + if model.model_name == "sum": + components = [_model_to_gstools(component, gstools=gstools, dim=dim) for component in model.components] + return gstools.SumModel(*components, nugget=model.nugget) + if model.model_name == "product": + raise NotImplementedError("Product covariance conversion is not supported by the GSTools adapter.") + if model.effective_range is None or model.partial_sill is None: + raise AssertionError("A base variogram model must define its range and partial sill.") + + # Pass the shared model's variance and effective range to the matching GSTools class + common = {"dim": dim, "var": model.partial_sill, "nugget": model.nugget, "len_scale": model.effective_range} + if model.model_name == "spherical": + return gstools.Spherical(**common) + if model.model_name == "exponential": + return gstools.Exponential(rescale=3.0, **common) + if model.model_name == "gaussian": + return gstools.Gaussian(rescale=2.0, **common) + if model.model_name == "cubic": + return gstools.Cubic(**common) + if model.model_name == "stable": + if model.shape is None: + raise AssertionError("A stable variogram model must define its shape.") + return gstools.Stable(alpha=model.shape, rescale=float(3 ** (1 / model.shape)), **common) + + # Pass Matérn smoothness through GSTools' matching parameter name + if model.model_name == "matern": + return gstools.Matern(nu=model.smoothness, rescale=4.0, **common) + raise NotImplementedError(f"Variogram model {model.model_name!r} has no GSTools adapter.") + + +def _model_to_gpytorch( + model: VariogramModel, *, gpytorch: Any, active_dims: tuple[int, ...] | None +) -> tuple[Any, float]: + """ + Build an equivalent GPyTorch kernel for models both packages represent exactly. + + The active_dims override is described by Variogram.to_gpytorch(). + + :param model: Portable fitted model to convert, including any summed or multiplied components. + :param gpytorch: GPyTorch module already imported by Variogram.to_gpytorch(). + :returns: Native GPyTorch covariance kernel and separate nugget variance for the observation likelihood. + """ + + # Convert nested components and join them by the model's sum or product rule + if model.model_name == "sum": + converted = [ + _model_to_gpytorch(component, gpytorch=gpytorch, active_dims=active_dims) for component in model.components + ] + kernel = converted[0][0] + for component_kernel, _ in converted[1:]: + kernel = kernel + component_kernel + return kernel, model.nugget + + if model.model_name == "product": + converted = [ + _model_to_gpytorch(component, gpytorch=gpytorch, active_dims=active_dims) for component in model.components + ] + kernel = converted[0][0] + for component_kernel, _ in converted[1:]: + kernel = kernel * component_kernel + return kernel, model.nugget + + # Convert effective range to the matching GPyTorch kernel length scale + parameters = _model_to_gpytorch_parameters(model, active_dims=active_dims) + resolved_active_dims = parameters["active_dims"] + if model.partial_sill is None: + raise AssertionError("A base variogram model must define its partial sill.") + + if model.model_name == "gaussian": + base_kernel = gpytorch.kernels.RBFKernel(active_dims=resolved_active_dims) + elif model.model_name == "exponential": + base_kernel = gpytorch.kernels.MaternKernel(nu=0.5, active_dims=resolved_active_dims) + elif model.model_name == "matern": + base_kernel = gpytorch.kernels.MaternKernel(nu=model.smoothness, active_dims=resolved_active_dims) + else: + raise NotImplementedError(f"Variogram model {model.model_name!r} has no exact GPyTorch adapter.") + + # Apply the model variance around GPyTorch's base correlation kernel + kernel = gpytorch.kernels.ScaleKernel(base_kernel) + kernel.base_kernel.lengthscale = parameters["lengthscale"] + kernel.outputscale = float(model.partial_sill) + return kernel, model.nugget + + +def _model_to_gpytorch_parameters( + model: VariogramModel, + *, + active_dims: tuple[int, ...] | None = None, +) -> dict[str, Any]: + """ + Describe the GPyTorch conversion with plain data without importing GPyTorch. + + Model and active_dims inputs follow _model_to_gpytorch(). + + :returns: Kernel name, length scale, variance, selected dimensions, and nugget; nested component dictionaries + preserve summed or multiplied structures. + """ + + # Describe nested models recursively so callers can inspect them without GPyTorch + if model.model_name in _COMPOSITE_MODELS: + return { + "combination": model.model_name, + "components": [ + _model_to_gpytorch_parameters(component, active_dims=active_dims) for component in model.components + ], + "noise": model.nugget, + } + if model.effective_range is None or model.partial_sill is None: + raise AssertionError("A base variogram model must define its range and partial sill.") + + # Convert effective range to the matching GPyTorch kernel length scale + if model.model_name == "gaussian": + kernel_name = "RBF" + lengthscale = float(model.effective_range) / (2 * np.sqrt(2)) + smoothness = None + elif model.model_name == "exponential": + kernel_name = "Matern" + lengthscale = float(model.effective_range) / 3 + smoothness = 0.5 + elif model.model_name == "matern": + if model.smoothness not in (0.5, 1.5, 2.5): + raise NotImplementedError("GPyTorch Matérn kernels support smoothness values 0.5, 1.5 and 2.5.") + kernel_name = "Matern" + lengthscale = float(model.effective_range) / (2 * np.sqrt(2)) + smoothness = model.smoothness + else: + raise NotImplementedError(f"Variogram model {model.model_name!r} has no exact GPyTorch adapter.") + + # Keep observation noise separate because GPyTorch applies it in the likelihood + return { + "kernel_name": kernel_name, + "lengthscale": lengthscale, + "outputscale": float(model.partial_sill), + "smoothness": smoothness, + "active_dims": model.active_dims if active_dims is None else active_dims, + "noise": model.nugget, + } diff --git a/geoutils/vector/base.py b/geoutils/vector/base.py index dd5affa6b..e5144b9be 100644 --- a/geoutils/vector/base.py +++ b/geoutils/vector/base.py @@ -60,7 +60,8 @@ VectorBaseType = TypeVar("VectorBaseType", bound="VectorBase") -VectorBaseLike = Union["VectorBase", gpd.GeoDataFrame] +# Accept Vector subclasses and accessors, as well as GeoDataFrames +VectorLike = Union["VectorBase", gpd.GeoDataFrame] def _as_geodataframe(obj: Any) -> gpd.GeoDataFrame: @@ -293,7 +294,7 @@ def info(self, verbose: bool = True) -> str | None: def plot( self, - ref_crs: RasterLike | VectorBaseLike | CRS | int | None = None, + ref_crs: RasterLike | VectorLike | CRS | int | None = None, cmap: matplotlib.colors.Colormap | str | None = None, vmin: float | int | None = None, vmax: float | int | None = None, @@ -407,7 +408,7 @@ def active_geometry_name(self) -> str: @overload def crop( self: VectorBaseType, - bbox: RasterLike | VectorBaseLike | tuple[float, float, float, float], + bbox: RasterLike | VectorLike | tuple[float, float, float, float], clip: bool, *, inplace: Literal[False] = False, @@ -417,7 +418,7 @@ def crop( @overload def crop( self: VectorBaseType, - bbox: RasterLike | VectorBaseLike | tuple[float, float, float, float], + bbox: RasterLike | VectorLike | tuple[float, float, float, float], clip: bool, *, inplace: Literal[True], @@ -427,7 +428,7 @@ def crop( @overload def crop( self: VectorBaseType, - bbox: RasterLike | VectorBaseLike | tuple[float, float, float, float], + bbox: RasterLike | VectorLike | tuple[float, float, float, float], clip: bool, *, inplace: bool = False, @@ -437,7 +438,7 @@ def crop( @profiler.profile("geoutils.vector.base.crop", memprof=True) def crop( self: VectorBaseType, - bbox: RasterLike | VectorBaseLike | tuple[float, float, float, float] = None, + bbox: RasterLike | VectorLike | tuple[float, float, float, float] = None, clip: bool = False, *, inplace: bool = False, @@ -467,7 +468,7 @@ def crop( @overload def reproject( self: VectorBaseType, - ref: RasterLike | VectorBaseLike | None = None, + ref: RasterLike | VectorLike | None = None, crs: CRS | str | int | None = None, *, inplace: Literal[False] = False, @@ -476,7 +477,7 @@ def reproject( @overload def reproject( self: VectorBaseType, - ref: RasterLike | VectorBaseLike | None = None, + ref: RasterLike | VectorLike | None = None, crs: CRS | str | int | None = None, *, inplace: Literal[True], @@ -485,7 +486,7 @@ def reproject( @overload def reproject( self: VectorBaseType, - ref: RasterLike | VectorBaseLike | None = None, + ref: RasterLike | VectorLike | None = None, crs: CRS | str | int | None = None, *, inplace: bool = False, @@ -494,7 +495,7 @@ def reproject( @profiler.profile("geoutils.vector.base.reproject", memprof=True) def reproject( self: VectorBaseType, - ref: RasterLike | VectorBaseLike | None = None, + ref: RasterLike | VectorLike | None = None, crs: CRS | str | int | None = None, inplace: bool = False, ) -> VectorBaseType | gpd.GeoDataFrame | None: @@ -699,7 +700,7 @@ def rasterize( @classmethod def from_bounds_projected( - cls, raster_or_vector: RasterType | VectorBaseLike, out_crs: CRS | None = None, densify_points: int = 5000 + cls, raster_or_vector: RasterType | VectorLike, out_crs: CRS | None = None, densify_points: int = 5000 ) -> VectorBaseType | gpd.GeoDataFrame: """Create a vector polygon from projected bounds of a raster or vector. diff --git a/geoutils/vector/transformation.py b/geoutils/vector/transformation.py index 15325bf2f..607f15d5f 100644 --- a/geoutils/vector/transformation.py +++ b/geoutils/vector/transformation.py @@ -69,12 +69,11 @@ def _crop(source_vector: Any, bbox: Any, clip: bool) -> Any: return _crop_geodataframe(source_vector.ds, bounds=bounds, clip=clip) -def _reproject( - source_vector: Any, +def _get_reproject_crs( ref: RasterLike | VectorLike | None = None, crs: CRS | str | int | None = None, -) -> gpd.GeoDataFrame: - """Reproject a vector. See Vector.reproject() for more details.""" +) -> CRS: + """Resolve a target CRS from exactly one reference object or explicit CRS.""" # Check that either ref or crs is provided if (ref is not None and crs is not None) or (ref is None and crs is None): @@ -91,6 +90,17 @@ def _reproject( # Determine user-input target CRS crs = CRS.from_user_input(crs) - new_ds = source_vector.ds.to_crs(crs=crs) + return crs + + +def _reproject( + source_vector: Any, + ref: RasterLike | VectorLike | None = None, + crs: CRS | str | int | None = None, +) -> gpd.GeoDataFrame: + """Reproject a vector. See Vector.reproject() for more details.""" + + target_crs = _get_reproject_crs(ref=ref, crs=crs) + new_ds = source_vector.ds.to_crs(crs=target_crs) return new_ds diff --git a/geoutils/vector/vector.py b/geoutils/vector/vector.py index 90d8923ce..78c3c6443 100644 --- a/geoutils/vector/vector.py +++ b/geoutils/vector/vector.py @@ -34,7 +34,6 @@ Literal, Sequence, TypeVar, - Union, ) import geopandas as gpd @@ -49,13 +48,13 @@ from geoutils import profiler from geoutils._misc import copy_doc from geoutils.vector.base import VectorBase +from geoutils.vector.base import VectorLike as VectorLike # noqa: F401 if TYPE_CHECKING: from geoutils.raster.base import RasterType # This is a generic Vector-type (if subclasses are made, this will change appropriately) VectorType = TypeVar("VectorType", bound="Vector") -VectorLike = Union["Vector", gpd.GeoDataFrame] class Vector(VectorBase): diff --git a/setup.cfg b/setup.cfg index 2532d373a..7ecd26763 100644 --- a/setup.cfg +++ b/setup.cfg @@ -59,11 +59,19 @@ opt = psutil plotly dask-geopandas +flox = + flox benchmark = %(opt)s asv distributed pytest>=7,<8 +geostat = + scikit-gstat>=1.0.23 + gstools>=1.3 +gp = + gpytorch>=1.11 + torch>=2 test = gdal pytest>=7,<8 @@ -94,7 +102,10 @@ dev = %(benchmark)s %(test)s %(doc)s + %(flox)s psutil plotly all = %(dev)s + %(geostat)s + %(gp)s diff --git a/tests/test_benchmark/test_benchmark_tools.py b/tests/test_benchmark/test_benchmark_tools.py index 0bcb8d746..6f895b891 100644 --- a/tests/test_benchmark/test_benchmark_tools.py +++ b/tests/test_benchmark/test_benchmark_tools.py @@ -3,10 +3,15 @@ from __future__ import annotations import json +import os from pathlib import Path +from typing import Literal +import numpy as np +import pandas as pd import pytest +from benchmarks.asv_suite import comparisons as benchmark_comparisons from benchmarks.asv_suite.comparisons import ( BENCHMARK_CASE_BY_CLASS, BENCHMARK_CASES, @@ -40,12 +45,23 @@ _warp_memory_limit_mb, build_gdal_command, ) +from benchmarks.workflows.grouped_reference import ( + compute_grouped_reference, + prepare_grouped_reference, +) from benchmarks.workflows.registry import ( OPERATION_METHODS, OPERATION_STRATEGIES, split_operation_case, ) from benchmarks.workflows.runner import BenchmarkConfig +from benchmarks.workflows.variography import ( + prepare_pair_pointcloud, + prepare_pair_raster, + prepare_variogram_pairs, +) +from geoutils import Variogram +from geoutils._misc import import_optional class TestComparisonReport: @@ -66,8 +82,13 @@ def test_collect_comparison_measurements(self) -> None: if record.series_dimension != "method": assert record.method == comparison.method if record.external_reference is not None: - assert record.series_label == GDAL_CLI_LABEL - assert record.execution_mode is None + if record.external_reference == "gdal_cli": + assert record.series_label == GDAL_CLI_LABEL + assert record.execution_mode is None + else: + assert record.external_reference == "flox" + assert record.execution_mode is not None + assert record.series_label == f"Flox ({EXECUTION_MODE_LABELS[record.execution_mode]})" assert record.calculation_engine is None else: assert record.execution_mode is not None @@ -94,6 +115,7 @@ def test_benchmark_dimension_registry(self) -> None: "Number of interpolated points", "Number of source points per axis", "Number of sampled values", + "Number of groups per axis", } method_by_key = { (specification.operation, specification.method): specification for specification in OPERATION_METHODS @@ -125,7 +147,7 @@ def test_benchmark_dimension_registry(self) -> None: cases = [ BENCHMARK_CASE_BY_CLASS[class_name] for label, class_name in comparison.series - if label != GDAL_CLI_LABEL + if class_name in BENCHMARK_CASE_BY_CLASS ] assert {case.operation for case in cases} == {comparison.operation} if comparison.series_dimension != "method": @@ -202,8 +224,9 @@ def test_select_complete_result__skip_incomplete(self) -> None: assert _select_complete_result((complete, incomplete)) is complete - def test_render_comparisons(self, tmp_path: Path) -> None: - """Checks that rendering a complete result writes navigation, numeric exports and every plot.""" + @pytest.mark.parametrize("include_flox", [True, False]) + def test_render_comparisons(self, tmp_path: Path, include_flox: bool) -> None: + """Checks that report navigation, exports and plots work with or without optional Flox measurements.""" pytest.importorskip("matplotlib") @@ -211,7 +234,13 @@ def test_render_comparisons(self, tmp_path: Path) -> None: asv_directory = tmp_path / "asv" asv_directory.mkdir() (asv_directory / "index.html").write_text("Native ASV report", encoding="utf-8") - render_comparisons(_PreviewResult(), tmp_path) + result = _PreviewResult() + if not include_flox: + for class_name, case in EXTERNAL_REFERENCE_CASE_BY_CLASS.items(): + if case.external_reference == "flox": + for method in ("time_operation", "track_end_to_end_time_s", "track_peak_process_tree_mem_mb"): + result.values[f"asv_suite.comparisons.{class_name}.{method}"] = [float("nan")] * 3 + render_comparisons(result, tmp_path) assert (tmp_path / "index.html").is_file() root_page = (tmp_path / "index.html").read_text(encoding="utf-8") assert "comparisons/index.html" in root_page @@ -224,7 +253,7 @@ def test_render_comparisons(self, tmp_path: Path) -> None: assert "GeoUtils calculation engine" in root_page assert "GeoUtils execution mode" in root_page assert "External reference" in root_page - for option in ("Rasterio/GDAL", "SciPy", "Numba", "GDAL CLI", "Eager", "Dask", "Multiprocessing"): + for option in ("Rasterio/GDAL", "SciPy", "Numba", "NumPy", "GDAL CLI", "Eager", "Dask", "Multiprocessing"): assert f">{option}" in root_page assert 'Fastest' not in root_page assert "MB" in root_page @@ -246,7 +275,7 @@ def test_render_comparisons(self, tmp_path: Path) -> None: assert "Prepared inputs through completed output" in root_page assert "data larger than memory" in root_page assert "Ratios" not in root_page - assert 'GeoUtils calculation engine' in root_page + assert 'GeoUtils calculation engine' in root_page assert 'GeoUtils execution mode' in root_page operation_groups = ( "Raster ⟶ Raster", @@ -291,7 +320,7 @@ def test_render_comparisons(self, tmp_path: Path) -> None: assert 'Memory' in comparison_page assert "Ratios" in comparison_page assert comparison_page.count('class="operation-group"') >= len(operation_groups) - assert 'class="group-heading" colspan="3" scope="colgroup">GeoUtils calculation engine' in comparison_page + assert 'class="group-heading" colspan="4" scope="colgroup">GeoUtils calculation engine' in comparison_page assert 'class="group-heading" colspan="3" scope="colgroup">GeoUtils execution mode' in comparison_page assert 'class="choice-table"' in comparison_page assert comparison_page.count('class="workload-heading"') == 2 @@ -307,6 +336,7 @@ def test_render_comparisons(self, tmp_path: Path) -> None: "Scaling with the number of interpolated points", "Scaling with the number of source points", "Scaling with the number of sampled values", + "Scaling with the number of groups", ): assert heading in scaling_page assert 'class="plot-toc"' in scaling_page @@ -344,8 +374,8 @@ def test_performance_change_markdown(self, tmp_path: Path) -> None: # Doubling each unique baseline GeoUtils result should produce a twofold normalized improvement baseline_keys = set() for comparison in COMPARISONS: - for series_label, class_name in comparison.series: - if series_label == GDAL_CLI_LABEL: + for _series_label, class_name in comparison.series: + if class_name in EXTERNAL_REFERENCE_CASE_BY_CLASS: continue key = f"asv_suite.comparisons.{class_name}.track_end_to_end_time_s" baseline_keys.add(key) @@ -390,6 +420,144 @@ def test_render_documentation_snapshot(self, tmp_path: Path) -> None: assert snapshot["metadata"]["machine"]["machine"] == "preview-machine" +class TestGroupedReferenceChunked: + """Checks equivalent GeoUtils/Flox results and real multiprocessing worker reuse for prepared arrays.""" + + @pytest.mark.parametrize("execution_mode", ["eager", "dask"]) + def test_grouped_reference__matching_statistics(self, execution_mode: Literal["eager", "dask"]) -> None: + """Checks that both benchmark implementations return the same masked counts and population moments.""" + + # Skip the optional Flox comparison when it is not installed + try: + import_optional("flox", extra_name="benchmark") + except ImportError as exc: + pytest.skip(str(exc)) + inputs = prepare_grouped_reference(32, 4, "interleaved", execution_mode) + values, groups, mask, categories = prepare_grouped_reference(32, 4, "interleaved", "eager") + + # Compute complete public GeoUtils and Flox tables from identical prepared inputs + geoutils_result = compute_grouped_reference(*inputs, implementation="geoutils") + flox_result = compute_grouped_reference(*inputs, implementation="flox") + expected = compute_grouped_reference(values, groups, mask, categories, implementation="geoutils") + pd.testing.assert_frame_equal(geoutils_result, expected, rtol=1e-12, atol=1e-12) + pd.testing.assert_frame_equal(geoutils_result, flox_result, rtol=1e-12, atol=1e-12) + if execution_mode == "dask": + import dask.array as da + + assert isinstance(inputs[0], da.Array) + assert isinstance(inputs[1], da.Array) + assert isinstance(inputs[2], da.Array) + + # Independently check every finite count, mean and ddof=0 standard deviation with NumPy + for value_index, name in enumerate(("first", "second")): + for group in categories: + members = values[value_index][(groups == group) & mask] + finite = members[np.isfinite(members)] + expected = [finite.size, finite.mean(), finite.std(ddof=0)] + np.testing.assert_allclose(geoutils_result.loc[group, name], expected, rtol=1e-12, atol=1e-12) + + def test_grouped_reference__multiprocessing_reuses_worker(self) -> None: + """Checks that repeated grouped calculations reuse one real worker and match eager values exactly.""" + + # Use the registered benchmark setup so its fixed process count, tile size and pool lifetime are checked + case = next( + case + for case in BENCHMARK_CASES + if case.comparison_group == "grouped-flox-raster-size" and case.execution_mode == "multiprocessing" + ) + benchmark = getattr(benchmark_comparisons, case.benchmark_class)() + # A 513-square input has nine 256-square tiles, so two calls exceed the pool's normal ten-task lifetime + try: + benchmark.setup(513) + assert benchmark.mp_cluster is not None + worker_pids = benchmark.mp_cluster.worker_pids() + assert len(worker_pids) == 1 + assert worker_pids[0] != os.getpid() + assert benchmark.mp_config.chunks == (256, 256) + + # Reduce the same prepared arrays twice using serialization and the initialized worker + expected = compute_grouped_reference(*benchmark.inputs, implementation="geoutils") + for _ in range(2): + result = compute_grouped_reference( + *benchmark.inputs, implementation="geoutils", mp_config=benchmark.mp_config + ) + pd.testing.assert_frame_equal(result, expected, rtol=1e-12, atol=1e-12) + assert benchmark.mp_cluster.worker_pids() == worker_pids + finally: + # Clean up the worker even when result or lifetime validation fails + benchmark.teardown(513) + + +class TestVariographyWorkflows: + """Checks prepared variography fixtures against independent distance and semivariance calculations. + + ASV covers execution and scaling of the benchmark classes; these tests check their inputs and full public outputs. + """ + + def test_variogram_pairs__independent_bin_statistics(self) -> None: + """Checks that prepared pair data gives the expected finite counts, mean lags and semivariances.""" + + # Use explicit edges and a simple mean-square formula so the fixture check needs no optional estimator package + pairs = prepare_variogram_pairs(1_000) + edges = np.geomspace(1, 1024, 9) + result = Variogram.from_pairs( + pairs, bins=edges, estimator=lambda differences: float(np.mean(differences**2) / 2) + ) + + # Calculate each distance interval independently from endpoint differences and include every prepared pair + distances = pairs["distance"].values + differences = np.diff(pairs["value"].values, axis=1).ravel() + for index, (lower, upper) in enumerate(zip(edges[:-1], edges[1:])): + selected = (distances > lower) & (distances <= upper) + assert result.counts[index] == np.count_nonzero(selected) + np.testing.assert_allclose(result.lags[index], distances[selected].mean()) + np.testing.assert_allclose(result.semivariance[index], np.mean(differences[selected] ** 2) / 2) + assert result.counts.sum() == pairs.sizes["pair"] + + def test_pair_pointcloud__original_values_and_distances(self) -> None: + """Checks that the irregular point fixture returns its smooth values at the sampled endpoint coordinates.""" + + # Use about one point per square map unit and request pairs well inside the source extent + points = prepare_pair_pointcloud(1_000) + pairs = points.pairsample(n_pairs=200, min_distance=1, max_distance=15, random_state=42) + + # Calculate signal values and Euclidean distances directly from the reported endpoint coordinates + x, y = pairs["x"].values, pairs["y"].values + np.testing.assert_allclose(pairs["value"], np.sin(x / 31) + np.cos(y / 53)) + expected_distances = np.hypot(np.diff(x, axis=1), np.diff(y, axis=1)).ravel() + np.testing.assert_allclose(pairs["distance"], expected_distances, rtol=1e-6) + assert pairs.sizes["pair"] == 200 + + +class TestPairRasterChunked: + """Checks prepared Dask raster pairs against the same eager pair sample.""" + + def test_pair_raster__finite_values_and_distances(self) -> None: + """Checks that Dask raster pairs match eager and exclude the deliberate nodata cells.""" + + # Draw the same public pair sample from eager and one-chunk Dask fixtures + import dask.array as da + + expected = prepare_pair_raster(32, "eager").pairsample( + n_pairs=200, min_distance=1, max_distance=16, random_state=42 + ) + raster = prepare_pair_raster(32, "dask") + pairs = raster.pairsample(n_pairs=200, min_distance=1, max_distance=16, random_state=42) + assert isinstance(raster.data, da.Array) and not raster._obj._in_memory + assert not pairs.chunks + for name in pairs.variables: + assert np.array_equal(pairs[name], expected[name]) + + # Recover the analytic signal from endpoint positions and check the pattern used to introduce nodata values + rows, columns = pairs["row"].values, pairs["column"].values + expected_values = (np.sin(columns / 31) + np.cos(rows / 53)).astype(np.float32) + np.testing.assert_array_equal(pairs["value"], expected_values) + assert not np.any((rows * 32 + columns) % 17 == 0) + expected_distances = np.hypot(np.diff(rows, axis=1), np.diff(columns, axis=1)).ravel() + np.testing.assert_allclose(pairs["distance"], expected_distances) + assert pairs.sizes["pair"] == 200 + + class TestGdalCommands: """Verify GDAL CLI commands match GeoUtils inputs, outputs and resource limits.""" diff --git a/tests/test_dispatch.py b/tests/test_dispatch.py index 107b81977..f346f7b87 100644 --- a/tests/test_dispatch.py +++ b/tests/test_dispatch.py @@ -26,6 +26,7 @@ _grid_from_bounds_res, _grid_from_bounds_shape, _grid_from_coords, + get_geo_interface, ) from geoutils.exceptions import ( IgnoredGridWarning, @@ -38,6 +39,28 @@ ) +class TestGeoInterface: + """Checks that spatial metadata come from the same object or accessor as its geospatial operations.""" + + @pytest.mark.parametrize("representation", ["raster", "accessor", "dataarray"]) + def test_get_geo_interface__raster_metadata(self, representation: str) -> None: + """Checks that raster band count and spatial shape stay distinct from native Xarray methods and dimensions.""" + + # Include two bands so native Xarray shape differs from the two-dimensional raster grid + raster = gu.Raster.from_array(np.ones((2, 3, 4)), rio.transform.from_origin(0, 3, 1, 1), crs=32633) + array = raster.to_xarray() + source = {"raster": raster, "accessor": array.rst, "dataarray": array}[representation] + + # Resolve one interface and read both metadata fields from it + interface = get_geo_interface(source, "ij2xy", accessors=("rst",)) + assert interface.count == 2 + assert interface.shape == (3, 4) + assert interface.transform == raster.transform + + # Plain arrays have no spatial interface and cannot supply a reference grid + assert get_geo_interface(np.ones((3, 4)), "ij2xy", accessors=("rst",)) is None + + class TestDispatchLevelZero: """Level zero user-input checks: CRS, bounds, shape, resolution, regular coordinates, point coordinates.""" diff --git a/tests/test_filters.py b/tests/test_filters.py index 8456344bc..f57d4c41d 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -2,7 +2,7 @@ from collections.abc import Callable from importlib.util import find_spec -from typing import Any +from typing import Any, Literal import numpy as np import pytest @@ -17,6 +17,78 @@ from geoutils.raster import get_array_and_mask +class TestPatchFilters: + """Checks convolution and mean filtering for stacks of raster patches.""" + + @pytest.mark.parametrize("shape", [(3, 3), (4, 4), (3, 4)]) + def test_stacked_convolution_and_kernel_orientation(self, shape: tuple[int, int]) -> None: + """Checks that stacked convolution agrees with SciPy for odd, even and asymmetric kernels.""" + + # Use distinct image and kernel values so reversed or transposed kernels give different results + images = np.arange(2 * 8 * 9, dtype=float).reshape(2, 8, 9) + kernels = np.arange(2 * np.prod(shape), dtype=float).reshape(2, *shape) + + # Compute the reference separately for every image and kernel using SciPy + expected = np.array( + [ + [scipy.ndimage.convolve(image, kernel, mode="constant", cval=np.nan) for kernel in kernels] + for image in images + ] + ) + + # Check the stack implementation, including NaN padding at image edges + actual = gu.filters.convolution(images, kernels) + np.testing.assert_allclose(actual, expected, equal_nan=True) + + # Check the optional compiled implementation against the same independent reference + if find_spec("numba") is not None: + accelerated = gu.filters.convolution(images, kernels, engine="numba") + np.testing.assert_allclose(accelerated, expected, equal_nan=True) + + @pytest.mark.parametrize("kernel_shape, expected_count", [("square", 25), ("circular", 9)]) + def test_patch_mean_filter_valid_counts( + self, kernel_shape: Literal["square", "circular"], expected_count: int + ) -> None: + """Checks that patch filtering excludes missing values from the mean and reports finite and total counts.""" + + # Remove the center of a constant image so the finite window count decreases by exactly one + values = np.ones((11, 11), dtype=float) + values[5, 5] = np.nan + + # Compute complete patch means and counts without keeping the missing center + mean, counts, kernel_count = gu.filters.mean_filter( + values, + 5, + kernel_shape=kernel_shape, + preserve_nodata=False, + boundless=False, + return_counts=True, + ) + + # Check the full kernel area, remaining finite cells and unchanged mean of one + assert kernel_count == expected_count + assert counts[5, 5] == expected_count - 1 + assert mean[5, 5] == 1 + + # Windows extending beyond the image keep the existing NaN edge convention + assert np.isnan(mean[0, 0]) + + # Check that the compiled engine produces the same means and counts + if find_spec("numba") is not None: + accelerated_mean, accelerated_counts, accelerated_kernel_count = gu.filters.mean_filter( + values, + 5, + kernel_shape=kernel_shape, + engine="numba", + preserve_nodata=False, + boundless=False, + return_counts=True, + ) + np.testing.assert_allclose(accelerated_mean, mean, equal_nan=True) + np.testing.assert_allclose(accelerated_counts, counts, equal_nan=True) + assert accelerated_kernel_count == kernel_count + + class TestGaussianFilter: """Tests for the Gaussian filter applied to raster data.""" @@ -43,9 +115,10 @@ def test_gauss(self) -> None: assert np.nanmax(raster_with_nans) >= np.nanmax(raster_sm) # 3D arrays - array_3d = np.vstack((raster_array[np.newaxis, :], raster_array[np.newaxis, :])) + array_3d = np.stack((raster_array, raster_array + 100)) raster_sm = gu.filters.gaussian_filter(array_3d, sigma=5) - assert array_3d.shape == raster_sm.shape + expected = np.stack([gu.filters.gaussian_filter(band, sigma=5) for band in array_3d]) + np.testing.assert_allclose(raster_sm, expected) # 1D array should raise data = raster_array[:, 0] @@ -53,7 +126,7 @@ def test_gauss(self) -> None: class TestStatisticalFilters: - """Tests for statistical filters: mean, median, min, max.""" + """Checks statistical filters and SciPy/Numba consistency across all built-in filters.""" landsat_data = gu.Raster(gu.examples.get_path("everest_landsat_b4")).astype(np.float32) @@ -111,24 +184,45 @@ def test_filters(self, name: str, filter_func: Callable[[NDArrayNum], NDArrayNum assert np.max(raster_filtered) == np.nanmax(raster_with_nans_filtered) if name != "mean": - array_3d = np.vstack((raster_array[np.newaxis, :], raster_array[np.newaxis, :])) + array_3d = np.stack((raster_array, raster_array + 100)) raster_filtered = filter_func(array_3d) - assert array_3d.shape == raster_filtered.shape + expected = np.stack([filter_func(band) for band in array_3d]) + np.testing.assert_allclose(raster_filtered, expected) data = raster_array[:, 0] pytest.raises(ValueError, filter_func, data) - def test_median_filter_nan_consistency(self) -> None: - """Test that different median filter engines return consistent results with NaNs.""" + @pytest.mark.parametrize( + "filter_func, kwargs", + [ + (gu.filters.gaussian_filter, {"sigma": 1}), + (gu.filters.median_filter, {"size": 3}), + (gu.filters.mean_filter, {"size": 3}), + (gu.filters.min_filter, {"size": 3}), + (gu.filters.max_filter, {"size": 3}), + (gu.filters.distance_filter, {"sigma": 1, "outlier_threshold": 2}), + ], + ) + def test_filter_engines_consistent(self, filter_func: Callable[..., NDArrayNum], kwargs: dict[str, Any]) -> None: + """Checks that the SciPy and Numba engines return consistent results with missing values.""" pytest.importorskip("numba") + # Filter the same array with SciPy and Numba arr = np.array([[1, 2, np.nan], [4, np.nan, 6], [7, 8, 9]], dtype=np.float32) - filtered_scipy = gu.filters.median_filter(arr, size=3, engine="scipy") - filtered_numba = gu.filters.median_filter(arr, size=3, engine="numba") + filtered_scipy = filter_func(arr, engine="scipy", **kwargs) + filtered_numba = filter_func(arr, engine="numba", **kwargs) + # Check equal shapes, values and missing pixels assert filtered_scipy.shape == arr.shape assert filtered_numba.shape == arr.shape - assert np.allclose(filtered_scipy, filtered_numba, equal_nan=True) + np.testing.assert_allclose(filtered_scipy, filtered_numba, equal_nan=True, rtol=1e-6, atol=1e-6) + + # Filters that support stacks must also agree while keeping bands independent + if filter_func is not gu.filters.mean_filter: + array_3d = np.stack((arr, arr + 10)) + filtered_scipy = filter_func(array_3d, engine="scipy", **kwargs) + filtered_numba = filter_func(array_3d, engine="numba", **kwargs) + np.testing.assert_allclose(filtered_scipy, filtered_numba, equal_nan=True, rtol=1e-6, atol=1e-6) @pytest.mark.skipif(find_spec("numba") is not None, reason="Only runs if numba is missing.") def test_filter_numba__missing_dep(self) -> None: diff --git a/tests/test_interface/test_gridding.py b/tests/test_interface/test_gridding.py index b20379cdd..273227504 100644 --- a/tests/test_interface/test_gridding.py +++ b/tests/test_interface/test_gridding.py @@ -777,7 +777,9 @@ def test_grid__point_raster_input_combinations( raster_dask: bool, tmp_path: Path, ) -> None: - """Grid every combination of eager and Dask point-cloud and raster reference inputs.""" + """ + Checks that grid() returns a lazy output when the point source or raster reference uses Dask. + """ pytest.importorskip("dask_geopandas") import dask.array as da @@ -785,9 +787,10 @@ def test_grid__point_raster_input_combinations( # Write both inputs so their Dask variants use the same values and georeferencing point_file = tmp_path / "points.gpkg" self.points.to_file(point_file) + # Match the requested grid coordinates to the source points so nearest neighbors have no ties reference = Raster.from_array( np.zeros((3, 3), dtype=np.uint8), - transform=rio.transform.from_origin(-0.5, 2.5, 1, 1), + transform=rio.transform.from_origin(0, 2, 1, 1), crs=self.points.crs, ) raster_file = tmp_path / "reference.tif" @@ -806,7 +809,7 @@ def test_grid__point_raster_input_combinations( dist_nodata_pixel=2, ) - # The point-cloud backend controls whether the output itself is eager or Dask + # Either lazy input selects Dask output, while two eager inputs return an eager raster output = ( points.pc.grid(ref=raster, resampling="nearest", dist_nodata_pixel=2) if point_dask @@ -816,24 +819,27 @@ def test_grid__point_raster_input_combinations( dist_nodata_pixel=2, ) ) - if point_dask: + if point_dask or raster_dask: assert isinstance(output, xr.DataArray) assert isinstance(output.data, da.Array) assert not output._in_memory computed_output = output.compute() - assert not points.pc.is_loaded + if point_dask: + assert not points.pc.is_loaded else: assert isinstance(output, Raster) assert output.is_loaded computed_output = output - # A Dask reference supplies only grid metadata and chunks, and must stay lazy + # A Dask reference supplies its spatial chunks without reading any of its raster values if raster_dask: assert isinstance(raster.data, da.Array) assert not raster._in_memory - if point_dask: - assert output.data.chunks == raster.data.chunks + assert output.data.chunks == raster.data.chunks + # Grid coordinates coincide with source points; raster rows reverse the points' increasing Y order + expected_values = self.points["z"].to_numpy().reshape(3, 3)[::-1] + np.testing.assert_array_equal(computed_output.data, expected_values) assert expected.raster_equal(computed_output, warn_failure_reason=True, strict_masked=False) @pytest.mark.parametrize("resampling", ["nearest", "idw", "mean"]) @@ -913,3 +919,25 @@ def test_grid__dask_multiprocessing_error(self, tmp_path: Path) -> None: mp_config=MultiprocConfig(chunks=(2, 2), outfile=str(tmp_path / "grid-error.tif")), ) assert not points.pc.is_loaded + + def test_grid__error_dask_reference_with_multiprocessing(self, tmp_path: Path) -> None: + """ + Checks that an eager point source with a Dask raster reference rejects multiprocessing before computing. + """ + + # Place source points at the requested grid coordinates and open the reference with spatial chunks + points = PointCloud(self.points, data_column="z") + reference_file = tmp_path / "reference.tif" + reference = Raster.from_array( + np.zeros((3, 3), dtype=np.uint8), rio.transform.from_origin(0, 2, 1, 1), self.points.crs + ) + reference.to_file(reference_file) + lazy_reference = gu.open_raster(str(reference_file), chunks={"x": 2, "y": 2}) + + # Reject competing schedulers without loading the reference or creating a multiprocessing output + output_file = tmp_path / "grid-error.tif" + config = MultiprocConfig(chunks=(2, 2), outfile=str(output_file)) + with pytest.raises(ValueError, match="Cannot use Multiprocessing and Dask simultaneously"): + points.grid(ref=lazy_reference, mp_config=config) + assert not lazy_reference._in_memory + assert not output_file.exists() diff --git a/tests/test_interface/test_interpolation.py b/tests/test_interface/test_interpolation.py index 5869257f2..7ec70a59c 100644 --- a/tests/test_interface/test_interpolation.py +++ b/tests/test_interface/test_interpolation.py @@ -4,7 +4,7 @@ import re import tempfile from pathlib import Path -from typing import Literal +from typing import Any, Literal import geopandas as gpd import numpy as np @@ -16,7 +16,7 @@ import geoutils as gu from geoutils import examples, open_raster -from geoutils._misc import silence_rasterio_message +from geoutils._misc import import_optional, silence_rasterio_message from geoutils.interface._nodata import NodataPropagation from geoutils.interface.interpolation import ( _get_dist_nodata_spread, @@ -314,16 +314,10 @@ def test_interp_points__synthetic(self, tag_aop: Literal["Area", "Point"] | None # Check the bilinear interpolation matches the mean value of those 4 points (equivalent as its the middle) assert raster_points_in[i] == np.mean([arr[xlow, ylow], arr[xupp, ylow], arr[xupp, yupp], arr[xlow, yupp]]) - # Check bilinear extrapolation for points at 1 spacing outside from the input grid - points_out = ( - [(-1, i) for i in np.arange(1, 4)] - + [(i, -1) for i in np.arange(1, 4)] - + [(4, i) for i in np.arange(1, 4)] - + [(i, 4) for i in np.arange(4, 1)] - ) - points_out_xy = tuple(zip(*points_out)) + # Select points beyond the outer half pixels for every pixel interpretation + points_out_xy = raster.ij2xy([-2, 1, 4, 1], [1, -2, 1, 4], shift_area_or_point=shift_aop) with pytest.warns(UserWarning, match="All provided points were outside of raster bounds"): - raster_points_out = raster.interp_points(points_out_xy, as_array=True) + raster_points_out = raster.interp_points(points_out_xy, shift_area_or_point=shift_aop, as_array=True) assert all(~np.isfinite(raster_points_out)) # To use cubic or quintic, we need a larger grid (minimum 6x6, but let's aim bigger with 50x50) @@ -358,7 +352,7 @@ def test_interp_points__synthetic(self, tag_aop: Literal["Area", "Point"] | None # see https://github.com/GlacioHack/geoutils/issues/533 assert np.allclose(raster_points_mapcoords, raster_points_interpn) - # Check that, outside the edge, the interpolation fails and returns a NaN + # Nearest and linear include the outer half pixels; splines keep their stricter coordinate bounds index_x_edge_rand = [-0.5, -0.5, -0.5, 25, 25, 49.5, 49.5, 49.5] index_y_edge_rand = [-0.5, 25, 49.5, -0.5, 49.5, -0.5, 25, 49.5] @@ -383,8 +377,13 @@ def test_interp_points__synthetic(self, tag_aop: Literal["Area", "Point"] | None as_array=True, ) - assert all(~np.isfinite(raster_points_mapcoords_edge)) - assert all(~np.isfinite(raster_points_interpn_edge)) + finite = ( + [True, True, False, True, False, False, False, False] + if method in {"nearest", "linear"} + else [False] * 8 + ) + np.testing.assert_array_equal(np.isfinite(raster_points_mapcoords_edge), finite) + np.testing.assert_array_equal(np.isfinite(raster_points_interpn_edge), finite) @pytest.mark.parametrize("shape", [(3, 7), (7, 3)]) # landscape and portrait exercise different bounds axes def test_interp_points__nonsquare(self, shape: tuple[int, int]) -> None: @@ -691,6 +690,48 @@ def test_interp_point__nodata_propag( assert np.allclose(vals, vals_near, equal_nan=False, rtol=10e-4) assert np.allclose(vals2, vals2_near, equal_nan=False, rtol=10e-4) + @pytest.mark.parametrize("method", ["interp_points", "reduce_points"]) + @pytest.mark.parametrize("point_input_type", ["pointcloud", "accessor", "geodataframe", "latlon"]) + @pytest.mark.parametrize("all_outside", [False, True]) + def test_methods__point_output_coordinates(self, method: str, point_input_type: str, all_outside: bool) -> None: + """Checks that interpolation and window reduction return points in the raster CRS, including outside points.""" + + # Place the first point inside cell (1, 1) for both sampling methods and the second beyond the grid + values = np.arange(36, dtype=float).reshape(6, 6) + raster = gu.Raster.from_array(values, rio.transform.from_origin(500_000, 4_100_000, 10, 10), crs=32610) + expected_x = np.array([500_012.5, 501_042.5]) + expected_y = np.array([4_099_987.5, 4_099_957.5]) + expected_values = np.array([values[1, 1], np.nan]) + if all_outside: + expected_x[0] += 1_000 + expected_values[0] = np.nan + + # Express the same query locations as spatial objects or longitude/latitude arrays + longitude, latitude = reproject_to_latlon((expected_x, expected_y), raster.crs) + pointcloud = gu.PointCloud.from_xyz(longitude, latitude, np.zeros(2), crs=4326) + point_inputs = { + "pointcloud": pointcloud, + "accessor": pointcloud.ds.pc, + "geodataframe": pointcloud.ds, + "latlon": (longitude, latitude), + } + point_input = point_inputs[point_input_type] + + # Request point output through the public API; interpolation warns when every query is outside + sample = getattr(raster, method) + options = {"method": "nearest"} if method == "interp_points" else {} + if method == "interp_points" and all_outside: + with pytest.warns(UserWarning, match="All provided points were outside of raster bounds"): + result = sample(point_input, input_latlon=point_input_type == "latlon", **options) + else: + result = sample(point_input, input_latlon=point_input_type == "latlon", **options) + + # Check coordinates and values in input order; longitude/latitude conversion rounds to centimeter accuracy + assert result.crs == raster.crs + np.testing.assert_allclose(result.geometry.x, expected_x, rtol=0, atol=1e-2) + np.testing.assert_allclose(result.geometry.y, expected_y, rtol=0, atol=1e-2) + np.testing.assert_allclose(result.data, expected_values, rtol=0, atol=0, equal_nan=True) + def test_reduce_points(self) -> None: """ Test reduce points. @@ -829,6 +870,257 @@ def test_reduce_points(self) -> None: class TestInterpPointsChunked: """Compare point interpolation across eager, Dask and Multiprocessing backends.""" + @pytest.mark.parametrize("backend", ["eager", "dask", "multiprocessing"]) + @pytest.mark.parametrize("dtype", ["int16", "float64"]) + @pytest.mark.parametrize( + "method,propagation,spread", + [ + ("nearest", "gdal", None), + ("linear", "ignore", None), + ("linear", "gdal", 0), + ("linear", "propagate", 1), + ("cubic", "gdal", None), + ("cubic", "gdal", 0), + ], + ) + def test_interp_points__validity_matches_explicit_source( + self, + backend: str, + dtype: str, + method: str, + propagation: NodataPropagation, + spread: int | None, + ) -> None: + """ + Checks that validity interpolation matches an explicit one/NaN raster with each backend and nodata rule. + """ + + # 1/ Prepare two bands with nodata in different cells and a separate validity reference + # Integer masks and floating NaNs must both describe unavailable cells in the selected second band + data = np.arange(2 * 20 * 24).reshape(2, 20, 24).astype(dtype) + invalid = np.zeros(data.shape, dtype=bool) + invalid[0, 3, 4] = True + invalid[1, 7:9, 8:10] = True + if dtype == "float64": + data[1, 12, 14] = np.nan + raster = gu.Raster.from_array(np.ma.array(data, mask=invalid), Affine(1, 0, 0, 0, -1, 20), 32632, nodata=-9999) + + # Construct the previous cosampling validity layer independently of the private interpolation option + finite = np.isfinite(np.ma.getdata(raster.data[1])) & ~np.ma.getmaskarray(raster.data[1]) + validity = gu.Raster.from_array( + np.where(finite, 1, np.nan).astype(np.float32), raster.transform, raster.crs, nodata=np.nan + ) + rows = np.array([-2, 0, 3, 7, 8, 10, 12, 16, 19, 23]) + columns = np.array([1, 0, 4, 8, 9, 11, 14, 18, 23, 1]) + x, y = raster.ij2xy(rows, columns) + points = (x + 0.2, y - 0.3) + + # 2/ Evaluate both representations with the same interpolation and backend settings + # Use multiple tiles so the validity conversion must run inside the shared block kernel + options: dict[str, Any] = { + "points": points, + "method": method, + "nodata_propagation": propagation, + "dist_nodata_spread": spread, + } + raster_input: Any = raster + validity_input: Any = validity + if backend == "dask": + import_optional("dask") + raster_input = raster.to_xarray().chunk({"x": 12, "y": 10}).rst + validity_input = validity.to_xarray().chunk({"x": 12, "y": 10}).rst + elif backend == "multiprocessing": + options["mp_config"] = MultiprocConfig(chunks=(10, 12)) + result = raster_input.interp_points(band=2, as_array=True, _validity_only=True, **options) + expected = validity_input.interp_points(as_array=True, **options) + if backend == "dask": + result, expected = import_optional("dask").compute(result, expected) + + # 3/ Check the exact finite locations and values, including points near holes and outside bounds + assert result.dtype == np.float32 + np.testing.assert_array_equal(result, expected) + assert np.any(np.isfinite(result)) + assert np.any(np.isnan(result)) + + @pytest.mark.parametrize("validity_only", [False, True]) + def test_interp_points__validity_does_not_load_multiprocessing_source( + self, validity_only: bool, tmp_path: Path + ) -> None: + """ + Checks that real workers interpolate selected values or validity while the complete raster stays unloaded. + """ + + from geoutils.multiproc.cluster import MpCluster + + # Store a raster whose nodata center is distinguishable from finite and out-of-bounds points + data = np.arange(240, dtype=np.float64).reshape(2, 10, 12) + data[0, 1, 1] = np.nan + data[1, 4, 5] = np.nan + raster = gu.Raster.from_array(data, Affine(1, 0, 0, 0, -1, 10), 32632, nodata=-9999) + filename = tmp_path / "validity-source.tif" + raster.to_file(filename) + unloaded = gu.Raster(filename, load_data=False) + points = raster.ij2xy(np.array([1, 4, 8, -2]), np.array([1, 5, 10, 1])) + + # Let worker tiles read the file and build only their local one/NaN arrays + assert not unloaded.is_loaded + with MpCluster({"nb_workers": 2}) as cluster: + result = unloaded.interp_points( + points, + method="nearest", + band=2, + as_array=True, + _validity_only=validity_only, + mp_config=MultiprocConfig(chunks=(5, 6), cluster=cluster), + ) + + # Preserve the original storage and report validity at the sampled locations as floating values + assert not unloaded.is_loaded + assert unloaded.bands == (1, 2) + assert result.dtype == (np.float32 if validity_only else np.float64) + expected = [1, np.nan, 1, np.nan] if validity_only else [133, np.nan, 226, np.nan] + np.testing.assert_array_equal(result, expected) + + @pytest.mark.parametrize("validity_only", [False, True]) + @pytest.mark.parametrize("as_array", [False, True]) + def test_interp_points__dask_points_defer_interpolation( + self, validity_only: bool, as_array: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + """ + Checks that sizing a lazy point result does not interpolate values or discard duplicate labels and geometry. + """ + + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + import geoutils.interface.interpolation as interpolation + from geoutils.pointcloud.pd_accessor import _register_dask_pointcloud_accessor + + # Synthetic Dask point tables bypass open_pointcloud(), which normally registers the optional accessor + _register_dask_pointcloud_accessor() + + # Use duplicate labels and a nodata cell so row order and optional validity conversion are both visible + data = np.arange(30, dtype=np.float64).reshape(5, 6) + data[2, 3] = np.nan + raster = gu.Raster.from_array(data, Affine(1, 0, 0, 0, -1, 5), 32632, nodata=-9999) + x, y = raster.ij2xy(np.array([1, 2, 3]), np.array([1, 3, 4])) + points = gpd.GeoDataFrame(geometry=gpd.points_from_xy(x, y), index=["a", "a", "b"], crs=raster.crs) + lazy_points = dgpd.from_geopandas(points, npartitions=2, sort=False) + + # Reject kernel execution while constructing the graph; computing input partition lengths is allowed + def fail_interpolation(*args: Any, **kwargs: Any) -> None: + """Reject interpolation before the caller computes the lazy result.""" + + raise AssertionError("Interpolation ran while constructing the lazy output.") + + with monkeypatch.context() as patch: + patch.setattr(interpolation, "_interp_points_base", fail_interpolation) + output = raster.interp_points( + lazy_points, method="nearest", as_array=as_array, _validity_only=validity_only + ) + + # Compute only after restoring the kernel, then compare values and the complete point output when requested + computed = output.compute() + expected = [1, np.nan, 1] if validity_only else [7, np.nan, 22] + if as_array: + np.testing.assert_array_equal(computed, expected) + else: + np.testing.assert_array_equal(computed["z"], expected) + np.testing.assert_array_equal(computed.geometry.to_numpy(), points.geometry.to_numpy()) + assert computed.crs == points.crs + np.testing.assert_array_equal(computed.index, points.index) + assert output.pc.data_column == "z" + + def test_interp_points__boolean_outside_bounds(self) -> None: + """Checks that nearest interpolation of booleans returns NaNs outside the raster with every backend.""" + + import_optional("dask") + + # Alternate true and false pixels and interleave interior points with points beyond the raster + rows, columns = np.indices((5, 6)) + values = (rows + columns) % 2 == 0 + raster = gu.Raster.from_array(values, Affine(1, 0, 0, 0, -1, 5), 32632) + x, y = raster.ij2xy(np.array([1, -2, 2, 7]), np.array([1, 1, 1, 1])) + expected = np.array([1, np.nan, 0, np.nan], dtype=np.float32) + + # Convert through integers because Raster.to_xarray() uses GDAL, which cannot store booleans + lazy = raster.astype("uint8").to_xarray().astype(bool).chunk({"x": 3, "y": 2}) + + # Interpolate the same boolean raster with eager, Dask and multiprocessing inputs + options = {"points": (x, y), "method": "nearest", "as_array": True} + results = [ + raster.interp_points(**options), + lazy.rst.interp_points(**options).compute(), + raster.interp_points(**options, mp_config=MultiprocConfig(chunks=(2, 3))), + ] + + # Return exact boolean values and represent out-of-bounds samples as nodata values + for result in results: + assert result.dtype == np.float32 + np.testing.assert_array_equal(result, expected) + + @pytest.mark.parametrize("area_or_point", ["Area", "Point"]) + @pytest.mark.parametrize("method", ["nearest", "linear"]) + @pytest.mark.parametrize("chunks", [(13, 17), (19, 11)]) + def test_interp_points__fractional_projected_coordinates_exact_backends( + self, area_or_point: str, method: str, chunks: tuple[int, int] + ) -> None: + """Checks that fractional projected coordinates give exact interpolation values regardless of raster tiling.""" + + da = pytest.importorskip("dask.array") + + # Large coordinates and fractional pixels expose rounding from recalculating a tile's local transform + rows, cols = np.indices((35, 43)) + values = (900 + rows * 1.3 + cols * 1.7 + 5 * np.sin(rows / 3)).astype(np.float32) + values[10:13, 14:17] = np.nan + transform = Affine(20, 0, 500000, 0, -20, 8600000) + raster = gu.Raster.from_array(values, transform, 32633, nodata=-9999, area_or_point=area_or_point) + xx, yy = raster.coords(grid=True) + points = (xx.ravel()[::7] + 3, yy.ravel()[::7] - 4) + options = {"points": points, "method": method, "as_array": True} + + # Interpolate the same values eagerly and with two independent rectangular block layouts + expected = raster.interp_points(**options) + lazy = raster.to_xarray().chunk({"y": chunks[0], "x": chunks[1]}) + dask_result = lazy.rst.interp_points(**options) + multiproc_result = raster.interp_points(**options, mp_config=MultiprocConfig(chunks=(11, 15))) + + # Pixel weights and missing values must be identical for every backend + assert isinstance(dask_result, da.Array) + np.testing.assert_array_equal(dask_result.compute(), expected) + np.testing.assert_array_equal(multiproc_result, expected) + + @pytest.mark.parametrize("method", ["nearest", "linear"]) + def test_interp_points__outer_half_pixels(self, method: Literal["nearest", "linear"]) -> None: + """Checks that every backend keeps points on the raster's outer half pixels.""" + + pytest.importorskip("dask.array") + + # Locate points on all outer half pixels and just beyond the raster + raster = gu.Raster.from_array(np.arange(30.0).reshape(5, 6), Affine(1, 0, 0, 0, -1, 5), 32632) + rows = np.array([-0.25, -0.25, 2, 4.25, -0.75, 4.75]) + columns = np.array([0.25, -0.25, 5.25, 5.25, 0, 0]) + x, y = raster.ij2xy(rows, columns) + + # The first four points remain inside the pixel footprint; the last two fall outside it + expected = np.array([0 if method == "nearest" else 0.25, 0, 17, 29, np.nan, np.nan]) + + # Interpolate the same points with eager, Dask and multiprocessing inputs + lazy = raster.to_xarray().chunk({"x": 3, "y": 2}) + results = [ + raster.interp_points((x, y), method=method, as_array=True), + lazy.rst.interp_points((x, y), method=method, as_array=True).compute(), + raster.interp_points((x, y), method=method, as_array=True, mp_config=MultiprocConfig(chunks=(2, 3))), + ] + + # Check that every backend applies the same half-pixel boundary rule + for result in results: + np.testing.assert_allclose(result, expected, equal_nan=True) + + # Cosampling keeps exactly the four points with finite interpolated values + points = gu.PointCloud.from_xyz(x, y, np.ones(len(x)), crs=raster.crs) + sample = raster.cosample(points, resample_method=method) + np.testing.assert_array_equal(sample.ds.index, np.arange(4)) + np.testing.assert_allclose(sample.ds["self"], expected[:4]) + @pytest.mark.parametrize("path_index", [0, 2]) @pytest.mark.parametrize("method", ["nearest", "linear"]) @pytest.mark.parametrize("ninterp", [2, 100]) diff --git a/tests/test_interface/test_raster_point.py b/tests/test_interface/test_raster_point.py index 5e41c244c..b63110d4f 100644 --- a/tests/test_interface/test_raster_point.py +++ b/tests/test_interface/test_raster_point.py @@ -10,6 +10,7 @@ import geoutils as gu from geoutils import examples +from geoutils._misc import import_optional class TestRasterPointInterface: @@ -236,3 +237,39 @@ def test_from_pointcloud(self) -> None: ValueError, match="Either grid coordinates or both geotransform and shape must be provided." ): gu.Raster.from_pointcloud_regular(pc1) + + +class TestToPointcloudChunked: + """ + Compare to_pointcloud() outputs from eager and Dask rasters. + + These tests cover the currently eager point outputs and keep the source Dask array. Expand them to cover lazy + outputs when to_pointcloud() returns lazy point data. + """ + + @pytest.mark.parametrize("subsample", [1, 11]) + @pytest.mark.parametrize("as_array", [False, True]) + def test_to_pointcloud__eager_samples_keep_lazy_source(self, subsample: int, as_array: bool) -> None: + """Checks that point sampling returns exact eager values without loading or replacing the Dask source.""" + + import_optional("dask") + import dask.array as da + + # Include a missing cell and uneven chunks to check the mask and deterministic sample order + values = np.arange(63, dtype=np.float32).reshape((7, 9)) + values[2, 3] = np.nan + transform = rio.transform.from_origin(500000, 8600000, 20, 20) + eager = gu.Raster.from_array(values, transform, 32633, nodata=-9999) + source = gu.RasterAccessor.from_array(da.from_array(values, chunks=(3, 4)), transform, 32633, nodata=-9999) + source_array = source.data + options = {"subsample": subsample, "as_array": as_array, "random_state": 42} + + # Sample eager point values while keeping the input array available for lazy operations + expected = eager.to_pointcloud(**options) + actual = source.rst.to_pointcloud(**options) + if as_array: + np.testing.assert_array_equal(expected, actual) + else: + assert expected.pointcloud_equal(actual) + assert source.data is source_array + assert not source._in_memory diff --git a/tests/test_interface/test_rasterization.py b/tests/test_interface/test_rasterization.py index 3338bc7d8..84d6d5c75 100644 --- a/tests/test_interface/test_rasterization.py +++ b/tests/test_interface/test_rasterization.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, Literal import geopandas as gpd import numpy as np @@ -12,6 +12,7 @@ import geoutils as gu from geoutils import examples +from geoutils._misc import import_optional from geoutils.exceptions import InvalidGridError from geoutils.interface import rasterization from geoutils.multiproc import MultiprocConfig @@ -57,6 +58,46 @@ class TestRasterVectorInterface: aster_dem_path = gu.examples.get_path_test("exploradores_aster_dem") aster_outlines_path = gu.examples.get_path_test("exploradores_rgi_outlines") + @pytest.mark.parametrize("method", ["rasterize", "create_mask"]) + @pytest.mark.parametrize("input_type", ["raster", "xarray", "dask"]) + def test_rasterize_create_mask__multi_band( + self, + method: Literal["rasterize", "create_mask"], + input_type: Literal["raster", "xarray", "dask"], + ) -> None: + """Checks that vector rasterization returns one spatial layer from every multi-band reference type.""" + + import_optional("dask") + import dask.array as da + + # Create equivalent native, Xarray and chunked Xarray references with two bands + raster = gu.Raster.from_array(np.zeros((2, 4, 4)), (1, 0, 8, 0, -1, 12), 4326) + if input_type == "raster": + reference = raster + else: + reference = raster.to_xarray() + if input_type == "dask": + reference = reference.chunk({"band": 1, "y": 3, "x": 2}) + # Rasterize the vector or create a mask from the selected reference + output = getattr(self.vector, method)(ref=reference) + + # Check that the result has one spatial layer and keeps Dask chunk sizes when applicable + assert output.shape == (4, 4) + if input_type == "dask": + assert isinstance(output, xr.DataArray) + assert isinstance(output.data, da.Array) + assert isinstance(reference, xr.DataArray) + assert isinstance(reference.data, da.Array) + assert output.data.chunks == reference.data.chunks[-2:] + else: + assert isinstance(output, gu.Raster) + + # The unit polygon occupies exactly the cell centred on (10.5, 10.5) + expected = np.zeros((4, 4)) + expected[1, 2] = 1 + actual = output.to_numpy() if isinstance(output, xr.DataArray) else output.data + np.testing.assert_array_equal(actual, expected) + def test_rasterize(self) -> None: """Test rasterizing an EPSG:3426 dataset into a projection.""" diff --git a/tests/test_multiproc/test_chunked.py b/tests/test_multiproc/test_chunked.py new file mode 100644 index 000000000..bfec456bc --- /dev/null +++ b/tests/test_multiproc/test_chunked.py @@ -0,0 +1,106 @@ +"""Tests for dividing arrays and georeferenced grids into chunks.""" + +from typing import Any + +import pytest +from rasterio.crs import CRS +from rasterio.transform import from_origin + +from geoutils.multiproc.chunked import ( + ChunkedGeoGrid, + GeoGrid, + iter_chunk_slices, + normalize_chunks, +) + + +class TestArrayChunks: + """ + Tests for manipulating array chunks: accepted chunk sizes, edge chunks and the order of the resulting slices. + """ + + @pytest.mark.parametrize( + "chunks,expected", + [ + (3, ((3, 3, 1), (3, 2))), + ((3, 2), ((3, 3, 1), (2, 2, 1))), + (((2, 5), (1, 3, 1)), ((2, 5), (1, 3, 1))), + ], + ) + def test_normalize_chunks__supported_forms( + self, + chunks: int | tuple[int, int] | tuple[tuple[int, ...], tuple[int, ...]], + expected: tuple[tuple[int, ...], tuple[int, ...]], + ) -> None: + """Checks that square, rectangular and explicit chunks cover the complete array shape.""" + + # Try square, rectangular and fully specified chunks on the same array + normalized = normalize_chunks(chunks, shape=(7, 5)) + assert normalized == expected + + @pytest.mark.parametrize( + "chunks,error", + [ + (0, ValueError), + ((3, 0), ValueError), + (((2, 5), (2, 2)), ValueError), + ((2, (2, 3)), TypeError), + ((2, 3, 4), ValueError), + ], + ) + def test_normalize_chunks__error_invalid_forms(self, chunks: Any, error: type[Exception]) -> None: + """Checks that zero sizes, incorrect totals and extra axes all raise an error.""" + + with pytest.raises(error): + normalize_chunks(chunks, shape=(7, 5)) + + def test_iter_chunk_slices__row_order_and_clipped_edges(self) -> None: + """Checks that array slices follow row order and stop at the array edges.""" + + # Divide a shape that is not evenly divisible by either chunk size + slices = list(iter_chunk_slices(shape=(5, 7), chunks=(2, 3))) + locations = [tuple((part.start, part.stop) for part in chunk) for chunk in slices] + + # Visit each column chunk before moving to the next row chunk, clipping the final slices to the shape + expected = [ + ((0, 2), (0, 3)), + ((0, 2), (3, 6)), + ((0, 2), (6, 7)), + ((2, 4), (0, 3)), + ((2, 4), (3, 6)), + ((2, 4), (6, 7)), + ((4, 5), (0, 3)), + ((4, 5), (3, 6)), + ((4, 5), (6, 7)), + ] + assert locations == expected + + +class TestChunkedGeoGrid: + """Checks how a georeferenced grid is split into spatial blocks.""" + + def test_chunked_geogrid__block_shapes_and_locations(self) -> None: + """Checks that uneven blocks have the full grid's resolution and occupy their matching locations.""" + + # Create a five-row, seven-column grid divided into six uneven blocks + grid = GeoGrid(transform=from_origin(100, 200, 10, 20), shape=(5, 7), crs=CRS.from_epsg(32633)) + chunked = ChunkedGeoGrid(grid, chunks=((2, 3), (3, 3, 1))) + blocks = chunked.get_blocks_as_geogrids() + + # Check block order, shapes, upper-left coordinates + expected = [ + ((2, 3), (100, 200)), + ((2, 3), (130, 200)), + ((2, 1), (160, 200)), + ((3, 3), (100, 160)), + ((3, 3), (130, 160)), + ((3, 1), (160, 160)), + ] + actual = [(block.shape, (block.transform.c, block.transform.f)) for block in blocks] + assert actual == expected + assert all(block.res == grid.res and block.crs == grid.crs for block in blocks) + assert chunked.flat_block_index((1, 2)) == 5 + + # Check the class raises an error for a row or column block positions outside this grid + with pytest.raises(IndexError): + chunked.flat_block_index((2, 0)) diff --git a/tests/test_multiproc/test_cluster.py b/tests/test_multiproc/test_cluster.py index a3cf3d400..56c499473 100644 --- a/tests/test_multiproc/test_cluster.py +++ b/tests/test_multiproc/test_cluster.py @@ -16,7 +16,7 @@ def sample_function(x: float, y: float) -> float: return x + y -# Function to simulate a long-running task +# Function to simulate a long task def long_running_task(x: float) -> float: time.sleep(0.01) return x * 2 diff --git a/tests/test_multiproc/test_mparray.py b/tests/test_multiproc/test_mparray.py index b86780af3..3aa934d30 100644 --- a/tests/test_multiproc/test_mparray.py +++ b/tests/test_multiproc/test_mparray.py @@ -5,10 +5,12 @@ import os import warnings from multiprocessing import cpu_count +from pathlib import Path from typing import Any import numpy as np import pytest +import rasterio as rio import scipy from numpy import floating @@ -61,6 +63,24 @@ def _custom_func_mask(raster: RasterType) -> gu.Raster: return gu.Raster.from_array(mask_array, raster.transform, raster.crs) +def _custom_func_bands(raster: Raster, n_bands: int) -> Raster: + """ + Return the first source band with known offsets, custom metadata and the worker's process identifier. + """ + + # Add a distinct constant to each output band so nodata or misordered bands are visible + first_band = raster.data[0] if raster.count > 1 else raster.data + bands = np.ma.stack([first_band + index for index in range(n_bands)]) + tags: dict[str, Any] = { + "operation": "band offsets", + "long_name": tuple(f"band_{index}" for index in range(n_bands)), + } + tags["worker_pid"] = os.getpid() + + # Use the tile's grid but deliberately change its pixel interpretation to check result metadata + return Raster.from_array(bands, raster.transform, raster.crs, nodata=-99999, area_or_point="Point", tags=tags) + + class TestTiling: landsat_b4_path = examples.get_path_test("everest_landsat_b4") @@ -200,6 +220,37 @@ def test_multiproc_config_rectangular_chunks(self) -> None: with pytest.raises(TypeError, match="integer or a tuple of two integers"): MultiprocConfig(chunks=(40, 25.0)) # type: ignore + def test_multiproc_config__temporary_outputs(self, tmp_path: Path) -> None: + """ + Checks that temporary configurations isolate and clean files without replacing the original worker cluster. + """ + + # Use a final output path separate from two intermediates needed by the same operation + config = MultiprocConfig(chunks=(4, 5), outfile=str(tmp_path / "final.tif")) + original = config.copy() + with pytest.raises(RuntimeError, match="operation failed"): + with config.temporary() as first, config.temporary() as second: + paths = [Path(first.outfile), Path(second.outfile)] + assert first.outfile != second.outfile + for temporary in (first, second): + assert temporary.chunks == config.chunks + assert temporary.driver == config.driver + assert temporary.cluster is config.cluster + + # Simulate a driver writing a main file and a sidecar before the enclosing operation fails + for path in paths: + path.write_text("intermediate") + path.with_suffix(".aux.xml").write_text("metadata") + raise RuntimeError("operation failed") + + # Remove both complete directories and leave the original configuration and cluster usable + assert all(not path.parent.exists() for path in paths) + assert config.outfile == original.outfile + assert config.chunks == original.chunks + assert config.driver == original.driver + assert config.cluster is original.cluster + assert config.cluster.submit(abs, -3) == 3 + def test_deprecated_map_names_preserve_signatures(self, tmp_path: Any) -> None: """Forward the former top-level map functions with explicit deprecation warnings.""" @@ -385,3 +436,54 @@ def test_map_blocks_rectangular_chunks(self) -> None: tiled_mean = np.nansum([stats["mean"] * stats["valid_count"] for stats in list_stats]) / tiled_count assert abs(total_stats["mean"] - tiled_mean) < tiled_mean * 1e-5 assert total_stats["valid_count"] == tiled_count + + +class TestMapOverlapChunked: + """Checks map_overlap() loading behavior and exact equality with eager raster calculations.""" + + @pytest.mark.parametrize("source_bands, output_bands", [(1, 3), (3, 1)]) + @pytest.mark.parametrize("execution_mode", ["basic", "multiprocessing"]) + def test_map_overlap__changed_bands_and_metadata( + self, tmp_path: Path, source_bands: int, output_bands: int, execution_mode: str + ) -> None: + """ + Checks that map_overlap() changes band count, returns the expected metadata and does not load the source. + """ + + # 1/ Write a small raster with uneven edge tiles and metadata distinct from the function's output + base = np.arange(99, dtype=np.float32).reshape(9, 11) + values = np.stack([base + index * 100 for index in range(source_bands)]) + transform = rio.transform.from_origin(500_000, 4_500_000, 10, 10) + source_file = tmp_path / "source.tif" + output_file = tmp_path / "mapped.tif" + Raster.from_array(values, transform, 32633, area_or_point="Area", tags={"operation": "source"}).to_file( + source_file + ) + source = Raster(source_file) + expected = _custom_func_bands(Raster(source_file), output_bands) + + # 2/ Run synchronous and real process workers through the same shared block writer + with ClusterGenerator(execution_mode, nb_workers=1) as cluster: + config = MultiprocConfig(chunks=(4, 5), outfile=str(output_file), cluster=cluster) + result = map_overlap(_custom_func_bands, source, config, output_bands) + assert not source.is_loaded + assert not result.is_loaded + + # Check metadata before reading data, including types preserved on the returned object + assert result.count == output_bands + assert result.area_or_point == "Point" + assert result.tags["operation"] == "band offsets" + assert result.tags["long_name"] == tuple(f"band_{index}" for index in range(output_bands)) + assert (result.tags["worker_pid"] == os.getpid()) == (execution_mode == "basic") + assert result.raster_equal(expected) + + # 3/ Verify every output band and inspect the actual file metadata independently of the returned Raster + expected = np.stack([base + index for index in range(output_bands)]) + np.testing.assert_array_equal(result.data, expected.squeeze()) + assert result.transform == transform + with rio.open(output_file) as dataset: + assert dataset.count == output_bands + assert dataset.tags()["AREA_OR_POINT"] == "Point" + assert dataset.tags()["operation"] == "band offsets" + assert dataset.nodata == -99999 + np.testing.assert_array_equal(dataset.read(), expected) diff --git a/tests/test_multiproc/test_readers.py b/tests/test_multiproc/test_readers.py new file mode 100644 index 000000000..eeceea45e --- /dev/null +++ b/tests/test_multiproc/test_readers.py @@ -0,0 +1,164 @@ +"""Tests reusable raster and point cloud readers used by the Multiproc backend.""" + +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from typing import Any + +import geopandas as gpd +import numpy as np +import pytest +from numpy.typing import NDArray +from pyproj import CRS +from rasterio.transform import from_origin + +import geoutils as gu +from geoutils._misc import import_optional +from geoutils.multiproc import ClusterGenerator, MultiprocConfig +from geoutils.multiproc.readers import ( + _read_values, + _reader_from_source, + _ValueReader, +) + + +class TestValueReaderChunked: + """ + Checks values read from raster and point cloud files without loading the source objects. + + - Raster tests cover the reading of bands, downsampling, nodata and masks. + - Point cloud tests cover the reading of GeoPackage columns, geometry heights and LAS attributes. + """ + + @pytest.mark.parametrize( + "source_bands,band,downsample,workers", + [(None, 2, 1, False), ([3, 1], 1, 1, True), ([3, 1], 2, 2, False)], + ) + def test_value_reader__raster_windows( + self, source_bands: list[int] | None, band: int, downsample: int, workers: bool, tmp_path: Path + ) -> None: + """Checks that Multiproc readers load only the requested raster band and window, and compares with eager.""" + + # Write a raster file to disk with three distinct bands, one nodata cell, and values preserved by downsampling + base = np.arange(80, dtype=np.int16).reshape(8, 10) + values = np.ma.array(np.stack((base, base + 100, base + 200)), mask=False) + values.mask[:, 2, 3] = True + raster = gu.Raster.from_array(values, from_origin(0, 8, 1, 1), 32633, nodata=-9999) + filename = tmp_path / "values.tif" + raster.to_file(filename) + source = gu.Raster(filename, bands=source_bands, downsample=downsample, load_data=False) + reference = gu.Raster(filename, bands=source_bands, downsample=downsample, load_data=True) + + # Read one small window either directly or in a worker, without loading the source Raster + reader = _reader_from_source(source, band, source, MultiprocConfig(chunks=2)) + assert reader is not None and not source.is_loaded + tile = (slice(1, 3), slice(2, 5)) + block = reader.block(tile) + assert reader.shape == source.shape + assert reader.size == np.prod(source.shape) + assert reader.dtype == values.dtype + assert not source.is_loaded + with ClusterGenerator("multi" if workers else "basic", nb_workers=2) as cluster: + result = cluster.compute(cluster.submit(_read_values, block)) + + # Compare the selected values and nodata mask with an independent small eager read + expected = reference.data[band - 1][tile] + assert np.array_equal(np.ma.getdata(result), np.ma.getdata(expected)) + assert np.array_equal(np.ma.getmaskarray(result), np.ma.getmaskarray(expected)) + assert result.shape == (2, 3) + assert not source.is_loaded + + @pytest.mark.parametrize("file_mask", [False, True]) + def test_value_reader__raster_nodata_and_mask(self, file_mask: bool, tmp_path: Path) -> None: + """Checks that a raster nodata cell and an array or raster mask are applied to one column.""" + + # Write an integer raster file to disk with one nodata cell + values = np.ma.array(np.arange(24, dtype=np.int16).reshape(4, 6), mask=False) + values.mask[1, 2] = True + transform = from_origin(0, 4, 1, 1) + filename = tmp_path / "values.tif" + gu.Raster.from_array(values, transform, 32633, nodata=-9999).to_file(filename) + source = gu.Raster(filename) + keep = np.ones(values.shape, dtype=bool) + keep[2, 2] = False + mask: NDArray[Any] | _ValueReader = keep + # Write a boolean raster mask to disk with a different excluded cell + if file_mask: + mask_filename = tmp_path / "mask.tif" + gu.Raster.from_array(keep, transform, 32633).to_file(mask_filename) + mask_source = gu.Raster(mask_filename, is_mask=True) + mask = _ValueReader(mask_source) + + # Apply the array or raster mask while reading one column from three rows + reader = replace(_ValueReader(source), mask=mask) + tile = (slice(1, 4), slice(2, 3)) + result = _read_values(reader.block(tile)) + + # Nodata and the user mask exclude different cells, so the remaining integer stays unchanged + assert result.shape == (3, 1) and result.dtype == values.dtype + assert np.array_equal(np.ma.getmaskarray(result), [[True], [True], [False]]) + assert result[2, 0] == values[3, 2] + assert not source.is_loaded + if file_mask: + assert not mask_source.is_loaded + + @pytest.mark.parametrize("column", ["height", "weight", None]) + def test_value_reader__geopackage_rows(self, column: str | None, tmp_path: Path) -> None: + """Checks that row slices select point columns or geometry Z without loading the GeoPackage source.""" + + # Write a GeoPackage file to disk with geometry elevations differing from its height/weight columns + positions = np.arange(6) + dataframe = gpd.GeoDataFrame( + {"height": positions.astype(float) + 10, "weight": positions.astype(np.int32) + 100}, + geometry=gpd.points_from_xy(500000 + positions, 5100000 + positions, positions + 1000), + crs=32633, + ) + filename = tmp_path / "points.gpkg" + dataframe.to_file(filename, index=False) + source = gu.PointCloud(filename) + + # Read the second through fourth rows from height, weight, or geometry Z, using a worker process for weight + reader = _reader_from_source(source, column, source, MultiprocConfig(chunks=2)) + assert reader is not None and not source.is_loaded + with ClusterGenerator("multi" if column == "weight" else "basic", nb_workers=2) as cluster: + result = cluster.compute(cluster.submit(_read_values, reader.block(slice(1, 4)))) + + # Read the same rows as point geometries + point_rows = reader.read_points(slice(1, 4)) + + # Check file order, selected values, and that an empty row range returns no data + expected = dataframe.geometry.z if column is None else dataframe[column] + assert np.array_equal(result, expected.iloc[1:4]) + assert np.array_equal(point_rows.geometry.x, dataframe.geometry.x.iloc[1:4]) + assert reader.shape == (len(dataframe),) + assert reader.read(slice(0, 0)).size == 0 + assert reader.read_points(slice(0, 0)).empty + assert not source.is_loaded + + @pytest.mark.parametrize("column", ["Z", "intensity"]) + def test_value_reader__las_rows(self, column: str, tmp_path: Path) -> None: + """Checks that LAS row slices return exact elevations and integer attributes while the source is not loaded.""" + + laspy = import_optional("laspy") + + # Write a LAS file to disk with scaled Z elevations and separate integer intensity values + header = laspy.LasHeader(point_format=6, version="1.4") + header.scales = np.array([0.01, 0.01, 0.01]) + header.add_crs(CRS.from_epsg(32633)) + records = laspy.LasData(header) + positions = np.arange(6) + records.x, records.y, records.z = 500000 + positions, 5100000 + positions, positions / 4 + 10 + records.intensity = positions + 100 + filename = tmp_path / "points.las" + records.write(filename) + source = gu.PointCloud(filename) + + # Read the third through fifth rows from either the Z elevations or the intensity attribute + reader = _ValueReader(source, column) + result = _read_values(reader.block(slice(2, 5))) + expected = np.asarray(records.z if column == "Z" else records.intensity)[2:5] + assert np.array_equal(result, expected) + assert reader.dtype == result.dtype + assert reader.read_points(slice(2, 5)).shape[0] == 3 + assert not source.is_loaded diff --git a/tests/test_pointcloud/test_base.py b/tests/test_pointcloud/test_base.py index bb1c49f73..3f645dcc2 100644 --- a/tests/test_pointcloud/test_base.py +++ b/tests/test_pointcloud/test_base.py @@ -12,6 +12,7 @@ import numpy as np import pandas as pd import pytest +import xarray as xr from geopandas.testing import assert_geodataframe_equal from pandas.testing import assert_frame_equal from pyproj import CRS @@ -70,6 +71,22 @@ def assert_output_equal(output_pc: Any, output_ds: Any, use_allclose: bool = Fal elif isinstance(output_pc, gpd.GeoDataFrame): assert_geodataframe_equal(output_pc, output_ds) + # For tabular statistics + elif isinstance(output_pc, pd.DataFrame): + assert_frame_equal(output_pc, output_ds) + + # For lightweight variogram records + elif isinstance(output_pc, gu.Variogram): + assert isinstance(output_ds, gu.Variogram) + assert np.allclose(output_pc.lags, output_ds.lags) + assert np.allclose(output_pc.semivariance, output_ds.semivariance, equal_nan=True) + assert np.array_equal(output_pc.counts, output_ds.counts) + assert output_pc.model == output_ds.model + + # For labelled pair samples + elif isinstance(output_pc, xr.Dataset): + assert output_pc.identical(output_ds) + # For any other object type else: assert output_pc == output_ds @@ -114,19 +131,46 @@ def test_properties__equality_and_loading(self, prop: str) -> None: methods_and_kwargs = [ ("set_data_column", {"new_data_column": "b2"}), ("copy", {}), + ("reproject", {"crs": 4326}), ("to_xyz", {}), ("to_array", {}), ("to_tuples", {}), ("pointcloud_equal", {"other": "self"}), ("pointcloud_allclose", {"other": "self"}), ("georeferenced_coords_equal", {"pc": "self"}), + ("stats", {}), + ("stats", {"by": {"group": "b2"}, "bins": {"group": 2}, "statistics": "mean"}), ("get_stats", {}), ("subsample", {"subsample": 2, "random_state": 42}), + ("cosample", {"other": "self", "subsample": 2, "random_state": 42}), + ( + "pairsample", + {"n_pairs": 4, "min_distance": 0.5, "max_distance": 2, "strategy": "kdtree", "random_state": 42}, + ), + ( + "variogram", + { + "n_pairs": 4, + "n_lags": 2, + "min_lag": 0.5, + "max_lag": 2, + "strategy": "kdtree", + "random_state": 42, + }, + ), ("to_geoutils", {}), ( "grid", {"grid_coords": (np.array([0.0, 1.0]), np.array([0.0, 1.0])), "resampling": "nearest"}, ), + ( + "grid", + { + "grid_coords": (np.array([0.0, 1.0]), np.array([0.0, 1.0])), + "resampling": "nearest", + "data_column": "b2", + }, + ), ] @pytest.mark.parametrize("method, kwargs", [(f, k) for f, k in methods_and_kwargs]) @@ -138,6 +182,8 @@ def test_methods__equality_and_loading(self, method: str, kwargs: dict[str, Any] pc = PointCloud(self.ds, data_column="b1") ds = self.ds.copy() ds.pc.set_data_column("b1") + if method == "variogram": + pytest.importorskip("skgstat") args_pc = kwargs.copy() args_ds = kwargs.copy() @@ -268,7 +314,7 @@ def test_point_preserving_vector_methods__return_pointcloud(self, method: str, k def test_shared_methods_and_arithmetic_ownership(self) -> None: """Check that shared operations live in the base while arithmetic remains exclusive to PointCloud.""" - shared_methods = {"from_xyz", "pointcloud_equal", "pointcloud_allclose", "get_stats", "grid"} + shared_methods = {"from_xyz", "pointcloud_equal", "pointcloud_allclose", "stats", "get_stats", "grid"} assert shared_methods <= set(PointCloudBase.__dict__) assert shared_methods.isdisjoint(PointCloud.__dict__) assert "__add__" not in PointCloudBase.__dict__ diff --git a/tests/test_pointcloud/test_dataframe.py b/tests/test_pointcloud/test_dataframe.py new file mode 100644 index 000000000..a50bbeadc --- /dev/null +++ b/tests/test_pointcloud/test_dataframe.py @@ -0,0 +1,64 @@ +"""Tests for shared point dataframe assignment and row selection.""" + +from __future__ import annotations + +import geopandas as gpd +import numpy as np +import pytest + +from geoutils._misc import import_optional + + +class TestPointRowsChunked: + """ + Checks shared row operations on Dask point tables. + + Duplicate labels must remain at their original positions, and known partition lengths should be reused without + running the Dask graph. Public cosample() behavior is covered in test_cosampling.py. + """ + + @pytest.mark.parametrize("partitions", [2, 5]) + @pytest.mark.parametrize("native_series", [False, True]) + def test_point_rows__reuse_partition_layout(self, partitions: int, native_series: bool) -> None: + """Checks that duplicate labels and known row counts do not start Dask computation.""" + + import_optional("dask") + import dask.array as da + from dask.callbacks import Callback + + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.dataframe import ( + _assign_point_values, + _point_partition_lengths, + _select_point_rows, + ) + + # Create duplicate labels and uneven partitions so label-based alignment would change selected rows + positions = np.arange(40) + frame = gpd.GeoDataFrame( + {"height": positions}, geometry=gpd.points_from_xy(positions, positions + 1), crs=32632 + ) + frame.index = np.tile(["a", "b", "a", "c"], 10) + lazy = dgpd.from_geopandas(frame, npartitions=partitions, sort=False) + geometry = lazy[["geometry"]] + lengths = None if native_series else _point_partition_lengths(geometry) + + # Matching Series need no counts; arrays with other chunks can reuse one earlier partition-length summary + values = lazy["height"] if native_series else da.from_array(positions, chunks=7) + mask = lazy["height"] % 3 == 0 if native_series else da.from_array(positions % 3 == 0, chunks=9) + tasks = [] + with Callback(pretask=lambda *args: tasks.append(args[0])): + assigned = _assign_point_values(geometry, {"self": values, "other": values * 2}, partition_lengths=lengths) + result = _select_point_rows(assigned, mask, partition_lengths=lengths) + assert tasks == [] + assert isinstance(result, type(lazy)) + assert not lazy.pc.is_loaded and not result.pc.is_loaded + + # Compare original positions and geometry after computing only the requested lazy result + output = result.compute() + expected = frame.iloc[positions % 3 == 0] + assert np.array_equal(output.index, expected.index) + assert np.array_equal(output.geometry, expected.geometry) + assert np.array_equal(output["self"], expected["height"]) + assert np.array_equal(output["other"], expected["height"] * 2) + assert not lazy.pc.is_loaded and not result.pc.is_loaded diff --git a/tests/test_pointcloud/test_pd_accessor_pointcloud.py b/tests/test_pointcloud/test_pd_accessor_pointcloud.py index b47fe7051..6b0679d18 100644 --- a/tests/test_pointcloud/test_pd_accessor_pointcloud.py +++ b/tests/test_pointcloud/test_pd_accessor_pointcloud.py @@ -5,15 +5,19 @@ import os.path import tempfile from importlib.util import find_spec +from pathlib import Path +from typing import Literal import geopandas as gpd import numpy as np import pytest import xarray as xr from geopandas.testing import assert_geodataframe_equal +from pyproj import CRS import geoutils as gu import geoutils.vector.pd_accessor as vector_pd_accessor +from geoutils._misc import import_optional from geoutils.multiproc import MultiprocConfig @@ -30,6 +34,41 @@ class TestPointCloudAccessor: ) fn_las = gu.examples.get_path_test("coromandel_lidar") + @pytest.mark.parametrize("suffix", [".las", ".laz"]) + @pytest.mark.parametrize("columns", ["main", "all", ["Z", "intensity"]]) + def test_open_pointcloud__empty_las_dask( + self, tmp_path: Path, suffix: str, columns: Literal["main", "all"] | list[str] + ) -> None: + """Checks that empty LAS/LAZ files open lazily with the same columns, dtypes and CRS as eager reading.""" + + laspy = pytest.importorskip("laspy") + if suffix == ".laz": + pytest.importorskip("lazrs") + dgpd = pytest.importorskip("dask_geopandas") + from dask.callbacks import Callback + + # Write a valid LAS header without any point records + # LasPy provides an independent fixture for the GeoUtils reader, including optional LAZ compression + path = tmp_path / ("empty" + suffix) + header = laspy.LasHeader(point_format=6, version="1.4") + header.add_crs(CRS.from_epsg(32633)) + laspy.LasData(header).write(path) + expected = gu.open_pointcloud(str(path), columns=columns) + + # Build a lazy collection and inspect metadata without executing a partition + tasks = [] + with Callback(pretask=lambda *args: tasks.append(args[0])): + source = gu.open_pointcloud(str(path), columns=columns, chunks=3) + assert isinstance(source, dgpd.GeoDataFrame) + assert source.pc.point_count == 0 + assert source.pc.crs == expected.crs + assert tasks == [] + graph = source.expr + + # Compute the empty collection and check its columns, types and unchanged lazy source + assert_geodataframe_equal(source.compute(), expected) + assert source.expr is graph and not source.pc.is_loaded + def test_accessor(self) -> None: """Expose point-cloud metadata, values and conversion through the accessor.""" @@ -168,9 +207,10 @@ def test_reproject_pointcloud__dask_geopandas(self) -> None: ], ) def test_geometric_methods__dask_geopandas(self, method: str, kwargs: dict[str, object]) -> None: - """Keep copied, cropped and translated point partitions lazy and equal to eager GeoPandas.""" + """Checks that lazy copies have the same location metadata while cropping and translation recalculate it.""" - dgpd = pytest.importorskip("dask_geopandas") + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from dask.callbacks import Callback # Open one on-disk source lazily and build the expected result through the eager accessor temp_dir = tempfile.TemporaryDirectory() @@ -178,16 +218,42 @@ def test_geometric_methods__dask_geopandas(self, method: str, kwargs: dict[str, self.gdf.to_file(temp_file) ds = gu.open_pointcloud(temp_file, data_column="z", chunks=5) expected = getattr(self.gdf.pc, method)(**kwargs) + source_count, source_bounds = ds.pc.point_count, ds.pc.bounds # Each dataframe operation should add work without evaluating any point partition - output = getattr(ds.pc, method)(**kwargs) + tasks = [] + with Callback(pretask=lambda *args: tasks.append(args[0])): + output = getattr(ds.pc, method)(**kwargs) + assert output.pc.crs == expected.crs + assert output.pc.data_column == "z" + if method == "copy": + # Changing only values does not move points, so copies can reuse counts and bounds + assert output.pc.bounds == source_bounds + assert output.pc.point_count == source_count + value_copy = ds.pc.copy(new_array=ds.pc.data * 2) + assert value_copy.pc.bounds == source_bounds + assert value_copy.pc.point_count == source_count + else: + assert output.pc.bounds is None + assert tasks == [] assert isinstance(output, dgpd.GeoDataFrame) assert not ds.pc.is_loaded assert not output.pc.is_loaded + + # Recount selected rows only on request; the original file's cached count and bounds are unchanged + assert output.pc.point_count == len(expected) + assert ds.pc.point_count == source_count + assert ds.pc.bounds == source_bounds assert_geodataframe_equal(output.compute(), expected) assert not ds.pc.is_loaded assert not output.pc.is_loaded + # Compute replacement values only on request and preserve every original point coordinate + if method == "copy": + expected_values = self.gdf.copy() + expected_values["z"] *= 2 + assert_geodataframe_equal(value_copy.compute(), expected_values) + def test_to_file__dask_geopandas(self) -> None: """Write a lazy point cloud to a regular GeoPandas-supported vector file.""" @@ -306,3 +372,94 @@ def test_load_las__multiprocessing(self) -> None: # Chunk scheduling must not change point order, values or metadata assert pc_chunked.pointcloud_equal(pc) + + +class TestPointCloudElevationMetadata: + """ + Checks elevation column and CRS metadata owned by the Pandas point cloud accessor. + + The tests cover direct Dask GeoDataFrames with empty and populated partitions, independent CRS metadata after + reprojection, and explicit use of 3D geometry when auxiliary numeric columns are present. Dask checks also keep + the original source graph and avoid computing partitions for metadata-only operations. + """ + + @pytest.mark.parametrize("point_count", [0, 1, 5]) + def test_crs__dask_geometry_metadata(self, point_count: int) -> None: + """Checks that Dask point accessors read existing geometry CRS without a file metadata cache or computation.""" + + dgpd = pytest.importorskip("dask_geopandas") + from dask.callbacks import Callback + + from geoutils.pointcloud.pd_accessor import _register_dask_pointcloud_accessor + + # Build a Dask GeoDataFrame directly, including empty and single-point collections + coordinates = np.arange(point_count, dtype=float) + frame = gu.PointCloudAccessor.from_xyz( + 500000 + coordinates * 20, + 8600000 + coordinates * 20, + coordinates, + crs="EPSG:32633+5703", + data_column="height", + ) + frame["intensity"] = coordinates + 100 + _register_dask_pointcloud_accessor() + source = dgpd.from_geopandas(frame, chunksize=2) + graph = source.expr + + # Read CRS and select elevations using metadata alone + # File readers populate a GeoUtils CRS cache, but direct dataframe construction must also work + tasks = [] + with Callback(pretask=lambda *args: tasks.append(args[0])): + assert source.pc.crs == frame.crs + source.pc.set_data_column("height") + assert source.pc.crs == frame.crs + + # Check that metadata reads keep the original CRS and source graph without loading any partitions + assert tasks == [] + assert source.expr is graph and not source.pc.is_loaded + assert source.crs == frame.crs + + @pytest.mark.parametrize("lazy", [False, True]) + def test_reproject__metadata_is_independent(self, lazy: bool) -> None: + """Checks that point cloud reprojection changes only the result CRS and keeps the source metadata.""" + + # A small projected point cloud makes both the reference coordinates and metadata deterministic + frame = gu.PointCloudAccessor.from_xyz([500000.0, 500020.0], [8600000.0, 8600020.0], [10.0, 20.0], crs=32633) + source = frame + if lazy: + dgpd = pytest.importorskip("dask_geopandas") + from geoutils.pointcloud.pd_accessor import ( + _register_dask_pointcloud_accessor, + ) + + _register_dask_pointcloud_accessor() + source = dgpd.from_geopandas(frame, npartitions=2) + original_crs = source.pc.crs + assert original_crs == frame.crs + + # Reprojection must not reuse a mutable CRS cache belonging to the source accessor + result = source.pc.reproject(crs=32632) + computed = result.compute() if lazy else result + assert_geodataframe_equal(computed, frame.to_crs(32632)) + assert source.pc.crs == original_crs + assert result.pc.crs == computed.crs + if lazy: + assert not source.pc.is_loaded and not result.pc.is_loaded + + def test_data_column__explicit_geometry_elevations(self) -> None: + """Checks that 3D geometry stays selected for elevations when auxiliary columns are present.""" + + # Keep elevations in 3D geometry and a distinct auxiliary column that must not become the main data + frame = gu.PointCloudAccessor.from_xyz([1.0, 2.0], [3.0, 4.0], [10.0, 20.0], crs=32633, use_z=True) + frame["intensity"] = np.array([2, 4], dtype=np.uint16) + assert frame.pc.data_column is None + np.testing.assert_array_equal(frame.pc.data, [10.0, 20.0]) + + # Two accessors over the same dataframe must observe the same current elevation selection + other = gu.PointCloudAccessor(frame) + with pytest.warns(UserWarning, match="Overriding 3D points"): + other.set_data_column("intensity") + assert frame.pc.data_column == "intensity" + other.set_data_column(None) + assert frame.pc.data_column is None + np.testing.assert_array_equal(frame.pc.copy().pc.data, [10.0, 20.0]) diff --git a/tests/test_pointcloud/test_pointcloud.py b/tests/test_pointcloud/test_pointcloud.py index 5142391e7..665185cbc 100644 --- a/tests/test_pointcloud/test_pointcloud.py +++ b/tests/test_pointcloud/test_pointcloud.py @@ -106,6 +106,20 @@ def test_init_from_file__lazy(self) -> None: assert np.array_equal(pc.data, self.gdf1["b1"].values) assert pc.is_loaded + def test_has_z__unloaded_3d_file(self) -> None: + """Checks that _has_z detects 3D file metadata without loading point geometries.""" + + # Write a point file with 3D geometry + with tempfile.TemporaryDirectory() as temp_dir: + filename = os.path.join(temp_dir, "points_3d.gpkg") + self.gdf3.to_file(filename) + + # Read the geometry type from file metadata without loading the points + point_cloud = PointCloud(filename) + + assert point_cloud._has_z + assert not point_cloud.is_loaded + def test_init_las(self) -> None: # Import optional laspy or skip test pytest.importorskip("laspy") diff --git a/tests/test_pointcloud/test_transformation.py b/tests/test_pointcloud/test_transformation.py new file mode 100644 index 000000000..3f977ce45 --- /dev/null +++ b/tests/test_pointcloud/test_transformation.py @@ -0,0 +1,420 @@ +"""Tests for point cloud reprojection across eager, Dask and multiprocessing backends.""" + +from __future__ import annotations + +from pathlib import Path + +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +from geopandas.testing import assert_geodataframe_equal +from pyproj import CRS + +import geoutils as gu +from geoutils._misc import import_optional +from geoutils.multiproc import MultiprocConfig +from geoutils.multiproc.cluster import MpCluster + + +@pytest.mark.filterwarnings("ignore:Overriding 3D points with with data column 'intensity':UserWarning") +class TestReprojectChunked: + """ + Checks reproject() for point clouds. + + - Eager, Dask and Multiproc outputs have the same coordinates and attributes, with file results not loaded. + - LAS, LAZ and GeoPackage outputs preserve the values their formats can represent. + - Empty inputs, in-place calls and invalid output choices are checked separately. + """ + + # Give heights values distinct from active intensity values so writing the wrong quantity as LAS Z is visible + positions = np.arange(11) + heights = 20 + positions / 8 + points = gpd.GeoDataFrame( + { + "intensity": (100 + positions).astype(np.int32), + "quality": positions / 16, + "row_id": positions.astype(np.int32), + }, + geometry=gpd.points_from_xy(500000 + 3 * positions, 5100000 + 2 * positions, heights), + crs=32633, + ) + + @pytest.mark.parametrize("chunks", [4, 6]) + @pytest.mark.parametrize("loaded", [False, True]) + def test_reproject__chunked_backends_equal(self, chunks: int, loaded: bool, tmp_path: Path) -> None: + """ + Checks that every backend returns the same coordinates and attributes. + + File inputs and the Multiproc output are not loaded. + """ + + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + + # 1/ Prepare independent eager and file sources with an incomplete final row chunk + # Eleven points split unevenly for both chunk sizes, exposing lost or duplicated rows at chunk edges + filename = tmp_path / "points.gpkg" + self.points.to_file(filename, index=False) + source = gu.PointCloud(self.points.copy(), data_column="intensity") + accessor = self.points.copy() + accessor.pc.set_data_column("intensity") + lazy = gu.open_pointcloud(str(filename), data_column="intensity", chunks=chunks) + multiproc = gu.PointCloud(filename, data_column="intensity") + if loaded: + multiproc.load() + assert multiproc.is_loaded == loaded + assert not lazy.pc.is_loaded + + # 2/ Reproject each interface to the same neighboring UTM zone + # GeoPandas provides an independent coordinate reference and preserves the original geometry heights + target_crs = CRS.from_epsg(32632) + expected = self.points.to_crs(target_crs) + eager_result = source.reproject(crs=target_crs) + accessor_result = accessor.pc.reproject(crs=target_crs) + lazy_result = lazy.pc.reproject(crs=target_crs) + outfile = tmp_path / "projected.gpkg" + with MpCluster({"nb_workers": 2}) as cluster: + configuration = MultiprocConfig(chunks=chunks, outfile=str(outfile), cluster=cluster) + multiproc_result = multiproc.reproject(crs=target_crs, mp_config=configuration) + + # 3/ Check metadata and loading before reading the partitioned results + assert isinstance(eager_result, gu.PointCloud) and eager_result.is_loaded + assert isinstance(accessor_result, gpd.GeoDataFrame) + assert isinstance(lazy_result, dgpd.GeoDataFrame) and not lazy_result.pc.is_loaded + assert isinstance(multiproc_result, gu.PointCloud) and not multiproc_result.is_loaded + assert multiproc_result.crs == target_crs + assert multiproc_result.point_count == len(self.points) + assert multiproc_result.data_column == "intensity" + np.testing.assert_allclose(multiproc_result.bounds, expected.total_bounds, rtol=0, atol=1e-9) + assert not multiproc_result.is_loaded + assert multiproc.is_loaded == loaded + + # 4/ Compare complete rows, including order, attributes and geometry Z after output files are closed + # GeoPackage can normalize integer widths, so compare attribute values without requiring identical dtypes + computed_lazy = lazy_result.compute() + for frame in (eager_result.ds, accessor_result, computed_lazy, multiproc_result.ds): + assert_geodataframe_equal(frame, expected, check_dtype=False) + np.testing.assert_array_equal(frame.geometry.z, self.heights) + np.testing.assert_array_equal(frame["intensity"], self.points["intensity"]) + + # Reading an output must not load its original source or replace the Dask graph with eager values + assert multiproc.is_loaded == loaded + assert not lazy.pc.is_loaded and not lazy_result.pc.is_loaded + assert source.crs == multiproc.crs == lazy.pc.crs == self.points.crs + + @pytest.mark.parametrize("source_suffix,output_suffix", [(".gpkg", ".las"), (".las", ".laz")]) + def test_reproject__las_coordinates_and_attributes( + self, source_suffix: str, output_suffix: str, tmp_path: Path + ) -> None: + """Checks that LAS and LAZ output contain the same elevations and attributes for any active value column.""" + + laspy = import_optional("laspy") + if output_suffix == ".laz": + import_optional("lazrs") + + # Write either 3D vector geometry or native LAS elevations independently of the reprojection implementation + filename = tmp_path / ("source" + source_suffix) + if source_suffix == ".gpkg": + self.points.to_file(filename, index=False) + else: + header = laspy.LasHeader(point_format=6, version="1.4") + header.scales = np.array([0.001, 0.001, 0.001]) + header.offsets = np.array([500000.0, 5100000.0, 0.0]) + header.add_crs(self.points.crs) + header.add_extra_dim(laspy.ExtraBytesParams(name="quality", type=np.float64)) + header.add_extra_dim(laspy.ExtraBytesParams(name="row_id", type=np.int32)) + records = laspy.LasData(header) + records.x = self.points.geometry.x.to_numpy() + records.y = self.points.geometry.y.to_numpy() + records.z = self.heights + records.intensity = self.points["intensity"].to_numpy() + records.quality, records.row_id = self.points["quality"].to_numpy(), self.positions + records.write(filename) + + # Project uneven row chunks while retaining the active intensity column and unloaded source + source = gu.PointCloud(filename, data_column="intensity") + expected = self.points.to_crs(32632) + outfile = tmp_path / ("projected" + output_suffix) + with MpCluster({"nb_workers": 2}) as cluster: + configuration = MultiprocConfig(chunks=4, outfile=str(outfile), cluster=cluster) + result = source.reproject(crs=32632, mp_config=configuration) + assert not source.is_loaded and not result.is_loaded + assert result.data_column == "intensity" + assert result.crs == expected.crs and result.point_count == len(expected) + + # Read the written LAS records independently and allow only the file's coordinate quantization error + records = laspy.read(outfile) + tolerance = records.header.scales / 2 + 1e-8 + np.testing.assert_allclose(records.x, expected.geometry.x, rtol=0, atol=tolerance[0]) + np.testing.assert_allclose(records.y, expected.geometry.y, rtol=0, atol=tolerance[1]) + np.testing.assert_allclose(records.z, self.heights, rtol=0, atol=tolerance[2]) + np.testing.assert_array_equal(records.intensity, self.points["intensity"]) + np.testing.assert_array_equal(records.quality, self.points["quality"]) + np.testing.assert_array_equal(records.row_id, self.positions) + assert records.header.parse_crs() == expected.crs + + # Loading active values uses intensity rather than the independent elevations stored in native LAS Z + np.testing.assert_array_equal(result.data, self.points["intensity"]) + assert not source.is_loaded + + @pytest.mark.parametrize("target_crs", [32633, 4326]) + def test_reproject__reference_and_default_output_format(self, target_crs: int, tmp_path: Path) -> None: + """Checks that reference reprojection writes an unloaded GeoPackage even when the target CRS is unchanged.""" + + # Use a separate reference object and file-backed source so neither needs a full read to choose its CRS + filename = tmp_path / "source.gpkg" + self.points.to_file(filename, index=False) + source = gu.PointCloud(filename, data_column="intensity") + expected = self.points.to_crs(target_crs) + reference = gu.Vector(expected) + outfile = tmp_path / "projected" + + # Infer GeoPackage for an output path without a suffix and inspect its unloaded result metadata + configuration = MultiprocConfig(chunks=4, outfile=str(outfile)) + result = source.reproject(ref=reference, mp_config=configuration) + assert outfile.exists() + assert result is not source + assert not result.is_loaded and not source.is_loaded + assert result.crs == reference.crs + + # Reuse the extensionless file as an unloaded source for another projection before inspecting its data + second_configuration = MultiprocConfig(chunks=4, outfile=str(tmp_path / "reprojected.gpkg")) + second_result = result.reproject(crs=self.points.crs, mp_config=second_configuration) + assert not result.is_loaded and not second_result.is_loaded + assert_geodataframe_equal(second_result.ds, expected.to_crs(self.points.crs), check_dtype=False) + assert not result.is_loaded + + # Both written outputs agree with the equivalent GeoPandas transformations, leaving the source unloaded + assert_geodataframe_equal(result.ds, expected, check_dtype=False) + assert not source.is_loaded + + @pytest.mark.parametrize("loaded", [False, True]) + def test_reproject__empty_point_cloud(self, loaded: bool, tmp_path: Path) -> None: + """Checks that empty inputs write point files with the same attributes and CRS without loading them.""" + + # Write a typed empty point layer so file metadata identifies its geometry without any point records + frame = self.points.iloc[:0].copy() + filename = tmp_path / "empty.gpkg" + frame.to_file(filename, geometry_type="Point", index=False) + source = gu.PointCloud(frame, data_column="intensity") if loaded else gu.PointCloud(filename, "intensity") + expected = frame.to_crs(32632) + configuration = MultiprocConfig(chunks=4, outfile=str(tmp_path / "projected.gpkg")) + + # Reprojection must write an empty point schema instead of skipping output or inventing a placeholder row + result = source.reproject(crs=32632, mp_config=configuration) + assert not result.is_loaded and source.is_loaded == loaded + assert result.point_count == 0 + assert result.crs == expected.crs and result.data_column == "intensity" + assert list(result.columns) == list(frame.columns) + assert not result.is_loaded + + # Reading the output returns every empty attribute column while the original source is not loaded + assert_geodataframe_equal(result.ds, expected, check_dtype=False) + assert source.is_loaded == loaded + + def test_reproject__error_inplace_with_chunked_execution(self, tmp_path: Path) -> None: + """Checks that multiprocessing rejects in-place replacement while eager reprojection still supports it.""" + + # Use an unloaded MP source so rejection happens before its data or the output file are touched + filename = tmp_path / "source.gpkg" + self.points.to_file(filename, index=False) + source = gu.PointCloud(filename, data_column="intensity") + outfile = tmp_path / "projected.gpkg" + configuration = MultiprocConfig(chunks=4, outfile=str(outfile)) + with pytest.raises(ValueError, match="inplace|in place"): + source.reproject(crs=32632, inplace=True, mp_config=configuration) + assert not source.is_loaded and not outfile.exists() + assert source.crs == self.points.crs + + # Without multiprocessing, the same public option updates the source and returns None + eager = gu.PointCloud(self.points.copy(), data_column="intensity") + expected = self.points.to_crs(32632) + assert eager.reproject(crs=32632, inplace=True) is None + assert_geodataframe_equal(eager.ds, expected) + + def test_reproject__multiprocessing_accessor_dataframe_output(self, tmp_path: Path) -> None: + """Checks that multiprocessing through a GeoDataFrame accessor returns the same dataframe family.""" + + # Store active values in a named column separate from the source's three-dimensional geometry + source = self.points.copy() + source.pc.set_data_column("intensity") + expected = self.points.to_crs(32632) + configuration = MultiprocConfig(chunks=4, outfile=str(tmp_path / "projected.gpkg")) + + # The accessor reads the completed point file into a GeoDataFrame with the original active column + result = source.pc.reproject(crs=32632, mp_config=configuration) + assert isinstance(result, gpd.GeoDataFrame) + assert result.pc.data_column == "intensity" + assert_geodataframe_equal(result, expected, check_dtype=False) + assert source.pc.crs == self.points.crs + + @pytest.mark.parametrize("data_column", ["intensity", None]) + def test_reproject__las_accessor_attributes_and_active_values( + self, data_column: str | None, tmp_path: Path + ) -> None: + """Checks that LAS accessor output contains every input attribute and selects intensity or native Z.""" + + laspy = import_optional("laspy") + + # Select either an auxiliary attribute or geometry heights without removing the other point values + source = self.points.copy() + source.pc.set_data_column(data_column) + outfile = tmp_path / "projected.las" + configuration = MultiprocConfig(chunks=4, outfile=str(outfile)) + result = source.pc.reproject(crs=32632, mp_config=configuration) + + # LAS stores geometry heights in its native Z column, which becomes active when geometry was selected + assert isinstance(result, gpd.GeoDataFrame) + expected_column = "Z" if data_column is None else data_column + assert result.pc.data_column == expected_column + for column in ("intensity", "quality", "row_id"): + np.testing.assert_array_equal(result[column], self.points[column]) + + # Allow only the written Z scale's rounding error and preserve the caller's original active selection + with laspy.open(outfile) as reader: + tolerance = reader.header.scales[2] / 2 + 1e-8 + expected_values = self.heights if data_column is None else self.points["intensity"] + np.testing.assert_allclose(result["Z"], self.heights, rtol=0, atol=tolerance) + np.testing.assert_allclose(result.pc.data, expected_values, rtol=0, atol=tolerance) + assert source.pc.data_column == data_column + + def test_reproject__error_dask_with_multiprocessing(self, tmp_path: Path) -> None: + """Checks that Dask and multiprocessing cannot be combined or execute partitions before rejection.""" + + import_optional("dask_geopandas", package_name="dask-geopandas") + from dask.callbacks import Callback + + # Build a file-backed Dask source with several row partitions + filename = tmp_path / "source.gpkg" + self.points.to_file(filename, index=False) + source = gu.open_pointcloud(str(filename), data_column="intensity", chunks=4) + outfile = tmp_path / "projected.gpkg" + configuration = MultiprocConfig(chunks=4, outfile=str(outfile)) + + # Reject competing execution backends before computing any of the source partitions + tasks = [] + with Callback(pretask=lambda *args: tasks.append(args[0])): + with pytest.raises(ValueError, match="Dask"): + source.pc.reproject(crs=32632, mp_config=configuration) + assert tasks == [] + assert not source.pc.is_loaded and not outfile.exists() + + @pytest.mark.parametrize( + "invalid_option", ["driver", "suffix", "driver_suffix", "chunks", "extensionless_las", "extensionless_laz"] + ) + def test_reproject__error_invalid_output_options(self, invalid_option: str, tmp_path: Path) -> None: + """Checks that unsupported output formats and raster-shaped chunks fail before reading the source.""" + + # Use a valid point source and vary only the output option that is incompatible with point reprojection + filename = tmp_path / "source.gpkg" + self.points.to_file(filename, index=False) + source = gu.PointCloud(filename, data_column="intensity") + outfile = tmp_path / ("projected.tif" if invalid_option == "suffix" else "projected.gpkg") + if invalid_option.startswith("extensionless"): + outfile = outfile.with_suffix("") + driver = { + "driver": "GTiff", + "driver_suffix": "LAS", + "extensionless_las": "LAS", + "extensionless_laz": "LAZ", + }.get(invalid_option) + configuration = MultiprocConfig( + chunks=(4, 4) if invalid_option == "chunks" else 4, + outfile=str(outfile), + driver=driver, + ) + + # No output should be created and the input metadata must still describe the original unloaded file + with pytest.raises(ValueError): + source.reproject(crs=32632, mp_config=configuration) + assert not outfile.exists() and not source.is_loaded + assert source.crs == self.points.crs + + def test_reproject__error_las_without_elevations(self, tmp_path: Path) -> None: + """Checks that LAS output rejects two-dimensional points whose active values do not define elevations.""" + + import_optional("laspy") + + # Remove geometry heights but leave intensity present; it must not silently become the output LAS Z + frame = gpd.GeoDataFrame( + self.points.drop(columns="geometry"), + geometry=gpd.points_from_xy(self.points.geometry.x, self.points.geometry.y), + crs=self.points.crs, + ) + source = gu.PointCloud(frame, data_column="intensity") + outfile = tmp_path / "projected.las" + configuration = MultiprocConfig(chunks=4, outfile=str(outfile)) + + # Require a geometry height or native LAS Z column before a point file can be written + with pytest.raises(ValueError, match="elevation|Z|3D"): + source.reproject(crs=32632, mp_config=configuration) + assert not outfile.exists() + np.testing.assert_array_equal(source.data, self.points["intensity"]) + + @pytest.mark.parametrize( + "attribute_kind", + ["nullable_integer", "submillisecond_datetime", "fractional_intensity", "out_of_range_intensity"], + ) + def test_reproject__error_lossy_attribute_storage(self, attribute_kind: str, tmp_path: Path) -> None: + """Checks that unrepresentable attribute values raise before replacing an existing output file.""" + + # Give three points values that the destination format would otherwise round or wrap without an error + frame = self.points.iloc[:3].copy() + if attribute_kind == "nullable_integer": + # Integer columns with nodata can force a float conversion that cannot represent values above 2**53 exactly + column, suffix = "identifier", ".gpkg" + frame[column] = pd.Series([2**53 + 1, pd.NA, 2**53 + 3], dtype="Int64") + elif attribute_kind == "submillisecond_datetime": + # GeoPackage datetime storage cannot represent these nanoseconds below the millisecond boundary + column, suffix = "observed_at", ".gpkg" + frame[column] = pd.date_range("2024-01-01T00:00:00.123456789", periods=3, freq="s") + else: + import_optional("laspy") + column, suffix = "intensity", ".las" + if attribute_kind == "fractional_intensity": + frame[column] = np.array([0.5, 1.5, 2.5]) + else: + frame[column] = np.array([0, 65536, 70000], dtype=np.uint32) + + # Stage one row per worker task to expose conversions that depend on an individual row's nodata value + source = gu.PointCloud(frame, data_column="intensity") + original_values = frame[column].copy() + outfile = tmp_path / ("projected" + suffix) + original_output = b"original" + outfile.write_bytes(original_output) + configuration = MultiprocConfig(chunks=1, outfile=str(outfile)) + + # A failed encoding must identify its attribute and leave both the existing output and source unchanged + with pytest.raises(ValueError, match=column): + source.reproject(crs=32632, mp_config=configuration) + assert outfile.read_bytes() == original_output + pd.testing.assert_series_equal(source.ds[column], original_values) + + @pytest.mark.parametrize("attribute_kind", ["nullable_integer", "millisecond_datetime"]) + def test_reproject__gpkg_representable_attributes(self, attribute_kind: str, tmp_path: Path) -> None: + """Checks that GeoPackage output contains exact small integers and millisecond datetime values.""" + + # Use values representable by the file format so precision checks do not reject valid point attributes + frame = self.points.iloc[:3].copy() + if attribute_kind == "nullable_integer": + column = "identifier" + frame[column] = pd.Series([1, pd.NA, 3], dtype="Int64") + else: + column = "observed_at" + frame[column] = pd.date_range("2024-01-01T00:00:00.123", periods=3, freq="ms") + source = gu.PointCloud(frame, data_column="intensity") + configuration = MultiprocConfig(chunks=1, outfile=str(tmp_path / "projected.gpkg")) + + # Use single-row chunks and place a nodata value in the middle chunk to check schema consistency + result = source.reproject(crs=32632, mp_config=configuration) + assert not result.is_loaded + actual = result.ds[column] + expected = frame[column] + + # Compare value precision independently of the reader's integer-null or datetime dtype representation + if attribute_kind == "nullable_integer": + actual_values = actual.to_numpy(dtype=float, na_value=np.nan) + expected_values = expected.to_numpy(dtype=float, na_value=np.nan) + np.testing.assert_array_equal(actual_values, expected_values) + else: + np.testing.assert_array_equal(actual.to_numpy(dtype="datetime64[ns]"), expected.to_numpy()) diff --git a/tests/test_raster/test_base.py b/tests/test_raster/test_base.py index c2f79b0f8..7aabee65a 100644 --- a/tests/test_raster/test_base.py +++ b/tests/test_raster/test_base.py @@ -14,7 +14,14 @@ from pandas.testing import assert_frame_equal from pyproj import CRS -from geoutils import PointCloud, Raster, Vector, examples, open_raster +from geoutils import ( + PointCloud, + Raster, + Variogram, + Vector, + examples, + open_raster, +) from geoutils.raster import MultiprocConfig from geoutils.raster.base import RasterBase from geoutils.raster.xr_accessor import RasterAccessor @@ -69,6 +76,21 @@ def assert_output_equal(output1: Any, output2: Any, use_allclose: bool = False, df1 = pd.DataFrame(index=[0], data=output1) df2 = pd.DataFrame(index=[0], data=output2) assert_frame_equal(df1, df2, check_dtype=False) + + # For tabular statistics + elif isinstance(output1, pd.DataFrame): + assert_frame_equal(output1, output2) + + # For lightweight variogram records + elif isinstance(output1, Variogram): + assert isinstance(output2, Variogram) + assert np.allclose(output1.lags, output2.lags) + assert np.allclose(output1.semivariance, output2.semivariance, equal_nan=True) + assert np.array_equal(output1.counts, output2.counts) + assert output1.model == output2.model + # For labelled pair samples + elif isinstance(output1, xr.Dataset): + assert output1.identical(output2) # For any other object type else: assert output1 == output2 @@ -245,9 +267,14 @@ def test_properties__equality_and_loading(self, path_index: int, prop: str, lazy ("to_pointcloud", {"subsample": 1, "random_state": 42}), ("polygonize", {"target_values": "all"}), ("subsample", {"subsample": 1000, "random_state": 42}), + ("cosample", {"other": "self", "subsample": 1_000, "random_state": 42}), + ("pairsample", {"n_pairs": 1_000, "random_state": 42}), + ("variogram", {"n_pairs": 1_000, "n_lags": 6, "random_state": 42}), ("filter", {"method": "median", "size": 7}), ("sieve", {"size": 7}), ("fill_nodata", {"max_search_distance": 3}), + ("stats", {}), + ("stats", {"by": {"group": 1}, "bins": {"group": 2}, "statistics": "mean"}), ("get_stats", {}), # 2.2. In-place methods ("load", {}), @@ -272,6 +299,8 @@ def test_methods__equality_and_loading( # Open both objects ds = open_raster(path_raster) raster = Raster(path_raster) + if method == "variogram": + pytest.importorskip("skgstat") # Sieve follows GDAL and accepts integer categories rather than continuous values if method == "sieve": @@ -314,11 +343,15 @@ def test_methods__equality_and_loading( args.update({"bbox": bbox}) elif method in ["raster_equal", "raster_allclose", "georeferenced_grid_equal", "intersection"]: args.update({"other": ds.copy(deep=False)}) + elif method == "cosample": + args.update({"other": raster}) elif method == "copy" and "new_array" in args: args.update({"new_array": np.ones(ds.shape)}) # Apply method for each class output_raster = getattr(raster, method)(**args) + if method == "cosample": + args.update({"other": ds.copy(deep=False)}) output_ds = getattr(ds.rst, method)(**args) # Determine if operation was in-place or not @@ -417,6 +450,7 @@ def test_methods__test_coverage(self) -> None: chunked_methods_and_args = ( ("reproject", {"crs": CRS.from_epsg(4326)}), ("interp_points", {"points": "random", "as_array": True}), + ("cosample", {"other": "self", "subsample": 100, "strategy": "topk", "random_state": 42}), ( "subsample", {"subsample": 100, "strategy": "topk"}, @@ -469,11 +503,16 @@ def test_chunked_methods__equality_loading_laziness( interp_y = raster.bounds.bottom + (rng.choice(raster.shape[1], ninterp) + rng.random(ninterp)) * res[1] kwargs.update({"points": (interp_x, interp_y)}) - # Apply method for each - output_raster = getattr(raster, method)(**kwargs, mp_config=mp_config) - output_ds = getattr(ds.rst, method)(**kwargs) - output_raster2 = getattr(raster2, method)(**kwargs) - output_ds2 = getattr(ds2.rst, method)(**kwargs) + # Apply the same method through each backend, pairing cosample with its own input representation + outputs = [] + for source, backend in ((raster, mp_config), (ds.rst, None), (raster2, None), (ds2.rst, None)): + options = kwargs.copy() + if method == "cosample": + options["other"] = source + if backend is not None: + options["mp_config"] = backend + outputs.append(getattr(source, method)(**options)) + output_raster, output_ds, output_raster2, output_ds2 = outputs # For a raster-type output (reprojection, rasterize, create_mask, proximity, etc...) if isinstance(output_raster, Raster): diff --git a/tests/test_raster/test_raster.py b/tests/test_raster/test_raster.py index 08b87bc2d..203313832 100644 --- a/tests/test_raster/test_raster.py +++ b/tests/test_raster/test_raster.py @@ -383,6 +383,40 @@ def test_to_rio_dataset(self, example: str) -> None: assert np.array_equal(rst.data.data, rio_ds.read().squeeze()) assert np.array_equal(rst.data.mask, rio_ds.read(masked=True).mask.squeeze()) + @pytest.mark.parametrize("method", ["to_rio_dataset", "to_xarray"]) + @pytest.mark.parametrize("loaded", [False, True]) + @pytest.mark.parametrize("area_or_point", ["Area", "Point"]) + def test_to_rio_dataset_to_xarray__loading_metadata( + self, tmp_path: pathlib.Path, method: str, loaded: bool, area_or_point: str + ) -> None: + """Checks that exports load the source and maintain exact values, custom tags and pixel interpretation.""" + + # Write test file with missing pixel and custom metadata to test conversion + values = np.arange(35, dtype=np.float32).reshape(5, 7) + values[2, 3] = np.nan + reference = gu.Raster.from_array( + values, + rio.transform.from_origin(500000, 8600000, 20, 20), + 32633, + nodata=-9999, + area_or_point=area_or_point, + tags={"survey": "synthetic"}, + ) + path = tmp_path / "conversion.tif" + reference.to_file(path) + source = gu.Raster(path, load_data=loaded) + assert source.is_loaded is loaded + + # Both to_rio_dataset and to_xarray use an in-memory Rasterio dataset and therefore load raster values + result = getattr(source, method)() + assert source.is_loaded + converted = gu.Raster(result) if method == "to_rio_dataset" else result.rst.to_geoutils() + + # Every original tag and georeferenced value must survive (even though Xarray adds encoding attributes) + assert source.raster_equal(converted, strict_masked=False, warn_failure_reason=True) + for name, value in source.tags.items(): + assert converted.tags[name] == value + @pytest.mark.parametrize("example", [landsat_b4_path, aster_dem_path, landsat_rgb_path]) def test_to_xarray(self, example: str) -> None: """Test the export to a xarray dataset""" diff --git a/tests/test_raster/test_transformations_raster.py b/tests/test_raster/test_transformations_raster.py index 6849db06b..e37f7dbb7 100644 --- a/tests/test_raster/test_transformations_raster.py +++ b/tests/test_raster/test_transformations_raster.py @@ -22,7 +22,7 @@ from geoutils.raster.transformation import ( _resampling_method_from_str, ) -from geoutils.stats.sampling import _subsample_numpy +from geoutils.sampling.subsampling import _subsample_numpy DO_PLOT = False @@ -729,6 +729,38 @@ def test_crop(self, mask: gu.Raster) -> None: mask_orig_pix.icrop(bbox2_pixel, inplace=True) assert mask_orig.raster_equal(mask_orig_pix) + @pytest.mark.parametrize("method", ["crop", "icrop"]) + def test_crop__unloaded_mask_boolean_values(self, tmp_path: Any, method: str) -> None: + """ + Checks that crop() and icrop() return exact boolean values and nodata cells without loading the source mask. + """ + + # Write integer mask values with one nodata cell inside the window being cropped + values = (np.arange(30).reshape(5, 6) % 2).astype("uint8") + values[2, 2] = 255 + masked_values = np.ma.masked_equal(values, 255) + transform = rio.transform.from_origin(100, 200, 10, 10) + path = tmp_path / "mask.tif" + gu.Raster.from_array(masked_values, transform, 32633, nodata=255).to_file(path) + unloaded = gu.Raster(path, is_mask=True) + loaded = gu.Raster(path, is_mask=True, load_data=True) + + # Select rows 1:4 and columns 1:5 through both public coordinate conventions + if method == "crop": + output = unloaded.crop((110, 160, 150, 190)) + expected = loaded.crop((110, 160, 150, 190)) + else: + output = unloaded.icrop((1, 1, 5, 4)) + expected = loaded.icrop((1, 1, 5, 4)) + + # Match the full-load path and verify logical values and nodata cells independently + assert not unloaded.is_loaded + assert output.is_mask + assert output.data.dtype == bool + assert output.raster_equal(expected, strict_masked=True) + np.testing.assert_array_equal(output.data.data, values[1:4, 1:5].astype(bool)) + np.testing.assert_array_equal(np.ma.getmaskarray(output.data), values[1:4, 1:5] == 255) + @pytest.mark.parametrize("mask", [mask_landsat_b4, mask_aster_dem, mask_everest]) def test_reproject(self, mask: gu.Raster) -> None: # Test 1: with a classic resampling (bilinear) @@ -799,6 +831,43 @@ class TestReprojectChunked: pytest.importorskip("dask") import dask.array as da + @pytest.mark.parametrize("load_source", [False, True]) + def test_reproject__multiprocessing_logical_mask(self, tmp_path: Any, load_source: bool) -> None: + """ + Checks that multiprocessing reprojection returns exact mask values and nodata cells on a shifted grid. + """ + + # Write both boolean states and one nodata cell, then shift the target to leave an uncovered column + rows, columns = np.indices((6, 7)) + values = ((rows + columns) % 2).astype("uint8") + values[2, 3] = 255 + transform = rio.transform.from_origin(0, 6, 1, 1) + source_path = tmp_path / "mask_source.tif" + data = np.ma.masked_equal(values, 255) + gu.Raster.from_array(data, transform, 32633, nodata=255).to_file(source_path) + reference = gu.Raster.from_array(np.zeros((6, 7)), rio.transform.from_origin(1, 6, 1, 1), 32633) + + # Compare a complete in-memory reprojection with windows read through the multiprocessing backend + loaded = gu.Raster(source_path, is_mask=True, load_data=True) + expected = loaded.reproject(ref=reference, resampling="nearest") + source = gu.Raster(source_path, is_mask=True, load_data=load_source) + config = MultiprocConfig(chunks=(3, 4), outfile=str(tmp_path / "mask_reprojected.tif")) + output = source.reproject(ref=reference, resampling="nearest", mp_config=config) + + # Check that the input loading state is unchanged and restore boolean interpretation before reading the output + assert source.is_loaded == load_source + assert not output.is_loaded + assert output.is_mask + assert output.data.dtype == bool + assert output.transform == reference.transform + + # Check the internal hole and uncovered column independently, then compare the known true/false values + expected_missing = np.zeros((6, 7), dtype=bool) + expected_missing[2, 2] = True + expected_missing[:, -1] = True + np.testing.assert_array_equal(np.ma.getmaskarray(output.data), expected_missing) + np.testing.assert_array_equal(output.data.compressed(), expected.data.compressed()) + def test_reproject__small_chunked_grid_matches_base(self, tmp_path: Any) -> None: """Regression test for small-grid Dask/Multiprocessing block placement during reprojection.""" diff --git a/tests/test_raster/test_xr_accessor.py b/tests/test_raster/test_xr_accessor.py index 957e1c68b..d816e81c3 100644 --- a/tests/test_raster/test_xr_accessor.py +++ b/tests/test_raster/test_xr_accessor.py @@ -32,6 +32,25 @@ class TestAccessor: def test_open_raster(self) -> None: pass + @pytest.mark.parametrize("shape", [(1, 3), (3, 1), (1, 1)]) + @pytest.mark.parametrize("bands", [1, 2]) + def test_open_raster__single_row_or_column(self, tmp_path: Path, shape: tuple[int, int], bands: int) -> None: + """Checks that opening a raster returns spatial dimensions of length one and every requested band.""" + + # A single row or column is still a two-dimensional grid, including for several bands + values = np.arange(bands * np.prod(shape), dtype=np.float32).reshape((bands, *shape)) + path = tmp_path / "narrow.tif" + transform = from_origin(0, 3, 1, 1) + gu.Raster.from_array(values, transform, 32631).to_file(path) + + # Opening removes only a single band dimension, while retaining the grid coordinates and values + result = open_raster(str(path)) + expected = values[0] if bands == 1 else values + assert result.dims == (("y", "x") if bands == 1 else ("band", "y", "x")) + assert result.rst.shape == shape + assert result.rst.transform == transform + np.testing.assert_array_equal(result.data, expected) + @pytest.mark.parametrize("path_raster", [landsat_b4_path, aster_dem_path]) def test_copy(self, path_raster: str) -> None: @@ -43,6 +62,36 @@ def test_copy(self, path_raster: str) -> None: assert ds.rst.crs == ds_copy.rst.crs assert ds.rst.nodata == ds_copy.rst.nodata + @pytest.mark.parametrize("lazy", [False, True]) + def test_to_geoutils__loading_laziness(self, tmp_path: Path, lazy: bool) -> None: + """Checks that native conversion loads exact values while keeping a Dask source lazy.""" + + # Write a test file with a missing pixel and Point metadata + values = np.arange(35, dtype=np.float32).reshape(5, 7) + values[2, 3] = np.nan + reference = gu.Raster.from_array( + values, from_origin(500000, 8600000, 20, 20), 32633, nodata=-9999, area_or_point="Point" + ) + path = tmp_path / "conversion.tif" + reference.to_file(path) + if lazy: + pytest.importorskip("dask.array") + source = open_raster(str(path), chunks={"y": 3, "x": 4} if lazy else None) + graph = source.data if lazy else None + assert not source._in_memory + + # Convert to a loaded Raster while keeping the caller's Dask array lazy + result = source.rst.to_geoutils() + assert isinstance(result, gu.Raster) and result.is_loaded + assert source._in_memory is not lazy + if lazy: + assert source.data is graph + + # Check exact values, missing pixels and the complete spatial reference + assert reference.raster_equal(result, strict_masked=False, warn_failure_reason=True) + if lazy: + assert source.data is graph and not source._in_memory + @pytest.mark.parametrize("path_raster", [landsat_b4_path, aster_dem_path]) def test_open__loaded(self, path_raster: str) -> None: """ diff --git a/tests/test_sampling/test_cosampling.py b/tests/test_sampling/test_cosampling.py new file mode 100644 index 000000000..b18169454 --- /dev/null +++ b/tests/test_sampling/test_cosampling.py @@ -0,0 +1,1539 @@ +"""Tests for cosampling at the same support on raster grids and point geometries.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext +from pathlib import Path +from typing import Any + +import geopandas as gpd +import numpy as np +import pytest +import xarray as xr +from geopandas.testing import assert_geodataframe_equal +from rasterio.transform import from_origin +from shapely.geometry import box + +import geoutils as gu +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.multiproc import MultiprocConfig + + +def _raster(data: NDArrayNum, *, x_origin: float = 0) -> gu.Raster: + """Simplify creating a raster for tests below, to only need a one line call.""" + + return gu.Raster.from_array(data, transform=from_origin(x_origin, data.shape[-2], 1, 1), crs=32633, nodata=-99999) + + +class TestCosample: + """ + Checks cosample() calls that combine raster and point data, with eager inputs. + + See TestCosampleChunked further below for Dask/Multiprocessing tests. + + Here, we test that: + - Point values can be adequately placed on a raster grid common support (with gridding). + - Raster values can be adequately placed on a point locations common support (with interpolation). + - Inputs with different CRS or nodata definition (Xarray with NaNs versus masked arrays) behaves properly during + cosampling. + """ + + @pytest.mark.parametrize("caller", ["raster", "points"]) + @pytest.mark.parametrize("accessor", [False, True]) + @pytest.mark.parametrize("explicit_at", [False, True]) + def test_cosample__grid_points(self, caller: str, accessor: bool, explicit_at: bool) -> None: + """Checks that gridding points to a common raster support behaves as expected.""" + + # Place two points around each grid cell so their circular mean is known + expected = np.arange(20, dtype=float).reshape(4, 5) + raster = _raster(expected + 100) + rows, columns = np.indices(raster.shape) + x, y = raster.ij2xy(rows.ravel(), columns.ravel()) + point_values = np.repeat(expected.ravel(), 2) + np.tile([-2, 2], expected.size) + points = gu.PointCloud.from_xyz( + np.repeat(x, 2) + np.tile([-0.2, 0.2], expected.size), + np.repeat(y, 2), + point_values, + crs=raster.crs, + ) + keep = np.ones(raster.shape, dtype=bool) + keep[1, 2] = False + + # Select raster output through either the conversion mode or an explicit grid + first, second = (raster, points) if caller == "raster" else (points, raster) + source = first + if accessor: + source = first.to_xarray().rst if caller == "raster" else first.ds.pc + second = points.ds if caller == "raster" else raster.to_xarray() + target_grid = raster.to_xarray() if accessor else raster + target = {"at": target_grid} if explicit_at else {"raster_point_mode": "grid_points"} + result = source.cosample( + second, + **target, + grid_method="mean", + grid_kwargs={"dist_nodata_pixel": 0.4, "min_points": 2}, + auxiliary={"offset": point_values + 10}, + auxiliary_at="other" if caller == "raster" else "self", + mask=keep, + ) + + # Check the means, added point values, and common mask against arrays + assert isinstance(result, xr.DataArray if accessor else gu.Raster) + output = result.rst if accessor else result + data = result.values if accessor else result.data.filled(np.nan) + expected_bands = [expected + 100, expected] if caller == "raster" else [expected, expected + 100] + expected_bands.append(expected + 10) + assert output.transform == raster.transform + np.testing.assert_allclose(data, np.where(keep, np.stack(expected_bands), np.nan)) + np.testing.assert_array_equal(points.data, point_values) + + @pytest.mark.parametrize("caller", ["raster", "points"]) + @pytest.mark.parametrize("method", ["nearest", "linear"]) + def test_cosample__resample_raster(self, caller: str, method: str) -> None: + """Checks that raster resampling properly interpolates at irregular point coordinates.""" + + # Create a ramp raster whose bilinear values can be calculated exactly + rows, columns = np.indices((7, 9)) + raster = _raster((10 * rows + 2 * columns).astype(float)) + target_rows = np.array([1.2, 2.3, 4.1]) + target_columns = np.array([1.1, 3.2, 5.4]) + x, y = raster.ij2xy(target_rows, target_columns) + points = gu.PointCloud.from_xyz(x, y, np.array([5.0, 6.0, 7.0]), crs=raster.crs) + + # Read the raster at point locations through either object + first, second = (raster, points) if caller == "raster" else (points, raster) + result = first.cosample(second, raster_point_mode="resample_raster", resample_method=method) + + # Check linear values on the slope and nearest values from the closest cells + expected = 10 * target_rows + 2 * target_columns + if method == "nearest": + expected = 10 * np.rint(target_rows) + 2 * np.rint(target_columns) + raster_column = "self" if caller == "raster" else "other" + np.testing.assert_allclose(result.ds[raster_column], expected) + assert result.ds.geometry.equals(points.ds.geometry) + + def test_cosample__single_point_auxiliary_on_raster(self) -> None: + """Checks edge case of a single point auxiliary input, and that it stays an array when gridded onto a raster.""" + + # Place one observation at the only cell center so every output band has one known value + raster = _raster(np.array([[10.0]])) + x, y = raster.ij2xy(np.array([0]), np.array([0])) + points = gu.PointCloud.from_xyz(x, y, np.array([2.0]), crs=raster.crs) + auxiliary = np.array([7.0]) + + # Grid point value and auxiliary onto the explicitly selected raster + result = points.cosample( + raster, auxiliary={"extra": auxiliary}, auxiliary_at="self", at=raster, grid_method="nearest" + ) + + # Check we still have same spatial dimensions and the original point value + expected = np.array([[[2.0]], [[10.0]], [[7.0]]]) + np.testing.assert_array_equal(result.data.filled(np.nan), expected) + np.testing.assert_array_equal(points.data, [2.0]) + + def test_cosample__grid_crs_alignment(self) -> None: + """ + Checks that gridding works with different CRS when 'align' user argument is 'reproject', + otherwise raises an error (for default option 'raises'). + """ + + # Repoject points from raster into a different CRS + raster = _raster(np.arange(20, dtype=float).reshape(4, 5)) + points = raster.to_pointcloud().reproject(crs=4326) + + # Test with/without 'align' argument allowing reprojection + with pytest.raises(ValueError, match="support CRS"): + raster.cosample(points, raster_point_mode="grid_points", grid_method="nearest") + result = raster.cosample(points, raster_point_mode="grid_points", grid_method="nearest", align="reproject") + + # Check exact equality after cosample reprojection back to original CRS + np.testing.assert_allclose(result.data[0], result.data[1]) + + @pytest.mark.parametrize("accessor", [False, True]) + def test_cosample__different_point_locations_share_grid(self, accessor: bool) -> None: + """ + Checks that point clouds with different X/Y coordinates can share a common support after gridding + (functionality mostly useful for dense point cloud inputs). + """ + + # Place one observation near each grid center in two point clouds with deliberately different X coordinates, + # but all within nearest neighbour range of the initial grid center + values = np.arange(20, dtype=float).reshape(4, 5) + raster = _raster(values) + rows, columns = np.indices(raster.shape) + x, y = raster.ij2xy(rows.ravel(), columns.ravel()) + first = gu.PointCloud.from_xyz(x, y, values.ravel(), crs=raster.crs) + second = gu.PointCloud.from_xyz(x + 0.1, y, 2 * values.ravel(), crs=raster.crs) + + # Grid each point set onto the same support and return every input with the caller's container type + source = first.ds.pc if accessor else first + other = second.ds if accessor else second + support = raster.to_xarray() if accessor else raster + result = source.cosample(other, at=support, grid_method="nearest") + + # Check all values are almost equal, as the small coordinate offset makes each observation nearest to its + # original grid cell + output = result.values if accessor else result.data.filled(np.nan) + np.testing.assert_allclose(output, np.stack((values, 2 * values))) + + def test_cosample__resampling_nodata_options(self) -> None: + """Checks nodata options passed through resampling kwargs.""" + + # Put one missing cell beside a point whose other interpolation neighbors equal one + values = np.ones((5, 6), dtype=float) + values[2, 2] = np.nan + raster = _raster(values) + x, y = raster.ij2xy(np.array([1.25, 3.0]), np.array([1.25, 3.0])) + points = gu.PointCloud.from_xyz(x, y, np.array([10.0, 20.0]), crs=raster.crs) + + # Read the same points with strict and lenient missing data settings + ignored = raster.cosample(points, resample_kwargs={"nodata_propagation": "ignore"}) + propagated = raster.cosample(points, resample_kwargs={"nodata_propagation": "propagate"}) + + # Check that only the strict setting drops the point beside the missing cell + assert list(ignored.ds.index) == [0, 1] + assert list(propagated.ds.index) == [1] + np.testing.assert_allclose(ignored.ds["self"], 1) + np.testing.assert_allclose(propagated.ds["self"], 1) + + +class TestRasterCosampleSupport: + """Checks common validity, masks, bands and output shapes when cosample() returns a raster.""" + + @pytest.mark.parametrize("accessor", [False, True]) + @pytest.mark.parametrize("input_type", ["raster", "numpy", "xarray"]) + @pytest.mark.parametrize("auxiliary_at", ["self", "other"]) + def test_cosample__common_validity_and_auxiliary(self, accessor: bool, input_type: str, auxiliary_at: str) -> None: + """Checks that all values and the user mask determine the valid output cells.""" + + # Give each input and the user mask a different cell to exclude + first = np.arange(20, dtype=float).reshape(4, 5) + second, auxiliary = 10 * first, 100 * first + first[0, 0], second[1, 1], auxiliary[2, 2] = np.nan, np.nan, np.nan + mask = np.ones(first.shape, dtype=bool) + mask[3, 3] = False + raster = _raster(first) + + # Run the same public call through a Raster and an Xarray accessor + source = raster.to_xarray().rst if accessor else raster + other: Any = _raster(second) + added: Any = auxiliary + selected_mask: Any = mask + if input_type == "numpy": + other = second + elif input_type == "xarray": + # Plain DataArrays carry dimensions but inherit their coordinates from the geospatial input + other = xr.DataArray(second, dims=("y", "x")) + added = xr.DataArray(auxiliary, dims=("y", "x")) + selected_mask = xr.DataArray(mask, dims=("y", "x")) + elif accessor: + other = other.to_xarray() + result = source.cosample(other, auxiliary={"aux": added}, auxiliary_at=auxiliary_at, mask=selected_mask) + assert isinstance(result, xr.DataArray if accessor else gu.Raster) + output = result.rst if accessor else result + data = result.to_numpy() if accessor else result.data.filled(np.nan) + + # Check the common support grid, validity mask, and documented band order + assert output.shape == raster.shape + assert output.transform == raster.transform + assert output.crs == raster.crs + assert output.tags["long_name"] == ("self", "other", "aux") + expected = mask & np.isfinite(first) & np.isfinite(second) & np.isfinite(auxiliary) + assert np.array_equal(np.isfinite(data), np.broadcast_to(expected, data.shape)) + + # Check every selected value directly, including valid zeros + assert np.array_equal(data[0, expected], first[expected]) + assert np.array_equal(data[1, expected], second[expected]) + assert np.array_equal(data[2, expected], auxiliary[expected]) + + @pytest.mark.parametrize("at", ["self", "other", "explicit"]) + def test_cosample__raster_support_and_alignment(self, at: str) -> None: + """Checks that cosample() uses the selected raster grid and requires permission to reproject.""" + + # Shift the second raster by one cell so the grids overlap but do not match + first = _raster(np.arange(12, dtype=float).reshape(3, 4)) + second = _raster(np.arange(12, dtype=float).reshape(3, 4), x_origin=1) + common_support = first if at == "self" else second + selected_at = second if at == "explicit" else at + + # Require the caller to allow reprojection onto the selected grid + with pytest.raises(ValueError, match="does not share"): + first.cosample(second, at=selected_at) + result = first.cosample(second, at=selected_at, align="reproject") + + # Check the chosen grid and the cells outside the overlap + assert isinstance(result, gu.Raster) + assert result.transform == common_support.transform + assert result.shape == common_support.shape + assert result.count == 2 + assert np.count_nonzero(~np.ma.getmaskarray(result.data[0])) == 9 + + def test_cosample__selected_bands_and_auxiliary_sources(self) -> None: + """Checks that selected bands and auxiliary arrays use their specified input grid.""" + + # Give each requested band distinct values so a wrong band is easy to detect + base = np.arange(20, dtype=float).reshape(4, 5) + first = _raster(np.stack((base, base + 100))) + second = _raster(np.stack((2 * base, 2 * base + 200))) + + # Specify whether each auxiliary array shares the first or second raster grid + result = first.cosample( + second, + band=2, + other_band=2, + auxiliary={"first_aux": base + 1, "second_aux": 2 * base + 2}, + auxiliary_at={"first_aux": "self", "second_aux": "other"}, + ) + + # Check the band order and values through the public Raster API + bands = result.split_bands() + expected = (base + 100, 2 * base + 200, base + 1, 2 * base + 2) + for band, values in zip(bands, expected): + assert np.array_equal(band.data, values) + assert result.tags["long_name"] == ("self", "other", "first_aux", "second_aux") + + @pytest.mark.parametrize("shape", [(1, 5), (5, 1), (1, 1)]) + @pytest.mark.parametrize("accessor", [False, True]) + def test_cosample__singleton_spatial_dimensions(self, shape: tuple[int, int], accessor: bool) -> None: + """Checks that one row or column remains a spatial dimension in a combined multiband raster.""" + + # Select a few cells from a grid with only one row or one column + values = np.arange(np.prod(shape), dtype=float).reshape(shape) + raster = _raster(values) + source = gu.RasterAccessor.from_array(values, raster.transform, raster.crs).rst if accessor else raster + mask = values % 2 == 0 + result = source.cosample(source, mask=mask) + + # Check that both spatial dimensions remain separate from the two bands + data = result.to_numpy() if accessor else result.data.filled(np.nan) + assert data.shape == (2, *shape) + assert np.array_equal(np.isfinite(data[0]), mask) + assert np.array_equal(data[0, mask], values[mask]) + + +class TestPointCosampleSupport: + """Checks mixed inputs, masks, labels and sampling when cosample() returns points.""" + + @pytest.mark.parametrize("caller", ["raster", "pointcloud"]) + @pytest.mark.parametrize("accessor", [False, True]) + def test_cosample__mixed_raster_and_point_inputs(self, caller: str, accessor: bool) -> None: + """Checks that raster and point values are returned together at the point locations.""" + + # Give the raster and point cloud a nodata value at different locations + values = np.arange(30, dtype=float).reshape(5, 6) + values[1, 2] = np.nan + raster = _raster(values) + rows, columns = np.array([0, 1, 3, 4]), np.array([1, 2, 4, 5]) + x, y = raster.ij2xy(rows, columns) + points = gu.PointCloud.from_xyz(x, y, np.array([4.0, 5.0, 6.0, np.nan]), crs=raster.crs) + points.ds.index = ["a", "b", "c", "d"] + source_column, source_bounds = points.data_column, points.bounds + + # Use the point locations as common support through either spatial object or accessor + source, other = (raster, points) if caller == "raster" else (points, raster) + if accessor: + source = source.to_xarray().rst if caller == "raster" else source.ds.pc + other = points.ds if caller == "raster" else raster.to_xarray() + result = source.cosample(other, resample_method="nearest") + assert isinstance(result, gpd.GeoDataFrame if accessor else gu.PointCloud) + output = result if accessor else result.ds + output_points = result.pc if accessor else result + + # Check the original point labels, geometry, and two named value columns + assert np.array_equal(output.index, ["a", "c"]) + assert output.geometry.equals(points.ds.geometry.iloc[[0, 2]]) + assert output_points.data_column == "self" + raster_values = values[rows[[0, 2]], columns[[0, 2]]] + expected = (raster_values, [4.0, 6.0]) if caller == "raster" else ([4.0, 6.0], raster_values) + assert np.array_equal(output["self"], expected[0]) + assert np.array_equal(output["other"], expected[1]) + + # Check that point count and bounds describe the selected points while the original metadata stays unchanged + assert output_points.point_count == 2 + assert np.array_equal(output_points.bounds, points.ds.iloc[[0, 2]].total_bounds) + assert points.point_count == 4 + assert points.data_column == source_column + assert points.bounds == source_bounds + + @pytest.mark.parametrize("at", [None, "self", "other", "explicit", "raw_self", "raw_other"]) + @pytest.mark.parametrize("auxiliary_type", ["numpy", "xarray"]) + def test_cosample__point_auxiliaries_and_labels(self, at: str | None, auxiliary_type: str) -> None: + """Checks that point labels, geometry and auxiliary columns are unchanged in the result.""" + + # Create three-dimensional points with duplicate row labels + positions = np.arange(8, dtype=float) + first = gu.PointCloud.from_xyz(positions, positions**2, positions, crs=32633, use_z=True) + second = gu.PointCloud.from_xyz(positions, positions**2, 2 * positions, crs=32633, use_z=True) + first.ds.index = ["a", "b", "a", "c", "d", "e", "f", "g"] + second.ds.index = np.arange(8) + 10 + auxiliary = 3 * positions + auxiliary[3] = np.nan + + # Choose either point input as the common support + selected_at = second if at == "explicit" else at + other: Any = second + added: Any = xr.DataArray(auxiliary, dims="point") if auxiliary_type == "xarray" else auxiliary + auxiliary_at = "self" + if at in {"raw_self", "raw_other"}: + # A plain second array uses the first point locations even when auxiliary_at selects the other input + other = second.data + selected_at = "other" if at == "raw_other" else None + auxiliary_at = "other" if at == "raw_other" else "self" + result = first.cosample(other, auxiliary={"weight": added}, auxiliary_at=auxiliary_at, at=selected_at) + support = second if at in {"other", "explicit"} else first + expected = np.array([0, 1, 2, 4, 5, 6, 7]) + + # Check the chosen labels, 3D geometry, values, and column order + assert result.ds.geometry.equals(support.ds.geometry.iloc[expected]) + assert list(result.ds.columns) == ["self", "other", "weight", "geometry"] + assert np.array_equal(result.ds["other"], 2 * result.ds["self"]) + assert np.array_equal(result.ds["weight"], 3 * result.ds["self"]) + + @pytest.mark.parametrize("accessor", [False, True]) + @pytest.mark.parametrize("common_support", ["raster", "points"]) + def test_cosample__auxiliary_point_column(self, accessor: bool, common_support: str) -> None: + """Checks that a named point column can supply auxiliary values for raster or point output.""" + + # Give each point an active data column and a separate auxiliary column at a known raster cell + values = np.arange(20, dtype=float).reshape(4, 5) + raster = _raster(values) + rows, columns = np.indices(raster.shape) + x, y = raster.ij2xy(rows.ravel(), columns.ravel()) + points = gu.PointCloud.from_xyz(x, y, values.ravel(), crs=raster.crs) + weights = 3 * values.ravel() + 100 + points.ds["weight"] = weights + original_column = points.data_column + + # Select the auxiliary column without changing the active values of either primary input + source = points.ds.pc if accessor else points + other = points.ds if accessor else points + support = raster.to_xarray() if accessor else raster + selected_at = support if common_support == "raster" else "self" + result = source.cosample(other, auxiliary={"chosen": (other, "weight")}, at=selected_at, grid_method="nearest") + + # Check the selected column against its own values and preserve the original active column + if common_support == "raster": + output = result.values if accessor else result.data.filled(np.nan) + assert np.array_equal(output, np.stack((values, values, weights.reshape(values.shape)))) + else: + output = result if accessor else result.ds + assert np.array_equal(output["chosen"], weights) + assert np.array_equal(output["self"], values.ravel()) + assert output.geometry.equals(points.ds.geometry) + assert points.data_column == original_column + assert np.array_equal(points.data, values.ravel()) + + def test_cosample__plain_auxiliaries_at_reprojected_points(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Checks that plain auxiliary arrays follow their point cloud through reprojection.""" + + # Choose the output coordinates by projecting the original points once, avoiding roundtrip rounding differences + values = np.arange(5, dtype=float) + points = gu.PointCloud.from_xyz(10 + values / 100, 45 + values / 100, values, crs=4326) + support = points.reproject(crs=32632) + auxiliaries = {"scaled": 3 * values, "offset": values + 100} + + # Count coordinate comparisons to check that arrays tied to the same point cloud share one comparison + comparisons: list[tuple[Any, Any]] = [] + original_comparison = gu.PointCloud.georeferenced_coords_equal + + def record_comparison(first: gu.PointCloud, second: Any) -> bool: + """Count coordinate comparisons while preserving their ordinary result.""" + + comparisons.append((first, second)) + return original_comparison(first, second) + + monkeypatch.setattr(gu.PointCloud, "georeferenced_coords_equal", record_comparison) + result = points.cosample(2 * values, auxiliary=auxiliaries, auxiliary_at="self", at=support, align="reproject") + + # Check all original values at the projected coordinates after comparing their source point cloud once + assert len(comparisons) == 1 + expected = np.column_stack((values, 2 * values, 3 * values, values + 100)) + assert np.array_equal(result.ds[["self", "other", "scaled", "offset"]], expected) + assert result.ds.geometry.equals(support.ds.geometry) + assert points.crs.to_epsg() == 4326 + + @pytest.mark.parametrize("mask_type,mask_mode", [("vector", "inside"), ("vector", "outside"), ("raster", "inside")]) + def test_cosample__mask_matches_statistics(self, mask_type: str, mask_mode: str) -> None: + """Checks that cosample() and stats() apply spatial masks to the same points.""" + + # 1/ Prepare points and equivalent spatial masks + # Place two points in each half of a raster, with distinct values and no nodata values + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + rows, columns = np.array([0, 1, 3, 4]), np.array([1, 2, 4, 5]) + x, y = raster.ij2xy(rows, columns) + mask: gu.Vector | gu.Raster + if mask_type == "vector": + # Put the fifth point exactly on the polygon's right edge, where create_mask() must exclude it + x, y = np.append(x, 3.0), np.append(y, 2.0) + # Extend the polygon above the raster so the first point lies inside rather than on its boundary + mask = gu.Vector(gpd.GeoDataFrame({"zone": ["west"]}, geometry=[box(0, 0, 3, 6)], crs=raster.crs)) + else: + keep = np.indices(raster.shape)[1] < 3 + mask = gu.Raster.from_array(keep, raster.transform, raster.crs) + # A point beyond the grid must receive a nodata mask value and remain excluded + x, y = np.append(x, 100.0), np.append(y, 100.0) + points = gu.PointCloud.from_xyz(x, y, np.arange(1, len(x) + 1, dtype=float) * 10, crs=raster.crs) + + # 2/ Apply the same mask through each public operation + # Compare point cosampling, one summary, and a single group with its complete group mask + sampled = points.cosample(points, mask=mask, mask_mode=mask_mode) + summary = points.stats(["mean", "validinliercount"], mask=mask, mask_mode=mask_mode) + grouped, masks = points.stats( + "mean", + values={"height": None}, + by={"zone": np.zeros(len(x), dtype=int)}, + categories={"zone": [0]}, + mask=mask, + mask_mode=mask_mode, + return_masks=True, + ) + + # 3/ Check the selected point positions and values + # Follow create_mask(): polygon boundaries belong to the outside selection + if isinstance(mask, gu.Vector): + inside = np.asarray(mask.create_mask(ref=points, as_array=True), dtype=bool) + assert np.array_equal(inside, [True, True, False, False, False]) + expected = np.flatnonzero(inside if mask_mode == "inside" else ~inside) + else: + expected = np.array([0, 1]) + expected_mean = np.asarray(points.data)[expected].mean() + assert np.array_equal(sampled.ds.index, expected) + assert np.array_equal(np.flatnonzero(masks[0].data), expected) + assert summary == pytest.approx({"mean": expected_mean, "validinliercount": len(expected)}) + np.testing.assert_allclose(grouped["height"], [[len(expected), expected_mean]]) + + # Categorical vector zones use intersections and therefore include the boundary point + if isinstance(mask, gu.Vector) and mask_mode == "inside": + zonal = points.stats("mean", values={"height": None}, by={"zone": (mask, "zone")}) + np.testing.assert_allclose(zonal["height"], [[3, (10 + 20 + 50) / 3]]) + + def test_cosample__masked_auxiliary_and_mask(self) -> None: + """Checks that missing added point values and mask values remove their matching rows.""" + + # Exclude different points through a missing added value, a missing mask value, and a false mask value + positions = np.arange(5, dtype=float) + points = gu.PointCloud.from_xyz(positions, positions, positions, crs=32633) + auxiliary = np.ma.array(np.arange(5), mask=[False, True, False, False, False]) + mask = np.ma.array([True, True, True, False, True], mask=[False, False, True, False, False]) + result = points.cosample(points, auxiliary={"aux": auxiliary}, auxiliary_at="self", mask=mask) + + # Check that only the first and last points remain in their original order + assert np.array_equal(result.ds.index, [0, 4]) + assert np.array_equal(result.ds["aux"], [0, 4]) + + @pytest.mark.parametrize("singleton_band", [False, True]) + def test_cosample__raster_auxiliary_interpolation_options(self, singleton_band: bool) -> None: + """Checks that a plain raster auxiliary uses the requested interpolation options.""" + + # Place points around a nodata auxiliary cell while leaving both primary inputs fully valid + raster = _raster(np.arange(99, dtype=float).reshape(9, 11)) + auxiliary = raster.data.filled(np.nan) + 1000 + auxiliary[4, 5] = np.nan + # The primary raster's nodata sentinel is a valid number in this independent auxiliary array + auxiliary[1, 1] = -99999 + positions = np.arange(40) + x, y = raster.ij2xy(1 + positions % 7, 1 + positions % 9) + points = gu.PointCloud.from_xyz(x, y, positions.astype(float), crs=raster.crs) + + # Use Raster.interp_points() to identify which points remain valid around the nodata auxiliary cell + auxiliary_raster = gu.Raster.from_array( + np.ma.masked_invalid(auxiliary), raster.transform, raster.crs, nodata=None + ) + expected = auxiliary_raster.interp_points((x, y), method="nearest", as_array=True, dist_nodata_spread=2) + kept = np.flatnonzero(np.isfinite(expected)) + raw_auxiliary = auxiliary[np.newaxis, ...] if singleton_band else auxiliary + result = points.cosample( + raster, + auxiliary={"offset": raw_auxiliary}, + auxiliary_at="other", + resample_method="nearest", + resample_kwargs={"dist_nodata_spread": 2}, + ) + + # Check the exact finite positions returned by interpolation for two-dimensional and one-band inputs + assert 0 < kept.size < positions.size + assert -99999 in expected[kept] + assert np.array_equal(result.ds.index, kept) + assert np.array_equal(result.ds["offset"], expected[kept]) + + @pytest.mark.parametrize("accessor", [False, True]) + def test_cosample__selected_band_validity(self, accessor: bool) -> None: + """Checks that only nodata cells in the requested raster band remove points.""" + + # Put nodata cells at different point locations in the two raster bands + data = np.arange(30, dtype=float).reshape(5, 6) + bands = np.stack((data, data + 100)) + bands[0, 1, 1], bands[1, 3, 3] = np.nan, np.nan + raster = _raster(bands) + x, y = raster.ij2xy(np.array([1, 2, 3]), np.array([1, 2, 3])) + points = gu.PointCloud.from_xyz(x, y, np.array([1.0, 2.0, 3.0]), crs=raster.crs) + + # Request only the second band through a Raster and an Xarray accessor + source = raster.to_xarray().rst if accessor else raster + other = points.ds if accessor else points + result = source.cosample(other, band=2, resample_method="nearest") + output = result if accessor else result.ds + + # Check that a point over nodata in only the unused first band remains + assert np.array_equal(output.index, [0, 1]) + assert np.array_equal(output["self"], [107, 114]) + assert np.array_equal(output["other"], [1, 2]) + + def test_cosample__subsampling_and_independent_result(self) -> None: + """Checks that subsampling returns points in source order and data independent of its inputs.""" + + # Use duplicate row labels so sampling must follow row positions + positions = np.arange(20, dtype=float) + points = gu.PointCloud.from_xyz(positions, positions, positions, crs=32633) + points.ds.index = np.tile(["z", "a"], 10) + original = points.ds.copy(deep=True) + result = points.cosample(points, subsample=5, random_state=42) + repeated = points.cosample(points, subsample=5, random_state=42) + + # Check repeatable selection in the original point order + assert len(result.ds) == 5 + assert np.all(np.diff(result.ds["self"]) > 0) + assert_geodataframe_equal(result.ds, repeated.ds) + + # Check that changing the returned PointCloud does not change either input + result.ds["self"] = -1 + assert_geodataframe_equal(points.ds, original) + + +class TestCosampleChunked: + """Checks cosample() loading behavior and exact results with Dask and Multiproc inputs.""" + + @pytest.mark.parametrize("auxiliary_type", ["numpy", "dask", "column"]) + def test_cosample__dask_point_auxiliaries_on_raster(self, auxiliary_type: str) -> None: + """Checks that lazy point gridding assigns array or column auxiliaries to the correct raster cells.""" + + import_optional("dask") + import dask.array as da + + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import _register_dask_pointcloud_accessor + + # Give every cell one point and a distinct auxiliary value, using duplicate labels to expose index alignment + _register_dask_pointcloud_accessor() + values = np.arange(12, dtype=float).reshape(3, 4) + raster = _raster(values + 100) + rows, columns = np.indices(raster.shape) + x, y = raster.ij2xy(rows.ravel(), columns.ravel()) + points = gu.PointCloud.from_xyz(x, y, values.ravel(), crs=raster.crs) + auxiliary_values = 3 * values.ravel() + 7 + points.ds["weight"] = auxiliary_values + points.ds.index = np.tile(["a", "b"], 6) + + # Partition points independently from Dask arrays, or select the existing point column + lazy_points = dgpd.from_geopandas(points.ds, npartitions=3, sort=False) + lazy_points.pc.data_column = points.data_column + target = raster.to_xarray().chunk({"y": 2, "x": 3}) + auxiliary: Any = auxiliary_values + if auxiliary_type == "dask": + auxiliary = da.from_array(auxiliary_values, chunks=5) + elif auxiliary_type == "column": + auxiliary = (lazy_points, "weight") + result = lazy_points.pc.cosample( + target, auxiliary={"extra": auxiliary}, auxiliary_at="self", at=target, grid_method="nearest" + ) + + # Calculate the same result from the eager point cloud and raster + eager_auxiliary: Any = (points, "weight") if auxiliary_type == "column" else auxiliary_values + expected = points.cosample( + raster, + auxiliary={"extra": eager_auxiliary}, + auxiliary_at="self", + at=raster, + grid_method="nearest", + ) + + # Check that neither input nor the result is computed until their values are requested + assert target.chunks is not None and result.chunks is not None + assert not lazy_points.pc.is_loaded + assert lazy_points.pc.data_column == points.data_column + assert np.array_equal(result.compute().values, expected.data.filled(np.nan), equal_nan=True) + assert not lazy_points.pc.is_loaded + + @pytest.mark.parametrize("chunks", [(7, 11), (32, 47), (256, 256)]) + @pytest.mark.parametrize("caller", ["raster", "points"]) + def test_cosample__dask_gridding_chunks(self, chunks: tuple[int, int], caller: str, tmp_path: Path) -> None: + """Checks that Dask point gridding returns the same seeded sample as eager gridding for each chunk size.""" + + # Place one point at each grid cell so the nearest-neighbor gridding is exact + pytest.importorskip("dask.array") + pytest.importorskip("dask_geopandas") + values = np.arange(65 * 97, dtype=float).reshape(65, 97) + raster = _raster(values) + points = raster.to_pointcloud() + lazy_raster = raster.to_xarray().chunk({"y": chunks[0], "x": chunks[1]}) + + # Write the points and reopen them as lazy partitions + point_file = tmp_path / "observations.gpkg" + points.to_file(point_file) + lazy_points = gu.open_pointcloud(str(point_file), data_column=points.data_column, chunks=1400) + + # Grid lazy point data with two partition sizes and both public calling objects + source, other = (lazy_raster.rst, lazy_points) if caller == "raster" else (lazy_points.pc, lazy_raster) + result = source.cosample( + other, + raster_point_mode="grid_points", + grid_method="nearest", + grid_kwargs={"chunksizes": chunks}, + subsample=200, + random_state=42, + strategy="topk", + ) + expected = raster.cosample(raster, subsample=200, random_state=42, strategy="topk") + + # Check that all runs stay lazy and select the same cells and values as the eager call + assert result.data.chunks is not None + assert not lazy_points.pc.is_loaded + assert np.array_equal(result.compute().values, expected.data.filled(np.nan), equal_nan=True) + assert not lazy_points.pc.is_loaded + + @pytest.mark.parametrize("reproject", [False, True]) + def test_cosample__multiproc_point_columns_loading(self, tmp_path: Path, reproject: bool) -> None: + """Checks that aligning and gridding point columns doesn't load their file.""" + + # Store one point at every grid center so both selected columns have exact expected raster values + values = np.arange(30, dtype=float).reshape(5, 6) + raster = _raster(values + 100) + points = _raster(values + 10).to_pointcloud() + points.ds["weight"] = 3 * values.ravel() + 7 + if reproject: + points = points.reproject(crs=4326) + expected = raster.cosample( + points, + auxiliary={"weight": (points, "weight")}, + at=raster, + grid_method="nearest", + align="reproject", + ) + point_file = tmp_path / "observations.gpkg" + points.to_file(point_file) + unloaded = gu.PointCloud(point_file, data_column=points.data_column) + assert not unloaded.is_loaded + + # Read both point columns into raster tiles without copying or loading the point source in the parent + outfile = tmp_path / "cosampled.tif" + result = raster.cosample( + unloaded, + auxiliary={"weight": (unloaded, "weight")}, + at=raster, + grid_method="nearest", + align="reproject", + mp_config=MultiprocConfig(chunks=(2, 3), outfile=str(outfile)), + ) + + # Check that the original active column is unchanged, loading behaviour and that the result matches eager + assert not unloaded.is_loaded + assert unloaded.data_column == points.data_column + assert not result.is_loaded + assert outfile.exists() + assert np.array_equal(result.data.filled(np.nan), expected.data.filled(np.nan), equal_nan=True) + assert not unloaded.is_loaded + + def test_cosample__multiproc_preserves_point_locations(self, tmp_path: Path) -> None: + """Checks that Multiproc alignment returns exact point locations for values and plain auxiliaries.""" + + # Project irregular geographic coordinates so rounding them through LAS would change their exact locations + positions = np.arange(11, dtype=float) + points = gu.PointCloud.from_xyz(10 + positions / 31, 45 + positions / 47, positions + 1, crs=4326) + support = points.reproject(crs=32632) + point_file = tmp_path / "geographic.gpkg" + points.to_file(point_file) + unloaded = gu.PointCloud(point_file, data_column=points.data_column) + + # Pass rectangular chunks and a raster output path to check point output leaves both options unchanged + config = MultiprocConfig(chunks=(2, 3), outfile=str(tmp_path / "unused-output.tif")) + + # Compare eager and Multiproc alignment of unloaded point values and two arrays at reprojected point locations + expected = points.cosample( + 2 * positions, + auxiliary={"offset": positions + 10}, + auxiliary_at="self", + at=support, + align="reproject", + ) + result = unloaded.cosample( + 2 * positions, + auxiliary={"offset": positions + 10}, + auxiliary_at="self", + at=support, + align="reproject", + mp_config=config, + ) + + # Point matching requires identical ordered X/Y values, no source load or config change is needed + assert result.georeferenced_coords_equal(support) + assert np.array_equal(result.ds[["self", "other", "offset"]], expected.ds[["self", "other", "offset"]]) + assert not unloaded.is_loaded + assert config.chunks == (2, 3) + assert config.driver is None + assert not Path(config.outfile).exists() + + def test_cosample__multiproc_distinct_tempfiles(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """ + Checks that point gridding and shifted raster inputs use separate temporary files and preserve eager values. + """ + + # 1/ Prepare distinguishable values on a target grid and two shifted auxiliary grids + # Different offsets expose an intermediate temporary raster files being overwritten by a later spatial operation + rows, columns = np.indices((12, 15)) + values = (10 * rows + columns).astype(float) + raster = _raster(values) + points = _raster(values + 1000).to_pointcloud() + auxiliaries = {"east": _raster(values + 2000, x_origin=1), "west": _raster(values + 3000, x_origin=-1)} + raster_mask = gu.Raster.from_array(columns % 4 != 0, auxiliaries["west"].transform, raster.crs) + options: dict[str, Any] = {"raster_point_mode": "grid_points", "grid_method": "nearest", "align": "reproject"} + expected = raster.cosample(points, auxiliary=auxiliaries, mask=raster_mask, **options) + + # Reopen all raster inputs unloaded so every spatial stage must read its own source file + unloaded = {} + for name, source in {"reference": raster, **auxiliaries, "mask": raster_mask}.items(): + filename = tmp_path / f"{name}.tif" + source.to_file(filename) + unloaded[name] = gu.Raster(filename, load_data=False, is_mask=name == "mask") + + # 2/ Record temporary destinations while leaving their normal creation and cleanup unchanged + intermediate_paths: list[Path] = [] + original_temporary = MultiprocConfig.temporary + + @contextmanager + def record_temporary(config: MultiprocConfig) -> Iterator[MultiprocConfig]: + """ + Record each intermediate filename without changing its normal context lifetime. + """ + with original_temporary(config) as temporary: + intermediate_paths.append(Path(temporary.outfile)) + yield temporary + + monkeypatch.setattr(MultiprocConfig, "temporary", record_temporary) + outfile = tmp_path / "cosampled.tif" + result = unloaded["reference"].cosample( + points, + auxiliary={"east": unloaded["east"], "west": unloaded["west"]}, + mask=unloaded["mask"], + mp_config=MultiprocConfig(chunks=(5, 6), outfile=str(outfile)), + **options, + ) + + # 3/ Check that the final file is still readable after all distinct intermediate files have been removed + assert len(intermediate_paths) >= 4 + assert len(set(intermediate_paths)) == len(intermediate_paths) + assert outfile not in intermediate_paths + assert all(not path.exists() for path in intermediate_paths) + assert outfile.exists() + assert not result.is_loaded + assert all(not source.is_loaded for source in unloaded.values()) + assert result.transform == raster.transform + assert result.tags["long_name"] == ("self", "other", "east", "west") + np.testing.assert_allclose(result.data.filled(np.nan), expected.data.filled(np.nan)) + + @pytest.mark.parametrize("lazy", [False, True]) + @pytest.mark.parametrize("raster_mask", [False, True]) + def test_cosample__masked_integers_and_boolean_masks(self, lazy: bool, raster_mask: bool) -> None: + """Checks that nodata and boolean masks do not convert valid integers.""" + + # Place source nodata, masked auxiliary data, a masked user-mask cell, and a false mask cell separately + data = np.ma.array(np.arange(30, dtype=np.int32).reshape(5, 6), mask=False) + data.mask[0, 0] = True + auxiliary = np.ma.array(2 * data.data, mask=False) + auxiliary.mask[1, 1] = True + mask = np.ma.array(np.ones(data.shape, dtype=bool), mask=False) + mask.mask[2, 2], mask[3, 3] = True, False + first = _raster(data) + selected_mask = first.from_array(mask, first.transform, first.crs) if raster_mask else mask + expected_output = first.cosample( + first, auxiliary={"aux": auxiliary}, auxiliary_at="self", mask=selected_mask + ).data.filled(np.nan) + + # Cosample eager or Dask rasters with the user mask supplied as an array or raster + source = first + if lazy: + da = pytest.importorskip("dask.array") + source = gu.RasterAccessor.from_array( + da.from_array(data.astype(float).filled(np.nan), chunks=(2, 3)), first.transform, first.crs + ).rst + if raster_mask: + # Convert the raster mask to Xarray with its masked cell excluded and its values still boolean + selected_mask = gu.RasterAccessor.from_array( + selected_mask.data.filled(False), selected_mask.transform, selected_mask.crs + ) + result = source.cosample(source, auxiliary={"aux": auxiliary}, auxiliary_at="self", mask=selected_mask) + output = result.to_numpy() if lazy else result.data.filled(np.nan) + + # Compare every output value with eager and check the common validity mask directly + if lazy: + assert result.chunks is not None + assert np.array_equal(output, expected_output, equal_nan=True) + expected = ~data.mask & ~auxiliary.mask & mask.filled(False) + assert np.array_equal(np.isfinite(output[0]), expected) + assert np.array_equal(output[0, expected], data.data[expected]) + assert np.array_equal(output[2, expected], auxiliary.data[expected]) + + @pytest.mark.parametrize("subsample", [1, 12, 0.25]) + def test_cosample__dask_raster_chunks(self, subsample: int | float) -> None: + """Checks that Dask output stays lazy and topk does not depend on chunk sizes.""" + + # Create the two Dask rasters with different chunk layouts + da = pytest.importorskip("dask.array") + array = np.arange(63, dtype=float).reshape(7, 9) + transform = from_origin(0, 7, 1, 1) + first = gu.RasterAccessor.from_array(da.from_array(array, chunks=(2, 4)), transform, 32633) + second = gu.RasterAccessor.from_array(da.from_array(2 * array, chunks=(4, 3)), transform, 32633) + + # Repeat the same seeded sample after changing one chunk layout + result = first.rst.cosample(second, subsample=subsample, random_state=42, strategy="topk") + changed = first.chunk({"y": 4, "x": 3}).rst.cosample( + second, subsample=subsample, random_state=42, strategy="topk" + ) + assert isinstance(result, xr.DataArray) + assert isinstance(result.data, da.Array) + assert isinstance(first.data, da.Array) + + # Check that the result remains Dask-backed and matches eager selected cells and values + eager = _raster(array).cosample(_raster(2 * array), subsample=subsample, random_state=42, strategy="topk") + output = result.compute().to_numpy() + assert np.array_equal(output, changed.to_numpy(), equal_nan=True) + assert np.array_equal(output, eager.data.filled(np.nan), equal_nan=True) + count = array.size if subsample == 1 else int(subsample * array.size) if subsample < 1 else subsample + assert np.count_nonzero(np.isfinite(output[0])) == count + np.testing.assert_allclose(output[1], 2 * output[0], equal_nan=True) + + @pytest.mark.parametrize( + "mask_type,chunks,workers", + [("array", (4, 5), False), ("raster", (5, 4), False), ("vector", (4, 5), True)], + ) + def test_cosample__multiproc_raster_matches_eager_and_dask( + self, mask_type: str, chunks: tuple[int, int], workers: bool, tmp_path: Path + ) -> None: + """ + Checks that Multiproc does not load raster inputs and matches eager and Dask results for each mask. + """ + + import_optional("dask") + from geoutils.multiproc.cluster import MpCluster + + # 1/ Create selected bands, auxiliary values, and masks with nodata at different locations + # Add unused first bands with large offsets so the result reveals if a worker reads the wrong band + values = np.arange(180, dtype=float).reshape(12, 15) + first = _raster(np.stack((values + 1000, values))) + second = _raster(np.stack((values + 2000, 2 * values))) + auxiliary = _raster(np.stack((values + 3000, 3 * values))) + first.data[0, 1, 1] = np.ma.masked + first.data[1, 2, 3] = np.ma.masked + second.data[1, 4, 5] = np.ma.masked + auxiliary.data[1, 6, 7] = np.ma.masked + raw_auxiliary = 4 * values + raw_auxiliary[8, 9] = np.nan + + # Use a regular boolean pattern for array/raster masks, and a polygon for the vector case + keep = values % 7 != 0 + raster_mask = gu.Raster.from_array(keep, first.transform, first.crs) + mask: Any = keep + if mask_type == "raster": + mask = raster_mask + elif mask_type == "vector": + mask = gu.Vector(gpd.GeoDataFrame(geometry=[box(0, 0, 10, 12)], crs=first.crs)) + options: dict[str, Any] = { + "band": 2, + "other_band": 2, + "auxiliary_at": {"raw": "self"}, + "subsample": 17, + "random_state": 42, + "strategy": "topk", + } + + # 2/ Calculate eager and differently chunked Dask references through the public method + eager = first.cosample(second, auxiliary={"scaled": (auxiliary, 2), "raw": raw_auxiliary}, mask=mask, **options) + expected = eager.data.filled(np.nan) + for rows, columns in ((4, 5), (5, 4)): + lazy_first = first.to_xarray().chunk({"y": rows, "x": columns}) + lazy_second = second.to_xarray().chunk({"y": columns, "x": rows}) + lazy_auxiliary = auxiliary.to_xarray().chunk({"y": rows, "x": columns}) + lazy_mask = mask + if mask_type == "raster": + lazy_mask = gu.RasterAccessor.from_array(mask.data, mask.transform, mask.crs) + lazy = lazy_first.rst.cosample( + lazy_second, + auxiliary={"scaled": (lazy_auxiliary, 2), "raw": raw_auxiliary}, + mask=lazy_mask, + **options, + ) + assert lazy.data.chunks is not None + assert np.array_equal(lazy.compute().values, expected, equal_nan=True) + + # Write three raster files to disk with the selected bands and nodata patterns used by the eager reference + unloaded = [] + for name, raster in (("first", first), ("second", second), ("auxiliary", auxiliary)): + filename = tmp_path / f"{name}.tif" + raster.to_file(filename) + unloaded.append(gu.Raster(filename, load_data=False)) + mp_mask = mask + if mask_type == "raster": + # Write the boolean raster mask to disk so Multiproc reads it by tile + mask_file = tmp_path / "mask.tif" + raster_mask.to_file(mask_file) + mp_mask = gu.Raster(mask_file, is_mask=True, load_data=False) + + # 3/ Write the cosampled raster to disk by tile, using worker processes for the vector-mask case + outfile = tmp_path / "cosampled.tif" + with MpCluster({"nb_workers": 2}) if workers else nullcontext(None) as cluster: + configuration = MultiprocConfig(chunks=chunks, outfile=str(outfile), cluster=cluster) + result = unloaded[0].cosample( + unloaded[1], + auxiliary={"scaled": (unloaded[2], 2), "raw": raw_auxiliary}, + mask=mp_mask, + mp_config=configuration, + **options, + ) + assert outfile.exists() + assert not result.is_loaded + assert all(not source.is_loaded for source in unloaded) + # Check that processing Multiproc tiles does not change the caller's original boolean mask + assert np.array_equal(keep, values % 7 != 0) + if mask_type == "raster": + assert not mp_mask.is_loaded + + # Check complete output bands and the exact common sample after temporary worker files have been cleaned up + assert np.array_equal(result.data.filled(np.nan), expected, equal_nan=True) + assert np.count_nonzero(np.isfinite(expected[0])) == 17 + assert result.georeferenced_grid_equal(first) + assert tuple(result.tags["long_name"]) == ("self", "other", "scaled", "raw") + + @pytest.mark.parametrize("lazy", [False, True]) + def test_cosample__no_replacement_after_interpolation(self, lazy: bool) -> None: + """Checks that points rejected during interpolation are not replaced in the sample.""" + + # Place points around one raster nodata cell, including neighbors rejected by slinear's default nodata spread + raster = _raster(np.ones((11, 11), dtype=float)) + raster.data[5, 5] = np.ma.masked + rows, columns = np.meshgrid(np.arange(3, 8), np.arange(3, 8), indexing="ij") + x, y = raster.ij2xy(rows.ravel(), columns.ravel()) + points = gu.PointCloud.from_xyz(x, y, np.arange(x.size, dtype=float), crs=raster.crs).ds + options = {"subsample": 23, "random_state": 42} + + # Select almost every initially valid point so final interpolation removes some selected nodata neighbors + coverage_values = raster.interp_points((x, y), method="slinear", dist_nodata_spread=0, as_array=True) + candidate_validity = np.isfinite(coverage_values) + selected = points.pc.cosample(points, mask=candidate_validity, **options) + interpolated_values = raster.interp_points((x, y), method="slinear", as_array=True) + selected_indices = selected.index.to_numpy() + expected_indices = selected_indices[np.isfinite(interpolated_values[selected_indices])] + expected = points.pc.cosample(raster.to_xarray(), resample_method="slinear", **options) + assert candidate_validity.sum() == 24 + assert 0 < len(expected_indices) < options["subsample"] + + # Apply the same seeded selection with raster values read eagerly or through Dask + source: Any = points + other: Any = raster.to_xarray() + if lazy: + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import ( + _register_dask_pointcloud_accessor, + ) + + _register_dask_pointcloud_accessor() + source = dgpd.from_geopandas(points, npartitions=3, sort=False) + other = raster.to_xarray().chunk({"y": 6, "x": 6}) + result = source.pc.cosample(other, resample_method="slinear", **options) + if lazy: + assert not source.pc.is_loaded + output = result.compute() if lazy else result + + # Return the originally selected finite rows in their original order, with no replacement points + assert_geodataframe_equal(output, expected) + assert np.array_equal(output.index.to_numpy(), expected_indices) + np.testing.assert_allclose(output["other"], interpolated_values[expected_indices]) + assert np.array_equal(output.geometry.to_numpy(), points.geometry.iloc[expected_indices].to_numpy()) + if lazy: + assert not source.pc.is_loaded and not result.pc.is_loaded + + @pytest.mark.parametrize("method", ["nearest", "linear"]) + def test_cosample__multiproc_raster_inputs_not_loaded(self, method: str, tmp_path: Path) -> None: + """Checks that Multiproc point output matches eager while raster inputs stay unloaded.""" + + # 1/ Prepare an affine raster, a boolean raster mask, and irregular point observations + # Fractional pixel locations make nearest and linear interpolation produce distinct, predictable values + rows, columns = np.indices((12, 15)) + values = (10 * rows + 2 * columns).astype(float) + raster = _raster(np.stack((values + 1000, values))) + raster.data[1, 4, 5] = np.ma.masked + raster.data[0, 2, 3] = np.ma.masked + keep = columns % 4 != 0 + raster_mask = gu.Raster.from_array(keep, raster.transform, raster.crs) + positions = np.arange(24) + x, y = raster.ij2xy(1.2 + positions % 8, 1.3 + positions % 11) + points = gu.PointCloud.from_xyz(x, y, positions.astype(float), crs=raster.crs) + points.ds.index = np.repeat(np.arange(12), 2) + auxiliary = 3 * positions.astype(float) + auxiliary[3] = np.nan + + # 2/ Calculate the eager reference from loaded inputs + options: dict[str, Any] = { + "other_band": 2, + "auxiliary": {"scaled": auxiliary}, + "auxiliary_at": "self", + "resample_method": method, + "subsample": 7, + "random_state": 42, + } + expected = points.cosample(raster, mask=raster_mask, **options) + + # Write the value raster and boolean raster mask to disk, then reopen both without loading their values + raster_file, mask_file = tmp_path / "values.tif", tmp_path / "mask.tif" + raster.to_file(raster_file) + raster_mask.to_file(mask_file) + unloaded = gu.Raster(raster_file, load_data=False) + unloaded_mask = gu.Raster(mask_file, load_data=False, is_mask=True) + + # Read the unloaded raster and mask in Multiproc tiles while interpolating values at the selected points + result = points.cosample( + unloaded, + mask=unloaded_mask, + mp_config=MultiprocConfig(chunks=(5, 6), outfile=str(tmp_path / "unused-point-output.tif")), + **options, + ) + + # 3/ Check that both rasters stay unloaded and the output preserves values, row order, and duplicate labels + assert not unloaded.is_loaded + assert not unloaded_mask.is_loaded + assert len(result.ds) == 7 + assert_geodataframe_equal(result.ds, expected.ds) + + @pytest.mark.parametrize("subsample", [1, 8, 0.3]) + @pytest.mark.parametrize("mask_type", ["array", "raster", "vector"]) + def test_cosample__dask_point_partitions(self, subsample: int | float, mask_type: str) -> None: + """Checks that values, labels and 3D geometry match eager across Dask point partitions.""" + + import_optional("dask") + import dask.array as da + + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import _register_dask_pointcloud_accessor + + # 1/ Prepare point columns and independent geometry elevations on known raster cells + # Give geometry Z coordinates and data-column values different numbers so both can be checked independently + _register_dask_pointcloud_accessor() + positions = np.arange(40, dtype=float) + raster = _raster(np.arange(120, dtype=float).reshape(10, 12)) + raster.data[3, 4] = np.ma.masked + x, y = raster.ij2xy(1 + positions.astype(int) % 7, 1 + positions.astype(int) % 9) + geometry = gpd.points_from_xy(x, y, z=positions + 500) + first = gpd.GeoDataFrame({"height": positions.copy()}, geometry=geometry, crs=raster.crs) + second = gpd.GeoDataFrame({"height": 2 * positions}, geometry=geometry, crs=raster.crs) + first.index = np.tile(["a", "b", "a", "c"], 10) + second.index = np.arange(40) + 100 + first.iloc[1, 0], second.iloc[2, 0] = np.nan, np.nan + auxiliary = 3 * positions + auxiliary[3] = np.nan + mask: Any = positions.astype(int) % 5 != 0 + if mask_type == "raster": + keep = np.indices(raster.shape)[1] % 4 != 0 + mask = gu.Raster.from_array(keep, raster.transform, raster.crs) + elif mask_type == "vector": + mask = gu.Vector(gpd.GeoDataFrame(geometry=[box(0, 0, 8, 10)], crs=raster.crs)) + + # 2/ Calculate an eager reference, then vary each Dask input's partitions or chunks + options: dict[str, Any] = { + "auxiliary_at": {"scaled": "self"}, + "resample_method": "nearest", + "subsample": subsample, + "random_state": 42, + } + eager_mask = mask + if mask_type == "raster": + eager_mask = gu.RasterAccessor.from_array(mask.data, mask.transform, mask.crs) + expected = first.pc.cosample( + second, auxiliary={"grid": raster.to_xarray(), "scaled": auxiliary}, mask=eager_mask, **options + ) + for partitions in (2, 5): + lazy_first = dgpd.from_geopandas(first, npartitions=partitions, sort=False) + lazy_second = dgpd.from_geopandas(second, npartitions=partitions + 1, sort=False) + lazy_raster = raster.to_xarray().chunk({"y": partitions + 1, "x": 4}) + lazy_auxiliary = da.from_array(auxiliary, chunks=7) + # Give an array mask independent Dask chunks, or apply the spatial mask to each Dask point partition + lazy_mask = mask + if mask_type == "array": + lazy_mask = da.from_array(mask, chunks=9) + elif mask_type == "raster": + lazy_mask = _raster(keep.astype(float)).to_xarray().astype(bool).chunk({"y": 3, "x": 5}) + result = lazy_first.pc.cosample( + lazy_second, + auxiliary={"grid": lazy_raster, "scaled": lazy_auxiliary}, + mask=lazy_mask, + **options, + ) + + # 3/ Check that the output remains partitioned and its metadata is unchanged before comparing values + assert isinstance(result, dgpd.GeoDataFrame) + assert not result.pc.is_loaded + assert result.pc.data_column == "self" + assert not lazy_first.pc.is_loaded + output = result.compute() + assert list(output.columns) == ["self", "other", "grid", "scaled", "geometry"] + assert np.array_equal(output.index.to_numpy(), expected.index.to_numpy()) + assert np.array_equal(output.geometry.to_numpy(), expected.geometry.to_numpy()) + assert np.array_equal(output.geometry.z.to_numpy(), expected.geometry.z.to_numpy()) + assert np.array_equal(output.drop(columns="geometry").to_numpy(), expected.drop(columns="geometry")) + assert output.crs == expected.crs + + @pytest.mark.parametrize("subsample", [1, 8]) + def test_cosample__dask_points_defer_value_interpolation( + self, subsample: int, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Checks that Dask reads raster validity first and waits to interpolate the selected values.""" + + import_optional("dask") + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + import geoutils.interface.interpolation as interpolation + from geoutils.pointcloud.pd_accessor import _register_dask_pointcloud_accessor + + # 1/ Prepare duplicate point labels and one raster nodata cell to check validity and row order + _register_dask_pointcloud_accessor() + positions = np.arange(24, dtype=float) + raster = _raster(np.arange(120, dtype=float).reshape(10, 12)) + raster.data[3, 4] = np.ma.masked + x, y = raster.ij2xy(1 + positions.astype(int) % 7, 1 + positions.astype(int) % 9) + points = gu.PointCloud.from_xyz(x, y, positions, crs=raster.crs).ds + points.index = np.tile(["b", "a", "b"], 8) + options: dict[str, Any] = {"resample_method": "nearest", "subsample": subsample, "random_state": 42} + expected = points.pc.cosample(raster.to_xarray(), **options) + lazy_points = dgpd.from_geopandas(points, npartitions=3, sort=False) + lazy_raster = raster.to_xarray().chunk({"y": 4, "x": 5}) + + # 2/ Check that building the Dask result reads raster validity without interpolating raster values + calls: list[bool] = [] + original_interpolation = interpolation._interp_points_base + + def track_interpolation(*args: Any, **kwargs: Any) -> Any: + """ + Record whether interpolation is checking raster validity or calculating raster values. + """ + validity_only = kwargs.get("_validity_only", False) + calls.append(validity_only) + return original_interpolation(*args, **kwargs) + + monkeypatch.setattr(interpolation, "_interp_points_base", track_interpolation) + result = lazy_points.pc.cosample(lazy_raster, **options) + assert isinstance(result, dgpd.GeoDataFrame) + assert not result.pc.is_loaded + assert result.pc.data_column == "self" + assert result.pc.bounds is None + assert lazy_points.pc.data_column == points.pc.data_column + assert calls and all(calls) + + # 3/ Compute values only on request, preserving the eager labels, coordinates and selected values + output = result.compute() + assert any(not validity_only for validity_only in calls) + assert np.array_equal(output.index.to_numpy(), expected.index.to_numpy()) + assert np.array_equal(output.geometry.to_numpy(), expected.geometry.to_numpy()) + assert np.array_equal(output.drop(columns="geometry").to_numpy(), expected.drop(columns="geometry")) + + # Asking for a row count may compute the result, but must describe selected rows rather than the source + assert result.pc.point_count == len(expected) + assert lazy_points.pc.point_count == len(points) + + @pytest.mark.parametrize("lazy_input", ["raster", "points"]) + @pytest.mark.parametrize("caller", ["raster", "points"]) + def test_cosample__mixed_eager_and_dask_inputs(self, lazy_input: str, caller: str) -> None: + """Checks that eager and Dask raster and point inputs can be used together.""" + + import_optional("dask") + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import _register_dask_pointcloud_accessor + + # Place observations on known raster cells so interpolation and point values have exact references + _register_dask_pointcloud_accessor() + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + rows, columns = np.array([1, 2, 3]), np.array([1, 2, 3]) + x, y = raster.ij2xy(rows, columns) + values = np.array([10.0, 20.0, 30.0]) + points = gu.PointCloud.from_xyz(x, y, values, crs=raster.crs).ds + + # Use accessor-compatible raster and point inputs while making only one of them Dask-backed + raster_input = raster.to_xarray() + point_input = points + if lazy_input == "raster": + raster_input = raster_input.chunk({"y": 2, "x": 3}) + else: + point_input = dgpd.from_geopandas(points, npartitions=2, sort=False) + source, other = (raster_input.rst, point_input) if caller == "raster" else (point_input.pc, raster_input) + result = source.cosample(other, resample_method="nearest") + + # Calculate the same caller order from eager accessor inputs + eager_raster = raster.to_xarray() + eager_source, eager_other = (eager_raster.rst, points) if caller == "raster" else (points.pc, eager_raster) + expected = eager_source.cosample(eager_other, resample_method="nearest") + + # Check input and output loading before comparing every row with eager + if lazy_input == "raster": + da = pytest.importorskip("dask.array") + assert isinstance(raster_input.data, da.Array) + else: + assert not point_input.pc.is_loaded and not result.pc.is_loaded + output = result.compute() if lazy_input == "points" else result + assert_geodataframe_equal(output, expected) + raster_column, point_column = ("self", "other") if caller == "raster" else ("other", "self") + assert np.array_equal(output[raster_column], raster.data[rows, columns]) + assert np.array_equal(output[point_column], values) + assert output.geometry.equals(points.geometry) + + +class TestCosampleErrors: + """ + Checks cosample() raise proper errors for inputs that cannot produce a clear result. + + We test the following cases: + - Invalid band indexes or columns names raise an error, without loading objects, + - Geospatial inputs with incompatible types (Xarray/Pandas vs GeoUtils objects) and empty common support + (no intersection, different CRS with reprojection option turned off) both raise an error. + - Consistency of common support definition through `at` and `raster_point_mode`. + """ + + def test_cosample__error_basic_inputs(self) -> None: + """Checks that wrong inputs raise clear errors.""" + + # Create two grids and one point set so a conversion direction alone cannot choose one grid + raster = _raster(np.arange(20, dtype=float).reshape(4, 5)) + other = _raster(np.ones((4, 5)), x_origin=1) + points = raster.to_pointcloud() + + # Check that raster_point_mode inputs that are ambiguous or conflict with "at" raise proper errors + with pytest.raises(ValueError, match="unambiguous"): + raster.cosample(other, raster_point_mode="grid_points") + with pytest.raises(ValueError, match="conflicts"): + raster.cosample(points, at=points, raster_point_mode="grid_points") + with pytest.raises(ValueError, match="Argument ``raster_point_mode`` must"): + raster.cosample(points, raster_point_mode="unknown") + + # Check error that reduce points is not yet implemented + with pytest.raises(NotImplementedError, match="revision of Raster.reduce_points"): + raster.cosample(points, raster_point_mode="resample_raster", resample_method="reduce") + + # Check errors for passing similarly named arguments of grid/resample kwargs and cosample() + # (some must be passed directly to cosample(), not kwargs) + with pytest.raises(ValueError, match="outside ``grid_kwargs``"): + raster.cosample(points, grid_kwargs={"resampling": "nearest"}) + with pytest.raises(ValueError, match="outside ``resample_kwargs``"): + raster.cosample(points, resample_kwargs={"points": (np.ones(1), np.ones(1))}) + + @pytest.mark.parametrize("argument", ["self", "other", "auxiliary"]) + @pytest.mark.parametrize("output_support", ["raster", "points"]) + @pytest.mark.parametrize("invalid_band,error", [("height", TypeError), (0, ValueError), (3, ValueError)]) + def test_cosample__error_invalid_bands_do_not_load_rasters( + self, argument: str, output_support: str, invalid_band: Any, error: type[Exception], tmp_path: Path + ) -> None: + """Checks that invalid primary or auxiliary bands raise error from raster metadata before loading.""" + + # Write two band raster, so we can request band 0 and 3 and check for errors below + values = np.arange(20, dtype=float).reshape(4, 5) + raster = _raster(np.stack((values, values + 100))) + path = tmp_path / "two_bands.tif" + raster.to_file(path) + x, y = raster.ij2xy(np.array([1, 2]), np.array([1, 2])) + points = gu.PointCloud.from_xyz(x, y, np.ones(2), crs=raster.crs) + source = gu.Raster(path, load_data=False) + + # Put the invalid selection in one argument while all other bands and output locations remain valid + options: dict[str, Any] = {"at": points if output_support == "points" else "self"} + if argument == "auxiliary": + options["auxiliary"] = {"extra": (source, invalid_band)} + else: + options["band" if argument == "self" else "other_band"] = invalid_band + with pytest.raises(error, match="band|selector"): + source.cosample(source, **options) + + # The error must be raised before reading the input + assert not source.is_loaded + + @pytest.mark.parametrize( + "auxiliary_at", [{"missing": "self"}, {"extra": "unknown"}, {"extra": None}, {"extra": 1}, "unknown"] + ) + def test_cosample__error_invalid_auxiliary_locations(self, auxiliary_at: Any) -> None: + """Checks that unknown auxiliary names and invalid location choices raise errors.""" + + # Use an auxiliary for which the coordinates make auxiliary_at unnecessary + raster = _raster(np.arange(20, dtype=float).reshape(4, 5)) + auxiliary = {"extra": (raster, 1)} + + # Validate every supplied location choice so a typo cannot silently survive input preparation + with pytest.raises(ValueError, match="auxiliary_at"): + raster.cosample(raster, auxiliary=auxiliary, auxiliary_at=auxiliary_at) + + @pytest.mark.parametrize("output_support", ["raster", "points"]) + @pytest.mark.parametrize( + "native_support,shape", + [ + ("raster", (12,)), + ("raster", (4, 3)), + ("raster", (2, 3, 4)), + ("points", (12, 1)), + ("points", (1, 12)), + ("points", (11,)), + ], + ) + def test_cosample__error_raw_auxiliary_shape( + self, output_support: str, native_support: str, shape: tuple[int, ...] + ) -> None: + """Checks that raw arrays match their native grid or point shape before any conversion to output locations.""" + + # Use twelve grid cells and twelve corresponding points (so that equal counts cannot hide the wrong shape) + raster = _raster(np.arange(12, dtype=float).reshape(3, 4)) + points = raster.to_pointcloud() + source = raster if native_support == "raster" else points + support = raster if output_support == "raster" else points + auxiliary = np.ones(shape) + + # Raise error for flattened grids, extra point dimensions, multiple raw bands and arrays with the wrong size + with pytest.raises(ValueError, match="(?i)shape|native|point"): + source.cosample(source, auxiliary={"extra": auxiliary}, auxiliary_at="self", at=support) + + @pytest.mark.parametrize("caller", ["raster", "points"]) + @pytest.mark.parametrize("accessor", [False, True]) + @pytest.mark.parametrize("input_type", ["raster", "points"]) + @pytest.mark.parametrize("argument", ["other", "auxiliary", "selected_auxiliary", "at", "mask"]) + def test_cosample__error_mixed_spatial_container_families( + self, caller: str, accessor: bool, input_type: str, argument: str + ) -> None: + """Checks that inputs all use the same "family" (Xarray/Pandas or GeoUtils objects).""" + + # Give the caller valid point locations that can also be used by each input + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + points = raster.to_pointcloud() + source = raster if caller == "raster" else points + if accessor: + source = raster.to_xarray().rst if caller == "raster" else points.ds.pc + other = points.ds if accessor else points + incompatible: Any = raster if input_type == "raster" else points + + # Use boolean values for masks so only their container family is invalid + if argument == "mask": + if input_type == "raster": + incompatible = gu.Raster.from_array(np.ones(raster.shape, dtype=bool), raster.transform, raster.crs) + else: + incompatible = points.copy(new_array=np.ones(points.point_count, dtype=bool)) + if not accessor: + if input_type == "raster": + incompatible = gu.RasterAccessor.from_array(incompatible.data, incompatible.transform, incompatible.crs) + else: + incompatible = incompatible.ds + + # Put the incompatible input in one argument while all other inputs use the caller's container type + options: dict[str, Any] = {} + if argument == "other": + other = incompatible + elif argument == "auxiliary": + options["auxiliary"] = {"extra": incompatible} + elif argument == "selected_auxiliary": + selector = 1 if input_type == "raster" else None + options["auxiliary"] = {"extra": (incompatible, selector)} + else: + options[argument] = incompatible + with pytest.raises(TypeError, match="mix"): + source.cosample(other, **options) + + @pytest.mark.parametrize("mismatch", ["shifted", "reordered", "shortened"]) + def test_cosample__error_point_locations_before_raster_preparation( + self, mismatch: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Checks that incompatible auxiliary points raise before any raster alignment or interpolation occurs.""" + + # Use point support in a different CRS so processing the raster would require reprojection + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + points = raster.to_pointcloud().reproject(crs=4326) + x, y = points.ds.geometry.x.to_numpy(), points.ds.geometry.y.to_numpy() + values = np.arange(points.point_count, dtype=float) + if mismatch == "shifted": + x = x + 0.1 + elif mismatch == "reordered": + x, y, values = x[::-1], y[::-1], values[::-1] + else: + x, y, values = x[:-1], y[:-1], values[:-1] + auxiliary = gu.PointCloud.from_xyz(x, y, values, crs=points.crs) + + # Fail if expensive spatial preparation starts before all point locations have been checked + def unexpected_raster_operation(*args: Any, **kwargs: Any) -> Any: + """Report raster work that must not occur for incompatible point inputs.""" + + pytest.fail("Point coordinates must be checked before raster preparation.") + + monkeypatch.setattr(gu.Raster, "reproject", unexpected_raster_operation) + monkeypatch.setattr(gu.Raster, "interp_points", unexpected_raster_operation) + with pytest.raises(ValueError, match="ordered support coordinates"): + raster.cosample(points, auxiliary={"extra": auxiliary}, align="reproject") + + @pytest.mark.parametrize("lazy_input", ["self", "other", "auxiliary", "mask", "at"]) + def test_cosample__error_multiproc_with_dask_inputs(self, lazy_input: str, tmp_path: Path) -> None: + """Checks that Dask inputs cannot be combined with Multiproc output.""" + + # Pass a Dask input in each input position while the other arguments are eager + import_optional("dask") + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + eager = raster.to_xarray() + lazy = eager.chunk({"y": 2, "x": 3}) + source = lazy.rst if lazy_input == "self" else eager.rst + other = lazy if lazy_input == "other" else eager + options: dict[str, Any] = {} + if lazy_input == "auxiliary": + options["auxiliary"] = {"extra": lazy} + elif lazy_input == "mask": + options["mask"] = lazy > 0 + elif lazy_input == "at": + options["at"] = lazy + + # Raise an error without writing a Multiproc result (ensures files are properly cleaned at any breakpoints) + outfile = tmp_path / "should-not-exist.tif" + with pytest.raises(ValueError, match="Multiprocessing and Dask"): + source.cosample(other, mp_config=MultiprocConfig(chunks=(2, 3), outfile=str(outfile)), **options) + assert not outfile.exists() + + @pytest.mark.parametrize( + "option,argument", + [ + ("grid_kwargs", "mp_config"), + ("grid_kwargs", "data_column"), + ("resample_kwargs", "mp_config"), + ("resample_kwargs", "_validity_only"), + ], + ) + def test_cosample__error_reserved_backend_and_validity_options(self, option: str, argument: str) -> None: + """Checks that internal options cannot be passed through interpolation keywords.""" + + # Use only valid raster values so the misplaced option is the sole cause of failure + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + options: dict[str, Any] = {option: {argument: None}} + with pytest.raises(ValueError, match=f"outside ``{option}``"): + raster.cosample(raster, **options) + + @pytest.mark.parametrize("name", ["self", "other", "geometry"]) + def test_cosample__error_reserved_auxiliary_names(self, name: str) -> None: + """Checks that auxiliary names cannot replace a primary value or the point geometry column.""" + + # Use an otherwise valid call so only the added column name is invalid + raster = _raster(np.arange(12, dtype=float).reshape(3, 4)) + with pytest.raises(ValueError, match="cannot be"): + raster.cosample(raster, auxiliary={name: raster}) + + def test_cosample__error_conflicting_raster_point_mode(self) -> None: + """Checks that an explicit point resampling mode rejects a raster output target.""" + + # Request point output while explicitly selecting a raster grid + raster = _raster(np.arange(12, dtype=float).reshape(3, 4)) + points = raster.to_pointcloud() + with pytest.raises(ValueError, match="conflicts"): + raster.cosample(points, at="self", raster_point_mode="resample_raster") + + def test_cosample__error_empty_common_support(self) -> None: + """Checks that an empty common sample raises before creating an unusable spatial output.""" + + # Remove every available raster cell or point through the user mask + raster = _raster(np.arange(12, dtype=float).reshape(3, 4)) + points = raster.to_pointcloud() + with pytest.raises(ValueError, match="no finite data common"): + raster.cosample(raster, mask=np.zeros(raster.shape, dtype=bool)) + with pytest.raises(ValueError, match="no finite data common"): + points.cosample(points, mask=np.zeros(len(points.ds), dtype=bool)) diff --git a/tests/test_sampling/test_pairsampling.py b/tests/test_sampling/test_pairsampling.py new file mode 100644 index 000000000..667637246 --- /dev/null +++ b/tests/test_sampling/test_pairsampling.py @@ -0,0 +1,480 @@ +"""Tests for sampling raster cell and point row pairs without loading unnecessary data.""" + +from __future__ import annotations + +from typing import Any + +import geopandas as gpd +import numpy as np +import pytest +import xarray as xr +from rasterio.transform import from_origin +from shapely.geometry import box + +import geoutils as gu +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.sampling.pairsampling import _RegularPairSampler + + +@pytest.fixture +def raster() -> gu.Raster: + """Return a finite raster with a small nodata region.""" + + array = np.arange(900, dtype=float).reshape(30, 30) + array[2:5, 4:8] = np.nan + return gu.Raster.from_array(array, from_origin(0, 30, 2, 3), 32633, nodata=-99999) + + +class TestRasterPairSampling: + """ + Checks pairsample() on raster grids. + + Dask behavior is covered in TestPairSampleChunked further below. + + This module is checking the following: + - All sampling "strategies" return the correct shape of outputs. + - Behaviour of user masks, min/max sampling distance and duplicate pairs is respected. + """ + + @pytest.mark.parametrize("strategy", ["independent", "anchors", "chunk_anchors", "anchor_batched"]) + def test_pairsample__raster_strategies_return_pair_dataset(self, raster: gu.Raster, strategy: str) -> None: + """Checks that every strategy returns the requested pairs and labelled values.""" + + # Draw raster pairs with one strategy and min/max + pairs = raster.pairsample( + n_pairs=200, + min_distance=2, + max_distance=40, + strategy=strategy, + random_state=42, + anchors_per_round=100, + distances_per_anchor=3, + angles_per_distance=3, + ) + + # Check the common Xarray layout and that every pair follows the min/max distance and skips nodata (we + # must have a finite output) + assert isinstance(pairs, xr.Dataset) + assert pairs.sizes == {"pair": 200, "endpoint": 2} + assert set(pairs.data_vars) == {"index", "value", "distance", "row", "column", "x", "y"} + assert np.all(np.isfinite(pairs.value)) + assert np.all((pairs.distance >= 2) & (pairs.distance <= 40)) + assert np.array_equal(pairs.value, raster.data.data.ravel()[pairs["index"]]) + expected_distance = np.hypot(np.diff(pairs.x, axis=1), np.diff(pairs.y, axis=1)).ravel() + np.testing.assert_allclose(pairs.distance, expected_distance) + + def test_pairsample__raster_candidate_batch_limit(self, raster: gu.Raster, monkeypatch: pytest.MonkeyPatch) -> None: + """Checks that pairsample() never samples more pairs at once than batch_pairs allows.""" + + # Isolate the method that samples pairs before invalid pairs are removed + candidate_counts = [] + original_candidates = _RegularPairSampler._candidates + + def candidates(sampler: _RegularPairSampler, count: int) -> tuple[NDArrayNum, NDArrayNum]: + """Record how many pairs were requested, then generate those pairs with the original method.""" + candidate_counts.append(count) + return original_candidates(sampler, count) + + # During this test, send every call to _candidates() through the recording function above + monkeypatch.setattr(_RegularPairSampler, "_candidates", candidates) + + # Request 200 pairs while allowing at most 37 candidates at once, which forces several generator calls + pairs = raster.pairsample(n_pairs=200, batch_pairs=37, strategy="independent", random_state=8) + + # Check that all requested pairs are returned and every candidate batch stays within the limit + assert pairs.sizes["pair"] == 200 + assert len(candidate_counts) > 1 + assert max(candidate_counts) <= 37 + + @pytest.mark.parametrize("option", ["batch_pairs", "max_rounds", "chunks_per_round", "angles_per_distance"]) + def test_pairsample__error_raster_invalid_batch_controls(self, raster: gu.Raster, option: str) -> None: + """Checks that zero sampling controls fail before creating batches or indexing empty anchors.""" + + with pytest.raises(ValueError, match="controls must be"): + raster.pairsample(n_pairs=20, **{option: 0}) + + def test_pairsample__raster_reproducible_and_globally_unique(self, raster: gu.Raster) -> None: + """Checks that a fixed seed returns the same unique raster pairs in the same order.""" + + # Draw the same globally unique sample twice with one seed + first = raster.pairsample(n_pairs=300, deduplicate="global", random_state=4) + second = raster.pairsample(n_pairs=300, deduplicate="global", random_state=4) + indexes = np.sort(first["index"].values, axis=1) + + # Check exact repeatability and treat reversed endpoint order as the same pair + assert first.identical(second) + assert len(np.unique(indexes, axis=0)) == first.sizes["pair"] + + def test_pairsample__raster_random_xy_and_mask(self, raster: gu.Raster) -> None: + """Checks that independent raster endpoints stay inside an aligned boolean mask.""" + + # Allow pairs only in the upper half of the raster + mask = np.zeros(raster.shape, dtype=bool) + mask[:15] = True + pairs = raster.pairsample(n_pairs=100, sampling="random_xy", mask=mask, random_state=8) + + # Check both endpoints and the sampling method recorded in the result + assert np.all(pairs["row"] < 15) + assert pairs.attrs["sampling"] == "random_xy" + + def test_pairsample__raster_geodataframe_mask(self, raster: gu.Raster) -> None: + """Checks that a GeoDataFrame mask places both sampled raster endpoints inside its polygon.""" + + # Mask the raster with one polygon that covers only its upper-left area + mask = gpd.GeoDataFrame(geometry=[box(0, -15, 30, 30)], crs=raster.crs) + pairs = raster.pairsample(n_pairs=80, mask=mask, max_distance=20, random_state=2) + + # Check the map coordinates of both endpoints against the polygon bounds + assert np.all(pairs.x < 30) + assert np.all(pairs.y > -15) + + @pytest.mark.parametrize("raster_mask", [False, True]) + def test_pairsample__raster_masked_integers_and_mask_cells(self, raster_mask: bool) -> None: + """Checks that masked integer values and masked boolean cells never enter raster pairs.""" + + # Exclude different cells through the integer data, the mask's own mask, and a false mask value + data = np.ma.array(np.arange(100, dtype=np.int32).reshape(10, 10), mask=False) + data.mask[0, 0] = True + mask = np.ma.array(np.ones(data.shape, dtype=bool), mask=False) + mask.mask[1, 1] = True + mask[2, 2] = False + raster = gu.Raster.from_array(data, from_origin(0, 10, 1, 1), 32633, nodata=-9999) + selected_mask = raster.from_array(mask, raster.transform, raster.crs) if raster_mask else mask + + # Draw pairs with either the boolean array or its Raster form + pairs = raster.pairsample(n_pairs=200, mask=selected_mask, random_state=3) + + # Check each endpoint against the combined mask and original integer values + eligible = ~data.mask & mask.filled(False) + indexes = pairs["index"].values + assert pairs.sizes["pair"] == 200 + assert np.all(eligible.ravel()[indexes]) + assert np.array_equal(pairs["value"].values, data.data.ravel()[indexes]) + + +class TestPointPairSampling: + """ + Checks pairsample() on point clouds. + + - Search strategies return pairs at the requested distances with their original row indexes. + - Array and spatial masks select the same points. + - Exact searches cover reused anchors and the upper distance limit. + """ + + @pytest.mark.parametrize("strategy", ["kdtree", "hashgrid", "nn_logvector"]) + def test_pairsample__point_strategies(self, strategy: str) -> None: + """Checks that every point strategy returns the requested distances and original rows.""" + + # Create a regular set of points with values that vary in both directions + y, x = np.mgrid[:20, :20] + values = np.sin(x.ravel() / 3) + np.cos(y.ravel() / 4) + points = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), values, crs=32633) + # Draw pairs with one nearby point search strategy + pairs = points.pairsample( + n_pairs=150, + min_distance=1, + max_distance=15, + strategy=strategy, + anchors_per_round=200, + nn_tolerance=0.6, + random_state=3, + ) + + # Check the pair count, distances, and coordinates from the original rows + assert pairs.sizes["pair"] == 150 + assert np.all((pairs.distance >= 1) & (pairs.distance <= 15)) + assert np.array_equal(pairs["x"], x.ravel()[pairs["index"]]) + assert np.array_equal(pairs["y"], y.ravel()[pairs["index"]]) + + @pytest.mark.parametrize("mask_form", ["array", "masked", "xarray", "dask"]) + def test_pairsample__point_random_pairs_and_mask(self, mask_form: str) -> None: + """Checks that independent point pairs return original row numbers after nodata and mask filtering.""" + + # Create points with one missing value and allow only the left half + y, x = np.mgrid[:12, :12] + values = (x + y).astype(float).ravel() + values[5] = np.nan + points = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), values, crs=32633) + eligible = x.ravel() < 6 + mask: Any = eligible.copy() + + # Exclude masked entries and accept other array layouts with one boolean value per point + if mask_form == "masked": + mask = np.ma.array(mask, mask=False) + mask.mask[13] = True + eligible[13] = False + elif mask_form == "xarray": + mask = xr.DataArray(mask.reshape(12, 12), dims=("row", "column")) + elif mask_form == "dask": + mask = import_optional("dask.array").array.from_array(mask, chunks=17) + + # Draw independent endpoints with smaller output number types + pairs = points.pairsample( + n_pairs=100, + sampling="random_xy", + min_distance=1, + max_distance=10, + mask=mask, + random_state=9, + index_dtype=np.int16, + distance_dtype=np.float32, + ) + + # Check output types, removal of the missing row, and the mask boundary + assert pairs["index"].dtype == np.int16 + assert pairs["distance"].dtype == np.float32 + assert not np.any(pairs["index"] == 5) + assert np.all(pairs.x < 6) + assert np.all(eligible[pairs["index"]]) + + @pytest.mark.parametrize("mask_form", ["raster", "pointcloud", "vector", "geodataframe"]) + def test_pairsample__point_spatial_masks(self, mask_form: str) -> None: + """Checks that spatial masks select the same point pairs as their equivalent boolean array.""" + + # Select the left half of a point grid, with raster samples at the same integer coordinates + y, x = np.mgrid[:12, :12] + points = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), (x + y).ravel(), crs=32633) + eligible = x < 6 + mask: Any + if mask_form == "raster": + mask = gu.Raster.from_array(eligible[::-1], from_origin(0, 11, 1, 1), points.crs) + elif mask_form == "pointcloud": + mask = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), eligible.ravel(), crs=points.crs) + else: + # Place the first excluded column on the polygon boundary to check that it stays ineligible + mask = gpd.GeoDataFrame(geometry=[box(-0.5, -0.5, 6, 11.5)], crs=points.crs) + if mask_form == "vector": + mask = gu.Vector(mask) + + # Use identical random draws so both masks must return the same rows, values and distances + options: dict[str, Any] = {"n_pairs": 100, "sampling": "random_xy", "random_state": 9} + expected = points.pairsample(mask=eligible.ravel(), **options) + result = points.pairsample(mask=mask, **options) + assert result.identical(expected) + + @pytest.mark.parametrize("mask_form", ["wrong_count", "numeric", "numeric_pointcloud", "different_coordinates"]) + def test_pairsample__error_point_invalid_masks(self, mask_form: str) -> None: + """Checks that point masks require boolean values at the same number of ordered source locations.""" + + # Use only finite source values so the invalid mask is the sole reason pair sampling fails + y, x = np.mgrid[:5, :5] + points = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), (x + y).ravel(), crs=32633) + mask: Any = np.ones(25, dtype=bool) + error = "Argument ``mask`` must be boolean and contain one value per input location" + if mask_form == "wrong_count": + mask = mask[:-1] + elif mask_form == "numeric": + mask = mask.astype(float) + elif mask_form == "numeric_pointcloud": + mask = points + error = "point support mask must contain boolean values" + else: + mask = gu.PointCloud.from_xyz(x.ravel() + 1, y.ravel(), mask, crs=points.crs) + error = "does not share the ordered support coordinates" + + # Report the mask problem before choosing any pair endpoints + with pytest.raises(ValueError, match=error): + points.pairsample(n_pairs=10, mask=mask, random_state=9) + + def test_pairsample__point_exact_sampling_reuses_anchors(self) -> None: + """Checks that exact point searches can reuse first endpoints when one round requests more than exist.""" + + # Create fewer points than the requested number of first endpoints per round + y, x = np.mgrid[:5, :5] + points = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), (x + y).ravel(), crs=32633) + # Request enough pairs to require reuse during the single allowed round + pairs = points.pairsample( + n_pairs=50, + min_distance=1, + max_distance=5, + strategy="kdtree", + anchors_per_round=100, + attempts_per_anchor=2, + max_rounds=1, + random_state=1, + ) + + # Check that reuse still fills the requested sample + assert pairs.sizes["pair"] == 50 + + @pytest.mark.parametrize("strategy", ["kdtree", "hashgrid"]) + def test_pairsample__point_exact_maximum_distance(self, strategy: str) -> None: + """Checks that exact point searches include pairs at the requested maximum distance.""" + + # Use two points whose only nonzero separation is exactly the upper boundary + points = gu.PointCloud.from_xyz(np.array([0, 1]), np.array([0, 0]), np.array([3, 5]), crs=32633) + pairs = points.pairsample( + n_pairs=10, strategy=strategy, n_bins=1, min_distance=0.5, max_distance=1, random_state=2 + ) + + # Both exact search methods must find the available pair, with either endpoint order + assert np.array_equal(pairs.distance, np.ones(10)) + assert np.array_equal(np.sort(pairs["index"], axis=1), np.tile([0, 1], (10, 1))) + + +class TestPairSampleChunked: + """Checks pairsample() loading behavior and exact results with chunked inputs.""" + + @pytest.mark.parametrize("strategy", ["independent", "anchors", "chunk_anchors", "anchor_batched"]) + def test_pairsample__raster_uneven_local_chunks(self, strategy: str) -> None: + """Checks that local pairs stay in the same actual chunk when interior chunk sizes vary.""" + + # Use irregular row and column boundaries that differ from a repeated first-chunk grid + da = import_optional("dask.array").array + chunks = ((4, 9, 7), (5, 3, 12)) + eager = np.arange(400, dtype=float).reshape(20, 20) + array = da.from_array(eager, chunks=chunks) + raster = gu.RasterAccessor.from_array(array, from_origin(0, 20, 1, 1), 32633) + + # Draw only local pairs so both endpoints must belong to one original Dask chunk + pairs = raster.rst.pairsample( + n_pairs=100, strategy=strategy, hybrid_local_fraction=1, min_distance=1, max_distance=6, random_state=4 + ) + + # The Dask input stays lazy and pairs are returned as an eager dataset + assert isinstance(raster.data, da.Array) + assert not raster._in_memory + assert not pairs.chunks + + # Find chunks independently from their cumulative boundaries and check each endpoint pair exactly + rows = np.searchsorted(np.cumsum(chunks[0]), pairs.row, side="right") + columns = np.searchsorted(np.cumsum(chunks[1]), pairs.column, side="right") + assert np.array_equal(rows[:, 0], rows[:, 1]) + assert np.array_equal(columns[:, 0], columns[:, 1]) + assert np.array_equal(pairs.value, eager[pairs.row, pairs.column]) + assert pairs.sizes["pair"] == 100 + + def test_pairsample__raster_dask_endpoint_reads(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Checks that a Dask chunk is read once for both endpoints of a pair.""" + + # Count reads of one source chunk so separate endpoint computations would be visible + dask = import_optional("dask") + da = import_optional("dask.array").array + reads, candidate_counts = [], [] + + @dask.delayed + def load_values() -> NDArrayNum: + """Record each source read and return a small raster with finite values.""" + reads.append(1) + return np.arange(400, dtype=float).reshape(20, 20) + + # Count sampling rounds separately from the initial finite count and final output read + original_candidates = _RegularPairSampler._candidates + + def candidates(sampler: _RegularPairSampler, count: int) -> tuple[NDArrayNum, NDArrayNum]: + """Record candidate generation before checking endpoint values.""" + candidate_counts.append(count) + return original_candidates(sampler, count) + + monkeypatch.setattr(_RegularPairSampler, "_candidates", candidates) + array = da.from_delayed(load_values(), shape=(20, 20), dtype=float) + raster = gu.RasterAccessor.from_array(array, from_origin(0, 20, 1, 1), 32633) + with dask.config.set(scheduler="synchronous"): + pairs = raster.rst.pairsample(n_pairs=80, strategy="independent", random_state=2) + dask_candidate_count = len(candidate_counts) + expected = gu.Raster.from_array( + np.arange(400, dtype=float).reshape(20, 20), from_origin(0, 20, 1, 1), 32633 + ).pairsample(n_pairs=80, strategy="independent", random_state=2) + + # Expect one count read, one read per round, and one output read regardless of two endpoints + assert isinstance(raster.data, da.Array) and not raster._in_memory + assert not pairs.chunks + xr.testing.assert_equal(pairs, expected) + assert pairs.sizes["pair"] == 80 + assert len(reads) == dask_candidate_count + 2 + assert np.array_equal(pairs.value, np.asarray(pairs["index"], dtype=float)) + + def test_pairsample__raster_dask_source_is_lazy(self) -> None: + """Checks that raster pair sampling reads selected Dask cells without loading the source.""" + + # Create one lazy raster chunk so sampling order also has an exact eager reference + da = pytest.importorskip("dask.array") + array = np.arange(600, dtype=float).reshape(24, 25) + raster = gu.RasterAccessor.from_array( + da.from_array(array, chunks=array.shape), from_origin(0, 24, 2, 2), 32633, nodata=None + ) + + # Draw pairs through the Xarray accessor without loading the complete array + pairs = raster.rst.pairsample(n_pairs=250, hybrid_local_fraction=0, random_state=42) + expected = gu.Raster.from_array(array, from_origin(0, 24, 2, 2), 32633).pairsample( + n_pairs=250, hybrid_local_fraction=0, random_state=42 + ) + assert pairs.sizes["pair"] == 250 + assert isinstance(raster.data, da.Array) + assert not raster._in_memory + assert not pairs.chunks + xr.testing.assert_equal(pairs, expected) + + def test_pairsample__raster_dask_local_chunks_and_dtypes(self) -> None: + """Checks that nearby Dask pairs stay in one chunk and use the requested number types.""" + + # Create a lazy raster whose row and column chunks have different sizes + da = pytest.importorskip("dask.array") + array = np.arange(576, dtype=float).reshape(24, 24) + raster = gu.RasterAccessor.from_array( + da.from_array(array, chunks=(6, 8)), from_origin(0, 24, 1, 1), 32633, nodata=None + ) + # Request only nearby pairs and smaller output number types + pairs = raster.rst.pairsample( + n_pairs=200, + min_distance=1, + max_distance=6, + hybrid_local_fraction=1, + random_state=7, + index_dtype=np.int16, + distance_dtype=np.float32, + ) + + # Check that both endpoints share a chunk and that output types match the request + assert isinstance(raster.data, da.Array) and not raster._in_memory + assert not pairs.chunks + first_chunk = np.column_stack((pairs.row[:, 0] // 6, pairs.column[:, 0] // 8)) + second_chunk = np.column_stack((pairs.row[:, 1] // 6, pairs.column[:, 1] // 8)) + assert np.array_equal(first_chunk, second_chunk) + assert np.array_equal(pairs.value, array[pairs.row, pairs.column]) + assert pairs["index"].dtype == np.int16 + assert pairs["distance"].dtype == np.float32 + + @pytest.mark.parametrize("mask_form", ["vector", "raster", "pointcloud"]) + def test_pairsample__point_masks_reuse_loaded_coordinates(self, mask_form: str) -> None: + """Checks that spatial masking reads each Dask source partition once during eager pair sampling.""" + + # Count reads from three source partitions so a second coordinate load would be visible + dask = import_optional("dask") + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import _register_dask_pointcloud_accessor + + _register_dask_pointcloud_accessor() + y, x = np.mgrid[:12, :12] + dataframe = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), (x + y).ravel(), crs=32633).ds + reads = [] + + def read_partition(partition: gpd.GeoDataFrame) -> gpd.GeoDataFrame: + """Record a source read before returning the point partition.""" + reads.append(1) + return partition + + lazy_points = dgpd.from_geopandas(dataframe, npartitions=3, sort=False).map_partitions( + read_partition, meta=dataframe.iloc[:0] + ) + + # Define the same left-half mask, matching integer raster coordinates to the point grid + mask: Any + if mask_form == "raster": + mask = gu.Raster.from_array((x < 6)[::-1], from_origin(0, 11, 1, 1), dataframe.crs) + elif mask_form == "pointcloud": + mask = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), (x < 6).ravel(), crs=dataframe.crs) + else: + mask = gpd.GeoDataFrame(geometry=[box(-0.5, -0.5, 5.5, 11.5)], crs=dataframe.crs) + + # Evaluate pairs immediately, reusing the point table that was loaded for the pair search + with dask.config.set(scheduler="synchronous"): + pairs = lazy_points.pc.pairsample(n_pairs=100, sampling="random_xy", mask=mask, random_state=9) + expected = dataframe.pc.pairsample(n_pairs=100, sampling="random_xy", mask=mask, random_state=9) + + assert len(reads) == 3 + assert not lazy_points.pc.is_loaded and not pairs.chunks + xr.testing.assert_equal(pairs, expected) + assert pairs.sizes["pair"] == 100 + assert np.all(pairs.x < 6) diff --git a/tests/test_sampling/test_stratified.py b/tests/test_sampling/test_stratified.py new file mode 100644 index 000000000..e6547d958 --- /dev/null +++ b/tests/test_sampling/test_stratified.py @@ -0,0 +1,228 @@ +"""Tests 'stratified' subsampling, i.e. subsampling independent within integer group IDs.""" + +from __future__ import annotations + +import warnings +from typing import Literal +from unittest.mock import patch + +import numpy as np +import pytest + +from geoutils._misc import import_optional +from geoutils.multiproc import MultiprocConfig +from geoutils.sampling.stratified import _stratified_subsample_indices +from geoutils.sampling.subsampling import _splitmix64 + + +class TestStratifiedSubsample: + """ + Test module for stratified sampling (i.e. independent subsampling within integer groups). + + Tests covering Dask/Multiprocessing backends are located further below in TestStratifiedSubsampleChunked. + Tests using the parent function stats(by=, subsample_per_group=True) which uses stratified sampling are in + test_grouping.py, here we test the subfunctions instead. + + For eager arrays, we check that: + - Each group receives the requested number or fraction of samples, while negative group IDs are properly excluded. + - A fixed seed selects the same samples when group IDs or the array shape change. + - Groups smaller than the requested sample size contribute all their samples without warnings. + - Zero, negative, infinite and NaN sample sizes raise an error. + """ + + def test_stratified_subsample_indices__sample_size_per_group(self) -> None: + """Checks that a fixed sample size is applied separately to each valid group ID.""" + + # Select two positions from each valid group and exclude locations marked with an ID of -1 + groups = np.array([0, 0, 0, -1, 1, 1, 1, 1]) + selected = _stratified_subsample_indices(groups, 2, random_state=42) + + # Check the basic output shape and selected group counts + assert selected.shape == (4,) + assert np.array_equal(np.bincount(groups[selected]), [2, 2]) + assert np.all(groups[selected] >= 0) + + @pytest.mark.parametrize("subsample", [1, 2, 5, 100, 0.1, 0.5]) + def test_stratified_subsample_indices__topk_reference(self, subsample: int | float) -> None: + """Checks that topk selects positions with the lowest seeded random scores within each group.""" + + # Use unequal groups and excluded locations so fractions, maximum counts and empty samples are distinguishable + groups = np.array([0, -1, 0, 1, 1, -1, 1, 1, 1, 1, 1, 1, 2, -1, 2]) + selected = _stratified_subsample_indices(groups, subsample, random_state=42) + + # Build a reference for each group by sorting seeded random scores from the original positions + expected = [] + for group in (0, 1, 2): + positions = np.flatnonzero(groups == group) + sample_size = int(subsample * len(positions)) if subsample <= 1 else min(int(subsample), len(positions)) + keys = _splitmix64(np.uint64(42) ^ positions.astype(np.uint64)) + expected.extend(positions[np.argsort(keys)[:sample_size]]) + + # Check the documented output order and confirm that no location is duplicated or selected from ID -1 + expected_indices = np.asarray(expected, dtype=np.int64) + if subsample == 1: + expected_indices.sort() + else: + keys = _splitmix64(np.uint64(42) ^ expected_indices.astype(np.uint64)) + expected_indices = expected_indices[np.argsort(keys)] + assert np.array_equal(selected, expected_indices) + assert len(np.unique(selected)) == len(selected) + assert np.all(groups[selected] >= 0) + + def test_stratified_subsample_indices__original_positions(self) -> None: + """ + Checks that "topk" strategy uses original positions across group IDs, array shapes for a given random seed. + """ + + # Assign different positive IDs to the same groups and reshape their locations + groups = np.repeat([0, 1, 2], [2, 8, 5]) + expected = _stratified_subsample_indices(groups, 3, random_state=42) + remapped = _stratified_subsample_indices(100 - groups, 3, random_state=42) + reshaped = _stratified_subsample_indices(groups.reshape(3, 5), 3, random_state=42) + assert np.array_equal(remapped, expected) + assert np.array_equal(reshaped, expected) + + # Check that one seed is drawn from a supplied random generator for the complete sampling call + actual_rng = np.random.default_rng(9) + reference_rng = np.random.default_rng(9) + seed = int(reference_rng.integers(0, np.iinfo(np.uint32).max, dtype=np.uint32)) + selected = _stratified_subsample_indices(groups, 3, random_state=actual_rng) + assert np.array_equal(selected, _stratified_subsample_indices(groups, 3, random_state=seed)) + assert actual_rng.integers(100000) == reference_rng.integers(100000) + + def test_stratified_subsample_indices__per_group_selection(self) -> None: + """ + Checks that partial selection sorts groups and includes every position from smaller groups without warnings. + """ + + # Use groups of very different sizes so each partial sort reveals exactly which group slice it received + groups = np.repeat([10, 20, 30], [2, 20, 200]) + with warnings.catch_warnings(record=True) as recorded: + with patch("numpy.argpartition", wraps=np.argpartition) as partition: + selected = _stratified_subsample_indices(groups, 5, random_state=42) + + # The two larger groups are partially selected; the first group already fits within its sample size + assert [call.args[0].size for call in partition.call_args_list] == [20, 200] + assert np.array_equal(np.unique(groups[selected], return_counts=True)[1], [2, 5, 5]) + assert not recorded + + @pytest.mark.parametrize("subsample", [0, -1, np.inf, np.nan]) + def test_stratified_subsample_indices__error_invalid_subsample(self, subsample: int | float) -> None: + """Checks that invalid sampling amounts fail even when all locations are excluded.""" + + with pytest.raises(ValueError, match="positive finite"): + _stratified_subsample_indices(np.full(5, -1), subsample) + + +class TestStratifiedSubsampleChunked: + """Checks stratified subsampling against eager results for Dask arrays and Multiproc chunks.""" + + @pytest.mark.parametrize( + "shape,chunks", [((120,), (7,)), ((10, 12), (3, 5)), ((10, 12), (4, 4)), ((10, 12), (12, 12))] + ) + @pytest.mark.parametrize("subsample", [1, 3, 0.15, 0.5]) + @pytest.mark.parametrize("strategy", ["topk", "sequential"]) + def test_stratified_subsample_indices__eager_dask_and_multiproc( + self, + shape: tuple[int, ...], + chunks: tuple[int, ...], + subsample: int | float, + strategy: Literal["topk", "sequential"], + ) -> None: + """ + Checks that stratified subsampling selects exactly the same positions (for "topk" strategy) and behave the + same depending on chunks (for "sequential" strategy). + """ + + # Create repeating group IDs, exclude every seventh location, and add one group containing a single location + import_optional("dask") + import dask.array as da + + groups = (np.arange(120) % 4).reshape(shape) + groups.ravel()[::7] = -1 + groups.ravel()[1] = 4 + lazy = da.from_array(groups, chunks=chunks) + config = MultiprocConfig(chunks=chunks[0] if len(chunks) == 1 else (chunks[0], chunks[1])) + + # Sample the same group IDs through the eager, Dask and Multiproc backends + eager = _stratified_subsample_indices(groups, subsample, random_state=42, strategy=strategy) + dask_result = _stratified_subsample_indices(lazy, subsample, random_state=42, strategy=strategy) + tile_result = _stratified_subsample_indices( + groups, subsample, random_state=42, strategy=strategy, mp_config=config + ) + + # Inputs remain chunked, while this index-selection helper returns eager NumPy positions + assert isinstance(lazy, da.Array) + assert isinstance(dask_result, np.ndarray) and isinstance(tile_result, np.ndarray) + + # Fractions must be rounded from the full group sizes, rather than once in each chunk + _, counts = np.unique(groups[groups >= 0], return_counts=True) + expected_counts = (counts * subsample).astype(int) if subsample <= 1 else np.minimum(counts, int(subsample)) + for selected in (eager, dask_result, tile_result): + assert np.array_equal(np.bincount(groups.ravel()[selected], minlength=5), expected_counts) + assert len(np.unique(selected)) == len(selected) + assert np.array_equal(dask_result, tile_result) + + # Topk is independent of chunks; sequential sampling repeats when the chunk layout stays the same + if strategy == "topk": + assert np.array_equal(dask_result, eager) + else: + repeated = _stratified_subsample_indices(lazy, subsample, random_state=42, strategy=strategy) + assert np.array_equal(dask_result, repeated) + + @pytest.mark.parametrize("strategy", ["topk", "sequential"]) + def test_stratified_subsample_indices__multiproc_workers( + self, + strategy: Literal["topk", "sequential"], + ) -> None: + """Checks that Multiproc workers combine several task batches and match the positions selected with Dask.""" + + # Use more than eight tiles so Multiproc processes several task batches and combines their samples + from geoutils.multiproc.cluster import MpCluster + + import_optional("dask") + import dask.array as da + + groups = (np.arange(360) % 7).reshape(18, 20) + groups[::3, ::4] = -1 + lazy = da.from_array(groups, chunks=(3, 4)) + expected = _stratified_subsample_indices(lazy, 0.3, 42, strategy) + eager = _stratified_subsample_indices(groups, 0.3, 42, strategy) + + # Run the same tile layout through two Multiproc worker processes and compare the selected positions + with MpCluster({"nb_workers": 2, "max_tasks_per_child": None}) as cluster: + config = MultiprocConfig(chunks=(3, 4), cluster=cluster) + selected = _stratified_subsample_indices(groups, 0.3, 42, strategy, mp_config=config) + assert np.array_equal(selected, expected) + assert isinstance(lazy, da.Array) and isinstance(selected, np.ndarray) + if strategy == "topk": + assert np.array_equal(selected, eager) + + @pytest.mark.parametrize("strategy", ["topk", "sequential"]) + def test_stratified_subsample_indices__empty_groups(self, strategy: Literal["topk", "sequential"]) -> None: + """Checks that ID -1 and sample sizes rounded to zero return no positions for eager and Dask inputs.""" + + # Create one input with no valid group ID and another with one valid location in its final Dask block + import_optional("dask") + import dask.array as da + + all_excluded = np.full((4, 6), -1) + tiny_group = all_excluded.copy() + tiny_group[-1, -1] = 0 + + # 50% of a one-member group rounds to zero (same floor rule that we use for ordinary subsampling) + for groups in (all_excluded, tiny_group): + expected = _stratified_subsample_indices(groups, 0.5, random_state=42, strategy=strategy) + lazy = da.from_array(groups, chunks=(2, 3)) + selected = _stratified_subsample_indices(lazy, 0.5, random_state=42, strategy=strategy) + assert isinstance(lazy, da.Array) and isinstance(selected, np.ndarray) + assert np.array_equal(selected, expected) + assert selected.size == 0 + assert np.issubdtype(selected.dtype, np.integer) + + # Check that a sample size of two selects the only valid location despite the other empty Dask blocks + expected = _stratified_subsample_indices(tiny_group, 2, random_state=42, strategy=strategy) + lazy = da.from_array(tiny_group, chunks=(2, 3)) + selected = _stratified_subsample_indices(lazy, 2, random_state=42, strategy=strategy) + assert np.array_equal(selected, expected) + assert np.array_equal(selected, [23]) diff --git a/tests/test_sampling/test_subsampling.py b/tests/test_sampling/test_subsampling.py new file mode 100644 index 000000000..d3fe5bc8a --- /dev/null +++ b/tests/test_sampling/test_subsampling.py @@ -0,0 +1,726 @@ +"""Test internal raster and point cloud subsampling tools.""" + +from __future__ import annotations + +import warnings +from pathlib import Path +from typing import Any, Literal +from unittest.mock import patch + +import geopandas as gpd +import numpy as np +import pytest +import xarray as xr +from rasterio.transform import from_origin +from shapely.geometry import box + +import geoutils as gu +from geoutils import open_raster +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.multiproc import MultiprocConfig +from geoutils.raster.array import get_mask_from_array +from geoutils.sampling.subsampling import ( + _sample_valid_indices, + _subsample_numpy, +) + + +class TestSubsample: + """ + Base tests for eager subsampling with one-, two-, and three-dimensional masked arrays. + + Tests specific to point/raster inputs are available in other modules below, as well as for the testing across + Dask and Multiprocessing backends! + """ + + # One-dimensional array with one masked value + array1D = np.ma.masked_array(np.arange(10), mask=np.zeros(10)) + array1D.mask[3] = True + assert np.ndim(array1D) == 1 + assert np.count_nonzero(array1D.mask) > 0 + + # Two-dimensional array with one masked value + array2D = np.ma.masked_array(np.arange(9).reshape((3, 3)), mask=np.zeros((3, 3))) + array2D.mask[0, 1] = True + assert np.ndim(array2D) == 2 + assert np.count_nonzero(array2D.mask) > 0 + + # Three-dimensional array with one masked value + array3D = np.ma.masked_array(np.arange(9).reshape((1, 3, 3)), mask=np.zeros((1, 3, 3))) + array3D = np.ma.vstack((array3D, array3D + 10)) + array3D.mask[0, 0, 1] = True + assert np.ndim(array3D) == 3 + assert np.count_nonzero(array3D.mask) > 0 + + @pytest.mark.parametrize("array", [array1D, array2D, array3D]) + def test_subsample(self, array: NDArrayNum) -> None: + """Checks that counts, fractions, returned indexes, and random seeds follow the public sampling rules.""" + + warnings.filterwarnings("ignore", message=".*larger than the number of valid pixels.*", category=UserWarning) + + # Check every requested count below the input size + for npts in np.arange(2, np.size(array)): + random_values = _subsample_numpy(array, subsample=npts) + assert np.ndim(random_values) == 1 + assert np.size(random_values) == npts + assert np.count_nonzero(random_values.mask) == 0 + + # Check that a count above the available values returns every available value + random_values = _subsample_numpy(array, subsample=np.size(array) + 3) + assert np.all(np.sort(random_values) == array[~array.mask]) + + # Check that 1 returns every available value in the original order + random_values = _subsample_numpy(array, subsample=1) + assert np.all(np.sort(random_values) == array[~array.mask]) + + random_values_2 = _subsample_numpy(array, subsample=1) + assert np.array_equal(random_values, random_values_2) + + # Check that a fraction between 0 and 1 returns the right amount of valid values + random_values = _subsample_numpy(array, subsample=0.5) + assert np.size(random_values) == int(np.count_nonzero(~array.mask) * 0.5) + + # Check returned indexes against the input dimensions and requested fraction + indices = _subsample_numpy(array, subsample=0.3, return_indices=True) + assert np.ndim(indices) == 2 + assert len(indices) == np.ndim(array) + assert np.ndim(array[indices]) == 1 + assert np.size(array[indices]) == int(np.count_nonzero(~array.mask) * 0.3) + + # Check that an integer seed and the matching NumPy generator select the same values + sub42 = _subsample_numpy(array, subsample=10, random_state=42) + rng = np.random.default_rng(42) + sub42_gen = _subsample_numpy(array, subsample=10, random_state=rng) + assert np.array_equal(sub42, sub42_gen) + + +class TestPointSubsample: + """ + Checks subsample() on point values. + + We only check behaviour with input masks on point data here. Other tests are done directly in TestSubsample, or in + TestSubsampleChunked. + """ + + @pytest.mark.parametrize("mask_form", ["array", "masked", "xarray", "vector", "raster", "pointcloud"]) + def test_subsample__point_masks(self, mask_form: str) -> None: + """Checks all types of input inlier masks.""" + + # Give the points duplicate labels so only positional indexes can identify the sampled rows + y, x = np.mgrid[:6, :6] + values = np.arange(36, dtype=np.int16) + points = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), values, crs=32633) + points.ds.index = np.arange(36) % 3 + original = points.ds.copy() + eligible = (x < 3).ravel() + mask: Any = eligible.copy() + + # Build the same mask with different input types, covering only half of the input + if mask_form == "masked": + mask = np.ma.array(mask, mask=False) + mask.mask[7] = True + eligible[7] = False + elif mask_form == "xarray": + mask = xr.DataArray(mask.reshape(6, 6), dims=("row", "column")) + elif mask_form == "vector": + mask = gpd.GeoDataFrame(geometry=[box(-0.5, -0.5, 2.5, 5.5)], crs=points.crs) + elif mask_form == "raster": + mask = gu.Raster.from_array((x < 3)[::-1], from_origin(0, 5, 1, 1), points.crs) + elif mask_form == "pointcloud": + mask = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), eligible, crs=points.crs) + + # Draw expected sample + expected_rng = np.random.default_rng(42) + expected_indices = expected_rng.choice(np.flatnonzero(eligible), int(eligible.sum() * 0.5), replace=False) + actual_rng = np.random.default_rng(42) + indices = points.subsample(0.5, return_indices=True, mask=mask, random_state=actual_rng) + sampled = points.subsample(0.5, mask=mask, random_state=42) + + # Compare exact equality + np.testing.assert_array_equal(indices[0], expected_indices) + np.testing.assert_array_equal(sampled, values[expected_indices]) + assert sampled.dtype == values.dtype + assert actual_rng.integers(100000) == expected_rng.integers(100000) + assert points.ds.equals(original) + + @pytest.mark.parametrize("wrong_count", [False, True]) + def test_subsample__error_invalid_point_mask(self, wrong_count: bool) -> None: + """Checks that point subsampling rejects non-boolean masks, or masks with a different shape.""" + + # We raise an error for non-boolean masks, or array masks with a wrong shape + values = np.arange(6) + points = gu.PointCloud.from_xyz(values, np.zeros(6), values, crs=32633) + mask = np.ones(5, dtype=bool) if wrong_count else np.ones(6, dtype=int) + with pytest.raises( + ValueError, match="Argument ``mask`` must be boolean and contain one value per input location" + ): + points.subsample(1, mask=mask) + + +class TestRasterSubsample: + """Checks subsample() on raster values and masks. + + Selected bands return exact values and original grid indexes. Empty masks return the expected output shape and type. + Nonempty Dask and Multiproc comparisons are covered in TestSubsampleChunked. + """ + + @pytest.mark.parametrize("mask_form", ["array", "masked", "xarray", "vector", "raster"]) + def test_subsample__raster_masks(self, mask_form: str) -> None: + """Checks that masks restrict the sampled band and return its exact values and original grid indices.""" + + # Use distinct integer bands and a nodata second-band cell to expose band selection or mask mistakes + values = np.arange(36, dtype=np.int16).reshape(6, 6) + data = np.ma.array(np.stack([values + 100, values]), mask=False) + data.mask[1, 1, 1] = True + raster = gu.Raster.from_array(data, from_origin(0, 6, 1, 1), crs=32633, nodata=-9999) + original = raster.data.copy() + eligible = np.indices(raster.shape)[1] < 3 + mask: Any = eligible.copy() + + # Select the left half of the grid, excluding nodata entries of an explicitly masked boolean array + if mask_form == "masked": + mask = np.ma.array(mask, mask=False) + mask.mask[2, 1] = True + eligible[2, 1] = False + elif mask_form == "xarray": + mask = xr.DataArray(mask[None], dims=("band", "row", "column")) + elif mask_form == "vector": + mask = gpd.GeoDataFrame(geometry=[box(0, 0, 3, 6)], crs=raster.crs) + elif mask_form == "raster": + mask = gu.Raster.from_array(mask, raster.transform, raster.crs) + eligible &= ~data.mask[1] + + # Draw from original flat grid positions to check the count, random order and generator advancement + expected_rng = np.random.default_rng(42) + expected_flat = expected_rng.choice(np.flatnonzero(eligible), int(eligible.sum() * 0.5), replace=False) + expected_indices = np.unravel_index(expected_flat, raster.shape) + actual_rng = np.random.default_rng(42) + indices = raster.subsample(0.5, band=2, return_indices=True, mask=mask, random_state=actual_rng) + sampled = raster.subsample(0.5, band=2, mask=mask, random_state=42) + + # Check that sampled integers match the selected band and the source data and nodata mask are unchanged + np.testing.assert_array_equal(indices, expected_indices) + np.testing.assert_array_equal(sampled, values[expected_indices]) + assert sampled.dtype == values.dtype + assert actual_rng.integers(100000) == expected_rng.integers(100000) + np.testing.assert_array_equal(raster.data.data, original.data) + np.testing.assert_array_equal(raster.data.mask, original.mask) + + @pytest.mark.parametrize("mask_form", ["numeric", "different_shape", "flat", "different_grid"]) + def test_subsample__error_invalid_raster_mask(self, mask_form: str) -> None: + """Checks that raster masks must contain booleans on the same two-dimensional grid as the source.""" + + # Arrays with the same cell count still need the source shape; spatial masks also need matching coordinates + raster = gu.Raster.from_array(np.ones((3, 4)), from_origin(0, 3, 1, 1), crs=32633) + mask: Any = np.ones(raster.shape, dtype=int) + message = "Argument ``mask`` must be boolean" + if mask_form == "different_shape": + mask = np.ones((2, 6), dtype=bool) + message = "match the support grid" + elif mask_form == "flat": + mask = np.ones(12, dtype=bool) + message = "match the support grid" + elif mask_form == "different_grid": + mask = gu.Raster.from_array(np.ones(raster.shape, dtype=bool), from_origin(1, 3, 1, 1), raster.crs) + message = "does not share the selected support grid" + + # Fail during mask validation instead of sampling a silently coerced or reshaped population + with pytest.raises(ValueError, match=message): + raster.subsample(1, mask=mask) + + +class TestSubsampleChunked: + """Checks subsample() across eager, Dask and Multiproc inputs. + + We check that backends return the requested sample size and exactly the same values (for "topk" strategy). + We also check behaviour with selected bands, mask input and chunk size. + """ + + # Strategies supported by _subsample() + subsample_strategies = ("sequential", "topk") + + @pytest.mark.parametrize("path_index", [0, 2]) + @pytest.mark.parametrize("strategy", subsample_strategies) + @pytest.mark.parametrize("return_indices", [False, True]) + @pytest.mark.parametrize("subsample", [2, 100, 0.05]) # int size and fraction + def test_subsample__backends( + self, + path_index: int, + strategy: Literal["sequential", "topk"], + return_indices: bool, + subsample: int | float, + lazy_test_files_tiny: list[str], + ) -> None: + """Checks that all storage paths follow the same sampling, loading, and repeatability rules.""" + + pytest.importorskip("dask") + import dask.array as da + + warnings.filterwarnings("ignore", category=UserWarning, message="Argument ``subsample`` with value*") + + # 1/ Open matching inputs for NumPy, Dask, and multiprocessing + path_raster = lazy_test_files_tiny[path_index] + + # The two NumPy calls use loaded Raster and Xarray inputs + raster_base = gu.Raster(path_raster) + raster_base.load() + assert raster_base.is_loaded + + ds_base = open_raster(path_raster) + ds_base.load() + assert ds_base._in_memory + + # Worker processes read tiles from a Raster that stays linked to its file + raster_mp = gu.Raster(path_raster) + assert not raster_mp.is_loaded + + # The Xarray accessor keeps the chunked Dask input lazy + ds_dask = open_raster(path_raster, chunks={"x": 10, "y": 10}) + assert not ds_dask._in_memory + assert isinstance(ds_dask.data, da.Array) + assert ds_dask.data.chunks is not None + + # 2/ Run every storage path with the same seed + seed = 42 + mp_config = MultiprocConfig(chunks=(10, 7)) + + # NumPy through Raster + out_raster = raster_base.subsample( + subsample=subsample, + return_indices=return_indices, + random_state=seed, + strategy=strategy, + ) + + # NumPy through an Xarray accessor + out_xr = ds_base.rst.subsample( + subsample=subsample, + return_indices=return_indices, + random_state=seed, + strategy=strategy, + ) + + # Dask through an Xarray accessor + out_dask = ds_dask.rst.subsample( + subsample=subsample, + return_indices=return_indices, + random_state=seed, + strategy=strategy, + ) + + # Worker processes through Raster + out_mp = raster_mp.subsample( + subsample=subsample, + return_indices=return_indices, + random_state=seed, + strategy=strategy, + mp_config=mp_config, + ) + + # 3/ Check that Dask and file-backed inputs stay unloaded + assert not ds_dask._in_memory + assert isinstance(ds_dask.data, da.Array) + + assert not raster_mp.is_loaded + + # Dask outputs also stay lazy until values are compared + if return_indices: + assert isinstance(out_dask, tuple) and len(out_dask) == 2 + assert isinstance(out_dask[0], da.Array) + assert isinstance(out_dask[1], da.Array) + else: + assert isinstance(out_dask, da.Array) + + # 4/ Load each small result as NumPy arrays for comparison + + def _as_numpy( + out: object, + ) -> NDArrayNum | tuple[NDArrayNum, NDArrayNum]: + """Convert any returned values or indexes to NumPy arrays.""" + if isinstance(out, tuple): + r, c = out + if hasattr(r, "compute"): + r = r.compute() + if hasattr(c, "compute"): + c = c.compute() + return (np.asarray(r), np.asarray(c)) + else: + if hasattr(out, "compute"): + out = out.compute() + return np.asarray(out) + + out_raster_np = _as_numpy(out_raster) + out_xr_np = _as_numpy(out_xr) + out_dask_np = _as_numpy(out_dask) + out_mp_np = _as_numpy(out_mp) + + # 5/ Check the result size and values shared by every storage path + # _subsample() uses band one by default + arr = raster_base.data if raster_base.data.ndim == 2 else raster_base.data[0, :, :] + assert arr.ndim == 2 + + mask = get_mask_from_array(arr) + n_valid = int(np.count_nonzero(~mask)) + + # Calculate the expected sample size independently from the implementation helper + if isinstance(subsample, float): + expected = int(subsample * n_valid) + else: + expected = min(int(subsample), n_valid) + + def _check_output(out_np: NDArrayNum | tuple[NDArrayNum, NDArrayNum]) -> None: + """Check one output's size and that every returned position is available.""" + if isinstance(out_np, tuple): + rr, cc = out_np + assert rr.shape == cc.shape + assert rr.ndim == 1 and cc.ndim == 1 + assert len(rr) == expected + # Returned rows and columns must lie inside the raster + assert np.all((0 <= rr) & (rr < arr.shape[0])) + assert np.all((0 <= cc) & (cc < arr.shape[1])) + # Every returned cell must be available according to the shared raster mask + assert np.all(~mask[rr, cc]) + else: + assert out_np.ndim == 1 + assert len(out_np) == expected + assert np.all(np.isfinite(out_np)) + + _check_output(out_raster_np) + _check_output(out_xr_np) + _check_output(out_dask_np) + _check_output(out_mp_np) + + # 6/ Check exact agreement for `topk` and repeatability for `sequential` + if strategy == "topk": + assert np.array_equal(out_raster_np, out_xr_np) + assert np.array_equal(out_raster_np, out_dask_np) + assert np.array_equal(out_raster_np, out_mp_np) + else: + # Sequential sampling depends on chunk order, so repeat each path with the same seed + out_raster_np_2 = _as_numpy( + raster_base.subsample( + subsample=subsample, + return_indices=return_indices, + random_state=seed, + strategy=strategy, + ) + ) + out_dask_np_2 = _as_numpy( + ds_dask.rst.subsample( + subsample=subsample, + return_indices=return_indices, + random_state=seed, + strategy=strategy, + ) + ) + out_mp_np_2 = _as_numpy( + raster_mp.subsample( + subsample=subsample, + return_indices=return_indices, + random_state=seed, + strategy=strategy, + mp_config=mp_config, + ) + ) + + assert np.array_equal(out_raster_np, out_raster_np_2) + assert np.array_equal(out_dask_np, out_dask_np_2) + assert np.array_equal(out_mp_np, out_mp_np_2) + + # 7/ Check that returned indexes select the same values as value mode + if return_indices: + rr, cc = out_raster_np + vals_from_indices = arr[rr, cc] + vals_raster = raster_base.subsample( + subsample=subsample, + return_indices=False, + random_state=seed, + strategy=strategy, + ) + vals_raster_np = _as_numpy(vals_raster) + assert np.array_equal(np.asarray(vals_from_indices), np.asarray(vals_raster_np)) + + @pytest.mark.parametrize("subsample", [1, 5, 0.25, 0.001]) + @pytest.mark.parametrize("return_indices", [False, True]) + @pytest.mark.parametrize("mask_form", ["none", "array", "vector"]) + def test_subsample__lazy_point_values(self, subsample: int | float, return_indices: bool, mask_form: str) -> None: + """Checks that lazy point sampling gathers the requested values in the same seeded order as NumPy.""" + + # Give several points the same labels so returned indexes must refer to row positions + import_optional("dask_geopandas") + import dask_geopandas as dgpd + + from geoutils.pointcloud.pd_accessor import _register_dask_pointcloud_accessor + + values = np.arange(40, dtype=float) - 1 + values[::9] = np.nan + values[2] = np.inf + dataframe = gpd.GeoDataFrame( + {"height": values}, geometry=gpd.points_from_xy(np.arange(40), np.zeros(40)), crs=32632 + ) + dataframe.index = np.arange(40) % 3 + _register_dask_pointcloud_accessor() + lazy = dgpd.from_geopandas(dataframe, npartitions=6, sort=False) + assert not lazy.pc.is_loaded + + # Restrict the population before applying counts or fractions, without changing the point data + eligible = np.ones(40, dtype=bool) + mask: Any = None + if mask_form == "array": + eligible = np.arange(40) % 4 != 0 + mask = eligible + elif mask_form == "vector": + eligible = (np.arange(40) >= 10) & (np.arange(40) < 30) + mask = gpd.GeoDataFrame(geometry=[box(9.5, -0.5, 29.5, 0.5)], crs=dataframe.crs) + expected_values = np.ma.array(values, mask=~eligible) + + # Compare exact random order and generator advancement with the established finite NumPy sampler + actual_rng = np.random.default_rng(42) + expected_rng = np.random.default_rng(42) + expected_indices = _subsample_numpy(expected_values, subsample, return_indices=True, random_state=expected_rng) + expected = expected_indices if return_indices else values[expected_indices] + series_type = type(lazy["height"]) + with patch.object(series_type, "compute", autospec=True, side_effect=series_type.compute) as compute: + result = lazy.pc.subsample(subsample, return_indices=return_indices, random_state=actual_rng, mask=mask) + + # Only row-count summaries may be collected as Series; the full original data column stays partitioned + assert all(call.args[0].name != "height" for call in compute.call_args_list) + np.testing.assert_array_equal(result, expected) + assert not lazy.pc.is_loaded + assert actual_rng.integers(100000) == expected_rng.integers(100000) + if return_indices: + assert isinstance(result, tuple) and isinstance(result[0], np.ndarray) + else: + assert isinstance(result, np.ndarray) + + @pytest.mark.parametrize("lazy", [False, True]) + @pytest.mark.parametrize("tiny_fraction", [False, True]) + def test_subsample__point_empty_masks(self, lazy: bool, tiny_fraction: bool) -> None: + """Checks that empty point populations and fractions rounded to zero return empty values and positions.""" + + # Either exclude every point or select one point whose half-sample rounds down to zero + values = np.arange(6, dtype=np.int16) + points = gu.PointCloud.from_xyz(values, np.zeros(6), values, crs=32633) + mask = np.zeros(6, dtype=bool) + subsample = 1.0 + if tiny_fraction: + mask[-1] = True + subsample = 0.5 + expected_sample = points.subsample(subsample, mask=mask, random_state=42) + expected_indices = points.subsample(subsample, return_indices=True, mask=mask, random_state=42) + source = points + if lazy: + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import ( + _register_dask_pointcloud_accessor, + ) + + _register_dask_pointcloud_accessor() + source = dgpd.from_geopandas(points.ds, npartitions=2, sort=False).pc + assert not source.is_loaded + + # Preserve the ordinary value dtype and the one-dimensional positional-index layout + sampled = source.subsample(subsample, mask=mask, random_state=42) + indices = source.subsample(subsample, return_indices=True, mask=mask, random_state=42) + assert sampled.size == 0 and sampled.dtype == values.dtype + assert len(indices) == 1 and indices[0].size == 0 + assert np.issubdtype(indices[0].dtype, np.integer) + assert np.array_equal(sampled, expected_sample) + assert np.array_equal(indices, expected_indices) + if lazy: + assert not source.is_loaded + + @pytest.mark.parametrize("strategy", ["sequential", "topk"]) + @pytest.mark.parametrize("shape", [(30,), (5, 6)]) + def test_sample_valid_indices__boolean_mask( + self, strategy: Literal["sequential", "topk"], shape: tuple[int, ...] + ) -> None: + """Checks that boolean False locations are excluded and topk preserves positions across array chunks.""" + + # Make both accepted and excluded positions occur in every chunk + import_optional("dask") + import dask.array as da + + valid = (np.arange(30) % 4 != 0).reshape(shape) + lazy = da.from_array(valid, chunks=3) + eager = _sample_valid_indices(valid, subsample=5, random_state=42, strategy=strategy) + result = _sample_valid_indices(lazy, subsample=5, random_state=42, strategy=strategy) + + # Sequential grid traversal can differ across layouts; point rows and topk have identical positions + assert isinstance(lazy, da.Array) + assert all(isinstance(index, np.ndarray) for index in result) + assert len(result) == len(shape) + assert len(result[0]) == 5 + assert np.all(valid[result]) + assert len(np.unique(np.ravel_multi_index(result, shape))) == 5 + if strategy == "topk" or len(shape) == 1: + np.testing.assert_array_equal(result, eager) + + @pytest.mark.parametrize("backend", ["numpy", "dask", "multiprocessing"]) + @pytest.mark.parametrize("tiny_fraction", [False, True]) + @pytest.mark.parametrize("strategy", ["sequential", "topk"]) + def test_subsample__raster_empty_masks( + self, backend: str, tiny_fraction: bool, strategy: Literal["sequential", "topk"], tmp_path: Path + ) -> None: + """Checks that empty raster samples have the source value dtype and expected index dimensions.""" + + # Select no cells or one cell whose half-sample rounds to zero, crossing several worker tiles + values = np.arange(24, dtype=np.int16).reshape(4, 6) + raster = gu.Raster.from_array(values, from_origin(0, 4, 1, 1), crs=32633) + mask = np.zeros(raster.shape, dtype=bool) + subsample = 1.0 + if tiny_fraction: + mask[-1, -1] = True + subsample = 0.5 + expected_sample = raster.subsample(subsample, mask=mask, random_state=42, strategy=strategy) + expected_indices = raster.subsample( + subsample, return_indices=True, mask=mask, random_state=42, strategy=strategy + ) + + # Use an unloaded file for multiprocessing and the usual lazy accessor for Dask + source: Any = raster + config = None + if backend != "numpy": + path = tmp_path / "empty_sample.tif" + raster.to_file(path) + if backend == "dask": + import_optional("dask") + source = open_raster(str(path), chunks={"x": 3, "y": 2}).rst + else: + source = gu.Raster(path) + config = MultiprocConfig(chunks=(2, 3)) + + # Return empty samples and two empty coordinate arrays without changing the source value dtype + sampled = source.subsample(subsample, mask=mask, random_state=42, strategy=strategy, mp_config=config) + indices = source.subsample( + subsample, return_indices=True, mask=mask, random_state=42, strategy=strategy, mp_config=config + ) + assert sampled.size == 0 and sampled.dtype == np.dtype(source.dtype) + assert len(indices) == 2 and all(index.size == 0 for index in indices) + assert all(np.issubdtype(index.dtype, np.integer) for index in indices) + assert np.array_equal(sampled, expected_sample) + assert np.array_equal(indices, expected_indices) + if backend == "dask": + da = pytest.importorskip("dask.array") + assert isinstance(source.data, da.Array) and not source._obj._in_memory + if backend == "multiprocessing": + assert not source.is_loaded + + @pytest.mark.parametrize("opened_bands,selected_band", [(None, 2), ([2], 1)]) + def test_subsample__multiproc_selected_band( + self, opened_bands: list[int] | None, selected_band: int, tmp_path: Path + ) -> None: + """Checks that masked multiprocessing samples the requested disk band without changing the source band list.""" + + # Write distinct integer bands with one nodata value in the second band's eligible half + values = np.arange(24, dtype=np.int16).reshape(4, 6) + data = np.ma.array(np.stack([values + 100, values]), mask=False) + data.mask[1, 1, 4] = True + raster = gu.Raster.from_array(data, from_origin(0, 4, 1, 1), crs=32633, nodata=-9999) + path = tmp_path / "multiband_sample.tif" + raster.to_file(path) + + # Selecting an already restricted Raster counts from its available bands, not the original disk bands + source = gu.Raster(path, bands=opened_bands) + original_bands = source.bands + mask = np.indices(source.shape)[1] >= 3 + eligible = mask & ~data.mask[1] + config = MultiprocConfig(chunks=(2, 3)) + + # Select all eligible cells so the expected population is independent of random draws and worker tile order + indices = source.subsample(1, band=selected_band, return_indices=True, mask=mask, mp_config=config) + sampled = source.subsample(1, band=selected_band, mask=mask, mp_config=config) + assert len(indices[0]) == eligible.sum() + assert np.all(eligible[indices]) + np.testing.assert_array_equal(sampled, values[indices]) + assert sampled.dtype == values.dtype + assert source.bands == original_bands and not source.is_loaded + + @pytest.mark.parametrize("strategy", subsample_strategies) + @pytest.mark.parametrize("mask_form", ["array", "masked", "raster", "vector"]) + def test_subsample__masked_backends( + self, + strategy: Literal["sequential", "topk"], + mask_form: str, + lazy_test_files_tiny: list[str], + tmp_path: Path, + ) -> None: + """Checks that masks and finite data jointly define fractional samples across eager, Dask and worker tiles.""" + + import_optional("dask") + import dask.array as da + + from geoutils.multiproc.cluster import MpCluster + + # Prepare a common mask and compute the eligible population from the original raster values + path = lazy_test_files_tiny[0] + raster = gu.Raster(path, load_data=True) + values = raster.data + allowed = np.indices(raster.shape)[1] < raster.width // 2 + mask: Any = allowed.copy() + + # Use equivalent spatial selections, with extra nodata entries for the masked-array case + if mask_form == "masked": + mask = np.ma.array(mask, mask=False) + mask.mask[::3, ::4] = True + allowed &= ~mask.mask + elif mask_form == "raster": + mask = gu.Raster.from_array(mask, raster.transform, raster.crs) + elif mask_form == "vector": + left, bottom, _, top = raster.bounds + middle = left + (raster.width // 2) * raster.res[0] + mask = gpd.GeoDataFrame(geometry=[box(left, bottom, middle, top)], crs=raster.crs) + + # Apply the fraction after excluding source nodata values and mask entries that are nodata or False + eligible = allowed & ~get_mask_from_array(values) + expected_size = int(eligible.sum() * 0.25) + assert expected_size > 0 + + # Sample with different chunk layouts and check exact source grid positions while inputs remain lazy + lazy = open_raster(path, chunks={"x": 9, "y": 7}) + unloaded = gu.Raster(path) + lazy_mask = da.from_array(mask, chunks=(5, 8)) if mask_form == "masked" else mask + worker_mask = mask + + # Leave the boolean mask file unloaded so workers also read its cells by tile + if mask_form == "raster": + mask_path = tmp_path / "sampling_mask.tif" + mask.to_file(mask_path) + worker_mask = gu.Raster(mask_path, is_mask=True) + results = [] + + # Real processes check that both count and selection workers receive the same sliced mask + with MpCluster({"nb_workers": 2, "max_tasks_per_child": None}) as cluster: + config = MultiprocConfig(chunks=(6, 10), cluster=cluster) + sources = [(raster, mask, None), (lazy.rst, lazy_mask, None), (unloaded, worker_mask, config)] + for source, source_mask, source_config in sources: + options: dict[str, Any] = { + "mask": source_mask, + "random_state": 42, + "strategy": strategy, + "mp_config": source_config, + } + indices = source.subsample(0.25, return_indices=True, **options) + sampled = source.subsample(0.25, **options) + + # Compute only the small sample results; the original raster objects are not loaded + indices = tuple( + np.asarray(index.compute() if hasattr(index, "compute") else index) for index in indices + ) + sampled = sampled.compute() if hasattr(sampled, "compute") else sampled + assert len(indices) == 2 and len(indices[0]) == expected_size + assert np.all(eligible[indices]) + assert len(np.unique(np.ravel_multi_index(indices, raster.shape))) == expected_size + np.testing.assert_array_equal(sampled, values[indices]) + assert sampled.dtype == values.dtype + results.append(indices) + + # Check that inputs are not loaded and topk returns the same sample for every chunk layout + assert not lazy._in_memory and isinstance(lazy.data, da.Array) + assert not unloaded.is_loaded + if mask_form == "raster": + assert not worker_mask.is_loaded + if strategy == "topk": + np.testing.assert_array_equal(results[0], results[1]) + np.testing.assert_array_equal(results[0], results[2]) diff --git a/tests/test_sampling/test_support.py b/tests/test_sampling/test_support.py new file mode 100644 index 000000000..fcb30acb4 --- /dev/null +++ b/tests/test_sampling/test_support.py @@ -0,0 +1,392 @@ +"""Tests for defining a common geospatial support and reprojecting point/raster values on it.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import geopandas as gpd +import numpy as np +import pytest +import xarray as xr +from rasterio.transform import from_origin +from shapely.geometry import box + +import geoutils as gu +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.multiproc import MultiprocConfig +from geoutils.sampling.support import ( + _mask_at_support, + _sampling_support, + _values_at_support, +) + + +def _raster(data: NDArrayNum, *, x_origin: float = 0) -> gu.Raster: + """Simplify creating a raster for the tests below.""" + + return gu.Raster.from_array(data, transform=from_origin(x_origin, data.shape[-2], 1, 1), crs=32633, nodata=-99999) + + +class TestSupport: + """ + Test the functions to choose a common support and reproject point/raster on them for eager inputs. + + Dask inputs are tested further below in TestSupportChunked, while Multiproc integration is covered in + test_cosampling.py. + + Here, we specifically test: + - _sampling_support() chooses point locations by default, or an explicitly requested raster grid. + - _values_at_support() aligns rasters and reads raster bands, arrays, or point columns at that support. + - _mask_at_support() places vector, raster, or array masks on raster grids and point locations. + """ + + def test_sampling_support__default_and_explicit(self) -> None: + """Checks that points take precedence by default and an explicit raster takes precedence when requested.""" + + # Create a raster and points on three of its cells + raster = _raster(np.arange(20, dtype=float).reshape(4, 5)) + x, y = raster.ij2xy(np.array([0, 1, 2]), np.array([1, 2, 3])) + points = gu.PointCloud.from_xyz(x, y, np.arange(3, dtype=float), crs=raster.crs) + + # Choose the common locations with and without an explicit request + point_support = _sampling_support((raster, points), None) + raster_support = _sampling_support((points, raster), raster) + + # Selected support should be the original object + assert point_support is points + assert raster_support is raster + + def test_sampling_support__error_array(self) -> None: + """Checks that an array cannot define spatial locations by itself.""" + + # Pass an array without a raster grid or point coordinates + values = np.arange(6).reshape(2, 3) + + # Using it as support should raise an error + with pytest.raises(TypeError, match="must select raster or point cloud support"): + _sampling_support((values,), None) + + def test_values_at_support__raster_grid_alignment(self) -> None: + """Checks that a shifted raster requires alignment and then returns values on the selected grid.""" + + # Shift the value raster by one cell from the selected raster grid + support = _raster(np.zeros((3, 4), dtype=float)) + source_values = np.arange(12, dtype=float).reshape(3, 4) + source = _raster(source_values, x_origin=1) + + # Require explicit permission before moving values to the selected grid + with pytest.raises(ValueError, match="does not share"): + _values_at_support( + source, + None, + input_support=source, + support=support, + support_dataframe=None, + name="source", + interpolation="nearest", + align="raise", + mp_config=None, + ) + result = _values_at_support( + source, + None, + input_support=source, + support=support, + support_dataframe=None, + name="source", + interpolation="nearest", + align="reproject", + mp_config=None, + ) + + # Check the three overlapping columns and the nodata values outside the source grid + expected = np.full(support.shape, np.nan) + expected[:, 1:] = source_values[:, :3] + assert np.array_equal(np.ma.filled(result, np.nan), expected, equal_nan=True) + + def test_values_at_support__selected_raster_band(self) -> None: + """Checks that a raster band selector returns the requested values on the same grid.""" + + # Give both raster bands distinct values on the same grid + base = np.arange(20, dtype=float).reshape(4, 5) + raster = _raster(np.stack((base, base + 100))) + + # Read the second band on the raster support + result = _values_at_support( + raster, + 2, + input_support=raster, + support=raster, + support_dataframe=None, + name="values", + interpolation="nearest", + align="raise", + mp_config=None, + ) + + # The selected values should come from the second band + assert np.array_equal(result, base + 100) + + @pytest.mark.parametrize("input_type", ["numpy", "xarray"]) + def test_values_at_support__raster_array(self, input_type: str) -> None: + """Checks that an array uses its raster input grid and preserves masked cells.""" + + # Create an independent array on the raster grid with one nodata value + raster = _raster(np.ones((4, 5), dtype=float)) + values = np.ma.array(np.arange(20, dtype=float).reshape(4, 5), mask=False) + values.mask[1, 2] = True + source: Any = values + if input_type == "xarray": + source = xr.DataArray(values.filled(np.nan), dims=("y", "x")) + + # Place the values directly on the grid supplied by their raster input + result = _values_at_support( + source, + None, + input_support=raster, + support=raster, + support_dataframe=None, + name="values", + interpolation="nearest", + align="raise", + mp_config=None, + ) + + # Values should be exactly equal + assert np.array_equal(result, values.filled(np.nan), equal_nan=True) + + def test_values_at_support__raster_at_points(self) -> None: + """Checks that raster values are read at the selected point locations.""" + + # Choose point locations at three known raster cells + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + rows, columns = np.array([1, 2, 3]), np.array([1, 2, 3]) + x, y = raster.ij2xy(rows, columns) + support = gu.PointCloud.from_xyz(x, y, np.zeros(3), crs=raster.crs) + + # Interpolate the raster at the point support + result = _values_at_support( + raster, + None, + input_support=raster, + support=support, + support_dataframe=support.ds, + name="values", + interpolation="nearest", + align="raise", + mp_config=None, + ) + + # Extract the same cells directly from the raster to verify the interpolated values + expected = raster.data.filled(np.nan)[rows, columns] + assert np.array_equal(result, expected) + + def test_values_at_support__point_column(self) -> None: + """Checks that a selected point column is returned in the original point order.""" + + # Create point values with duplicate row labels and a separate numeric column + positions = np.arange(5, dtype=float) + points = gu.PointCloud.from_xyz(positions, positions, positions + 10, crs=32633) + points.ds["weight"] = 2 * positions + points.ds.index = ["a", "b", "a", "c", "b"] + + # Read the selected column at the same ordered point locations + result = _values_at_support( + points, + "weight", + input_support=points, + support=points, + support_dataframe=points.ds, + name="weight", + interpolation="nearest", + align="raise", + mp_config=None, + ) + + # The returned array should follow dataframe row order (even when several rows have the same index label) + expected = points.ds["weight"].to_numpy() + assert np.array_equal(result, expected) + + def test_mask_at_support__vector_on_raster(self) -> None: + """Checks that a vector mask properly rasterizes inside its geometry.""" + + # Mask the first two columns of a raster grid with one polygon + raster = _raster(np.arange(20, dtype=float).reshape(4, 5)) + geometry = gpd.GeoDataFrame(geometry=[box(0, 0, 2, 4)], crs=raster.crs) + + # Place the vector mask on the raster support + result = _mask_at_support(geometry, raster) + assert result is not None + + # Check that only cells inside the polygon are masked + expected = np.zeros(raster.shape, dtype=bool) + expected[:, :2] = True + assert np.array_equal(result, expected) + + @pytest.mark.parametrize("mask_mode", ["inside", "outside"]) + def test_mask_at_support__vector_on_points(self, mask_mode: str) -> None: + """Checks that a vector input (without input feature) perform inside/outside geometry masking.""" + + # Place two points inside a polygon and two points outside it + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + rows, columns = np.array([0, 1, 3, 4]), np.array([1, 2, 4, 5]) + x, y = raster.ij2xy(rows, columns) + points = gu.PointCloud.from_xyz(x, y, np.arange(4, dtype=float), crs=raster.crs) + mask = gu.Vector(gpd.GeoDataFrame(geometry=[box(0, 3, 3, 6)], crs=raster.crs)) + + # Create the polygon mask on the point support with inside/outside mode + result = _mask_at_support(mask, points, support_dataframe=points.ds, mask_mode=mask_mode) + assert result is not None + + # Check exact boolean output is as expected + expected = np.array([True, True, False, False]) + if mask_mode == "outside": + expected = ~expected + assert np.array_equal(result, expected) + + def test_mask_at_support__raster_on_points(self) -> None: + """Checks that a raster mask selects points in true cells and excludes points outside its grid.""" + + # Set the first three raster columns to true and add one point beyond the raster + raster = _raster(np.arange(30, dtype=float).reshape(5, 6)) + allowed = np.indices(raster.shape)[1] < 3 + mask = gu.Raster.from_array(allowed, raster.transform, raster.crs) + rows, columns = np.array([0, 1, 3, 4]), np.array([1, 2, 4, 5]) + x, y = raster.ij2xy(rows, columns) + x, y = np.append(x, 100.0), np.append(y, 100.0) + points = gu.PointCloud.from_xyz(x, y, np.arange(5, dtype=float), crs=raster.crs) + + # Read the boolean raster at each point location + result = _mask_at_support(mask, points, support_dataframe=points.ds) + assert result is not None + + # Check true cells, false cells, and the point outside the raster + assert np.array_equal(result, [True, True, False, False, False]) + + def test_mask_at_support__masked_point_array(self) -> None: + """Checks that nodata and false entries in a point mask are both excluded.""" + + # Give one point a nodata mask value and another point a false value + positions = np.arange(5, dtype=float) + points = gu.PointCloud.from_xyz(positions, positions, positions, crs=32633) + mask = np.ma.array([True, True, True, False, True], mask=[False, False, True, False, False]) + + # Place the plain mask array on the ordered point support + result = _mask_at_support(mask, points, support_dataframe=points.ds) + assert result is not None + + # Check that only true, finite mask entries remain selected + assert np.array_equal(result, [True, True, False, False, True]) + + +class TestSupportChunked: + """ + Test module checking that the support helpers respect Dask/MP chunked execution and that their values exactly + match eager results. + """ + + def test_values_at_support__backends(self, tmp_path: Path) -> None: + """Checks that Dask and Multiproc raster alignment exactly matches eager without loading the inputs.""" + + import_optional("dask") + import dask.array as da + + # Create a two-band raster with one nodata cell, and shift the selected support by one column + values = np.arange(30, dtype=float).reshape(5, 6) + source = _raster(np.stack((values + 100, values))) + source.data[1, 2, 3] = np.ma.masked + support = _raster(np.zeros(source.shape, dtype=float), x_origin=1) + + # Calculate the eager reference from the selected second band + expected = _values_at_support( + source, + 2, + input_support=source, + support=support, + support_dataframe=None, + name="values", + interpolation="nearest", + align="reproject", + mp_config=None, + ) + + # Reproject the same values through Dask chunks and Multiproc tiles + lazy_source = source.to_xarray().chunk({"band": 1, "y": 2, "x": 3}) + dask_result = _values_at_support( + lazy_source, + 2, + input_support=lazy_source, + support=support, + support_dataframe=None, + name="values", + interpolation="nearest", + align="reproject", + mp_config=None, + ) + source_file = tmp_path / "values.tif" + source.to_file(source_file) + multiproc_source = gu.Raster(source_file, load_data=False) + outfile = tmp_path / "aligned_values.tif" + multiproc_result = _values_at_support( + multiproc_source, + 2, + input_support=multiproc_source, + support=support, + support_dataframe=None, + name="values", + interpolation="nearest", + align="reproject", + mp_config=MultiprocConfig(chunks=(2, 3), outfile=str(outfile)), + ) + + # Check Dask laziness and the file loading state before comparing every output value + assert isinstance(lazy_source.data, da.Array) + assert isinstance(dask_result, da.Array) + assert not multiproc_source.is_loaded + assert outfile.exists() + expected_values = np.ma.filled(expected, np.nan) + assert np.array_equal(dask_result.compute(), expected_values, equal_nan=True) + assert np.array_equal(np.ma.filled(multiproc_result, np.nan), expected_values, equal_nan=True) + assert isinstance(lazy_source.data, da.Array) + assert not multiproc_source.is_loaded + + def test_mask_at_support__backends(self, tmp_path: Path) -> None: + """Checks that Dask and Multiproc vector masks exactly match eager without loading the support grids.""" + + import_optional("dask") + import dask.array as da + + # Create a raster support and a polygon that covers its first three columns + support = _raster(np.zeros((5, 6), dtype=float)) + mask = gu.Vector(gpd.GeoDataFrame(geometry=[box(0, 0, 3, 5)], crs=support.crs)) + + # Calculate the eager reference, then place the same vector mask through Dask chunks + expected = _mask_at_support(mask, support, align="reproject") + lazy_support = support.to_xarray().chunk({"y": 2, "x": 3}) + dask_result = _mask_at_support(mask, lazy_support.rst, align="reproject") + + # Write the support grid to disk and place the vector mask through Multiproc tiles + support_file = tmp_path / "support.tif" + support.to_file(support_file) + multiproc_support = gu.Raster(support_file, load_data=False) + outfile = tmp_path / "aligned_mask.tif" + multiproc_result = _mask_at_support( + mask, + multiproc_support, + align="reproject", + mp_config=MultiprocConfig(chunks=(2, 3), outfile=str(outfile)), + ) + + # Check Dask laziness and the file loading state before comparing every output mask value + assert expected is not None + assert dask_result is not None and isinstance(dask_result, da.Array) + assert multiproc_result is not None + assert isinstance(lazy_support.data, da.Array) + assert not multiproc_support.is_loaded + assert outfile.exists() + assert np.array_equal(dask_result.compute(), expected) + assert np.array_equal(multiproc_result, expected) + assert isinstance(lazy_support.data, da.Array) + assert not multiproc_support.is_loaded diff --git a/tests/test_stats/test_estimators.py b/tests/test_stats/test_estimators.py index 4b73d4a33..af1c1577e 100644 --- a/tests/test_stats/test_estimators.py +++ b/tests/test_stats/test_estimators.py @@ -82,7 +82,7 @@ def test_rmse(self) -> None: assert rmse_data == pytest.approx(3.0276503540974917) def test_sum_square(self) -> None: - """Test Sum Square functionality runs on any type of input""" + """Test sum square functionality runs on any type of input""" test_data = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) # Test masked arrays with invalid data (should ignore NaNs/masked values) diff --git a/tests/test_stats/test_grouping.py b/tests/test_stats/test_grouping.py new file mode 100644 index 000000000..e0e837e26 --- /dev/null +++ b/tests/test_stats/test_grouping.py @@ -0,0 +1,1297 @@ +"""Tests for grouping statistics by continuous and categorical variables.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +from affine import Affine +from shapely.geometry import box + +import geoutils as gu +from geoutils._misc import import_optional +from geoutils.multiproc import MultiprocConfig + + +class TestGroupedStats: + """ + Test module for grouped statistics, i.e. stats(by=), with eager inputs. + + Dask and Multiproc inputs are covered in TestGroupedStatsChunked further below. + + The tests cover group definitions, sampling, vector zones, returned tables and masks, plotting, and using the + optional Flox backend. + """ + + def test_stats__interval_counts_and_masks(self) -> None: + """ + Checks that group masks include selected locations with nodata values while counts exclude them. + + Contrary to cosample() that wants to sample the common valid locations, for stats we want to keep + nodata as it is valuable information for the user, and is described directly by valid/total counts. + + """ + + # Create two values with different nodata patterns and a user mask with excluded points + values = { + "first": np.array([1.0, 2.0, np.nan, 4.0, 5.0, 6.0]), + "second": np.arange(6, dtype=float), + } + grouper = np.arange(6, dtype=float) + user_mask = np.array([True, False, True, True, True, True]) + + # Group values into two explicit intervals, and return their final mask + table, masks = gu.stats.stats( + values, + "median", + by={"slope": grouper}, + bins={"slope": [0, 3, 6]}, + mask=user_mask, + return_masks=True, + ) + + # Check interval labels, separate value counts, and returned mask keys + assert isinstance(table.index, pd.IntervalIndex) + assert list(table.columns.names) == ["value", "statistic"] + assert table[("first", "count")].tolist() == [1, 3] + assert table[("second", "count")].tolist() == [2, 3] + assert isinstance(masks, Mapping) + assert list(masks) == list(table.index) + + # Check that masks split locations allowed by the user mask without removing selected nodata values + group_masks = [np.asarray(masks[key]) for key in masks] + assert [int(np.count_nonzero(group_mask)) for group_mask in group_masks] == [2, 3] + assert np.array_equal(np.logical_or.reduce(group_masks), user_mask) + assert not np.any(group_masks[0] & group_masks[1]) + + def test_stats__combines_categories_and_empty_groups(self) -> None: + """Checks that two grouping variables return labels in the input order, and include empty combinations.""" + + # Create numeric bins and named categories with one combination absent in the data + values = np.arange(6, dtype=float) + continuous = np.array([0, 0, 1, 1, 2, 2], dtype=float) + categorical = np.array(["forest", "forest", "forest", "ice", "ice", "ice"]) + + # Request every defined category and bin combination, including the empty one + table = gu.stats.stats( + values, + by={"elevation": continuous, "surface": categorical}, + bins={"elevation": [0, 1, 2, 3]}, + categories={"surface": ["forest", "ice"]}, + statistics="mean", + observed=False, + ) + + # Check index types, order, and the empty group's zero count and NaN mean + assert isinstance(table.index, pd.MultiIndex) + assert isinstance(table.index.levels[0], pd.IntervalIndex) + assert isinstance(table.index.levels[1], pd.CategoricalIndex) + assert table.index.levels[1].ordered + assert len(table) == 6 + assert table.loc[(pd.Interval(0, 1, closed="left"), "ice"), ("value", "count")] == 0 + assert np.isnan(table.loc[(pd.Interval(0, 1, closed="left"), "ice"), ("value", "mean")]) + + def test_stats__respects_interval_closure_and_nonfinite_values(self) -> None: + """Checks that right-closed intervals and selected nodata values affect groups and counts separately.""" + + # Define right-closed intervals with one value below all intervals and one infinite selected value + intervals = pd.IntervalIndex.from_breaks([0, 1, 2], closed="right", name="distance") + + # Calculate all statistics and return the complete interval masks + table, masks = gu.stats.stats( + np.array([100.0, 1.0, np.inf]), + by={"distance": np.array([0.0, 1.0, 2.0])}, + bins={"distance": intervals}, + statistics="all", + return_masks=True, + ) + + # Check the interval edge rules, finite counts, total counts, means, and locations in each group mask + assert table.index.equals(intervals) + assert table[("value", "count")].tolist() == [1, 0] + assert table[("value", "totalcount")].tolist() == [1, 1] + assert table.loc[intervals[0], ("value", "mean")] == 1 + assert np.isnan(table.loc[intervals[1], ("value", "mean")]) + assert not np.asarray(masks[intervals[0]])[0] + + def test_stats__subsampling_does_not_change_masks(self) -> None: + """Checks that subsampling limits statistic counts without modifying returned group masks.""" + + # Split twenty values into two intervals + values = np.arange(20, dtype=float) + groups = np.arange(20, dtype=float) + # Calculate statistics from six sampled locations and request complete masks + table, masks = gu.stats.stats( + values, + by={"distance": groups}, + bins={"distance": [0, 10, 20]}, + statistics="mean", + subsample=6, + random_state=42, + return_masks=True, + ) + + # Check the sampled statistic count and the full twenty-location mask count + assert int(table[("value", "count")].sum()) == 6 + assert sum(int(np.count_nonzero(masks[key])) for key in masks) == 20 + + def test_stats__raster_mask_type(self, tmp_path: Path) -> None: + """Checks that raster group masks have the same grid and boolean type when written and reopened.""" + + # Create a georeferenced raster and split its cells into two numeric intervals + transform = Affine(10, 0, 100, 0, -10, 200) + raster = gu.Raster.from_array(np.arange(1, 7, dtype=float).reshape(2, 3), transform, 32631) + grouper = np.arange(6, dtype=float).reshape(2, 3) + table, masks = raster.stats( + "mean", + by={"slope": grouper}, + bins={"slope": [0, 3, 6]}, + return_masks=True, + ) + + # Check that the first returned mask is a boolean Raster on the source grid + first_mask = masks[table.index[0]] + assert isinstance(first_mask, gu.Raster) + assert first_mask.is_mask + assert first_mask.georeferenced_grid_equal(raster) + # Write the first group mask to disk, reopen it, and compare its type and values + output_path = tmp_path / "group_mask.tif" + first_mask.to_file(output_path) + reopened = gu.Raster(output_path, is_mask=True, load_data=True) + assert reopened.is_mask + assert np.array_equal(reopened.data, first_mask.data) + + @pytest.mark.parametrize("source_type", ["raster", "xarray", "pointcloud", "dataframe"]) + def test_stats__mask_types(self, source_type: str) -> None: + """Checks that every eager spatial input returns a boolean mask on its common support.""" + + # Create the same six values and two groups as either raster cells or points + values = np.arange(1, 7, dtype=float) + groups = np.arange(6) % 2 + if source_type in {"raster", "xarray"}: + values = values.reshape(2, 3) + groups = groups.reshape(2, 3) + raster = gu.Raster.from_array(values, Affine(1, 0, 0, 0, -1, 2), 32631) + source: Any = raster if source_type == "raster" else raster.to_xarray().rst + else: + pointcloud = gu.PointCloud.from_xyz(np.arange(6), np.zeros(6), values, crs=32631) + source = pointcloud if source_type == "pointcloud" else pointcloud.ds.pc + + # Request the complete mask for each group + table, masks = source.stats("mean", by={"zone": groups}, categories={"zone": [0, 1]}, return_masks=True) + mask = masks[table.index[0]] + + # Check that the mask has the matching spatial type, common support, and group values + if source_type == "raster": + assert isinstance(mask, gu.Raster) + assert mask.is_mask + assert mask.georeferenced_grid_equal(raster) + mask_values = mask.data + elif source_type == "xarray": + assert mask.rst.is_mask + assert mask.rst.georeferenced_grid_equal(raster) + mask_values = mask.data + else: + assert isinstance(mask, gu.PointCloud if source_type == "pointcloud" else gpd.GeoDataFrame) + interface = mask if source_type == "pointcloud" else mask.pc + assert interface.is_mask + assert interface.georeferenced_coords_equal(pointcloud) + mask_values = interface.data + assert np.array_equal(np.asarray(mask_values).reshape(groups.shape), groups == 0) + + @pytest.mark.parametrize( + "by,bins,categories,message", + [ + ({}, None, None, "at least one named grouper"), + ({"zone": np.zeros((2, 2))}, {"missing": 2}, None, "do not match"), + ({"zone": np.zeros((2, 2))}, {"zone": 2}, {"zone": [0]}, "cannot define both"), + ({"zone": np.zeros((2, 2))}, {"zone": [0, 0, 1]}, None, "strictly increasing"), + ({"zone": np.zeros((2, 2))}, None, {"zone": [0, 0]}, "unique"), + ], + ) + def test_stats__error_invalid_group_definitions( + self, by: Any, bins: Any, categories: Any, message: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Checks that invalid group definitions fail before an unloaded raster reads any data.""" + + # Write a raster file to disk with four valid cells, then make any attempt to load all cells fail + raster = gu.Raster.from_array(np.ones((2, 2)), Affine(1, 0, 0, 0, -1, 2), 32631) + path = tmp_path / "group_validation.tif" + raster.to_file(path) + source = gu.Raster(path) + + def fail_load(*args: Any, **kwargs: Any) -> None: + """Reject data access before group definition validation finishes.""" + raise AssertionError("Group definitions must be checked before reading values.") + + monkeypatch.setattr(source, "load", fail_load) + + # Raise the group definition error without consulting the source's values + with pytest.raises(ValueError, match=message): + source.stats("mean", by=by, bins=bins, categories=categories) + assert not source.is_loaded + + @pytest.mark.parametrize("source_type", ["raster", "xarray", "pointcloud", "geopandas"]) + def test_stats__subsample_per_group_spatial(self, source_type: str) -> None: + """Checks that raster and point cloud stats() include observed groups that receive a zero-size sample.""" + + # Include a single-location group whose quarter sample rounds down to zero + values = np.arange(12, dtype=float).reshape(3, 4) + groups = np.array([0] * 7 + [1] * 4 + [2]).reshape(values.shape) + if source_type in {"raster", "xarray"}: + source = gu.Raster.from_array(values, Affine(1, 0, 0, 0, -1, 3), 32631) + if source_type == "xarray": + source = source.to_xarray().rst + else: + source = gu.PointCloud.from_xyz(np.arange(12), np.zeros(12), values.ravel(), crs=32631) + groups = groups.ravel() + if source_type == "geopandas": + source = source.ds.pc + + # Leave observed=True so an unsampled group must be distinguished from an absent category + table = source.stats( + "mean", + by={"category": groups}, + categories={"category": [0, 1, 2, 3]}, + subsample=0.25, + subsample_per_group=True, + random_state=42, + ) + + # Check that all three original groups are present, with no estimate for the unsampled one + assert table.index.tolist() == [0, 1, 2] + assert np.array_equal(table.xs("count", level="statistic", axis=1).iloc[:, 0], [1, 1, 0]) + assert np.isnan(table.xs("mean", level="statistic", axis=1).iloc[2, 0]) + + @pytest.mark.parametrize("subsampling_strategy", ["topk", "sequential"]) + def test_stats__separate_sampling_and_reduction(self, subsampling_strategy: str) -> None: + """Checks that the sampling option stays separate from the grouped calculation strategy.""" + + # Sample ten values while requesting masks for both complete groups + values = np.arange(100, dtype=float) + result, masks = gu.stats.stats( + values, + by={"zone": values % 2 == 0}, + statistics="mean", + subsample=10, + random_state=42, + strategy="sparse", + subsampling_strategy=subsampling_strategy, + return_masks=True, + ) + + # Check the sampled count, full masks, and recorded sampling option + assert result[("value", "count")].sum() == 10 + assert sum(np.count_nonzero(mask) for mask in masks.values()) == 100 + assert result.attrs["grouped_stats"]["subsampling_strategy"] == subsampling_strategy + + def test_stats__vector_union_and_feature_ids(self) -> None: + """Checks that a vector alone creates inside/outside groups while feature IDs create separate zones.""" + + # Create two one-cell vector features separated by uncovered raster cells + raster = gu.Raster.from_array(np.arange(1, 9, dtype=float).reshape(2, 4), Affine(1, 0, 0, 0, -1, 2), 32631) + zones = gu.Vector( + gpd.GeoDataFrame({"id": ["first", "second"]}, geometry=[box(0, 1, 1, 2), box(3, 0, 4, 1)], crs=32631) + ) + + # Group once by all vector coverage and once by each feature name + union = raster.stats("mean", by={"inside": zones}) + features = raster.stats("mean", by={"zone": (zones, "id")}) + + # Check that coverage includes outside cells while named features include only their own cells + assert union[("band_1", "count")].tolist() == [6, 2] + assert features[("band_1", "count")].tolist() == [1, 1] + assert features[("band_1", "mean")].tolist() == [1, 8] + + @pytest.mark.parametrize("source_type", ["pointcloud", "dataframe"]) + def test_stats__geometry_z_is_unchanged_by_masks(self, source_type: str) -> None: + """Checks that masks of geometry Z values add a boolean column without discarding original elevations.""" + + # Existing attribute names must remain unchanged when the default new mask-column name is already taken + points: Any = gu.PointCloud.from_xyz(np.arange(3), np.zeros(3), np.arange(3) + 100, crs=32631, use_z=True) + points.ds["group_mask"] = [10, 11, 12] + original = points.ds.copy() + if source_type == "dataframe": + points = points.ds.pc + + # Make the separate boolean mask column active without changing any geometry Z elevation + _, masks = points.stats("mean", by={"zone": np.array([True, False, True])}, return_masks=True) + result = masks[True] + output = result.ds if source_type == "pointcloud" else result + interface = result if source_type == "pointcloud" else result.pc + assert interface.data_column == "_group_mask" + assert interface.is_mask + pd.testing.assert_series_equal(output.geometry, original.geometry) + pd.testing.assert_series_equal(output.group_mask, original.group_mask) + assert np.array_equal(output._group_mask, [True, False, True]) + + def test_plot_grouped_stats__one_and_two_dimensions(self) -> None: + """Checks that plotting one or two grouping variables creates the expected panels.""" + + # Load the optional plotting package only for this plotting test + pytest.importorskip("matplotlib") + import matplotlib.pyplot as plt + + # Plot one grouping variable and check its count and statistic panels + one_dimensional = gu.stats.stats( + np.arange(6, dtype=float), + by={"x": np.arange(6, dtype=float)}, + bins={"x": [0, 3, 6]}, + statistics="mean", + ) + axes_1d = gu.stats.plot_grouped_stats(one_dimensional, statistic="mean") + assert set(axes_1d) == {"count", "statistic"} + + # Plot two grouping variables and check row counts, column counts, statistic, and color scale + two_dimensional = gu.stats.stats( + np.arange(6, dtype=float), + by={"x": np.array([0, 0, 1, 1, 2, 2]), "surface": np.array(["a", "b", "a", "b", "a", "b"])}, + bins={"x": [0, 1, 2, 3]}, + categories={"surface": ["a", "b"]}, + statistics="mean", + ) + axes_2d = gu.stats.plot_grouped_stats(two_dimensional, statistic="mean") + assert set(axes_2d) == {"count_x", "count_y", "statistic", "colorbar"} + plt.close("all") + + @pytest.mark.parametrize("kind", ["integer", "boolean", "string"]) + def test_stats__masked_values_and_categories(self, kind: str) -> None: + """Checks that masked values and masked category labels are excluded independently.""" + + # Mask one selected value and a different group label for three category types + values = np.ma.array([1, 2, 3, 4, 5, 6], mask=[False, True, False, False, False, False]) + group_values = { + "integer": [0, 0, 0, 1, 1, 1], + "boolean": [False, False, False, True, True, True], + "string": ["a", "a", "a", "b", "b", "b"], + } + groups = np.ma.array(group_values[kind], mask=[False, False, True, False, False, False]) + categories = {"integer": {"group": [0, 1]}, "boolean": None, "string": {"group": ["a", "b", "N/A"]}} + + # Calculate group means and request masks for the remaining category locations + table, masks = gu.stats.stats( + values, + by={"group": groups}, + categories=categories[kind], + statistics="mean", + return_masks=True, + ) + + # Check value counts, means, and the sizes of complete group masks separately + assert table[("value", "count")].tolist() == [1, 3] + assert table[("value", "mean")].tolist() == [1, 5] + assert [int(np.count_nonzero(masks[key])) for key in masks] == [2, 3] + + def test_raster_stats__masked_integer_data_and_boolean_mask(self) -> None: + """Checks that raster value masks and boolean user masks affect counts and group masks separately.""" + + # Mask one integer value and exclude two different cells through a boolean Raster mask + data = np.ma.array([[1, 2, 3], [4, 5, 6]], mask=[[False, True, False], [False, False, False]]) + raster = gu.Raster.from_array(data, Affine(1, 0, 0, 0, -1, 2), 32631, nodata=-9999) + mask = raster.from_array( + np.ma.array([[True, True, True], [True, False, True]], mask=[[False, False, True], [False, False, False]]), + raster.transform, + raster.crs, + ) + + # Calculate one group mean and request all locations allowed by the user mask + table, masks = raster.stats( + "mean", + by={"group": np.zeros(data.shape, dtype=int)}, + categories={"group": [0]}, + mask=mask, + return_masks=True, + ) + + # Check three available values in four group locations + assert table[("band_1", "count")].tolist() == [3] + assert table[("band_1", "mean")].tolist() == [pytest.approx(11 / 3)] + assert int(np.count_nonzero(masks[0].data)) == 4 + + def test_stats__flox_categories(self) -> None: + """Checks that Flox returns the same ordered category counts, means and standard deviations as GeoUtils.""" + + # Create three categories in a different order from their declaration, with one selected nodata value + pytest.importorskip("flox") + values = np.array([1.0, 2.0, np.nan, 4.0, 5.0, 6.0]) + groups = np.array([2, 2, 0, 0, 1, 1]) + options = { + "statistics": ["mean", "std"], + "by": {"zone": groups}, + "categories": {"zone": [2, 0, 1, 3]}, + "observed": False, + } + + # Calculate the same complete table with the built-in and Flox reducers + expected = gu.stats.stats(values, backend="geoutils", **options) + result = gu.stats.stats(values, backend="flox", **options) + + # Check the declared row order, empty category and every reduced value + pd.testing.assert_frame_equal(result, expected) + assert result.attrs["grouped_stats"]["strategy"] == "flox" + + def test_stats__flox_bins_mask_and_subsample(self) -> None: + """Checks that Flox matches GeoUtils bin boundaries, masking and global subsampling.""" + + # Combine interval bins and categories, then exclude some locations before a reproducible global sample + pytest.importorskip("flox") + values = { + "first": np.arange(24, dtype=float), + "second": np.arange(24, dtype=float) * 2, + } + values["first"][5] = np.nan + distance = np.arange(24, dtype=float) % 6 + surface = np.array(["ice", "rock"] * 12) + options = { + "statistics": ["mean", "sum", "totalcount", "percentagevalidpoints"], + "by": {"distance": distance, "surface": surface}, + "bins": {"distance": [0, 2, 4, 6]}, + "categories": {"surface": ["rock", "ice", "water"]}, + "mask": np.arange(24) % 5 != 0, + "subsample": 9, + "random_state": 42, + "observed": False, + } + + # Compare every sampled value and empty declared group with the built-in reducer + expected = gu.stats.stats(values, backend="geoutils", **options) + result = gu.stats.stats(values, backend="flox", **options) + pd.testing.assert_frame_equal(result, expected) + + def test_stats__error_flox_options(self) -> None: + """Checks that Flox rejects options and statistics that its grouped path cannot reproduce.""" + + # Use one ordinary category input so each call reaches Flox-specific validation + pytest.importorskip("flox") + values = np.arange(6, dtype=float) + grouping = {"by": {"zone": np.arange(6) % 2}, "categories": {"zone": [0, 1]}} + + # Reject global statistics, group masks, sampling within groups, multiprocessing and GeoUtils strategies + with pytest.raises(ValueError, match="requires grouped statistics"): + gu.stats.stats(values, "mean", backend="flox") + for options in ( + {"return_masks": True}, + {"subsample_per_group": True}, + {"mp_config": MultiprocConfig(chunks=2)}, + {"strategy": "dense"}, + ): + with pytest.raises(ValueError, match="Flox backend requires"): + gu.stats.stats(values, "mean", backend="flox", **grouping, **options) + with pytest.raises(ValueError, match="does not support"): + gu.stats.stats(values, "nmad", backend="flox", **grouping) + + def test_raster_stats__flox_loading_warning(self) -> None: + """Checks that a Raster input warns that the Flox backend loads its values.""" + + # Create a small Raster whose cells belong to two boolean categories + pytest.importorskip("flox") + raster = gu.Raster.from_array(np.arange(6, dtype=float).reshape(2, 3), Affine.identity(), 32631) + groups = np.arange(6).reshape(2, 3) % 2 == 0 + + # Calculate the grouped mean and check the warning is raised before spatial values are selected + with pytest.warns(UserWarning, match="loads Raster and PointCloud inputs"): + result = raster.stats("mean", by={"zone": groups}, backend="flox") + assert result[("band_1", "count")].tolist() == [3, 3] + + +class TestGroupedStatsChunked: + """ + Tests grouped statistics from stats(by=) with Dask and Multiproc inputs. + + The tests compare group definitions, common support, sampling, vector zones, returned masks, and the optional Flox + backend with eager results. Numerical reduction strategies are covered in test_reduction.py. + """ + + @pytest.mark.parametrize("strategy", ["dense", "sparse", "groupwise"]) + def test_stats__empty_selection(self, strategy: str) -> None: + """Checks that fully masked chunks return an empty table and mask mapping.""" + + # Mask every location while requesting both boolean groups in the result + da = pytest.importorskip("dask.array") + values = da.ones((5, 6), chunks=2) + table, masks = gu.stats.stats( + values, + by={"zone": np.ones((5, 6), dtype=bool)}, + mask=np.zeros((5, 6), dtype=bool), + statistics="mean", + strategy=strategy, + return_masks=True, + ) + eager_table, _ = gu.stats.stats( + np.ones((5, 6)), + by={"zone": np.ones((5, 6), dtype=bool)}, + mask=np.zeros((5, 6), dtype=bool), + statistics="mean", + strategy=strategy, + return_masks=True, + ) + + # Check the empty table columns and mask mapping + assert table.empty + pd.testing.assert_frame_equal(table, eager_table, check_exact=True) + assert list(table.columns) == [("value", "count"), ("value", "mean")] + assert len(masks) == 0 + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + def test_stats__empty_arrays(self, backend: str) -> None: + """Checks that empty inputs include requested groups with zero counts.""" + + # Calculate statistics for one requested group from empty Eager, Dask, or Multiproc inputs + values = np.empty(0) + expected = gu.stats.stats( + values, + by={"zone": np.empty(0)}, + categories={"zone": [0]}, + statistics=["mean", "totalcount", "validcount"], + observed=False, + ) + config = MultiprocConfig(chunks=4) if backend == "multiproc" else None + if backend == "dask": + da = pytest.importorskip("dask.array") + values = da.from_array(values, chunks=4) + result = gu.stats.stats( + values, + by={"zone": np.empty(0)}, + categories={"zone": [0]}, + statistics=["mean", "totalcount", "validcount"], + observed=False, + mp_config=config, + ) + pd.testing.assert_frame_equal(result, expected, check_exact=True) + + # Check zero counts and a NaN mean for the empty group + assert result.loc[0, ("value", "count")] == 0 + assert result.loc[0, ("value", "totalcount")] == 0 + assert result.loc[0, ("value", "validcount")] == 0 + assert np.isnan(result.loc[0, ("value", "mean")]) + + def test_stats__dask_lazy_masks_match_eager(self) -> None: + """Checks that the Dask backend matches eager calculation and creates group masks only when requested.""" + + # Calculate an eager reference from three numeric intervals + da = pytest.importorskip("dask.array") + values = np.arange(12, dtype=float).reshape(3, 4) + grouper = np.arange(12, dtype=float).reshape(3, 4) + expected = gu.stats.stats(values, by={"x": grouper}, bins={"x": [0, 4, 8, 12]}, statistics="mean") + + # Repeat with different Dask chunks for values and group labels + table, masks = gu.stats.stats( + da.from_array(values, chunks=(2, 2)), + by={"x": da.from_array(grouper, chunks=(1, 4))}, + bins={"x": [0, 4, 8, 12]}, + statistics="mean", + return_masks=True, + ) + + # Check the complete table and load only the first returned mask + pd.testing.assert_frame_equal(table, expected) + first_mask = masks[next(iter(masks))] + assert isinstance(first_mask, da.Array) + assert int(first_mask.sum().compute()) == 4 + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + @pytest.mark.parametrize("kind", ["boolean", "categorical"]) + @pytest.mark.parametrize("bins", [[-0.5, 1.5], 1]) + def test_stats__explicit_bins_override_category_dtype(self, backend: str, kind: str, bins: Any) -> None: + """Checks that explicit numeric bins take precedence over boolean and Pandas categorical group types.""" + + # Put both distinct labels in one numeric interval, with an unused Pandas category available for inference + groups: Any = np.array([False, True, False, True]) + if kind == "categorical": + groups = pd.Series(pd.Categorical([0, 1, 0, 1], categories=[0, 1, 2])) + expected = gu.stats.stats( + np.arange(4, dtype=float), by={"group": groups}, bins={"group": bins}, statistics="mean" + ) + config = MultiprocConfig(chunks=3) if backend == "multiproc" else None + if backend == "dask": + import_optional("dask") + import dask.array as da + import dask.dataframe as dd + + groups = dd.from_pandas(groups, npartitions=2) if kind == "categorical" else da.from_array(groups, chunks=3) + + # Check that explicit bins produce one interval for every input type and existing category definition + table = gu.stats.stats( + np.arange(4, dtype=float), by={"group": groups}, bins={"group": bins}, statistics="mean", mp_config=config + ) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + assert isinstance(table.index, pd.IntervalIndex) + assert table[("value", "count")].tolist() == [4] + assert table[("value", "mean")].tolist() == [1.5] + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + def test_stats__explicit_categories_override_existing_order(self, backend: str) -> None: + """Checks that explicit categories replace the existing order and preserve requested absent categories.""" + + # Exclude one existing label from the requested categories and put a different absent label first + groups: Any = pd.Series(pd.Categorical(["west", "east", "west"], categories=["west", "east", "north"])) + categories = ["south", "west"] + expected = gu.stats.stats( + np.array([1, 2, 3]), + by={"zone": groups}, + categories={"zone": categories}, + statistics="sum", + observed=False, + ) + config = MultiprocConfig(chunks=2) if backend == "multiproc" else None + if backend == "dask": + import_optional("dask") + import dask.dataframe as dd + + groups = dd.from_pandas(groups, npartitions=2) + + # Pass the requested categories as an iterator and check that it is consumed only once + table = gu.stats.stats( + np.array([1, 2, 3]), + by={"zone": groups}, + categories={"zone": iter(categories)}, + statistics="sum", + observed=False, + mp_config=config, + ) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + assert table.index.tolist() == ["south", "west"] + assert table[("value", "count")].tolist() == [0, 2] + assert table[("value", "sum")].iloc[1] == 4 + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + @pytest.mark.parametrize("masked", [False, True]) + def test_stats__masked_integer_labels_are_exact(self, backend: str, masked: bool) -> None: + """Checks that masked integer categories above float precision stay distinct and exclude masked labels.""" + + # Adjacent integers above 2**53 collapse if a nodata label first promotes the array to float + labels = [2**53, 2**53 + 1] + groups: Any = np.ma.array([labels[0], labels[1], labels[1]], mask=[False, False, masked]) + expected, expected_masks = gu.stats.stats( + np.array([1, 2, 4]), + by={"zone": groups}, + categories={"zone": labels}, + statistics="sum", + return_masks=True, + ) + config = MultiprocConfig(chunks=2) if backend == "multiproc" else None + if backend == "dask": + import_optional("dask") + import dask.array as da + + groups = da.from_array(groups, chunks=2) + + # Read both the independent category counts and all locations in each group + table, masks = gu.stats.stats( + np.array([1, 2, 4]), + by={"zone": groups}, + categories={"zone": labels}, + statistics="sum", + return_masks=True, + mp_config=config, + ) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + for label in labels: + assert np.array_equal(np.asarray(masks[label]), np.asarray(expected_masks[label])) + assert table.index.tolist() == labels + assert table[("value", "count")].tolist() == [1, 1 if masked else 2] + assert table[("value", "sum")].tolist() == [1, 2 if masked else 6] + assert np.array_equal(np.asarray(masks[labels[1]]), [False, True, not masked]) + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + def test_stats__tuple_categories_are_single_labels(self, backend: str) -> None: + """Checks that hashable tuple categories form one categorical level instead of a Pandas MultiIndex.""" + + # Assign tuples into a one-dimensional object array so each pair is treated as one input label + groups: Any = np.empty(3, dtype=object) + groups[:] = [(1, 2), (3, 4), (1, 2)] + labels = [(3, 4), (1, 2), (5, 6)] + expected, expected_masks = gu.stats.stats( + np.array([1, 2, 3]), + by={"zone": groups}, + categories={"zone": labels}, + statistics="sum", + observed=False, + return_masks=True, + ) + config = MultiprocConfig(chunks=2) if backend == "multiproc" else None + if backend == "dask": + import_optional("dask") + import dask.array as da + + groups = da.from_array(groups, chunks=2) + + # Request one absent tuple and check that the specified order also works for mask lookup + table, masks = gu.stats.stats( + np.array([1, 2, 3]), + by={"zone": groups}, + categories={"zone": labels}, + statistics="sum", + observed=False, + return_masks=True, + mp_config=config, + ) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + for label in masks: + assert np.array_equal(np.asarray(masks[label]), np.asarray(expected_masks[label])) + assert isinstance(table.index, pd.CategoricalIndex) + assert table.index.tolist() == labels + assert table[("value", "count")].tolist() == [1, 2, 0] + assert np.array_equal(np.asarray(masks[(1, 2)]), [True, False, True]) + + @pytest.mark.parametrize("source_type", ["pointcloud", "dataframe", "dask"]) + def test_stats__point_column_existing_categories(self, source_type: str) -> None: + """Checks that point columns return categories in their existing order, including absent categories.""" + + # Select an ordered category column with an absent category preceding the observed labels + dataframe = gpd.GeoDataFrame( + { + "height": [1, 2, 3], + "zone": pd.Categorical(["west", "east", "west"], categories=["north", "east", "west"]), + }, + geometry=gpd.points_from_xy(np.arange(3), np.zeros(3)), + crs=32631, + ) + points: Any = gu.PointCloud(dataframe, data_column="height") + expected = points.stats("mean", by={"zone": "zone"}, observed=False) + if source_type == "dataframe": + points = points.ds.pc + elif source_type == "dask": + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import ( + _register_dask_pointcloud_accessor, + ) + + _register_dask_pointcloud_accessor() + points = dgpd.from_geopandas(points.ds, npartitions=2, sort=False).pc + points.data_column = "height" + + # Infer the same definition from the named column as from its original Pandas categorical values + table = points.stats("mean", by={"zone": "zone"}, observed=False) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + if source_type == "dask": + assert not points.is_loaded + assert table.index.tolist() == ["north", "east", "west"] + assert table[("height", "count")].tolist() == [0, 1, 2] + np.testing.assert_allclose(table[("height", "mean")], [np.nan, 2, 2], equal_nan=True) + + def test_stats__known_dask_categories(self) -> None: + """Checks that known Dask categories return their specified order and absent groups like Pandas categories.""" + + # Declare an unused category so discovering only observed labels would lose part of the requested groups + import_optional("dask") + import dask.dataframe as dd + + values = np.array([1.0, 3.0, 5.0, 7.0]) + categories = ["low", "high", "unused"] + groups = pd.Series(pd.Categorical(["low", "low", "high", "high"], categories=categories, ordered=True)) + lazy_groups = dd.from_pandas(groups, npartitions=2) + + # Infer the same category order from the Pandas and Dask grouping variables + expected = gu.stats.stats(values, "mean", by={"zone": groups}, observed=False) + result = gu.stats.stats(values, "mean", by={"zone": lazy_groups}, observed=False) + + # Check that the unused category has zero count and a NaN estimate in the original category order + pd.testing.assert_frame_equal(result, expected) + assert result.index.tolist() == categories + expected_values = np.array([[2, 2], [2, 6], [0, np.nan]], dtype=float) + np.testing.assert_allclose(result["value"], expected_values, equal_nan=True) + + def test_stats__external_values_on_common_support(self, tmp_path: Path) -> None: + """Checks that Dask point inputs use the same common support and return the same result as eager points.""" + + # Place points at known raster cells and give two polygons different numeric values + raster = gu.Raster.from_array(np.arange(16, dtype=float).reshape(4, 4), Affine(1, 0, 0, 0, -1, 4), 32631) + x, y = raster.ij2xy([0, 1, 2, 3], [0, 0, 3, 3]) + frame = gpd.GeoDataFrame( + {"height": [100.0, np.nan, 102.0, 103.0], "zone": ["west", "west", "east", "east"]}, + geometry=gpd.points_from_xy(x, y), + crs=32631, + ) + points = gu.PointCloud(frame, data_column="height") + features = gu.Vector( + gpd.GeoDataFrame( + {"weight": [2.0, 4.0]}, + geometry=[box(0, 0, 2, 4), box(2, 0, 4, 4)], + crs=32631, + ) + ) + expected_table, expected_masks = raster.stats( + by={"zone": (points, "zone")}, + categories={"zone": ["west", "east"]}, + values={"raster": raster, "points": points, "weight": (features, "weight")}, + statistics="mean", + interpolation="nearest", + return_masks=True, + ) + + # Write a point file to disk with height and zone columns, then reopen it as Dask partitions + pytest.importorskip("dask_geopandas") + filename = tmp_path / "points.gpkg" + points.ds.to_file(filename) + points = gu.open_pointcloud(str(filename), chunks=2, data_column="height").pc + assert not points.is_loaded + + # Let the external point data define the common support for all selected values and grouping variables + table, masks = raster.stats( + by={"zone": (points, "zone")}, + categories={"zone": ["west", "east"]}, + values={"raster": raster, "points": points, "weight": (features, "weight")}, + statistics="mean", + interpolation="nearest", + return_masks=True, + ) + pd.testing.assert_frame_equal(table, expected_table, check_exact=True) + + # Check each mean and count separately, then check the complete western point mask + np.testing.assert_allclose(table[("raster", "mean")], [2.0, 13.0]) + np.testing.assert_allclose(table[("points", "mean")], [100.0, 102.5]) + np.testing.assert_allclose(table[("weight", "mean")], [2.0, 4.0]) + assert table[("points", "count")].tolist() == [1, 2] + mask = masks["west"] + mask_values = mask.pc.data + dd = pytest.importorskip("dask.dataframe") + assert isinstance(mask_values, dd.Series) + assert not mask.pc.is_loaded + mask_values = mask_values.compute() + assert np.array_equal(mask_values, [True, True, False, False]) + assert np.array_equal(mask_values, expected_masks["west"].data) + assert not points.is_loaded and not mask.pc.is_loaded + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + def test_stats__vector_numeric_bins_and_values(self, backend: str) -> None: + """Checks that a numeric vector column can define bins and also appear as a selected value.""" + + # Give two separated polygons numeric values and leave one raster column uncovered + raster = gu.Raster.from_array(np.arange(12, dtype=float).reshape(3, 4), Affine(1, 0, 0, 0, -1, 3), 32631) + features = gu.Vector( + gpd.GeoDataFrame( + {"slope": [5.0, 15.0]}, + geometry=[box(0, 0, 1, 3), box(2, 0, 4, 3)], + crs=32631, + ) + ) + expected = raster.stats( + by={"slope": (features, "slope")}, + bins={"slope": [0, 10, 20]}, + values={"raster": 1, "slope": (features, "slope")}, + statistics="mean", + ) + if backend == "dask": + raster = raster.to_xarray().chunk({"x": 2, "y": 2}).rst + + # Use one Multiproc tile size to read vector values, assign group IDs, and calculate statistics + config = MultiprocConfig(chunks=(2, 3)) if backend == "multiproc" else None + + # Use the vector column as numeric bins and as a selected output value + table = raster.stats( + by={"slope": (features, "slope")}, + bins={"slope": [0, 10, 20]}, + values={"raster": 1, "slope": (features, "slope")}, + statistics="mean", + mp_config=config, + ) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + + # Check group counts and both means only where a polygon covers the raster + assert table[("raster", "count")].tolist() == [3, 6] + np.testing.assert_allclose(table[("raster", "mean")], [4.0, 6.5]) + np.testing.assert_allclose(table[("slope", "mean")], [5.0, 15.0]) + + def test_stats__vector_projection_shared_by_point_partitions(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Checks that vectors are reprojected once and later features provide values where polygons overlap.""" + + # Place points inside each polygon, inside their overlap, and outside both, away from polygon boundaries + features = gpd.GeoDataFrame( + {"weight": [10.0, 20.0]}, geometry=[box(-1, -1, 2, 1), box(1, -1, 3, 1)], crs=4326, index=[9, 9] + ) + dataframe = gpd.GeoDataFrame( + {"height": np.ones(4)}, + geometry=gpd.points_from_xy([-0.5, 1.5, 2.5, 4.0], np.zeros(4)), + crs=4326, + index=[7, 7, 2, 2], + ).to_crs(3857) + + expected = dataframe.pc.stats( + "mean", + values={"weight": (features, "weight")}, + by={"point": np.arange(4)}, + categories={"point": range(4)}, + observed=False, + ) + + # Divide the projected points into Dask partitions without changing duplicate labels or row order + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import ( + _register_dask_pointcloud_accessor, + ) + + _register_dask_pointcloud_accessor() + points = dgpd.from_geopandas(dataframe, npartitions=3, sort=False).pc + assert not points.is_loaded + + # Count feature projections so each partition cannot repeat the same coordinate conversion + projections = [] + to_crs = gpd.GeoDataFrame.to_crs + + def record_projection(dataframe: gpd.GeoDataFrame, *args: Any, **kwargs: Any) -> gpd.GeoDataFrame: + """Record each feature projection while applying the usual coordinate transformation.""" + + projections.append(dataframe.crs) + return to_crs(dataframe, *args, **kwargs) + + monkeypatch.setattr(gpd.GeoDataFrame, "to_crs", record_projection) + + # Use one group per input row to check point order independently of repeated dataframe labels + table = points.stats( + "mean", + values={"weight": (features, "weight")}, + by={"point": np.arange(4)}, + categories={"point": range(4)}, + observed=False, + ) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + + # The later feature supplies the overlap value, while the uncovered point has no finite observation + assert projections == [features.crs] + np.testing.assert_allclose(table[("weight", "mean")], [10.0, 20.0, 20.0, np.nan], equal_nan=True) + assert np.array_equal(table[("weight", "count")], [1, 1, 1, 0]) + assert not points.is_loaded + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + @pytest.mark.parametrize("sampling_strategy", ["topk", "sequential"]) + @pytest.mark.parametrize("subsample", [0.25, 3, 1]) + def test_stats__subsample_per_group(self, backend: str, sampling_strategy: str, subsample: int | float) -> None: + """Checks that each combined group receives its own sample size while masks include every eligible location.""" + + # 1/ Prepare unequal groups, one absent category/bin combination, and values with different finite counts + groups = np.repeat(np.arange(5), [20, 14, 8, 5, 1]).reshape(6, 8) + position = np.arange(groups.size, dtype=float).reshape(groups.shape) + missing = position.copy() + missing[groups == 1] = np.nan + values = {"position": position, "doubled": 2 * position + 1, "missing": missing} + by = {"category": groups // 2, "height": groups % 2} + keep = np.ones(groups.shape, dtype=bool) + keep.flat[[0, 22]] = False + expected_table, expected_masks = gu.stats.stats( + values, + ["mean", "totalcount"], + by=by, + categories={"category": [0, 1, 2]}, + bins={"height": [-0.5, 0.5, 1.5]}, + mask=keep, + subsample=subsample, + subsample_per_group=True, + subsampling_strategy=sampling_strategy, + random_state=42, + observed=False, + return_masks=True, + ) + + # Cross group boundaries with uneven chunks, including differently chunked Dask grouping variables + config = None + if backend == "dask": + import_optional("dask") + import dask.array as da + + values = {name: da.from_array(array, chunks=(2, 3)) for name, array in values.items()} + by = {name: da.from_array(array, chunks=(3, 2)) for name, array in by.items()} + elif backend == "multiproc": + config = MultiprocConfig(chunks=(2, 3)) + + # 2/ Apply the sample size to each category and bin combination and request its complete group mask + table, masks = gu.stats.stats( + values, + ["mean", "totalcount"], + by=by, + categories={"category": [0, 1, 2]}, + bins={"height": [-0.5, 0.5, 1.5]}, + mask=keep, + subsample=subsample, + subsample_per_group=True, + subsampling_strategy=sampling_strategy, + random_state=42, + observed=False, + return_masks=True, + mp_config=config, + ) + if sampling_strategy == "topk" or subsample == 1: + pd.testing.assert_frame_equal(table, expected_table, check_exact=True) + for key in masks: + assert np.array_equal(np.asarray(masks[key]), np.asarray(expected_masks[key])) + + # 3/ Check sample sizes independently from group sizes, including the one-location and absent groups + for number, key in enumerate(table.index): + members = (groups == number) & keep + size = np.count_nonzero(members) + expected_count = int(size * subsample) if subsample <= 1 else min(int(subsample), size) + assert table.loc[key, ("position", "count")] == expected_count + assert table.loc[key, ("doubled", "count")] == expected_count + assert table.loc[key, ("missing", "totalcount")] == expected_count + assert np.array_equal(np.asarray(masks[key]), members) + + # Selected nodata values do not change the sample size or locations used for other columns + assert table.loc[key, ("missing", "count")] == (0 if number == 1 else expected_count) + if expected_count: + expected_mean = 2 * table.loc[key, ("position", "mean")] + 1 + assert table.loc[key, ("doubled", "mean")] == pytest.approx(expected_mean) + else: + assert np.isnan(table.loc[key, ("position", "mean")]) + assert table.attrs["grouped_stats"]["subsample_per_group"] is True + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + @pytest.mark.parametrize("rows", [1, 2]) + def test_stats__vector_integer_labels_outside_coverage(self, backend: str, rows: int) -> None: + """Checks that adjacent large vector labels remain distinct when some cells fall outside every feature.""" + + # Use integer labels that cannot both be represented as floats and leave the final raster column uncovered + labels = [2**53, 2**53 + 1] + zones = gpd.GeoDataFrame({"label": labels}, geometry=[box(0, 0, 1, 2), box(1, 0, 2, 2)], crs=32631) + values = np.array([[1, 2, 4], [5, 6, 8]])[:rows] + source: Any = gu.Raster.from_array(values, Affine(1, 0, 0, 0, -1, 2), 32631) + expected_table, expected_masks = source.stats("sum", by={"zone": (zones, "label")}, return_masks=True) + config = MultiprocConfig(chunks=(1, 2)) if backend == "multiproc" else None + if backend == "dask": + import_optional("dask") + import dask.array as da + + source = source.to_xarray().chunk({"y": 1, "x": 2}).rst + + # Infer ordered feature categories and check that uncovered cells contribute to neither group + table, masks = source.stats("sum", by={"zone": (zones, "label")}, return_masks=True, mp_config=config) + pd.testing.assert_frame_equal(table, expected_table, check_exact=True) + assert table.index.tolist() == labels + assert table[("band_1", "count")].tolist() == [rows, rows] + assert table[("band_1", "sum")].tolist() == values[:, :2].sum(axis=0).tolist() + mask = masks[labels[1]] + mask_values = mask.rst.data if backend == "dask" else mask.data + expected_mask = expected_masks[labels[1]].data + if backend == "dask": + assert isinstance(mask_values, da.Array) + mask_values = mask_values.compute() + assert np.array_equal(np.asarray(mask_values), np.asarray(expected_mask)) + assert np.array_equal(np.asarray(mask_values).reshape(-1), np.tile([False, True, False], rows)) + + @pytest.mark.parametrize("grouping", ["categories", "bins"]) + def test_stats__vector_groups_file_points_not_loaded(self, grouping: str, tmp_path: Path) -> None: + """Checks that vector categories and bins match eager while a point file remains unloaded.""" + + # Write a point file to disk with four values and the final point outside both vector features + dataframe = gpd.GeoDataFrame( + {"height": [1, 2, 3, 4]}, geometry=gpd.points_from_xy(np.arange(4) + 0.5, np.full(4, 0.5)), crs=32631 + ) + path = tmp_path / "vector_group_points.gpkg" + dataframe.to_file(path) + source = gu.PointCloud(path, data_column="height") + zones = gpd.GeoDataFrame( + {"label": ["west", "east"], "number": [10, 20]}, + geometry=[box(0, 0, 2, 1), box(2, 0, 3, 1)], + crs=32631, + ) + column = "label" if grouping == "categories" else "number" + bins = None if grouping == "categories" else {"zone": [5, 15, 25]} + expected = gu.PointCloud(dataframe, data_column="height").stats("sum", by={"zone": (zones, column)}, bins=bins) + + # Read coordinates and values in two-point blocks when assigning vector groups + table = source.stats("sum", by={"zone": (zones, column)}, bins=bins, mp_config=MultiprocConfig(chunks=2)) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + + # The first feature contains two points, the second contains one, and the source stays unloaded + assert not source.is_loaded + assert table[("height", "count")].tolist() == [2, 1] + assert table[("height", "sum")].tolist() == [3, 3] + + def test_stats__vector_coverage_includes_point_boundaries(self, tmp_path: Path) -> None: + """Checks that vector coverage groups include boundary points for eager and unloaded point values.""" + + # Write a point file to disk with one point on the feature edge, one inside, and one outside + dataframe = gpd.GeoDataFrame( + {"height": [1, 2, 4]}, geometry=gpd.points_from_xy([0, 0.5, 2], [0.5, 0.5, 0.5]), crs=32631 + ) + path = tmp_path / "coverage_points.gpkg" + dataframe.to_file(path) + source = gu.PointCloud(path, data_column="height") + outline = gu.Vector(gpd.GeoDataFrame(geometry=[box(0, 0, 1, 1)], crs=32631)) + expected = gu.PointCloud(dataframe, data_column="height").stats("sum", by={"covered": outline}) + + # Group points by vector coverage and include the point exactly on the polygon boundary + table = source.stats("sum", by={"covered": outline}, mp_config=MultiprocConfig(chunks=2)) + pd.testing.assert_frame_equal(table, expected, check_exact=True) + assert table.index.tolist() == [False, True] + assert table[("height", "count")].tolist() == [1, 2] + assert table[("height", "sum")].tolist() == [4, 3] + assert not source.is_loaded + + @pytest.mark.parametrize("source_type", ["raster", "xarray", "dask", "multiproc", "pointcloud", "geopandas"]) + @pytest.mark.parametrize("vector_object", [False, True]) + def test_stats__feature_zones(self, source_type: str, vector_object: bool) -> None: + """Checks that vector feature names define zones with separate nodata counts and complete masks.""" + + # Build repeated named zones, a gap between features, and one zone outside the data + elevation = np.arange(1, 17, dtype=float).reshape(4, 4) + other = elevation + 100 + elevation[0, 0], other[2, 3] = np.nan, np.nan + zones = gpd.GeoDataFrame( + {"id": ["west", "east", "west", "empty"]}, + geometry=[box(0, 0, 1, 2), box(3, 0, 4, 4), box(0, 2, 1, 4), box(10, 10, 11, 11)], + crs=32631, + ) + vector = gu.Vector(zones) if vector_object else zones + + # Store the same two value arrays as raster bands, Xarray data, or point columns + selected_values: dict[str, int] | dict[str, str] + if source_type in {"raster", "xarray", "dask", "multiproc"}: + source = gu.Raster.from_array(np.stack((elevation, other)), Affine(1, 0, 0, 0, -1, 4), 32631, nodata=np.nan) + selected_values = {"elevation": 1, "other": 2} + expected_table, _ = source.stats( + by={"zone": (vector, "id")}, + values=selected_values, + statistics=("mean", "min", "max"), + observed=False, + return_masks=True, + ) + if source_type in {"xarray", "dask"}: + xarray_source = source.to_xarray() + if source_type == "dask": + da = pytest.importorskip("dask.array") + xarray_source = xarray_source.chunk({"band": 1, "y": 3, "x": 2}) + source = xarray_source.rst + else: + x, y = np.meshgrid(np.arange(4) + 0.5, 3.5 - np.arange(4)) + frame = gpd.GeoDataFrame( + {"elevation": elevation.ravel(), "other": other.ravel()}, + geometry=gpd.points_from_xy(x.ravel(), y.ravel()), + crs=32631, + ) + source = gu.PointCloud(frame, data_column="elevation") + selected_values = {"elevation": "elevation", "other": "other"} + expected_table, _ = source.stats( + by={"zone": (vector, "id")}, + values=selected_values, + statistics=("mean", "min", "max"), + observed=False, + return_masks=True, + ) + if source_type == "geopandas": + source = source.ds.pc + + # Use uneven Multiproc tiles so vector placement and calculation cross tile boundaries + config = MultiprocConfig(chunks=(3, 2)) if source_type == "multiproc" else None + + # Group by the vector names, include the outside zone, and request group masks + table, masks = source.stats( + by={"zone": (vector, "id")}, + values=selected_values, + statistics=("mean", "min", "max"), + observed=False, + return_masks=True, + mp_config=config, + ) + pd.testing.assert_frame_equal(table, expected_table, check_exact=True) + assert isinstance(table.index, pd.CategoricalIndex) + assert table.index.tolist() == ["west", "east", "empty"] + + # Check each value's nodata separately and leave the gap outside every group + assert table[("elevation", "count")].tolist() == [3, 4, 0] + assert table[("other", "count")].tolist() == [4, 3, 0] + np.testing.assert_allclose(table[("elevation", "mean")], [9, 10, np.nan], equal_nan=True) + np.testing.assert_allclose(table[("elevation", "min")], [5, 4, np.nan], equal_nan=True) + np.testing.assert_allclose(table[("elevation", "max")], [13, 16, np.nan], equal_nan=True) + + # Check the full western zone mask, including its location with a selected nodata value + expected = np.zeros((4, 4), dtype=bool) + expected[:, 0] = True + west_mask = masks["west"] + mask_data = west_mask.pc.data if source_type == "geopandas" else west_mask.data + assert np.array_equal(np.asarray(mask_data).reshape(4, 4), expected) + if source_type == "dask": + assert isinstance(source.data, da.Array) + assert isinstance(mask_data, da.Array) + + @pytest.mark.parametrize("source_type", ["pointcloud", "dataframe", "dask"]) + @pytest.mark.parametrize("size", [1, 3]) + def test_stats__point_mask_rows_and_geometry(self, source_type: str, size: int) -> None: + """Checks that point masks have the same 3D geometry and row order, including singleton inputs.""" + + # Use duplicate row labels and unrelated attributes so assigning by index would change the result + dataframe = gpd.GeoDataFrame( + {"height": np.arange(size), "attribute": np.arange(size) + 10}, + geometry=gpd.points_from_xy(np.arange(size), np.zeros(size), np.arange(size) + 100), + crs=32631, + index=np.zeros(size, dtype=int), + ) + dataframe.attrs["data_column"] = "height" + if source_type == "pointcloud": + with pytest.warns(UserWarning, match="Overriding 3D points"): + points: Any = gu.PointCloud(dataframe, data_column="height") + elif source_type == "dataframe": + points = dataframe.pc + else: + dgpd = import_optional("dask_geopandas", package_name="dask-geopandas") + from geoutils.pointcloud.pd_accessor import ( + _register_dask_pointcloud_accessor, + ) + + _register_dask_pointcloud_accessor() + points = dgpd.from_geopandas(dataframe, npartitions=min(size, 2), sort=False).pc + points.data_column = "height" + + # Request complete masks and replace only the active point data column with boolean mask values + table, masks = points.stats( + "mean", by={"zone": np.arange(size) % 2}, categories={"zone": [0, 1]}, return_masks=True + ) + result = masks[0] + output = result.ds if source_type == "pointcloud" else result + if source_type == "dask": + assert output.height.dtype == bool + output = output.compute() + + # Group zero contains every second original row, with all coordinates and other attributes unchanged + pd.testing.assert_series_equal(output.geometry, dataframe.geometry) + pd.testing.assert_series_equal(output.attribute, dataframe.attribute) + assert np.array_equal(output.height, np.arange(size) % 2 == 0) + assert output.height.dtype == bool + assert table[("height", "count")].sum() == size + + def test_stats__flox_matches_eager(self) -> None: + """Checks that Flox computes a Pandas table equal to eager GeoUtils output while inputs remain Dask arrays.""" + + # Create two interleaved groupers and give the selected values a different Dask chunk layout + pytest.importorskip("flox") + da = pytest.importorskip("dask.array") + values = np.arange(48, dtype=float).reshape(6, 8) + values[0, 0] = np.nan + rows, columns = np.indices(values.shape) + row_groups = rows % 2 + column_groups = columns % 3 + keep = columns != 1 + options = { + "statistics": ["mean", "std", "sum", "sumofsquares", "rmse", "totalcount"], + "by": {"row": row_groups, "column": column_groups}, + "categories": {"row": [1, 0, 2], "column": [2, 0, 1, 3]}, + "mask": keep, + "observed": False, + } + expected = gu.stats.stats(values, backend="geoutils", **options) + lazy_values = da.from_array(values, chunks=(2, 4)) + lazy_groups = { + "row": da.from_array(row_groups, chunks=(3, 2)), + "column": da.from_array(column_groups, chunks=(3, 2)), + } + + # Run the Dask Flox backend, which computes only the small grouped result table + result = gu.stats.stats( + lazy_values, + statistics=options["statistics"], + by=lazy_groups, + categories=options["categories"], + mask=da.from_array(keep, chunks=(3, 2)), + observed=False, + backend="flox", + ) + + # Check input laziness, computed output type and every result against the eager calculation + assert isinstance(lazy_values, da.Array) + assert all(isinstance(grouper, da.Array) for grouper in lazy_groups.values()) + assert isinstance(result, pd.DataFrame) + pd.testing.assert_frame_equal(result, expected) diff --git a/tests/test_stats/test_reduction.py b/tests/test_stats/test_reduction.py new file mode 100644 index 000000000..7f64f328d --- /dev/null +++ b/tests/test_stats/test_reduction.py @@ -0,0 +1,773 @@ +"""Tests shared global and grouped statistic reductions.""" + +from __future__ import annotations + +import warnings +from functools import partial +from pathlib import Path +from typing import Any, cast + +import numpy as np +import pandas as pd +import pytest +from affine import Affine + +import geoutils as gu +from geoutils._misc import import_optional +from geoutils._typing import NDArrayNum +from geoutils.multiproc import MultiprocConfig +from geoutils.multiproc.readers import _ValueReader +from geoutils.stats.reduction import ( + _aggregate_collected, + _collect_block, + _finalize_blocks, + _merge_blocks, + _normalize_statistics, + _reduce_block, + _reduce_global_values, + _reduce_values, + _resolve_strategy, + _Statistics, + _statistics_dask, +) + + +class TestReduction: + """ + Tests statistic requests and reductions of eager arrays. + + We check statistic names and aliases, global and grouped calculations, complete-group statistics, nodata values, + and integer calculations. + """ + + def test_normalize_statistics__names_and_aliases(self) -> None: + """Checks that strings, partial functions and callable objects get clear names and internal aliases.""" + + # Build a callable example + class Span: + def __call__(self, values: Any) -> Any: + return np.nanmax(values) - np.nanmin(values) + + # Include a synonym, a partial without its own name, and a callable object + percentile = partial(np.nanpercentile, q=75) + span = Span() + statistics = _normalize_statistics(["standard deviation", percentile, span], grouped=True) + + # Check the original requests, output names, aliases used by reducers, and mandatory grouped count + assert isinstance(statistics, _Statistics) + assert statistics.requested == ["standard deviation", percentile, span] + assert statistics.names == ["standard deviation", "nanpercentile", "Span"] + assert statistics.aliases == ["std", None, None] + assert statistics.output_names == ["count", "standard deviation", "nanpercentile", "Span"] + + # A single global statistic returns its own value without a grouped count column + summary = _normalize_statistics("mean", grouped=False) + assert summary.single + assert summary.output_names == ["mean"] + + def test_normalize_statistics__defaults(self) -> None: + """Checks that default global and grouped requests contain their expected statistics and display names.""" + + # Normalize the default request for global output, grouped output, and a masked global calculation + summary = _normalize_statistics(None, grouped=False) + grouped = _normalize_statistics(None, grouped=True) + masked = _normalize_statistics("all", grouped=False, masked=True) + + # Global defaults use display names, while grouped columns keep the names requested by stats() + assert summary.names == [ + "Min", + "Max", + "Mean", + "Standard deviation", + "Valid count", + "Total count", + "Percentage valid points", + ] + assert grouped.output_names == [ + "count", + "min", + "max", + "mean", + "std", + "totalcount", + "percentagevalidpoints", + ] + assert masked.aliases[-4:] == [ + "validinliercount", + "totalinliercount", + "percentagevalidinlierpoints", + "percentageinlierpoints", + ] + + def test_normalize_statistics__error_names(self) -> None: + """Checks that ambiguous, reserved, unknown and invalid statistic requests are rejected.""" + + # Two partials have the same output name even though their percentile arguments differ + percentiles = [partial(np.nanpercentile, q=25), partial(np.nanpercentile, q=75)] + with pytest.raises(ValueError, match="unique"): + _normalize_statistics(percentiles) + + # Grouped output reserves count for a callable and rejects names without a known reducer + def count(values: Any) -> int: + """Return the input size under the reserved grouped output name.""" + + return len(values) + + with pytest.raises(ValueError, match="reserved"): + _normalize_statistics([count]) + with pytest.raises(ValueError, match="Unknown statistic names"): + _normalize_statistics(["made_up"]) + with pytest.raises(ValueError, match="cannot be combined"): + _normalize_statistics(["all", "mean"]) + + # Every request must be either a recognized name or a callable function + with pytest.raises(TypeError, match="names or callable"): + _normalize_statistics(cast(Any, [object()])) + + def test_reduce_values__global(self) -> None: + """Checks basic global statistics for multiple eager arrays with separate nodata values.""" + + # Give two values different nodata locations on the same four positions + first = np.array([1.0, np.nan, 3.0, 4.0]) + second = np.array([10.0, 20.0, np.nan, 40.0]) + statistics = _normalize_statistics( + ["mean", "std", "sum", "validcount", "totalcount", "percentagevalidpoints"], grouped=False + ) + + # Reduce the complete arrays as one implicit group + table, strategy = _reduce_values([first, second], statistics) + + # Check each value independently against NumPy and the known input size + assert strategy == "dense" + assert table.index.tolist() == [0] + for index, values in enumerate([first, second]): + finite = values[np.isfinite(values)] + assert table.loc[0, (index, "count")] == finite.size + assert table.loc[0, (index, "mean")] == pytest.approx(finite.mean()) + assert table.loc[0, (index, "std")] == pytest.approx(finite.std()) + assert table.loc[0, (index, "sum")] == pytest.approx(finite.sum()) + assert table.loc[0, (index, "validcount")] == finite.size + assert table.loc[0, (index, "totalcount")] == values.size + assert table.loc[0, (index, "percentagevalidpoints")] == 75 + + def test_reduce_values__groups(self) -> None: + """Checks grouped eager statistics from integer group IDs, including excluded and nodata locations.""" + + # Use three groups, one excluded position, and a different nodata location in each selected value + group_ids = np.array([[0, 1, 0, 1], [2, 2, -1, 1]]) + first = np.array([[1.0, 2.0, np.nan, 4.0], [5.0, 7.0, 8.0, 10.0]]) + second = np.array([[10.0, np.nan, 30.0, 40.0], [50.0, 70.0, 80.0, 100.0]]) + statistics = _normalize_statistics( + ["mean", "std", "sum", "min", "max", "rmse", "validcount", "totalcount", "percentagevalidpoints"] + ) + + # Calculate all groups directly from their integer IDs + table, strategy = _reduce_values( + [first, second], statistics, group_ids=group_ids, total_groups=4, strategy="auto" + ) + + # Compare every observed group and value with the corresponding NumPy calculation + assert strategy == "dense" + assert table.index.tolist() == [0, 1, 2] + for index, values in enumerate([first, second]): + for group_id in table.index: + members = values[group_ids == group_id] + finite = members[np.isfinite(members)] + expected = [ + finite.size, + finite.mean(), + finite.std(), + finite.sum(), + finite.min(), + finite.max(), + np.sqrt(np.mean(finite**2)), + finite.size, + members.size, + 100 * finite.size / members.size, + ] + np.testing.assert_allclose(table.loc[group_id, index], expected) + + def test_reduce_values__complete_group_values(self) -> None: + """Checks that medians, NMAD and custom functions receive every value from each eager group.""" + + # Interleave groups and include nodata so complete group collection and total size are both visible + values = np.arange(18, dtype=float) + values[::5] = np.nan + group_ids = np.arange(values.size) % 3 + statistics = _normalize_statistics(["median", "nmad", np.size]) + + # Exact statistics use the complete values from each group + table, strategy = _reduce_values([values], statistics, group_ids=group_ids, total_groups=3, strategy="auto") + + # Check all results directly from the original groups + assert strategy == "groupwise" + for group_id in range(3): + members = values[group_ids == group_id] + median = np.nanmedian(members) + expected = [ + np.isfinite(members).sum(), + median, + 1.4826 * np.nanmedian(np.abs(members - median)), + members.size, + ] + np.testing.assert_allclose(table.loc[group_id, 0], expected) + + @pytest.mark.parametrize("masked", [False, True]) + def test_reduce_values__integer_squares(self, masked: bool) -> None: + """Checks that eager sums of squares and RMSE do not overflow the input integer type.""" + + # Squared elevations exceed int16, with one optional location excluded by a NumPy mask + values = np.array([[2000, 3000], [5000, 4000]], dtype=np.int16) + source: Any = np.ma.array(values, mask=[[False, True], [False, False]]) if masked else values + reference = source.astype(np.float64) + statistics = _normalize_statistics(["sumofsquares", "rmse"], grouped=False) + + # Compare the reducer with arithmetic performed after conversion to floating point + table, _ = _reduce_values([source], statistics) + assert table.loc[0, (0, "sumofsquares")] == pytest.approx(np.ma.sum(reference**2)) + assert table.loc[0, (0, "rmse")] == pytest.approx(np.sqrt(np.ma.mean(reference**2))) + + @pytest.mark.parametrize("strategy", ["dense", "sparse", "groupwise"]) + @pytest.mark.parametrize("unsigned", [False, True]) + def test_reduce_values__large_integer_extrema(self, strategy: str, unsigned: bool) -> None: + """Checks that eager grouped extrema remain exact beyond floating-point integer precision.""" + + # Neighboring values above 2**53 would become indistinguishable in a floating-point intermediate array + dtype = np.uint64 if unsigned else np.int64 + offset = 2**63 + 10 if unsigned else 2**60 + values = np.array([offset + step for step in [1, 3, 5, 7, 9, 11]], dtype=dtype) + source = np.ma.array(values, mask=[False, True, False, False, True, True]) + group_ids = np.array([0, 0, 1, 1, 2, 2]) + names = ["min", "max", "median" if strategy == "groupwise" else "mean"] + statistics = _normalize_statistics(names) + + # Reduce two populated groups and one group containing only nodata values + table, _ = _reduce_values([source], statistics, group_ids=group_ids, total_groups=4, strategy=strategy) + + # Compare extrema as integers so floating-point rounding cannot hide an incorrect result + assert int(table.loc[0, (0, "min")]) == int(values[0]) + assert int(table.loc[0, (0, "max")]) == int(values[0]) + assert int(table.loc[1, (0, "min")]) == int(values[2]) + assert int(table.loc[1, (0, "max")]) == int(values[3]) + assert pd.isna(table.loc[2, (0, "min")]) + assert pd.isna(table.loc[2, (0, "max")]) + assert np.array_equal(table[(0, "count")], [1, 2, 0]) + + @pytest.mark.parametrize("nodata", ["empty", "nan", "masked"]) + def test_reduce_values__empty_sums(self, nodata: str) -> None: + """Checks that eager sums and squared sums are undefined when no valid values remain.""" + + # Cover an empty array, explicit NaNs, and ordinary values excluded by a NumPy mask + source: Any = np.empty(0) if nodata == "empty" else np.full(4, np.nan) + if nodata == "masked": + source = np.ma.array(np.arange(4), mask=True) + statistics = _normalize_statistics(["sum", "sumofsquares", "validcount"], grouped=False) + + # A zero valid count must stay distinct from a genuine sum of zero + with pytest.warns(UserWarning, match="Empty raster"): + table, _ = _reduce_values([source], statistics) + assert np.isnan(table.loc[0, (0, "sum")]) + assert np.isnan(table.loc[0, (0, "sumofsquares")]) + assert table.loc[0, (0, "validcount")] == 0 + + def test_reduce_global_values__output_forms(self) -> None: + """Checks that global reduction restores scalar and named output for one or several selected values.""" + + # Prepare the same unmasked arrays in the form returned by global selection + first = np.array([1.0, 2.0, np.nan, 4.0]) + second = first + 10 + selected = {"first": (first, None), "second": (second, None)} + + # One statistic returns scalar values, while a list returns a dictionary for each selected value + single = _reduce_global_values( + {"first": selected["first"]}, _normalize_statistics("mean", grouped=False), strategy="auto", mp_config=None + ) + multiple = _reduce_global_values( + selected, + _normalize_statistics(["mean", "validcount", "totalcount"], grouped=False), + strategy="auto", + mp_config=None, + ) + + # Check both established global output forms + assert single == pytest.approx(7 / 3) + assert multiple == { + "first": {"mean": pytest.approx(7 / 3), "validcount": 3, "totalcount": 4}, + "second": {"mean": pytest.approx(37 / 3), "validcount": 3, "totalcount": 4}, + } + + +class TestReductionChunked: + """ + Tests reductions split across Dask chunks or Multiproc tiles. + + Backend results are compared exactly with eager calculations, while Dask inputs remain lazy and file inputs remain + unloaded. Additional tests cover the calculation strategies and their numerical edge cases: + - "auto" selects another strategy from the requested statistics and number of groups. + - "dense" includes every possible group in the summary from each chunk. + - "sparse" includes only the groups found in each chunk. + - "groupwise" gathers all values from each group before calculating its statistics. + """ + + @pytest.mark.parametrize("strategy", ["auto", "dense", "sparse", "groupwise"]) + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + def test_reduce_values__backend_equality(self, strategy: str, backend: str) -> None: + """Checks that every chunk strategy matches eager results for Dask and Multiproc inputs.""" + + # Spread three groups and separate nodata patterns across a rectangular array + shape = (9, 10) + group_ids: Any = (np.arange(np.prod(shape)) % 3).reshape(shape) + first = np.arange(np.prod(shape), dtype=float).reshape(shape) + second = first * 2 + 10 + first.flat[::7] = np.nan + second.flat[::11] = np.inf + group_ids.flat[::13] = -1 + values: list[Any] = [first, second] + statistics = _normalize_statistics(["mean", "std", "sum", "min", "max", "rmse", "totalcount"]) + expected, _ = _reduce_values(values, statistics, group_ids=group_ids, total_groups=4, strategy=strategy) + + # Store the same values in Dask chunks or ask Multiproc to read NumPy tiles + config = None + if backend == "dask": + da = pytest.importorskip("dask.array") + values = [da.from_array(first, chunks=(2, 4)), da.from_array(second, chunks=(3, 2))] + group_ids = da.from_array(group_ids, chunks=(3, 3)) + else: + config = MultiprocConfig(chunks=(2, 4)) + + # Check the computed result and confirm that Dask inputs are still lazy collections + result, resolved = _reduce_values( + values, + statistics, + group_ids=group_ids, + total_groups=4, + strategy=strategy, + mp_config=config, + ) + assert isinstance(result, pd.DataFrame) + pd.testing.assert_frame_equal(result, expected) + assert resolved == ("dense" if strategy == "auto" else strategy) + if backend == "dask": + assert all(isinstance(value, da.Array) for value in values) + assert isinstance(group_ids, da.Array) + + def test_reduce_values__multiproc_reader_loading(self, tmp_path: Path) -> None: + """Checks that Multiproc reduces raster blocks without loading the complete source.""" + + # Write a raster file to disk with nodata values spread across several reduction tiles + values = np.arange(35, dtype=float).reshape(5, 7) + values[1, 2] = np.nan + values[3, 5] = np.nan + path = tmp_path / "reduction.tif" + gu.Raster.from_array(values, Affine(1, 0, 0, 0, -1, 5), 32631, nodata=np.nan).to_file(path) + source = gu.Raster(path) + reader = _ValueReader(source, selector=1) + group_ids = np.indices(values.shape).sum(axis=0) % 3 + statistics = _normalize_statistics(["mean", "std", "sum", "min", "max"]) + expected, _ = _reduce_values([values], statistics, group_ids=group_ids, total_groups=3) + + # Read only raster windows in worker processes and return a computed table + from geoutils.multiproc.cluster import MpCluster + + with MpCluster({"nb_workers": 2}) as cluster: + result, _ = _reduce_values( + [reader], + statistics, + group_ids=group_ids, + total_groups=3, + mp_config=MultiprocConfig(chunks=(2, 3), cluster=cluster), + ) + + # Compare every result while the original Raster still has no loaded data + pd.testing.assert_frame_equal(result, expected) + assert not source.is_loaded + + @pytest.mark.parametrize("dense", [True, False]) + def test_reduce_block__dense_and_sparse(self, dense: bool) -> None: + """Checks that block summaries merge into the same table as a complete eager reduction.""" + + # Divide two values and three observed group IDs between two blocks + group_ids = np.array([0, 2, -1, 0, 1, 2, 1, 2]) + first = np.array([1.0, 2.0, 30.0, np.nan, 5.0, 7.0, 9.0, 11.0]) + second = np.array([10.0, np.nan, 300.0, 40.0, 50.0, 70.0, 90.0, 110.0]) + statistics = _normalize_statistics(["mean", "std", "sum", "min", "max", "rmse", "validcount", "totalcount"]) + aliases = {alias for alias in statistics.aliases if alias is not None} + + # Summarize each block, merge their small arrays, and construct the final table + summaries = [ + _reduce_block([first[block], second[block]], group_ids[block], 4, dense, aliases) + for block in [slice(0, 4), slice(4, 8)] + ] + summary = _merge_blocks(summaries) + result = _finalize_blocks(summary, statistics) + expected, _ = _reduce_values([first, second], statistics, group_ids=group_ids, total_groups=4, strategy="dense") + + # Dense blocks reserve every group; sparse blocks store only the IDs present in each block + expected_labels = np.arange(4) if dense else np.array([0, 2]) + assert np.array_equal(summaries[0][0], expected_labels) + assert set(summaries[0][2]) == {"count", "mean", "m2", "sum", "min", "max", "sumofsquares"} + pd.testing.assert_frame_equal(result, expected) + + def test_aggregate_collected__complete_groups(self) -> None: + """Checks that values collected from separate blocks reconstruct complete groups in their original order.""" + + # Split two interleaved groups and one nodata value across three blocks + group_ids = np.array([0, 1, 0, 1, 0, 1, 0, 1, -1]) + values = np.array([9.0, 2.0, np.nan, 4.0, 5.0, 6.0, 3.0, 8.0, 100.0]) + statistics = _normalize_statistics(["median", "nmad", np.size]) + + # Select both groups from each block, then calculate from their joined values + blocks = [ + _collect_block([values[block]], group_ids[block], [0, 1]) + for block in [slice(0, 3), slice(3, 6), slice(6, 9)] + ] + result = _aggregate_collected(blocks, statistics) + expected, _ = _reduce_values([values], statistics, group_ids=group_ids, total_groups=2, strategy="groupwise") + + # The excluded final value is absent, while group order and nodata values are preserved + assert np.array_equal(np.concatenate([block[0] for block in blocks]), group_ids[group_ids >= 0]) + pd.testing.assert_frame_equal(result, expected, check_exact=True) + + @pytest.mark.parametrize( + "aliases,total_groups,expected,mergeable", + [ + (["mean"], 4096, "dense", True), + (["mean"], 4097, "sparse", True), + (["median"], 2, "groupwise", False), + ([None], 2, "groupwise", False), + ], + ) + def test_resolve_strategy__automatic( + self, aliases: list[str | None], total_groups: int, expected: str, mergeable: bool + ) -> None: + """Checks that auto selects a strategy from the statistics and number of groups.""" + + strategy, can_merge = _resolve_strategy(aliases, "auto", total_groups, chunked=True) + assert strategy == expected + assert can_merge is mergeable + + def test_resolve_strategy__error_invalid(self) -> None: + """Checks that unknown strategies and incomplete-group calculations are rejected for chunked inputs.""" + + # Reject an unknown option independently of the requested statistics + with pytest.raises(ValueError, match="must be 'auto', 'dense', 'sparse' or 'groupwise'"): + _resolve_strategy(["mean"], "topk", 2, chunked=True) + + # Median and custom functions need complete groups rather than dense or sparse summaries + with pytest.raises(ValueError, match="require ``strategy``='groupwise'"): + _resolve_strategy(["median"], "dense", 2, chunked=True) + with pytest.raises(ValueError, match="require ``strategy``='groupwise'"): + _resolve_strategy([None], "sparse", 2, chunked=True) + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + def test_reduce_values__groupwise_backends(self, backend: str) -> None: + """Checks that chunked medians, NMAD and custom functions exactly match complete eager groups.""" + + # Spread each group and its nodata values across several chunks + values: Any = np.arange(35, dtype=float) + values[::6] = np.nan + group_ids: Any = np.arange(35) % 4 + statistics = _normalize_statistics(["median", "nmad", np.size]) + expected, _ = _reduce_values([values], statistics, group_ids=group_ids, total_groups=4, strategy="groupwise") + config = MultiprocConfig(chunks=6) if backend == "multiproc" else None + if backend == "dask": + da = pytest.importorskip("dask.array") + values = da.from_array(values, chunks=6) + group_ids = da.from_array(group_ids, chunks=5) + + # Gather complete groups through the selected backend + result, strategy = _reduce_values( + [values], + statistics, + group_ids=group_ids, + total_groups=4, + strategy="auto", + mp_config=config, + ) + + # The result is computed, exactly equal, and leaves Dask inputs as lazy arrays + assert strategy == "groupwise" + pd.testing.assert_frame_equal(result, expected, check_exact=True) + if backend == "dask": + assert isinstance(values, da.Array) + assert isinstance(group_ids, da.Array) + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + @pytest.mark.parametrize("masked", [False, True]) + def test_reduce_values__integer_squares(self, backend: str, masked: bool) -> None: + """Checks that chunked integer squares exactly match the eager floating-point calculation.""" + + # Squared elevations exceed int16, with one optional location excluded by a NumPy mask + values = np.array([[2000, 3000], [5000, 4000]], dtype=np.int16) + source: Any = np.ma.array(values, mask=[[False, True], [False, False]]) if masked else values + statistics = _normalize_statistics(["sumofsquares", "rmse"], grouped=False) + expected, _ = _reduce_values([source], statistics) + config = MultiprocConfig(chunks=(1, 2)) if backend == "multiproc" else None + if backend == "dask": + da = pytest.importorskip("dask.array") + source = da.from_array(source, chunks=(1, 2)) + + # Reduce separate blocks without squaring in the original integer data type + result, _ = _reduce_values([source], statistics, mp_config=config) + pd.testing.assert_frame_equal(result, expected, check_exact=True) + if backend == "dask": + assert isinstance(source, da.Array) + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + @pytest.mark.parametrize("strategy", ["dense", "sparse", "groupwise"]) + @pytest.mark.parametrize("unsigned", [False, True]) + def test_reduce_values__large_integer_extrema(self, backend: str, strategy: str, unsigned: bool) -> None: + """Checks that chunked extrema remain exact for signed and unsigned integers beyond float precision.""" + + # Include two populated groups and one group containing only masked integer values + dtype = np.uint64 if unsigned else np.int64 + offset = 2**63 + 10 if unsigned else 2**60 + values = np.array([offset + step for step in [1, 3, 5, 7, 9, 11]], dtype=dtype) + source: Any = np.ma.array(values, mask=[False, True, False, False, True, True]) + group_ids: Any = np.array([0, 0, 1, 1, 2, 2]) + names = ["min", "max", "median" if strategy == "groupwise" else "mean"] + statistics = _normalize_statistics(names) + expected, _ = _reduce_values([source], statistics, group_ids=group_ids, total_groups=4, strategy=strategy) + config = MultiprocConfig(chunks=2) if backend == "multiproc" else None + if backend == "dask": + da = pytest.importorskip("dask.array") + source = da.from_array(source, chunks=2) + group_ids = da.from_array(group_ids, chunks=3) + + # Compare the complete chunked table and exact extrema with the eager reducer + result, _ = _reduce_values( + [source], + statistics, + group_ids=group_ids, + total_groups=4, + strategy=strategy, + mp_config=config, + ) + pd.testing.assert_frame_equal(result, expected) + assert int(result.loc[0, (0, "min")]) == int(values[0]) + assert int(result.loc[1, (0, "max")]) == int(values[3]) + assert pd.isna(result.loc[2, (0, "min")]) + if backend == "dask": + assert isinstance(source, da.Array) + assert isinstance(group_ids, da.Array) + + @pytest.mark.parametrize("backend", ["dask", "multiproc"]) + @pytest.mark.parametrize("nodata", ["empty", "nan", "masked"]) + def test_reduce_values__empty_sums(self, backend: str, nodata: str) -> None: + """Checks that chunked sums are undefined when empty, NaN or masked inputs have no valid values.""" + + # Build the requested empty representation and its eager reference + source: Any = np.empty(0) if nodata == "empty" else np.full(4, np.nan) + if nodata == "masked": + source = np.ma.array(np.arange(4), mask=True) + statistics = _normalize_statistics(["sum", "sumofsquares", "validcount"], grouped=False) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="Empty raster") + expected, _ = _reduce_values([source], statistics) + config = MultiprocConfig(chunks=2) if backend == "multiproc" else None + if backend == "dask": + da = pytest.importorskip("dask.array") + source = da.from_array(source, chunks=2) + + # Keep undefined sums and a zero valid count after reducing separate chunks + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="Empty raster") + result, _ = _reduce_values([source], statistics, mp_config=config) + pd.testing.assert_frame_equal(result, expected, check_exact=True) + assert np.isnan(result.loc[0, (0, "sum")]) + assert result.loc[0, (0, "validcount")] == 0 + if backend == "dask": + assert isinstance(source, da.Array) + + @pytest.mark.parametrize("strategy", ["dense", "sparse"]) + def test_reduce_values__variance_with_large_offset(self, strategy: str) -> None: + """Checks that merging Dask chunks preserves small variation around a large base value.""" + + # Values near 1e8 expose unstable variance formulas that subtract two large squared totals + da = pytest.importorskip("dask.array") + values = 1e8 + np.random.default_rng(42).normal(scale=0.1, size=3000) + group_ids = np.arange(values.size) % 3 + statistics = _normalize_statistics("std") + expected, _ = _reduce_values([values], statistics, group_ids=group_ids, total_groups=3, strategy=strategy) + + # Compare chunked pairwise variance with the eager reducer and NumPy + result, _ = _reduce_values( + [da.from_array(values, chunks=101)], + statistics, + group_ids=group_ids, + total_groups=3, + strategy=strategy, + ) + pd.testing.assert_frame_equal(result, expected) + direct = [np.std(values[group_ids == group_id]) for group_id in range(3)] + np.testing.assert_allclose(result[(0, "std")], direct) + + def test_reduce_values__groupwise_block_reads(self) -> None: + """Checks that groupwise Dask reduction reads only blocks containing a requested group.""" + + # Record value blocks read after group locations are known and reject the excluded final block + dask = pytest.importorskip("dask") + da = pytest.importorskip("dask.array") + reads = [] + + def read_values(block: int) -> NDArrayNum: + """Record each value block read and reject the block outside every group.""" + + assert block < 2 + reads.append(block) + return np.arange(4, dtype=float) + 4 * block + + blocks = [da.from_delayed(dask.delayed(read_values)(block), shape=(4,), dtype=float) for block in range(3)] + values = da.concatenate(blocks) + group_ids = np.array([0, 1, 0, 1, 2, 2, 2, 2, -1, -1, -1, -1]) + statistics = _normalize_statistics("median") + + # Calculate medians from only the blocks containing the three groups + with dask.config.set(scheduler="synchronous"): + result, _ = _reduce_values([values], statistics, group_ids=group_ids, total_groups=3, strategy="groupwise") + + # Every needed block is read once and the excluded block is never loaded + assert sorted(reads) == [0, 1] + np.testing.assert_allclose(result[(0, "median")], [1.0, 2.0, 5.5]) + + def test_statistics_dask__matches_shared_reducer(self) -> None: + """Checks that the optional native Dask reductions are lazy and match the shared GeoUtils reducer.""" + + # Include nodata values across uneven chunks and request every numerical estimator + import_optional("dask") + import dask + import dask.array as da + + values = np.arange(35, dtype=float).reshape(5, 7) + values[1, 2] = np.nan + values[3, 5] = np.nan + source = da.from_array(values, chunks=(2, 3)) + aliases = { + "mean", + "median", + "min", + "max", + "sum", + "sumofsquares", + "90thpercentile", + "iqr", + "le90", + "nmad", + "rmse", + "std", + } + statistics = _normalize_statistics(sorted(aliases), grouped=False) + expected, _ = _reduce_values([values], statistics) + + # Build both Dask calculations before computing their small results + native, native_count = _statistics_dask(source, aliases) + assert all(isinstance(value, da.Array) for value in native.values()) + assert isinstance(native_count, da.Array) + shared, _ = _reduce_values([source], statistics) + native_values, count = dask.compute(native, native_count) + + # Match both implementations with the eager reducer while the input remains a Dask array + pd.testing.assert_frame_equal(shared, expected) + assert count == expected.loc[0, (0, "count")] + for name, value in native_values.items(): + assert value == pytest.approx(expected.loc[0, (0, name)]) + assert isinstance(source, da.Array) + + @pytest.mark.parametrize("strategy", ["dense", "sparse", "groupwise"]) + def test_reduce_values__multiproc_workers(self, strategy: str) -> None: + """Checks that real Multiproc workers return the same grouped table as the eager reducer.""" + + # Split a rectangular array so every group crosses several worker tiles + from geoutils.multiproc.cluster import MpCluster + + shape = (31, 47) + values = np.arange(np.prod(shape), dtype=float).reshape(shape) + group_ids = np.indices(shape).sum(axis=0) % 3 + statistics = _normalize_statistics("mean") + expected, _ = _reduce_values([values], statistics, group_ids=group_ids, total_groups=3, strategy=strategy) + + # Calculate the same groups with two worker processes + with MpCluster({"nb_workers": 2}) as cluster: + result, _ = _reduce_values( + [values], + statistics, + group_ids=group_ids, + total_groups=3, + strategy=strategy, + mp_config=MultiprocConfig(chunks=(7, 11), cluster=cluster), + ) + + # Check the complete computed table + pd.testing.assert_frame_equal(result, expected) + + @pytest.mark.parametrize( + "shape,chunks,layout", + [ + ((131, 197), (37, 61), "interleaved"), + ((257, 509), (128, 193), "local"), + ((131, 197), (2048, 2048), "interleaved"), + ], + ) + @pytest.mark.parametrize("strategy", ["dense", "sparse"]) + def test_reduce_values__chunk_layouts( + self, shape: tuple[int, int], chunks: tuple[int, int], layout: str, strategy: str + ) -> None: + """Checks mergeable statistics across local, interleaved, uneven and oversized Dask chunks.""" + + # Create exact quarter-step values with separate nodata patterns + da = pytest.importorskip("dask.array") + rows, columns = np.indices(shape) + positions = rows * shape[1] + columns + first = 20 + (positions % 97) * 0.25 + second = -2 * first + positions % 3 + first[positions % 17 == 0] = np.nan + second[positions % 29 == 0] = np.inf + + # Place groups in separate areas or spread them across the array, then exclude selected locations + group_ids = positions % 8 + if layout == "local": + group_ids = (rows * 2 // shape[0]) * 4 + columns * 4 // shape[1] + keep = (rows + columns) % 19 != 0 + keep[: shape[0] // 4, : shape[1] // 4] = False + group_ids[~keep | (positions % 31 == 0)] = -1 + statistics = _normalize_statistics(["mean", "std", "sum", "min", "max", "rmse", "totalcount"]) + expected, _ = _reduce_values( + [first, second], statistics, group_ids=group_ids, total_groups=9, strategy=strategy + ) + + # Give values and group IDs different chunk layouts + values = [da.from_array(first, chunks=chunks), da.from_array(second, chunks=chunks[::-1])] + lazy_ids = da.from_array(group_ids, chunks=(chunks[0] + 3, chunks[1] + 5)) + result, _ = _reduce_values(values, statistics, group_ids=lazy_ids, total_groups=9, strategy=strategy) + + # Compare every computed column and confirm that all inputs remain Dask arrays + pd.testing.assert_frame_equal(result, expected) + assert all(isinstance(value, da.Array) for value in values) + assert isinstance(lazy_ids, da.Array) + + def test_reduce_values__sparse_group_ids(self) -> None: + """Checks that automatic sparse reduction keeps widely separated group IDs distinct.""" + + # Use three group IDs from a space just above the automatic dense threshold + da = pytest.importorskip("dask.array") + total_groups = 4097 + labels = np.array([0, total_groups // 2, total_groups - 1]) + values = np.arange(30, dtype=float) + group_ids = labels[np.arange(values.size) % 3] + statistics = _normalize_statistics(["mean", "std"]) + + # Let auto select sparse summaries for the Dask chunks + result, strategy = _reduce_values( + [da.from_array(values, chunks=7)], + statistics, + group_ids=da.from_array(group_ids, chunks=5), + total_groups=total_groups, + strategy="auto", + ) + + # Check the selected strategy, exact group IDs, and values from each complete group + assert strategy == "sparse" + assert result.index.tolist() == labels.tolist() + for label in labels: + members = values[group_ids == label] + np.testing.assert_allclose(result.loc[label], [members.size, members.mean(), members.std()]) diff --git a/tests/test_stats/test_sampling.py b/tests/test_stats/test_sampling.py deleted file mode 100644 index b28729113..000000000 --- a/tests/test_stats/test_sampling.py +++ /dev/null @@ -1,318 +0,0 @@ -"""Test sampling statistical tools.""" - -from __future__ import annotations - -import warnings -from typing import Literal - -import numpy as np -import pytest - -import geoutils as gu -from geoutils import open_raster -from geoutils._typing import NDArrayNum -from geoutils.multiproc import MultiprocConfig -from geoutils.raster.array import get_mask_from_array -from geoutils.stats.sampling import _subsample_numpy - - -class TestSampling: - """ - Different examples of 1D to 3D arrays with masked values for testing. - """ - - # Case 1 - 1D array, 1 masked value - array1D = np.ma.masked_array(np.arange(10), mask=np.zeros(10)) - array1D.mask[3] = True - assert np.ndim(array1D) == 1 - assert np.count_nonzero(array1D.mask) > 0 - - # Case 2 - 2D array, 1 masked value - array2D = np.ma.masked_array(np.arange(9).reshape((3, 3)), mask=np.zeros((3, 3))) - array2D.mask[0, 1] = True - assert np.ndim(array2D) == 2 - assert np.count_nonzero(array2D.mask) > 0 - - # Case 3 - 3D array, 1 masked value - array3D = np.ma.masked_array(np.arange(9).reshape((1, 3, 3)), mask=np.zeros((1, 3, 3))) - array3D = np.ma.vstack((array3D, array3D + 10)) - array3D.mask[0, 0, 1] = True - assert np.ndim(array3D) == 3 - assert np.count_nonzero(array3D.mask) > 0 - - @pytest.mark.parametrize("array", [array1D, array2D, array3D]) - def test_subsample(self, array: NDArrayNum) -> None: - """ - Test gu.stats.subsample_array. - """ - - warnings.filterwarnings("ignore", message=".*larger than the number of valid pixels.*", category=UserWarning) - - # Test that subsample > 1 works as expected, i.e. output 1D array, with no masked values, or selected size - for npts in np.arange(2, np.size(array)): - random_values = _subsample_numpy(array, subsample=npts) - assert np.ndim(random_values) == 1 - assert np.size(random_values) == npts - assert np.count_nonzero(random_values.mask) == 0 - - # Test if subsample > number of valid values => return all - random_values = _subsample_numpy(array, subsample=np.size(array) + 3) - assert np.all(np.sort(random_values) == array[~array.mask]) - - # Test if subsample = 1 => return all valid values - random_values = _subsample_numpy(array, subsample=1) - assert np.all(np.sort(random_values) == array[~array.mask]) - - # Check that order is preserved for subsample = 1 (no random sampling, simply returns valid mask) - random_values_2 = _subsample_numpy(array, subsample=1) - assert np.array_equal(random_values, random_values_2) - - # Test if subsample < 1 - random_values = _subsample_numpy(array, subsample=0.5) - assert np.size(random_values) == int(np.count_nonzero(~array.mask) * 0.5) - - # Test with optional argument return_indices - indices = _subsample_numpy(array, subsample=0.3, return_indices=True) - assert np.ndim(indices) == 2 - assert len(indices) == np.ndim(array) - assert np.ndim(array[indices]) == 1 - assert np.size(array[indices]) == int(np.count_nonzero(~array.mask) * 0.3) - - # Check that we can pass an integer to fix the random state - sub42 = _subsample_numpy(array, subsample=10, random_state=42) - # Check by passing a generator directly - rng = np.random.default_rng(42) - sub42_gen = _subsample_numpy(array, subsample=10, random_state=rng) - # Both should be equal - assert np.array_equal(sub42, sub42_gen) - - -class TestSubsampleChunked: - - # Strategies supported by _subsample - subsample_strategies = ("sequential", "topk") - - @pytest.mark.parametrize("path_index", [0, 2]) - @pytest.mark.parametrize("strategy", subsample_strategies) - @pytest.mark.parametrize("return_indices", [False, True]) - @pytest.mark.parametrize("subsample", [2, 100, 0.05]) # int size and fraction - def test_subsample__backends( - self, - path_index: int, - strategy: Literal["sequential", "topk"], - return_indices: bool, - subsample: int | float, - lazy_test_files_tiny: list[str], - ) -> None: - """ - Test that subsample behaves consistently across backends: - - NumPy backend through Raster (in-memory), - - NumPy backend through Xarray DataArray (in-memory), - - Dask backend through Xarray accessor (lazy), - - Multiprocessing backend through Raster (lazy input; eager output by design). - - Notes: - - "topk" strategy is intended to be chunk-invariant -> outputs should match across backends. - - "sequential" strategy is chunk/order-dependent -> we do NOT require cross-backend equality, - but we still validate output properties and determinism for a fixed random_state. - """ - - pytest.importorskip("dask") - import dask.array as da - - warnings.filterwarnings("ignore", category=UserWarning, message="Subsample value*") - - # Get filepath of on-disk (for laziness) test file - path_raster = lazy_test_files_tiny[path_index] - - # 1/ Prepare inputs for each backend - - # Base raster input (in-memory -> NumPy backend) - raster_base = gu.Raster(path_raster) - raster_base.load() - assert raster_base.is_loaded - - # Base data array input (in-memory -> NumPy backend through Xarray) - ds_base = open_raster(path_raster) - ds_base.load() - assert ds_base._in_memory - - # Multiprocessing input (keep lazy until mp backend reads) - raster_mp = gu.Raster(path_raster) - assert not raster_mp.is_loaded - - # Dask input (lazy chunked xarray) - ds_dask = open_raster(path_raster, chunks={"x": 10, "y": 10}) - assert not ds_dask._in_memory - assert isinstance(ds_dask.data, da.Array) - assert ds_dask.data.chunks is not None - - # 2/ Run subsample across backends (fixed seed for determinism) - seed = 42 - mp_config = MultiprocConfig(chunks=(10, 7)) - - # NumPy backend via Raster - out_raster = raster_base.subsample( - subsample=subsample, - return_indices=return_indices, - random_state=seed, - strategy=strategy, - ) - - # NumPy backend via Xarray DataArray - out_xr = ds_base.rst.subsample( - subsample=subsample, - return_indices=return_indices, - random_state=seed, - strategy=strategy, - ) - - # Dask backend via Xarray accessor: should return lazy dask arrays - out_dask = ds_dask.rst.subsample( - subsample=subsample, - return_indices=return_indices, - random_state=seed, - strategy=strategy, - ) - - # Multiprocessing backend via Raster: output is eager by design, input raster stays lazy - out_mp = raster_mp.subsample( - subsample=subsample, - return_indices=return_indices, - random_state=seed, - strategy=strategy, - mp_config=mp_config, - ) - - # 3/ Laziness checks - - # Dask input stays unloaded and lazy - assert not ds_dask._in_memory - assert isinstance(ds_dask.data, da.Array) - - # Multiprocessing should not load the raster object itself (still points to disk) - assert not raster_mp.is_loaded - - # Dask output type checks (lazy) - if return_indices: - assert isinstance(out_dask, tuple) and len(out_dask) == 2 - assert isinstance(out_dask[0], da.Array) - assert isinstance(out_dask[1], da.Array) - else: - assert isinstance(out_dask, da.Array) - - # 4/ Normalize outputs to comparable NumPy representations - - def _as_numpy( - out: object, - ) -> NDArrayNum | tuple[NDArrayNum, NDArrayNum]: - """Convert backend outputs to NumPy arrays for comparison.""" - if isinstance(out, tuple): - r, c = out - if hasattr(r, "compute"): - r = r.compute() - if hasattr(c, "compute"): - c = c.compute() - return (np.asarray(r), np.asarray(c)) - else: - if hasattr(out, "compute"): - out = out.compute() - return np.asarray(out) - - out_raster_np = _as_numpy(out_raster) - out_xr_np = _as_numpy(out_xr) - out_dask_np = _as_numpy(out_dask) - out_mp_np = _as_numpy(out_mp) - - # 5/ Generic output validity checks (all backends) - - # Mirror _subsample() array selection (band=1 default) - arr = raster_base.data if raster_base.data.ndim == 2 else raster_base.data[0, :, :] - assert arr.ndim == 2 - - # Mirror _subsample_numpy() validity definition - mask = get_mask_from_array(arr) # True where invalid - n_valid = int(np.count_nonzero(~mask)) # valid pixels - - # Mirror _get_subsample_size_from_user_input() - if isinstance(subsample, float): - expected = int(subsample * n_valid) - else: - expected = min(int(subsample), n_valid) - - def _check_output(out_np: NDArrayNum | tuple[NDArrayNum, NDArrayNum]) -> None: - """Validate length, finiteness and index/value consistency.""" - if isinstance(out_np, tuple): - rr, cc = out_np - assert rr.shape == cc.shape - assert rr.ndim == 1 and cc.ndim == 1 - assert len(rr) == expected - # Indices should be within bounds - assert np.all((0 <= rr) & (rr < arr.shape[0])) - assert np.all((0 <= cc) & (cc < arr.shape[1])) - # Indices must point to finite values - # Indices must point to valid values per get_mask_from_array - assert np.all(~mask[rr, cc]) - else: - assert out_np.ndim == 1 - assert len(out_np) == expected - assert np.all(np.isfinite(out_np)) - - _check_output(out_raster_np) - _check_output(out_xr_np) - _check_output(out_dask_np) - _check_output(out_mp_np) - - # 6/ Backend equivalence logic - - if strategy == "topk": - # Chunk-invariant strategy: require exact equality across backends. - assert np.array_equal(out_raster_np, out_xr_np) - assert np.array_equal(out_raster_np, out_dask_np) - assert np.array_equal(out_raster_np, out_mp_np) - else: - # Sequential: do not require equality across backends (order/chunk dependent), - # but require determinism per backend for the same seed. - out_raster_np_2 = _as_numpy( - raster_base.subsample( - subsample=subsample, - return_indices=return_indices, - random_state=seed, - strategy=strategy, - ) - ) - out_dask_np_2 = _as_numpy( - ds_dask.rst.subsample( - subsample=subsample, - return_indices=return_indices, - random_state=seed, - strategy=strategy, - ) - ) - out_mp_np_2 = _as_numpy( - raster_mp.subsample( - subsample=subsample, - return_indices=return_indices, - random_state=seed, - strategy=strategy, - mp_config=mp_config, - ) - ) - - assert np.array_equal(out_raster_np, out_raster_np_2) - assert np.array_equal(out_dask_np, out_dask_np_2) - assert np.array_equal(out_mp_np, out_mp_np_2) - - # 7/ Indices versus values consistency check for return_indices=True - if return_indices: - rr, cc = out_raster_np # use any backend; for topk they match; for sequential we validate per-backend above - vals_from_indices = arr[rr, cc] - # Compare to "values mode" from same backend and same seed/strategy - vals_raster = raster_base.subsample( - subsample=subsample, - return_indices=False, - random_state=seed, - strategy=strategy, - ) - vals_raster_np = _as_numpy(vals_raster) - assert np.array_equal(np.asarray(vals_from_indices), np.asarray(vals_raster_np)) diff --git a/tests/test_stats/test_selection.py b/tests/test_stats/test_selection.py new file mode 100644 index 000000000..41c16d750 --- /dev/null +++ b/tests/test_stats/test_selection.py @@ -0,0 +1,663 @@ +"""Tests selecting statistic values and masks on a common support.""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +import rasterio as rio +from pyproj import CRS +from rasterio.transform import from_origin +from shapely.geometry import box + +import geoutils as gu +from geoutils._misc import import_optional +from geoutils.multiproc import ClusterGenerator, MultiprocConfig + + +@pytest.fixture(params=["raster", "gpkg", "las"]) +def stats_file(request: pytest.FixtureRequest, tmp_path: Path) -> tuple[Any, Any, dict[str, Any], Any]: + """Write raster or point files with two value columns and one grouping variable to disk.""" + + # Create two value columns with nodata values at different locations and one integer grouping column + positions = np.arange(48) + first = positions.astype(float) + 1 + second = positions.astype(float) ** 2 + 100 + first[[3, 19]] = np.nan + second[[8, 27, 43]] = np.nan + groups = positions % 8 + if request.param == "raster": + # Write a raster file to disk with the two value columns and grouping column stored as separate bands + filename = tmp_path / "statistics.tif" + values = np.stack((first, second, groups)).reshape(3, 6, 8) + gu.Raster.from_array(values, from_origin(500000, 5100006, 1, 1), 32633, nodata=np.nan).to_file(filename) + source = gu.Raster(filename) + reference = gu.Raster(filename, load_data=True) + return source, reference, {"first": 1, "second": 2}, 3 + + filename = tmp_path / f"statistics.{request.param}" + if request.param == "las": + # Write a LAS file to disk with values in attributes separate from the point coordinates + laspy = import_optional("laspy") + header = laspy.LasHeader(point_format=6, version="1.4") + header.add_crs(CRS.from_epsg(32633)) + for name, dtype in (("first", "float64"), ("second", "float64"), ("group", "int32")): + header.add_extra_dim(laspy.ExtraBytesParams(name=name, type=dtype)) + records = laspy.LasData(header) + records.x, records.y, records.z = 500000 + positions, 5100000 + positions, positions + 1000 + records.first, records.second, records.group = first, second, groups + records.write(filename) + else: + # Write a GeoPackage file to disk with both value columns and the grouping column + dataframe = gpd.GeoDataFrame( + {"first": first, "second": second, "group": groups}, + geometry=gpd.points_from_xy(500000 + positions, 5100000 + positions), + crs=32633, + ) + dataframe.to_file(filename, index=False) + + # Open one unloaded source for Multiproc and one eager source for the expected result + source = gu.PointCloud(filename, data_column="first") + reference = gu.PointCloud(filename, data_column="first") + reference.load(columns="all") + return source, reference, {"first": "first", "second": "second"}, "group" + + +class TestSelection: + """ + Checks the values and masks accepted by stats(). + + - Rasters and point clouds can select bands, columns or plain Xarrays. + - NumPy and tabular masks use positions on the common support. + - Dask and Multiproc behavior is covered by TestSelectionChunked below. + """ + + @pytest.mark.parametrize("grouped", [False, True]) + def test_stats__band_selection(self, grouped: bool) -> None: + """Checks that the values argument selects the requested raster band with and without groups.""" + + # Give the two bands different means so selecting the default band would fail + values = np.arange(1, 7, dtype=float).reshape(2, 3) + raster = gu.Raster.from_array( + np.stack((values, values + 100)), + transform=rio.transform.from_origin(0, 2, 1, 1), + crs=4326, + ) + options = {} + if grouped: + options = {"by": {"zone": np.array([[0, 0, 0], [1, 1, 1]])}, "categories": {"zone": [0, 1]}} + + # Select the second band through the raster method + selected = raster.stats("mean", values=2, **options) + + # Check the selected band's whole mean or separate row means + if grouped: + np.testing.assert_allclose(selected[("band_2", "mean")], [102, 105]) + else: + assert selected == pytest.approx(103.5) + + @pytest.mark.parametrize("source_type", ["array", "raster", "pointcloud"]) + @pytest.mark.parametrize("grouped", [False, True]) + def test_stats__mask_selection(self, source_type: str, grouped: bool) -> None: + """Checks that mask selects the same values and preserves total and valid counts across input types.""" + + # Place one nodata value inside the mask to distinguish total and valid inlier counts + values = np.array([[1.0, np.nan, 3.0], [4.0, 5.0, 6.0]]) + keep = np.array([[True, True, False], [True, False, True]]) + groups = np.array([[0, 0, 0], [1, 1, 1]]) + source: Any = values + if source_type == "raster": + source = gu.Raster.from_array(values, rio.transform.from_origin(0, 2, 1, 1), 4326, nodata=np.nan) + elif source_type == "pointcloud": + source = gu.PointCloud.from_xyz(np.arange(values.size), np.zeros(values.size), values.ravel(), crs=4326) + keep, groups = keep.ravel(), groups.ravel() + + # Apply the same mask to either a complete summary or the same two groups + options: dict[str, Any] = {} + statistics: str | list[str] = "all" + if grouped: + options = {"by": {"zone": groups}, "categories": {"zone": [0, 1]}} + statistics = ["mean", "totalcount"] + masked = gu.stats.stats(source, statistics, mask=keep, **options) + + # Check counts before and after the mask for the summary and for each group + if grouped: + np.testing.assert_allclose(masked.xs("count", axis=1, level="statistic").iloc[:, 0], [1, 2]) + np.testing.assert_allclose(masked.xs("totalcount", axis=1, level="statistic").iloc[:, 0], [2, 2]) + np.testing.assert_allclose(masked.xs("mean", axis=1, level="statistic").iloc[:, 0], [1, 5]) + else: + assert masked["Mean"] == pytest.approx(11 / 3) + assert masked["Valid count"] == 5 + assert masked["Total count"] == 6 + assert masked["Valid inlier count"] == 3 + assert masked["Total inlier count"] == 4 + assert masked["Percentage inlier points"] == 60 + assert masked["Percentage valid inlier points"] == 75 + + def test_stats__point_column_selection(self) -> None: + """Checks that direct point cloud summaries use the active or explicitly selected column.""" + + # Give the additional column a distinct scale so active-column fallback is detectable + pointcloud = gu.PointCloud.from_xyz(np.arange(4), np.zeros(4), np.arange(1, 5), crs=4326) + pointcloud.ds["temperature"] = [10.0, 20.0, 30.0, 40.0] + + # Calculate default and selected summaries through both public entry points + default = gu.stats.stats(pointcloud) + selected = gu.stats.stats(pointcloud, "mean", values="temperature") + method_selected = pointcloud.stats("mean", values="temperature") + method_masked = pointcloud.stats("mean", values="temperature", mask=np.array([True, True, False, False])) + + # Match the active-column and auxiliary-column means independently + assert default["Mean"] == pytest.approx(2.5) + assert selected == method_selected == pytest.approx(25) + assert method_masked == pytest.approx(15) + + @pytest.mark.parametrize("grouped", [False, True]) + def test_stats__point_inputs_from_plain_xarrays(self, grouped: bool) -> None: + """Checks that point values, masks and groups in plain Xarrays follow the source's point order.""" + + import xarray as xr + + # Use a nonspatial dimension so the point cloud supplies all location information + pointcloud = gu.PointCloud.from_xyz(np.arange(4), np.zeros(4), np.arange(4), crs=4326) + values = xr.DataArray([10.0, 20.0, 30.0, 40.0], dims="point") + keep = xr.DataArray([True, False, True, True], dims="point") + groups = xr.DataArray([0, 0, 1, 1], dims="point") + options = {"by": {"zone": groups}, "categories": {"zone": [0, 1]}} if grouped else {} + + # Use plain Xarrays for selected values, the user mask, and the optional grouping variable + result = pointcloud.stats("mean", values={"temperature": values}, mask=keep, **options) + + # Exclude the second point, leaving one value in the first group and two in the second + if grouped: + np.testing.assert_allclose(result["temperature"], [[1, 10], [2, 35]]) + else: + assert result == pytest.approx((10 + 30 + 40) / 3) + + @pytest.mark.parametrize("grouped", [False, True]) + @pytest.mark.parametrize("source_type", ["array", "raster"]) + def test_stats__grid_inputs_from_plain_xarrays(self, grouped: bool, source_type: str) -> None: + """Checks that Xarray dimensions named x and y do not require spatial coordinates for statistics.""" + + import xarray as xr + + # Give values, groups and the mask spatial dimension names without defining coordinates or a CRS + values = xr.DataArray([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], dims=("y", "x")) + keep = xr.DataArray([[True, False, True], [False, True, True]], dims=("y", "x")) + groups = xr.DataArray([[0, 0, 0], [1, 1, 1]], dims=("y", "x")) + options = {"by": {"zone": groups}, "categories": {"zone": [0, 1]}} if grouped else {} + + # Use these arrays directly or select them as values on an existing raster grid + if source_type == "array": + result = gu.stats.stats(values, "mean", mask=keep, **options) + value_name = "value" + else: + raster = gu.Raster.from_array(np.zeros((2, 3)), rio.transform.from_origin(0, 2, 1, 1), crs=4326) + result = raster.stats("mean", values={"temperature": values}, mask=keep, **options) + value_name = "temperature" + + # The mask leaves two values in each row, giving row means of 20 and 55 and a whole mean of 37.5 + if grouped: + np.testing.assert_allclose(result[value_name], [[2, 20], [2, 55]]) + else: + assert result == pytest.approx(37.5) + + @pytest.mark.parametrize("container", ["series", "dataframe", "dask_series"]) + def test_stats__tabular_values_and_masks(self, container: str) -> None: + """Checks that tabular inputs apply masks by position and return the expected shape and counts.""" + + # 1/ Prepare tabular values, masks and groups + # Use different row labels to check positional selection, with one nodata value inside the mask + labels = np.arange(6) + 10 + source: Any = pd.Series([1.0, np.nan, 3.0, 4.0, 5.0, 6.0], index=labels) + keep: Any = pd.Series([True, True, False, True, False, True], index=labels[::-1]) + groups: Any = pd.Series([0, 0, 0, 1, 1, 1], index=labels + 100) + if container == "dataframe": + source, keep, groups = (array.to_frame() for array in (source, keep, groups)) + elif container == "dask_series": + import_optional("dask") + import dask.array as da + import dask.dataframe as dd + + source, keep, groups = ( + dd.from_pandas(array, npartitions=2, sort=False) for array in (source, keep, groups) + ) + + # 2/ Calculate masked summaries and groups + # Request the callable's input shape as well as counts before and after selection + statistics: list[str | Callable[[Any], Any]] = [ + "mean", + "validcount", + "validinliercount", + "totalinliercount", + np.shape, + ] + summary = gu.stats.stats(source, statistics, mask=keep) + grouped = gu.stats.stats( + source, + ["mean", "totalcount"], + by={"zone": groups}, + categories={"zone": [0, 1]}, + mask=keep, + ) + + # 3/ Check shape, calculation result and the selected observations + # Dask Series return a computed summary; DataFrame summaries use their original two-dimensional shape + if container == "dask_series": + assert not isinstance(summary["mean"], da.Array) + assert summary["shape"] == ((6, 1) if container == "dataframe" else (6,)) + assert summary["mean"] == pytest.approx(11 / 3) + assert summary["validcount"] == 5 + assert summary["validinliercount"] == 3 + assert summary["totalinliercount"] == 4 + + # The first group has one finite value in two selected locations; the second has two finite values + np.testing.assert_allclose(grouped["value"], [[1, 1, 2], [2, 5, 2]]) + + @pytest.mark.parametrize("grouped", [False, True]) + def test_stats__error_tabular_inputs(self, grouped: bool) -> None: + """Checks that tabular statistics reject invalid masks and mixed Dask and Multiproc backends.""" + + # A DataFrame is a two-dimensional value array, so each grouper must use that same shape + values = pd.DataFrame(np.arange(6, dtype=float).reshape(2, 3)) + options: dict[str, Any] = ( + {"by": {"zone": np.zeros(values.shape, dtype=int)}, "categories": {"zone": [0]}} if grouped else {} + ) + + # Reject a row-only mask and a numeric mask rather than broadcasting or treating nonzero values as True + with pytest.raises( + ValueError, match="Argument ``mask`` must be boolean and contain one value per input location" + ): + gu.stats.stats(values, "mean", mask=pd.Series([True, False]), **options) + with pytest.raises( + ValueError, match="Argument ``mask`` must be boolean and contain one value per input location" + ): + gu.stats.stats(values, "mean", mask=pd.DataFrame(np.ones(values.shape, dtype=int)), **options) + + # Check that a Dask Series is rejected before the Multiproc backend starts any workers + import_optional("dask") + import dask.dataframe as dd + + lazy_values = dd.from_pandas(pd.Series(np.arange(6, dtype=float)), npartitions=2) + options = {"by": {"zone": np.zeros(6, dtype=int)}, "categories": {"zone": [0]}} if grouped else {} + with pytest.raises(ValueError, match="Dask inputs cannot be combined with Multiprocessing"): + gu.stats.stats(lazy_values, "mean", mp_config=MultiprocConfig(chunks=2), **options) + + +class TestSelectionChunked: + """ + Checks stats() input selection with Dask and Multiproc inputs. + + Dask and Multiproc results are compared with eager calculations. Raster and point cloud files stay unloaded while + their values, masks and common support are read in chunks. + """ + + def test_stats__masked_dask_summary_is_computed(self) -> None: + """Checks that a masked Dask raster gives computed values and the same counts as eager data.""" + + # Use mismatched data and mask chunks so masking also checks automatic chunk alignment + import_optional("dask") + import dask.array as da + import xarray as xr + + values = np.array([[1.0, np.nan, 3.0], [4.0, 5.0, 6.0]]) + keep = np.array([[True, True, False], [True, False, True]]) + raster = gu.Raster.from_array(values, rio.transform.from_origin(0, 2, 1, 1), 4326, nodata=np.nan) + lazy_raster = raster.to_xarray().chunk({"x": 2, "y": 1}).rst + lazy_keep = xr.DataArray(da.from_array(keep, chunks=(2, 1)), dims=("y", "x")) + statistics = ["mean", "std", "validcount", "validinliercount", "totalinliercount"] + assert isinstance(lazy_raster.data, da.Array) and isinstance(lazy_keep.data, da.Array) + + # Check that stats() computes the Dask summary immediately, as it does for grouped results + result = lazy_raster.stats(statistics, mask=lazy_keep) + assert not isinstance(result["mean"], da.Array) + + # Compare the result with the same masked eager raster + expected = raster.stats(statistics, mask=keep) + assert result == pytest.approx(expected) + assert isinstance(lazy_raster.data, da.Array) and isinstance(lazy_keep.data, da.Array) + + def test_stats__dask_point_vector_mask(self, tmp_path: Path) -> None: + """Checks that eager and Dask points agree on points lying along a polygon boundary.""" + + # Write a point file to disk with two points inside the polygon, one on its edge, and one outside it + import geopandas as gpd + from shapely.geometry import box + + import_optional("dask_geopandas", package_name="dask-geopandas") + points = gu.PointCloud.from_xyz( + np.array([0.5, 1.5, 2.0, 2.5]), np.ones(4), np.array([10.0, 20.0, 30.0, 40.0]), crs=32633 + ) + mask = gu.Vector(gpd.GeoDataFrame(geometry=[box(0, 0, 2, 2)], crs=points.crs)) + filename = tmp_path / "masked_points.gpkg" + points.to_file(filename) + lazy_points = gu.open_pointcloud(str(filename), data_column=points.data_column, chunks=2).pc + assert not lazy_points.is_loaded + + # Calculate the same masked summary from Dask partitions and eager points + statistics = ["mean", "validinliercount"] + expected = points.stats(statistics, mask=mask) + result = lazy_points.stats(statistics, mask=mask) + + # Only the first two values contribute; the point on the polygon edge remains excluded + assert result == pytest.approx(expected) + assert result == pytest.approx({"mean": 15, "validinliercount": 2}) + assert not lazy_points.is_loaded + + def test_stats__mixed_eager_and_dask_values(self) -> None: + """Checks that one stats() call returns computed summaries for mixed eager and Dask values.""" + + # Select the same locations from eager and Dask values with distinct scales + import_optional("dask") + import dask.array as da + + eager = np.array([1.0, np.nan, 3.0, 4.0, 5.0, 6.0]) + lazy = da.from_array(eager * 10, chunks=2) + keep = np.array([True, True, False, True, False, True]) + + # Calculate both value arrays in the same Dask call + result = gu.stats.stats({"eager": eager, "lazy": lazy}, ["mean", "validinliercount"], mask=keep) + assert isinstance(result["eager"]["mean"], float) + assert isinstance(result["lazy"]["mean"], float) + + # Both values select three finite locations, with means differing by the known factor of ten + assert result["eager"] == pytest.approx({"mean": 11 / 3, "validinliercount": 3}) + assert result["lazy"] == pytest.approx({"mean": 110 / 3, "validinliercount": 3}) + assert isinstance(lazy, da.Array) + + @pytest.mark.parametrize("workers", [False, True]) + @pytest.mark.parametrize("grouped", [False, True]) + @pytest.mark.parametrize("exact", [False, True]) + def test_stats__file_values_match_eager( + self, + stats_file: tuple[Any, Any, dict[str, Any], Any], + workers: bool, + grouped: bool, + exact: bool, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Checks that Multiproc statistics match eager while the complete source stays on disk.""" + + # Select both value columns so their separate nodata locations produce different counts and statistics + source, reference, values, group = stats_file + # Include statistics that require every finite value from a group at once + statistics = ["median", np.nanmean] if exact else ["mean", "std", "min", "max", "sumofsquares"] + options = {"by": {"range": group}, "bins": {"range": [-0.5, 2.5, 5.5, 7.5]}} if grouped else {} + expected = reference.stats(statistics, values=values, **options) + + # Prevent full-source reads so file statistics must use raster windows or point row ranges + def reject_load(*args: Any, **kwargs: Any) -> None: + """Reject a full read while file statistics should be reading only the requested blocks.""" + + raise AssertionError("Statistics must read file blocks without calling source.load().") + + with monkeypatch.context() as guarded: + guarded.setattr(type(source), "load", reject_load) + with ClusterGenerator("multi" if workers else "basic", nb_workers=2) as cluster: + config = MultiprocConfig(chunks=(2, 3), cluster=cluster) + result = source.stats(statistics, values=values, mp_config=config, **options) + + # Compare the complete grouped table or every global result with the eager stats() call + if grouped: + pd.testing.assert_frame_equal(result, expected, check_dtype=False) + else: + assert result.keys() == expected.keys() + for name in expected: + assert result[name] == pytest.approx(expected[name]) + assert not source.is_loaded + + @pytest.mark.parametrize("workers", [False, True]) + def test_stats__automatic_bins_sampling_and_masks( + self, stats_file: tuple[Any, Any, dict[str, Any], Any], workers: bool + ) -> None: + """Checks that automatic bins and sampling return full group masks while the source is not loaded.""" + + # Exclude the highest group coordinate so automatic edges must be inferred after applying the mask + source, reference, values, group = stats_file + keep = np.arange(48) % 8 != 7 + if isinstance(source, gu.Raster): + keep = keep.reshape(source.shape) + options = { + "values": values, + "by": {"range": group}, + "bins": {"range": 3}, + "mask": keep, + "subsample": 12, + "random_state": 42, + "return_masks": True, + } + expected, expected_masks = reference.stats(["mean", "std"], **options) + + # Find automatic bin edges, select a repeatable sample, and calculate its statistics in file blocks + with ClusterGenerator("multi" if workers else "basic", nb_workers=2) as cluster: + config = MultiprocConfig(chunks=(2, 3), cluster=cluster) + result, masks = source.stats(["mean", "std"], mp_config=config, **options) + pd.testing.assert_frame_equal(result, expected, check_dtype=False) + assert not source.is_loaded + + # Requested masks contain every group location, including locations outside the twelve-point sample + group_location_count = 0 + for key in result.index: + mask = masks[key] + assert np.array_equal(mask.data, expected_masks[key].data) + group_location_count += int(np.count_nonzero(mask.data)) + assert not source.is_loaded + assert group_location_count == int(keep.sum()) == 42 + + @pytest.mark.parametrize("workers", [False, True]) + def test_stats__file_mask_summary_counts(self, workers: bool, tmp_path: Path) -> None: + """Checks that unloaded boolean masks preserve original valid counts and selected inlier counts.""" + + # Write a value raster and boolean mask raster to disk, with one value nodata cell inside and outside the mask + values = np.arange(1, 25, dtype=float).reshape(4, 6) + values[0, 0], values[3, 5] = np.nan, np.nan + keep = np.ones(values.shape, dtype=bool) + keep[2:, 4:] = False + transform = from_origin(0, 4, 1, 1) + value_filename, mask_filename = tmp_path / "values.tif", tmp_path / "mask.tif" + gu.Raster.from_array(values, transform, 32633, nodata=np.nan).to_file(value_filename) + gu.Raster.from_array(keep, transform, 32633).to_file(mask_filename) + source, mask_source = gu.Raster(value_filename), gu.Raster(mask_filename, is_mask=True) + statistics = ["mean", "validcount", "totalcount", "validinliercount", "totalinliercount"] + expected = gu.stats.stats(values, statistics, mask=keep) + + # Count the original finite values separately from values retained by the file mask + with ClusterGenerator("multi" if workers else "basic", nb_workers=2) as cluster: + config = MultiprocConfig(chunks=(3, 5), cluster=cluster) + result = source.stats(statistics, mask=mask_source, mp_config=config) + + # The edge windows have one row or column, and all twenty selected cells include one nodata cell + assert result == pytest.approx(expected) + assert result["validcount"] == 22 and result["totalcount"] == 24 + assert result["validinliercount"] == 19 and result["totalinliercount"] == 20 + assert not source.is_loaded and not mask_source.is_loaded + + @pytest.mark.parametrize("nullable_integer", [False, True]) + def test_stats__sampled_point_missing_values(self, nullable_integer: bool, tmp_path: Path) -> None: + """Checks that sampled point summaries read floating and nullable integer columns containing nodata.""" + + # Write a point file to disk with a nodata value in the final block and an optional nullable integer column + values = ( + pd.array([1, 2, 3, 4, 5, None], dtype="Int64") if nullable_integer else np.array([1, 2, 3, 4, 5, np.nan]) + ) + dataframe = gpd.GeoDataFrame( + {"value": values}, geometry=gpd.points_from_xy(np.arange(6), np.arange(6)), crs=32633 + ) + filename = tmp_path / "missing_points.gpkg" + dataframe.to_file(filename, index=False) + source = gu.PointCloud(filename, data_column="value") + + # Selecting all six rows includes the row with nodata in the sample without changing the mean + statistics = ["mean", "min", "max", "sum", "validcount", "totalcount"] + config = MultiprocConfig(chunks=2) + result = source.stats(statistics, subsample=6, random_state=0, mp_config=config) + + # The five finite values sum to fifteen; the file's nodata row contributes only to the total count + expected = {"mean": 3, "min": 1, "max": 5, "sum": 15, "validcount": 5, "totalcount": 6} + assert result == pytest.approx(expected) + assert not source.is_loaded + + @pytest.mark.parametrize("subsample", [1, 4]) + def test_stats__point_infinity_summary(self, subsample: int, tmp_path: Path) -> None: + """Checks that point summaries include infinity in estimates but count only finite values as valid.""" + + # Write a point file to disk with one unmasked infinite value and three finite values + dataframe = gpd.GeoDataFrame( + {"value": [1, np.inf, 3, 4]}, geometry=gpd.points_from_xy(np.arange(4), np.arange(4)), crs=32633 + ) + filename = tmp_path / "infinite_points.gpkg" + dataframe.to_file(filename, index=False) + source = gu.PointCloud(filename, data_column="value") + + # A sample size of four reads every row through sampling, while one uses the complete summary directly + statistics = ["mean", "min", "max", "sum", "validcount", "totalcount"] + result = source.stats(statistics, subsample=subsample, random_state=0, mp_config=MultiprocConfig(chunks=2)) + + # NumPy's NaN-aware estimators include positive infinity, and only three observations are finite + expected = {"mean": np.inf, "min": 1, "max": np.inf, "sum": np.inf, "validcount": 3, "totalcount": 4} + assert result == expected + assert not source.is_loaded + + @pytest.mark.parametrize("kind", ["raster", "points"]) + @pytest.mark.parametrize("mask_mode", ["inside", "outside"]) + def test_stats__vector_masks_and_values(self, kind: str, mask_mode: str, tmp_path: Path) -> None: + """Checks that file statistics apply vector masks and attributes while preserving all summary counts.""" + + # Create twelve values on both sides of a polygon, with one nonfinite value on each side + values = np.arange(12, dtype=float).reshape(3, 4) + values[0, 0], values[2, 3] = np.nan, np.nan + zones = gpd.GeoDataFrame({"rating": [10.0]}, geometry=[box(0, 0, 2, 3)], crs=32631) + source: Any + reference: Any + if kind == "raster": + # Write a raster file to disk with nodata cells inside and outside the polygon + path = tmp_path / "masked.tif" + gu.Raster.from_array(values, from_origin(0, 3, 1, 1), 32631, nodata=np.nan).to_file(path) + source = gu.Raster(path) + reference = gu.Raster(path, load_data=True) + else: + # Write a point file to disk with nodata values inside and outside the polygon + x, y = np.meshgrid(np.arange(4) + 0.5, np.arange(3) + 0.5) + dataframe = gpd.GeoDataFrame( + {"height": values.ravel()}, geometry=gpd.points_from_xy(x.ravel(), y.ravel()), crs=32631 + ) + path = tmp_path / "masked.gpkg" + dataframe.to_file(path, index=False) + source = gu.PointCloud(path, data_column="height") + reference = gu.PointCloud(dataframe, data_column="height") + + # Use edge tiles with one row or column and retain counts from before the vector mask + statistics = ["mean", "validcount", "totalcount", "validinliercount", "totalinliercount"] + expected = reference.stats(statistics, mask=zones, mask_mode=mask_mode) + with ClusterGenerator("multi", nb_workers=2) as cluster: + config = MultiprocConfig(chunks=(2, 3), cluster=cluster) + result = source.stats(statistics, mask=zones, mask_mode=mask_mode, mp_config=config) + attribute = source.stats("mean", values={"rating": (zones, "rating")}, mp_config=config) + + # Both selections contain five finite values, while the vector attribute is constant where defined + assert result == pytest.approx(expected) + assert result["validcount"] == 10 and result["validinliercount"] == 5 + assert result["totalinliercount"] == 6 and attribute == 10 + assert not source.is_loaded + + def test_stats__raster_values_at_file_points(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Checks that Multiproc interpolates raster values at file point locations while both sources stay unloaded.""" + + # Create a raster and points on distant, exactly known raster cells + values = np.arange(48, dtype=float).reshape(6, 8) + rows, columns = np.array([1, 4, 2, 3, 1, 4]), np.array([1, 6, 4, 2, 5, 3]) + raster_filename, point_filename = tmp_path / "values.tif", tmp_path / "locations.gpkg" + raster = gu.Raster.from_array(values, from_origin(500000, 5100006, 1, 1), 32633) + x, y = raster.ij2xy(rows, columns) + points = gpd.GeoDataFrame( + {"height": np.zeros(len(rows))}, + geometry=gpd.points_from_xy(x, y), + crs=32633, + ) + + # Write the raster and point files to disk with values at the matching locations + raster.to_file(raster_filename) + points.to_file(point_filename, index=False) + source, support = gu.Raster(raster_filename), gu.PointCloud(point_filename, data_column="height") + expected = raster.stats("mean", at=gu.PointCloud(points, data_column="height"), interpolation="nearest") + + # Reject full-file loading so Multiproc must read raster windows and matching point rows + def reject_load(*args: Any, **kwargs: Any) -> None: + """Reject full loading when only selected point locations are needed.""" + + raise AssertionError("Spatial statistics must read bounded file blocks.") + + with monkeypatch.context() as guarded: + guarded.setattr(gu.Raster, "load", reject_load) + guarded.setattr(gu.PointCloud, "load", reject_load) + with ClusterGenerator("multi", nb_workers=2) as cluster: + config = MultiprocConfig(chunks=(2, 3), cluster=cluster) + result = source.stats("mean", at=support, interpolation="nearest", mp_config=config) + + # Compare the Multiproc mean with eager stats() and the exact raster values at the point locations + assert result == pytest.approx(expected) + assert result == pytest.approx(values[rows, columns].mean()) + assert not source.is_loaded and not support.is_loaded + + @pytest.mark.parametrize("reordered", [False, True]) + def test_stats__point_files_on_common_support(self, reordered: bool, tmp_path: Path) -> None: + """Checks that distinct point files must have the same ordered coordinates before their values are reduced.""" + + # Create two point datasets with matching coordinates and optionally reorder the common support + positions = np.arange(8) + values = gpd.GeoDataFrame( + {"height": positions + 10.0}, + geometry=gpd.points_from_xy(500000 + positions, 5100000 + positions), + crs=32633, + ) + locations = values.copy() + locations["height"] = 0.0 + if reordered: + locations = locations.iloc[[0, 1, 3, 2, 4, 5, 6, 7]] + + # Write both point datasets to disk with unrelated values in their data columns + value_filename, point_filename = tmp_path / "values.gpkg", tmp_path / "locations.gpkg" + values.to_file(value_filename, index=False) + locations.to_file(point_filename, index=False) + source = gu.PointCloud(value_filename, data_column="height") + support = gu.PointCloud(point_filename, data_column="height") + + # Compare all ordered coordinates in Multiproc row blocks without loading either complete file + with ClusterGenerator("multi", nb_workers=2) as cluster: + config = MultiprocConfig(chunks=3, cluster=cluster) + if reordered: + with pytest.raises(ValueError, match="ordered support coordinates"): + source.stats("mean", at=support, mp_config=config) + else: + result = source.stats("mean", at=support, mp_config=config) + assert result == pytest.approx(values.height.mean()) + assert not source.is_loaded and not support.is_loaded + + def test_stats__file_raster_alignment(self, tmp_path: Path) -> None: + """Checks that Multiproc aligns an unloaded raster to another raster before calculating statistics.""" + + # Write two raster files to disk with a non-default value band and a coarser common support grid + first = np.arange(96, dtype=float).reshape(8, 12) + values = np.stack((first, first * 3 + 100)) + source_filename, reference_filename = tmp_path / "values.tif", tmp_path / "reference.tif" + raster = gu.Raster.from_array(values, from_origin(500000, 5100008, 1, 1), 32633, nodata=-9999) + reference = gu.Raster.from_array(np.zeros((4, 6)), from_origin(500000, 5100008, 2, 2), 32633) + raster.to_file(source_filename) + reference.to_file(reference_filename) + source, support = gu.Raster(source_filename), gu.Raster(reference_filename) + expected = raster.stats(["mean", "std"], values=2, at=reference, align="reproject") + + # Reproject to a temporary file and read its selected band in separate reduction workers + with ClusterGenerator("multi", nb_workers=2) as cluster: + config = MultiprocConfig(chunks=(3, 4), cluster=cluster) + result = source.stats(["mean", "std"], values=2, at=support, align="reproject", mp_config=config) + + # The completed statistics match eager alignment, and original sources remain available on disk + assert result == pytest.approx(expected) + assert not source.is_loaded and not support.is_loaded + assert source_filename.exists() and reference_filename.exists() diff --git a/tests/test_stats/test_stats.py b/tests/test_stats/test_stats.py index b2ec5b859..1cb6a871f 100644 --- a/tests/test_stats/test_stats.py +++ b/tests/test_stats/test_stats.py @@ -5,13 +5,15 @@ from typing import Any import numpy as np +import pandas as pd import pytest import rasterio as rio import geoutils as gu from geoutils import examples from geoutils._typing import NDArrayNum -from geoutils.stats.stats import ( +from geoutils.multiproc import MultiprocConfig +from geoutils.stats.reduction import ( _STATS_ALIAS_ALL, _STATS_ALIAS_CALLABLE, _STATS_ALIAS_GEN, @@ -37,16 +39,147 @@ class TestStats: landsat_rgb_path = examples.get_path_test("everest_landsat_rgb") aster_dem_path = examples.get_path_test("exploradores_aster_dem") + def test_stats__summary_and_grouped_routes(self) -> None: + """Checks that by chooses between a summary and a grouped table.""" + + # Create a small raster and two declared land-cover categories + values = np.array([[1.0, 2.0], [3.0, 4.0]]) + landcover = np.array([[0, 0], [1, 1]]) + raster = gu.Raster.from_array( + values, + transform=rio.transform.from_origin(0, 2, 1, 1), + crs=4326, + ) + + # Calculate the default, scalar, and selected summaries through the common method + summary = raster.stats() + mean = raster.stats("mean") + selected = raster.stats(["mean", "std", "nmad"]) + array_mean = gu.stats.stats(values, "mean") + + # Calculate the same selected statistics separately for each category + grouped = raster.stats( + ["mean", "std"], + by={"landcover": landcover}, + categories={"landcover": [0, 1]}, + ) + array_grouped = gu.stats.stats( + values, + ["mean", "std"], + by={"landcover": landcover}, + categories={"landcover": [0, 1]}, + ) + default_grouped = raster.stats( + by={"landcover": landcover}, + categories={"landcover": [0, 1]}, + ) + + # Check the summary, grouped routes, and shared default statistic selection + assert summary["Mean"] == pytest.approx(values.mean()) + assert mean == pytest.approx(values.mean()) + assert array_mean == pytest.approx(values.mean()) + assert selected == pytest.approx({"mean": values.mean(), "std": values.std(), "nmad": 1.4826}) + np.testing.assert_allclose(grouped["band_1"], [[2, 1.5, 0.5], [2, 3.5, 0.5]]) + np.testing.assert_allclose(array_grouped["value"], grouped["band_1"]) + expected_grouped_statistics = ["count", *[name for name in _STATS_LIST_MIN if name != "validcount"]] + assert default_grouped.columns.get_level_values("statistic").tolist() == expected_grouped_statistics + + @pytest.mark.parametrize("subsampling_strategy", ["topk", "sequential"]) + def test_stats__subsample_per_group_without_by(self, subsampling_strategy: str) -> None: + """Checks that per-group sampling without groups behaves like ordinary summary sampling.""" + + # Use a shared mask so the sample must exclude locations before selecting any values + values = pd.Series(np.arange(30, dtype=float)) + mask = values.to_numpy() % 3 != 0 + options = {"mask": mask, "subsample": 7, "random_state": 42, "subsampling_strategy": subsampling_strategy} + + # Compare the option against the established summary path with the same seed + expected = gu.stats.stats(values, ["mean", "std", "totalcount"], **options) + result = gu.stats.stats(values, ["mean", "std", "totalcount"], subsample_per_group=True, **options) + assert result == expected + assert result["totalcount"] == 7 + + @pytest.mark.parametrize("grouped", [False, True]) + @pytest.mark.parametrize("invalid", ["False", 1, None]) + def test_stats__error_subsample_per_group_validation(self, grouped: bool, invalid: object) -> None: + """Checks that both stats() routes reject non-boolean per-group sampling options.""" + + # A string such as 'False' must not accidentally enable sampling within groups + values = pd.Series(np.arange(4, dtype=float)) + by = {"group": values > 1} if grouped else None + with pytest.raises(TypeError, match="Argument ``subsample_per_group`` must be a boolean"): + gu.stats.stats(values, "mean", by=by, subsample_per_group=invalid) + + @pytest.mark.parametrize("as_list", [False, True]) + def test_stats__empty_mask_callable_result(self, as_list: bool) -> None: + """Checks that a callable's result stays masked when no selected values remain.""" + + # NumPy Masked mean returns a masked scalar for an entirely excluded array + values = np.arange(6, dtype=float).reshape(2, 3) + keep = np.zeros(values.shape, dtype=bool) + statistics = [np.ma.mean] if as_list else np.ma.mean + + # Check the scalar and named-result routes without allowing scalar conversion to expose masked storage + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="Empty raster") + result = gu.stats.stats(values, statistics, mask=keep) + assert np.ma.is_masked(result["mean"] if as_list else result) + + @pytest.mark.parametrize("grouped", [False, True]) + def test_stats__generator_for_multiple_values(self, grouped: bool) -> None: + """Checks that a generator requests the same statistics for every selected value array.""" + + # Give the second named array a distinct mean and the exact same spread + first = np.arange(6, dtype=float) + values = {"first": first, "second": first + 100} + options = {} + if grouped: + options = {"by": {"zone": np.array([0, 0, 0, 1, 1, 1])}, "categories": {"zone": [0, 1]}} + + # Pass a generator that can only be consumed once while selecting both arrays + statistics = (name for name in ["mean", "std"]) + result = gu.stats.stats(values, statistics, **options) + + # Check each array independently so an exhausted request cannot silently omit its statistics + for name, array in values.items(): + if grouped: + expected = [[3, array[:3].mean(), array[:3].std()], [3, array[3:].mean(), array[3:].std()]] + np.testing.assert_allclose(result[name], expected) + else: + assert result[name] == pytest.approx({"mean": array.mean(), "std": array.std()}) + + @pytest.mark.parametrize("source_type", ["raster", "pointcloud"]) + def test_get_stats__deprecated_forwarding(self, source_type: str) -> None: + """Checks that get_stats() warns about deprecation and forwards its selections to stats().""" + + # Select a second raster band and a mask, or the active values of a point cloud + values = np.arange(1, 7, dtype=float).reshape(2, 3) + statistics = ["mean", "std"] + old_options: dict[str, Any] = {"stats_name": statistics} + new_options: dict[str, Any] = {"statistics": statistics} + if source_type == "raster": + source = gu.Raster.from_array(np.stack((values, values + 100)), rio.transform.from_origin(0, 2, 1, 1), 4326) + keep = np.array([[True, True, False], [True, False, True]]) + old_options.update(band=2, inlier_mask=keep) + new_options.update(values=2, mask=keep) + else: + source = gu.PointCloud.from_xyz(np.arange(values.size), np.zeros(values.size), values.ravel(), crs=4326) + + # Require an explicit deprecation warning while comparing with the canonical API + expected = source.stats(**new_options) + with pytest.deprecated_call(match="get_stats"): + deprecated = source.get_stats(**old_options) + + # Check that the requested statistic names and all numerical results are unchanged + compare_dict(expected, deprecated) + @pytest.mark.parametrize("example", [landsat_b4_path, aster_dem_path]) - def test_get_stats_raster_one_band(self, example: str) -> None: - """ - Verify get_stats() method for a raster, especially output stats for different inputs - parameters and some stats. - """ + def test_stats__raster_one_band(self, example: str) -> None: + """Checks raster stats for various statistic names, masks, callables and empty data.""" raster = gu.Raster(example) # Default stats - stats = raster.get_stats() + stats = raster.stats() assert len(stats) == len(_STATS_LIST_MIN) assert list(stats.keys()) == [_STATS_ALIAS_ALL[key] for key in _STATS_LIST_MIN] for name in _STATS_LIST_MIN: @@ -54,7 +187,7 @@ def test_get_stats_raster_one_band(self, example: str) -> None: assert isinstance(stats.get(_STATS_ALIAS_ALL[name]), stat_types) # Full stats - stats = raster.get_stats("all") + stats = raster.stats("all") assert len(stats) == len(_STATS_ALIAS_GEN) for name in _STATS_ALIAS_GEN.values(): assert name in stats @@ -62,7 +195,7 @@ def test_get_stats_raster_one_band(self, example: str) -> None: # With mask (inlier=True) inlier_mask = ~raster.get_mask() - stats_masked = raster.get_stats("all", inlier_mask=inlier_mask) + stats_masked = raster.stats("all", mask=inlier_mask) assert len(stats_masked) == len(_STATS_ALIAS_ALL) assert list(stats_masked.keys()) == [_STATS_ALIAS_ALL[key] for key in _STATS_ALIAS_ALL] for name in _STATS_ALIAS_MASK.values(): @@ -71,38 +204,32 @@ def test_get_stats_raster_one_band(self, example: str) -> None: assert stats_masked == stats # Print of the values - stats = raster.get_stats("all", inlier_mask=inlier_mask) + stats = raster.stats("all", mask=inlier_mask) for stat in stats: assert not isinstance(stat, np.generic) for stat in _STATS_ALIAS_ALL: - assert not isinstance(raster.get_stats(stat, inlier_mask=inlier_mask), np.generic) + assert not isinstance(raster.stats(stat, mask=inlier_mask), np.generic) # With mask (inlier=True) and default list - stats_masked = raster.get_stats(inlier_mask=inlier_mask) + stats_masked = raster.stats(mask=inlier_mask) assert len(stats_masked) == len(_STATS_LIST_MIN) assert list(stats_masked.keys()) == [_STATS_ALIAS_ALL[key] for key in _STATS_LIST_MIN] for name in _STATS_LIST_MIN: assert _STATS_ALIAS_ALL[name] in stats_masked # Test case sensitive + space/underscore possibilities - stats_masked = raster.get_stats(inlier_mask=inlier_mask) + stats_masked = raster.stats(mask=inlier_mask) name = "Standard deviation" - assert stats_masked["Standard deviation"] == raster.get_stats( - stats_name="standard deviation", inlier_mask=inlier_mask - ) - assert stats_masked["Standard deviation"] == raster.get_stats( - stats_name="standarddeviation", inlier_mask=inlier_mask - ) - assert stats_masked["Standard deviation"] == raster.get_stats( - stats_name="standard_deviation", inlier_mask=inlier_mask - ) - assert stats_masked[name] == raster.get_stats(stats_name="standard_deviation", inlier_mask=inlier_mask) + assert stats_masked["Standard deviation"] == raster.stats(statistics="standard deviation", mask=inlier_mask) + assert stats_masked["Standard deviation"] == raster.stats(statistics="standarddeviation", mask=inlier_mask) + assert stats_masked["Standard deviation"] == raster.stats(statistics="standard_deviation", mask=inlier_mask) + assert stats_masked[name] == raster.stats(statistics="standard_deviation", mask=inlier_mask) # Empty mask (=False) empty_mask = np.zeros_like(inlier_mask) with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="Empty raster") - stats_masked = raster.get_stats("all", inlier_mask=empty_mask) + stats_masked = raster.stats("all", mask=empty_mask) assert len(stats_masked) == len(_STATS_ALIAS_ALL) for name in _STATS_ALIAS_CALLABLE.values(): assert np.isnan(stats_masked.get(name)) @@ -114,26 +241,26 @@ def test_get_stats_raster_one_band(self, example: str) -> None: for stat in _STATS_ALIAS_ALL: with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="Empty raster") - stats_masked = raster.get_stats(inlier_mask=empty_mask, stats_name=stat) + stats_masked = raster.stats(mask=empty_mask, statistics=stat) with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="Empty raster") - stats_masked = raster.get_stats(inlier_mask=empty_mask, stats_name="mean") + stats_masked = raster.stats(mask=empty_mask, statistics="mean") assert np.isnan(stats_masked) with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="Empty raster") - stats_masked = raster.get_stats(inlier_mask=empty_mask, stats_name="valid_count") + stats_masked = raster.stats(mask=empty_mask, statistics="valid_count") assert stats_masked == stats.get("Valid count") with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="Empty raster") - stats_masked = raster.get_stats(inlier_mask=empty_mask, stats_name="Valid inlier count") + stats_masked = raster.stats(mask=empty_mask, statistics="Valid inlier count") assert stats_masked == 0 with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="Empty raster") - stats_masked = raster.get_stats("all", inlier_mask=inlier_mask) + stats_masked = raster.stats("all", mask=inlier_mask) for name in stats_masked: - assert stats_masked[name] == raster.get_stats(stats_name=name.lower(), inlier_mask=inlier_mask) - assert stats_masked[name] == raster.get_stats(stats_name="".join(name.split()), inlier_mask=inlier_mask) + assert stats_masked[name] == raster.stats(statistics=name.lower(), mask=inlier_mask) + assert stats_masked[name] == raster.stats(statistics="".join(name.split()), mask=inlier_mask) # Empty DEM dem_empty = gu.Raster.from_array( @@ -143,7 +270,7 @@ def test_get_stats_raster_one_band(self, example: str) -> None: ) with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="Empty raster") - stats_empty = dem_empty.get_stats("all") + stats_empty = dem_empty.stats("all") assert len(stats_empty) == len(_STATS_ALIAS_GEN) for name in _STATS_ALIAS_CALLABLE.values(): assert np.isnan(stats_empty.get(name)) @@ -153,14 +280,14 @@ def test_get_stats_raster_one_band(self, example: str) -> None: # Single stat for name in _STATS_ALIAS_GEN: - stat = raster.get_stats(stats_name=name) + stat = raster.stats(statistics=name) assert np.isfinite(stat) for name in _STATS_ALIAS_MASK: - stat = raster.get_stats(stats_name=name) + stat = raster.stats(statistics=name) assert np.isnan(stat) # Alias stat - assert raster.get_stats(stats_name="Valid count") == raster.get_stats(stats_name="valid_count") + assert raster.stats(statistics="Valid count") == raster.stats(statistics="valid_count") # Callable def percentile_95(data: NDArrayNum) -> np.floating[Any]: @@ -168,14 +295,12 @@ def percentile_95(data: NDArrayNum) -> np.floating[Any]: data = data.compressed() return np.nanpercentile(data, 95) - stat = raster.get_stats(stats_name=percentile_95) + stat = raster.stats(statistics=percentile_95) assert isinstance(stat, np.floating) # Selected stats and callable stats_name = ["mean", "max", "std", "validinliercount", "percentile_95"] - stats = raster.get_stats( - stats_name=["mean", "max", "std", "validinliercount", percentile_95], inlier_mask=inlier_mask - ) + stats = raster.stats(statistics=["mean", "max", "std", "validinliercount", percentile_95], mask=inlier_mask) assert len(stats) == len(stats_name) for name in stats_name: assert name in stats @@ -184,46 +309,45 @@ def percentile_95(data: NDArrayNum) -> np.floating[Any]: # Non-existing stats with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Statistic name 80 percentile is not recognized") - stat = raster.get_stats(stats_name="80 percentile") + stat = raster.stats(statistics="80 percentile") assert isnan(stat) with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Statistic name 42 is a not recognized string") - stat = raster.get_stats(stats_name=42) + stat = raster.stats(statistics=42) assert stat is None # IQR (scipy) validation with numpy nan_arr = raster.get_nanarray() if nan_arr.ndim == 3: nan_arr = nan_arr[0, :, :] - assert raster.get_stats(stats_name="iqr") == pytest.approx( + assert raster.stats(statistics="iqr") == pytest.approx( np.nanpercentile(nan_arr, 75) - np.nanpercentile(nan_arr, 25) ) @pytest.mark.parametrize("example", [landsat_rgb_path]) - def test_get_stats_multi_bands(self, example: str) -> None: + def test_stats__multiple_raster_bands(self, example: str) -> None: + """Checks that stats() calculates each band of a multiband raster separately.""" + raster = gu.Raster(example) - stats = raster.get_stats() + stats = raster.stats() assert list(stats.keys()) == ["band 1", "band 2", "band 3"] data = raster.get_nanarray() for band in range(1, raster.count + 1): assert stats["band " + str(band)]["Mean"] == pytest.approx(np.nanmean(data[band - 1])) - stats = raster.get_stats("mean") + stats = raster.stats("mean") for band in range(1, raster.count + 1): assert stats["band " + str(band)] == pytest.approx(np.nanmean(data[band - 1])) @pytest.mark.parametrize("example", [landsat_b4_path, aster_dem_path]) - def test_get_stats_raster_pointcloud(self, example: str) -> None: - """ - Verify get_stats() method for a raster converted to pointcloud, especially output stats for different inputs - parameters. - """ + def test_stats__raster_pointcloud(self, example: str) -> None: + """Checks statistics for a raster converted to a point cloud.""" raster = gu.Raster(example) pointcloud = raster.to_pointcloud() # Default stats - stats = pointcloud.get_stats() + stats = pointcloud.stats() assert len(stats) == len(_STATS_LIST_MIN) assert list(stats.keys()) == [_STATS_ALIAS_ALL[key] for key in _STATS_LIST_MIN] for name in _STATS_LIST_MIN: @@ -231,7 +355,7 @@ def test_get_stats_raster_pointcloud(self, example: str) -> None: assert isinstance(stats.get(_STATS_ALIAS_GEN[name]), stat_types) # Full stats - stats = pointcloud.get_stats("all") + stats = pointcloud.stats("all") assert len(stats) == len(_STATS_ALIAS_GEN) assert list(stats.keys()) == [_STATS_ALIAS_GEN[key] for key in _STATS_ALIAS_GEN] for name in _STATS_ALIAS_GEN.values(): @@ -240,18 +364,18 @@ def test_get_stats_raster_pointcloud(self, example: str) -> None: # Single stat for name in _STATS_ALIAS_GEN: - stat = pointcloud.get_stats(stats_name=name) + stat = pointcloud.stats(statistics=name) assert np.isfinite(stat) for name in _STATS_ALIAS_MASK: - stat = pointcloud.get_stats(stats_name=name) + stat = pointcloud.stats(statistics=name) assert np.isnan(stat) # Print of the values - stats = pointcloud.get_stats("all") + stats = pointcloud.stats("all") for stat in stats: assert not isinstance(stat, np.generic) for stat in _STATS_ALIAS_ALL: - assert not isinstance(pointcloud.get_stats(stat), np.generic) + assert not isinstance(pointcloud.stats(stat), np.generic) # Callable def percentile_95(data: NDArrayNum) -> np.floating[Any]: @@ -261,7 +385,7 @@ def percentile_95(data: NDArrayNum) -> np.floating[Any]: # Selected stats and callable stats_name = ["mean", "max", "std", "percentile_95"] - stats = pointcloud.get_stats(stats_name=["mean", "max", "std", percentile_95]) + stats = pointcloud.stats(statistics=["mean", "max", "std", percentile_95]) assert len(stats) == len(stats_name) for name in stats_name: assert name in stats @@ -270,12 +394,12 @@ def percentile_95(data: NDArrayNum) -> np.floating[Any]: # Non-existing stats with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Statistic name 80 percentile is not recognized") - stat = pointcloud.get_stats(stats_name="80 percentile") + stat = pointcloud.stats(statistics="80 percentile") assert isnan(stat) with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Statistic name 42 is a not recognized string") - stat = pointcloud.get_stats(stats_name=42) + stat = pointcloud.stats(statistics=42) assert stat is None # Empty mask (=False) @@ -286,7 +410,7 @@ def percentile_95(data: NDArrayNum) -> np.floating[Any]: pointcloud = raster.to_pointcloud() with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Empty raster") - stats_masked = pointcloud.get_stats("all") + stats_masked = pointcloud.stats("all") assert len(stats_masked) == len(_STATS_ALIAS_GEN) for name in _STATS_ALIAS_CALLABLE.values(): @@ -295,10 +419,8 @@ def percentile_95(data: NDArrayNum) -> np.floating[Any]: assert stats_masked.get("Total count") == 0 assert isnan(stats_masked.get("Percentage valid points")) - def test_raster_get_stats_values(self) -> None: - """ - Verify the output statistics values of a raster. - """ + def test_stats__raster_values(self) -> None: + """Checks the output statistics values of a raster.""" filename_rast = gu.examples.get_path("everest_landsat_b4") filename_vect = gu.examples.get_path("everest_rgi_outlines") rast = gu.Raster(filename_rast) @@ -312,18 +434,18 @@ def test_raster_get_stats_values(self) -> None: "Max": np.uint8(255), "Min": np.uint8(13), "Sum": np.uint64(75479373), - "Sum of squares": np.uint64(44549637), + "Sum of squares": np.uint64(14179501317), "90th percentile": np.float64(255.0), "LE90": np.float64(218.0), "IQR": np.float64(164.0), "NMAD": np.float64(94.8864), - "RMSE": np.float64(9.220541807365446), + "RMSE": np.float64(164.49959579638966), "Standard deviation": np.float64(79.44349437534403), "Valid count": 524000, "Total count": 524000, "Percentage valid points": np.float64(100.0), } - compare_dict(res_stats, rast.get_stats("all")) + compare_dict(res_stats, rast.stats("all")) # Verify raster stats with a mask res_stats_mask = { @@ -332,12 +454,12 @@ def test_raster_get_stats_values(self) -> None: "Max": np.uint8(255), "Min": np.uint8(13), "Sum": np.uint64(26650493), - "Sum of squares": np.uint64(24696991), + "Sum of squares": np.uint64(3963154847), "90th percentile": np.float64(225.0), "LE90": np.float64(223.0), "IQR": np.float64(83.0), "NMAD": np.float64(54.856199999999994), - "RMSE": np.float64(10.118943490060417), + "RMSE": np.float64(128.18395566310244), "Standard deviation": np.float64(64.98157041836747), "Valid count": 524000, "Total count": 524000, @@ -347,7 +469,7 @@ def test_raster_get_stats_values(self) -> None: "Percentage inlier points": np.float64(46.03015267175572), "Percentage valid inlier points": np.float64(100.0), } - compare_dict(res_stats_mask, rast.get_stats("all", inlier_mask=inlier_mask)) + compare_dict(res_stats_mask, rast.stats("all", mask=inlier_mask)) # Verify cropped raster nrows, ncols = rast.shape @@ -358,18 +480,18 @@ def test_raster_get_stats_values(self) -> None: "Max": np.uint8(255), "Min": np.uint8(14), "Sum": np.uint64(40594831), - "Sum of squares": np.uint64(22875807), + "Sum of squares": np.uint64(7754447263), "90th percentile": np.float64(255.0), "LE90": np.float64(218.0), "IQR": np.float64(166.0), "NMAD": np.float64(105.26459999999999), - "RMSE": np.float64(9.153915273540871), + "RMSE": np.float64(168.53655012767328), "Standard deviation": np.float64(79.32951386752386), "Valid count": 273000, "Total count": 273000, "Percentage valid points": np.float64(100.0), } - compare_dict(res_stats_crop, rast_crop.get_stats("all")) + compare_dict(res_stats_crop, rast_crop.stats("all")) # Verify reprojected raster with warnings.catch_warnings(): @@ -382,18 +504,18 @@ def test_raster_get_stats_values(self) -> None: "Max": np.uint8(254), "Min": np.uint8(14), "Sum": np.uint64(24919216), - "Sum of squares": np.uint64(22814334), + "Sum of squares": np.uint64(3757165438), "90th percentile": np.float64(218.0), "LE90": np.float64(204.0), "IQR": np.float64(93.0), "NMAD": np.float64(66.717), - "RMSE": np.float64(10.38534653788665), + "RMSE": np.float64(133.27455905093126), "Standard deviation": np.float64(62.319986152883956), "Valid count": 211527, "Total count": 524000, "Percentage valid points": np.float64(40.36774809160305), } - compare_dict(res_stats_crop_proj, rast_crop_proj.get_stats("all")) + compare_dict(res_stats_crop_proj, rast_crop_proj.stats("all")) # Verify stats of a masked raster rast.set_mask(inlier_mask) @@ -403,18 +525,18 @@ def test_raster_get_stats_values(self) -> None: "Max": np.uint8(255), "Min": np.uint8(15), "Sum": np.uint64(48828880), - "Sum of squares": np.uint64(19852646), + "Sum of squares": np.uint64(10216346470), "90th percentile": np.float64(255.0), "LE90": np.float64(209.0), "IQR": np.float64(156.0), "NMAD": np.float64(99.3342), - "RMSE": np.float64(8.37853254688833), + "RMSE": np.float64(190.06693359773863), "Standard deviation": np.float64(79.45825061580675), "Valid count": 282802, "Total count": 524000, "Percentage valid points": np.float64(53.96984732824428), } - compare_dict(stats_masked_rast, rast.get_stats("all")) + compare_dict(stats_masked_rast, rast.stats("all")) # Verify stats of a masked raster with the other part covered by the inler_mask (=> empty raster) stats_masked_rast_masked = { @@ -440,12 +562,10 @@ def test_raster_get_stats_values(self) -> None: } with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="Empty raster") - compare_dict(stats_masked_rast_masked, rast.get_stats("all", inlier_mask=inlier_mask)) + compare_dict(stats_masked_rast_masked, rast.stats("all", mask=inlier_mask)) - def test_pointcloud_get_stats_values(self) -> None: - """ - Verify the output statistics values of a pointcloud. - """ + def test_stats__pointcloud_values(self) -> None: + """Checks the output statistics values of a pointcloud.""" filename_rast = gu.examples.get_path("everest_landsat_b4") rast = gu.Raster(filename_rast) @@ -458,18 +578,18 @@ def test_pointcloud_get_stats_values(self) -> None: "Max": np.uint8(255), "Min": np.uint8(13), "Sum": np.uint64(75479373), - "Sum of squares": np.uint64(44549637), + "Sum of squares": np.uint64(14179501317), "90th percentile": np.float64(255.0), "LE90": np.float64(218.0), "IQR": np.float64(164.0), "NMAD": np.float64(94.8864), - "RMSE": np.float64(9.220541807365446), + "RMSE": np.float64(164.49959579638966), "Standard deviation": np.float64(79.44349437534403), "Valid count": 524000, "Total count": 524000, "Percentage valid points": np.float64(100.0), } - compare_dict(rast_stats_pc, rast_pc.get_stats("all")) + compare_dict(rast_stats_pc, rast_pc.stats("all")) # Verify cropped raster pc nrows, ncols = rast.shape @@ -482,18 +602,18 @@ def test_pointcloud_get_stats_values(self) -> None: "Max": np.uint8(255), "Min": np.uint8(14), "Sum": np.uint64(40594831), - "Sum of squares": np.uint64(22875807), + "Sum of squares": np.uint64(7754447263), "90th percentile": np.float64(255.0), "LE90": np.float64(218.0), "IQR": np.float64(166.0), "NMAD": np.float64(105.26459999999999), - "RMSE": np.float64(9.153915273540871), + "RMSE": np.float64(168.53655012767328), "Standard deviation": np.float64(79.32951386752386), "Valid count": 273000, "Total count": 273000, "Percentage valid points": np.float64(100.0), } - compare_dict(rast_stats_crop_pc, rast_crop_pc.get_stats("all")) + compare_dict(rast_stats_crop_pc, rast_crop_pc.stats("all")) # Verify reprojected raster pc with warnings.catch_warnings(): @@ -508,15 +628,75 @@ def test_pointcloud_get_stats_values(self) -> None: "Max": np.uint8(254), "Min": np.uint8(14), "Sum": np.uint64(24919216), - "Sum of squares": np.uint64(22814334), + "Sum of squares": np.uint64(3757165438), "90th percentile": np.float64(218.0), "LE90": np.float64(204.0), "IQR": np.float64(93.0), "NMAD": np.float64(66.717), - "RMSE": np.float64(10.38534653788665), + "RMSE": np.float64(133.27455905093126), "Standard deviation": np.float64(62.319986152883956), "Valid count": 211527, "Total count": 211527, "Percentage valid points": np.float64(100.0), } - compare_dict(rast_stats_crop_proj_pc, rast_crop_proj_pc.get_stats("all")) + compare_dict(rast_stats_crop_proj_pc, rast_crop_proj_pc.stats("all")) + + +class TestStatsChunked: + """Checks stats() loading behavior and exact results with chunked inputs.""" + + @pytest.mark.parametrize("masked", [False, True]) + def test_stats__multiprocessing_summary(self, masked: bool) -> None: + """Checks that stats() combines ungrouped array tiles and optional masks in worker processes.""" + + # Include nodata values and uneven edge tiles in both the array and raster routes + from geoutils.multiproc.cluster import MpCluster + + values = np.arange(35, dtype=float).reshape(5, 7) + values[1, 2] = np.nan + values[3, 5] = np.nan + raster = gu.Raster.from_array( + values, + transform=rio.transform.from_origin(0, 5, 1, 1), + crs=4326, + nodata=np.nan, + ) + keep = np.indices(values.shape).sum(axis=0) % 3 != 0 if masked else None + expected_summary = gu.stats.stats(values, mask=keep) + expected_all = gu.stats.stats(values, "all", mask=keep) + expected_reductions = gu.stats.stats(values, ["mean", "std", "sum", "validcount"], mask=keep) + + # Run exact default statistics and mergeable reductions through real worker processes + with MpCluster({"nb_workers": 2}) as cluster: + config = MultiprocConfig(chunks=(2, 3), cluster=cluster) + array_summary = gu.stats.stats(values, mask=keep, mp_config=config) + array_all = gu.stats.stats(values, "all", mask=keep, mp_config=config) + raster_summary = raster.stats(mask=keep, mp_config=config) + reductions = raster.stats(["mean", "std", "sum", "validcount"], mask=keep, mp_config=config) + + # Match the eager values and preserve integer count results in both public APIs + compare_dict(expected_summary, array_summary) + compare_dict(expected_all, array_all) + compare_dict(expected_summary, raster_summary) + compare_dict(expected_reductions, reductions) + assert isinstance(array_summary["Valid count"], int) + assert isinstance(array_summary["Total count"], int) + + def test_stats__custom_summary_preserves_shape(self) -> None: + """Checks that Multiproc custom summaries match eager results and see the complete array shape.""" + + # Use a rectangular array whose full shape differs from every worker tile and from a flattened array + values = np.arange(12, dtype=float).reshape(3, 4) + config = MultiprocConfig(chunks=(2, 3)) + + # Request the same custom summary eagerly and through Multiproc chunks + expected_shape = gu.stats.stats(values, np.shape) + expected_combined = gu.stats.stats(values, [np.shape, "mean"]) + shape = gu.stats.stats(values, np.shape, mp_config=config) + combined = gu.stats.stats(values, [np.shape, "mean"], mp_config=config) + + # Preserve the complete shape and exactly match both eager output forms + assert shape == expected_shape + assert combined == expected_combined + assert shape == combined["shape"] == values.shape + assert combined["mean"] == pytest.approx(values.mean()) diff --git a/tests/test_stats/test_variography.py b/tests/test_stats/test_variography.py new file mode 100644 index 000000000..dd24cfa66 --- /dev/null +++ b/tests/test_stats/test_variography.py @@ -0,0 +1,403 @@ +"""Tests for estimating, fitting, saving, and converting variograms.""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +import xarray as xr +from rasterio.transform import from_origin + +import geoutils as gu +from geoutils._typing import NDArrayNum +from geoutils.stats.variography import VariogramModel + + +@pytest.fixture(autouse=True) +def _writable_matplotlib_config(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Keep optional plotting imports inside the test workspace.""" + + monkeypatch.setenv("MPLCONFIGDIR", str(tmp_path)) + + +class TestVariogramStorage: + """Checks the compact Variogram result and its stored distance bins. + + The methods cover optional imports, immutable arrays, serialization, pair release, and stable bin limits. + """ + + def test_importing_geoutils_does_not_load_variogram_backends(self) -> None: + """Checks that importing GeoUtils does not import optional variogram packages.""" + + # Build a fresh Python process that imports only GeoUtils + code = ( + "import sys\n" + "import geoutils\n" + "assert not {'skgstat', 'gstools', 'gpytorch', 'torch'}.intersection(sys.modules)\n" + ) + # Check inside that process that none of the optional packages was loaded + subprocess.run([sys.executable, "-c", code], check=True) + + def test_variogram_is_small_immutable_and_serializable(self) -> None: + """Checks that Variogram owns read-only arrays and survives JSON and Xarray conversion.""" + + # Create a complete measured and fitted result from small arrays and plain model details + result = gu.Variogram( + lags=np.array([1.0, 2.0]), + semivariance=np.array([0.2, 0.5]), + counts=np.array([10, 8]), + semivariance_error=np.array([0.02, 0.03]), + bin_lower_edges=np.array([0.5, 1.5]), + bin_edges=np.array([1.5, 2.5]), + fitted_semivariance=np.array([0.25, 0.45]), + model=VariogramModel("gaussian", effective_range=4, partial_sill=0.8, nugget=0.1), + estimator="matheron", + ) + + # Check that callers cannot change stored distance values in place + with pytest.raises(ValueError, match="read-only"): + result.lags[0] = 3 + + # Convert through JSON and check the restored model and optional error values + restored = gu.Variogram.from_dict(json.loads(json.dumps(result.to_dict()))) + assert restored.model == result.model + assert restored.semivariance_error is not None and result.semivariance_error is not None + assert np.array_equal(restored.semivariance_error, result.semivariance_error) + # Check that Xarray contains every per-distance value + assert set(restored.to_xarray().data_vars) == { + "semivariance", + "semivariance_error", + "count", + "bin_lower_edge", + "bin_edge", + "fitted_semivariance", + } + + def test_from_pairs_discards_pair_data(self) -> None: + """Checks that from_pairs() keeps per-distance results and releases individual pairs.""" + + # Create twenty pairs with a constant endpoint difference of two + values = np.column_stack((np.arange(20, dtype=float), np.arange(20, dtype=float) + 2)) + pairs = xr.Dataset( + {"value": (("pair", "endpoint"), values), "distance": ("pair", np.linspace(1, 20, 20))}, + coords={"pair": np.arange(20), "endpoint": ["first", "second"]}, + ) + + # Reduce the pairs into four log-spaced distance bins + result = gu.Variogram.from_pairs(pairs, estimator="matheron", bins="log", n_lags=4) + + # Check the bin count, total pair count, semivariance, and released source object + assert len(result.lags) == 4 + assert np.sum(result.counts) == 20 + assert result.backend_object is None + assert result.semivariance == pytest.approx(np.full(4, 2.0)) + + def test_from_pairs_uses_sampled_distance_limits_for_stable_bins(self) -> None: + """Checks that from_pairs() uses requested distance limits as repeatable bin edges.""" + + # Store requested limits that extend beyond the distances drawn in this pair sample + pairs = xr.Dataset( + { + "value": (("pair", "endpoint"), np.column_stack((np.zeros(6), np.arange(1, 7)))), + "distance": ("pair", np.linspace(2, 8, 6)), + }, + attrs={"min_distance": 1.0, "max_distance": 10.0}, + ) + + # Build three log-spaced bins from the stored limits + result = gu.Variogram.from_pairs(pairs, bins="log", n_lags=3) + + # Check both outer edges and that every sampled pair enters one bin + assert result.bin_lower_edges is not None and result.bin_edges is not None + assert result.bin_lower_edges[0] == 1 + assert result.bin_edges[-1] == 10 + assert np.sum(result.counts) == 6 + + def test_from_pairs__bin_edges_missing_values_and_order(self) -> None: + """ + Checks that from_pairs() handles exact edges and pairs containing nodata without changing within-bin order. + """ + + # Mix endpoint order, exact boundaries, empty bins, and distances outside the requested range + distances = np.array([4, 1, 2, 1.5, 6, 0.5, 7, np.nan, np.inf, 0, 3, 2.5], dtype=float) + differences = np.arange(1, len(distances) + 1, dtype=float) + differences[-2:] = np.nan + pairs = xr.Dataset( + { + "value": (("pair", "endpoint"), np.column_stack((np.zeros(len(distances)), differences))), + "distance": ("pair", distances), + } + ) + edges = np.array([1, 2, 3, 4, 6], dtype=float) + + # Use an order-sensitive estimator so sorting must preserve each bin's original pair order + def weighted_difference(values: NDArrayNum) -> float: + """Weight each difference by its position within the supplied bin.""" + return float(np.dot(values, np.arange(1, len(values) + 1))) + + result = gu.Variogram.from_pairs(pairs, estimator=weighted_difference, bins=edges) + + # Check independently selected right-closed bins, with the first lower edge included + for index, (lower, upper) in enumerate(zip(edges[:-1], edges[1:])): + above_lower = distances >= lower if index == 0 else distances > lower + selected = above_lower & (distances <= upper) & np.isfinite(differences) + assert result.counts[index] == np.count_nonzero(selected) + if np.any(selected): + assert result.semivariance[index] == weighted_difference(differences[selected]) + assert result.lags[index] == np.mean(distances[selected]) + else: + assert np.isnan(result.semivariance[index]) + assert np.isnan(result.lags[index]) + + def test_from_pairs__constant_distances_with_explicit_bins(self) -> None: + """Checks that explicit bins accept pairs sharing one distance and include empty bins.""" + + # Place every finite pair at one distance with an endpoint difference of two + pairs = xr.Dataset( + {"value": (("pair", "endpoint"), np.tile([1.0, 3.0], (4, 1))), "distance": ("pair", np.full(4, 2.0))} + ) + result = gu.Variogram.from_pairs(pairs, estimator=np.mean, bins=[1, 2, 3]) + + # The first bin includes its upper boundary; the unoccupied bin remains NaN + np.testing.assert_array_equal(result.counts, [4, 0]) + np.testing.assert_allclose(result.semivariance, [2, np.nan], equal_nan=True) + np.testing.assert_allclose(result.lags, [2, np.nan], equal_nan=True) + + def test_from_pairs__distance_dimensions(self) -> None: + """Checks that from_pairs() rejects distances that do not provide one value per pair.""" + + # Give distances an extra dimension that would otherwise broadcast against endpoint differences + pairs = xr.Dataset( + {"value": (("pair", "endpoint"), np.zeros((3, 2))), "distance": (("pair", "extra"), np.ones((3, 1)))} + ) + with pytest.raises(ValueError, match="Variable 'distance' in argument ``pairs`` must have dimensions"): + gu.Variogram.from_pairs(pairs) + + +class TestVariogramEstimation: + """Checks variogram estimation and fitting through raster and point cloud methods. + + The methods cover repeated samples, validation, model names, functions, and the shared pair API. + """ + + def test_object_variogram_aggregates_runs_and_fits_summed_model(self) -> None: + """Checks that Raster.variogram() combines repeated samples and fits a summed model.""" + + # Create a smooth raster with variation in both grid directions + y, x = np.mgrid[:35, :35] + raster = gu.Raster.from_array( + np.sin(x / 4) + 0.5 * np.cos(y / 10), from_origin(0, 35, 2, 2), 32633, nodata=None + ) + # Combine three pair samples and fit Gaussian plus spherical components + result = raster.variogram( + n_pairs=1_000, + n_lags=8, + n_runs=3, + model=["gaussian", "spherical"], + random_state=42, + ) + + # Check both fitted components, sampling errors, released pairs, and run count + assert result.model is not None and result.model.model_name == "sum" + assert [component.model_name for component in result.model.components] == ["gaussian", "spherical"] + assert np.any(np.isfinite(result.semivariance_error)) + assert result.backend_object is None + assert result.attrs["n_runs"] == 3 + + @pytest.mark.parametrize("n_runs", [1, 3]) + def test_variogram_repetitions_match_independent_samples(self, n_runs: int) -> None: + """Checks that repeated pair samples combine their values, counts, and sampling errors.""" + + # Reproduce each pair sample separately with fixed distance bin edges + y, x = np.mgrid[:20, :20] + raster = gu.Raster.from_array(np.sin(x / 4) + np.cos(y / 5), from_origin(0, 20, 1, 1), 32633) + edges = np.linspace(0, 25, 6) + seeds = np.random.default_rng(42).integers(0, np.iinfo(np.int32).max, n_runs) + samples = [ + gu.Variogram.from_pairs(raster.pairsample(n_pairs=500, random_state=int(seed)), bins=edges) + for seed in seeds + ] + + # Check the public result against the mean values and summed counts from separate samples + result = gu.stats.variogram(raster, n_pairs=500, bins=edges, n_runs=n_runs, random_state=42) + empirical = np.stack([sample.semivariance for sample in samples]) + np.testing.assert_allclose(result.semivariance, np.nanmean(empirical, axis=0), equal_nan=True) + np.testing.assert_array_equal(result.counts, np.sum([sample.counts for sample in samples], axis=0)) + assert result.attrs["n_runs"] == n_runs + + # Check that sampling error is absent for one run and follows the usual standard error for repeated runs + if n_runs == 1: + assert np.all(np.isnan(result.semivariance_error)) + else: + expected = np.nanstd(empirical, ddof=1, axis=0) / np.sqrt(np.isfinite(empirical).sum(axis=0)) + np.testing.assert_allclose(result.semivariance_error, expected, equal_nan=True) + assert result.attrs["pair_count"] == int(result.counts.sum()) + + def test_variogram__repeated_explicit_bin_generator(self) -> None: + """Checks that repeated variogram samples reuse explicit boundaries supplied as a generator.""" + + # Build a raster and fixed bins that cover every sampled distance + array = np.arange(400, dtype=float).reshape(20, 20) + raster = gu.Raster.from_array(array, from_origin(0, 20, 1, 1), 32633) + edges = np.linspace(0, 30, 7) + options = {"n_pairs": 100, "n_runs": 2, "random_state": 7, "estimator": np.mean} + + # Check that every run receives the same edges even when the input can be iterated only once + expected = raster.variogram(bins=edges, **options) + result = raster.variogram(bins=(edge for edge in edges), **options) + np.testing.assert_array_equal(result.counts, expected.counts) + np.testing.assert_allclose(result.semivariance, expected.semivariance, equal_nan=True) + + @pytest.mark.parametrize("n_runs", [0, -1, 1.5, True]) + def test_variogram_rejects_invalid_repetitions(self, n_runs: int | float) -> None: + """Checks that invalid repetition counts are rejected before sampling.""" + + # Create a small valid raster because this option should fail before pair sampling + raster = gu.Raster.from_array(np.arange(16, dtype=float).reshape(4, 4), from_origin(0, 4, 1, 1), 32633) + + # Reject zero, negative, fractional, and boolean run counts + with pytest.raises(ValueError, match="Argument ``n_runs`` must be a positive integer"): + raster.variogram(n_runs=n_runs) + + def test_fit_accepts_short_names_and_skgstat_model_functions(self) -> None: + """Checks that fit() accepts SciKit-GStat functions and the short model names used by xDEM.""" + + # Create exact Gaussian values with the SciKit-GStat function + skgstat = pytest.importorskip("skgstat") + lags = np.linspace(1, 20, 12) + empirical = gu.Variogram( + lags=lags, + semivariance=skgstat.models.gaussian(lags, 10, 2), + counts=np.full(12, 100), + ) + + # Fit one function and one short string name as a summed model + fitted = empirical.fit([skgstat.models.gaussian, "Sph"]) + + # Check the standard component names stored in order + assert fitted.model is not None + assert [component.model_name for component in fitted.model.components] == ["gaussian", "spherical"] + + def test_pointcloud_variogram_and_advanced_pairs_share_api(self) -> None: + """Checks that PointCloud exposes both pair samples and their reduced variogram values.""" + + # Create point values that vary smoothly in both coordinate directions + y, x = np.mgrid[:18, :18] + pointcloud = gu.PointCloud.from_xyz(x.ravel(), y.ravel(), (np.sin(x / 3) + np.cos(y / 5)).ravel(), crs=32633) + + # Request either the individual pairs or six reduced distance bins + pairs = pointcloud.pairsample(n_pairs=300, min_distance=1, max_distance=15, random_state=2) + result = pointcloud.variogram(n_pairs=300, min_lag=1, max_lag=15, n_lags=6, random_state=2) + + # Check the requested sizes and that no model is fitted by default + assert pairs.sizes == {"pair": 300, "endpoint": 2} + assert len(result.lags) == 6 + assert result.model is None + + +class TestVariogramConversion: + """Checks fitted model evaluation and conversion to supported optional packages. + + The methods cover covariance composition and equivalent GPyTorch, SciKit-GStat, and GSTools parameters. + """ + + def test_model_evaluation_and_gpytorch_parameters(self) -> None: + """Checks that a Gaussian model gives the expected zero-distance values and GPyTorch length scale.""" + + # Create a Gaussian model with no measured values on the first two coordinate columns + result = gu.Variogram.from_model("gaussian", effective_range=12, partial_sill=3, nugget=0.2, active_dims=(0, 1)) + parameters = result.gpytorch_parameters() + + # Check the converted length scale and model values at zero distance + assert parameters["kernel_name"] == "RBF" + assert parameters["lengthscale"] == pytest.approx(12 / (2 * np.sqrt(2))) + assert result.variogram(0) == 0 + assert result.correlation(0) == 1 + assert result.correlation(np.zeros((2, 3))).shape == (2, 3) + + def test_product_model_multiplies_covariances(self) -> None: + """Checks that a product model multiplies component covariances and keeps one nugget.""" + + # Combine spatial and temporal models that use different coordinate columns + spatial = gu.Variogram.from_model("gaussian", effective_range=10, partial_sill=2, active_dims=(0, 1)) + temporal = gu.Variogram.from_model("exponential", effective_range=3, partial_sill=4, active_dims=(2,)) + combined = gu.Variogram.combine(spatial, temporal, combination="product", nugget=0.5) + + # Check the combined sill, zero-distance values, coordinate columns, and observation noise + assert combined.model is not None + assert combined.model.sill == 8.5 + assert combined.covariance(0) == pytest.approx(8.5) + assert combined.variogram(0) == 0 + parameters = combined.gpytorch_parameters() + assert [component["active_dims"] for component in parameters["components"]] == [(0, 1), (2,)] + assert parameters["noise"] == 0.5 + + def test_skgstat_estimation_can_discard_or_keep_backend(self) -> None: + """Checks that direct SciKit-GStat estimation keeps its source object only when requested.""" + + # Estimate the same coordinate values with and without keeping the SciKit-GStat object + coordinates = np.linspace(0, 10, 40)[:, np.newaxis] + values = np.sin(coordinates[:, 0]) + result = gu.Variogram.estimate(coordinates, values, model="gaussian", n_lags=6, normalize=False) + kept = gu.Variogram.estimate( + coordinates, values, model="gaussian", n_lags=6, normalize=False, keep_backend=True + ) + + # Check both storage choices and the method that releases a kept object + assert result.backend_object is None + assert kept.backend_object is not None + assert kept.without_backend().backend_object is None + + @pytest.mark.parametrize("model_name,smoothness", [("gaussian", None), ("exponential", None), ("matern", 1.5)]) + def test_gstools_conversion_matches_skgstat(self, model_name: str, smoothness: float | None) -> None: + """Checks that GSTools conversion matches SciKit-GStat values and nugget behavior.""" + + # Build one model with no measured values and convert it to GSTools + gstools = pytest.importorskip("gstools") + skgstat = pytest.importorskip("skgstat") + result = gu.Variogram.from_model( + model_name, effective_range=8, partial_sill=2, nugget=0.1, smoothness=smoothness + ) + converted = result.to_gstools(dim=1) + lags = np.array([0.0, 0.25, 1.0, 4.0, 8.0]) + # Calculate expected semivariances with the matching SciKit-GStat function + model_function = getattr(skgstat.models, model_name) + expected = ( + model_function(lags, r=8, c0=2, b=0.1) + if smoothness is None + else model_function(lags, r=8, c0=2, s=smoothness, b=0.1) + ) + + # Check the GSTools model values and unchanged coordinate column selection + assert isinstance(converted.model, gstools.CovModel) + assert converted.model.vario_axis(lags) == pytest.approx(expected) + assert converted.active_dims is None + + def test_gstools_conversion_keeps_sum_and_common_active_dims(self) -> None: + """Checks that summed GSTools models share coordinate columns and one parent nugget.""" + + # Sum two models that use the same two coordinate columns + first = gu.Variogram.from_model("gaussian", 8, 2, active_dims=(0, 1)) + second = gu.Variogram.from_model("exponential", 20, 3, active_dims=(0, 1)) + converted = gu.Variogram.combine(first, second, nugget=0.1).to_gstools(dim=2) + + # Check the shared columns, summed variance, and parent nugget + assert converted.active_dims == (0, 1) + assert converted.model.var == pytest.approx(5) + assert converted.model.nugget == pytest.approx(0.1) + + def test_gstools_rejects_component_specific_dimensions(self) -> None: + """Checks that GSTools conversion rejects components that use different coordinate columns.""" + + # Combine spatial and temporal models that select different columns + spatial = gu.Variogram.from_model("gaussian", 8, 2, active_dims=(0, 1)) + temporal = gu.Variogram.from_model("exponential", 3, 1, active_dims=(2,)) + combined = gu.Variogram.combine(spatial, temporal) + + # Check the clear error for this unsupported conversion + with pytest.raises(NotImplementedError, match="different dimensions"): + combined.to_gstools(dim=3)