Resampling with pyresample (S3 storage, Zarr V3 store, Zarr reader with icechunk)¶

import argparse

import xarray as xr
from icechunk import IcechunkStore, StorageConfig
from pyresample.area_config import create_area_def
from pyresample.gradient import block_nn_interpolator, gradient_resampler_indices_block
from pyresample.resampler import resample_blocks

def warp_resample(dataset, zoom=0):
    from common import earthaccess_args, target_extent

    te = target_extent[zoom]

    # Define filepath, driver, and variable information
    args = earthaccess_args[dataset]
    # Open dataset
    storage = StorageConfig.s3_from_env(
        bucket="nasa-veda-scratch",
        prefix=f"resampling/icechunk/{dataset}",
        region="us-west-2",
    )
    store = IcechunkStore.open_existing(storage=storage, mode="r")
    da = xr.open_zarr(store, zarr_format=3, consolidated=False)[args["variable"]]
    # Define source and target projection
    dstSRS = "EPSG:3857"
    srcSRS = "EPSG:4326"
    width = height = 256
    # Rechunk MURSST to operate on fewer chunks
    if dataset == "mursst":
        da = da.chunk({"lat": 4000, "lon": 4000})
    elif dataset == "gpm_imerg":
        # Transpose dims to align with pyresample expectations
        da = da.squeeze()
    # Create area definition for the target dataset
    target_area_def = create_area_def(
        area_id=1,
        projection=dstSRS,
        shape=(height, width),
        area_extent=te,
    )
    # Create area definition for the source dataset
    source_area_def = create_area_def(
        area_id=2,
        projection=srcSRS,
        shape=(da.sizes["lat"], da.sizes["lon"]),
        area_extent=[-179.995, 89.995, 180.005, -89.995],
    )
    # Compute indices for resampling
    indices_xy = resample_blocks(
        gradient_resampler_indices_block,
        source_area_def,
        [],
        target_area_def,
        chunk_size=(1, height, width),
        dtype=float,
    )
    # Apply resampler
    resampled = resample_blocks(
        block_nn_interpolator,
        source_area_def,
        [da.data],
        target_area_def,
        dst_arrays=[indices_xy],
        chunk_size=(1, height, width),
        dtype=da.dtype,
    )
    # Reproject dataset
    return resampled.compute()

if __name__ == "__main__":
    if "get_ipython" in dir():
        # Just call warp_resample if running as a Jupyter Notebook
        da = warp_resample("gpm_imerg")
    else:
        # Configure dataset via argpase if running via CLI
        parser = argparse.ArgumentParser(description="Set environment for the script.")
        parser.add_argument(
            "--dataset",
            default="mursst",
            help="Dataset to resample.",
            choices=["gpm_imerg", "mursst"],
        )
        parser.add_argument(
            "--zoom",
            default=0,
            help="Zoom level for tile extent.",
        )
        user_args = parser.parse_args()
        da = warp_resample(user_args.dataset, int(user_args.zoom))