From f1d51f958ab2c5b66ee033908d82b41a3779f40d Mon Sep 17 00:00:00 2001 From: Panu Lahtinen Date: Tue, 20 Oct 2026 18:20:15 +0000 Subject: [PATCH 1/3] Add on-disk Zarr caching of gradient search target coordinates Add a generic ZarrCache for resampling look-up tables, enabled with the new 'cache_resampling_luts' config option and stored under 'cache_dir'. The gradient search resampler uses it for area-to-area resampling to cache the target area projection coordinates transformed to the source projection. The cache is namespaced and versioned so that the kd-tree based resamplers can later store their indices in it too. Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 9 +- docs/source/howtos/configuration.rst | 37 ++++++- pyresample/_caching.py | 139 ++++++++++++++++++++++++++- pyresample/_config.py | 1 + pyresample/gradient/__init__.py | 83 ++++++++++++++-- pyresample/test/conftest.py | 1 + pyresample/test/test_caching.py | 128 ++++++++++++++++++++++++ pyresample/test/test_gradient.py | 115 ++++++++++++++++++++++ 8 files changed, 502 insertions(+), 11 deletions(-) create mode 100644 pyresample/test/test_caching.py diff --git a/AGENTS.md b/AGENTS.md index aa3daf735..0021ca8e7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -118,7 +118,11 @@ anything. Known divergences: `BaseResampler` (only bilinear really uses it; `DaskEWAResampler.precompute` accepts it and warns that it is ignored), bilinear's own zarr `save_resampling_info`/ `load_resampling_info`, and `pyresample.config["cache_dir"]`, which is used only by - `_caching.py` for geometry slices. + `_caching.py`: for geometry slices (`cache_geometry_slices`, JSON) and for resampling + LUTs (`cache_resampling_luts`, `ZarrCache`). `ZarrCache` is the one meant to grow: + today only gradient search uses it (area-to-area target coordinates in the source CRS), + but it is designed so the kd-tree nearest/bilinear index caches can move onto it from + Satpy, each under its own namespace. Gradient search ignores the `cache_dir=` kwarg. - Only `nearest` is registered in the `pyresample.resamplers` entry point group, so `future.list_resamplers()` does not see bilinear, EWA, gradient or bucket. @@ -158,8 +162,9 @@ path from the old one; if that isn't possible, say so rather than extending the `pyresample.config` is a [donfig](https://donfig.readthedocs.io/) object defined in `pyresample/_config.py`. Keys and defaults: -- `cache_dir` — platformdirs user cache dir; used only for geometry-slice caching. +- `cache_dir` — platformdirs user cache dir; used only by `_caching.py`. - `cache_geometry_slices` — `False`. +- `cache_resampling_luts` — `False`. - `features.future_geometries` — `False`. Environment variables are `PYRESAMPLE_` prefixed with `__` for nesting diff --git a/docs/source/howtos/configuration.rst b/docs/source/howtos/configuration.rst index 6a79acf10..16090a9f3 100644 --- a/docs/source/howtos/configuration.rst +++ b/docs/source/howtos/configuration.rst @@ -94,7 +94,8 @@ the `platformdirs ` Some resampling algorithms provide caching functionality when the user provides a directory to cache to. These resamplers do not currently use this - configuration option. + configuration option. Resamplers that support the + ``cache_resampling_luts`` option (see below) do use it. .. _config_cache_sensor_angles_setting: @@ -131,6 +132,40 @@ intersection logic this caching option should hopefully not be needed. When setting this as an environment variable, this should be set with the string equivalent of the Python boolean values ``="True"`` or ``="False"``. +.. warning:: + + This caching does not limit the number of entries nor does it expire old + entries. It is up to the user to manage the contents of the cache + directory. + +Cache Resampling LUTs +^^^^^^^^^^^^^^^^^^^^^ + +* **Environment variable**: ``PYRESAMPLE_CACHE_RESAMPLING_LUTS`` +* **YAML/Config Key**: ``cache_resampling_luts`` +* **Default**: ``False`` + +Whether or not resampling look-up tables (LUTs), the expensive intermediate +arrays a resampler computes from the source and target geometries alone, are +cached to disk as Zarr. At the time of writing this is only done by the +gradient search resampler +(:func:`~pyresample.gradient.create_gradient_search_resampler`) when +resampling from an area to an area: the target area projection coordinates +transformed to the projection of the source area are cached, so they only need +to be computed once for each target area and source projection. They are +reused for all source areas sharing that projection, for example different +sub-areas or resolutions of the same geostationary satellite. + +The LUTs are stored in the ``resampling_luts`` subdirectory of ``cache_dir`` +(see above), in a separate subdirectory for each resampler and cache format +version. The first time a LUT is needed, it is computed for the whole target +area and written to disk during the resampler's ``precompute`` step. Later +uses read it lazily. The cached gradient search coordinates take up to 16 bytes +per target pixel. + +When setting this as an environment variable, this should be set with the +string equivalent of the Python boolean values ``="True"`` or ``="False"``. + .. warning:: This caching does not limit the number of entries nor does it expire old diff --git a/pyresample/_caching.py b/pyresample/_caching.py index 08b326e0c..4ca0d2e42 100644 --- a/pyresample/_caching.py +++ b/pyresample/_caching.py @@ -8,15 +8,20 @@ import hashlib import json +import logging import os +import shutil +import tempfile import warnings from functools import update_wrapper from glob import glob from pathlib import Path -from typing import Any, Callable +from typing import Any, Callable, Mapping import pyresample +logger = logging.getLogger(__name__) + class JSONCacheHelper: """Decorator class to cache results to a JSON file on-disk.""" @@ -126,3 +131,135 @@ def _decorator(func: Callable) -> Callable: return wrapper return _decorator + + +class ZarrCache: + """On-disk cache of named, geometry-only arrays stored as Zarr groups. + + This is meant for the expensive, reusable look-up tables (LUTs) that + resamplers compute from their source and target geometries alone, for + example target coordinates transformed to the source projection or + neighbour indices. Each cache entry is a Zarr group holding one or more + named arrays and is identified by a hexadecimal key, typically created with + :func:`hash_cache_key` from everything the arrays depend on. Entries are + stored in ``/resampling_luts/_v/.zarr`` + where ``cache_dir`` is ``pyresample.config["cache_dir"]`` unless given + explicitly. Each user of the cache should have its own ``namespace``, and + increase its ``version`` whenever the content or layout of what it stores + changes, so that stale entries are never read. + + Entries are written to a temporary directory first and then renamed into + place, so a partially written entry is never visible to readers and + concurrent writers of the same entry do not corrupt each other. + + Whether a resampler should use this cache at all is controlled by the + ``cache_resampling_luts`` configuration key, see :meth:`is_enabled`. + + Like the other caching in Pyresample, this cache does not limit its size nor + expire old entries. It is up to the user to manage the cache directory. + """ + + config_key = "cache_resampling_luts" + + def __init__(self, namespace: str, version: int = 1, cache_dir: str | os.PathLike | None = None): + self._namespace = namespace + self._version = version + self._cache_dir = cache_dir + + @classmethod + def is_enabled(cls) -> bool: + """Check whether caching resampling LUTs is enabled in ``pyresample.config``.""" + return bool(pyresample.config.get(cls.config_key, False)) + + @property + def directory(self) -> Path: + """Get the directory holding the entries of this namespace and version.""" + cache_dir = self._cache_dir or pyresample.config.get("cache_dir") + if cache_dir is None: + raise RuntimeError("Can't use Zarr caching. No 'cache_dir' configured.") + return Path(cache_dir) / "resampling_luts" / f"{self._namespace}_v{self._version}" + + def get_path(self, key: str) -> Path: + """Get the path of the Zarr group for the cache entry *key*.""" + return self.directory / f"{key}.zarr" + + def load(self, key: str, chunks: Mapping[str, Any] | None = None) -> dict[str, Any] | None: + """Load the arrays of a cache entry lazily as dask arrays. + + Args: + key: Key of the cache entry. + chunks: Mapping of array name to the chunks the loaded dask array + should have. Arrays not in the mapping use the chunking they + were stored with. + + Returns: + Mapping of array names to dask arrays, or ``None`` if the entry + does not exist. + """ + import dask.array as da + import zarr + + path = self.get_path(key) + if not path.is_dir(): + return None + chunks = chunks or {} + group = zarr.open_group(str(path), mode="r") + logger.debug("Loading cached resampling LUTs from %s", path) + return {name: da.from_zarr(group[name], chunks=chunks.get(name, None), inline_array=True) + for name in sorted(group.array_keys())} + + def save(self, key: str, arrays: Mapping[str, Any]) -> None: + """Compute and store arrays as a cache entry. + + Dask arrays are computed, in parallel, while being written. If the + entry already exists it is left untouched. + + Args: + key: Key of the cache entry. + arrays: Mapping of array names to numpy or dask arrays. + """ + import dask + import dask.array as da + + path = self.get_path(key) + if path.is_dir(): + return + path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = Path(tempfile.mkdtemp(prefix=f".{key}-", suffix=".tmp", dir=path.parent)) + try: + writes = [da.to_zarr(da.asarray(arr), str(tmp_path), component=name, compute=False) + for name, arr in arrays.items()] + dask.compute(*writes) + logger.debug("Saving resampling LUTs to %s", path) + os.rename(tmp_path, path) + except OSError: + if not path.is_dir(): + raise + # Another process wrote the same entry first, keep theirs. + finally: + if tmp_path.exists(): + shutil.rmtree(tmp_path, ignore_errors=True) + + def clear(self) -> None: + """Remove all entries of this namespace and version.""" + shutil.rmtree(self.directory, ignore_errors=True) + + +def hash_cache_key(*parts: Any) -> str: + """Create a cache key from geometries, strings, numbers and numpy arrays. + + Geometries are hashed with their ``update_hash`` method, so the key only + depends on the geometry's values (CRS, shape and extent for areas) and not + on metadata such as its name. + """ + import numpy as np + + the_hash = hashlib.sha1() # nosec: B324 + for part in parts: + if hasattr(part, "update_hash"): + part.update_hash(the_hash) + elif isinstance(part, np.ndarray): + the_hash.update(np.ascontiguousarray(part)) + else: + the_hash.update(json.dumps(part).encode("utf8")) + return the_hash.hexdigest() diff --git a/pyresample/_config.py b/pyresample/_config.py index 9a4a2b1ae..9867f4289 100644 --- a/pyresample/_config.py +++ b/pyresample/_config.py @@ -37,6 +37,7 @@ defaults=[{ "cache_dir": platformdirs.user_cache_dir("pyresample", "pytroll"), "cache_geometry_slices": False, + "cache_resampling_luts": False, "features": { "future_geometries": False, }, diff --git a/pyresample/gradient/__init__.py b/pyresample/gradient/__init__.py index e35fb8687..6304e059e 100644 --- a/pyresample/gradient/__init__.py +++ b/pyresample/gradient/__init__.py @@ -36,12 +36,16 @@ from shapely.geometry import Polygon from pyresample import CHUNK_SIZE +from pyresample._caching import ZarrCache, hash_cache_key from pyresample.geometry import AreaDefinition, SwathDefinition, get_geostationary_bounding_box_in_lonlats from pyresample.gradient._gradient_search import one_step_gradient_indices, one_step_gradient_search -from pyresample.resampler import BaseResampler, resample_blocks +from pyresample.resampler import BaseResampler, _normalize_chunks_for_area, resample_blocks logger = logging.getLogger(__name__) +# Bump the version whenever the content or layout of the cached target coordinates changes +_DST_COORDS_CACHE = ZarrCache("gradient_search_dst_coords", version=1) + def GradientSearchResampler(source_geo_def, target_geo_def): """Create a gradient search resampler.""" @@ -245,11 +249,49 @@ def __init__(self, source_geo_def, target_geo_def): self.indices_xy = None def precompute(self, **kwargs): - """Precompute resampling parameters.""" + """Precompute resampling parameters. + + If the ``cache_resampling_luts`` configuration option is set and both the source and + the target are areas, the target projection coordinates transformed to the source + projection are read from, or on first use computed and written to, an on-disk Zarr + cache in ``cache_dir``. Note that this ``cache_dir`` is the one in ``pyresample.config``, + the ``cache_dir`` keyword argument is ignored. + """ if self.indices_xy is None: - self.indices_xy = resample_blocks(gradient_resampler_indices_block, - self.source_geo_def, [], self.target_geo_def, - chunk_size=(2, CHUNK_SIZE, CHUNK_SIZE), dtype=float) + chunk_size = (2, CHUNK_SIZE, CHUNK_SIZE) + dst_coords = self._get_cached_dst_coords(chunk_size) + if dst_coords is None: + self.indices_xy = resample_blocks(gradient_resampler_indices_block, + self.source_geo_def, [], self.target_geo_def, + chunk_size=chunk_size, dtype=float) + else: + self.indices_xy = resample_blocks(_gradient_resampler_indices_block_from_dst_coords, + self.source_geo_def, [], self.target_geo_def, + dst_arrays=[dst_coords], + chunk_size=chunk_size, dtype=float) + + def _get_cached_dst_coords(self, chunk_size): + """Get the target coordinates in the source projection from the on-disk cache. + + Returns ``None`` if caching is disabled or not applicable. Only area-to-area resampling is + cached: for swath sources the search happens in the target projection, so there is nothing + to transform, and swath targets can't be identified cheaply. + """ + if not ZarrCache.is_enabled(): + return None + if not is_area_to_area(self.source_geo_def, self.target_geo_def): + return None + src_crs_wkt = self.source_geo_def.crs_wkt + chunks, _ = _normalize_chunks_for_area(self.target_geo_def, chunk_size, float) + key = hash_cache_key(self.target_geo_def, src_crs_wkt) + cached = _DST_COORDS_CACHE.load(key, chunks={"dst_xy": chunks}) + if cached is None: + logger.info("Computing and caching target coordinates in the source projection to %s", + _DST_COORDS_CACHE.get_path(key)) + dst_xy = _get_dst_coords_in_crs(self.target_geo_def, src_crs_wkt, chunks) + _DST_COORDS_CACHE.save(key, {"dst_xy": dst_xy}) + cached = _DST_COORDS_CACHE.load(key, chunks={"dst_xy": chunks}) + return cached["dst_xy"] @ensure_data_array def compute(self, data, method="bilinear", cache_id=None, **kwargs): @@ -316,6 +358,29 @@ def gradient_resampler_indices_block(block_info, **kwargs): return gradient_resampler_indices(source_area, target_area, block_info, **kwargs) +def _gradient_resampler_indices_block_from_dst_coords(dst_xy, block_info, **kwargs): + """Do the gradient search for an area source, using target coordinates already in the source projection.""" + source_area = block_info[0]["area"] + src_coords, src_gradients = _get_area_coordinates_and_gradients(source_area) + indices_xy = _gradient_resample_indices(*src_coords, *src_gradients, dst_xy[0], dst_xy[1]) + _offset_indices_to_full_source(indices_xy, block_info) + return indices_xy + + +def _get_dst_coords_in_crs(target_area, crs_wkt, chunks): + """Get the projection coordinates of the target area transformed to *crs_wkt*, stacked as a lazy (2, y, x) array.""" + dst_x, dst_y = target_area.get_proj_coords(chunks=chunks[-2:]) + return da.map_blocks(_transform_coords_block, dst_x, dst_y, + src_crs_wkt=target_area.crs_wkt, dst_crs_wkt=crs_wkt, + new_axis=0, chunks=((2,),) + tuple(chunks[-2:]), dtype=float, + meta=np.array((), dtype=float)) + + +def _transform_coords_block(x, y, src_crs_wkt, dst_crs_wkt): + transformer = pyproj.Transformer.from_crs(src_crs_wkt, dst_crs_wkt, always_xy=True) + return np.stack(transformer.transform(x, y)) + + def gradient_resampler_indices(source_area, target_area, block_info=None, **kwargs): """Do the gradient search resampling, returning the resulting indices.""" dst_coords, src_gradients, src_coords = _get_coordinates_in_same_projection(source_area, target_area) @@ -328,13 +393,17 @@ def gradient_resampler_indices(source_area, target_area, block_info=None, **kwar src_gradient_yl, src_gradient_yp, dst_x, dst_y) + _offset_indices_to_full_source(indices_xy, block_info) + return indices_xy + + +def _offset_indices_to_full_source(indices_xy, block_info): + """Shift indices relative to the cropped source to be relative to the full source, in place.""" if block_info: y_slice, x_slice = block_info[0]["array-location"][-2:] indices_xy[0, :, :] += x_slice.start indices_xy[1, :, :] += y_slice.start - return indices_xy - def _get_coordinates_in_same_projection(source_area, target_area): target_crs = target_area.crs diff --git a/pyresample/test/conftest.py b/pyresample/test/conftest.py index 0bce5f63f..9e8544330 100644 --- a/pyresample/test/conftest.py +++ b/pyresample/test/conftest.py @@ -39,6 +39,7 @@ def reset_pyresample_config(tmpdir): """Set pyresample config to logical defaults for tests.""" test_config = { "cache_geometry_slices": False, + "cache_resampling_luts": False, "features": { "future_geometries": False, }, diff --git a/pyresample/test/test_caching.py b/pyresample/test/test_caching.py new file mode 100644 index 000000000..08e828519 --- /dev/null +++ b/pyresample/test/test_caching.py @@ -0,0 +1,128 @@ +# Copyright (C) 2026 Pyresample developers +# +# This program is free software: you can redistribute it and/or modify it under +# the terms of the GNU Lesser General Public License as published by the Free +# Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, but WITHOUT +# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS +# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more +# details. +# +# You should have received a copy of the GNU Lesser General Public License along +# with this program. If not, see . +"""Tests for the on-disk Zarr cache.""" + +import os +from unittest import mock + +import dask.array as da +import numpy as np +import pytest + +import pyresample +from pyresample._caching import ZarrCache, hash_cache_key +from pyresample.test.utils import assert_maximum_dask_computes + + +@pytest.fixture +def cache(tmp_path): + """Create a cache in a temporary directory.""" + return ZarrCache("test_luts", version=3, cache_dir=tmp_path) + + +def test_is_enabled_follows_config(): + assert not ZarrCache.is_enabled() + with pyresample.config.set(cache_resampling_luts=True): + assert ZarrCache.is_enabled() + + +def test_directory_uses_cache_dir_from_config(tmp_path): + cache = ZarrCache("test_luts", version=2) + with pyresample.config.set(cache_dir=tmp_path): + assert cache.directory == tmp_path / "resampling_luts" / "test_luts_v2" + assert cache.get_path("abc") == tmp_path / "resampling_luts" / "test_luts_v2" / "abc.zarr" + + +def test_directory_without_cache_dir_raises(): + with pyresample.config.set(cache_dir=None), pytest.raises(RuntimeError, match="cache_dir"): + _ = ZarrCache("test_luts").directory + + +def test_load_missing_entry_returns_none(cache): + assert cache.load("missing") is None + + +def test_save_and_load_roundtrip(cache): + indices = np.arange(12, dtype=np.int64).reshape(3, 4) + distances = da.linspace(0, 1, 30, chunks=7).reshape(2, 15) + + cache.save("key", {"indices": indices, "distances": distances}) + with assert_maximum_dask_computes(0): + loaded = cache.load("key", chunks={"distances": (2, 5)}) + + assert set(loaded) == {"indices", "distances"} + assert all(isinstance(arr, da.Array) for arr in loaded.values()) + assert loaded["distances"].chunks == ((2,), (5, 5, 5)) + assert loaded["indices"].dtype == indices.dtype + np.testing.assert_array_equal(loaded["indices"].compute(), indices) + np.testing.assert_array_equal(loaded["distances"].compute(), distances.compute()) + + +def test_save_does_not_overwrite_existing_entry(cache): + cache.save("key", {"arr": np.zeros(3)}) + cache.save("key", {"arr": np.ones(3)}) + + np.testing.assert_array_equal(cache.load("key")["arr"].compute(), np.zeros(3)) + + +def test_save_leaves_no_temporary_files(cache): + cache.save("key", {"arr": np.zeros(3)}) + + assert os.listdir(cache.directory) == ["key.zarr"] + + +def test_failed_save_leaves_no_entry(cache): + def _fail(*args, **kwargs): + raise ValueError("Boom") + + failing = da.zeros(3, chunks=1).map_blocks(_fail, dtype=float) + with pytest.raises(ValueError, match="Boom"): + cache.save("key", {"arr": failing}) + + assert cache.load("key") is None + assert os.listdir(cache.directory) == [] + + +def test_concurrently_written_entry_is_kept(cache): + """Test that losing the race to write an entry keeps the winner's entry.""" + def _other_process_wins(src, dst): + os.mkdir(dst) + raise OSError("Directory not empty") + + with mock.patch("pyresample._caching.os.rename", side_effect=_other_process_wins): + cache.save("key", {"arr": np.zeros(3)}) + + assert os.listdir(cache.directory) == ["key.zarr"] + + +def test_clear(cache): + cache.save("key", {"arr": np.zeros(3)}) + cache.clear() + + assert not cache.directory.exists() + + +def test_hash_cache_key_depends_on_values_not_names(create_test_area): + proj_dict = {"proj": "stere", "lat_0": 90, "lon_0": 0, "ellps": "WGS84"} + extent = (-1000, -1000, 1000, 1000) + area1 = create_test_area(proj_dict, 10, 10, extent, attrs={"name": "one"}) + area2 = create_test_area(proj_dict, 10, 10, extent, attrs={"name": "two"}) + area3 = create_test_area(proj_dict, 10, 11, extent) + + assert hash_cache_key(area1, "wkt") == hash_cache_key(area2, "wkt") + assert hash_cache_key(area1, "wkt") != hash_cache_key(area3, "wkt") + assert hash_cache_key(area1, "wkt") != hash_cache_key(area1, "other wkt") + assert hash_cache_key("a", "b") != hash_cache_key("ab") + assert hash_cache_key(np.arange(3), 1.5) == hash_cache_key(np.arange(3), 1.5) diff --git a/pyresample/test/test_gradient.py b/pyresample/test/test_gradient.py index 8f1ff1adc..fe1a3d77d 100644 --- a/pyresample/test/test_gradient.py +++ b/pyresample/test/test_gradient.py @@ -684,3 +684,118 @@ def test_resampling_geos_edge_to_mercator(): res = gradient_resampler_indices(source_area, dest_area, fill_value=np.nan) assert not np.any(np.isnan(res[:, :, -1])) + + +class TestTargetCoordinateCaching: + """Test the on-disk caching of the target coordinates in the source projection.""" + + @pytest.fixture(autouse=True) + def _small_chunks(self, monkeypatch): + """Use small chunks so that the target is split in several blocks.""" + monkeypatch.setattr("pyresample.gradient.CHUNK_SIZE", 40) + + @pytest.fixture + def src_area(self, create_test_area): + """Create a geostationary source area.""" + return create_test_area({'ellps': 'WGS84', 'h': '35785831', 'proj': 'geos'}, + 100, 100, (5550000.0, 5550000.0, -5550000.0, -5550000.0)) + + @pytest.fixture + def dst_area(self, create_test_area): + """Create a stereographic target area.""" + return create_test_area({'proj': 'stere', 'lon_0': 14.0, 'lat_0': 90.0, 'lat_ts': 60.0, 'ellps': 'bessel'}, + 102, 97, (-2717181.7304994687, -5571048.14031214, + 1378818.2695005313, -1475048.1403121399)) + + @pytest.fixture + def data(self, src_area): + """Create source data.""" + return xr.DataArray(da.from_array(np.arange(np.prod(src_area.shape), dtype=np.float64).reshape(src_area.shape), + chunks=40), dims=("y", "x")) + + @staticmethod + def _cache_entries(cache_dir): + return sorted(cache_dir.glob("resampling_luts/gradient_search_dst_coords_v*/*.zarr")) + + @staticmethod + def _resample(src_area, dst_area, data, **kwargs): + resampler = create_gradient_search_resampler(src_area, dst_area) + return resampler.resample(data, **kwargs).values + + def test_caching_is_disabled_by_default(self, src_area, dst_area, data, tmp_path): + import pyresample + + with pyresample.config.set(cache_dir=tmp_path): + self._resample(src_area, dst_area, data) + + assert not self._cache_entries(tmp_path) + + @pytest.mark.parametrize("method", ["bilinear", "nn"]) + def test_cached_results_equal_uncached(self, src_area, dst_area, data, tmp_path, method): + import pyresample + + expected = self._resample(src_area, dst_area, data, method=method) + with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): + cold = self._resample(src_area, dst_area, data, method=method) + assert len(self._cache_entries(tmp_path)) == 1 + warm = self._resample(src_area, dst_area, data, method=method) + + # The target coordinates are generated for the whole area instead of per block, which can differ in the last bit + assert not np.all(np.isnan(expected)) + np.testing.assert_allclose(cold, expected, rtol=1e-12) + np.testing.assert_array_equal(warm, cold) + + def test_cached_coordinates_are_chunked_like_the_indices(self, src_area, dst_area, tmp_path): + import pyresample + + with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): + resampler = create_gradient_search_resampler(src_area, dst_area) + resampler.precompute() + + assert resampler.indices_xy.chunks == ((2,), (40, 40, 17), (40, 40, 22)) + + def test_warm_cache_is_read_lazily_without_recomputing(self, src_area, dst_area, tmp_path): + import pyresample + from pyresample.test.utils import assert_maximum_dask_computes + + with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): + create_gradient_search_resampler(src_area, dst_area).precompute() + resampler = create_gradient_search_resampler(src_area, dst_area) + with mock.patch("pyresample.gradient._get_dst_coords_in_crs") as get_coords, \ + assert_maximum_dask_computes(0): + resampler.precompute() + + get_coords.assert_not_called() + + def test_cache_is_shared_by_sources_with_the_same_projection(self, src_area, dst_area, tmp_path): + import pyresample + + with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): + create_gradient_search_resampler(src_area, dst_area).precompute() + create_gradient_search_resampler(src_area[10:60, 20:80], dst_area).precompute() + + assert len(self._cache_entries(tmp_path)) == 1 + + def test_cache_depends_on_source_projection(self, src_area, dst_area, create_test_area, tmp_path): + import pyresample + + other_src_area = create_test_area({'ellps': 'WGS84', 'h': '35785831', 'proj': 'geos', 'lon_0': 9.5}, + 100, 100, (5550000.0, 5550000.0, -5550000.0, -5550000.0)) + with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): + create_gradient_search_resampler(src_area, dst_area).precompute() + create_gradient_search_resampler(other_src_area, dst_area).precompute() + + assert len(self._cache_entries(tmp_path)) == 2 + + def test_swath_source_is_not_cached(self, dst_area, tmp_path): + import pyresample + + lons, lats = np.meshgrid(np.linspace(0, 20, 50), np.linspace(45, 66, 50)) + lons = xr.DataArray(da.from_array(lons, chunks=25), dims=["y", "x"]) + lats = xr.DataArray(da.from_array(lats, chunks=25), dims=["y", "x"]) + src_swath = SwathDefinition(lons, lats) + data = xr.DataArray(da.ones(src_swath.shape, chunks=25), dims=["y", "x"]) + with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path), np.errstate(invalid="ignore"): + self._resample(src_swath, dst_area, data) + + assert not self._cache_entries(tmp_path) From 7721c7723c4df4096e95af98b06eabab870bba16 Mon Sep 17 00:00:00 2001 From: Panu Lahtinen Date: Tue, 20 Oct 2026 19:45:52 +0000 Subject: [PATCH 2/3] Cache gradient search indices instead of target coordinates Caching the indices skips the search entirely on a cache hit. The cache key covers both the source and the target area and the block layout, so source areas of different resolution or extent each get their own entry, and the cached indices are identical to uncached ones. Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 2 +- docs/source/howtos/configuration.rst | 13 ++-- pyresample/_caching.py | 4 +- pyresample/gradient/__init__.py | 88 ++++++++-------------------- pyresample/test/test_gradient.py | 35 ++++++----- 5 files changed, 55 insertions(+), 87 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 0021ca8e7..8695a7b93 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -120,7 +120,7 @@ anything. Known divergences: `load_resampling_info`, and `pyresample.config["cache_dir"]`, which is used only by `_caching.py`: for geometry slices (`cache_geometry_slices`, JSON) and for resampling LUTs (`cache_resampling_luts`, `ZarrCache`). `ZarrCache` is the one meant to grow: - today only gradient search uses it (area-to-area target coordinates in the source CRS), + today only gradient search uses it (area-to-area `indices_xy`, keyed by both areas), but it is designed so the kd-tree nearest/bilinear index caches can move onto it from Satpy, each under its own namespace. Gradient search ignores the `cache_dir=` kwarg. - Only `nearest` is registered in the `pyresample.resamplers` entry point group, so diff --git a/docs/source/howtos/configuration.rst b/docs/source/howtos/configuration.rst index 16090a9f3..f89ff9796 100644 --- a/docs/source/howtos/configuration.rst +++ b/docs/source/howtos/configuration.rst @@ -150,17 +150,18 @@ arrays a resampler computes from the source and target geometries alone, are cached to disk as Zarr. At the time of writing this is only done by the gradient search resampler (:func:`~pyresample.gradient.create_gradient_search_resampler`) when -resampling from an area to an area: the target area projection coordinates -transformed to the projection of the source area are cached, so they only need -to be computed once for each target area and source projection. They are -reused for all source areas sharing that projection, for example different -sub-areas or resolutions of the same geostationary satellite. +resampling from an area to an area: the source pixel indices of each target +pixel are cached, so the search is only done once for each pair of source and +target areas. Source areas differing in resolution or extent, like the +channels of different resolution of geostationary imagers, each get their own +cache entry. The chunk size used (``PYTROLL_CHUNK_SIZE``) is also part of the +cache key. The LUTs are stored in the ``resampling_luts`` subdirectory of ``cache_dir`` (see above), in a separate subdirectory for each resampler and cache format version. The first time a LUT is needed, it is computed for the whole target area and written to disk during the resampler's ``precompute`` step. Later -uses read it lazily. The cached gradient search coordinates take up to 16 bytes +uses read it lazily. The cached gradient search indices take up to 16 bytes per target pixel. When setting this as an environment variable, this should be set with the diff --git a/pyresample/_caching.py b/pyresample/_caching.py index 4ca0d2e42..eb0782fd1 100644 --- a/pyresample/_caching.py +++ b/pyresample/_caching.py @@ -138,8 +138,8 @@ class ZarrCache: This is meant for the expensive, reusable look-up tables (LUTs) that resamplers compute from their source and target geometries alone, for - example target coordinates transformed to the source projection or - neighbour indices. Each cache entry is a Zarr group holding one or more + example the gradient search indices or the kd-tree neighbour + indices. Each cache entry is a Zarr group holding one or more named arrays and is identified by a hexadecimal key, typically created with :func:`hash_cache_key` from everything the arrays depend on. Entries are stored in ``/resampling_luts/_v/.zarr`` diff --git a/pyresample/gradient/__init__.py b/pyresample/gradient/__init__.py index 6304e059e..681a7b190 100644 --- a/pyresample/gradient/__init__.py +++ b/pyresample/gradient/__init__.py @@ -43,8 +43,8 @@ logger = logging.getLogger(__name__) -# Bump the version whenever the content or layout of the cached target coordinates changes -_DST_COORDS_CACHE = ZarrCache("gradient_search_dst_coords", version=1) +# Bump the version whenever the gradient search or the layout of the cached indices changes +_INDICES_CACHE = ZarrCache("gradient_search_indices", version=1) def GradientSearchResampler(source_geo_def, target_geo_def): @@ -252,46 +252,37 @@ def precompute(self, **kwargs): """Precompute resampling parameters. If the ``cache_resampling_luts`` configuration option is set and both the source and - the target are areas, the target projection coordinates transformed to the source - projection are read from, or on first use computed and written to, an on-disk Zarr - cache in ``cache_dir``. Note that this ``cache_dir`` is the one in ``pyresample.config``, - the ``cache_dir`` keyword argument is ignored. + the target are areas, the indices are read from, or on first use computed and written to, + an on-disk Zarr cache in ``cache_dir``. Note that this ``cache_dir`` is the one in + ``pyresample.config``, the ``cache_dir`` keyword argument is ignored. """ if self.indices_xy is None: chunk_size = (2, CHUNK_SIZE, CHUNK_SIZE) - dst_coords = self._get_cached_dst_coords(chunk_size) - if dst_coords is None: - self.indices_xy = resample_blocks(gradient_resampler_indices_block, - self.source_geo_def, [], self.target_geo_def, - chunk_size=chunk_size, dtype=float) + if ZarrCache.is_enabled() and is_area_to_area(self.source_geo_def, self.target_geo_def): + self.indices_xy = self._get_cached_indices(chunk_size) else: - self.indices_xy = resample_blocks(_gradient_resampler_indices_block_from_dst_coords, - self.source_geo_def, [], self.target_geo_def, - dst_arrays=[dst_coords], - chunk_size=chunk_size, dtype=float) + self.indices_xy = self._get_indices(chunk_size) - def _get_cached_dst_coords(self, chunk_size): - """Get the target coordinates in the source projection from the on-disk cache. + def _get_indices(self, chunk_size): + return resample_blocks(gradient_resampler_indices_block, + self.source_geo_def, [], self.target_geo_def, + chunk_size=chunk_size, dtype=float) - Returns ``None`` if caching is disabled or not applicable. Only area-to-area resampling is - cached: for swath sources the search happens in the target projection, so there is nothing - to transform, and swath targets can't be identified cheaply. + def _get_cached_indices(self, chunk_size): + """Get the indices from the on-disk cache, computing and caching them first if needed. + + Only area-to-area resampling is cached, as swath geometries usually change from one + dataset to the next. The source is cropped per target block, so the block layout is + part of the key too, which keeps the cached indices identical to uncached ones. """ - if not ZarrCache.is_enabled(): - return None - if not is_area_to_area(self.source_geo_def, self.target_geo_def): - return None - src_crs_wkt = self.source_geo_def.crs_wkt chunks, _ = _normalize_chunks_for_area(self.target_geo_def, chunk_size, float) - key = hash_cache_key(self.target_geo_def, src_crs_wkt) - cached = _DST_COORDS_CACHE.load(key, chunks={"dst_xy": chunks}) + key = hash_cache_key(self.source_geo_def, self.target_geo_def, chunks) + cached = _INDICES_CACHE.load(key, chunks={"indices_xy": chunks}) if cached is None: - logger.info("Computing and caching target coordinates in the source projection to %s", - _DST_COORDS_CACHE.get_path(key)) - dst_xy = _get_dst_coords_in_crs(self.target_geo_def, src_crs_wkt, chunks) - _DST_COORDS_CACHE.save(key, {"dst_xy": dst_xy}) - cached = _DST_COORDS_CACHE.load(key, chunks={"dst_xy": chunks}) - return cached["dst_xy"] + logger.info("Computing and caching gradient search indices to %s", _INDICES_CACHE.get_path(key)) + _INDICES_CACHE.save(key, {"indices_xy": self._get_indices(chunk_size)}) + cached = _INDICES_CACHE.load(key, chunks={"indices_xy": chunks}) + return cached["indices_xy"] @ensure_data_array def compute(self, data, method="bilinear", cache_id=None, **kwargs): @@ -358,29 +349,6 @@ def gradient_resampler_indices_block(block_info, **kwargs): return gradient_resampler_indices(source_area, target_area, block_info, **kwargs) -def _gradient_resampler_indices_block_from_dst_coords(dst_xy, block_info, **kwargs): - """Do the gradient search for an area source, using target coordinates already in the source projection.""" - source_area = block_info[0]["area"] - src_coords, src_gradients = _get_area_coordinates_and_gradients(source_area) - indices_xy = _gradient_resample_indices(*src_coords, *src_gradients, dst_xy[0], dst_xy[1]) - _offset_indices_to_full_source(indices_xy, block_info) - return indices_xy - - -def _get_dst_coords_in_crs(target_area, crs_wkt, chunks): - """Get the projection coordinates of the target area transformed to *crs_wkt*, stacked as a lazy (2, y, x) array.""" - dst_x, dst_y = target_area.get_proj_coords(chunks=chunks[-2:]) - return da.map_blocks(_transform_coords_block, dst_x, dst_y, - src_crs_wkt=target_area.crs_wkt, dst_crs_wkt=crs_wkt, - new_axis=0, chunks=((2,),) + tuple(chunks[-2:]), dtype=float, - meta=np.array((), dtype=float)) - - -def _transform_coords_block(x, y, src_crs_wkt, dst_crs_wkt): - transformer = pyproj.Transformer.from_crs(src_crs_wkt, dst_crs_wkt, always_xy=True) - return np.stack(transformer.transform(x, y)) - - def gradient_resampler_indices(source_area, target_area, block_info=None, **kwargs): """Do the gradient search resampling, returning the resulting indices.""" dst_coords, src_gradients, src_coords = _get_coordinates_in_same_projection(source_area, target_area) @@ -393,17 +361,13 @@ def gradient_resampler_indices(source_area, target_area, block_info=None, **kwar src_gradient_yl, src_gradient_yp, dst_x, dst_y) - _offset_indices_to_full_source(indices_xy, block_info) - return indices_xy - - -def _offset_indices_to_full_source(indices_xy, block_info): - """Shift indices relative to the cropped source to be relative to the full source, in place.""" if block_info: y_slice, x_slice = block_info[0]["array-location"][-2:] indices_xy[0, :, :] += x_slice.start indices_xy[1, :, :] += y_slice.start + return indices_xy + def _get_coordinates_in_same_projection(source_area, target_area): target_crs = target_area.crs diff --git a/pyresample/test/test_gradient.py b/pyresample/test/test_gradient.py index fe1a3d77d..b49390206 100644 --- a/pyresample/test/test_gradient.py +++ b/pyresample/test/test_gradient.py @@ -686,8 +686,8 @@ def test_resampling_geos_edge_to_mercator(): assert not np.any(np.isnan(res[:, :, -1])) -class TestTargetCoordinateCaching: - """Test the on-disk caching of the target coordinates in the source projection.""" +class TestIndicesCaching: + """Test the on-disk caching of the gradient search indices.""" @pytest.fixture(autouse=True) def _small_chunks(self, monkeypatch): @@ -715,7 +715,7 @@ def data(self, src_area): @staticmethod def _cache_entries(cache_dir): - return sorted(cache_dir.glob("resampling_luts/gradient_search_dst_coords_v*/*.zarr")) + return sorted(cache_dir.glob("resampling_luts/gradient_search_indices_v*/*.zarr")) @staticmethod def _resample(src_area, dst_area, data, **kwargs): @@ -740,12 +740,11 @@ def test_cached_results_equal_uncached(self, src_area, dst_area, data, tmp_path, assert len(self._cache_entries(tmp_path)) == 1 warm = self._resample(src_area, dst_area, data, method=method) - # The target coordinates are generated for the whole area instead of per block, which can differ in the last bit assert not np.all(np.isnan(expected)) - np.testing.assert_allclose(cold, expected, rtol=1e-12) - np.testing.assert_array_equal(warm, cold) + np.testing.assert_array_equal(cold, expected) + np.testing.assert_array_equal(warm, expected) - def test_cached_coordinates_are_chunked_like_the_indices(self, src_area, dst_area, tmp_path): + def test_cached_indices_are_chunked_like_uncached(self, src_area, dst_area, tmp_path): import pyresample with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): @@ -754,36 +753,40 @@ def test_cached_coordinates_are_chunked_like_the_indices(self, src_area, dst_are assert resampler.indices_xy.chunks == ((2,), (40, 40, 17), (40, 40, 22)) - def test_warm_cache_is_read_lazily_without_recomputing(self, src_area, dst_area, tmp_path): + def test_warm_cache_is_read_lazily_without_searching(self, src_area, dst_area, tmp_path): import pyresample from pyresample.test.utils import assert_maximum_dask_computes with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): create_gradient_search_resampler(src_area, dst_area).precompute() resampler = create_gradient_search_resampler(src_area, dst_area) - with mock.patch("pyresample.gradient._get_dst_coords_in_crs") as get_coords, \ + with mock.patch("pyresample.gradient.gradient_resampler_indices") as search, \ assert_maximum_dask_computes(0): resampler.precompute() + resampler.indices_xy.compute() - get_coords.assert_not_called() + search.assert_not_called() - def test_cache_is_shared_by_sources_with_the_same_projection(self, src_area, dst_area, tmp_path): + def test_each_source_resolution_and_extent_has_its_own_cache(self, src_area, dst_area, create_test_area, + tmp_path): import pyresample + high_res_src_area = create_test_area(src_area.crs, 200, 200, src_area.area_extent) with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): create_gradient_search_resampler(src_area, dst_area).precompute() + create_gradient_search_resampler(high_res_src_area, dst_area).precompute() create_gradient_search_resampler(src_area[10:60, 20:80], dst_area).precompute() + create_gradient_search_resampler(src_area, dst_area).precompute() - assert len(self._cache_entries(tmp_path)) == 1 + assert len(self._cache_entries(tmp_path)) == 3 - def test_cache_depends_on_source_projection(self, src_area, dst_area, create_test_area, tmp_path): + def test_cache_depends_on_chunk_size(self, src_area, dst_area, tmp_path, monkeypatch): import pyresample - other_src_area = create_test_area({'ellps': 'WGS84', 'h': '35785831', 'proj': 'geos', 'lon_0': 9.5}, - 100, 100, (5550000.0, 5550000.0, -5550000.0, -5550000.0)) with pyresample.config.set(cache_resampling_luts=True, cache_dir=tmp_path): create_gradient_search_resampler(src_area, dst_area).precompute() - create_gradient_search_resampler(other_src_area, dst_area).precompute() + monkeypatch.setattr("pyresample.gradient.CHUNK_SIZE", 50) + create_gradient_search_resampler(src_area, dst_area).precompute() assert len(self._cache_entries(tmp_path)) == 2 From cf0df1b56abcccfdb2a30cc4ec20458de2bb2eb4 Mon Sep 17 00:00:00 2001 From: Panu Lahtinen Date: Wed, 21 Oct 2026 12:28:14 +0300 Subject: [PATCH 3/3] Add optional dependencies for caching --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c2713ecae..1368c690f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ rasterio = ["rasterio"] dask = ["dask>=0.16.1"] cf = ["xarray"] gradient_search = ["shapely"] +caching = ["dask", "zarr"] xarray_bilinear = ["xarray", "dask", "zarr"] odc-geo = ["odc-geo"] tests = [ @@ -59,7 +60,7 @@ tests = [ "odc-geo", ] all = [ - "pyresample[numexpr,quicklook,rasterio,dask,cf,gradient_search,xarray_bilinear,odc-geo,tests]", + "pyresample[numexpr,quicklook,rasterio,dask,cf,gradient_search,caching,xarray_bilinear,odc-geo,tests]", ] [project.entry-points."pyresample.resamplers"]