Repository: British-Oceanographic-Data-Centre/COAsT (language: python)

Issue:
Add new module with functionality for seasonal decomposition of time series
- Add functionality to decompose a time series into trend, seasonal and residual components.
- Add module to house related functionality.
Create dask lazy evaluation / parallelism wrapper for season decomposition functionality 
- for issue #404 a statsmodel package function will be used to decompose time series, however the function works on numpy eager evaluation and so it will be useful to create a dask lazy wrapper for it that 

1. make the function evaluate lazily
2. allows for parallel execution
3. returns xarray DataArrays based on the input DataArray, rather than returning numpy arrays.

- The wrapper will be a useful example/template for other similar functionality that can benefit from this treatment.


Discussion / hints:

Fix patch:
diff --git a/coast/__init__.py b/coast/__init__.py
index f1993d38..4e428bbc 100644
--- a/coast/__init__.py
+++ b/coast/__init__.py
@@ -24,6 +24,7 @@
 from .tidegauge_analysis import TidegaugeAnalysis
 from .config_parser import ConfigParser
 from .xesmf_convert import xesmf_convert
+from .process_data import Process_data
 
 # Set default for logging level when coast is imported
 import logging
diff --git a/coast/process_data.py b/coast/process_data.py
new file mode 100644
index 00000000..ba537bf1
--- /dev/null
+++ b/coast/process_data.py
@@ -0,0 +1,120 @@
+from .gridded import Gridded
+from .logging_util import info, warn, error
+import dask.array as da
+from dask import delayed
+import xarray as xr
+import numpy as np
+
+
+class Process_data:  # TODO All abstract methods should be implemented
+    """
+    A Python class containing methods for processing time series.
+    """
+
+    def __init__(self):
+        return  # TODO Super __init__ should be called at some point
+
+    def seasonal_decomp(self, ts_chunk, **kwargs):
+        """
+        Note this is intended as an inner functions to be called by a wrapper function
+        Process_data.seasonal_decomposition
+
+        This function is itself a wrapper function for statsmodel.seasonal_decompose
+        that accepts multiple timeseries distributed across multiple dimensions
+        as an xr.DataArray: ts_chunk = ts_chunk(t_dim,z_dim,x_dim,y_dim,...).
+        Returns the trend, seasonal and residual componenets of the time series.
+        Invalid points i.e. land points, should be specified as np.nans.
+
+        If called directly, evaluation is eager, i.e. numpy arrays are returned
+        for the trend, seasonal and residual components of the time series
+
+        ts_chunk is the input array or chunk
+        **kwargs are the kwargs for statsmodels.seasonal_decompose
+
+        Calling with dask.delayed:
+        dask.delayed(self.seasonal_decomp, nout=3)(ts_chunk, **kwargs)
+
+        Need statsmodels package installed
+        """
+        from statsmodels.tsa.seasonal import seasonal_decompose
+
+        # Get the indices of land points (points that are nan for all time steps)
+        active_ind = np.where(np.nan_to_num(ts_chunk.data).any(axis=0))[0]
+        # Remove them (statsmodel.seasonal_decompose will not accept them)
+        ts_chunk_valid = ts_chunk[:, active_ind]
+        # Call statsmodels.seasonal_decompose()
+        decomp_chunk = seasonal_decompose(ts_chunk_valid, **kwargs)
+        # Store time series components in seperate arrays and return
+        trend_chunk, seasonal_chunk, residual_chunk = (
+            np.full(ts_chunk.shape, np.nan),
+            np.full(ts_chunk.shape, np.nan),
+            np.full(ts_chunk.shape, np.nan),
+        )
+        trend_chunk[:, active_ind] = decomp_chunk.trend
+        seasonal_chunk[:, active_ind] = decomp_chunk.seasonal
+        residual_chunk[:, active_ind] = decomp_chunk.resid
+        return trend_chunk, seasonal_chunk, residual_chunk
+
+    def seasonal_decomposition(self, time_series: xr.DataArray, num_chunks=1, **kwargs):
+        """
+        Dask delayed wrapper for statsmodel.seasonal_decompose that accepts multiple timeseries
+        distributed across multiple dimensitons as an xr.DataArray:
+        time_series = time_series(t_dim,z_dim,x_dim,y_dim,...)
+
+        Accepts land points (z,x,y) as time_series=np.nan for all time.
+
+        num_chunks allows caller to specify the number of chunks in the delayed call, which enables
+        parallelism. i.e. num_chunks = 4 splits the timeserieses into 4 groups that will be decomposed
+        in parallel by the dask scheduler.
+
+        **kwargs are the kwargs for statsmodels.seasonal_decompose function
+
+        Returns coast.Gridded object containing an xarray.Dataset with the trend, seasonal, residual
+        timeseries components as xarray.DataArrays.
+
+        Need statsmodels package installed
+        """
+        # first dim needs to be time
+        time_series = time_series.transpose("t_dim", ...)
+        # stack all the other dims
+        ts_stacked = time_series.stack(space=time_series.dims[1:])
+        # calculate chunk size for stacked dimension and rechunk
+        chunk_size = ts_stacked["space"].size // num_chunks + 1
+        ts_stacked = ts_stacked.chunk({"t_dim": ts_stacked["t_dim"].size, "space": chunk_size})
+        # convert to list of dask delayed objects (e.g. 4 chunks => 4 delayed objects in list)
+        ts_delayed = ts_stacked.data.to_delayed().ravel()
+
+        # delayed call to statsmodel.seasonal_decomposition (with some other processing) that returns
+        # 3 lists:
+        # 1) contains the dask delayed chunks for trend (e.g. 4 delayed objects in list if 4 chunks),
+        # 2) contains the dask delayed chunks for seasonal,
+        # 3) contains the dask delayed chunks for residual,
+        trend_de, seasonal_de, resid_de = map(
+            list, zip(*[delayed(self.seasonal_decomp, nout=3)(ts_chunk, **kwargs) for ts_chunk in ts_delayed])
+        )
+
+        # convert to lists of dask arrays to allow array operations
+        trend_da, seasonal_da, resid_da = [], [], []
+        for chunk_idx, (tr_chunk, se_chunk, re_chunk) in enumerate(zip(trend_de, seasonal_de, resid_de)):
+            # When converting from delayed to array, you must know the shape of the
+            # array. Here we know this from the chunk sizes of the (stacked) DataArray
+            chunk_shape = (ts_stacked.chunks[0][0], ts_stacked.chunks[1][chunk_idx])
+            trend_da.append(da.from_delayed(tr_chunk, shape=chunk_shape, dtype=float))
+            seasonal_da.append(da.from_delayed(se_chunk, shape=chunk_shape, dtype=float))
+            resid_da.append(da.from_delayed(re_chunk, shape=chunk_shape, dtype=float))
+
+        # concatenate the array chunks together and reshape back to original shape
+        trend = da.reshape(da.concatenate(trend_da, axis=1), time_series.shape)
+        seasonal = da.reshape(da.concatenate(seasonal_da, axis=1), time_series.shape)
+        residual = da.reshape(da.concatenate(resid_da, axis=1), time_series.shape)
+        # create  gridded object for return. New DataArrays created from the original
+        # timeseries DataArray so it will have the smae coordinates and shape. Rechunk it
+        # to have same structure as the newly created dask arrays and then assign to gridded
+        gd = Gridded()
+        gd.dataset["trend"] = xr.full_like(time_series.chunk(trend.chunks), np.nan)
+        gd.dataset["trend"][:, :, :] = trend
+        gd.dataset["seasonal"] = xr.full_like(time_series.chunk(seasonal.chunks), np.nan)
+        gd.dataset["seasonal"][:, :, :] = seasonal
+        gd.dataset["residual"] = xr.full_like(time_series.chunk(residual.chunks), np.nan)
+        gd.dataset["residual"][:, :, :] = residual
+        return gd
diff --git a/example_scripts/dask_wrapper_TEMPLATE.txt b/example_scripts/dask_wrapper_TEMPLATE.txt
new file mode 100644
index 00000000..95ba2ba6
--- /dev/null
+++ b/example_scripts/dask_wrapper_TEMPLATE.txt
@@ -0,0 +1,58 @@
+# Template and instruction for create a dask wrapper
+
+# Motivation
+# Start with an xarray.DataArray, myDataArray, that we want to pass into a
+# function. That function will perform eager evaluation and return a numpy array,
+# but we want lazy evaluation with the possibility to allow dask parallelism
+# see worked example in Process_data.seasonal_decomposition
+
+# Imports
+import dask.array as da
+from dask import delayed
+import xarray as xr
+import numpy as np
+
+# Step 1. (optional: allows dask to distribute computation across multiple cores, if not interested see comment 2)
+# Partition data in myDataArray by chunking it up as desired. Note that chunking dimensions
+# need to make sense for your particular problem! Here we just chunk along dim_2
+myDataArray = myDataArray.chunk({"dim_1": myDataArray.dim_1.size, "dim_2": chunksize})  # can be more dimensions
+# Then create a list containing all the array chunks as dask.delayed objects
+# (e.g. 4 chunks => list contain 4 delayed objects)
+myDataArray_partitioned = myDataArray.data.to_delayed().ravel()
+
+# Comment 1
+# There are different ways to partition your data. For example, if you start off with a numpy array
+# rather than an xarray DataArray you can just iterate over the array and partition it that way
+# (the partitions do NOT need to be dask.delayed objects). For example see the very simple case here:
+# https://docs.dask.org/en/stable/delayed.html
+# The method described in 1 is just very convenient for DataArrays where the multi-dimensional
+# chunks may be the desired way to partition the data.
+
+# Step 2.
+# Call your eager evaluating function using dask.delayed and pass in your data. This returns
+# a list containing the outputs from the function as dask.delayed objects. The list will have the same length as
+# myDataArray_partitioned
+delayed_myFunction_output = [
+    delayed(myFunction)(aChunk, other_args_for_myFunction) for aChunk in myDataArray_partitioned
+]
+
+# Step 3.
+# convert the lists of delayed objects to lists of dask arrays to allow array operations
+# It's possible this step is not necessary!
+dask_array_list = []
+for chunk_idx, aChunk in enumerate(delayed_myFunction_output):
+    # When converting from dask.delayed to dask.array, you must know the shape of the
+    # array. In this example we know this from the chunk sizes of the original DataArray
+    chunk_shape = (myDataArray.chunks[0][0], myDataArray.chunks[1][chunk_idx])
+    dask_array_list.append(da.from_delayed(aChunk, shape=chunk_shape, dtype=float))
+
+# Step 4.
+# Concatenate the array chunks together to get a single dask.array. This can be assigned to a
+# new DataArray as desired.
+myOutputArray = da.concatenate(dask_array_list, axis=1)
+
+# Comment 2
+# If you skipped step 1., i.e. just want a lazy operation and no parallelism, you can just do this
+myOutputArray = da.from_delayed(
+    delayed(myFunction)(myDataArray, other_args_for_myFunction), shape=myDataArray.shape, dtype=float
+)


