R/partitionSites.R

In rLakeHabitat: Interpolate Bathymetry and Quantify Physical Aquatic Habitat

#' Random Stratified Site Selection
#'
#' Randomly assign locations across a waterbody using Halton Iterative Partitioning (HIP; Robertson et al. 2018, via the spbal package), stratified by depth. Depth bins are defined from an interpolated DEM (e.g. the output of interpBathy()), and a customizable number of spatially determined locations is drawn within each bin - so, for example, you can request more locations in a shallow littoral bin and fewer in a deep pelagic bin.
#'
#' @param dem a SpatRaster of interpolated bathymetry (e.g. from interpBathy()), used both to define depth bins and as the pool of candidate locations. Assumed to represent the waterbody at full pool/elevation unless 'water_level_drop' is used.
#' @param depth_bins numeric vector of depth bin edges, e.g. c(0, 5, 15, 30, Inf) for three bins: 0-5 m, 5-15 m, and 15-30+ m. Must have at least 2 values (i.e. at least 1 bin). Bins are half-open, [lower, upper).
#' @param n_per_bin numeric vector giving how many locations to draw from each bin, in the same order as 'depth_bins' implies (length must be length(depth_bins) - 1). A single value is also accepted and recycled across all bins (e.g. n_per_bin = 5 draws 5 locations from every bin).
#' @param water_level_drop optional single non-negative numeric value giving how far the water level has dropped, in the same depth units as 'dem', relative to the full-pool elevation 'dem' is assumed to represent. If greater than 0, the DEM is rebuilt before sampling: every cell's depth is reduced by this amount, and any cell whose adjusted depth is at or below zero (i.e., now exposed/dry) is excluded from sampling. Default = 0 (no water level adjustment; 'dem' is used as-is).
#' @param min_spacing optional numeric value giving the minimum allowed distance, in meters, between selected locations within a bin (passed to spbal::HIP()'s minRadius argument). Default = NULL (no minimum spacing enforced).
#' @param iterations numeric value giving the number of Halton partition levels used by spbal::HIP() (see spbal::HIP() for details), default = 7. If HIP fails for a given bin (e.g. too few candidate cells), the function automatically retries with fewer partition levels before giving up on that bin with a warning.
#' @param plot logical: should a map of the resulting locations, colored by depth bin and shown over the (adjusted) DEM, be drawn? Default = TRUE.
#' @param seed optional numeric value used to seed the random number generator, for reproducible site selection across runs. Default = NULL (not seeded).
#' @details
#' This function requires the 'spbal' and 'sf' packages to be installed (not hard dependencies of this package, since they are only needed for this function).
#' The function automatically detects whether 'dem' is in a geographic (decimal degree) or projected (meters) coordinate system; if geographic, sampling and 'min_spacing' are carried out in the waterbody's best-fit UTM zone (determined from 'dem's own extent) so that distances are measured in meters, and returned locations are projected back to the original CRS of 'dem'.
#' If a depth bin contains fewer candidate DEM cells than the number of locations requested for it, all available cells are used and a warning is issued.
#' No outline shapefile is required - a boundary for the plot (if requested) is derived automatically from the non-NA footprint of 'dem' (or the water-level-adjusted DEM, if 'water_level_drop' > 0).
#' @return a list with:
#' \describe{
#'   \item{locations}{a data frame of selected locations, with columns 'x', 'y' (in the original CRS of
#'   'dem'), 'depth_bin' (the bin label), and 'bin_index' (the bin's position in 'depth_bins'). Depth bins/labels
#'   reflect the water-level-adjusted DEM if 'water_level_drop' > 0.}
#'   \item{map}{a recorded base R plot (see grDevices::recordPlot()) of the (adjusted) DEM and selected locations
#'   by depth bin, or NULL if plot = FALSE.}
#' }
#' @author Tristan Blechinger, Department of Zoology & Physiology, University of Wyoming
#' @export
#' @rawNamespace import(terra, except = c(union,intersect, animate))
#' @examples
#' \donttest{
#' outline <- terra::vect(system.file("extdata", "example_outline.shp", package = 'rLakeHabitat'))
#' data <- read.csv(system.file("extdata", "example_depths.csv", package = 'rLakeHabitat'))
#' dem <- interpBathy(outline, data, "x", "y", "z", zeros = FALSE,
#' separation = 10, res = 10, nmax = 8, method = "IDW")
#' partitionSites(dem, depth_bins = c(0, 5, 15, 30, Inf),
#' n_per_bin = c(8, 6, 4, 2), seed = 123)
#'
#' # with a 2 m water level drop from full pool
#' partitionSites(dem, depth_bins = c(0, 5, 15, 30, Inf),
#' n_per_bin = c(8, 6, 4, 2), water_level_drop = 2, seed = 123)}

