Memory error after slicing xarray dataset x and y dimensions

[MorningGlory747]

What is your issue?

Hello,

Hope i'm doing this right.

I am opening 20 years of daily precipitation data from May to October on a 881x1201 grid. Each file contains one year of data.
The setup is the following

ds= xr.open_mfdataset(filenames,parallel=True)
print(ds)

<xarray.Dataset>
Dimensions:                  (time: 7518, y: 881, x: 1201)
Coordinates:
  * time                     (time) datetime64[ns] 2000-04-01 ... 2020-10-30
    lat                      (time, y, x) float32 dask.array<chunksize=(1, 881, 1201), meta=np.ndarray>
  * y                        (y) float32 1.268e+03 1.267e+03 ... 389.0 388.0
  * x                        (x) float32 1.409e+03 1.41e+03 ... 2.609e+03
    lon                      (time, y, x) float32 dask.array<chunksize=(1, 881, 1201), meta=np.ndarray>
Data variables:
    lambert_conformal_conic  (time) float64 -3.277e+04 -3.277e+04 ... -3.277e+04
    prcp                     (time, y, x) float32 dask.array<chunksize=(1, 881, 1201), meta=np.ndarray>
Attributes:
    start_year:          2000
    source:              Daymet Software Version 4.0
    Version_software:    Daymet Software Version 4.0
    Version_data:        Daymet Data Version 4.0
    Conventions:         CF-1.6
    citation:            Please see http://daymet.ornl.gov/ for current Dayme...
    references:          Please see http://daymet.ornl.gov/ for current infor...
    History:             Translated to CF-1.0 Conventions by Netcdf-Java CDM ...
    geospatial_lat_min:  41.149529826818124
    geospatial_lat_max:  52.67908572537513
    geospatial_lon_min:  -81.22510372987668
    geospatial_lon_max:  -61.64963532370754

I need to compute various indices and I like looking at the values from the data to make sure that the code is working well.
To reduce the computational load during these verifications, I select only a specific x, y index, such as follows :

ds_small = ds['prcp'].isel(y = 0, x = 0)
print(ds_small)

<xarray.DataArray 'prcp' (time: 7518)>
dask.array<getitem, shape=(7518,), dtype=float32, chunksize=(154,), chunktype=numpy.ndarray>
Coordinates:
  * time     (time) datetime64[ns] 2000-04-01 2000-04-02 ... 2020-10-30
    lat      (time) float32 dask.array<chunksize=(1,), meta=np.ndarray>
    y        float32 1.268e+03
    x        float32 1.409e+03
    lon      (time) float32 dask.array<chunksize=(1,), meta=np.ndarray>
Attributes:
    long_name:     daily total precipitation
    units:         mm/day
    grid_mapping:  lambert_conformal_conic
    cell_methods:  area: mean time: sum
    _ChunkSizes:   [   1 1000 1000]

However, when comes the time where I want to run manipulations of numpy or panda (may that be from using .to_numpy(), .to_dataframe(), .compute(), .load()), the process will run very slowly and the following memory error will appear :

ds_small.compute()

C:\temp\Lib\site-packages\dask\array\numpy_compat.py:39: RuntimeWarning: invalid value encountered in true_divide
  x = np.divide(x1, x2, out)