Test patch:
diff --git a/unit_testing/test_process_data.py b/unit_testing/test_process_data.py
new file mode 100644
index 00000000..2b1151a6
--- /dev/null
+++ b/unit_testing/test_process_data.py
@@ -0,0 +1,54 @@
+# IMPORT modules. Must have unittest, and probably coast.
+import coast
+import unittest
+import numpy as np
+import xarray as xr
+import unit_test_files as files
+
+
+class test_process_data_methods(unittest.TestCase):
+    def test_statsmodel_seasonal_decompose(self):
+        gd_t = coast.Gridded(files.fn_nemo_grid_t_dat, files.fn_nemo_dom, config=files.fn_config_t_grid)
+        fake_time = np.arange(
+            np.datetime64("2010-01-01"), np.datetime64("2014-01-01"), np.timedelta64(1, "M"), dtype="datetime64[M]"
+        ).astype("datetime64[s]")
+        fake_temp = (
+            (np.arange(0, 48) * 0.05)[:, np.newaxis, np.newaxis, np.newaxis]
+            + np.random.normal(0, 0.1, 48)[:, np.newaxis, np.newaxis, np.newaxis]
+            + np.tile(gd_t.dataset.temperature[:-1, :2, :, :], (8, 1, 1, 1))
+        )
+        fake_data_array = xr.DataArray(
+            fake_temp,
+            coords={
+                "t_dim": fake_time,
+                "depth_0": gd_t.dataset.depth_0[:2, :, :],
+                "longitude": gd_t.dataset.longitude,
+                "latitude": gd_t.dataset.latitude,
+            },
+            dims=["t_dim", "z_dim", "y_dim", "x_dim"],
+        )
+
+        # Coast version
+        proc_data = coast.Process_data()
+        grd = proc_data.seasonal_decomposition(fake_data_array, 4, model="additive", period=6, extrapolate_trend="freq")
+        # statsmodel version
+        from statsmodels.tsa.seasonal import seasonal_decompose
+
+        stm_result = seasonal_decompose(
+            fake_data_array.fillna(0).values.reshape((48, 2 * 375 * 297)),
+            model="additive",
+            period=6,
+            extrapolate_trend="freq",
+        )
+
+        check_trend = np.allclose(stm_result.trend.reshape((48, 2, 375, 297)), grd.dataset.trend.compute().fillna(0))
+        check_seasonal = np.allclose(
+            stm_result.seasonal.reshape((48, 2, 375, 297)), grd.dataset.seasonal.compute().fillna(0)
+        )
+        check_residual = np.allclose(
+            stm_result.resid.reshape((48, 2, 375, 297)), grd.dataset.residual.compute().fillna(0)
+        )
+
+        self.assertTrue(check_trend, "trends don't match")
+        self.assertTrue(check_seasonal, "seasonals don't match")
+        self.assertTrue(check_residual, "residuals don't match")
diff --git a/unit_testing/unit_test.py b/unit_testing/unit_test.py
index d6e0c88e..08cd6f4c 100644
--- a/unit_testing/unit_test.py
+++ b/unit_testing/unit_test.py
@@ -31,6 +31,7 @@
 from test_climatology import test_climatology
 from test_wod_read_data import test_wod_read_data
 from test_example_scripts import test_example_scripts
