Store processed data in Pooch cache

[mdtanker]

I'm writing a function to retrieve and apply some processing to a dataset using Pooch. While subsequent calls to the function don't re-download the data, they do result in re-running the processing steps, which takes some time. Is there a way for Pooch to recognize that the processing has already occurred, and to retrieve the processed file instead of the un-processed file?

Here's a simplified version of the function:

def ice_vel(
) -> xr.DataArray:
    """
    MEaSUREs Phase-Based Antarctica Ice Velocity Map, version 1: 
    https://nsidc.org/data/nsidc-0754/versions/1#anchor-1
    Data part of https://doi.org/10.1029/2019GL083826

    Returns
    -------
    xr.DataArray
        _description_
    """

    path = pooch.retrieve(
        url="https://n5eil01u.ecs.nsidc.org/MEASURES/NSIDC-0754.001/1996.01.01/antarctic_ice_vel_phase_map_v01.nc", # noqa
        downloader=EarthDataDownloader(),
        known_hash=None,
        progressbar=True,)

    grd = xr.load_dataset(path)
    
    vel = (grd.VX**2 + grd.VY**2)**0.5

    return vel

[leouieda]

Hi @mdtanker thanks for asking! This is where you'd want to make a custom processor. The processor allows you to insert custom code between the file download and the return of the path.

For example, this code downloads one of our sample grids, converts the units to m/s², saves the converted grid, and returns the path to it instead:

from pathlib import Path


def custom_processor(fname, action, pooch):
    "Load the data, calculate something, save it back."
    fname = Path(fname)
    # Rename the file to ***-processed.nc
    fname_processed = fname.with_stem(fname.stem + "-processed")
    # Only recalculate if making a new download or the processed file doesn't exist yet
    if action in ("download", "update") or not fname_processed.exists():
        grid = xr.load_dataarray(fname)
        # Convert from mGal to m/s²
        processed = grid * 1e-5
        # Save to disk
        processed.to_netcdf(fname_processed)
    return str(fname_processed)


# This is now the path to the processed grid (in m/s²)
path = pooch.retrieve(
    url="doi:10.5281/zenodo.5882207/earth-gravity-10arcmin.nc", 
    known_hash="md5:56df20e0e67e28ebe4739a2f0357c4a6",
    progressbar=True,
    processor=custom_processor,
)
grd = xr.load_dataset(path)

[mdtanker]

Hi @leouieda, I had a follow up question on this same topic. I need to use retrieve to download and unzip the contents of a folder. Is there a way to also apply a pre-processing step to a file in the unzipped folder? I'm envisioning passing both the pooch.Unzip and custom_processor function to the processor argument of pooch.retrieve. Is this possible?

[leouieda]

@mdtanker the solution there is to call the Unzip processor inside your custom_processor. Pooch only allows one processor but you can chain them together inside a custom one.

[mdtanker]

Thanks 😄 I assumed that would be the case, but my initial try didn't work. I was getting an error since pooch was passed as an argument to the custom processor function, when I tried including pooch.retrieve() within the function it was giving an error. Renaming the argument pooch to something different (like pooch2) fixed it. Here's an example if anyone else needs to use it.

def preprocessing(fname, action, pooch2):
            "Unzip the folder, grid the .dat file, and save it back as a .nc"
            path = pooch.retrieve(
                url="https://ftp.space.dtu.dk/pub/RF/4D-ANTARCTICA/ant4d_gravity.zip",
                known_hash=None,
                processor=pooch.Unzip(),
                progressbar=True,
            )
            fname = [p for p in path if p.endswith(".dat")][0]

            fname = Path(fname)
            # Rename to the file to ***_preprocessed.nc
            fname_pre = fname.with_stem(fname.stem + f"_{anomaly_type}_preprocessed")
            fname_processed = fname_pre.with_suffix('.nc')

            # Only recalculate if new download or the processed file doesn't exist yet
            if action in ("download", "update") or not fname_processed.exists():
                # load data
                df = pd.read_csv(
                    fname,
                    delim_whitespace=True,
                    skiprows=3,
                    names=["id", "lat", "lon", "FA", "Err", "DG", "BA"],
                )
                # re-project to polar stereographic
                transformer = Transformer.from_crs("epsg:4326", "epsg:3031")
                df["x"], df["y"] = transformer.transform(df.lat.tolist(), df.lon.tolist()) # noqa
                
                # block-median and grid the data
                df = pygmt.blockmedian(
                    df[["x", "y", anomaly_type]], 
                    spacing=initial_spacing, 
                    region=initial_region, 
                    registration=initial_registration,
                )
                processed = pygmt.surface(
                    data=df[["x", "y", anomaly_type]], 
                    spacing=initial_spacing, 
                    region=initial_region, 
                    registration=initial_registration,
                    M="1c", 
                )
                # Save to disk
                processed.to_netcdf(fname_processed)
            return str(fname_processed)

        path = pooch.retrieve(
            url="https://ftp.space.dtu.dk/pub/RF/4D-ANTARCTICA/ant4d_gravity.zip",
            known_hash=None,
            processor=preprocessing,
            progressbar=True,
        )

[leouieda]

Ah right! The call to pooch.retrieve in the preprocessing isn't needed. The download is handled before calling the processor, so the call inside preprocessing isn't downloading ever. What I meant to say was:

def preprocessing(fname, action, pooch2):
    "Unzip the folder, grid the .dat file, and save it back as a .nc"
    path = pooch.Unzip()(fname, action, pooch2)
    fname = [p for p in path if p.endswith(".dat")][0]

    fname = Path(fname)
    # Rename to the file to ***_preprocessed.nc
    fname_pre = fname.with_stem(fname.stem + f"_{anomaly_type}_preprocessed")
    fname_processed = fname_pre.with_suffix('.nc')

    # Only recalculate if new download or the processed file doesn't exist yet
    if action in ("download", "update") or not fname_processed.exists():
        # load data
        df = pd.read_csv(
            fname,
            delim_whitespace=True,
            skiprows=3,
            names=["id", "lat", "lon", "FA", "Err", "DG", "BA"],
        )
        # re-project to polar stereographic
        transformer = Transformer.from_crs("epsg:4326", "epsg:3031")
        df["x"], df["y"] = transformer.transform(df.lat.tolist(), df.lon.tolist()) # noqa
        
        # block-median and grid the data
        df = pygmt.blockmedian(
            df[["x", "y", anomaly_type]], 
            spacing=initial_spacing, 
            region=initial_region, 
            registration=initial_registration,
        )
        processed = pygmt.surface(
            data=df[["x", "y", anomaly_type]], 
            spacing=initial_spacing, 
            region=initial_region, 
            registration=initial_registration,
            M="1c", 
        )
        # Save to disk
        processed.to_netcdf(fname_processed)
    return str(fname_processed)

path = pooch.retrieve(
    url="https://ftp.space.dtu.dk/pub/RF/4D-ANTARCTICA/ant4d_gravity.zip",
    known_hash=None,
    processor=preprocessing,
    progressbar=True,
)

[mdtanker]

Ahh ok yes that makes sense. I wasn't sure how to call Unzip like that. Thanks again!