Fix `NotImplementedError: xarray can't set arrays with multiple array indices to dask yet.`

[tomvothecoder]

What versions of software are you using?

• Package Version: latest master

What are the steps to reproduce this issue?

import xcdat

# specify chunks to open with Dask
ds = xcdat.open_dataset("path/to/file.nc", var="tas", chunks={"time": 5})

# call spatial averaging and receive 
ts_n34 = ds.spatial.avg("tas", axis=["lat", "lon"], lat_bounds=(-5, 5), lon_bounds=(-170, -120))["tas"]

What happens? Any logs, error output, etc?

```python
NotImplementedError: xarray can't set arrays with multiple array indices to dask yet.
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
~/Documents/Repositories/tomvothecoder/xcdat/qa/PR98/refactor.py in <module>
      24 # ---------------------
      25 # NotImplementedError: xarray can't set arrays with multiple array indices to dask yet.
----> 26 avg = ds.spatial.avg("tas", axis, lat_bounds=lat_bounds, lon_bounds=lon_bounds)

~/Documents/Repositories/tomvothecoder/xcdat/xcdat/spatial_avg.py in avg(self, data_var_name, axis, weights, lat_bounds, lon_bounds)
    137             if lon_bounds is not None:
    138                 self._validate_region_bounds("lon", lon_bounds)
--> 139             weights = self._get_weights(axis, lat_bounds, lon_bounds)
    140 
    141         self._validate_weights(data_var, axis, weights)

~/Documents/Repositories/tomvothecoder/xcdat/xcdat/spatial_avg.py in _get_weights(self, axis, lat_bounds, lon_bounds)
    270                             domain_bounds, reg_bounds
    271                         )
--> 272                     dom_bounds = self._scale_domain_to_region(dom_bounds, reg_bounds)
    273 
    274                 if dim == "lat":

~/Documents/Repositories/tomvothecoder/xcdat/xcdat/spatial_avg.py in _scale_domain_to_region(self, domain_bounds, region_bounds)
    420             # Case 1: not wrapping around prime meridian.
    421             # Adjustments for above / right of region.
--> 422             d_bounds[d_bounds[:, 0] > r_bounds[1], 0] = r_bounds[1]
    423             d_bounds[d_bounds[:, 1] > r_bounds[1], 1] = r_bounds[1]
    424             # Adjustments for below / left of region.

/opt/miniconda3/envs/xcdat_dev/lib/python3.8/site-packages/xarray/core/dataarray.py in __setitem__(self, key, value)
    765                 for k, v in self._item_key_to_dict(key).items()
    766             }
--> 767             self.variable[key] = value
    768 
    769     def __delitem__(self, key: Any) -> None:

/opt/miniconda3/envs/xcdat_dev/lib/python3.8/site-packages/xarray/core/variable.py in __setitem__(self, key, value)
    852 
    853         indexable = as_indexable(self._data)
--> 854         indexable[index_tuple] = value
    855 
    856     @property

/opt/miniconda3/envs/xcdat_dev/lib/python3.8/site-packages/xarray/core/indexing.py in __setitem__(self, key, value)
   1244                 )
   1245                 if num_non_slices > 1:
-> 1246                     raise NotImplementedError(
   1247                         "xarray can't set arrays with multiple "
   1248                         "array indices to dask yet."

NotImplementedError: xarray can't set arrays with multiple array indices to dask yet.

### What were you expecting to happen?
<!-- A clear and concise description of what you expected to happen. -->
Parallelization of chunks should occur without issue.

### Any other comments?
<!-- Add any other context about the problem here. -->
Figure out a way to use xarray to perform the scaling.

[tomvothecoder]

Hey @pochedls, maybe you can help look into this issue if you have time?

[pochedls]

Hey @pochedls, maybe you can help look into this issue if you have time?

@tomvothecoder - I just tried this with the latest release and I didn't get an error. I am not sure if it is automatically parallelizing, but I don't seem to be getting an error. Do you get this error with the latest release?

[tomvothecoder]

@pochedls I've updated the example so that region selection is performed, which raises the error.
ts_n34 = ds.spatial.avg("tas", axis=["lat", "lon"], lat_bounds=(-5, 5), lon_bounds=(-170, -120))["tas"]

[pochedls]

@tomvothecoder - here is a proposed fix. I wonder if there are going to be a lot of little places where we need to add a .load().

Note that I get a different error than the one specified in this issue. I got this error:

TypeError: this variable's data is stored in a dask array, which does not support item assignment. To assign to this variable, you must first load it into memory explicitly using the .load() method or accessing its .values attribute.

This actually can speed up the code pretty substantially.

# With parallelization / specifying chunks [chunks={"time": 20}]
$ time python notImplementedError.py                                                                                      
...
real    0m1.037s
user    0m0.824s
sys     0m0.119s

# Without parallelization
$ time python notImplementedError.py 

real    0m5.090s
user    0m4.040s
sys     0m0.891s

[tomvothecoder]

Note that I get a different error than the one specified in this issue. I got this error:
TypeError: this variable's data is stored in a dask array, which does not support item assignment. To assign to this variable, you must first load it into memory explicitly using the .load() method or accessing its .values attribute.

I get NotImplementedError: xarray can't set arrays with multiple array indices to dask yet. when using the xcdat_test env (https://github.com/XCDAT/xcdat_test/blob/master/conda-env/test.yml), which includes xcdat=0.1.0, xarray=0.19.0, and dask=2021.9.1.

What versions are you running in your env? I am trying to reproduce your error.

@tomvothecoder - here is a proposed fix. I wonder if there are going to be a lot of little places where we need to add a .load().

Nice to see the speed gains! Have you tried using the .values attribute instead of .load()? I wonder if the speed gains are comparable using .values since it isn't loading the data into memory. Using .values will hopefully reduce the memory footprint if it is working on "the array's data as numpy.ndarray" directly.
https://xarray.pydata.org/en/stable/generated/xarray.DataArray.values.html

Also, I don't think we need to have .load()/.values unless we're attempting to work on the Dask arrays directly (e.g., scaling domain down to region using numpy instead of xarray). The TypeError will definitely let us know where we need to place those lines of code.

[pochedls]

I just re-created my environment and can reproduce the error message that you got. I'm not sure if I was using an old environment derived from dev.xml or if I had tweaked the dev environment for some reason. This commit does remove the NotImplementedError (even using the appropriate environment). Perhaps there is another way, though.

I did use .values, but this only works for dataarray objects (and we sometimes pass in a ndarray). I just updated the commit to remove .load() here. The speed gains are comparable (1.113s real).

[tomvothecoder]

The dev.yml was updated before the release of v0.1.0, so you may have been on an older version of xarray and/or dask. I forgot to bring it up at the meeting in case folks had an older env set up, my apologies.

I'm glad to hear that your fix still holds up with the new dependencies and with .values. You can open up a new PR whenever you're ready and I'll review it. Thanks!

[tomvothecoder]

Also, I'll compare .load() against .values() to see if there are any major differences in memory usage.

[pochedls]

@tomvothecoder - I'm working on the PR and adding some tests. I wanted to add a test that checks that spatial averaging works for chunked data, but the fixtures dataset cannot be chunked (TypeError: 'Frozen' object is not callable). Could I find a dataset online (e.g., via THREDDS) or could we include a small netCDF file somewhere for testing?

[tomvothecoder]

@pochedls How did you produce that error? generate_dataset() returns a normal, chunkable Dataset object.

This should work:

def setup(self):
    self.ds = generate_dataset(cf_compliant=True, has_bounds=True)

def some_test(self):
   ds = self.ds.copy()

   # Perform chunking
   ds = ds.chunk(5)

You can also set up your own Dataset object if the fixture isn't sufficient enough for the test(s). I often do this for specific tests cases.