diff --git a/CHANGELOG.md b/CHANGELOG.md index e570aba6e..f618eec03 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Compare large xarray job result arrays in memory-bounded chunks to avoid out-of-memory failures. ([#871](https://github.com/Open-EO/openeo-python-client/issues/871)) - `DataCube.resample_spatial()` now supports parameterized `resolution` and `projection` arguments. ([#897](https://github.com/Open-EO/openeo-python-client/issues/897)) - Sanitize asset download filenames (e.g. strip slashes, (semi)colon, hash, ...), instead of blindly using the asset key as filename. ([#820](https://github.com/Open-EO/openeo-python-client/issues/820)) - Support parameters in `DataCube` apply- and band-math operations ([#903](https://github.com/Open-EO/openeo-python-client/issues/903)) diff --git a/openeo/testing/results.py b/openeo/testing/results.py index f04ee702e..8991f6a9a 100644 --- a/openeo/testing/results.py +++ b/openeo/testing/results.py @@ -2,8 +2,10 @@ Assert functions for comparing actual (batch job) results against expected reference data. """ +import itertools import json import logging +import math import re import tempfile from pathlib import Path @@ -22,6 +24,7 @@ _DEFAULT_RTOL = 1e-6 _DEFAULT_ATOL = 1e-6 _DEFAULT_PIXELTOL = 0.0 +_DEFAULT_COMPARISON_MEMORY_BUDGET = 64 * 1024 * 1024 # https://paulbourke.net/dataformats/asciiart DEFAULT_GRAYSCALE_70_CHARACTERS = r"$@B%8&WM#*oahkbdpqwmZO0QLCJUYXzcvunxrjft/\|()1{}[]?-_+~<>i!lI;:,\"^`'. "[::-1] @@ -33,7 +36,96 @@ def _load_xarray_netcdf(path: Union[str, Path], **kwargs) -> xarray.Dataset: Load a netCDF file as Xarray Dataset """ _log.debug(f"_load_xarray_netcdf: {path!r}") - return xarray.load_dataset(path, **kwargs) + # Keep variables backed by the file so large arrays can be read in bounded chunks. + kwargs.setdefault("cache", False) + return xarray.open_dataset(path, **kwargs) + + +def _iter_chunk_indexers( + data: xarray.DataArray, + *, + bytes_per_cell: int, + memory_budget: Optional[int] = None, +): + """Yield rectangular ``isel`` indexers within an approximate memory budget.""" + memory_budget = memory_budget or _DEFAULT_COMPARISON_MEMORY_BUDGET + max_cells = max(1, memory_budget // max(1, bytes_per_cell)) + chunk_sizes = list(data.shape) + while math.prod(chunk_sizes) > max_cells: + axis = max(range(len(chunk_sizes)), key=chunk_sizes.__getitem__) + chunk_sizes[axis] = max(1, (chunk_sizes[axis] + 1) // 2) + + ranges = [range(0, size, chunk) for size, chunk in zip(data.shape, chunk_sizes)] + for starts in itertools.product(*ranges): + yield { + dim: slice(start, min(start + chunk, size)) + for dim, start, chunk, size in zip(data.dims, starts, chunk_sizes, data.shape) + } + + +def _compare_xarray_dataarray_chunked( + actual: xarray.DataArray, + expected: xarray.DataArray, + *, + rtol: float, + atol: float, + pixel_tolerance: float, +) -> List[str]: + """Compare compatible numerical arrays without materializing full-size intermediates.""" + bytes_per_cell = max(actual.dtype.itemsize, expected.dtype.itemsize, 8) * 8 + bad_count = 0 + total_count = 0 + finite_diff_count = 0 + diff_sum = 0.0 + diff_sum_squared = 0.0 + diff_min = math.inf + diff_max = -math.inf + + for indexers in _iter_chunk_indexers(actual, bytes_per_cell=bytes_per_cell): + actual_chunk = numpy.asarray(actual.isel(indexers).values) + expected_chunk = numpy.asarray(expected.isel(indexers).values) + with numpy.errstate(all="ignore"): + close = numpy.isclose(actual_chunk, expected_chunk, rtol=rtol, atol=atol, equal_nan=True) + bad = ~close + chunk_bad_count = int(numpy.count_nonzero(bad)) + bad_count += chunk_bad_count + total_count += int(bad.size) + + if chunk_bad_count: + with numpy.errstate(all="ignore"): + comparison_dtype = numpy.result_type(actual_chunk.dtype, expected_chunk.dtype, numpy.float64) + differences = numpy.abs( + actual_chunk.astype(comparison_dtype) - expected_chunk.astype(comparison_dtype) + )[bad] + finite_differences = differences[numpy.isfinite(differences)] + if finite_differences.size: + finite_diff_count += int(finite_differences.size) + diff_sum += float(finite_differences.sum(dtype=float)) + diff_sum_squared += float(numpy.square(finite_differences).sum(dtype=float)) + diff_min = min(diff_min, float(finite_differences.min())) + diff_max = max(diff_max, float(finite_differences.max())) + + bad_percentage = bad_count * 100 / total_count if total_count else 0.0 + if bad_percentage <= pixel_tolerance: + return [] + + if pixel_tolerance: + return [ + f"Fraction significantly differing pixels: {bad_percentage}% > {pixel_tolerance}% " + f"({bad_count}/{total_count} values)" + ] + + statistics = "" + if finite_diff_count: + diff_mean = diff_sum / finite_diff_count + diff_variance = max(0.0, diff_sum_squared / finite_diff_count - diff_mean**2) + statistics = ( + f", absolute difference min: {diff_min}, max: {diff_max}, " f"mean: {diff_mean}, variance: {diff_variance}" + ) + return [ + f"Data values are not close (rtol={rtol}, atol={atol}): " + f"{bad_count}/{total_count} values differ ({bad_percentage}%){statistics}" + ] def _load_rioxarray_geotiff(path: Union[str, Path], **kwargs) -> xarray.DataArray: @@ -240,6 +332,21 @@ def _compare_xarray_dataarray( issues.append(f"Shape mismatch: {actual.shape} != {expected.shape}") compatible = len(issues) == 0 is_numerical_data = numpy.issubdtype(actual.dtype, numpy.number) and numpy.issubdtype(expected.dtype, numpy.number) + use_chunked_comparison = ( + compatible and is_numerical_data and actual.nbytes + expected.nbytes > _DEFAULT_COMPARISON_MEMORY_BUDGET + ) + if use_chunked_comparison: + issues.extend( + _compare_xarray_dataarray_chunked( + actual=actual, + expected=expected, + rtol=rtol, + atol=atol, + pixel_tolerance=pixel_tolerance, + ) + ) + return issues + try: if pixel_tolerance and compatible and is_numerical_data: threshold = abs(expected * rtol) + atol diff --git a/tests/testing/test_results.py b/tests/testing/test_results.py index 04562abe7..86a01f308 100644 --- a/tests/testing/test_results.py +++ b/tests/testing/test_results.py @@ -1,5 +1,6 @@ import contextlib import json +import math import re from pathlib import Path from typing import List, Optional, Union @@ -9,6 +10,7 @@ import pytest import xarray +import openeo.testing.results from openeo.rest.job import DEFAULT_JOB_RESULTS_FILENAME from openeo.testing.results import ( _compare_xarray_dataarray, @@ -213,6 +215,69 @@ def test_nan_handling(self): dirty_equals.IsStr(regex=r"Left and right DataArray objects are not close.*", regex_flags=re.DOTALL), ] + def test_large_array_comparison_is_chunked(self, monkeypatch): + monkeypatch.setattr(openeo.testing.results, "_DEFAULT_COMPARISON_MEMORY_BUDGET", 1024) + expected = xarray.DataArray(numpy.arange(10_000, dtype=numpy.int16).reshape(100, 100), dims=["y", "x"]) + actual = expected.copy(deep=True) + original_isclose = numpy.isclose + compared_shapes = [] + + def recording_isclose(a, b, **kwargs): + compared_shapes.append(a.shape) + return original_isclose(a, b, **kwargs) + + monkeypatch.setattr(openeo.testing.results.numpy, "isclose", recording_isclose) + + assert _compare_xarray_dataarray(actual, expected) == [] + assert len(compared_shapes) > 1 + assert max(math.prod(shape) for shape in compared_shapes) <= 16 + + def test_large_array_comparison_reports_streaming_statistics(self, monkeypatch): + monkeypatch.setattr(openeo.testing.results, "_DEFAULT_COMPARISON_MEMORY_BUDGET", 1024) + expected = xarray.DataArray(numpy.zeros((100, 100)), dims=["y", "x"]) + actual = expected.copy(deep=True) + actual.values[10, 20] = 2 + actual.values[90, 80] = 4 + + issues = _compare_xarray_dataarray(actual, expected, rtol=0, atol=0.1) + + assert issues == [ + dirty_equals.IsStr( + regex=r"Data values are not close .*2/10000 values differ \(0.02%\).*min: 2.0, max: 4.0.*mean: 3.0.*" + ) + ] + + def test_large_array_pixel_tolerance(self, monkeypatch): + monkeypatch.setattr(openeo.testing.results, "_DEFAULT_COMPARISON_MEMORY_BUDGET", 1024) + expected = xarray.DataArray(numpy.zeros((100, 100)), dims=["y", "x"]) + actual = expected.copy(deep=True) + actual.values.flat[:100] = 1 + + assert _compare_xarray_dataarray(actual, expected, pixel_tolerance=1.0) == [] + assert _compare_xarray_dataarray(actual, expected, pixel_tolerance=0.5) == [ + dirty_equals.IsStr(regex=r"Fraction significantly differing pixels: 1.0% > 0.5% \(100/10000 values\)") + ] + + def test_large_complex_array_statistics(self, monkeypatch): + monkeypatch.setattr(openeo.testing.results, "_DEFAULT_COMPARISON_MEMORY_BUDGET", 32) + expected = xarray.DataArray(numpy.zeros(4, dtype=numpy.complex128), dims=["x"]) + actual = expected.copy(deep=True) + actual.values[2] = 3 + 4j + + assert _compare_xarray_dataarray(actual, expected, rtol=0, atol=0) == [ + dirty_equals.IsStr(regex=r".*1/4 values differ .*min: 5.0, max: 5.0, mean: 5.0.*") + ] + + def test_load_netcdf_is_lazy(self, tmp_path): + path = tmp_path / "large.nc" + xarray.Dataset({"data": (("y", "x"), numpy.ones((100, 100), dtype=numpy.int16))}).to_netcdf(path) + + dataset = openeo.testing.results._load_xarray_netcdf(path) + try: + assert dataset["data"].variable._in_memory is False + finally: + dataset.close() + @contextlib.contextmanager def raises_assertion_error_or_not(message: Union[None, str, re.Pattern]):