compute_nino.R 21.8 KB
Newer Older
compute_nino <- function(data, recipe, region, standardised = TRUE,
                         running_mean = NULL, plot_ts = TRUE, plot_sp = TRUE,
                         alpha = 0.5, save = 'all', na.rm = TRUE, logo = NULL,
                         ncores = NULL) {
  if (!is.null(data$fcst)) {
    warn(recipe$Run$logger, "Nino computed only for hindcast data.")
  }
  var <- recipe$Analysis$Variables$name
  if (!(var %in% c('tos', 'sst'))) {
    warn(recipe$Run$logger, "Variable name is not one of the expected sst or tos")
  }
  var_units <- data$hcst$attrs$Variable$metadata[[var]]$units
  nino_hcst <- WeightedMean(data = data$hcst$data, 
                       lon = as.vector(data$hcst$coords$longitude),
                       lat = as.vector(data$hcst$coords$latitude),
                       region = region,
                       londim = 'longitude',
                       latdim = 'latitude',
                       na.rm = na.rm,
                       ncores = ncores)
  nino_obs <- WeightedMean(data = data$obs$data,
                       lon = as.vector(data$hcst$coords$longitude),
                       lat = as.vector(data$hcst$coords$latitude),
                       region = region,
                       londim = 'longitude',
                       latdim = 'latitude',
                       na.rm = na.rm,
                       ncores = ncores)
  if (standardised) {
  nino_hcst <- Apply(list(nino_hcst), target_dims = c('syear', 'ensemble'),
                   fun = function(x) {
                          sd <- sqrt(var(as.vector(x), na.rm = TRUE))
                          means <- mean(as.vector(x), na.rm = TRUE)
                          res <- apply(x, c(1,2), function(x) {(x-means)/sd})},
                   ncores = recipe$Analysis$ncores)$output1
  nino_obs <- Apply(list(nino_obs), target_dims = c('syear', 'ensemble'),
                   fun = function(x) {
                          sd <- sqrt(var(as.vector(x), na.rm = TRUE))
                          means <- mean(as.vector(x), na.rm = TRUE)
                          res <- apply(x, c(1,2), function(x) {(x-means)/sd})},
                   ncores = recipe$Analysis$ncores)$output1
    var_units <- 'adim'
  }
  if (!is.null(running_mean)) {
    nino_hcst <- Smoothing(nino_hcst, 
                           runmeanlen = running_mean,
                           time_dim = 'time',
                           ncores = ncores)
    nino_obs <- Smoothing(nino_obs,
                          runmeanlen = running_mean,
                          time_dim = 'time',
                          ncores = ncores)
  }
  if (all(region == c(-90, -80, -10, 0))) {
    region_name <- "1+2"
    nino_name <- "nino12"
  } else if (all(region == c(-150, -90, -5, 5))) {
    region_name <- "3"
    nino_name <- "nino3"
  } else if (all(region == c(-170, -120, -5, 5))) {
    region_name <- "3.4"
    nino_name <- "nino34"
  } else if (all(region == c(160, -150, -5, 5))) {
    region_name <- "4"
    nino_name <- "nino4"
  } else {
    stop("Unknown nino region")
  }
  nino_hcst <- InsertDim(nino_hcst, posdim = 1, lendim = 1, name = 'region')
  nino_obs <- InsertDim(nino_obs, posdim = 1, lendim = 1, name = 'region')
  dims_dates_not_null <- dim(data$hcst$attrs$Dates)[which(dim(data$hcst$attrs$Dates) > 1)]
  hcst_dates <- Subset(data$hcst$attrs$Dates, along = names(dims_dates_not_null),
                       indices = lapply(dims_dates_not_null, function(x){1:x}),
                       drop = "non-selected")	
   if (!("time" %in% names(dim(hcst_dates)))) {
       hcst_dates <- InsertDim(hcst_dates, pos = 1, len = 1, name = 'time')
       ## TODO: recover dates format
       hcst_dates <- as.POSIXct(hcst_dates, origin = '1970-01-01', tz = 'UTC')
    }
  nino <- list(hcst = s2dv_cube(
                      data = nino_hcst,
                      varName = nino_name,
                      metadata = list(
                      region = list(name = paste("Nino", region_name, "region"),
                                    lats_range = paste(region[3:4]),
                                    lons_range = paste(region[1:2])),
                      time = data$hcst$attrs$Variable$metadata$time,
                      nino = list(units = var_units,
                             longname = paste("El Niño", region_name, "Index"))),
                      Dates = hcst_dates,
                      Dataset = recipe$Analysis$Datasets$System$name),
               obs = s2dv_cube(
                     data = nino_obs,
                     varName = nino_name,
                     metadata = list(
                     region = list(name = paste("Nino", region_name, "region"),
                                   lats_range = paste(region[3:4]),
                                   lons_range = paste(region[1:2])),
                     time = data$obs$attrs$Variable$metadata$time,
                     nino = list(units = var_units,
                            longname = paste("El Niño", region_name, "Index"))),
                     Dates = data$obs$attrs$Dates,
                     Dataset = recipe$Analysis$Datasets$Reference$name))

vagudets's avatar
vagudets committed
  if (save == 'all') {
    file_dest <- paste0(recipe$Run$output_dir, "/outputs/Indices/")
    # Use startdates param from SaveExp to correctly name the files:
    if (tolower(recipe$Analysis$Horizon) == "seasonal") { 
      if (length(data$hcst$attrs$source_files) == dim(data$hcst$data)['syear']) {
        file_dates <- Apply(data$hcst$attrs$source_files, target_dim = NULL,
                            fun = function(x) {
                            pos <- which(strsplit(x, "")[[1]] == ".")
                            res <- substr(x, pos-8, pos-1)
                            })$output1
      }
    } else if (tolower(recipe$Analysis$Horizon) == "decadal"){
      file_dates <- paste0('s',recipe$Analysis$Time$hcst_start : recipe$Analysis$Time$hcst_end)
    }
    nino$hcst$data <- .drop_indices_dims(nino_hcst)
    CST_SaveExp(data = nino$hcst,
                destination = file_dest,
                startdates = as.vector(file_dates),
                dat_dim = NULL, sdate_dim = 'syear',
                ftime_dim = 'time', var_dim = NULL,
                memb_dim = 'ensemble')
    res <- .drop_indices_dims(nino_obs)
    if (!("time" %in% names(dim(res)))) {
      res <- InsertDim(res, pos = 1, len = 1, name = 'time')
    }
    nino$obs$data <- res
    CST_SaveExp(data = nino$obs,
                destination = file_dest,
                startdates = as.vector(file_dates),
                dat_dim = NULL, sdate_dim = 'syear',
                ftime_dim = 'time', var_dim = NULL,
                memb_dim = NULL)
    nino$hcst$data <- nino_hcst
    nino$obs$data <- nino_obs
    res <- NULL
    gc()
  }
 # Read variable long_name to plot it
  conf <- yaml::read_yaml("conf/variable-dictionary.yml")
  var_name <- conf$vars[[which(names(conf$vars) ==
                               recipe$Analysis$Variables$name)]]$long
  if (plot_ts) {
    dir.create(paste0(recipe$Run$output_dir, "/plots/Indices/"),
               showWarnings = F, recursive = T)
    source("modules/Indices/R/plot_deterministic_forecast.R")
    for (tstep in 1:dim(nino$hcst$data)['time']) {
      mes <- as.numeric(substr(recipe$Analysis$Time$sdate, 1,2)) +
               (tstep - 1) + (recipe$Analysis$Time$ftime_min - 1)
      mes <- ifelse(mes > 12, mes - 12, mes)
      fmonth <- sprintf("%02d", tstep - 1 + recipe$Analysis$Time$ftime_min)
      obs <- Subset(nino$obs$data, along = 'time', ind = tstep)
      exp <- Subset(nino$hcst$data, along = 'time', ind = tstep)
      if (gsub(".", "", recipe$Analysis$Datasets$System$name) == "") {
        system <- recipe$Analysis$Datasets$System$name
      } else {
        system <- gsub(".", "", recipe$Analysis$Datasets$System$name)
      }
      if (tolower(recipe$Analysis$Horizon) == "seasonal") {
        toptitle <- paste("Ni\u00F1o", region_name, "SST Index\n",
                          month.abb[mes],
                          "/", recipe$Analysis$Time$hcst_start, "-",
                          recipe$Analysis$Time$hcst_end)
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/",
                           nino_name, "_",
                           system, "_", recipe$Analysis$Datasets$Reference$name,
                           "_s", recipe$Analysis$Time$sdate, "_ftime",
                           sprintf("%02d", tstep - 1 + recipe$Analysis$Time$ftime_min), 
                           ".png")
        caption <-  paste0("Nominal start date: 1st of ",
                           month.name[as.numeric(substr(recipe$Analysis$Time$sdate, 1,2))],
                           "\n",
                           "Forecast month: ", fmonth)
        xlabs = as.numeric(substr(file_dates, 1, 4))
      } else if (tolower(recipe$Analysis$Horizon) == "decadal") {
        toptitle <- paste("Ni\u00F1o", region_name, "SST Index\n",
                          "Lead time", fmonth,
                          "/", recipe$Analysis$Time$hcst_start, "-",
                          recipe$Analysis$Time$hcst_end)
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/", nino_name, "_",
                           system, "_", recipe$Analysis$Datasets$Reference$name,
                           "_ftime",
                           sprintf("%02d", tstep - 1 + recipe$Analysis$Time$ftime_min), ".png")
       caption <- paste0("Start date month: ",
                          month.name[get_archive(recipe)$System[[recipe$Analysis$Datasets$System$name]]$initial_month],
                          "\n",
                          "Lead time: ", fmonth, "\n")
        xlabs <- substr(file_dates, 2,5)
      } else {
        toptitle <- NULL
        warning("The plot title is not defined for this time horizon. ",
                "The plots will not have a title.")
      }
      plot_deterministic_forecast(obs, exp,
                                  time_dim = 'syear',
                                  member_dim = 'ensemble', style = 'boxplot',
                                  xlabs = xlabs,
                                  ylab = var_units, 
                                  title = toptitle, fileout = plotfile,
                                  caption =  caption,
                                  legend_text = c(
                                  recipe$Analysis$Datasets$Reference$name,
                                  recipe$Analysis$Datasets$System$name))
    }
  }
  if (plot_sp) {
    source("modules/Visualization/R/tmp/PlotRobinson.R")
    source("modules/Indices/R/correlation_eno.R")
    source("modules/Visualization/R/get_proj_code.R")
    lons <- data$hcst$coords$longitude
    lats <- data$hcst$coords$latitude
    # Get code for Robinson projection depending on GEOS/GDAL/PROJ version
    projection_code <- get_proj_code("robinson")
    # Avoid rewriten longitudes a shift is neeced:
    dir.create(paste0(recipe$Run$output_dir, "/plots/Indices/"),
               showWarnings = F, recursive = T)
    correl_obs <- Apply(list(data$obs$data, nino$obs$data),
                        target_dims = 'syear',
                        fun = .correlation_eno,
                        time_dim = 'syear', method = 'pearson', alpha = alpha,
                        test.type = 'two-sided', pval = FALSE,
                        ncores = recipe$Analysis$ncores)
    correl_hcst <- Apply(list(data$hcst$data, nino$hcst$data),
                         target_dims = c('syear', 'ensemble'),
                         fun = function(x, y) {
                         x <- apply(x, 1, mean, na.rm = TRUE)
                         y <- apply(y, 1, mean, na.rm = TRUE)
                         dim(y) <- c(syear = length(y))
                         dim(x) <- c(syear = length(x))
                         res <- .correlation_eno(x, y,
                           time_dim = 'syear', method = 'pearson', alpha = alpha,
                           test.type = 'two-sided', pval = FALSE)},
                           ncores = recipe$Analysis$ncores)
    correl_hcst_full <- Apply(list(data$hcst$data, nino$hcst$data),
                              target_dims = c('syear', 'ensemble'),
                              fun = function(x,y) {
                              dim(y) <- c(syear = length(y))
                              dim(x) <- c(syear = length(x))
                           res <- .correlation_eno(x, y,
                              time_dim = 'syear', method = 'pearson', alpha = alpha,
                              test.type = 'two-sided', pval = FALSE)},
                              ncores = recipe$Analysis$ncores)
    for (tstep in 1:dim(nino$obs$data)['time']) {
      map <- Subset(correl_obs$r, along = 'time', ind = tstep, drop = T)
      sig <- Subset(correl_obs$sig, along = 'time', ind = tstep, drop = T)
      if (tolower(recipe$Analysis$Horizon) == "seasonal") {
        mes <- as.numeric(substr(recipe$Analysis$Time$sdate, 1,2)) +
               (tstep - 1) + (recipe$Analysis$Time$ftime_min - 1)
        mes <- ifelse(mes > 12, mes - 12, mes)
        fmonth <- sprintf("%02d", tstep - 1 + recipe$Analysis$Time$ftime_min)
        toptitle <- paste(recipe$Analysis$Datasets$Reference$name, "\n",
                          "Ni\u00F1o", region_name, "SST Index -",var_name, "\n",
                          " Correlation /",
                          month.abb[mes],
                          "/", recipe$Analysis$Time$hcst_start, "-",
                          recipe$Analysis$Time$hcst_end)
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/", nino_name,
                            "_correlation_",
                           recipe$Analysis$Variable$name, "_",
                           recipe$Analysis$Datasets$Reference$name,
                           "_s", recipe$Analysis$Time$sdate,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("Nominal start date: 1st of ",
                month.name[as.numeric(substr(recipe$Analysis$Time$sdate, 1,2))],
                         "\n",
                         "Forecast month: ", fmonth, "\n",
                         "Pearson correlation ; alpha = ", alpha)
      } else if (tolower(recipe$Analysis$Horizon) == "decadal") {
        fmonth <- sprintf("%02d", tstep - 1 + recipe$Analysis$Time$ftime_min)
        toptitle <- paste(recipe$Analysis$Datasets$Reference$name, "\n",
                          "Ni\u00F1o", region_name, "SST Index -",var_name, "\n",
                          "Correlation / Start dates",
                          recipe$Analysis$Time$hcst_start, "-",
                          recipe$Analysis$Time$hcst_end)
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/",
                           nino_name, "_correlation_",
                           recipe$Analysis$Variable$name, "_",
                           recipe$Analysis$Datasets$Reference$name,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("Start date: month ",
                           month.name[get_archive(recipe)$System[[recipe$Analysis$Datasets$System$name]]$initial_month],
                           "\n",
                           "Forecast month: ", fmonth, "\n",
                           "Pearson correlation; alpha = ", alpha)     
      } else {
        toptitle <- NULL
        warning("The plot title is not defined for this time horizon. ",
                "The plots will not have a title.")
      }
      
      PlotRobinson(map, lon = lons, lat = lats, 
                   target_proj = projection_code, #"ESRI:54030",
                   lat_dim = 'latitude', lon_dim = 'longitude',
                   legend = 's2dv', style = 'point',
                   toptitle = toptitle, bar_extra_margin = c(4,0,4,0), 
                   caption = caption, mask = sig,
                   fileout = plotfile, width = 8, height = 6,
                   brks = seq(-1, 1, 0.2), cols = brewer.pal(10, 'PuOr'))
 
      ## Ensemble-mean
      map <- Subset(correl_hcst$r, along = 'time', ind = tstep)
      sig <- Subset(correl_hcst$sig, along = 'time', ind = tstep)
      if (gsub(".", "", recipe$Analysis$Datasets$System$name) == "") {
        system <- recipe$Analysis$Datasets$System$name
      } else {
        system <-gsub(".", "", recipe$Analysis$Datasets$System$name)
      }
      if (tolower(recipe$Analysis$Horizon) == "seasonal") {
       toptitle <- paste(recipe$Analysis$Datasets$System$name, "\n",
                          "Ni\u00F1o", region_name, "SST Index -",var_name, "\n",
                          "Correlation /",
                          month.abb[as.numeric(fmonth)],
                          "/", recipe$Analysis$Time$hcst_start, "-",
                          recipe$Analysis$Time$hcst_end)
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/",
                           nino_name, "_correlation_",
                           recipe$Analysis$Variable$name, "_ensmean_",
                           system,
                           "_s", recipe$Analysis$Time$sdate,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("Ensemble mean\n",
                           "Nominal start date: 1st of ",
                month.name[as.numeric(substr(recipe$Analysis$Time$sdate, 1,2))],
                           "\n",
                           "Forecast month: ", fmonth, "\n",
                         "Pearson correlation ; alpha = ", alpha)
      } else if (tolower(recipe$Analysis$Horizon) == "decadal"){
        toptitle <- paste(recipe$Analysis$Datasets$System$name, "\n",
                          "Ni\u00F1o", region_name, "SST Index -",var_name, "\n",
                          "Correlation / Start dates",
                          recipe$Analysis$Time$hcst_start, "-",
                          recipe$Analysis$Time$hcst_end)
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/",
                           nino_name, "_correlation_",
                           recipe$Analysis$Variable$name, "_ensmean_",
                           system,
                           recipe$Analysis$Datasets$Reference$name,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("Correlation ensemble mean\n",
                           "Start date month: ",
                           month.name[get_archive(recipe)$System[[recipe$Analysis$Datasets$System$name]]$initial_month],
                           "\n",
                           "Forecast month: ", fmonth, "\n",
                           "Pearson correlation; alpha = ", alpha)
      } else {
        toptitle <- NULL
        warning("The plot title is not defined for this time horizon. ",
                "The plots will not have a title.")
      }
      PlotRobinson(map, lon = lons, lat = lats,
                   target_proj = projection_code, #"ESRI:54030",
                   lat_dim = 'latitude', lon_dim = 'longitude',
                   legend = 's2dv', style = 'point',
                   toptitle = toptitle, bar_extra_margin = c(4,0,4,0),
                   caption = caption, mask = sig,
                   fileout = plotfile, width = 8, height = 6,
                   brks = seq(-1, 1, 0.2), cols = brewer.pal(10, 'PuOr'))

     # Full hcst corr
      map <- Subset(correl_hcst_full$r, along = 'time', ind = tstep)
      sig <- Subset(correl_hcst_full$sig, along = 'time', ind = tstep)
      if (tolower(recipe$Analysis$Horizon) == "seasonal") {
        toptitle <- paste(recipe$Analysis$Datasets$System$name, "\n",
                          "Ni\u00F1o", region_name, "SST Index -",var_name, "\n",
                          " Correlation /",
                          month.abb[as.numeric(fmonth)],
                          "/", recipe$Analysis$Time$hcst_start, "-",
                          recipe$Analysis$Time$hcst_end)
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/",
                           nino_name, "_correlation_",
                           recipe$Analysis$Variable$name, "_member_",
                           system,
                           "_s", recipe$Analysis$Time$sdate,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("Individual members\n",
                           "Nominal start date: 1st of ",
                month.name[as.numeric(substr(recipe$Analysis$Time$sdate, 1,2))],
                         "\n",
                         "Forecast month: ", fmonth, "\n",
                         "Pearson correlation ; alpha = ", alpha)
    } else if (tolower(recipe$Analysis$Horizon) == "decadal"){
      toptitle <- paste(recipe$Analysis$Datasets$System$name, "\n",
                        "Ni\u00F1o", region_name, "SST Index -",var_name, "\n",
                        "Correlation / Start dates",
                        recipe$Analysis$Time$hcst_start, "-",
                        recipe$Analysis$Time$hcst_end)
      plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/",
                         nino_name, "_correlation_",
                         recipe$Analysis$Variable$name, "_ensmean_",
                         system,
                         recipe$Analysis$Datasets$Reference$name,
                         "_ftime", fmonth, ".png")
      caption <-  paste0("Correlation ensemble mean\n",
                         "Start date month: ",
                         month.name[get_archive(recipe)$System[[recipe$Analysis$Datasets$System$name]]$initial_month],
                         "\n",
                         "Forecast month: ", fmonth, "\n",
                         "Pearson correlation; alpha = ", alpha)
    } else {
      toptitle <- NULL
      warning("The plot title is not defined for this time horizon. ",
              "The plots will not have a title.")
    }
    PlotRobinson(map, lon = lons, lat = lats,
                 target_proj = projection_code, #"ESRI:54030",
                 lat_dim = 'latitude', lon_dim = 'longitude',
                 legend = 's2dv', style = 'point',
                 toptitle = toptitle, bar_extra_margin = c(4,0,4,0),
                 caption = caption, mask = sig,
                 fileout = plotfile, width = 8, height = 6,
                 brks = seq(-1, 1, 0.2), cols = brewer.pal(10, 'PuOr'))

 
    }
  }  
  return(nino)
}