+from test_process_data import test_process_data_methods
 from test_bgc_gridded_initialisation import test_bgc_gridded_initialisation
 
 
@@ -57,6 +58,7 @@
     test_climatology,
     test_wod_read_data,
     test_bgc_gridded_initialisation,
+    test_process_data_methods,
     test_example_scripts,
 ]
 
diff --git a/unit_testing/unit_test_contents.txt b/unit_testing/unit_test_contents.txt
index b13c45dd..aa4c9471 100755
--- a/unit_testing/unit_test_contents.txt
+++ b/unit_testing/unit_test_contents.txt
@@ -18,12 +18,12 @@
       d. day_of_week
       e. getitem
 
-3. test_gridded_harmonics
-      a. combine_and_convert_harmonics
-
-4. test_xesmf_convert
+3. test_xesmf_convert
       a. basic_conversion_to_xesmf
 
+4. test_gridded_harmonics
+      a. combine_and_convert_harmonics
+
 5. test_diagnostic_methods
       a. compute_vertical_spatial_derivative
       b. construct_density
@@ -64,59 +64,62 @@
       a. compute_heofs
       b. compute_regular_eof
 
-11. test_contour_f_methods
+12. test_contour_f_methods
       a. calculate_flow_across_contour
       b. calculate_pressure_gradient_driven_flow
       c. extract_isobath_contour_between_two_points
 
