diff --git a/datashader/tests/test_raster.py b/datashader/tests/test_raster.py index 4072098c1..896d38a2b 100644 --- a/datashader/tests/test_raster.py +++ b/datashader/tests/test_raster.py @@ -6,6 +6,9 @@ import datashader as ds import xarray as xr import numpy as np +import dask.array as da + +from datashader.resampling import compute_chunksize BASE_PATH = path.split(__file__)[0] DATA_PATH = path.abspath(path.join(BASE_PATH, 'data')) @@ -358,6 +361,83 @@ def test_raster_single_pixel_range_with_padding(): assert np.allclose(agg.y.values, np.array([-0.40625, -0.21875, -0.03125, 0.15625])) +def test_raster_distributed_downsample(): + """ + Ensure that distributed regrid is equivalent to regular regrid. + """ + cvs = ds.Canvas(plot_height=2, plot_width=2) + + size = 4 + vs = np.linspace(-1, 1, size) + xs, ys = np.meshgrid(vs, vs) + arr = np.sin(xs*ys) + + darr = da.from_array(arr, (2, 2)) + xr_darr = xr.DataArray(darr, coords=[('y', range(size)), ('x', range(size))], name='z') + xr_arr = xr.DataArray(arr, coords=[('y', range(size)), ('x', range(size))], name='z') + + agg_arr = cvs.raster(xr_arr) + agg_darr = cvs.raster(xr_darr) + + assert np.allclose(agg_arr.data, agg_darr.data.compute()) + assert np.allclose(agg_arr.x.values, agg_darr.x.values) + assert np.allclose(agg_arr.y.values, agg_darr.y.values) + + +def test_raster_distributed_upsample(): + """ + Ensure that distributed regrid is equivalent to regular regrid. + """ + cvs = ds.Canvas(plot_height=8, plot_width=8) + + size = 4 + vs = np.linspace(-1, 1, size) + xs, ys = np.meshgrid(vs, vs) + arr = np.sin(xs*ys) + + darr = da.from_array(arr, (2, 2)) + xr_darr = xr.DataArray(darr, coords=[('y', range(size)), ('x', range(size))], name='z') + xr_arr = xr.DataArray(arr, coords=[('y', range(size)), ('x', range(size))], name='z') + + agg_arr = cvs.raster(xr_arr, interpolate='nearest') + agg_darr = cvs.raster(xr_darr, interpolate='nearest') + + assert np.allclose(agg_arr.data, agg_darr.data.compute()) + assert np.allclose(agg_arr.x.values, agg_darr.x.values) + assert np.allclose(agg_arr.y.values, agg_darr.y.values) + + +def test_raster_distributed_regrid_chunksize(): + """ + Ensure that distributed regrid respects explicit chunk size. + """ + cvs = ds.Canvas(plot_height=2, plot_width=2) + + size = 4 + vs = np.linspace(-1, 1, size) + xs, ys = np.meshgrid(vs, vs) + arr = np.sin(xs*ys) + + darr = da.from_array(arr, (2, 2)) + xr_darr = xr.DataArray(darr, coords=[('y', range(size)), ('x', range(size))], name='z') + + agg_darr = cvs.raster(xr_darr, chunksize=(1, 1)) + + assert agg_darr.data.chunksize == (1, 1) + + +def test_resample_compute_chunksize(): + """ + Ensure chunksize computation is correct. + """ + darr = da.from_array(np.zeros((100, 100)), (10, 10)) + + mem_limited_chunksize = compute_chunksize(darr, 10, 10, max_mem=2000) + assert mem_limited_chunksize == (2, 1) + + explicit_chunksize = compute_chunksize(darr, 10, 10, chunksize=(5, 4)) + assert explicit_chunksize == (5, 4) + def test_resample_methods(): """Assert that an error is raised when incorrect upsample and/or downsample