R/subset.R

In reconhub/incidence: Compute, Handle, Plot and Model Incidence of Dated Events

##' Subsetting 'incidence' objects
##'
##' Two functions can be used to subset incidence objects.  The function
##' `subset` permits to retain dates within a specified range and,
##' optionally, specific groups. The operator "[" can be used as for matrices,
##' using the syntax `x[i,j]` where 'i' is a subset of dates, and 'j' is a
##' subset of groups.
##'
##' @author Thibaut Jombart \email{thibautjombart@@gmail.com}
##'
##' @export
##'
##' @rdname subset
##'
##' @aliases "subset.incidence" "[.incidence"
##'
##' @seealso The [incidence()] function to generate the 'incidence'
##' objects.
##'
##' @param x An incidence object, generated by the function
##' [incidence()].
##'
##' @param from The starting date; data strictly before this date are discarded.
##'
##' @param to The ending date; data strictly after this date are discarded.
##'
##' @param groups (optional) The groups to retained, indicated as subsets of the
##' columns of x$counts.
##'
##' @param ... Further arguments passed to other methods (not used).
##'
##' @examples
##' ## example using simulated dataset
##' if(require(outbreaks)) { withAutoprint({
##'   onset <- ebola_sim$linelist$date_of_onset
##'
##'   ## weekly incidence
##'   inc <- incidence(onset, interval = 7)
##'   inc
##'   inc[1:10] # first 10 weeks
##'   plot(inc[1:10])
##'   inc[-c(11:15)] # remove weeks 11-15
##'   plot(inc[-c(11:15)])
##' })}
##'

subset.incidence <- function(x, ..., from = min(x$dates), to = max(x$dates),
                             groups = TRUE){

  ## We need to make sure the comparison with dates is going to work. As for the
  ## [ operator, 'from' and 'to' are assumed to be expressed in the same way as
  ## the x$dates.

  is_date      <- inherits(x$dates, "Date")
  numeric_from <- is.numeric(from)
  numeric_to   <- is.numeric(to)

  if (is_date && (numeric_from || numeric_to)) {
    the_intervals <- get_interval(x, integer = TRUE)
    if (length(the_intervals) == 1L) {
      the_intervals <- rep(the_intervals, length(x$dates))
    }
    the_intervals <- cumsum(c(0, the_intervals))
  }
  if (is_date && numeric_from) {
    if (from <= 0) {
      from <- 0L
    } else if (from >= length(the_intervals) - 1L) {
      from <- the_intervals[length(the_intervals) - 1L]
    } else {
      from <- the_intervals[from]
    }
    from <- min(x$dates) + from
  }

  if (is_date && numeric_to) {
    if (to <= 0) {
      to <- 0L
    } else if (to >= length(the_intervals) - 1L) {
      to <- the_intervals[length(the_intervals) - 1L]
    } else {
      to <- the_intervals[to]
    }
    to <- min(x$dates) + to
  }

  to.keep <- x$dates >= from & x$dates <= to

  if (sum(to.keep) < 1) {
    stop("No data retained.")
  }
  x[to.keep, groups]
}




##' @export
##' @rdname subset
##' @param i a subset of dates to retain
##' @param j a subset of groups to retain

"[.incidence" <- function(x, i, j){
    if (missing(i)) {
        i <- TRUE
    }

    if (missing(j)) {
        j <- TRUE
    }

    out <- x
    if (is.character(j) && !all(j %in% group_names(x))) {
      odd_names <- j[!j %in% group_names(x)]
      groups    <- if (length(odd_names) > 1) "groups do" else "group does"
      odd_names <- paste(j[!j %in% group_names(x)], collapse = "', '")
      msg       <- sprintf("The following %s not exist: '%s'", groups, odd_names)
      stop(msg)
    }
    out$counts <- out$counts[i, j, drop = FALSE]
    out$dates <- out$dates[i]
    if ("weeks" %in% names(x)) {
      out$weeks    <- out$weeks[i] 
      out$isoweeks <- out$isoweeks[i]
    }
    # Need to use 1L here to keep things type-stable:
    # double + integer = double
    # integer + integer = integer
    # Date + integer = Date
    out$timespan <-  diff(range(out$dates, na.rm = TRUE)) + 1L
    out$n <- sum(out$counts)
    out
}