compute_nao.R 23.4 KB
Newer Older

compute_nao <- function(data, recipe, obsproj, plot_ts, plot_sp,
                        alpha, logo = NULL) {
  ## TODO: if fcst object in data compute the nao too
  if (!is.null(data$fcst)) {
    warning("NAO computed only for hindcast data.")
  }
  # Check if subsetting the data for the region is required
  lons <- data$hcst$coords$longitude
  lats <- data$hcst$coords$latitude
  subset_needed <- FALSE
  nao_region <- c(lonmin = -80, lonmax = 40,
                  latmin = 20, latmax = 80)
  if (any(lons < 0)) {  #[-180, 180] 
    if (!(min(lons) > -90 & min(lons) < -70 & 
          max(lons) < 50 & max(lons) > 30)) {
      subset_needed <- TRUE
    }
  } else {  #[0, 360]
    if (any(lons >= 50 & lons <= 270)) {
      susbset_needed <- TRUE
    } else {
      lon_E <- lons[which(lons < 50)]
      lon_W <- lons[-which(lons < 50)]
      if (max(lon_E) < 30 | min(lon_W) > 290) {
        subset_needed <- TRUE
      }
    }
  }
  if (any(max(lats) > 80 & min(lats) < 20)) {
    subset_needed <- TRUE
  }
  if (subset_needed) {
    warning("The data is being subsetted for 20N-80N and 80W-40E region.")

      hcst1 <- ClimProjDiags::SelBox(data$hcst$data,
                                     lon = as.vector(lons),
                                     lat = as.vector(lats),
                                     region = nao_region,
                                     londim = "longitude",
                                     latdim = "latitude")
      obs1 <- ClimProjDiags::SelBox(data$obs$data,
                                    lon = as.vector(lons),
                                    lat = as.vector(lats),
                                    region = nao_region,
                                    londim = "longitude",
                                    latdim = "latitude")
      hcst <- s2dv_cube(data = hcst1$data, lat = hcst1$lat, lon = hcst1$lon,
                        Variable = c(data$hcst$Variable[1], level = 'surface'),
                        data$hcst$Va, Datasets = data$hcst$Datasets,
                        time_dims = c('syear', 'time'),
                        Dates = data$hcst$Dates)
      obs <- s2dv_cube(data = obs1$data, lat = obs1$lat, lon = obs1$lon,
                       Variable = c(data$obs$Variable[1], level = 'surface'),
                       Datasets = data$obs$Datasets, time_dims = c('syear', 'time'))
# TODO check and create data object for the next steps
     data <- list(hcst = hcst, obs = obs)
     lons <- data$hcst$coords$longitude
     lats <- data$hcst$coords$latitude
     obs1 <- hcst1 <- NULL
     gc()
  }
  nao <- NAO(exp = data$hcst$data, obs = data$obs$data,
             lat = data$hcst$coords$latitude,
             lon = data$hcst$coords$longitude,
             time_dim = 'syear',
             memb_dim = 'ensemble', space_dim = c('latitude', 'longitude'),
             ftime_dim = 'time', ftime_avg = NULL,
             obsproj = obsproj, ncores = recipe$Analysis$ncores)
  ## Standardisation:
  nao$exp <- Apply(list(nao$exp), 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
  nao$obs <- Apply(list(nao$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
  nao$exp <- InsertDim(nao$exp, posdim = 1, lendim = 1, name = 'region')
  nao$obs <- InsertDim(nao$obs, posdim = 1, lendim = 1, name = 'region')
  hcst_dates <- data$hcst$attrs$Dates
  hcst_dates <- drop(data$hcst$attrs$Dates)

  if (!("time" %in% names(dim(hcst_dates)))) {
     if (is.null(dim(hcst_dates))) {
       hcst_dates <- array(hcst_dates, c(syear = length(hcst_dates)))
     }
     hcst_dates <- InsertDim(hcst_dates, pos = 1, len = 1, name = 'time')
     hcst_dates <- as.POSIXct(hcst_dates, origin = '1970-01-01', tz = 'UTC')
  }
  nao <- list(hcst = s2dv_cube(
                       data = nao$exp, 
                       varName = "nao",
                       metadata = list(
                         region = list(name = "NAO region",
                                      lats_range = paste0(range(lats)),
                                      lons_range = paste0(range(lons))),
                         time = data$hcst$attrs$Variable$metadata$time,
                       nao = list(units = 'adim', 
                                  longname = 'North Atlantic Oscillation')),
                       Dates = hcst_dates,
                       Dataset = recipe$Analysis$Datasets$System$name),
               obs = s2dv_cube(
                       data = nao$obs, 
                       varName = "nao",
                       metadata = list(
                         region = list(name = "NAO region",
                                       lats_range = paste0(range(lats)),
                                       lons_range = paste0(range(lons))),
                         time = data$obs$attrs$Variable$metadata$time,
                         nao = list(units = 'adim', 
                                    longname = 'North Atlantic Oscillation')),
                         Dates = data$obs$attrs$Dates,
                         Dataset = recipe$Analysis$Datasets$Reference$name))
  if (recipe$Analysis$Workflow$Indices$NAO$save == 'all') {
    file_dest <- paste0(recipe$Run$output_dir, "/outputs/Indices/")
    if (tolower(recipe$Analysis$Horizon) == "seasonal") {
      # Use startdates param from SaveExp to correctly name the files:
      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)
    }
    # need to recover original dimensions after saving to make Skill module work
    nao_original_dims_hcst <- nao$hcst$data
    nao$hcst$data <- .drop_indices_dims(nao$hcst$data)
    CST_SaveExp(data = nao$hcst,
          destination = file_dest, 
          startdates = as.vector(file_dates),
          dat_dim = NULL, sdate_dim = 'syear', 
          ftime_dim = 'time', var_dim = NULL, 
          memb_dim = 'ensemble')
    nao_original_dims_obs <- nao$obs$data
    nao$obs$data <- .drop_indices_dims(nao$obs$data)
    CST_SaveExp(data = nao$obs, #res,
          destination = file_dest,
          startdates = as.vector(file_dates),
          dat_dim = NULL, sdate_dim = 'syear',
          ftime_dim = 'time', var_dim = NULL,
          memb_dim = NULL)
    nao$hcst$data <- nao_original_dims_hcst
    nao$obs$data <- nao_original_dims_obs
    nao_original_dims_hcst <- nao_original_dims_obs <- 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(nao$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(nao$obs$data, along = 'time', ind = tstep)
      exp <- Subset(nao$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("NAO Index\n",  
                          month.abb[mes],
                          "/", recipe$Analysis$Time$hcst_start, "-",
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/NAO_",
                           system, "_", recipe$Analysis$Datasets$Reference$name,
                           "_s", recipe$Analysis$Time$sdate, "_ftime",
                           sprintf("%02d", tstep - 1 + recipe$Analysis$Time$ftime_min), ".png") 
        caption <- paste0("NAO method: ",
                          ifelse(recipe$Analysis$Workflow$Indices$NAO$obsproj,
                                 "Pobs", "Pmod"), " (Doblas-Reyes et al., 2003)\n",
                          "Nominal start date: 1st of ",
                          month.name[as.numeric(substr(recipe$Analysis$Time$sdate, 1,2))],
                          "\n",
                          "Forecast month: ", fmonth, "\n")
        xlabs <- as.numeric(substr(file_dates, 1, 4))
      } else if (tolower(recipe$Analysis$Horizon) == "decadal"){
        toptitle <- paste("NAO Index\n",  
                          "Lead time", fmonth,
                          " / Start dates", recipe$Analysis$Time$hcst_start, "-",
        plotfile <- paste0(recipe$Run$output_dir, "/plots/Indices/NAO_",
                           system, "_", recipe$Analysis$Datasets$Reference$name,
                           "_ftime",
                           sprintf("%02d", tstep - 1 + recipe$Analysis$Time$ftime_min), ".png") 
        caption <- paste0("NAO method: ",
                          ifelse(recipe$Analysis$Workflow$Indices$NAO$obsproj,
                                 "Pobs", "Pmod"), " (Doblas-Reyes et al., 2003)\n",
                          "Start date month: ",
                          month.name[get_archive(recipe)$System[[recipe$Analysis$Datasets$System$name]]$initial_month],
                          "\n",
                          "Lead time: ", fmonth, "\n")
        xlabs <- file_dates
      } 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,
                                  title = toptitle, fileout = plotfile,
                                  caption = caption, caption_line = 6.5,
                                  legend_text = c(
                                  recipe$Analysis$Datasets$Reference$name,
                                  recipe$Analysis$Datasets$System$name))
    }
  }
  if (plot_sp) {
    ## TODO: To be removed when s2dv is released:
    source("modules/Visualization/R/tmp/PlotRobinson.R")
    source("modules/Indices/R/correlation_eno.R")
    source("modules/Visualization/R/get_proj_code.R")
    dir.create(paste0(recipe$Run$output_dir, "/plots/Indices/"),
               showWarnings = F, recursive = T)
    # Get correct code for stereographic projection
    projection_code <- get_proj_code(proj_name = "stereographic")
    correl_obs <- Apply(list(data$obs$data, nao$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, nao$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, nao$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(nao$obs$data)['time']) {
      fmonth <- sprintf("%02d", tstep - 1 + recipe$Analysis$Time$ftime_min)
      ## Observations
      map <- drop(Subset(correl_obs$r, along = 'time', ind = tstep))
      sig <- drop(Subset(correl_obs$sign, along = 'time', ind = tstep))
      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)
        toptitle <- paste(recipe$Analysis$Datasets$Reference$name, "\n",
                          "NAO 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/NAO_correlation_",
                           recipe$Analysis$Variable$name, "_",
                           recipe$Analysis$Datasets$Reference$name,
                           "_s", recipe$Analysis$Time$sdate,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("NAO method: ",
                           ifelse(recipe$Analysis$Workflow$Indices$NAO$obsproj,
                                  "Pobs", "Pmod"), " (Doblas-Reyes et al., 2003)\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$Reference$name, "\n",
                          "NAO 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/NAO_correlation_",
                           recipe$Analysis$Variable$name, "_",
                           recipe$Analysis$Datasets$Reference$name,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("NAO method: ",
                           ifelse(recipe$Analysis$Workflow$Indices$NAO$obsproj,
                                  "Pobs", "Pmod"), " (Doblas-Reyes et al., 2003)\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.")
      }
      if (gsub(".", "", recipe$Analysis$Datasets$System$name) == "") {
        system <- recipe$Analysis$Datasets$System$name
      } else {
        system <- gsub(".", "", recipe$Analysis$Datasets$System$name)
      }

      PlotRobinson(map, lon = lons, lat = lats, target_proj = projection_code,
                   lat_dim = 'latitude', lon_dim = 'longitude',
                   legend = 's2dv', style = 'polygon',
                   toptitle = toptitle, crop_coastlines = nao_region,
                   caption = caption, mask = sig, bar_extra_margin = c(4,0,4,0),
                   fileout = plotfile, width = 8, height = 6,
                   brks = seq(-1, 1, 0.2), cols = brewer.pal(10, 'PuOr'))
      ## Ensemble-mean
      map <- drop(Subset(correl_hcst$r, along = 'time', ind = tstep))
      sig <- drop(Subset(correl_hcst$sign, along = 'time', ind = tstep))
      if (tolower(recipe$Analysis$Horizon) == "seasonal") {
        toptitle <- paste(recipe$Analysis$Datasets$System$name, "\n",
                          "NAO 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/NAO_correlation_",
                           recipe$Analysis$Variable$name, "_ensmean_",
                           system,
                           "_s", recipe$Analysis$Time$sdate,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("NAO method: ",
                           ifelse(recipe$Analysis$Workflow$Indices$NAO$obsproj,
                                  "Pobs", "Pmod"), " (Doblas-Reyes et al., 2003)\n",
                           "Correlation 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",
                          "NAO 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/NAO_correlation_",
                           recipe$Analysis$Variable$name, "_ensmean_",
                           system, 
                           recipe$Analysis$Datasets$Reference$name,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("NAO method: ",
                           ifelse(recipe$Analysis$Workflow$Indices$NAO$obsproj,
                                  "Pobs", "Pmod"), " (Doblas-Reyes et al., 2003)\n",
                           "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,
                   lat_dim = 'latitude', lon_dim = 'longitude',
                   legend = 's2dv', bar_extra_margin = c(4,0,4,0),
                   toptitle = toptitle, style = 'polygon',
                   caption = caption, mask = sig, crop_coastline = nao_region,
                   fileout = plotfile, width = 8, height = 6,
                   brks = seq(-1, 1, 0.2), cols = brewer.pal(10, 'PuOr'))
      
      # Full hcst corr
      map <- drop(Subset(correl_hcst_full$r, along = 'time', ind = tstep))
      sig <- drop(Subset(correl_hcst_full$sign, along = 'time', ind = tstep))
      if (tolower(recipe$Analysis$Horizon) == "seasonal") {
        toptitle <- paste(recipe$Analysis$Datasets$System$name,"\n",
                          "NAO 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/NAO_correlation_",
                           recipe$Analysis$Variable$name, "_member_",
                           system,
                           "_s", recipe$Analysis$Time$sdate,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("NAO method: ",
                           ifelse(recipe$Analysis$Workflow$Indices$NAO$obsproj,
                                  "Pobs", "Pmod"), " (Doblas-Reyes et al., 2003)\n",
                           "Correlation all 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",
                          "NAO 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/NAO_correlation_",
                           recipe$Analysis$Variable$name, "_member_",
                           system, 
                           recipe$Analysis$Datasets$Reference$name,
                           "_ftime", fmonth, ".png")
        caption <-  paste0("NAO method: ",
                           ifelse(recipe$Analysis$Workflow$Indices$NAO$obsproj,
                                  "Pobs", "Pmod"), " (Doblas-Reyes et al., 2003)\n",
                           "Correlation all members\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,
                   lat_dim = 'latitude', lon_dim = 'longitude',
                   legend = 's2dv', bar_extra_margin = c(4,0,4,0),
                   toptitle = toptitle, style = 'polygon',
                   caption = caption, mask = sig, crop_coastline = nao_region,
                   fileout = plotfile, width = 8, height = 6,
                   brks = seq(-1, 1, 0.2), cols = brewer.pal(10, 'PuOr'))
    } # end tstep loop
  }
  return(nao)
}