From a52e5055492add3b044e7c7edb1f68b6768c85cd Mon Sep 17 00:00:00 2001 From: Robbi Bishop-Taylor Date: Tue, 14 Nov 2023 09:41:05 +1100 Subject: [PATCH] Pull pixel_tide resampling into function for easier re-use (#1149) * Pull pixel_tide resampling into function * Identify x and y dims inside func instead * Remove x, y dim params from function call --- Tools/dea_tools/coastal.py | 106 +++++++++++++++++++++++++++---------- 1 file changed, 78 insertions(+), 28 deletions(-) diff --git a/Tools/dea_tools/coastal.py b/Tools/dea_tools/coastal.py index 35ed6f31b..b25b83121 100644 --- a/Tools/dea_tools/coastal.py +++ b/Tools/dea_tools/coastal.py @@ -16,7 +16,7 @@ If you would like to report an issue with this script, you can file one on Github (https://github.com/GeoscienceAustralia/dea-notebooks/issues/new). -Last modified: September 2023 +Last modified: November 2023 """ @@ -712,6 +712,81 @@ def model_tides( return tide_df +def _pixel_tides_resample( + tides_lowres, + ds, + resample_method="bilinear", + dask_chunks="auto", + dask_compute=True, +): + """ + Resamples low resolution tides modelled by `pixel_tides` into the + geobox (e.g. spatial resolution and extent) of the original higher + resolution satellite dataset. + + Parameters: + ----------- + tides_lowres : xarray.DataArray + The low resolution tide modelling data array to be resampled. + ds : xarray.Dataset + The dataset whose geobox will be used as the template for the + resampling operation. This is typically the same satellite + dataset originally passed to `pixel_tides`. + resample_method : string, optional + The resampling method to use. Defaults to "bilinear"; valid + options include "nearest", "cubic", "min", "max", "average" etc. + dask_chunks : str or tuple, optional + Can be used to configure custom Dask chunking for the final + resampling step. The default of "auto" will automatically set + x/y chunks to match those in `ds` if they exist, otherwise will + set x/y chunks that cover the entire extent of the dataset. + For custom chunks, provide a tuple in the form `(y, x)`, e.g. + `(2048, 2048)`. + dask_compute : bool, optional + Whether to compute results of the resampling step using Dask. + If False, this will return `tides_highres` as a Dask array. + + Returns: + -------- + tides_highres, tides_lowres : tuple of xr.DataArrays + In addition to `tides_lowres` (see above), a high resolution + array of tide heights will be generated matching the + exact spatial resolution and extent of `ds`. + """ + # Determine spatial dimensions + y_dim, x_dim = ds.odc.spatial_dims + + # Convert array to Dask, using no chunking along y and x dims, + # and a single chunk for each timestep/quantile and tide model + tides_lowres_dask = tides_lowres.chunk( + {d: None if d in [y_dim, x_dim] else 1 for d in tides_lowres.dims} + ) + + # Automatically set Dask chunks for reprojection if set to "auto". + # This will either use x/y chunks if they exist in `ds`, else + # will cover the entire x and y dims) so we don't end up with + # hundreds of tiny x and y chunks due to the small size of + # `tides_lowres` (possible odc.geo bug?) + if dask_chunks == "auto": + if (y_dim in ds.chunks) & (x_dim in ds.chunks): + dask_chunks = (ds.chunks[y_dim], ds.chunks[x_dim]) + else: + dask_chunks = ds.odc.geobox.shape + + # Reproject into the GeoBox of `ds` using odc.geo and Dask + tides_highres = tides_lowres_dask.odc.reproject( + how=ds.odc.geobox, + chunks=dask_chunks, + resampling=resample_method, + ).rename("tide_m") + + # Optionally process and load into memory with Dask + if dask_compute: + tides_highres.load() + + return tides_highres, tides_lowres + + def pixel_tides( ds, times=None, @@ -985,34 +1060,9 @@ def pixel_tides( # Reproject into original high resolution grid if resample: print("Reprojecting tides into original array") - # Convert array to Dask, using no chunking along y and x dims, - # and a single chunk for each timestep/quantile and tide model - tides_lowres_dask = tides_lowres.chunk( - {d: None if d in [y_dim, x_dim] else 1 for d in tides_lowres.dims} + tides_highres, tides_lowres = _pixel_tides_resample( + tides_lowres, ds, resample_method, dask_chunks, dask_compute ) - - # Automatically set Dask chunks for reprojection if set to "auto". - # This will either use x/y chunks if they exist in `ds`, else - # will cover the entire x and y dims) so we don't end up with - # hundreds of tiny x and y chunks due to the small size of - # `tides_lowres` (possible odc.geo bug?) - if dask_chunks == "auto": - if (y_dim in ds.chunks) & (x_dim in ds.chunks): - dask_chunks = (ds.chunks[y_dim], ds.chunks[x_dim]) - else: - dask_chunks = ds.odc.geobox.shape - - # Reproject into the GeoBox of `ds` using odc.geo and Dask - tides_highres = tides_lowres_dask.odc.reproject( - how=ds.odc.geobox, - chunks=dask_chunks, - resampling=resample_method, - ).rename("tide_m") - - # Optionally process and load into memory with Dask - if dask_compute: - tides_highres.load() - return tides_highres, tides_lowres else: