# Author: An-Chi Ho
# Date: 10th September 2021
# ------------------------------------------------------------------
# This use case provides an example of transforming and chunking latitude and longitude 
# dimensions. If all other dimensions are used as target dimensions in the operation, 
# it would be good to have the option of chunking the spatial dimensions. However, the 
# errors due to transformation at borders may increase because chunking produces more 
# borders. There are many factors may impact the results of transformation or 
# transformation + chunking. See FAQ How-to-24 for related information(https://earth.bsc.es/gitlab/es/startR/-/blob/master/inst/doc/faq.md#24-do-both-interpolation-and-chunking-on-spatial-dimensions) 
# Here we provide some scripts that are more common and less error-prone.
# Common things to notice:
# 1. 'transform_extra_cells' is increased to 8 to avoid errors at borders. 
#    Fewer cells may be enough, depending on cases.
# 2. The 'crop' argument in 'transform_params' is defined by the borders of the region or FALSE.
#    TRUE may return wrong values, depending on cases.
# 3. CircularSort() is required to use even if the longitude fully falls in the range because
#    it tells startR that longitude dimension is circular and the extra cells should be got from
#    the other side if the border is reached. In the scripts below, CircularSort(0, 360) is used,
#    but it can also be replaced by CircularSort(-180, 180).
# ------------------------------------------------------------------

library(startR)

lons.min <- 0
lons.max <- 359.9
lats.min <- -90
lats.max <- 90
sdates <- paste0(1981:2011, '0101')

#---------------------------------
# Method 1:
# - Use list of 2 values to define longitude and latitude.
#---------------------------------

exp <- Start(dat = '/esarchive/exp/ecmwf/system5_m1/monthly_mean/$var$_f6h/$var$_$sdate$.nc',
             var = 'tas',
             sdate = sdates,
             ensemble = 'all',
             time = indices(1),
             latitude = values(list(lats.min, lats.max)),
             latitude_reorder = Sort(),
             longitude = values(list(lons.min, lons.max)),
             longitude_reorder = CircularSort(0, 360),
             transform = CDORemapper,
             transform_params = list(grid = 'r100x50', 
                                     method = 'con'),
             transform_vars = c('latitude', 'longitude'),
             transform_extra_cells = 8,
             synonims = list(latitude = c('latitude', 'lat'),
                             longitude = c('longitude', 'lon')),
             return_vars = list(latitude = NULL,
                                longitude = NULL,
                                time = 'sdate'),
             retrieve = F)

func <- function(x) {
  # x: [sdate, ensemble]
  ensemble_mean <- s2dv::MeanDims(x, 2)
  trend <- s2dv:::.Trend(ensemble_mean)$trend
  return(trend)
}
step <- Step(func,
             target_dims = c('sdate', 'ensemble'), output_dims = 'stats')
wf <- AddStep(exp, step)

res1 <- Compute(wf, chunks = list(latitude = 2, longitude = 2))$output1

dim(res1)
#    stats       dat       var      time  latitude longitude 
#        2         1         1         1        50       100 
summary(res1)
#     Min.   1st Qu.    Median      Mean   3rd Qu.      Max. 
# -0.09884   0.00879 117.19019 138.18788 279.98564 305.25010 



#---------------------------------
# Method 2:
# - Use vector of indices to define longitude and latitude.
# - The 'crop' argument in 'transform_params' is FALSE (but it can also be defined by the borders
#   of the region.)
#---------------------------------

exp <- Start(dat = '/esarchive/exp/ecmwf/system5_m1/monthly_mean/$var$_f6h/$var$_$sdate$.nc',
             var = 'tas',
             sdate = sdates,
             ensemble = 'all',
             time = indices(1),
             latitude = indices(1:640),
             latitude_reorder = Sort(),
             longitude = indices(1:1296),
             longitude_reorder = CircularSort(0, 360),
             transform = CDORemapper,
             transform_params = list(grid = 'r100x50', 
                                     method = 'con'),
             transform_vars = c('latitude', 'longitude'),
             transform_extra_cells = 8,
             synonims = list(latitude = c('latitude', 'lat'),
                             longitude = c('longitude', 'lon')),
             return_vars = list(latitude = NULL,
                                longitude = NULL,
                                time = 'sdate'),
             retrieve = F)

func <- function(x) {
  # x: [sdate, ensemble]
  ensemble_mean <- s2dv::MeanDims(x, 2)
  trend <- s2dv:::.Trend(ensemble_mean)$trend
  return(trend)
}
step <- Step(func,
             target_dims = c('sdate', 'ensemble'), output_dims = 'stats')
wf <- AddStep(exp, step)

res2 <- Compute(wf, chunks = list(latitude = 2, longitude = 2))$output1

identical(res1, res2)
#[1] TRUE

#---------------------------------
# Method 3:
# - Use 'all' to define longitude and latitude.
# - The 'crop' argument in 'transform_params' is FALSE (but it can also be defined by the borders
#   of the region.)
#---------------------------------
exp <- Start(dat = '/esarchive/exp/ecmwf/system5_m1/monthly_mean/$var$_f6h/$var$_$sdate$.nc',
             var = 'tas',
             sdate = sdates,
             ensemble = 'all',
             time = indices(1),
             latitude = 'all',
             latitude_reorder = Sort(),
             longitude = 'all',
             longitude_reorder = CircularSort(0, 360),
             transform = CDORemapper,
             transform_params = list(grid = 'r100x50', 
                                     method = 'con'),
             transform_vars = c('latitude', 'longitude'),
             transform_extra_cells = 8,
             synonims = list(latitude = c('latitude', 'lat'),
                             longitude = c('longitude', 'lon')),
             return_vars = list(latitude = NULL,
                                longitude = NULL,
                                time = 'sdate'),
             retrieve = F)

func <- function(x) {
  # x: [sdate, ensemble]
  ensemble_mean <- s2dv::MeanDims(x, 2)
  trend <- s2dv:::.Trend(ensemble_mean)$trend
  return(trend)
}
step <- Step(func,
             target_dims = c('sdate', 'ensemble'), output_dims = 'stats')
wf <- AddStep(exp, step)

res3 <- Compute(wf, chunks = list(latitude = 2, longitude = 2))$output1

identical(res2, res3)
#[1] TRUE
