R/add_ncdf_output.R

In aemon-j/LakeEnsemblR: Run Ensemble of Lake Models

#' Add output from model-specific model runs to  netcdf file
#'
#' Add data from lists of output, generated by run_ensemble to  a
#' netcdf file created by earlier
#' run of run_ensemble
#'
#' @param output_lists list; list containing lists of 
#' output (e.g. temperature, ice_height)
#' @param folder filepath; to folder which contains the model folders
#'  generated by export_config()
#' @param model vector; model to export driving data. Options 
#' include c("GOTM", "GLM", "Simstrat",
#' "FLake", "MyLake")
#' @param out_file filepath; to save netCDF file defaults to "ensemble_output.nc"
#' @import ncdf4
#'
#' @keywords internal

add_netcdf_output <- function(output_lists, folder = ".", model, out_file) {
  
 
  # create connection
  nc <- ncdf4::nc_open(file.path(folder, "output", out_file), write = TRUE)
  
  # make sure the connection is closed on exit
  on.exit({
   ncdf4::nc_close(nc)
  })

  # # check for models
  # Extract model names
  mod_names <- ncatt_get(nc, "model", "Model")$value
  mod_names <- strsplit(mod_names, ", ")[[1]]
  mod_names <- substring(mod_names, 5)
  mod_names <- mod_names[!mod_names %in% "Obs"] 
  
  # sort models so they match the attribute number
  #model <- c(model, "Obs")
  model <- model[match(model, mod_names)]
  
  # ncdf4::ncvar_get(nc, "model")
  deps <- abs(ncvar_get(nc, "z"))
  
  # Extract the time
  tim <- ncvar_get(nc, "time")
  tunits <- ncatt_get(nc, "time")
  tustr <- strsplit(tunits$units, " ")
  # step <- tustr[[1]][1]
  tdstr <- strsplit(unlist(tustr)[3], "-")
  tmonth <- as.integer(unlist(tdstr)[2])
  tday <- as.integer(unlist(tdstr)[3])
  tyear <- as.integer(unlist(tdstr)[1])
  tdstr <- strsplit(unlist(tustr)[4], ":")
  thour <- as.integer(unlist(tdstr)[1])
  tmin <- as.integer(unlist(tdstr)[2])
  origin <- as.POSIXct(paste0(tyear, "-", tmonth,
                              "-", tday, " ", thour, ":", tmin),
                       format = "%Y-%m-%d %H:%M", tz = "UTC")
  time <- as.POSIXct(tim, origin = origin, tz = "UTC")
  datetime <- data.frame(datetime = time)
  
  # check for variables
  vars_old <- names(nc$var)
  
  # Load in vars and check for empty members
  var <- ncvar_get(nc, vars_old[1])
  n_vals <- (nc$dim$lon$len * nc$dim$lat$len *
               nc$dim$model$len * nc$dim$time$len * nc$dim$z$len)
  # Check for empty member
  for(m in seq_len(nc$dim$member$len)) {
    nas <- sum(is.na(var[m, , , ]))
    if(nas == n_vals) {
      break
    }
  }
  if(m == nc$dim$member$len) {
    stop("All members are full!\nChange output:
         file: \nor write function to add new members in netCDF :P")
  }

  #set new member number to old + 1
  mem_num_new <- m
  
  # Extract missing value
  # miss_val <- lapply(seq_len(length(vars_old)),
                     # function(x)ncatt_get(nc, vars_old[x])$missing_value)
 
  
  # Loop through and add each variable
  #  different for 2D or 3D variables
  for(i in seq_len(length(output_lists))) {
    
    if(ncol(output_lists[[i]][[1]]) == 2) {
      # Add 2D variable
      for (md in seq_len(length(model))) {
        dat_add <- as.matrix(output_lists[[i]][[md]][, -1])
        ncdf4::ncvar_put(nc, vars_old[i], dat_add, start = c(1, 1, mem_num_new, md, 1),
                         count = c(1, 1, 1, 1, length(dat_add)))
      }

    } else if(ncol(output_lists[[i]][[1]]) > 2) {
      
      arr <- array(NA, dim = c(1, 1, 1, (nc$dim$model$len), (nc$dim$time$len), 
                               (nc$dim$z$len)))
      # Add 3D variable
      for (m in seq_len(length(model))) {
        
        # mat1 <- as.matrix(output_lists[[i]][[m]][, !colnames(output_lists[[i]][[m]])
        # %in% "datetime"])
        df <- merge(output_lists[[i]][[m]], datetime, by = 1, all.y = TRUE)
        
        mat1 <- matrix(NA, nrow = nc$dim$time$len, ncol = nc$dim$z$len)
        
        # vector of depths to input into the matrix
        deps_tmp <- get.offsets(df)
        
        mat <- as.matrix(df[, -1])
        
        for(k in seq_len(ncol(mat))) {
          col <- which(deps == deps_tmp[k])
          mat1[, col] <- mat[, k]
        }
        
        splitted_name <- strsplit(names(output_lists[[i]])[m], "_")[[1]]
        m_name <- splitted_name[1]
        idx <- which(mod_names == m_name)
        
        arr[1, 1, 1, idx, , ] <- mat1
        
      }
      # nc <- ncdf4::nc_open(file.path(folder, "output", out_file), write=TRUE)
      
      ncdf4::ncvar_put(nc = nc, vars_old[i], vals = arr,
                       start = c(1, 1, mem_num_new, 1, 1, 1),
                       count = c(dim(arr)[1], dim(arr)[2],
                                 dim(arr)[3], dim(arr)[4], dim(arr)[5], dim(arr)[6]))
      # nc_close(nc)
    }
  }

  
  message("Finished adding data to NetCDF file [", Sys.time(), "]")
  
}