Repository: corteva/rioxarray

Issue:
band_as_variable ignores VRT warping
```python
import rasterio as rio
from rasterio import crs
from rasterio.enums import Resampling
from rasterio.vrt import WarpedVRT
import rioxarray
import xarray as xr

resampling = 'average'
warp_args : dict = {
    'resampling': Resampling._member_map_[resampling],
    'crs': crs.CRS.from_epsg(4326), # Input raster must be converted to WGS84 (4326) for H3 indexing
    'warp_mem_limit': 12000
}
with rio.open(path_to_raster) as src:
    with WarpedVRT(src, src_crs=src.crs, **warp_args) as vrt:
        da : xr.Dataset = rioxarray.open_rasterio(
            vrt,
            band_as_variable=True # <---- !!
        ).chunk(**{'x':'auto','y':'auto'})
        logging.info(da['x'])
```

```
array([1831067.831059, 1831077.828418, 1831087.825776, ..., 1856771.040311,
       1856781.03767 , 1856791.035028])
```

In contast, the same code with `band_as_variable=False` gives:

```
array([175.666845, 175.666949, 175.667054, ..., 175.972549, 175.972654,
       175.972758])
```

The latter has correctly respected the WarpedVRT transformation to EPSG:4326, but the former has retained coordinates in my input projection and seems to ignore the VRT CRS. This is a problem because I need to transform to 4326 prior to performing an H3 index on cell data for later system integration.

#### Environment Information

```
python = "^3.10"
gdal = "^3.6.2"
geopandas = "^0.12.2"
h3pandas = "^0.2.3"
rioxarray = "^0.13.4"
dask-geopandas = "^0.3.0"
pyarrow = "^11.0.0"
dask = "^2023.3.0"
click = "^8.1.3"
```

Linux-5.15.0-41-generic-x86_64-with-glibc2.35

Installed with poetry (pyenv).


Discussion:

