How to apply a custom function on a dataset's chunks to reduce its size and save it on disk ?

[vdevauxchupin]

Hello,

This question is a clone of something I asked on StackOverflow.

I have a dataset (110 Gb) that is created by using combine_first on a list gathering several datasets (saved as zarr) on a remote server.
This dataset is chunked and too big for my RAM. I want to write it on my local disk, but I overload my RAM whatever the method I use (I have 96Gb of RAM).

In short: how can I write on disk a dataset that is too big to be stored in the RAM ?

Why am I doing that ? I want to apply a computationally intensive method that goes iteratively through every 2D pixel of my dataset (x,y) and works with their entire time dimension (mid_date). I thought the best way of doing that was to first download the dataset, then work on it locally. Otherwise accessing the pixels one by one takes too much time (100000 pixels, 5 seconds/pixel access). I also asked a question on StackOverflow concerning that, but got no satisfying answer [https://stackoverflow.com/questions/76384679/chunked-xarray-load-only-1-cell-in-memory-efficiently].

###############################################################

Here is a code reproducing what I am trying to do (also as a python script):

Dataset:

# import the necessary packages
import xarray as xr
import os
import sys
from dask.diagnostics import ProgressBar
import numpy as np

# Define some temporal boundaries
sdate = '2018-01-01'
edate = '2023-06-02'

# URLs of the datasets I want to merge as one
cubes = ['http://its-live-data.s3.amazonaws.com/datacubes/v02/N60W140/ITS_LIVE_vel_EPSG3413_G0120_X-3250000_Y350000.zarr',
 'http://its-live-data.s3.amazonaws.com/datacubes/v02/N50W130/ITS_LIVE_vel_EPSG3413_G0120_X-3350000_Y250000.zarr',
 'http://its-live-data.s3.amazonaws.com/datacubes/v02/N60W130/ITS_LIVE_vel_EPSG3413_G0120_X-3250000_Y250000.zarr',
 'http://its-live-data.s3.amazonaws.com/datacubes/v02/N50W140/ITS_LIVE_vel_EPSG3413_G0120_X-3350000_Y350000.zarr']

# Define spatial boundaries so you don't end up downloading tens of terabytes
xmin =-3340199.841439601
xmax =-3266124.863581888
ymin =273869.3657119706
ymax =364857.29614944384


# Create path to the save the merged dataset
pathsave = 'Datasets/'
# Create a subfolder
os.makedirs(pathsave, exist_ok = True)
# Here we drop all the variables except from those ones otherwise the cube is even larger
variables_to_keep = ['v', 'date_dt', 'acquisition_date_img1', 'acquisition_date_img2']

# Function to load the datasets
def load_datacubes(datacube_address, sdate, edate, xmax, xmin, ymax, ymin, variables_drop):

    # Grab the time values
    t = xr.open_dataset(datacube_address, engine='zarr').mid_date.values

    # Create a time mask, based on the validity of layers and the custom date-range
    t_mask = np.logical_and(t>np.datetime64(sdate), t<np.datetime64(edate))

    # Load datacube according to prerequisites (time, space and variables)
    ds = xr.open_zarr(datacube_address,
                chunks=({'mid_date': 3000, 'y': 300, 'x': 300}),
                drop_variables=variables_drop
                ).sel(  mid_date = t_mask,
                        x=slice(xmin, xmax),
                        y=slice(ymax, ymin)).drop_duplicates(dim='mid_date')
    ds.attrs={}
    return ds

# Add the dimensions to the download otherwise they are dropped
variables_to_keep += ['mid_date', 'x', 'y']

# List of variables to drop for the download (we drop everything but the variables in "variables to keep")
variables_drop = [ele for ele in list(
        xr.open_dataset(cubes[0], engine='zarr').variables
        ) if ele not in variables_to_keep
]

# Load the cubes, chunked
datacubes = [load_datacubes(datacube_address, sdate, edate, xmax, xmin, ymax, ymin, variables_drop).sortby(['mid_date','y','x']).chunk({'mid_date': 3000, 'y': 300, 'x': 300}) for datacube_address in cubes]


# Combine the datacubes into one (stitches the cubes spatially, and aligns them temporally when they have the same timestamps)
xrds = datacubes[0]
xrds = xrds.combine_first(datacubes[1])
xrds = xrds.combine_first(datacubes[2])
xrds = xrds.combine_first(datacubes[3])

#Option 1, did not work
#write_job = xrds.to_zarr(f"{pathsave}Test.zarr", mode='w', compute=False, consolidated=True)

#Option 2, did not work (with chunk size {'mid_date':3000, 'x':100, 'y':100})


#write_job = xrds.to_netcdf(f"test.nc",compute=False)
#with ProgressBar():
#    print(f"Writing to {pathsave}")
#    write_job = write_job.compute()                 

# Option 3, did not work. That's the option I took the screenshot from
# I force the chunks to be really small so I don't overload the memory
chunk_size = {'mid_date':1, 'y':xrds.y.shape[0], 'x':xrds.x.shape[0]}

with ProgressBar():
    for i, (key, chunk) in enumerate(xrds.chunk(chunk_size).items()):
        chunk_dataset = xr.Dataset({key: chunk})
        chunk_dataset.to_netcdf(f"{pathsave}/chunk_{i}.nc", mode="w", compute=True)
        print(f"Chunk {i+1} saved.")    

[dcherian]

Does any of the material in https://e-marshall.github.io/itslive/intro.html help?

cc @e-marshall

[e-marshall]

The tutorial didn't go into anything that required loading into memory IIRC

[e-marshall]

Hey @vdevauxchupin ! I've struggled with this type of thing of thing with ITSLIVE before but don't think I've figured out much of a solution, and haven't figured out a way to do directly what you're asking. A few thoughts:

• can you reduce the spatial size / x,y dimensions of the object any further ? (most of the glaciers I've been working with are relatively small which has worked to my advantage at times ) . In the past I have called compute() on an object (smaller than the one you're working with) right before a non-lazy operation and then del object right after and that seems to work. it hasn't been cleaned but there's an example of that in the def calc_glacier_medians_1granule function in this notebook (all the way at the bottom)
• what operation do you want to apply to each 2d xy array of the object? Is it one that could be implemented in a vectorized way with apply_ufunc, map_blocks() or reduce()

[vdevauxchupin]

Hi @e-marshall ! Thanks for your answer.

There are several ways I sometimes reduce the dataset size:

• Use the quality index of each slice (amount of NaNs per slice), and kick out every slice below a quality threshold. CON: you get rid of some potentially interesting data

• Use a glacier mask to download only the pixels on the ice. CON: annoying to automatize and you need yet another dataset for that (but it seems you've managed to do that ?)

• Invert a velocity function (to be applied on each 2d xy array) that takes what you have, but can be resampled in order to get 1 value per day (or 1 value per 5 days, or whatever you want). CON: takes a while to compute (9h for 11Gb of data, non-parallelized). But that function is the one that reduces the dataset's size the best way while ensuring you don't mix yearly or seasonal signals..

• Use the resample capabilities of xarray to force 1 value per day or less. CON: the result is wrong, as the velocities are generated from different time intervals (you end up mixing seasonal and yearly signals altogether). The outcome is the same as the inverse function, only this method is not scientifically correct.

I did find a way to make it work somehow:

• Not use combine_first(), but instead save every chunked dataset individually, apply the inverse function, and then combine the smaller datasets together. CON: Will take approximately 4 days of computation, unless I manage to parallelize my inverse function.

But in the end the issues remain the same: if you don't have a system with a lot of RAM (at least 64Gb), this becomes a complicated issue.

[e-marshall]

nice! let's add this to the list of things to talk about when we meet - this seems like an unavoidable step in these workflows so would be great to brainstorm together some. the inverse function seems like a neat approach. i've mostly been using the glacier mask approach and wrapping clipping by a glacier outline with some other processing steps, but its also very slow...

[vdevauxchupin]

Oh I see ! Sometimes I wish matrices could have other shapes than rectangles.
I let the use clip the dataset to their AOI with the mouse on the GUI, it might not be as precise as the glacier limits but it shouldn't change much in the end.
However the inversion could definitely be sped-up if it is applied to the glacier pixels only, instead of everything on the image

[TomNicholas]

I want to apply a computationally intensive method that goes iteratively through every 2D pixel of my dataset (x,y) and works with their entire time dimension (mid_date). I thought the best way of doing that was to first download the dataset, then work on it locally.

It sounds more like you should do this:

• Open the dataset lazily with a chunking scheme like {'x': 100, 'y': 100. 'time': -1} (so same as chunks on disk expect for opening all time chunks as contiguous).
• Then you want to write a function that can happily apply your analysis to a single chunk after loading just that chunk into memory. Because you want to apply elementwise (to each pixel) you'll need to look into what np.vectorize does.
• Then one that works for one chunk you can worry about using it on many chunks in one go. The function you wrote in the last step can be applied with xr.apply_ufunc, and now you can use dask='parallelized' to avoid the bigger than RAM issue. Of course this assumes that you didn't just try to download the whole thing locally. Do you have access to a jupyterhub or similar where you could run this using dask?

[vdevauxchupin]

It should be, but only over dask chunks, not within them.
Ok I see, thank you !

Here's a WIP preview of reworked apply_ufunc tutorial material just in case, though it looks like you got things working

Thank you for the link ! It;s super useful. Without tutorials I have trouble understanding otherwise.

[TomNicholas]

Great. Has your original question been answered now? If not you can always open another discussion for any follow-up questions too.

[vdevauxchupin]

I will close this question but maybe modify the title since the original question hasn't really been answered: how to write on disk a dataset too big for the RAM ?

[TomNicholas]

To answer your original question you should be able to follow the same strategy that you now are using to process the data chunk-by-chunk to instead write the data chunk-by-chunk. Basically instead of applying your analysis function you use to_zarr or to_netcdf instead, and run it with dask.

[vdevauxchupin]

Ok I will try that. Thank you for helping me through that !