R/nsink_generate_static_maps.R

In USEPA/nsink: Flow path nitrogen removal estimation

#' Generate N-Sink Static Maps for a given HUC
#'
#' @param input_data A list of input datasets created with
#'                   \code{\link{nsink_prep_data}}.
#' @param removal The removal raster stack or removal list, generated by
#'                \code{\link{nsink_calc_removal}}
#' @param samp_dens A value, in the units of the input data, divided by total area of
#'             the input HUC.  It is used to determine the number of points,
#'             determined through a regular sample, to calculate removal.  For
#'             instance, a value of 90 would roughly equate to a point per every
#'             90 meters.
#' @param ncpu Number of CPUs to use for calculating flowpath removal for larger
#'             (i.e. greater than 50) number of flowpaths.  Default is the number
#'             of cores available minus one.
#' @param seed Random seed to ensure reproducibility of sample point creation
#'             for transport maps. Default set to 23.
#' @return This function returns a list of rasters: nitrogen removal efficiency,
#'         nitrogen loading index, nitrogen transport index, and the
#'         nitrogen delivery index.
#'
#' @export
#' @examples
#' \dontrun{
#' library(nsink)
#' niantic_huc <- nsink_get_huc_id("Niantic River")$huc_12
#' niantic_data <- nsink_get_data(niantic_huc, data_dir = "nsink_data")
#' aea <- 5072
#' niantic_nsink_data <- nsink_prep_data(niantic_huc, projection = aea,
#'                                       data_dir = "nsink_data")
#' removal <- nsink_calc_removal(niantic_nsink_data)
#' static_maps <- nsink_generate_static_maps(niantic_nsink_data, removal,samp_dens = 900)
#' }
nsink_generate_static_maps <- function(input_data, removal, samp_dens,
                                       ncpu = future::availableCores() - 1,
                                       seed = 23) {

  # Create static rasters
  message("Creating removal efficiency map...")
  removal_map <- removal$raster_method[["removal"]]
  message("Creating the loading index map...")
  n_load_idx <- nsink_generate_n_loading_index(input_data)
  message("Creating the transport and delivery index maps...")
  n_delivery_heat <- 100 - nsink_generate_n_removal_heatmap(input_data,
    removal, samp_dens,
    ncpu = ncpu, seed
  )

  n_load_idx <- terra::resample(n_load_idx, input_data$raster_template,
                                 method = "near")
  n_delivery_heat <- terra::resample(n_delivery_heat, input_data$raster_template,
                                      method = "near")
  n_delivery_index <- n_load_idx * n_delivery_heat

  static_maps <- lapply(list(removal_effic = removal_map, loading_idx = n_load_idx,
              transport_idx = n_delivery_heat,delivery_idx = n_delivery_index),
         function(x) signif(x, 3))
  static_maps
}

#' Generates the Nitrogen Loading Index
#'
#' This function reclassifies the NLCD data to an index of Nitrogen loading.
#' The index ranges from 0 to 1.
#'
#' @param input_data List of input data
#' @keywords internal
nsink_generate_n_loading_index <- function(input_data) {

  nlcd <- input_data$nlcd
  rcl_m <- matrix(cbind(
    n_load_idx_lookup$codes,
    n_load_idx_lookup$n_loading_index
  ), ncol = 2)
  terra::classify(nlcd, rcl_m)
}

