From 46ca164bc71ca78a6a6b5eb8b8ef3a51b66cb445 Mon Sep 17 00:00:00 2001 From: Jeffrey Seifried Date: Thu, 15 Oct 2026 07:12:15 -0700 Subject: [PATCH 1/2] perform downscale as matrix multiplication in Dask, skipping COG materialization --- jdluc/__tests__/geo_test.py | 105 ++++++++++++++++++++++++ jdluc/geo.py | 158 ++++++++++++++++++++++++------------ 2 files changed, 212 insertions(+), 51 deletions(-) diff --git a/jdluc/__tests__/geo_test.py b/jdluc/__tests__/geo_test.py index 953fbdb..0f4f233 100644 --- a/jdluc/__tests__/geo_test.py +++ b/jdluc/__tests__/geo_test.py @@ -4,14 +4,19 @@ import numpy import pytest import rasterio +import rasterio.enums import rasterio.transform +import rasterio.warp import shapely import xarray +from rioxarray.rioxarray import affine_to_coords from jdluc.geo import ( clip_dset, convert_vector_to_flatgeobuf, + downscale_darray, get_chunk_size, + get_overlap_weights, get_overview_level, validate_geotiff, ) @@ -132,6 +137,106 @@ def test_clip_dset(geometry: shapely.Polygon, expected: int) -> None: assert int(result["var"].count()) == expected +@pytest.mark.parametrize( + ("fine_start", "fine_stop", "coarse_cells", "weights"), + ( + pytest.param( + 0, + 10, + slice(0, 3), + [ + [1, 1, 1, 1 / 3, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 2 / 3, 1, 1, 2 / 3, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 1 / 3, 1, 1, 1], + ], + id="the whole axis", + ), + pytest.param( + 3, + 7, + slice(0, 3), + [[1 / 3, 0, 0, 0], [2 / 3, 1, 1, 2 / 3], [0, 0, 0, 1 / 3]], + id="fine cells cut by two coarse cell edges", + ), + pytest.param( + 4, 6, slice(1, 2), [[1, 1]], id="fine cells inside one coarse cell" + ), + ), +) +def test_get_overlap_weights( + fine_start: int, fine_stop: int, coarse_cells: slice, weights: list[list[float]] +) -> None: + # 10 fine cells to 3 coarse ones, so each coarse cell covers 3 1/3 fine cells + result = get_overlap_weights( + coarse_size=3, fine_size=10, fine_start=fine_start, fine_stop=fine_stop + ) + assert result.coarse_cells == coarse_cells + numpy.testing.assert_allclose(result.weights, weights) + # Each fine cell's shares across the coarse cells add up to the whole cell + numpy.testing.assert_allclose(result.weights.sum(axis=0), 1) + + +@pytest.mark.parametrize( + ("fine_size", "coarse_size"), + ( + pytest.param( + 1_000, + 3, + id="coarse cell edges cut through fine cells, as 30 m to 1/12° does", + ), + pytest.param(900, 3, id="coarse cell edges fall on fine cell edges"), + ), +) +def test_downscale_darray_matches_a_gdal_average_warp( + coarse_size: int, fine_size: int +) -> None: + generator = numpy.random.default_rng(seed=0) + values = generator.random(dtype=numpy.float32, size=(fine_size, fine_size)) + values[generator.random(size=values.shape) < 0.3] = numpy.nan + # Every fine cell inside the top-left coarse cell is missing + values[: fine_size // 2, : fine_size // 2] = numpy.nan + fine_transform = rasterio.transform.from_origin( + north=20, west=-90, xsize=10 / fine_size, ysize=10 / fine_size + ) + coarse_transform = rasterio.transform.from_origin( + north=20, west=-90, xsize=10 / coarse_size, ysize=10 / coarse_size + ) + darray = xarray.DataArray( + coords=affine_to_coords( + affine=fine_transform, height=fine_size, width=fine_size + ), + data=values, + dims=("y", "x"), + name="var", + ).rio.write_crs(4326) + + expected = numpy.full( + dtype=numpy.float32, fill_value=numpy.nan, shape=(coarse_size, coarse_size) + ) + rasterio.warp.reproject( + destination=expected, + dst_crs=4326, + dst_nodata=numpy.nan, + dst_transform=coarse_transform, + resampling=rasterio.enums.Resampling.average, + source=values, + src_crs=4326, + src_nodata=numpy.nan, + src_transform=fine_transform, + ) + result = downscale_darray( + darray=darray, + epsg=4326, + height=coarse_size, + resampling=rasterio.enums.Resampling.average, + transform=coarse_transform.to_gdal(), + width=coarse_size, + ) + + assert numpy.isnan(result.values[0, 0]) + numpy.testing.assert_array_equal(result.values, expected) + + def get_dissolved_names( records: list[tuple[str, str, float]], tmp_path: pathlib.Path ) -> dict[str, str]: diff --git a/jdluc/geo.py b/jdluc/geo.py index 6dbb5cb..6bcf2f8 100644 --- a/jdluc/geo.py +++ b/jdluc/geo.py @@ -1,15 +1,13 @@ import collections.abc +import dataclasses import logging import math -import tempfile import typing import numpy import rasterio import rasterio.enums -import rasterio.shutil import rasterio.transform -import rasterio.vrt import shapely import xarray @@ -181,6 +179,35 @@ def exact_merge(*dsets: xarray.Dataset) -> xarray.Dataset: ).unify_chunks() +@dataclasses.dataclass +class AxisOverlaps: + coarse_cells: slice + weights: numpy.ndarray + + +def get_overlap_weights( + coarse_size: int, fine_size: int, fine_start: int, fine_stop: int +) -> AxisOverlaps: + """Along one axis, the share of each fine cell that falls inside each coarse cell. + + The axis is cut into `coarse_size` equal coarse cells and, separately, into `fine_size` equal + fine cells. Only the fine cells from `fine_start` up to `fine_stop` are looked at. The result + names the coarse cells those fine cells fall in, and for every pair of coarse and fine cell + gives the share of the fine cell inside the coarse cell: 1 when it lies wholly inside, 0 when + it lies outside, and a fraction when the coarse cell's edge cuts through it. + """ + first = fine_start * coarse_size // fine_size + last = (fine_stop * coarse_size - 1) // fine_size + edges = numpy.arange(first, last + 2) * (fine_size / coarse_size) + fine = numpy.arange(fine_start, fine_stop)[numpy.newaxis, :] + overlaps = numpy.minimum(fine + 1, edges[1:, numpy.newaxis]) - numpy.maximum( + fine, edges[:-1, numpy.newaxis] + ) + return AxisOverlaps( + coarse_cells=slice(first, last + 1), weights=overlaps.clip(min=0) + ) + + def downscale_darray( darray: xarray.DataArray, epsg: int, @@ -189,55 +216,84 @@ def downscale_darray( transform: tuple[float, float, float, float, float, float], width: int, ) -> xarray.DataArray: - import rioxarray + """Average `darray` onto a coarser grid that covers the same area. + + `transform`, `height` and `width` describe the coarser grid. Each coarse cell gets the mean of + the fine cells inside it, skipping missing (NaN) ones. A fine cell that the edge of a coarse + cell cuts through counts only by the share of it inside that coarse cell. A coarse cell with + no non-missing fine cell inside it is NaN. This matches GDAL's average resampling. + + The fine grid is read one chunk at a time. Each chunk adds its weighted sum and weighted count + to the coarse cells it overlaps, and the mean is taken once every chunk is added in. So only a + few chunks are in memory at once, never the whole fine grid. + """ + from rioxarray.rioxarray import affine_to_coords + + assert resampling == rasterio.enums.Resampling.average + assert darray.rio.crs.to_epsg() == epsg + assert darray.dims == ("y", "x") + affine = rasterio.transform.Affine.from_gdal(*transform) + assert numpy.allclose( + darray.rio.bounds(), + rasterio.transform.array_bounds(height=height, transform=affine, width=width), + ) logger.info(f"Downscaling {darray.name=:s}") - - import dask.diagnostics - - path_to_unscaled = tempfile.NamedTemporaryFile( # noqa: SIM115 - delete=False, suffix=".glad.tif" - ).name - chunk_size = 1 << 11 # 2k - logger.debug(f"Serializing to {path_to_unscaled=:s} with {chunk_size=:d}") - darray = darray.assign_attrs(long_name=darray.name).chunk(chunks=chunk_size) - with dask.diagnostics.ProgressBar(): - darray.rio.to_raster( - path_to_unscaled, - BIGTIFF="IF_SAFER", - ZSTD_LEVEL=1, - blockxsize=chunk_size, - blockysize=chunk_size, - compress="ZSTD", - driver="GTiff", - lock=True, - num_threads="all_cpus", - tiled=True, + fine_height, fine_width = darray.shape + + def get_partial_sums( + block: numpy.ndarray, block_info: dict[int | None, typing.Any] + ) -> numpy.ndarray: + (row_start, row_stop), (column_start, column_stop) = block_info[0][ + "array-location" + ] + rows = get_overlap_weights( + coarse_size=height, + fine_size=fine_height, + fine_start=row_start, + fine_stop=row_stop, + ) + columns = get_overlap_weights( + coarse_size=width, + fine_size=fine_width, + fine_start=column_start, + fine_stop=column_stop, + ) + valid = numpy.isfinite(block) + # For each coarse cell, the sum of the non-missing fine values inside it and how many + # there are, with each fine cell counted by the share of it inside the coarse cell + partial_sums = numpy.zeros(shape=(1, 1, 2, height, width)) + partial_sums[0, 0, :, rows.coarse_cells, columns.coarse_cells] = ( + rows.weights + @ numpy.stack([numpy.where(valid, block, 0), valid]).astype(numpy.float64) + @ columns.weights.T ) - with ( - rasterio.open(path_to_unscaled) as unscaled_fp, - rasterio.vrt.WarpedVRT( - unscaled_fp, - crs=epsg, - height=height, - resampling=resampling, - transform=rasterio.transform.Affine.from_gdal(*transform), - width=width, - ) as scaled_vrt, - tempfile.NamedTemporaryFile(delete=False, suffix=".warped.vrt") as scaled_fp, - ): - assert unscaled_fp.crs.to_epsg() == epsg - logger.debug(f"Writing warped VRT to {scaled_fp.name=:s}") - rasterio.shutil.copy(scaled_vrt, scaled_fp.name, driver="VRT") - ret = rioxarray.open_rasterio( - scaled_fp.name, - chunks=get_chunk_size( - dtypes=[numpy.dtype("float32")], - # 128 kiB ~ 1 GiB x (12 ÷ 4_000 resolution) ÷ 14 bands - max_bytes_per_chunk=1 << 17, - ), - lock=False, + return partial_sums + + data = darray.chunk(chunks=1 << 11).data # 2k + numerator, denominator = ( + data.map_blocks( + get_partial_sums, + chunks=( + (1,) * data.numblocks[0], + (1,) * data.numblocks[1], + (2,), + (height,), + (width,), + ), + dtype=numpy.float64, + new_axis=[2, 3, 4], + ) + .sum(axis=(0, 1), split_every=4) + .compute() + ) + mean = numpy.full(fill_value=numpy.nan, shape=(height, width)) + numpy.divide(numerator, denominator, out=mean, where=denominator > 0) + return unify_dtype_and_no_data( + darray=xarray.DataArray( + coords=affine_to_coords(affine=affine, height=height, width=width), + data=mean.astype(numpy.float32), + dims=("y", "x"), + name=darray.name, + ).chunk() ) - assert isinstance(ret, xarray.DataArray) - ret = ret.drop_vars("spatial_ref").isel(band=0, drop=True) - return unify_dtype_and_no_data(darray=ret.rename(ret.attrs.pop("long_name"))) From a453c1b678afbaac0578f300c3b999a70798c985 Mon Sep 17 00:00:00 2001 From: Jeffrey Seifried Date: Thu, 15 Oct 2026 10:00:48 -0700 Subject: [PATCH 2/2] address review feedback --- jdluc/__tests__/geo_test.py | 85 ++++++++++++++++++++++++++++++++++--- jdluc/geo.py | 33 +++++++++++--- 2 files changed, 108 insertions(+), 10 deletions(-) diff --git a/jdluc/__tests__/geo_test.py b/jdluc/__tests__/geo_test.py index 0f4f233..d364377 100644 --- a/jdluc/__tests__/geo_test.py +++ b/jdluc/__tests__/geo_test.py @@ -18,6 +18,7 @@ get_chunk_size, get_overlap_weights, get_overview_level, + get_weighted_sums, validate_geotiff, ) @@ -138,11 +139,12 @@ def test_clip_dset(geometry: shapely.Polygon, expected: int) -> None: @pytest.mark.parametrize( - ("fine_start", "fine_stop", "coarse_cells", "weights"), + ("fine_start", "fine_stop", "fine_size", "coarse_cells", "weights"), ( pytest.param( 0, 10, + 10, slice(0, 3), [ [1, 1, 1, 1 / 3, 0, 0, 0, 0, 0, 0], @@ -154,21 +156,37 @@ def test_clip_dset(geometry: shapely.Polygon, expected: int) -> None: pytest.param( 3, 7, + 10, slice(0, 3), [[1 / 3, 0, 0, 0], [2 / 3, 1, 1, 2 / 3], [0, 0, 0, 1 / 3]], id="fine cells cut by two coarse cell edges", ), pytest.param( - 4, 6, slice(1, 2), [[1, 1]], id="fine cells inside one coarse cell" + 4, 6, 10, slice(1, 2), [[1, 1]], id="fine cells inside one coarse cell" + ), + pytest.param( + 0, 4, 12, slice(0, 1), [[1] * 4], id="first coarse cell, ending on its edge" ), + pytest.param( + 4, + 8, + 12, + slice(1, 2), + [[1] * 4], + id="middle coarse cell, both ends on edges", + ), + pytest.param(11, 12, 12, slice(2, 3), [[1]], id="last fine cell"), ), ) def test_get_overlap_weights( - fine_start: int, fine_stop: int, coarse_cells: slice, weights: list[list[float]] + fine_start: int, + fine_stop: int, + fine_size: int, + coarse_cells: slice, + weights: list[list[float]], ) -> None: - # 10 fine cells to 3 coarse ones, so each coarse cell covers 3 1/3 fine cells result = get_overlap_weights( - coarse_size=3, fine_size=10, fine_start=fine_start, fine_stop=fine_stop + coarse_size=3, fine_size=fine_size, fine_start=fine_start, fine_stop=fine_stop ) assert result.coarse_cells == coarse_cells numpy.testing.assert_allclose(result.weights, weights) @@ -176,6 +194,28 @@ def test_get_overlap_weights( numpy.testing.assert_allclose(result.weights.sum(axis=0), 1) +@pytest.mark.parametrize( + ("block", "expected"), + ( + pytest.param( + [[numpy.nan, numpy.nan], [numpy.nan, numpy.nan]], [0, 0], id="all missing" + ), + pytest.param([[2, 4], [6, 8]], [10.5, 2.25], id="fractional weights"), + pytest.param([[2, numpy.nan], [6, 8]], [6.5, 1.25], id="unequal valid counts"), + ), +) +def test_get_weighted_sums(block: list[list[float]], expected: list[float]) -> None: + values = numpy.asarray(block) + actual = get_weighted_sums( + block=values, + column_weights=numpy.array([[0.5, 1.0]]), + row_weights=numpy.array([[1.0, 0.5]]), + valid=numpy.isfinite(values), + ) + # Assert the weighted sum and count before either is divided. + numpy.testing.assert_allclose(actual[:, 0, 0], expected) + + @pytest.mark.parametrize( ("fine_size", "coarse_size"), ( @@ -185,6 +225,7 @@ def test_get_overlap_weights( id="coarse cell edges cut through fine cells, as 30 m to 1/12° does", ), pytest.param(900, 3, id="coarse cell edges fall on fine cell edges"), + pytest.param(2_051, 3, id="multiple chunks with short final chunks"), ), ) def test_downscale_darray_matches_a_gdal_average_warp( @@ -195,6 +236,10 @@ def test_downscale_darray_matches_a_gdal_average_warp( values[generator.random(size=values.shape) < 0.3] = numpy.nan # Every fine cell inside the top-left coarse cell is missing values[: fine_size // 2, : fine_size // 2] = numpy.nan + # Set the short last chunks apart, so a chunk that is dropped or put in the wrong place + # changes the result (a no-op when the grid fits in one 2048-cell chunk) + values[2_048:, :] = 1_000 + values[:, 2_048:] = 2_000 fine_transform = rasterio.transform.from_origin( north=20, west=-90, xsize=10 / fine_size, ysize=10 / fine_size ) @@ -237,6 +282,36 @@ def test_downscale_darray_matches_a_gdal_average_warp( numpy.testing.assert_array_equal(result.values, expected) +@pytest.mark.parametrize("nodata", [numpy.nan, -9999.0]) +def test_downscale_darray_respects_source_nodata(nodata: float) -> None: + values = numpy.array( + [ + [nodata, nodata, 1, 3], + [nodata, nodata, 5, 7], + [2, nodata, 10, 20], + [4, 6, 30, 40], + ], + dtype=numpy.float32, + ) + source = ( + xarray.DataArray(values, dims=("y", "x"), name="var") + .rio.write_crs(4326) + .rio.write_transform(rasterio.transform.from_origin(0, 4, 1, 1)) + .rio.write_nodata(nodata) + ) + result = downscale_darray( + darray=source, + epsg=4326, + height=2, + resampling=rasterio.enums.Resampling.average, + transform=rasterio.transform.from_origin(0, 4, 2, 2).to_gdal(), + width=2, + ) + numpy.testing.assert_allclose( + result.values, [[numpy.nan, 4.0], [4.0, 25.0]], equal_nan=True + ) + + def get_dissolved_names( records: list[tuple[str, str, float]], tmp_path: pathlib.Path ) -> dict[str, str]: diff --git a/jdluc/geo.py b/jdluc/geo.py index 6bcf2f8..3badd3e 100644 --- a/jdluc/geo.py +++ b/jdluc/geo.py @@ -197,6 +197,8 @@ def get_overlap_weights( it lies outside, and a fraction when the coarse cell's edge cuts through it. """ first = fine_start * coarse_size // fine_size + # fine_stop is exclusive, so when the last fine cell ends exactly on a coarse cell edge, the + # coarse cell that starts at that edge isn't included last = (fine_stop * coarse_size - 1) // fine_size edges = numpy.arange(first, last + 2) * (fine_size / coarse_size) fine = numpy.arange(fine_start, fine_stop)[numpy.newaxis, :] @@ -208,6 +210,21 @@ def get_overlap_weights( ) +def get_weighted_sums( + block: numpy.ndarray, + column_weights: numpy.ndarray, + row_weights: numpy.ndarray, + valid: numpy.ndarray, +) -> numpy.ndarray: + """The weighted sum of the valid values in `block` and the weighted count of them, as a + (2, coarse rows, coarse columns) array, each fine cell counted by its share of the coarse cell. + """ + values_and_counts = numpy.stack([numpy.where(valid, block, 0), valid]).astype( + numpy.float64 + ) + return row_weights @ values_and_counts @ column_weights.T + + def downscale_darray( darray: xarray.DataArray, epsg: int, @@ -233,6 +250,8 @@ def downscale_darray( assert darray.rio.crs.to_epsg() == epsg assert darray.dims == ("y", "x") affine = rasterio.transform.Affine.from_gdal(*transform) + # Weights by row and column index need both grids to run north to south and west to east. + assert all(grid.a > 0 and grid.e < 0 for grid in (darray.rio.transform(), affine)) assert numpy.allclose( darray.rio.bounds(), rasterio.transform.array_bounds(height=height, transform=affine, width=width), @@ -240,6 +259,7 @@ def downscale_darray( logger.info(f"Downscaling {darray.name=:s}") fine_height, fine_width = darray.shape + source_nodata = darray.rio.nodata def get_partial_sums( block: numpy.ndarray, block_info: dict[int | None, typing.Any] @@ -260,13 +280,16 @@ def get_partial_sums( fine_stop=column_stop, ) valid = numpy.isfinite(block) - # For each coarse cell, the sum of the non-missing fine values inside it and how many - # there are, with each fine cell counted by the share of it inside the coarse cell + if source_nodata is not None: + valid &= block != source_nodata partial_sums = numpy.zeros(shape=(1, 1, 2, height, width)) partial_sums[0, 0, :, rows.coarse_cells, columns.coarse_cells] = ( - rows.weights - @ numpy.stack([numpy.where(valid, block, 0), valid]).astype(numpy.float64) - @ columns.weights.T + get_weighted_sums( + block=block, + column_weights=columns.weights, + row_weights=rows.weights, + valid=valid, + ) ) return partial_sums