R/count-tally-SpatVector.R

In tidyterra: 'tidyverse' Methods and 'ggplot2' Helpers for 'terra' Objects

#' Count the observations in each `SpatVector` group
#'
#' @description
#' `count()` lets you quickly count the unique values of one or more variables:
#' `df |> count(a, b)` is roughly equivalent to
#' `df |> group_by(a, b) |> summarise(n = n())`. `count()` is paired with
#' `tally()`, a lower-level helper that is equivalent to
#' `df |> summarise(n = n())`.  Supply `wt` to perform weighted counts,
#' switching the summary from `n = n()` to `n = sum(wt)`.
#'
#' `add_count()` is equivalent to `count()` but use [mutate()] instead of
#' [summarise()] so that it adds a new column with group-wise counts.
#'
#' @export
#' @rdname count.SpatVector
#' @name count.SpatVector
#'
#' @seealso [dplyr::count()], [dplyr::tally()]
#'
#' @family dplyr.groups
#' @family dplyr.methods
#'
#' @importFrom dplyr count
#'
#' @param x A `SpatVector`.
#' @param .drop `r lifecycle::badge("deprecated")` Argument not longer
#'   supported; empty groups are always removed (see [dplyr::count()],
#'   `.drop = TRUE` argument).
#' @inheritParams dplyr::count
#' @inheritParams summarise.SpatVector
#'
#' @return A `SpatVector` object with an additional attribute.
#'
#' @section \CRANpkg{terra} equivalent:
#'
#' [terra::aggregate()]
#'
#' @section Methods:
#'
#' Implementation of the **generic** [dplyr::count()] methods for `SpatVector`
#' objects.
#'
#' [tally()] will always return a disaggregated geometry while [count()] can
#' handle this. See also [summarise.SpatVector()].
#'
#' @examples
#' \donttest{
#'
#' library(terra)
#' f <- system.file("ex/lux.shp", package = "terra")
#' p <- vect(f)
#'
#' p |> count(NAME_1, sort = TRUE)
#'
#' p |> count(pop = ifelse(POP < 20000, "A", "B"))
#'
#' # tally() is a lower-level function that assumes you've done the grouping
#' p |> tally()
#'
#' p |>
#'   group_by(NAME_1) |>
#'   tally()
#'
#' # Dissolve geometries by default
#'
#' library(ggplot2)
#' p |>
#'   count(NAME_1) |>
#'   ggplot() +
#'   geom_spatvector(aes(fill = n))
#'
#' # Opt out
#' p |>
#'   count(NAME_1, .dissolve = FALSE, sort = TRUE) |>
#'   ggplot() +
#'   geom_spatvector(aes(fill = n))
#' }
count.SpatVector <- function(
  x,
  ...,
  wt = NULL,
  sort = FALSE,
  name = NULL,
  .drop = deprecated(),
  .dissolve = TRUE
) {
  if (lifecycle::is_present(.drop)) {
    lifecycle::deprecate_warn(
      when = "1.1.0",
      what = "tidyterra::count.SpatVector(.drop)",
      details = paste0(
        "Argument not longer supported; empty groups are always removed",
        "(see `dplyr::count()`, `.drop = TRUE` argument)."
      )
    )
  }

  # Maybe regroup
  if (!missing(...)) {
    out <- group_by(x, ..., .add = TRUE, .drop = TRUE)
  } else {
    out <- x
  }

  vend <- tally(out, sort = sort, name = name, wt = {{ wt }})

  # Prepare a template for groups
  template <- dplyr::count(
    as_tibble(x),
    ...,
    sort = sort,
    name = name,
    .drop = TRUE
  )

  # Dissolve if requested
  if (.dissolve) {
    keepdf <- as_tibble(vend)

    var_index <- make_safe_index("tterra_index", keepdf)
    vend[[var_index]] <- seq_len(nrow(keepdf))
    vend <- terra::aggregate(vend, by = var_index, dissolve = TRUE)
    vend <- cbind(vend[, 0], keepdf)
  }

  # Ensure groups
  vend <- ungroup(vend)

  # Re-group based on the template
  if (dplyr::is_grouped_df(template)) {
    gvars <- dplyr::group_vars(template)
    vend <- group_by(vend, across_all_of(gvars))
  }

  vend
}

#' @export
dplyr::count

#' @importFrom dplyr tally
#' @export
#' @name count.SpatVector
tally.SpatVector <- function(x, wt = NULL, sort = FALSE, name = NULL) {
  # Use terra method on ungrouped
  if (!is_grouped_spatvector(x)) {
    vargroup <- make_safe_index("tterra_index", x)
    x[[vargroup]] <- "UNIQUE"

    vend <- terra::aggregate(x, by = vargroup, dissolve = FALSE, count = TRUE)
    df <- tally(as_tibble(x), sort = sort, name = name, wt = {{ wt }})

    vend <- cbind(vend[, 0], df)
    return(vend)
  }

  # Get tibble and index of rows
  tblforindex <- as_tibble(x)
  # Get the df
  template <- dplyr::tally(
    tblforindex,
    sort = FALSE,
    name = name,
    wt = {{ wt }}
  )

  vargroup <- dplyr::group_vars(tblforindex)
  x <- x[, vargroup]
  vend <- terra::aggregate(x, by = vargroup, dissolve = FALSE, count = TRUE)
  # Keep and rename
  vend <- cbind(vend[, 0], template)

  newvar <- setdiff(names(template), vargroup)

  if (sort) {
    # Re-sort
    vend <- vend[order(template[[newvar]], decreasing = TRUE), ]
  }

  vend <- ungroup(vend)

  # Re-group based on the template
  if (dplyr::is_grouped_df(template)) {
    gvars <- dplyr::group_vars(template)
    vend <- group_by(vend, across_all_of(gvars))
  }

  vend
}

#' @export
dplyr::tally

#' @importFrom dplyr add_count
#' @export
#' @name count.SpatVector
add_count.SpatVector <- function(
  x,
  ...,
  wt = NULL,
  sort = FALSE,
  name = NULL,
  .drop = deprecated()
) {
  # Maybe regroup
  if (!missing(...)) {
    out <- group_by(x, ..., .add = TRUE, .drop = TRUE)
  } else {
    out <- x
  }
  x

  # Add count with dplyr method, no sorting yet
  tbl <- dplyr::add_count(
    as_tibble(out),
    sort = FALSE,
    name = name,
    wt = {{ wt }}
  )

  vend <- cbind(out[, 0], tbl)

  if (sort) {
    newvar <- rev(names(tbl))[1]

    vend <- vend[order(tbl[[newvar]], decreasing = TRUE), ]
  }

  # Prepare a template for groups
  template <- dplyr::add_count(
    as_tibble(x),
    ...,
    sort = sort,
    name = name
  )
  # Ensure groups
  vend <- ungroup(vend)

  # Re-group based on the template
  if (dplyr::is_grouped_df(template)) {
    gvars <- dplyr::group_vars(template)
    vend <- group_by(vend, across_all_of(gvars))
  }

  vend
}

#' @export
dplyr::add_count