Patch:
diff --git a/docs/history.rst b/docs/history.rst
index 1034c6f3..64392041 100644
--- a/docs/history.rst
+++ b/docs/history.rst
@@ -5,6 +5,7 @@ Latest
 ------
 - DEP: Drop Python 3.8 support (issue #582)
 - DEP: pin rasterio>=1.2 (pull #642)
+- BUG: Fix WarpedVRT in :func:`rioxarray.open_rasterio` when band_as_variable=True (issue #644)
 
 0.13.4
 ------
diff --git a/rioxarray/_io.py b/rioxarray/_io.py
index 0aa0c282..b9abe85f 100644
--- a/rioxarray/_io.py
+++ b/rioxarray/_io.py
@@ -5,7 +5,7 @@ Credits:
 This file was adopted from: https://github.com/pydata/xarray # noqa
 Source file: https://github.com/pydata/xarray/blob/1d7bcbdc75b6d556c04e2c7d7a042e4379e15303/xarray/backends/rasterio_.py # noqa
 """
-
+# pylint: disable=too-many-lines
 import contextlib
 import functools
 import importlib.metadata
@@ -42,18 +42,47 @@ RASTERIO_LOCK = SerializableLock()
 NO_LOCK = contextlib.nullcontext()
 
 
+def _ensure_warped_vrt(riods, vrt_params):
+    """
+    Ensuire the dataset is represented as a warped vrt
+    """
+    if vrt_params is None:
+        return riods
+    if isinstance(riods, SingleBandDatasetReader):
+        riods._create_vrt(vrt_params)
+    else:
+        riods = WarpedVRT(riods, **vrt_params)
+    return riods
+
+
 class SingleBandDatasetReader:
     """
     Hack to have a DatasetReader behave like it only has one band
     """
 
-    def __init__(self, riods, bidx) -> None:
+    def __init__(self, riods, bidx, vrt_params=None) -> None:
         self._riods = riods
         self._bidx = bidx
+        self._vrt_params = vrt_params
+        self._create_vrt(vrt_params=vrt_params)
 
     def __getattr__(self, __name: str) -> Any:
         return getattr(self._riods, __name)
 
+    def _create_vrt(self, vrt_params):
+        if vrt_params is not None and not isinstance(self._riods, WarpedVRT):
+            self._riods = WarpedVRT(self._riods, **vrt_params)
+        self._vrt_params = vrt_params
+
+    @property
+    def name(self):
+        """
+        str: name of the dataset. Usually the path.
+        """
+        if isinstance(self._riods, rasterio.vrt.WarpedVRT):
+            return self._riods.src_dataset.name
+        return self._riods.name
+
     @property
     def count(self):
         """
@@ -275,9 +304,7 @@ class RasterioArrayWrapper(BackendArray):
         self.mask_and_scale = mask_and_scale
 
         # cannot save riods as an attribute: this would break pickleability
-        riods = manager.acquire()
-        if vrt_params is not None:
-            riods = WarpedVRT(riods, **vrt_params)
+        riods = _ensure_warped_vrt(manager.acquire(), vrt_params)
         self.vrt_params = vrt_params
         self._shape = (riods.count, riods.height, riods.width)
 
@@ -395,9 +422,9 @@ class RasterioArrayWrapper(BackendArray):
             out = np.zeros(shape, dtype=self.dtype)
         else:
             with self.lock:
-                riods = self.manager.acquire(needs_lock=False)
-                if self.vrt_params is not None:
-                    riods = WarpedVRT(riods, **self.vrt_params)
+                riods = _ensure_warped_vrt(
+                    self.manager.acquire(needs_lock=False), self.vrt_params
+                )
                 out = riods.read(band_key, window=window, masked=self.masked)
                 if self._unsigned_dtype is not None:
                     out = out.astype(self._unsigned_dtype)
@@ -818,6 +845,7 @@ def _load_bands_as_variables(
     mask_and_scale: bool,
     decode_times: bool,
     decode_timedelta: Optional[bool],
+    vrt_params: Optional[dict],
     **open_kwargs,
 ) -> Union[Dataset, list[Dataset]]:
     """
@@ -829,6 +857,7 @@ def _load_bands_as_variables(
         band_riods = SingleBandDatasetReader(
             riods=riods,
             bidx=band - 1,
+            vrt_params=vrt_params,
         )
         band_name = f"band_{band}"
         data_vars[band_name] = (
@@ -1049,6 +1078,7 @@ def open_rasterio(
             _single_band_open,
             bidx=filename._bidx,
         )
+        vrt_params = filename._vrt_params
     if isinstance(filename, (rasterio.io.DatasetReader, SingleBandDatasetReader)):
         filename = filename.name
     elif isinstance(filename, rasterio.vrt.WarpedVRT):
@@ -1088,22 +1118,6 @@ def open_rasterio(
         riods = manager.acquire()
         captured_warnings = rio_warnings.copy()
 
-    if band_as_variable:
-        dataset_result = _load_bands_as_variables(
-            riods=riods,
-            parse_coordinates=parse_coordinates,
-            chunks=chunks,
-            cache=cache,
-            lock=lock,
-            masked=masked,
-            mask_and_scale=mask_and_scale,
-            decode_times=decode_times,
-            decode_timedelta=decode_timedelta,
-            **open_kwargs,
-        )
-        manager.close()
-        return dataset_result
-
     # raise the NotGeoreferencedWarning if applicable
     for rio_warning in captured_warnings:
         if not riods.subdatasets or not isinstance(
@@ -1130,12 +1144,28 @@ def open_rasterio(
         manager.close()
         return subdataset_result
 
-    if vrt_params is not None:
-        riods = WarpedVRT(riods, **vrt_params)
+    if band_as_variable:
+        dataset_result = _load_bands_as_variables(
+            riods=riods,
+            parse_coordinates=parse_coordinates,
+            chunks=chunks,
+            cache=cache,
+            lock=lock,
+            masked=masked,
+            mask_and_scale=mask_and_scale,
+            decode_times=decode_times,
+            decode_timedelta=decode_timedelta,
+            vrt_params=vrt_params,
+            **open_kwargs,
+        )
+        manager.close()
+        return dataset_result
 
     if cache is None:
         cache = chunks is None
 
+    riods = _ensure_warped_vrt(riods, vrt_params)
+
     # Get bands
     if riods.count < 1:
         raise ValueError("Unknown dims")


Test patch:
diff --git a/test/integration/test_integration__io.py b/test/integration/test_integration__io.py
index 83827ec0..72f5f95f 100644
--- a/test/integration/test_integration__io.py
+++ b/test/integration/test_integration__io.py
@@ -600,7 +600,6 @@ def test_indexing():
         8, 10, 3, transform_args=[1, 2, 0.5, 2.0], crs="EPSG:4326"
     ) as (tmp_file, expected):
         with rioxarray.open_rasterio(tmp_file, cache=False) as actual:
-
             # tests
             # assert_allclose checks all data + coordinates
             assert_allclose(actual, expected)
@@ -717,7 +716,6 @@ def test_caching():
     ) as (tmp_file, expected):
         # Cache is the default
         with rioxarray.open_rasterio(tmp_file) as actual:
-
             # This should cache everything
             assert_allclose(actual, expected)
 
@@ -866,7 +864,8 @@ def test_rasterio_environment():
                     assert_allclose(actual, expected)
 
 
-def test_rasterio_vrt():
+@pytest.mark.parametrize("band_as_variable", [True, False])
+def test_rasterio_vrt(band_as_variable):
     # tmp_file default crs is UTM: CRS({'init': 'epsg:32618'}
     with create_tmp_geotiff() as (tmp_file, expected):
         with rasterio.open(tmp_file) as src:
@@ -874,7 +873,12 @@ def test_rasterio_vrt():
                 # Value of single pixel in center of image
                 lon, lat = vrt.xy(vrt.width // 2, vrt.height // 2)
                 expected_val = next(vrt.sample([(lon, lat)]))
-                with rioxarray.open_rasterio(vrt) as rds:
+                with rioxarray.open_rasterio(
+                    vrt, band_as_variable=band_as_variable
+                ) as rds:
+                    if band_as_variable:
+                        rds = rds.band_1
+
                     actual_val = rds.sel(dict(x=lon, y=lat), method="nearest").data
 
                     assert_array_equal(rds.rio.shape, (vrt.height, vrt.width))
@@ -1413,7 +1417,6 @@ def test_cint16_dtype_masked(dtype, tmp_path):
 def test_cint16_promote_dtype(tmp_path):
     test_file = os.path.join(TEST_INPUT_DATA_DIR, "cint16.tif")
     with rioxarray.open_rasterio(test_file) as xds:
-
         tmp_output = tmp_path / "tmp_cfloat64.tif"
         with pytest.warns(NotGeoreferencedWarning):
             xds.rio.to_raster(str(tmp_output), dtype="complex64")