Traceback (most recent call last):
  File "C:\temp\Lib\site-packages\IPython\core\interactiveshell.py", line 3441, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-40-97898f43589b>", line 1, in <module>
    ds_small.compute()
  File "C:\temp\Lib\site-packages\xarray\core\dataarray.py", line 955, in compute
    return new.load(**kwargs)
  File "C:\temp\Lib\site-packages\xarray\core\dataarray.py", line 929, in load
    ds = self._to_temp_dataset().load(**kwargs)
  File "C:\temp\Lib\site-packages\xarray\core\dataset.py", line 865, in load
    evaluated_data = da.compute(*lazy_data.values(), **kwargs)
  File "C:\temp\Lib\site-packages\dask\base.py", line 568, in compute
    results = schedule(dsk, keys, **kwargs)
  File "C:\temp\Lib\site-packages\dask\threaded.py", line 79, in get
    results = get_async(
  File "C:\temp\Lib\site-packages\dask\local.py", line 517, in get_async
    raise_exception(exc, tb)
  File "C:\temp\Lib\site-packages\dask\local.py", line 325, in reraise
    raise exc
  File "C:\temp\Lib\site-packages\dask\local.py", line 223, in execute_task
    result = _execute_task(task, data)
  File "C:\temp\Lib\site-packages\dask\core.py", line 121, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "C:\temp\Lib\site-packages\dask\array\core.py", line 105, in getter
    c = np.asarray(c)
  File "C:\temp\Lib\site-packages\numpy\core\_asarray.py", line 102, in asarray
    return array(a, dtype, copy=False, order=order)
  File "C:\temp\Lib\site-packages\xarray\core\indexing.py", line 354, in __array__
    return np.asarray(self.array, dtype=dtype)
  File "C:\temp\Lib\site-packages\numpy\core\_asarray.py", line 102, in asarray
    return array(a, dtype, copy=False, order=order)
  File "C:\temp\Lib\site-packages\xarray\core\indexing.py", line 518, in __array__
    return np.asarray(self.array, dtype=dtype)
  File "C:\temp\Lib\site-packages\numpy\core\_asarray.py", line 102, in asarray
    return array(a, dtype, copy=False, order=order)
  File "C:\temp\Lib\site-packages\xarray\core\indexing.py", line 419, in __array__
    return np.asarray(array[self.key], dtype=None)
  File "C:\temp\Lib\site-packages\numpy\core\_asarray.py", line 102, in asarray
    return array(a, dtype, copy=False, order=order)
  File "C:\temp\Lib\site-packages\xarray\coding\variables.py", line 70, in __array__
    return self.func(self.array)
  File "C:\temp\Lib\site-packages\xarray\coding\variables.py", line 137, in _apply_mask
    data = np.asarray(data, dtype=dtype)
  File "C:\temp\Lib\site-packages\numpy\core\_asarray.py", line 102, in asarray
    return array(a, dtype, copy=False, order=order)
  File "C:\temp\Lib\site-packages\xarray\core\indexing.py", line 419, in __array__
    return np.asarray(array[self.key], dtype=None)
  File "C:\temp\Lib\site-packages\xarray\backends\netCDF4_.py", line 91, in __getitem__
    return indexing.explicit_indexing_adapter(
  File "C:\temp\Lib\site-packages\xarray\core\indexing.py", line 710, in explicit_indexing_adapter
    result = raw_indexing_method(raw_key.tuple)
  File "C:\temp\Lib\site-packages\xarray\backends\netCDF4_.py", line 104, in _getitem
    array = getitem(original_array, key)
  File "netCDF4\_netCDF4.pyx", line 4451, in netCDF4._netCDF4.Variable.__getitem__
  File "netCDF4\_netCDF4.pyx", line 5376, in netCDF4._netCDF4.Variable._get
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 860. MiB for an array with shape (213, 881, 1201) and data type float32

But this doesn't make sense to me, why would xarray still think I have a 881,1201 dataset shape after I've selected specific indices (in this case, x = 0 and y = 0)?
I feel like i'm not understanding something trivial about xarray and I'd greatly appreciate any help!

`INSTALLED VERSIONS

commit: None
python: 3.9.10 (tags/v3.9.10:f2f3f53, Jan 17 2022, 15:14:21) [MSC v.1929 64 bit (AMD64)]
python-bits: 64
OS: Windows
OS-release: 10
machine: AMD64
processor: Intel64 Family 6 Model 142 Stepping 9, GenuineIntel
byteorder: little
LC_ALL: None
LANG: None
LOCALE: ('French_Canada', '1252')
libhdf5: 1.10.6
libnetcdf: 4.7.3
xarray: 0.19.0
pandas: 1.3.2
numpy: 1.20.3
scipy: 1.7.1
netCDF4: 1.5.5.1
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: None
cftime: 1.3.1
nc_time_axis: None
PseudoNetCDF: None
rasterio: 1.2.10
cfgrib: 0.9.8.1
iris: None
bottleneck: 1.3.2
dask: 2021.08.1
distributed: None
matplotlib: 3.4.3
cartopy: 0.19.0.post1
seaborn: 0.11.2
numbagg: None
pint: 0.17
setuptools: 57.4.0
pip: 22.0.4
conda: None
pytest: None
IPython: 7.27.0
sphinx: None
`

[dcherian]

chunksize=(1, 881, 1201)

This means that it will load a 1x881x1201 into memory and then subset it. I would specify chunks in open_mfdataset. See https://docs.xarray.dev/en/stable/user-guide/dask.html#optimization-tips for some tips

[MorningGlory747]

Hi!
If I understand correctly about chunking, the important thing is to aim for chunks that are neither too big nor too small. Therefore I've tried with different parameters : ds = xr.open_mfdataset(filenames,parallel=True,engine='netcdf4',chunks= {'time':1,'y': 100, 'x':120}) and now i'm able to slice one index (x = 0, y = 0) without getting memory errors. I don't know if this is optimal, but it does the trick.

That being said, converting the dataset to an array or a dataframe, even if i'm just calling one pixel, still takes up to one minute to complete. Is this to be expected? Thank you so much!

[dcherian]

chunks= {'time':1,'y': 100, 'x':120})

This seems too small. You mentioned that a year of data is in a single file, and you're extracting time-series. This suggests using small chunk sizes in space and large chunksizes in time. Try -1 for time. You want chunksizes to be in the 100-200MB range usually

[MorningGlory747]

Wow amazing! The computational speed increase after specifying chunks= {'time':-1,'y': 300, 'x':250} is incredible. Thank you!

The notion of chunks was abstract to me and so I was just putting it away, but now I see how pivotal it is. That being said, why isn't slicing one pixel sufficient to increase the speed? Shouldn't slicing be independent of the way the chunks were specified? Unless what actually was slowing down the computation (after slicing) was how time was chunked and not the space?

Also what's the difference between specifying -1 instead of 1 for time? Apparently -1 does no chunking in that dimension, but how is that different from having one chunk?

Lastly, is there a way to see the mb range of the chunks?

Another edit sorry, apparently there is an "auto" feature for ideal chunk size and that appears to be (213, 339, 463)

[delgadom]

the difference between chunks of -1 and 1 is that specifying a chunk size of -1 will set the chunk size to the length of the dimension. For your data:

• specifying chunks={'time': 1, 'y': 300, 'x': 250} results in chunk size of --> (time: 1, y: 300, x: 250) (approx 293 kB per chunk)
• specifying chunks={'time': -1, 'y': 300, 'x': 250} results in chunk size of --> (time: 7518, y: 300, x: 250) (approx 2 GB per chunk)

The latter is actually a little on the large side - as dcherian mentioned you want these to be ~100-200 MB in size. You can back-of-the-envelope the size in memory by multiplying the dimensions and the byte size of the data type (float32 = 32 bits = 4 bytes). This is why the auto chunks suggested were (213, 339, 463) (~128MB). For your case, I'd probably recommend (-1, 50, 100) (~143MB) or something like that.

The performance increase comes because of a critical tradeoff: minimizing I/O and maximizing parallelism vs. increased overhead of dask tasks relative to numpy. Numpy will always be faster than dask+numpy when working on the same sized array in terms of CPU time. Think of dask user time as roughly (numpy + overhead * number of tasks) / number of threads. When you specify chunks of (time: 1, y: 300, x: 250), xarray creates one dask task for every timestep in your array! you then read in that chunk and throw away everything except one value, return that to dask, and dask has to assemble the time series. On the other hand, for chunks of (-1, 50, 100), you read in less data (because the time series you want makes up a larger share of each chunk) and then your whole time series is in memory already, and numpy just calculates whatever you need within that thread. There's no parallelism in the data processing in this case, and you could play with creating a few chunks along the time dimension to see if that speeds things up even more. But it looks like you're not doing any computation before loading the data, so the most important thing will probably be simply tuning your I/O to do as little extra work as possible.