Add on-disk Zarr caching of gradient search target coordinates

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

docs/source/howtos/configuration.rst

@@ -94,7 +94,8 @@ the `platformdirs <https://github.com/platformdirs/platformdirs#example-output>`
 
    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,41 @@ 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 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 indices 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

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"]

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 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 ``<cache_dir>/resampling_luts/<namespace>_v<version>/<key>.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()

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,
         },

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 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):
     """Create a gradient search resampler."""
@@ -245,11 +249,40 @@ 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 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:
-            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)
+            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 = self._get_indices(chunk_size)
+
+    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)
+
+    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.
+        """
+        chunks, _ = _normalize_chunks_for_area(self.target_geo_def, chunk_size, float)
+        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 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):

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,
         },