From fe35810874ed2920f6f344cee4c158dc11d31988 Mon Sep 17 00:00:00 2001 From: "Alan D. Snow" Date: Tue, 13 Apr 2021 15:05:38 -0500 Subject: [PATCH] TST: Add more tests for rasterio engine (#293) --- rioxarray/xarray_plugin.py | 12 +- test/conftest.py | 25 + test/integration/test_integration__io.py | 900 +++++++++--------- .../integration/test_integration_rioxarray.py | 140 ++- .../test_integration_xarray_plugin.py | 14 + 5 files changed, 597 insertions(+), 494 deletions(-) diff --git a/rioxarray/xarray_plugin.py b/rioxarray/xarray_plugin.py index c7787fa3..082c0a57 100644 --- a/rioxarray/xarray_plugin.py +++ b/rioxarray/xarray_plugin.py @@ -3,6 +3,7 @@ import xarray as xr from . import _io +from .exceptions import RioXarrayError CAN_OPEN_EXTS = { "asc", @@ -28,11 +29,13 @@ class RasterioBackend(xr.backends.common.BackendEntrypoint): def open_dataset( self, filename_or_obj, - drop_variables=None, - mask_and_scale=True, + drop_variables=None, # SKIP FROM XARRAY parse_coordinates=None, + chunks=None, + cache=None, lock=None, masked=False, + mask_and_scale=True, variable=None, group=None, default_name="band_data", @@ -53,6 +56,11 @@ def open_dataset( ) if isinstance(ds, xr.DataArray): ds = ds.to_dataset() + if not isinstance(ds, xr.Dataset): + raise RioXarrayError( + "Multiple resolution sets found. " + "Use 'variable' or 'group' to filter." + ) return ds def guess_can_open(self, filename_or_obj): diff --git a/test/conftest.py b/test/conftest.py index 3d48568f..86aeb227 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,12 +1,19 @@ import os +from distutils.version import LooseVersion +import pyproj +import pytest +import rasterio from numpy.testing import assert_almost_equal, assert_array_equal +import rioxarray from rioxarray.raster_array import UNWANTED_RIO_ATTRS TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), "test_data") TEST_INPUT_DATA_DIR = os.path.join(TEST_DATA_DIR, "input") TEST_COMPARE_DATA_DIR = os.path.join(TEST_DATA_DIR, "compare") +PYPROJ_LT_3 = LooseVersion(pyproj.__version__) < LooseVersion("3") +RASTERIO_LT_122 = LooseVersion(rasterio.__version__) < LooseVersion("1.2.2") # xarray.testing.assert_equal(input_xarray, compare_xarray) @@ -24,6 +31,8 @@ def _assert_attrs_equal(input_xr, compare_xr, decimal_precision): and attr not in UNWANTED_RIO_ATTRS and attr != "creation_date" and attr != "grid_mapping" + and attr != "coordinates" + and "#" not in attr ): try: assert_almost_equal( @@ -84,3 +93,19 @@ def _assert_xarrays_equal( assert input_xarray.rio.grid_mapping == compare_xarray.rio.grid_mapping for unwanted_attr in UNWANTED_RIO_ATTRS: assert unwanted_attr not in input_xarray.attrs + + +def open_rasterio_engine(file_name_or_object, **kwargs): + # FIXME: change to the next xarray version after release + xr = pytest.importorskip("xarray", minversion="0.17.1.dev0") + return xr.open_dataset(file_name_or_object, engine="rasterio", **kwargs) + + +@pytest.fixture( + params=[ + rioxarray.open_rasterio, + open_rasterio_engine, + ] +) +def open_rasterio(request): + return request.param diff --git a/test/integration/test_integration__io.py b/test/integration/test_integration__io.py index 84e0eed8..f08b205e 100644 --- a/test/integration/test_integration__io.py +++ b/test/integration/test_integration__io.py @@ -15,7 +15,7 @@ import rasterio import xarray as xr from affine import Affine -from numpy.testing import assert_almost_equal +from numpy.testing import assert_almost_equal, assert_array_equal from rasterio.errors import NotGeoreferencedWarning from rasterio.transform import from_origin from rasterio.warp import calculate_default_transform @@ -25,6 +25,7 @@ import rioxarray from rioxarray._io import build_subdataset_filter +from rioxarray.rioxarray import DEFAULT_GRID_MAP from test.conftest import ( TEST_COMPARE_DATA_DIR, TEST_INPUT_DATA_DIR, @@ -178,15 +179,15 @@ def test_build_subdataset_filter(subdataset, variable, group, match): ) == match -def test_open_variable_filter(): - with rioxarray.open_rasterio( +def test_open_variable_filter(open_rasterio): + with open_rasterio( os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc"), variable=["blue"] ) as rds: assert list(rds.data_vars) == ["blue"] -def test_open_group_filter__missing(): - with rioxarray.open_rasterio( +def test_open_group_filter__missing(open_rasterio): + with open_rasterio( os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc"), variable="blue", group=["non-existent"], @@ -208,8 +209,8 @@ def test_open_multiple_resolution(): assert rds_list[1].dims == {"y": 2400, "x": 2400, "band": 1} -def test_open_group_filter(): - with rioxarray.open_rasterio( +def test_open_group_filter(open_rasterio): + with open_rasterio( os.path.join( TEST_INPUT_DATA_DIR, "MOD09GA.A2008296.h14v17.006.2015181011753.hdf" ), @@ -230,11 +231,12 @@ def test_open_group_filter(): ] -def test_open_group_load_attrs(): - with rioxarray.open_rasterio( +def test_open_group_load_attrs(open_rasterio): + with open_rasterio( os.path.join( TEST_INPUT_DATA_DIR, "MOD09GA.A2008296.h14v17.006.2015181011753.hdf" ), + mask_and_scale=False, variable="sur_refl_b05_1", ) as rds: attrs = rds["sur_refl_b05_1"].attrs @@ -258,6 +260,8 @@ def test_open_rasterio_mask_chunk_clip(): chunks=True, default_name="dem", ) as xdi: + if isinstance(xdi, xr.Dataset): + xdi = xdi.dem assert xdi.name == "dem" assert str(xdi.dtype) == "float64" assert str(xdi.data.dtype) == "float64" @@ -385,6 +389,7 @@ def create_tmp_geotiff( s.write(data, **write_kwargs) dx, dy = s.res[0], -s.res[1] tt = s.transform + crs_wkt = s.crs.to_wkt() if s.crs else None if not transform_args: a, b, c, d = tt.c, tt.f, -tt.e, tt.a @@ -400,448 +405,455 @@ def create_tmp_geotiff( "x": np.arange(nx) * c + a + dx / 2, }, ) + if crs_wkt is not None: + expected.coords[DEFAULT_GRID_MAP] = xr.Variable((), 0) + expected.coords[DEFAULT_GRID_MAP].attrs["spatial_ref"] = crs_wkt yield tmp_file, expected -class TestRasterio: - def test_serialization(self): - with create_tmp_geotiff(additional_attrs={}) as (tmp_file, expected): - # Write it to a netcdf and read again (roundtrip) - with xr.open_rasterio(tmp_file) as rioda: - with create_tmp_file(suffix=".nc") as tmp_nc_file: - rioda.to_netcdf(tmp_nc_file) - with xr.open_dataarray(tmp_nc_file) as ncds: - assert_identical(rioda, ncds) - - def test_utm(self): - with create_tmp_geotiff() as (tmp_file, expected): - with xr.open_rasterio(tmp_file) as rioda: - assert_allclose(rioda, expected) - assert rioda.attrs["scales"] == (1.0, 1.0, 1.0) - assert rioda.attrs["offsets"] == (0.0, 0.0, 0.0) - assert rioda.attrs["descriptions"] == ("d1", "d2", "d3") - assert rioda.attrs["units"] == ("u1", "u2", "u3") - assert isinstance(rioda.attrs["crs"], str) - assert isinstance(rioda.attrs["res"], tuple) - assert isinstance(rioda.attrs["is_tiled"], np.uint8) - assert isinstance(rioda.rio._cached_transform(), Affine) - np.testing.assert_array_equal( - rioda.attrs["nodatavals"], [np.NaN, np.NaN, np.NaN] - ) - - # Check no parse coords - with xr.open_rasterio(tmp_file, parse_coordinates=False) as rioda: - assert "x" not in rioda.coords - assert "y" not in rioda.coords - - def test_platecarree(self): - with create_tmp_geotiff( - 8, - 10, - 1, - transform_args=[1, 2, 0.5, 2.0], - crs="+proj=latlong", - open_kwargs={"nodata": -9765}, - ) as (tmp_file, expected): - with xr.open_rasterio(tmp_file) as rioda: - assert_allclose(rioda, expected) - assert rioda.attrs["scales"] == (1.0,) - assert rioda.attrs["offsets"] == (0.0,) - assert isinstance(rioda.attrs["descriptions"], tuple) - assert isinstance(rioda.attrs["units"], tuple) - assert isinstance(rioda.attrs["crs"], str) - assert isinstance(rioda.attrs["res"], tuple) - assert isinstance(rioda.attrs["is_tiled"], np.uint8) - assert isinstance(rioda.rio._cached_transform(), Affine) - np.testing.assert_array_equal(rioda.attrs["nodatavals"], [-9765.0]) - - def test_notransform(self): - # regression test for https://github.com/pydata/xarray/issues/1686 - # Create a geotiff file - with warnings.catch_warnings(): - # rasterio throws a NotGeoreferencedWarning here, which is - # expected since we test rasterio's defaults in this case. - warnings.filterwarnings( - "ignore", - category=UserWarning, - message="Dataset has no geotransform set", - ) - with create_tmp_file(suffix=".tif") as tmp_file: - # data - nx, ny, nz = 4, 3, 3 - data = np.arange(nx * ny * nz, dtype=rasterio.float32).reshape( - nz, ny, nx - ) - with rasterio.open( - tmp_file, - "w", - driver="GTiff", - height=ny, - width=nx, - count=nz, - dtype=rasterio.float32, - ) as s: - s.descriptions = ("nx", "ny", "nz") - s.units = ("cm", "m", "km") - s.write(data) - - # Tests - expected = DataArray( - data, - dims=("band", "y", "x"), - coords={ - "band": [1, 2, 3], - "y": [0.5, 1.5, 2.5], - "x": [0.5, 1.5, 2.5, 3.5], - }, - ) - with xr.open_rasterio(tmp_file) as rioda: - assert_allclose(rioda, expected) - assert rioda.attrs["scales"] == (1.0, 1.0, 1.0) - assert rioda.attrs["offsets"] == (0.0, 0.0, 0.0) - assert rioda.attrs["descriptions"] == ("nx", "ny", "nz") - assert rioda.attrs["units"] == ("cm", "m", "km") - assert isinstance(rioda.attrs["res"], tuple) - assert isinstance(rioda.attrs["is_tiled"], np.uint8) - assert isinstance(rioda.rio._cached_transform(), Affine) - - def test_indexing(self): - with create_tmp_geotiff( - 8, 10, 3, transform_args=[1, 2, 0.5, 2.0], crs="+proj=latlong" - ) as (tmp_file, expected): - with xr.open_rasterio(tmp_file, cache=False) as actual: - - # tests - # assert_allclose checks all data + coordinates - assert_allclose(actual, expected) - assert not actual.variable._in_memory - - # Basic indexer - ind = {"x": slice(2, 5), "y": slice(5, 7)} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - ind = {"band": slice(1, 2), "x": slice(2, 5), "y": slice(5, 7)} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - ind = {"band": slice(1, 2), "x": slice(2, 5), "y": 0} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - # orthogonal indexer - ind = { - "band": np.array([2, 1, 0]), - "x": np.array([1, 0]), - "y": np.array([0, 2]), - } - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - ind = {"band": np.array([2, 1, 0]), "x": np.array([1, 0]), "y": 0} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - ind = {"band": 0, "x": np.array([0, 0]), "y": np.array([1, 1, 1])} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - # minus-stepped slice - ind = {"band": np.array([2, 1, 0]), "x": slice(-1, None, -1), "y": 0} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - ind = {"band": np.array([2, 1, 0]), "x": 1, "y": slice(-1, 1, -2)} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - # empty selection - ind = {"band": np.array([2, 1, 0]), "x": 1, "y": slice(2, 2, 1)} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - ind = {"band": slice(0, 0), "x": 1, "y": 2} - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - # vectorized indexer - ind = { - "band": DataArray([2, 1, 0], dims="a"), - "x": DataArray([1, 0, 0], dims="a"), - "y": np.array([0, 2]), - } - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - ind = { - "band": DataArray([[2, 1, 0], [1, 0, 2]], dims=["a", "b"]), - "x": DataArray([[1, 0, 0], [0, 1, 0]], dims=["a", "b"]), - "y": 0, - } - assert_allclose(expected.isel(**ind), actual.isel(**ind)) - assert not actual.variable._in_memory - - # Selecting lists of bands is fine - ex = expected.isel(band=[1, 2]) - ac = actual.isel(band=[1, 2]) - assert_allclose(ac, ex) - ex = expected.isel(band=[0, 2]) - ac = actual.isel(band=[0, 2]) - assert_allclose(ac, ex) - - # Integer indexing - ex = expected.isel(band=1) - ac = actual.isel(band=1) - assert_allclose(ac, ex) - - ex = expected.isel(x=1, y=2) - ac = actual.isel(x=1, y=2) - assert_allclose(ac, ex) - - ex = expected.isel(band=0, x=1, y=2) - ac = actual.isel(band=0, x=1, y=2) - assert_allclose(ac, ex) - - # Mixed - ex = actual.isel(x=slice(2), y=slice(2)) - ac = actual.isel(x=[0, 1], y=[0, 1]) - assert_allclose(ac, ex) - - ex = expected.isel(band=0, x=1, y=slice(5, 7)) - ac = actual.isel(band=0, x=1, y=slice(5, 7)) - assert_allclose(ac, ex) - - ex = expected.isel(band=0, x=slice(2, 5), y=2) - ac = actual.isel(band=0, x=slice(2, 5), y=2) - assert_allclose(ac, ex) - - # One-element lists - ex = expected.isel(band=[0], x=slice(2, 5), y=[2]) - ac = actual.isel(band=[0], x=slice(2, 5), y=[2]) - assert_allclose(ac, ex) - - def test_caching(self): - with create_tmp_geotiff( - 8, 10, 3, transform_args=[1, 2, 0.5, 2.0], crs="+proj=latlong" - ) as (tmp_file, expected): - # Cache is the default - with xr.open_rasterio(tmp_file) as actual: - - # This should cache everything - assert_allclose(actual, expected) - - # once cached, non-windowed indexing should become possible - ac = actual.isel(x=[2, 4]) - ex = expected.isel(x=[2, 4]) - assert_allclose(ac, ex) - - def test_chunks(self): - with create_tmp_geotiff( - 8, 10, 3, transform_args=[1, 2, 0.5, 2.0], crs="+proj=latlong" - ) as (tmp_file, expected): - # Chunk at open time - with xr.open_rasterio(tmp_file, 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() - ex = expected.mean() - assert_allclose(ac, ex) - - ac = actual.sel(band=1).mean(dim="x") - ex = expected.sel(band=1).mean(dim="x") - assert_allclose(ac, ex) - - def test_pickle_rasterio(self): - # regression test for https://github.com/pydata/xarray/issues/2121 - with create_tmp_geotiff() as (tmp_file, expected): - with xr.open_rasterio(tmp_file) as rioda: - temp = pickle.dumps(rioda) - with pickle.loads(temp) as actual: - assert_equal(actual, rioda) - - def test_ENVI_tags(self): - # Create an ENVI file with some tags in the ENVI namespace - # this test uses a custom driver, so we can't use create_tmp_geotiff - with create_tmp_file(suffix=".dat") as tmp_file: +def test_serialization(): + with create_tmp_geotiff(additional_attrs={}) as (tmp_file, expected): + # Write it to a netcdf and read again (roundtrip) + with rioxarray.open_rasterio(tmp_file, mask_and_scale=True) as rioda: + with create_tmp_file(suffix=".nc") as tmp_nc_file: + rioda.to_netcdf(tmp_nc_file) + with xr.open_dataarray(tmp_nc_file, decode_coords="all") as ncds: + assert_identical(rioda, ncds) + + +def test_utm(): + with create_tmp_geotiff() as (tmp_file, expected): + with rioxarray.open_rasterio(tmp_file) as rioda: + assert_allclose(rioda, expected) + assert rioda.attrs["scale_factor"] == 1.0 + assert rioda.attrs["add_offset"] == 0.0 + assert rioda.attrs["long_name"] == ("d1", "d2", "d3") + assert rioda.attrs["units"] == ("u1", "u2", "u3") + assert rioda.rio.crs == expected.rio.crs + assert_array_equal(rioda.rio.resolution(), expected.rio.resolution()) + assert isinstance(rioda.rio._cached_transform(), Affine) + assert rioda.rio.nodata is None + + # Check no parse coords + with rioxarray.open_rasterio(tmp_file, parse_coordinates=False) as rioda: + assert "x" not in rioda.coords + assert "y" not in rioda.coords + + +def test_platecarree(): + with create_tmp_geotiff( + 8, + 10, + 1, + transform_args=[1, 2, 0.5, 2.0], + crs="+proj=latlong", + open_kwargs={"nodata": -9765}, + ) as (tmp_file, expected): + with rioxarray.open_rasterio(tmp_file) as rioda: + assert_allclose(rioda, expected) + assert rioda.attrs["scale_factor"] == 1.0 + assert rioda.attrs["add_offset"] == 0.0 + assert rioda.attrs["long_name"] == "d1" + assert rioda.attrs["units"] == "u1" + assert rioda.rio.crs == expected.rio.crs + assert isinstance(rioda.rio.resolution(), tuple) + assert isinstance(rioda.rio._cached_transform(), Affine) + assert rioda.rio.nodata == -9765.0 + + +def test_notransform(): + # regression test for https://github.com/pydata/xarray/issues/1686 + # Create a geotiff file + with warnings.catch_warnings(): + # rasterio throws a NotGeoreferencedWarning here, which is + # expected since we test rasterio's defaults in this case. + warnings.filterwarnings( + "ignore", + category=UserWarning, + message="Dataset has no geotransform set", + ) + with create_tmp_file(suffix=".tif") as tmp_file: # data nx, ny, nz = 4, 3, 3 data = np.arange(nx * ny * nz, dtype=rasterio.float32).reshape(nz, ny, nx) - transform = from_origin(5000, 80000, 1000, 2000.0) with rasterio.open( tmp_file, "w", - driver="ENVI", + driver="GTiff", height=ny, width=nx, count=nz, - crs={ - "units": "m", - "no_defs": True, - "ellps": "WGS84", - "proj": "utm", - "zone": 18, - }, - transform=transform, dtype=rasterio.float32, ) as s: - s.update_tags( - ns="ENVI", - description="{Tagged file}", - wavelength="{123.000000, 234.234000, 345.345678}", - fwhm="{1.000000, 0.234000, 0.000345}", - ) + s.descriptions = ("nx", "ny", "nz") + s.units = ("cm", "m", "km") s.write(data) - dx, dy = s.res[0], -s.res[1] # Tests - coords = { - "band": [1, 2, 3], - "y": -np.arange(ny) * 2000 + 80000 + dy / 2, - "x": np.arange(nx) * 1000 + 5000 + dx / 2, - "wavelength": ("band", np.array([123, 234.234, 345.345678])), - "fwhm": ("band", np.array([1, 0.234, 0.000345])), - } - expected = DataArray(data, dims=("band", "y", "x"), coords=coords) + expected = DataArray( + data, + dims=("band", "y", "x"), + coords={ + "band": [1, 2, 3], + "y": [0.5, 1.5, 2.5], + "x": [0.5, 1.5, 2.5, 3.5], + }, + ) + expected.coords[DEFAULT_GRID_MAP] = xr.Variable((), 0) + expected.coords[DEFAULT_GRID_MAP].attrs[ + "GeoTransform" + ] = "0.0 1.0 0.0 0.0 0.0 1.0" - with xr.open_rasterio(tmp_file) as rioda: + with rioxarray.open_rasterio(tmp_file) as rioda: assert_allclose(rioda, expected) - assert isinstance(rioda.attrs["crs"], str) - assert isinstance(rioda.attrs["res"], tuple) - assert isinstance(rioda.attrs["is_tiled"], np.uint8) + assert rioda.attrs["scale_factor"] == 1.0 + assert rioda.attrs["add_offset"] == 0.0 + assert rioda.attrs["long_name"] == ("nx", "ny", "nz") + assert rioda.attrs["units"] == ("cm", "m", "km") + assert isinstance(rioda.rio.resolution(), tuple) assert isinstance(rioda.rio._cached_transform(), Affine) - # from ENVI tags - assert isinstance(rioda.attrs["description"], str) - assert isinstance(rioda.attrs["map_info"], str) - assert isinstance(rioda.attrs["samples"], str) - - def test_no_mftime(self): - # rasterio can accept "filename" urguments that are actually urls, - # including paths to remote files. - # In issue #1816, we found that these caused dask to break, because - # the modification time was used to determine the dask token. This - # tests ensure we can still chunk such files when reading with - # rasterio. - with create_tmp_geotiff( - 8, 10, 3, transform_args=[1, 2, 0.5, 2.0], crs="+proj=latlong" - ) as (tmp_file, expected): - with patch("os.path.getmtime", side_effect=OSError): - with xr.open_rasterio(tmp_file, chunks=(1, 2, 2)) as actual: - assert isinstance(actual.data, dask.array.Array) - assert_allclose(actual, expected) - @pytest.mark.xfail(reason="Network could be problematic") - def test_http_url(self): - # more examples urls here - # http://download.osgeo.org/geotiff/samples/ - url = "http://download.osgeo.org/geotiff/samples/made_up/ntf_nord.tif" - with xr.open_rasterio(url) as actual: - assert actual.shape == (1, 512, 512) - # make sure chunking works - with xr.open_rasterio(url, chunks=(1, 256, 256)) as actual: + +def test_indexing(): + with create_tmp_geotiff( + 8, 10, 3, transform_args=[1, 2, 0.5, 2.0], crs="+proj=latlong" + ) 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) + assert not actual.variable._in_memory + + # Basic indexer + ind = {"x": slice(2, 5), "y": slice(5, 7)} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + ind = {"band": slice(1, 2), "x": slice(2, 5), "y": slice(5, 7)} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + ind = {"band": slice(1, 2), "x": slice(2, 5), "y": 0} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + # orthogonal indexer + ind = { + "band": np.array([2, 1, 0]), + "x": np.array([1, 0]), + "y": np.array([0, 2]), + } + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + ind = {"band": np.array([2, 1, 0]), "x": np.array([1, 0]), "y": 0} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + ind = {"band": 0, "x": np.array([0, 0]), "y": np.array([1, 1, 1])} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + # minus-stepped slice + ind = {"band": np.array([2, 1, 0]), "x": slice(-1, None, -1), "y": 0} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + ind = {"band": np.array([2, 1, 0]), "x": 1, "y": slice(-1, 1, -2)} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + # empty selection + ind = {"band": np.array([2, 1, 0]), "x": 1, "y": slice(2, 2, 1)} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + ind = {"band": slice(0, 0), "x": 1, "y": 2} + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + # vectorized indexer + ind = { + "band": DataArray([2, 1, 0], dims="a"), + "x": DataArray([1, 0, 0], dims="a"), + "y": np.array([0, 2]), + } + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + ind = { + "band": DataArray([[2, 1, 0], [1, 0, 2]], dims=["a", "b"]), + "x": DataArray([[1, 0, 0], [0, 1, 0]], dims=["a", "b"]), + "y": 0, + } + assert_allclose(expected.isel(**ind), actual.isel(**ind)) + assert not actual.variable._in_memory + + # Selecting lists of bands is fine + ex = expected.isel(band=[1, 2]) + ac = actual.isel(band=[1, 2]) + assert_allclose(ac, ex) + ex = expected.isel(band=[0, 2]) + ac = actual.isel(band=[0, 2]) + assert_allclose(ac, ex) + + # Integer indexing + ex = expected.isel(band=1) + ac = actual.isel(band=1) + assert_allclose(ac, ex) + + ex = expected.isel(x=1, y=2) + ac = actual.isel(x=1, y=2) + assert_allclose(ac, ex) + + ex = expected.isel(band=0, x=1, y=2) + ac = actual.isel(band=0, x=1, y=2) + assert_allclose(ac, ex) + + # Mixed + ex = actual.isel(x=slice(2), y=slice(2)) + ac = actual.isel(x=[0, 1], y=[0, 1]) + assert_allclose(ac, ex) + + ex = expected.isel(band=0, x=1, y=slice(5, 7)) + ac = actual.isel(band=0, x=1, y=slice(5, 7)) + assert_allclose(ac, ex) + + ex = expected.isel(band=0, x=slice(2, 5), y=2) + ac = actual.isel(band=0, x=slice(2, 5), y=2) + assert_allclose(ac, ex) + + # One-element lists + ex = expected.isel(band=[0], x=slice(2, 5), y=[2]) + ac = actual.isel(band=[0], x=slice(2, 5), y=[2]) + assert_allclose(ac, ex) + + +def test_caching(): + with create_tmp_geotiff( + 8, 10, 3, transform_args=[1, 2, 0.5, 2.0], crs="+proj=latlong" + ) as (tmp_file, expected): + # Cache is the default + with rioxarray.open_rasterio(tmp_file) as actual: + + # This should cache everything + assert_allclose(actual, expected) + + # once cached, non-windowed indexing should become possible + ac = actual.isel(x=[2, 4]) + ex = expected.isel(x=[2, 4]) + assert_allclose(ac, ex) + + +def test_chunks(): + with create_tmp_geotiff( + 8, 10, 3, 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, 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() + ex = expected.mean() + assert_allclose(ac, ex) + + ac = actual.sel(band=1).mean(dim="x") + ex = expected.sel(band=1).mean(dim="x") + 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): + with rioxarray.open_rasterio(tmp_file) as rioda: + temp = pickle.dumps(rioda) + with pickle.loads(temp) as actual: + assert_equal(actual, rioda) + + +def test_ENVI_tags(): + # Create an ENVI file with some tags in the ENVI namespace + # this test uses a custom driver, so we can't use create_tmp_geotiff + with create_tmp_file(suffix=".dat") as tmp_file: + # data + nx, ny, nz = 4, 3, 3 + data = np.arange(nx * ny * nz, dtype=rasterio.float32).reshape(nz, ny, nx) + transform = from_origin(5000, 80000, 1000, 2000.0) + with rasterio.open( + tmp_file, + "w", + driver="ENVI", + height=ny, + width=nx, + count=nz, + crs={ + "units": "m", + "no_defs": True, + "ellps": "WGS84", + "proj": "utm", + "zone": 18, + }, + transform=transform, + dtype=rasterio.float32, + ) as s: + s.update_tags( + ns="ENVI", + description="{Tagged file}", + wavelength="{123.000000, 234.234000, 345.345678}", + fwhm="{1.000000, 0.234000, 0.000345}", + ) + s.write(data) + dx, dy = s.res[0], -s.res[1] + crs_wkt = s.crs.to_wkt() + + # Tests + coords = { + "band": [1, 2, 3], + "y": -np.arange(ny) * 2000 + 80000 + dy / 2, + "x": np.arange(nx) * 1000 + 5000 + dx / 2, + "wavelength": ("band", np.array([123, 234.234, 345.345678])), + "fwhm": ("band", np.array([1, 0.234, 0.000345])), + } + expected = DataArray(data, dims=("band", "y", "x"), coords=coords) + expected.coords[DEFAULT_GRID_MAP] = xr.Variable((), 0) + expected.coords[DEFAULT_GRID_MAP].attrs["crs_wkt"] = crs_wkt + + with rioxarray.open_rasterio(tmp_file) as rioda: + assert_allclose(rioda, expected) + assert rioda.rio.crs == crs_wkt + assert isinstance(rioda.rio._cached_transform(), Affine) + # from ENVI tags + assert isinstance(rioda.attrs["description"], str) + assert isinstance(rioda.attrs["map_info"], str) + assert isinstance(rioda.attrs["samples"], str) + + +def test_no_mftime(): + # rasterio can accept "filename" urguments that are actually urls, + # including paths to remote files. + # In issue #1816, we found that these caused dask to break, because + # the modification time was used to determine the dask token. This + # tests ensure we can still chunk such files when reading with + # rasterio. + with create_tmp_geotiff( + 8, 10, 3, transform_args=[1, 2, 0.5, 2.0], crs="+proj=latlong" + ) as (tmp_file, expected): + with patch("os.path.getmtime", side_effect=OSError): + with rioxarray.open_rasterio(tmp_file, chunks=(1, 2, 2)) as actual: + assert isinstance(actual.data, dask.array.Array) + assert_allclose(actual, expected) + - def test_rasterio_environment(self): - with create_tmp_geotiff() as (tmp_file, expected): - # Should fail with error since suffix not allowed - with pytest.raises(Exception): - with rasterio.Env(GDAL_SKIP="GTiff"): - with xr.open_rasterio(tmp_file) as actual: - assert_allclose(actual, expected) - - @pytest.mark.xfail( - rasterio.__version__ == "1.1.1", - reason="https://github.com/mapbox/rasterio/issues/1833", +@pytest.mark.xfail( + error=rasterio.errors.RasterioIOError, reason="Network could be problematic" +) +def test_http_url(): + # more examples urls here + # http://download.osgeo.org/geotiff/samples/ + url = ( + "https://github.com/corteva/rioxarray/blob/master/" + "test/test_data/input/cog.tif?raw=true" ) - def test_rasterio_vrt(self): - # 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: - with rasterio.vrt.WarpedVRT(src, crs="epsg:4326") as vrt: - expected_shape = (vrt.width, vrt.height) - expected_crs = vrt.crs - expected_res = vrt.res - # 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 xr.open_rasterio(vrt) as rds: - actual_shape = (rds.sizes["x"], rds.sizes["y"]) - actual_crs = rds.crs - actual_res = rds.res - actual_val = rds.sel(dict(x=lon, y=lat), method="nearest").data - - assert actual_crs == expected_crs - assert actual_res == expected_res - assert actual_shape == expected_shape - assert expected_val.all() == actual_val.all() - - def test_rasterio_vrt_with_transform_and_size(self): - # Test open_rasterio() support of WarpedVRT with transform, width and - # height (issue #2864) - with create_tmp_geotiff() as (tmp_file, expected): - with rasterio.open(tmp_file) as src: - # Estimate the transform, width and height - # for a change of resolution - # tmp_file initial res is (1000,2000) (default values) - trans, w, h = calculate_default_transform( - src.crs, src.crs, src.width, src.height, resolution=500, *src.bounds - ) - with rasterio.vrt.WarpedVRT( - src, transform=trans, width=w, height=h - ) as vrt: - expected_shape = (vrt.width, vrt.height) - expected_res = vrt.res - expected_transform = vrt.transform - with xr.open_rasterio(vrt) as rds: - actual_shape = (rds.sizes["x"], rds.sizes["y"]) - actual_res = rds.res - actual_transform = Affine(*rds.transform) - assert actual_res == expected_res - assert actual_shape == expected_shape - assert actual_transform == expected_transform - - @pytest.mark.xfail(reason="Network could be problematic") - def test_rasterio_vrt_network(self): - url = "https://storage.googleapis.com/\ - gcp-public-data-landsat/LC08/01/047/027/\ - LC08_L1TP_047027_20130421_20170310_01_T1/\ - LC08_L1TP_047027_20130421_20170310_01_T1_B4.TIF" - env = rasterio.Env( - GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", - CPL_VSIL_CURL_USE_HEAD=False, - CPL_VSIL_CURL_ALLOWED_EXTENSIONS="TIF", - ) - with env: - with rasterio.open(url) as src: - with rasterio.vrt.WarpedVRT(src, crs="epsg:4326") as vrt: - expected_shape = (vrt.width, vrt.height) - expected_crs = vrt.crs - expected_res = vrt.res - # 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 xr.open_rasterio(vrt) as rds: - actual_shape = (rds.sizes["x"], rds.sizes["y"]) - actual_crs = rds.crs - actual_res = rds.res - actual_val = rds.sel(dict(x=lon, y=lat), method="nearest").data - - assert_equal(actual_shape, expected_shape) - assert_equal(actual_crs, expected_crs) - assert_equal(actual_res, expected_res) - assert_equal(expected_val, actual_val) - - def test_rasterio_vrt_with_src_crs(self): - # Test open_rasterio() support of WarpedVRT with specified src_crs - # create geotiff with no CRS and specify it manually - with create_tmp_geotiff(crs=None) as (tmp_file, expected): - src_crs = rasterio.crs.CRS.from_epsg(32618) - with rasterio.open(tmp_file) as src: - assert src.crs is None - with rasterio.vrt.WarpedVRT(src, src_crs=src_crs) as vrt: - with rioxarray.open_rasterio(vrt) as rds: - assert rds.rio.crs == src_crs + with rioxarray.open_rasterio(url) as actual: + assert actual.shape == (1, 500, 500) + # make sure chunking works + with rioxarray.open_rasterio(url, chunks=(1, 256, 256)) as actual: + assert isinstance(actual.data, dask.array.Array) + + +def test_rasterio_environment(): + with create_tmp_geotiff() as (tmp_file, expected): + # Should fail with error since suffix not allowed + with pytest.raises(Exception): + with rasterio.Env(GDAL_SKIP="GTiff"): + with rioxarray.open_rasterio(tmp_file) as actual: + assert_allclose(actual, expected) + + +@pytest.mark.xfail( + rasterio.__version__ == "1.1.1", + reason="https://github.com/mapbox/rasterio/issues/1833", +) +def test_rasterio_vrt(): + # 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: + with rasterio.vrt.WarpedVRT(src, crs="epsg:4326") as 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: + actual_val = rds.sel(dict(x=lon, y=lat), method="nearest").data + + assert_array_equal(rds.rio.shape, (vrt.height, vrt.width)) + assert rds.rio.crs == vrt.crs + assert_array_equal( + tuple(abs(val) for val in rds.rio.resolution()), vrt.res + ) + assert expected_val.all() == actual_val.all() + + +def test_rasterio_vrt_with_transform_and_size(): + # Test open_rasterio() support of WarpedVRT with transform, width and + # height (issue #2864) + with create_tmp_geotiff() as (tmp_file, expected): + with rasterio.open(tmp_file) as src: + # Estimate the transform, width and height + # for a change of resolution + # tmp_file initial res is (1000,2000) (default values) + trans, w, h = calculate_default_transform( + src.crs, src.crs, src.width, src.height, resolution=500, *src.bounds + ) + with rasterio.vrt.WarpedVRT(src, transform=trans, width=w, height=h) as vrt: + with rioxarray.open_rasterio(vrt) as rds: + assert_array_equal(rds.rio.shape, (vrt.height, vrt.width)) + assert rds.rio.crs == vrt.crs + assert_array_equal( + tuple(abs(val) for val in rds.rio.resolution()), vrt.res + ) + assert rds.rio.transform() == vrt.transform + + +@pytest.mark.xfail(reason="Network could be problematic") +def test_rasterio_vrt_network(): + url = "https://storage.googleapis.com/\ + gcp-public-data-landsat/LC08/01/047/027/\ + LC08_L1TP_047027_20130421_20170310_01_T1/\ + LC08_L1TP_047027_20130421_20170310_01_T1_B4.TIF" + env = rasterio.Env( + GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", + CPL_VSIL_CURL_USE_HEAD=False, + CPL_VSIL_CURL_ALLOWED_EXTENSIONS="TIF", + ) + with env: + with rasterio.open(url) as src: + with rasterio.vrt.WarpedVRT(src, crs="epsg:4326") as 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: + actual_val = rds.sel(dict(x=lon, y=lat), method="nearest").data + assert_array_equal(rds.rio.shape, (vrt.height, vrt.width)) + assert rds.rio.crs == vrt.crs + assert_array_equal( + tuple(abs(val) for val in rds.rio.resolution()), vrt.res + ) + assert_equal(expected_val, actual_val) + + +def test_rasterio_vrt_with_src_crs(): + # Test open_rasterio() support of WarpedVRT with specified src_crs + # create geotiff with no CRS and specify it manually + with create_tmp_geotiff(crs=None) as (tmp_file, expected): + src_crs = rasterio.crs.CRS.from_epsg(32618) + with rasterio.open(tmp_file) as src: + assert src.crs is None + with rasterio.vrt.WarpedVRT(src, src_crs=src_crs) as vrt: + with rioxarray.open_rasterio(vrt) as rds: + assert rds.rio.crs == src_crs def test_open_cog(): @@ -852,10 +864,10 @@ def test_open_cog(): assert rdso.shape == (1, 250, 250) -def test_mask_and_scale(): +def test_mask_and_scale(open_rasterio): test_file = os.path.join(TEST_INPUT_DATA_DIR, "tmmx_20190121.nc") with pytest.warns(SerializationWarning): - with rioxarray.open_rasterio(test_file, mask_and_scale=True) as rds: + with open_rasterio(test_file, mask_and_scale=True) as rds: assert np.nanmin(rds.air_temperature.values) == 248.7 assert np.nanmax(rds.air_temperature.values) == 302.1 test_encoding = dict(rds.air_temperature.encoding) @@ -883,9 +895,11 @@ def test_mask_and_scale(): } -def test_no_mask_and_scale(): - with rioxarray.open_rasterio( - os.path.join(TEST_INPUT_DATA_DIR, "tmmx_20190121.nc"), masked=True +def test_no_mask_and_scale(open_rasterio): + with open_rasterio( + os.path.join(TEST_INPUT_DATA_DIR, "tmmx_20190121.nc"), + mask_and_scale=False, + masked=True, ) as rds: assert np.nanmin(rds.air_temperature.values) == 287 assert np.nanmax(rds.air_temperature.values) == 821 @@ -914,20 +928,18 @@ def test_no_mask_and_scale(): } -def test_notgeoreferenced_warning(): +def test_notgeoreferenced_warning(open_rasterio): with create_tmp_geotiff(transform_args=None) as (tmp_file, expected): with pytest.warns(NotGeoreferencedWarning): - rioxarray.open_rasterio(tmp_file) + open_rasterio(tmp_file) @pytest.mark.xfail( LooseVersion(rasterio.__gdal_version__) < LooseVersion("3.0.4"), reason="This was fixed in GDAL 3.0.4", ) -def test_nc_attr_loading(): - with rioxarray.open_rasterio( - os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc") - ) as rds: +def test_nc_attr_loading(open_rasterio): + with open_rasterio(os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc")) as rds: assert rds.dims == {"y": 10, "x": 10, "time": 2} assert rds.attrs == {"coordinates": "spatial_ref"} assert rds.y.attrs["units"] == "metre" @@ -942,14 +954,14 @@ def test_nc_attr_loading(): def test_lockless(): with rioxarray.open_rasterio( - os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc"), lock=False, chunk=True + os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc"), lock=False, chunks=True ) as rds: rds.mean().compute() def test_lock_true(): with rioxarray.open_rasterio( - os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc"), lock=True, chunk=True + os.path.join(TEST_INPUT_DATA_DIR, "PLANET_SCOPE_3D.nc"), lock=True, chunks=True ) as rds: rds.mean().compute() @@ -979,9 +991,9 @@ def test_non_rectilinear(): assert isinstance(rioda.rio._cached_transform(), Affine) -def test_non_rectilinear__load_coords(): +def test_non_rectilinear__load_coords(open_rasterio): test_file = os.path.join(TEST_INPUT_DATA_DIR, "2d_test.tif") - xds = rioxarray.open_rasterio(test_file) + xds = open_rasterio(test_file) assert xds.rio.shape == (10, 10) with rasterio.open(test_file) as rds: assert rds.transform == xds.rio.transform() @@ -990,9 +1002,9 @@ def test_non_rectilinear__load_coords(): assert_almost_equal(rds.xy(yi, xi), (subset.xc.item(), subset.yc.item())) -def test_non_rectilinear__skip_parse_coordinates(): +def test_non_rectilinear__skip_parse_coordinates(open_rasterio): test_file = os.path.join(TEST_INPUT_DATA_DIR, "2d_test.tif") - xds = rioxarray.open_rasterio(test_file, parse_coordinates=False) + xds = open_rasterio(test_file, parse_coordinates=False) assert "xc" not in xds.coords assert "yc" not in xds.coords assert xds.rio.shape == (10, 10) diff --git a/test/integration/test_integration_rioxarray.py b/test/integration/test_integration_rioxarray.py index 913ade54..ed4ae546 100644 --- a/test/integration/test_integration_rioxarray.py +++ b/test/integration/test_integration_rioxarray.py @@ -7,7 +7,6 @@ import dask.array as da import numpy -import pyproj import pytest import rasterio import scipy @@ -30,16 +29,22 @@ ) from rioxarray.rioxarray import _make_coords from test.conftest import ( + PYPROJ_LT_3, + RASTERIO_LT_122, TEST_COMPARE_DATA_DIR, TEST_INPUT_DATA_DIR, _assert_xarrays_equal, + open_rasterio_engine, ) -PYPROJ_LT_3 = LooseVersion(pyproj.__version__) < LooseVersion("3") -RASTERIO_LT_122 = LooseVersion(rasterio.__version__) < LooseVersion("1.2.2") - -@pytest.fixture(params=[xarray.open_dataset, xarray.open_dataarray]) +@pytest.fixture( + params=[ + xarray.open_dataset, + xarray.open_dataarray, + open_rasterio_engine, + ] +) def modis_reproject(request): return dict( input=os.path.join(TEST_INPUT_DATA_DIR, "MODIS_ARRAY.nc"), @@ -58,7 +63,13 @@ def modis_reproject_3d(): ) -@pytest.fixture(params=[xarray.open_dataset, xarray.open_dataarray]) +@pytest.fixture( + params=[ + xarray.open_dataset, + xarray.open_dataarray, + open_rasterio_engine, + ] +) def interpolate_na(request): return dict( input=os.path.join(TEST_INPUT_DATA_DIR, "MODIS_ARRAY.nc"), @@ -75,7 +86,13 @@ def interpolate_na_3d(): ) -@pytest.fixture(params=[xarray.open_dataset, xarray.open_dataarray]) +@pytest.fixture( + params=[ + xarray.open_dataset, + xarray.open_dataarray, + open_rasterio_engine, + ] +) def interpolate_na_filled(request): return dict( input=os.path.join(TEST_INPUT_DATA_DIR, "MODIS_ARRAY.nc"), @@ -95,7 +112,12 @@ def interpolate_na_veris(): @pytest.fixture( - params=[xarray.open_dataarray, rioxarray.open_rasterio, xarray.open_dataset] + params=[ + xarray.open_dataarray, + rioxarray.open_rasterio, + xarray.open_dataset, + open_rasterio_engine, + ] ) def interpolate_na_nan(request): return dict( @@ -111,6 +133,7 @@ def interpolate_na_nan(request): xarray.open_dataarray, rioxarray.open_rasterio, partial(rioxarray.open_rasterio, parse_coordinates=False), + open_rasterio_engine, ] ) def modis_reproject_match(request): @@ -123,7 +146,12 @@ def modis_reproject_match(request): @pytest.fixture( - params=[xarray.open_dataset, xarray.open_dataarray, rioxarray.open_rasterio] + params=[ + xarray.open_dataset, + xarray.open_dataarray, + rioxarray.open_rasterio, + open_rasterio_engine, + ] ) def modis_reproject_match_coords(request): return dict( @@ -161,6 +189,7 @@ def _del_attr(input_xr, attr): xarray.open_dataarray, rioxarray.open_rasterio, partial(rioxarray.open_rasterio, parse_coordinates=False), + open_rasterio_engine, ] ) def modis_clip(request, tmpdir): @@ -380,15 +409,16 @@ def test_slice_xy(modis_clip): clipped_ds.to_netcdf(modis_clip["output"]) +@pytest.mark.parametrize("from_disk", [True, False]) @pytest.mark.parametrize( - "open_func,from_disk", + "open_func", [ - (xarray.open_rasterio, False), - (rioxarray.open_rasterio, False), - (rioxarray.open_rasterio, True), - (partial(rioxarray.open_rasterio, parse_coordinates=False), False), - (partial(rioxarray.open_rasterio, parse_coordinates=False), True), - (partial(rioxarray.open_rasterio, parse_coordinates=False, masked=True), True), + rioxarray.open_rasterio, + partial(rioxarray.open_rasterio, parse_coordinates=False), + partial(rioxarray.open_rasterio, parse_coordinates=False, masked=True), + open_rasterio_engine, + partial(open_rasterio_engine, parse_coordinates=False), + partial(open_rasterio_engine, parse_coordinates=False, mask_and_scale=False), ], ) def test_clip_geojson(open_func, from_disk): @@ -398,6 +428,8 @@ def test_clip_geojson(open_func, from_disk): # get subset for testing subset = xdi.isel(x=slice(150, 160), y=slice(100, 150)) comp_subset = subset.isel(x=slice(1, None), y=slice(1, None)) + if isinstance(comp_subset, xarray.Dataset): + comp_subset = comp_subset.band_data # add transform for test comp_subset.rio.write_transform(inplace=True) # add grid mapping for test @@ -421,7 +453,7 @@ def test_clip_geojson(open_func, from_disk): ] # test data array clipped = xdi.rio.clip(geometries, from_disk=from_disk) - if from_disk: + if from_disk and not isinstance(clipped, xarray.Dataset): _assert_xarrays_equal(clipped[:, 1:, 1:], comp_subset) if comp_subset.rio.encoded_nodata is not None: assert numpy.isnan(clipped.values[:, 0, :]).all() @@ -430,11 +462,16 @@ def test_clip_geojson(open_func, from_disk): assert (clipped.values[:, 0, :] == comp_subset.rio.nodata).all() assert (clipped.values[:, :, 0] == comp_subset.rio.nodata).all() else: + if isinstance(clipped, xarray.Dataset): + clipped = clipped.band_data _assert_xarrays_equal(clipped, comp_subset) # test dataset - clipped_ds = xdi.to_dataset(name="test_data").rio.clip(geometries) - comp_subset_ds = comp_subset.to_dataset(name="test_data") + if isinstance(xdi, xarray.DataArray): + clipped_ds = xdi.to_dataset(name="band_data").rio.clip(geometries) + else: + clipped_ds = xdi.rio.clip(geometries) + comp_subset_ds = comp_subset.to_dataset(name="band_data") # This coordinate checking is skipped when parse_coordinates=False # as the auto-generated coordinates differ and can be ignored _assert_xarrays_equal( @@ -459,7 +496,6 @@ def test_clip_geojson(open_func, from_disk): @pytest.mark.parametrize( "open_func", [ - xarray.open_rasterio, rioxarray.open_rasterio, partial(rioxarray.open_rasterio, parse_coordinates=False), ], @@ -509,6 +545,8 @@ def test_clip_geojson__no_drop(open_func, invert, from_disk, expected_sum): xarray.open_dataarray, rioxarray.open_rasterio, partial(rioxarray.open_rasterio, parse_coordinates=False), + open_rasterio_engine, + partial(open_rasterio_engine, parse_coordinates=False), ], ) def test_transform_bounds(open_func): @@ -532,8 +570,8 @@ def test_transform_bounds(open_func): def test_reproject_with_shape(modis_reproject): new_shape = (9, 10) mask_args = ( - dict(masked=False) - if "open_rasterio" in str(modis_reproject["open"]) + dict(masked=False, mask_and_scale=False) + if "rasterio" in str(modis_reproject["open"]) else dict(mask_and_scale=False) ) with modis_reproject["open"](modis_reproject["input"], **mask_args) as mda: @@ -541,15 +579,15 @@ def test_reproject_with_shape(modis_reproject): # test if hasattr(mds_repr, "variables"): for var in mds_repr.rio.vars: - assert mds_repr[var].shape == new_shape + assert mds_repr[var].rio.shape == new_shape else: - assert mds_repr.shape == new_shape + assert mds_repr.rio.shape == new_shape def test_reproject(modis_reproject): mask_args = ( - dict(masked=False) - if "open_rasterio" in str(modis_reproject["open"]) + dict(masked=False, mask_and_scale=False) + if "rasterio" in str(modis_reproject["open"]) else dict(mask_and_scale=False) ) with modis_reproject["open"]( @@ -577,6 +615,8 @@ def test_reproject(modis_reproject): [ rioxarray.open_rasterio, partial(rioxarray.open_rasterio, parse_coordinates=False), + open_rasterio_engine, + partial(open_rasterio_engine, parse_coordinates=False), ], ) def test_reproject_3d(open_func, modis_reproject_3d): @@ -603,8 +643,8 @@ def test_reproject_3d(open_func, modis_reproject_3d): def test_reproject__grid_mapping(modis_reproject): mask_args = ( - dict(masked=False) - if "open_rasterio" in str(modis_reproject["open"]) + dict(masked=False, mask_and_scale=False) + if "rasterio" in str(modis_reproject["open"]) else dict(mask_and_scale=False) ) with modis_reproject["open"]( @@ -645,7 +685,7 @@ def test_reproject__no_transform(modis_reproject): with modis_reproject["open"](modis_reproject["input"]) as mda, modis_reproject[ "open" ](modis_reproject["compare"]) as mdc: - orig_trans = _get_attr(mda, "transform") + orig_trans = mda.rio.transform() _del_attr(mda, "transform") # reproject mds_repr = mda.rio.reproject(modis_reproject["to_proj"]) @@ -660,8 +700,8 @@ def test_reproject__no_transform(modis_reproject): def test_reproject__no_nodata(modis_reproject): mask_args = ( - dict(masked=False) - if "open_rasterio" in str(modis_reproject["open"]) + dict(masked=False, mask_and_scale=False) + if "rasterio" in str(modis_reproject["open"]) else dict(mask_and_scale=False) ) with modis_reproject["open"]( @@ -688,7 +728,7 @@ def test_reproject__no_nodata(modis_reproject): _assert_xarrays_equal(mds_repr, mdc) -@pytest.mark.parametrize("open_func", [xarray.open_rasterio, rioxarray.open_rasterio]) +@pytest.mark.parametrize("open_func", [rioxarray.open_rasterio, open_rasterio_engine]) def test_reproject__scalar_coord(open_func): with open_func( os.path.join(TEST_COMPARE_DATA_DIR, "small_dem_3m_merged.tif") @@ -711,8 +751,8 @@ def test_reproject__no_nodata_masked(modis_reproject): def test_reproject_match(modis_reproject_match): mask_args = ( - dict(masked=False) - if "open_rasterio" in str(modis_reproject_match["open"]) + dict(masked=False, mask_and_scale=False) + if "rasterio" in str(modis_reproject_match["open"]) else dict(mask_and_scale=False) ) with modis_reproject_match["open"]( @@ -744,7 +784,7 @@ def test_reproject_match(modis_reproject_match): def test_reproject_match__masked(modis_reproject_match): mask_args = ( dict(masked=True) - if "open_rasterio" in str(modis_reproject_match["open"]) + if "rasterio" in str(modis_reproject_match["open"]) else dict(mask_and_scale=True) ) with modis_reproject_match["open"]( @@ -763,7 +803,7 @@ def test_reproject_match__masked(modis_reproject_match): def test_reproject_match__no_transform_nodata(modis_reproject_match_coords): mask_args = ( dict(masked=True) - if "open_rasterio" in str(modis_reproject_match_coords["open"]) + if "rasterio" in str(modis_reproject_match_coords["open"]) else dict(mask_and_scale=True) ) with modis_reproject_match_coords["open"]( @@ -781,7 +821,7 @@ def test_reproject_match__no_transform_nodata(modis_reproject_match_coords): _assert_xarrays_equal(mds_repr, mdc) -@pytest.mark.parametrize("open_func", [xarray.open_rasterio, rioxarray.open_rasterio]) +@pytest.mark.parametrize("open_func", [rioxarray.open_rasterio, open_rasterio_engine]) def test_make_src_affine(open_func, modis_reproject): with xarray.open_dataarray(modis_reproject["input"]) as xdi, open_func( modis_reproject["input"] @@ -824,9 +864,10 @@ def test_make_src_affine__single_point(): "open_func", [ xarray.open_dataset, - xarray.open_rasterio, + open_rasterio_engine, rioxarray.open_rasterio, partial(rioxarray.open_rasterio, parse_coordinates=False), + partial(open_rasterio_engine, parse_coordinates=False), ], ) def test_make_coords__calc_trans(open_func, modis_reproject): @@ -864,9 +905,10 @@ def test_make_coords__calc_trans(open_func, modis_reproject): "open_func", [ xarray.open_dataset, - xarray.open_rasterio, rioxarray.open_rasterio, partial(rioxarray.open_rasterio, parse_coordinates=False), + open_rasterio_engine, + partial(open_rasterio_engine, parse_coordinates=False), ], ) def test_make_coords__attr_trans(open_func, modis_reproject): @@ -907,8 +949,8 @@ def test_make_coords__attr_trans(open_func, modis_reproject): def test_interpolate_na(interpolate_na): mask_args = ( - dict(masked=False) - if "open_rasterio" in str(interpolate_na["open"]) + dict(masked=False, mask_and_scale=False) + if "rasterio" in str(interpolate_na["open"]) else dict(mask_and_scale=False) ) with interpolate_na["open"]( @@ -944,8 +986,8 @@ def test_interpolate_na_3d(interpolate_na_3d): def test_interpolate_na__nodata_filled(interpolate_na_filled): mask_args = ( - dict(masked=False) - if "open_rasterio" in str(interpolate_na_filled["open"]) + dict(masked=False, mask_and_scale=False) + if "rasterio" in str(interpolate_na_filled["open"]) else dict(mask_and_scale=False) ) with interpolate_na_filled["open"]( @@ -965,12 +1007,11 @@ def test_interpolate_na__nodata_filled(interpolate_na_filled): def test_interpolate_na__all_nodata(interpolate_na_nan): - rio_opened = "open_rasterio" in str(interpolate_na_nan["open"]) - mask_args = dict(masked=True) if rio_opened else dict(mask_and_scale=True) + rio_opened = "open_rasterio " in str(interpolate_na_nan["open"]) with interpolate_na_nan["open"]( - interpolate_na_nan["input"], **mask_args + interpolate_na_nan["input"], mask_and_scale=True ) as mda, interpolate_na_nan["open"]( - interpolate_na_nan["compare"], **mask_args + interpolate_na_nan["compare"], mask_and_scale=True ) as mdc: if hasattr(mda, "variables"): for var in mda.rio.vars: @@ -981,7 +1022,6 @@ def test_interpolate_na__all_nodata(interpolate_na_nan): interpolated_ds = mda.rio.interpolate_na() if rio_opened and "__xarray_dataarray_variable__" in mdc: mdc = mdc["__xarray_dataarray_variable__"] - mdc.attrs.pop("coordinates") # test _assert_xarrays_equal(interpolated_ds, mdc) @@ -1056,6 +1096,7 @@ def test_geographic_resample_integer(): partial( rioxarray.open_rasterio, masked=True, chunks=True, lock=threading.Lock() ), + open_rasterio_engine, ], ) @pytest.mark.parametrize( @@ -1079,6 +1120,8 @@ def test_to_raster( tmp_raster = tmpdir.join("modis_raster.tif") test_tags = {"test": "1"} with open_method(os.path.join(TEST_INPUT_DATA_DIR, "MODIS_ARRAY.nc")) as mda: + if isinstance(mda, xarray.Dataset): + mda = mda.__xarray_dataarray_variable__ delayed = mda.rio.to_raster( str(tmp_raster), windowed=windowed, @@ -1128,6 +1171,7 @@ def test_to_raster( partial( rioxarray.open_rasterio, masked=True, chunks=True, lock=threading.Lock() ), + open_rasterio_engine, ], ) @pytest.mark.parametrize("windowed", [True, False]) diff --git a/test/integration/test_integration_xarray_plugin.py b/test/integration/test_integration_xarray_plugin.py index 2fbb06a8..2bfb921e 100644 --- a/test/integration/test_integration_xarray_plugin.py +++ b/test/integration/test_integration_xarray_plugin.py @@ -2,6 +2,7 @@ import pytest +from rioxarray.exceptions import RioXarrayError from test.conftest import TEST_INPUT_DATA_DIR # FIXME: change to the next xarray version after release @@ -24,3 +25,16 @@ def test_xarray_open_dataset(): ds = xr.open_dataset(cog_file) assert isinstance(ds, xr.Dataset) + + +def test_open_multiple_resolution(): + with pytest.raises( + RioXarrayError, + match="Multiple resolution sets found. Use 'variable' or 'group' to filter.", + ): + xr.open_dataset( + os.path.join( + TEST_INPUT_DATA_DIR, "MOD09GA.A2008296.h14v17.006.2015181011753.hdf" + ), + engine="rasterio", + )