Resampling with Rioxarray (S3 storage, NetCDF file, H5NetCDF driver, earthaccess auth)

import argparse

import earthaccess
import xarray as xr
from rasterio.warp import calculate_default_transform

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]
    input_uri = f'{args["folder"]}/{args["filename"]}'
    src = f's3://{args["bucket"]}/{input_uri}'
    # Define source and target projection
    dstSRS = "EPSG:3857"
    srcSRS = "EPSG:4326"
    width = height = 256
    # Authenticate with earthaccess
    fs = earthaccess.get_s3_filesystem(daac=args["daac"])
    # Specify fsspec caching since default options don't work well for raster data
    fsspec_caching = {
        "cache_type": "none",
    }
    with fs.open(src, **fsspec_caching) as f:
        # Open dataset
        da = xr.open_dataset(f, engine="h5netcdf", mask_and_scale=True)[
            args["variable"]
        ]
        if dataset == "gpm_imerg":
            # Transpose and rename dims to align with rioxarray expectations
            da = da.rename({"lon": "x", "lat": "y"}).transpose("time", "y", "x")
        # Set input dataset projection
        da = da.rio.write_crs(srcSRS)
        da = da.rio.clip_box(
            *te,
            crs=dstSRS,
        )
        # Define affine transformation from input to output dataset
        dst_transform, w, h = calculate_default_transform(
            srcSRS,
            dstSRS,
            da.rio.width,
            da.rio.height,
            *da.rio.bounds(),
            dst_width=width,
            dst_height=height,
        )
        # Reproject dataset
        return da.rio.reproject(dstSRS, shape=(h, w), transform=dst_transform)

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