def test_read_raster_multi_band(some_raster_path):
array = read_raster(some_raster_path, band=(1, 3))
assert isinstance(array, da.Array)
expected_array = da.stack([
read_raster_band(some_raster_path, band=1),
read_raster_band(some_raster_path, band=3)
])
assert array.shape == expected_array.shape
assert array.dtype == expected_array.dtype
assert_array_equal(array.compute(), expected_array.compute())
def test_read_raster_band_with_block_size(some_raster_path):
array = read_raster_band(some_raster_path)
array_4b = read_raster_band(some_raster_path, block_size=4)
assert array.shape == array_4b.shape
assert array.dtype == array_4b.dtype
assert_array_equal(array, array_4b)
with rasterio.open(some_raster_path) as src:
block_height, block_width = src.block_shapes[0]
height, width = src.shape
assert array.chunks[0][0] == block_height
assert array.chunks[1][0] == block_width
assert array_4b.chunks[0][0] == min(block_height * 4, height)
assert array_4b.chunks[1][0] == min(block_width * 4, width)
def test_read_raster_single_band(some_raster_path):
array = read_raster(some_raster_path, band=3)
assert isinstance(array, da.Array)
expected_array = read_raster_band(some_raster_path, band=3)
assert array.shape == expected_array.shape
assert array.dtype == expected_array.dtype
assert_array_equal(array.compute(), expected_array.compute())
def test_do_calcs_on_array(some_raster_path):
r_array = read_raster_band(some_raster_path, 1)
mean = np.mean(r_array)
assert isinstance(mean, da.Array)
with rasterio.open(some_raster_path) as src:
expected_array = src.read(1)
expected_mean = np.mean(expected_array)
assert mean.compute() == expected_mean
def test_read_raster_band(some_raster_path):
g_array = read_raster_band(some_raster_path, 2)
assert isinstance(g_array, da.Array)
with rasterio.open(some_raster_path) as src:
expected_array = src.read(2)
assert g_array.shape == expected_array.shape
assert g_array.dtype == expected_array.dtype
assert_array_equal(g_array.compute(), expected_array)
def test_write_raster_band(some_raster_path):
with tempfile.TemporaryDirectory(prefix='dask_rasterio_test_') as tmpdir:
# Read first band of raster
array = read_raster_band(some_raster_path, 1)
# Generate new data
new_array = array & (array > THRESHOLD)
# Build a profile for the new single-band GeoTIFF
prof = get_profile(some_raster_path)
prof.update(count=1)
# Write raster file
dst_path = os.path.join(tmpdir, 'test.tif')
write_raster(dst_path, new_array, **prof)
with rasterio.open(dst_path) as src:
assert src.count == 1
expected_new_array = src.read(1)
assert expected_new_array.dtype == new_array.dtype
assert_array_equal(new_array.compute(), expected_new_array)