-12. test_contour_t_methods
+13. test_contour_t_methods
       a. along_contour_2d_flow
       b. along_contour_flow
       c. calculate_pressure_along_contour
 
-13. test_profile_methods
+14. test_profile_methods
       a. compare_processed_profile_with_model
       b. load_process_and_compare_profile_data
 
-14. test_plot_utilities
+15. test_plot_utilities
       a. determine_clim_by_stdev
       b. determine_colorbar_extension
       c. geo_axes
       d. scatter_with_fit
 
-15. test_stats_utilities
+16. test_stats_utilities
       a. find_maxima
 
-16. test_maskmaker_methods
+17. test_maskmaker_methods
       a. fill_polygon_by_index
       b. fill_polygon_by_lonlat
 
-17. test_climatology
+18. test_climatology
       a. monthly_and_seasonal_climatology
 
-18. test_wod_read_data
+19. test_wod_read_data
       a. load_wod
       b. reshape_wod
 
-19. test_gridded_initialisation
+20. test_bgc_gridded_initialisation
       a. gridded_load_bgc_data
-      b. gridded_load_bgc_variables_correctly_renamed
+      b. gridded_load_bgc_data_and_domain
       c. gridded_load_bgc_dimensions_correctly_renamed
       d. gridded_load_bgc_domain_only
+      e. gridded_load_bgc_variables_correctly_renamed
+
+21. test_process_data_methods
+      a. statsmodel_seasonal_decompose
 
-20. test_example_scripts
+22. test_example_scripts
       a. altimetry_tutorial
       b. amm15_example_plot
       c. anchor_plots_of_nsea_wvel
       d. blz_example_plot
       e. contour_tutorial
-      f. wod_bgc_ragged_example
-      g. export_to_netcdf_tutorial
-      h. internal_tide_pycnocline_diagnostics
+      f. export_to_netcdf_tutorial
+      g. internal_tide_pycnocline_diagnostics
+      h. seasia_dic_example_plot
       i. seasia_r12_example_plot
-      j. tidegauge_tutorial
-      k. tidetable_tutorial
-      l. transect_tutorial
-      m. wcssp_india_example_plot
-      n. wod_bgc_ragged_example
-      o. seasia_dic_example_plot
+      j. tidetable_tutorial
+      k. transect_tutorial
+      l. wcssp_india_example_plot
+      m. wod_bgc_ragged_example
+