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

@@ -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: