I have just started using 'dask', and I am running into some memory issues. I would like to share my problem along with a dummy code snippet that illustrates what I’ve done. My goal is to read a large file containing high-resolution global discharge data and calculate the percentiles for each latitude-longitude cell. I am using 'dask' to chunk the data. However, I am cannot chunk it in time because I need to calculate the percentiles for the entire time period contained in the 'DataArray'. After creating a local Dask client, I chunk my 'DataArray' spatially and proceed to calculate the percentile for each chunk. Finally, I attempt to save the calculated percentiles as a 'NetCDF' file. I chose to use 'NetCDF4' for this task instead of 'xarray', as I recall that 'xarray' stores data in memory, while 'NetCDF4' allows for disk writing. The issue I’m facing is that while running the code, I notice that all my workers collectively consume more than '2GB' of memory, and the input DataArray itself is already '2GB'. My expectation was that the code would calculate the percentiles for a given chunk, save that chunk to a 'NetCDF file' using the 'writer function' I wrote, and then release that information from memory to make room for the next chunk. However, I think that the code is calculating the percentiles for all chunks before saving them as NetCDF, which may be causing the high memory usage. I would greatly appreciate any assistance in understanding what might be going wrong. Thank you! import numpy as np import netCDF4 as nc import xarray as xr import pandas as pd from dask.distributed import Client, LocalCluster def calculate_percentiles(da, percentile): return da.quantile(percentile / 100, dim="time", skipna=True) def write_percentiles_data( percentile_da: xr.DataArray, save_as: str, ) -> None: """ Save the calculated percentile into a single NetCDF file using NetCDF4. Parameters: - percentile_da: DataArray of the calculated percentile. - save_as: string, name of the NetCDF file that will be saved. """ # Open a new NetCDF file for writing with nc.Dataset(save_as, "w", format="NETCDF4") as dst: # Define dimensions dst.createDimension("lon", percentile_da.sizes["lon"]) dst.createDimension("lat", percentile_da.sizes["lat"]) # Create coordinate variables for lon and lat longitudes = dst.createVariable("lon", np.float64, "lon") latitudes = dst.createVariable("lat", np.float64, "lat") # Assign coordinate values longitudes[:] = percentile_da.coords["lon"].values latitudes[:] = percentile_da.coords["lat"].values # Define variable with dimensions (lat and lon) percentiles = dst.createVariable("percentile", np.float64, ("lat", "lon")) # Assign data directly without transposing, assuming this aligns correctly percentiles[:, :] = percentile_da.values # Use this directly # Assign attributes percentiles.units = "percentile" if __name__ == '__main__': #-------create dummy NetCDFfile time_dim = 50 lat_dim = 1400 lon_dim = 3600 # Generate time coordinates: a date range covering 120 hours from '1991-02-01' time_coords = pd.date_range("1991-02-01", periods=time_dim, freq="h") # Generate latitude and longitude coordinates lat_coords = np.linspace(83.95, -55.95, lat_dim) lon_coords = np.linspace(-179.9, 179.9, lon_dim) # Create random data data = np.random.rand(time_dim, lat_dim, lon_dim) * 100 # Create the DataArray da = xr.DataArray( data, coords=[time_coords, lat_coords, lon_coords], dims=["time", "lat", "lon"], attrs={"long_name": "routed streamflow", "unit": "m3 s-1"}, ) #-------Set up client cluster = LocalCluster( n_workers=6, threads_per_worker=4, memory_limit=3 * 1e9, # in bytes ) client = Client(cluster) #------- chunking chunks = {"time": -1, "lat": 200, "lon": 200} chunked_da = da.chunk(chunks) #-------- Calculate percentiles using dask percentile = 95 # getting 95th percentile percentile_da = calculate_percentiles(chunked_da, percentile) #-------- Save percentiles event_file_name = f"{percentile}th_percentile.nc" write_percentiles_data( percentile_da, event_file_name, ) Continue reading...
