Repository: corteva/rioxarray

Issue:
IndexError on compute() when using mask_and_scale=True and dask chunking along band dimension
#### Code Sample, a copy-pastable example if possible
[tmp.tif.gz](https://github.com/corteva/rioxarray/files/6920344/tmp.tif.gz)
```python

import xarray as xr
import rioxarray as rxr

rs = rxr.open_rasterio("tmp.tif", mask_and_scale=True).chunk((1, 100, 100))
print(rs)
print(rs.compute())

```
#### Problem description

I get an `IndexError` when I try to do anything that causes a compute operation. This only happens if I use `mask_and_scale=True` when I open the tiff and then chunk along the band dimension. If I chunk with (N, H, W) and have N equal to the number of bands i.e. (4, 100, 100), it works fine. Chunking along the other dimensions doesn't cause an error. Only the band dimension causes problems.

#### Expected Output
```text
<xarray.DataArray (band: 4, y: 100, x: 100)>
dask.array<xarray-<this-array>, shape=(4, 100, 100), dtype=float32, chunksize=(1, 100, 100), chunktype=numpy.ndarray>
Coordinates:
  * band         (band) int64 1 2 3 4
  * x            (x) float64 7.819e+05 7.819e+05 7.819e+05 ... 7.82e+05 7.82e+05
  * y            (y) float64 3.407e+06 3.407e+06 ... 3.407e+06 3.407e+06
    spatial_ref  int64 0
<xarray.DataArray (band: 4, y: 100, x: 100)>
array([[[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       ...

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]]], dtype=float32)
Coordinates:
  * band         (band) int64 1 2 3 4
  * x            (x) float64 7.819e+05 7.819e+05 7.819e+05 ... 7.82e+05 7.82e+05
  * y            (y) float64 3.407e+06 3.407e+06 ... 3.407e+06 3.407e+06
    spatial_ref  int64 0
```

#### Stack Trace
<details>

```text
Traceback (most recent call last):
  File "my.py", line 6, in <module>
    print(rs.compute())
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/dataarray.py", line 955, in compute
    return new.load(**kwargs)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/dataarray.py", line 929, in load
    ds = self._to_temp_dataset().load(**kwargs)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/dataset.py", line 865, in load
    evaluated_data = da.compute(*lazy_data.values(), **kwargs)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/dask/base.py", line 568, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/dask/threaded.py", line 79, in get
    results = get_async(
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/dask/local.py", line 514, in get_async
    raise_exception(exc, tb)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/dask/local.py", line 325, in reraise
    raise exc
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/dask/local.py", line 223, in execute_task
    result = _execute_task(task, data)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/dask/core.py", line 121, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/dask/array/core.py", line 104, in getter
    c = np.asarray(c)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/indexing.py", line 354, in __array__
    return np.asarray(self.array, dtype=dtype)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/indexing.py", line 548, in __array__
    self._ensure_cached()
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/indexing.py", line 545, in _ensure_cached
    self.array = NumpyIndexingAdapter(np.asarray(self.array))
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/indexing.py", line 518, in __array__
    return np.asarray(self.array, dtype=dtype)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/indexing.py", line 419, in __array__
    return np.asarray(array[self.key], dtype=None)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/rioxarray/_io.py", line 305, in __getitem__
    return indexing.explicit_indexing_adapter(
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/xarray/core/indexing.py", line 710, in explicit_indexing_adapter
    result = raw_indexing_method(raw_key.tuple)
  File "/home/$USER/anaconda3/envs/rstools/lib/python3.8/site-packages/rioxarray/_io.py", line 296, in _getitem
    out[band_iii] * riods.scales[band_iii]
IndexError: index 1 is out of bounds for axis 0 with size 1
```
</details>

#### Environment Information
```text
rioxarray (0.6.0) deps:
  rasterio: 1.2.6
    xarray: 0.19.0
      GDAL: 3.3.1

Other python deps:
     scipy: 1.7.0
    pyproj: 3.1.0

System:
    python: 3.8.10 | packaged by conda-forge | (default, May 11 2021, 07:01:05)  [GCC 9.3.0]
executable: /home/$USER/anaconda3/envs/rstools/bin/python
   machine: Linux-5.11.0-25-generic-x86_64-with-glibc2.10
```

#### Installation method
 - conda

#### Conda environment information (if you installed with conda):

<br/>
Environment (<code>conda list</code>):
<details>

```
$ conda list | grep -E "rasterio|xarray|gdal|dask"
dask                      2021.7.2           pyhd8ed1ab_0    conda-forge
dask-core                 2021.7.2           pyhd8ed1ab_0    conda-forge
dask-image                0.6.0              pyhd8ed1ab_0    conda-forge
libgdal                   3.3.1                h8f005ca_1    conda-forge
rasterio                  1.2.6            py38h84fce68_2    conda-forge
rioxarray                 0.6.0              pyhd8ed1ab_0    conda-forge
xarray                    0.19.0             pyhd8ed1ab_1    conda-forge
```
</details>

<br/>
Details about  <code>conda</code> and system ( <code>conda info</code> ):
<details>

```
$ conda info

     active environment : rstools
    active env location : /home/$USER/anaconda3/envs/rstools
            shell level : 3
       user config file : /home/$USER/.condarc
 populated config files : /home/$USER/.condarc
          conda version : 4.10.3
    conda-build version : not installed
         python version : 3.8.10.final.0
       virtual packages : __cuda=11.2=0
                          __linux=5.11.0=0
                          __glibc=2.33=0
                          __unix=0=0
                          __archspec=1=x86_64
       base environment : /home/$USER/anaconda3  (writable)
      conda av data dir : /home/$USER/anaconda3/etc/conda
  conda av metadata url : None
           channel URLs : https://conda.anaconda.org/conda-forge/linux-64
                          https://conda.anaconda.org/conda-forge/noarch
                          https://repo.anaconda.com/pkgs/main/linux-64
                          https://repo.anaconda.com/pkgs/main/noarch
                          https://repo.anaconda.com/pkgs/r/linux-64
                          https://repo.anaconda.com/pkgs/r/noarch
          package cache : /home/$USER/anaconda3/pkgs
                          /home/$USER/.conda/pkgs
       envs directories : /home/$USER/anaconda3/envs
                          /home/$USER/.conda/envs
               platform : linux-64
             user-agent : conda/4.10.3 requests/2.25.1 CPython/3.8.10 Linux/5.11.0-25-generic ubuntu/21.04 glibc/2.33
                UID:GID : 1000:1000
             netrc file : None
           offline mode : False
```
</details>

Comments:

Patch:
diff --git a/docs/history.rst b/docs/history.rst
index dfe640f..732d0f9 100644
--- a/docs/history.rst
+++ b/docs/history.rst
@@ -3,6 +3,7 @@ History
 
 Latest
 ------
+- BUG: Fix indexing error when `mask_and_scale=True` was combined with band dim chunking (issue #387, pull #388)
 
 0.6.0
 ------
diff --git a/rioxarray/_io.py b/rioxarray/_io.py
index 44c378d..e376254 100644
--- a/rioxarray/_io.py
+++ b/rioxarray/_io.py
@@ -288,11 +288,9 @@ class RasterioArrayWrapper(BackendArray):
                 if self.masked:
                     out = np.ma.filled(out.astype(self.dtype), self.fill_value)
                 if self.mask_and_scale:
-                    for band in np.atleast_1d(band_key):
-                        band_iii = band - 1
-                        out[band_iii] = (
-                            out[band_iii] * riods.scales[band_iii]
-                            + riods.offsets[band_iii]
+                    for iii, band_iii in enumerate(np.atleast_1d(band_key) - 1):
+                        out[iii] = (
+                            out[iii] * riods.scales[band_iii] + riods.offsets[band_iii]
                         )
 
         if squeeze_axis:


Test patch:
diff --git a/test/integration/test_integration__io.py b/test/integration/test_integration__io.py
index 0156f51..2b4da10 100644
--- a/test/integration/test_integration__io.py
+++ b/test/integration/test_integration__io.py
@@ -682,6 +682,23 @@ def test_chunks():
             assert_allclose(ac, ex)
 
 
+def test_chunks_with_mask_and_scale():
+    with create_tmp_geotiff(
+        10, 10, 4, transform_args=[1, 2, 0.5, 2.0], crs="+proj=latlong"
+    ) as (tmp_file, expected):
+        # Chunk at open time
+        with rioxarray.open_rasterio(
+            tmp_file, mask_and_scale=True, chunks=(1, 2, 2)
+        ) as actual:
+            assert isinstance(actual.data, dask.array.Array)
+            assert "open_rasterio" in actual.data.name
+
+            # do some arithmetic
+            ac = actual.mean().compute()
+            ex = expected.mean()
+            assert_allclose(ac, ex)
+
+
 def test_pickle_rasterio():
     # regression test for https://github.com/pydata/xarray/issues/2121
     with create_tmp_geotiff() as (tmp_file, expected):