#' Generate Nitrogen Removal Heatmap
#'
#' Generates the heatmap.  Only flowpaths that have one segment or more
#' completely contained within the HUC are included. Practically, this will only
#' remove flowpaths that do not intersect the stream network and extend beyond
#' the HUC.  This is rare.
#' @param input_data A list of input datasets created with
#'                   \code{\link{nsink_prep_data}}.
#' @param removal The removal raster stack or removal list, generated by
#'                \code{\link{nsink_calc_removal}}
#' @param samp_dens A value, in the units of the input data, divided by total area of
#'             the input HUC.  It is used to determine the number of points,
#'             determined through a regular sample, to calculate removal.  For
#'             instance, a value of 90 would roughly equate to a point per every
#'             90 meters.
#' @param ncpu Number of CPUs to use for calculating flowpath removal for larger
#'             (i.e. greater than 50) number of flowpaths.  Default is the number
#'             of cores available minus one.
#' @param seed Random seed to ensure reproducibility of sample point creation
#'             across runs.  Default set to 23.
#' @import future furrr
#' @importFrom sf st_area st_sample st_sf st_sfc st_crs
#' @importFrom utils txtProgressBar setTxtProgressBar capture.output
#' @importFrom methods as
#'
#' @keywords internal
nsink_generate_n_removal_heatmap <- function(input_data, removal, samp_dens,
                                             ncpu = future::availableCores() - 1,
                                             seed = 23) {


  initial_plan <- future::plan()
  set.seed(seed)


    num_pts <- as.numeric(round(sum(st_area(input_data$huc)) / (samp_dens * samp_dens)))
    num_pts <- sum(num_pts)
    if(num_pts <= 4){stop("Choose a smaller samp_dens.")}

    # While loop to make sure small number of points always pass something.
    sample_pts <- vector("numeric")
    cnt <- 1
    while(length(sample_pts) == 0 & cnt < 11){
      # on lower_west samples are outside of polygon...
      sample_pts <- st_sample(x = input_data$huc, size = num_pts, type = "regular")
      cnt <- cnt + 1
      seed <- seed + 1
      set.seed(seed)
    }

    if(length(sample_pts) == 0 & cnt == 11){stop("Choose a smaller samp_dens.")}

  fdr_check <- suppressWarnings(terra::extract(input_data$fdr,
                                               as(sample_pts, "SpatVector")))

  if(any(is.na(fdr_check))){
    sample_pts <- sample_pts[!is.na(fdr_check$fdr)]
    message("Note: NA values detected in flow direction grid.  Static maps still generated.")
  }

  # for fewer points, the interp sample is done serially
  # for more points, it is done in parallel
  message(paste0(" Running ", length(sample_pts), " sampled flowpaths..."))

  #if(length(sample_pts) < 50){
  if(TRUE){
    pb <- txtProgressBar(max = length(sample_pts), style = 3)
    xdf <- data.frame(fp_removal = vector("numeric", length(sample_pts)))

    suppressPackageStartupMessages({
    for(i in seq_along(st_geometry(sample_pts))){
      setTxtProgressBar(pb, i)
      #i<-4 hits issue - lots of missing stuf from fdr in up_west
      pt <- sample_pts[i]
      pt <- st_sf(st_sfc(pt, crs = st_crs(input_data$huc)))

      fp <- nsink_generate_flowpath(pt, input_data)
      if(is.null(fp$flowpath_ends)){
        xdf[i,] <- data.frame(fp_removal = NA)
      } else if(any(st_within(fp$flowpath_ends, input_data$huc, sparse = FALSE))){
        fp_summary <- nsink_summarize_flowpath(fp, removal)

        xdf[i,] <- data.frame(fp_removal = 100 - min(fp_summary$n_out))
      } else {
        xdf[i,] <- data.frame(fp_removal = NA)
      }
    }
    })
    close(pb)

    sample_pts_removal <- st_sf(sample_pts, data = xdf,
                                row.names = row.names(xdf))

  } else {
    fp_removal <- function(pt, input_data, removal) {
      #Make Terra

      input_data$fdr <- terra::unwrap(input_data$fdr)
      input_data$impervious <- terra::unwrap(input_data$impervious)
      input_data$nlcd <- terra::unwrap(input_data$nlcd)
      input_data$raster_template <- terra::unwrap(input_data$raster_template)
      removal$raster_method$removal <- terra::unwrap(removal$raster_method$removal)
      removal$raster_method$type <- terra::unwrap(removal$raster_method$type)
      pt <- st_sf(st_sfc(pt, crs = st_crs(input_data$huc)))
      fp <- nsink_generate_flowpath(pt, input_data)
      if(any(st_within(fp$flowpath_ends, input_data$huc, sparse = FALSE))){
        fp_summary <- nsink_summarize_flowpath(fp, removal)
        return(data.frame(fp_removal = 100 - min(fp_summary$n_out)))
      } else {
        return(data.frame(fp_removal = NA))
      }
    }

    # Prepping terra for parallel
    input_data$fdr <- terra::wrap(input_data$fdr)
    input_data$impervious <- terra::wrap(input_data$impervious)
    input_data$nlcd <- terra::wrap(input_data$nlcd)
    input_data$raster_template <- terra::wrap(input_data$raster_template)
    removal$raster_method$removal <- terra::wrap(removal$raster_method$removal)
    removal$raster_method$type <- terra::wrap(removal$raster_method$type)

    future::plan(future::multisession, workers = ncpu)

    sample_pts_removal <- st_sf(sample_pts,
      data = suppressPackageStartupMessages(furrr::future_map_dfr(
        sample_pts,
        function(x) {
          fp_removal(
            x, input_data,
            removal
          )

        }, .progress = TRUE, .options = furrr_options(seed = TRUE)))
    )

    future::plan(initial_plan)
    #future:::ClusterRegistry("stop")
  }

  sample_pts_removal <- dplyr::filter(sample_pts_removal, !is.na(fp_removal))
  message("\n Interpolating sampled flowpaths...")

  # Interp each poly separately
  st_agr(input_data$huc) <- "constant"
  huc_polygon <- st_cast(input_data$huc, "POLYGON")

  for(i in 1:nrow(huc_polygon)){

    num_pts <- round(as.numeric(units::set_units(st_area(huc_polygon[i,]), "m^2")) / (30 * 30))
    interp_points <- st_sample(huc_polygon[i,], as.numeric(num_pts),
                               type = "regular")

    raster_huc <- terra::rasterize(huc_polygon[i,], input_data$raster_template)

    sample_pts_removal_huc_poly <- sample_pts_removal[
      st_intersects(sample_pts_removal,huc_polygon[i,], sparse = FALSE),]
    if(length(sample_pts_removal_huc_poly$fp_removal)>0){
      interp_points <- st_transform(interp_points, st_crs(raster_huc))
      sample_pts_removal_huc_poly <- st_transform(sample_pts_removal_huc_poly,
                                                  st_crs(raster_huc))
      xy <- st_coordinates(sample_pts_removal_huc_poly)
      xy_remove <- mutate(data.frame(xy),
                          fp_removal = sample_pts_removal_huc_poly$fp_removal)
      xy_remove <- select(xy_remove, x = "X", y = "Y", fp_removal)
    }
    if(nrow(sample_pts_removal_huc_poly)>= 2){
      interpolated_pts <- gstat::gstat(formula = fp_removal ~ 1,
                                data = xy_remove,
                                locations = ~x+y,
                                nmin = 2, nmax = 10,
                                set = list(idp = 0.5))


      # From terra::interpolate examples
      #interpolate_gstat <- function(model, x, crs, ...) {
      #  v <- st_as_sf(x, coords=c("x", "y"), crs=crs)
      #  p <- predict(model, v, ...)
      #  as.data.frame(p)[,1:2]
      #}

      output <- assign(paste0("idw", i),
                    terra::mask(terra::interpolate(raster_huc, interpolated_pts,
                                                   crs = terra::crs(raster_huc),
                                                   ext = huc_polygon[i,],
                                                   debug.level = 0,
                                                   index = 1), huc_polygon[i,]))

      #Test interpIDW
      #output <- capture.output(assign(paste0("idw", i),
      #                                terra::mask(
      #                                  terra::interpolate(raster_huc,
      #                                  interpolated_pts,
      #                                  fun = interpolate_gstat,
      #                                  crs = terra::crs(raster_huc),
      #                                  ext = huc_polygon[i,], debug.level = 0,
      #                                  index = 1), huc_polygon[i,])))

    }
  }

  if(length(ls(pattern = "idw")) > 0){
    idw_list <- lapply(ls(pattern = "idw"), function(x) get(x))
  } else {
    stop("Too few sample points.  Decrease the samp_dens value.")
  }
  if(length(idw_list) > 1){
    idw_n_removal_heat_map <- do.call(terra::merge, idw_list)
  } else {
    idw_n_removal_heat_map <- idw_list[[1]]
  }

  idw_n_removal_heat_map_agg <- terra::aggregate(idw_n_removal_heat_map,
                                                  fun = max, fact = 3)
  idw_n_removal_heat_map_agg <- terra::mask(
    terra::crop(idw_n_removal_heat_map_agg, input_data$huc),input_data$huc)

  idw_n_removal_heat_map_agg

}