partitionSites <- function(dem, depth_bins, n_per_bin, water_level_drop = 0, min_spacing = NULL,
                           iterations = 7, plot = TRUE, seed = NULL){

  if(!requireNamespace("spbal", quietly = TRUE))
    stop("This function requires the 'spbal' package. Install it with install.packages('spbal').")
  if(!requireNamespace("sf", quietly = TRUE))
    stop("This function requires the 'sf' package. Install it with install.packages('sf').")

  if(!is.null(seed)){
    if(!is.numeric(seed))
      stop("seed must be numeric")
    set.seed(seed)
  }

  #checks
  if(!inherits(dem, "SpatRaster"))
    stop("dem must be a SpatRaster (e.g., the output of interpBathy())")
  if(!is.numeric(depth_bins) || length(depth_bins) < 2)
    stop("depth_bins must be a numeric vector of at least 2 bin edges")
  if(!is.numeric(n_per_bin) || !(length(n_per_bin) %in% c(1, length(depth_bins) - 1)))
    stop("n_per_bin must be a single value (recycled across all bins) or have length(depth_bins) - 1 elements")
  if(any(n_per_bin < 0))
    stop("n_per_bin values must be non-negative")
  if(length(n_per_bin) == 1){
    n_per_bin <- rep(n_per_bin, length(depth_bins) - 1)
  }
  if(!is.numeric(water_level_drop) || length(water_level_drop) != 1 || is.na(water_level_drop))
    stop("water_level_drop must be a single numeric value")
  if(water_level_drop < 0)
    stop("water_level_drop must be non-negative (dem is assumed to represent full pool)")

  # identify which layer holds depth (vs. e.g. an "error" layer from OK/UK) so
  # only depth values get shifted, and any other layers just get masked to match.
  # Uses positional indexing throughout (not name-based subsetting) since a
  # single-layer DEM (e.g. from IDW) may not have an explicit "depth" name.
  depth_idx <- if("depth" %in% names(dem)) which(names(dem) == "depth")[1] else 1L

  # rebuild the DEM for a lower water level, if requested: shift depth down by
  # the drop amount, and drop (NA) any cell that is now at or above the surface
  if(water_level_drop > 0){
    new_depth <- dem[[depth_idx]] - water_level_drop
    new_depth[new_depth <= 0] <- NA
    dry_mask <- is.na(new_depth)

    dem[[depth_idx]] <- new_depth
    if(terra::nlyr(dem) > 1){
      other_idx <- setdiff(seq_len(terra::nlyr(dem)), depth_idx)
      for(idx in other_idx){
        dem[[idx]][dry_mask] <- NA
      }
    }
  }

  original_crs <- terra::crs(dem)

  # Function to determine the best UTM CRS directly from the DEM's own extent
  # (no outline required)
  get_best_utm <- function(r) {
    centroid <- terra::centroids(terra::as.polygons(terra::ext(r), crs = terra::crs(r)))
    lon <- terra::crds(centroid)[1]
    lat <- terra::crds(centroid)[2]
    utm_zone <- base::floor((lon + 180) / 6) + 1
    hemisphere <- base::ifelse(lat >= 0, 32600, 32700)
    epsg_code <- hemisphere + utm_zone
    return(terra::crs(paste0("EPSG:", epsg_code)))
  }

  # determine a single working CRS (meters) for the whole analysis, so min_spacing is maintained as meters
  working_crs <- original_crs
  if(terra::is.lonlat(dem)){
    working_crs <- get_best_utm(dem)
  }

  # pull all valid candidate cell centers + depth values out of the DEM (original CRS)
  candidates <- as.data.frame(dem, xy = TRUE, na.rm = TRUE)
  names(candidates)[1:3] <- c("x", "y", "depth")

  # reproject candidate coordinates into working CRS, if needed
  if(!identical(working_crs, original_crs)){
    pts <- terra::vect(candidates, geom = c("x", "y"), crs = original_crs)
    pts <- terra::project(pts, working_crs)
    proj_coords <- terra::crds(pts)
    candidates$x <- proj_coords[, 1]
    candidates$y <- proj_coords[, 2]
  }

  # assign depth bins (half-open intervals [lower, upper))
  candidates$bin <- cut(candidates$depth, breaks = depth_bins, include.lowest = TRUE, right = FALSE)
  bin_labels <- levels(candidates$bin)

  #### helper: try HIP with progressively fewer partition levels if it errors
  run_hip_safely <- function(pop_sf, n, its){
    retry_vals <- its
    val <- its - 2
    while(val >= 2){
      retry_vals <- c(retry_vals, val)
      val <- val - 2
    }
    for(this_it in unique(retry_vals)){
      result <- tryCatch(
        spbal::HIP(population = pop_sf, n = n, iterations = this_it, minRadius = min_spacing, verbose = FALSE),
        error = function(e) NULL
      )
      if(!is.null(result)) return(result)
    }
    return(NULL)
  }

  all_samples <- list()

  pb <- utils::txtProgressBar(min = 0, max = length(bin_labels), style = 3)

  for(i in seq_along(bin_labels)){
    bin_lab <- bin_labels[i]
    bin_candidates <- candidates[!is.na(candidates$bin) & candidates$bin == bin_lab, ]
    n_target <- n_per_bin[i]

    if(n_target <= 0 || nrow(bin_candidates) == 0){
      utils::setTxtProgressBar(pb, i)
      next
    }
    if(nrow(bin_candidates) < n_target){
      warning("Depth bin '", bin_lab, "' has only ", nrow(bin_candidates),
              " candidate DEM cells, fewer than the requested ", n_target, " locations. Drawing ",
              nrow(bin_candidates), " instead.")
      n_target <- nrow(bin_candidates)
    }

    pop_sf <- sf::st_as_sf(bin_candidates[, c("x", "y")], coords = c("x", "y"), crs = as.character(working_crs))

    hip_result <- run_hip_safely(pop_sf, n_target, iterations)

    if(is.null(hip_result)){
      warning("HIP sampling failed for depth bin '", bin_lab, "'; skipping.")
      utils::setTxtProgressBar(pb, i)
      next
    }

    # if minRadius was requested, the correctly-spaced result lives in $minRadius, not $sample - must be truncated down to the first n_target rows
    sample_obj <- if(!is.null(min_spacing)) hip_result$minRadius else hip_result$sample

    if(is.null(sample_obj) || nrow(sample_obj) == 0){
      warning("HIP returned no locations for depth bin '", bin_lab, "'; skipping.")
      utils::setTxtProgressBar(pb, i)
      next
    }

    if(!is.null(min_spacing) && nrow(sample_obj) > n_target){
      sample_obj <- sample_obj[seq_len(n_target), ]
    }
    if(!is.null(min_spacing) && nrow(sample_obj) < n_target){
      warning("Depth bin '", bin_lab, "': only ", nrow(sample_obj), " of the requested ", n_target,
              " locations could be placed at least ", min_spacing, "m apart.")
    }

    sampled_xy <- sf::st_coordinates(sample_obj)
    all_samples[[i]] <- data.frame(x = sampled_xy[, 1], y = sampled_xy[, 2],
                                   depth_bin = bin_lab, bin_index = i)

    utils::setTxtProgressBar(pb, i)
  }
  close(pb)

  locations <- do.call(rbind, all_samples)
  if(is.null(locations) || nrow(locations) == 0)
    stop("No locations could be selected for any depth bin - check depth_bins/n_per_bin against the DEM's depth range.")

  # reproject selected locations back to original CRS
  if(!identical(working_crs, original_crs)){
    pts_out <- terra::vect(locations, geom = c("x", "y"), crs = working_crs)
    pts_out <- terra::project(pts_out, original_crs)
    proj_coords <- terra::crds(pts_out)
    locations$x <- proj_coords[, 1]
    locations$y <- proj_coords[, 2]
  }

  map_plot <- NULL

  if(plot){
    bin_colors <- grDevices::hcl.colors(length(bin_labels), "Zissou 1")
    point_colors <- bin_colors[locations$bin_index]

    # plot only the depth layer, not the whole (potentially multi-layer) dem,
    # so terra::plot() stays single-panel and add = TRUE / points() overlay correctly
    depth_layer <- dem[[depth_idx]]

    # derive a boundary purely for display from the DEM's own non-NA footprint,
    # since no outline is supplied to this function anymore
    derived_outline <- terra::as.polygons(!is.na(depth_layer), dissolve = TRUE)
    derived_outline <- derived_outline[terra::values(derived_outline)[[1]] == 1, ]
    plot_title <- if(water_level_drop > 0){
      paste0("Selected locations by depth bin (water level down ", water_level_drop, " units)")
    } else {
      "Selected locations by depth bin (full pool)"
    }

    terra::plot(depth_layer, col = grDevices::hcl.colors(100, "Blues"),
                main = plot_title)
    if(nrow(derived_outline) > 0){
      terra::plot(derived_outline, add = TRUE, border = "black", lwd = 1.5)
    }
    graphics::points(locations$x, locations$y, pch = 21, bg = point_colors, col = "black", cex = 1.2)
    graphics::legend("topright", legend = bin_labels, pt.bg = bin_colors, pch = 21, bty = "n", cex = 0.8)

    map_plot <- grDevices::recordPlot()
  }

  return(list(locations = locations, map = map_plot))
}