Compare large arrays in memory-bounded chunks

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))

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

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]):