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>

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/corteva/rioxarray/pull/388?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva) Report
> Merging [#388](https://codecov.io/gh/corteva/rioxarray/pull/388?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva) (a927e51) into [master](https://codecov.io/gh/corteva/rioxarray/commit/008c2a2c6f0147bd5e075fe6984eb8b636bd0bfc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva) (008c2a2) will **decrease** coverage by `0.20%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/corteva/rioxarray/pull/388/graphs/tree.svg?width=650&height=150&src=pr&token=C8QK7VO1OF&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva)](https://codecov.io/gh/corteva/rioxarray/pull/388?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva)

```diff
@@            Coverage Diff             @@
##           master     #388      +/-   ##
==========================================
- Coverage   95.86%   95.65%   -0.21%     
==========================================
  Files          13       13              
  Lines        1474     1473       -1     
==========================================
- Hits         1413     1409       -4     
- Misses         61       64       +3     
```


| [Impacted Files](https://codecov.io/gh/corteva/rioxarray/pull/388?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva) | Coverage Δ | |
|---|---|---|
| [rioxarray/\_io.py](https://codecov.io/gh/corteva/rioxarray/pull/388/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva#diff-cmlveGFycmF5L19pby5weQ==) | `95.12% <100.00%> (-0.02%)` | :arrow_down: |
| [rioxarray/raster\_array.py](https://codecov.io/gh/corteva/rioxarray/pull/388/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva#diff-cmlveGFycmF5L3Jhc3Rlcl9hcnJheS5weQ==) | `97.29% <0.00%> (-1.02%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/corteva/rioxarray/pull/388?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/corteva/rioxarray/pull/388?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva). Last update [008c2a2...a927e51](https://codecov.io/gh/corteva/rioxarray/pull/388?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=corteva).

snowman2: The latest test failure can be ignored. This looks good, thanks @fbunt :+1:. Any issues if I squash and merge?
fbunt: Great! Nope, squash away

Patch:
diff --git a/docs/history.rst b/docs/history.rst
index dfe640f7..732d0f9e 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 44c378d3..e376254f 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 0156f510..2b4da10f 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):